chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,252 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=missing-docstring
|
||||
# ruff: noqa: F401
|
||||
import subprocess
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
import tvm.libinfo
|
||||
import tvm.testing
|
||||
from tvm import relax
|
||||
from tvm.relax.frontend import nn
|
||||
from tvm.relax.frontend.nn import spec
|
||||
from tvm.relax.transform import AttachExternModules
|
||||
|
||||
|
||||
def _infer_scalar_add(x, y): # pylint: disable=invalid-name
|
||||
assert isinstance(x, nn.Tensor)
|
||||
assert isinstance(y, nn.Tensor)
|
||||
assert x.ndim == 0 and x.dtype == "float32"
|
||||
assert y.ndim == 0 and y.dtype == "float32"
|
||||
return nn.Tensor.placeholder(shape=(), dtype="float32")
|
||||
|
||||
|
||||
def _infer_test_sym(a, b): # pylint: disable=invalid-name
|
||||
def _var_equal(a, b): # pylint: disable=invalid-name
|
||||
return tvm_ffi.structural_equal(a, b, map_free_vars=True)
|
||||
|
||||
assert isinstance(a, nn.Tensor)
|
||||
assert isinstance(b, nn.Tensor)
|
||||
assert a.ndim == 3 and a.dtype == "float32" # [x, y, 1]
|
||||
assert b.ndim == 3 and b.dtype == "float32" # [y, z, 5]
|
||||
x, y, z = a.shape[0], b.shape[0], b.shape[1] # pylint: disable=invalid-name
|
||||
assert _var_equal(a.shape[0], x)
|
||||
assert _var_equal(a.shape[1], y)
|
||||
assert a.shape[2] == 1
|
||||
assert _var_equal(b.shape[0], y)
|
||||
assert _var_equal(b.shape[1], z)
|
||||
assert b.shape[2] == 5
|
||||
return nn.Tensor.placeholder(shape=(x, y, z, 9), dtype="float32")
|
||||
|
||||
|
||||
def _test_scalar_add(func):
|
||||
# pylint: disable=invalid-name
|
||||
x = tvm.runtime.tensor(np.array(1.0).astype("float32"))
|
||||
y = tvm.runtime.tensor(np.array(3.0).astype("float32"))
|
||||
z = func(x, y).numpy()
|
||||
# pylint: enable=invalid-name
|
||||
assert z.ndim == 0
|
||||
assert z.dtype == "float32"
|
||||
assert float(z) == 4.0
|
||||
|
||||
|
||||
def _test_infer_sym(func, x, y, z): # pylint: disable=invalid-name
|
||||
# pylint: disable=invalid-name
|
||||
a = tvm.runtime.tensor(np.random.uniform(size=(x, y, 1)).astype("float32"))
|
||||
b = tvm.runtime.tensor(np.random.uniform(size=(y, z, 5)).astype("float32"))
|
||||
c = func(a, b).numpy()
|
||||
# pylint: enable=invalid-name
|
||||
assert c.shape == (x, y, z, 9)
|
||||
|
||||
|
||||
def _check_ir_equality(mod):
|
||||
# pylint: disable=import-outside-toplevel
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
|
||||
# pylint: enable=import-outside-toplevel
|
||||
|
||||
@I.ir_module
|
||||
class ExpectedModule:
|
||||
@R.function
|
||||
def scalar_add(
|
||||
a: R.Tensor((), dtype="float32"), b: R.Tensor((), dtype="float32")
|
||||
) -> R.Tensor((), dtype="float32"):
|
||||
R.func_attr({"num_input": 2})
|
||||
with R.dataflow():
|
||||
ext_scalar_add = R.call_dps_packed(
|
||||
"ext_scalar_add", (a, b), out_ty=R.Tensor((), dtype="float32")
|
||||
)
|
||||
gv: R.Tensor((), dtype="float32") = ext_scalar_add
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
@R.function
|
||||
def test_sym(
|
||||
a: R.Tensor(("x", "y", 1), dtype="float32"), b: R.Tensor(("y", "z", 5), dtype="float32")
|
||||
) -> R.Tensor(("x", "y", "z", 9), dtype="float32"):
|
||||
x = T.int64()
|
||||
y = T.int64()
|
||||
z = T.int64()
|
||||
R.func_attr({"num_input": 2})
|
||||
with R.dataflow():
|
||||
ext_test_sym = R.call_dps_packed(
|
||||
"ext_test_sym", (a, b), out_ty=R.Tensor((x, y, z, 9), dtype="float32")
|
||||
)
|
||||
gv1: R.Tensor((x, y, z, 9), dtype="float32") = ext_test_sym
|
||||
R.output(gv1)
|
||||
return gv1
|
||||
|
||||
tvm.ir.assert_structural_equal(ExpectedModule, mod)
|
||||
|
||||
|
||||
def _compile_cc(src: Path, dst: Path):
|
||||
cmd = ["g++", str(src)]
|
||||
default_include_paths = [
|
||||
tvm.libinfo.find_include_path(),
|
||||
tvm_ffi.libinfo.find_include_path(),
|
||||
tvm_ffi.libinfo.find_dlpack_include_path(),
|
||||
]
|
||||
for include_path in default_include_paths:
|
||||
cmd += ["-I", include_path]
|
||||
cmd += [
|
||||
"-c",
|
||||
"-std=c++17",
|
||||
"-fPIC",
|
||||
"-o",
|
||||
str(dst),
|
||||
]
|
||||
with subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) as proc:
|
||||
(out, _) = proc.communicate()
|
||||
if proc.returncode != 0:
|
||||
msg = "Compilation error:\n"
|
||||
msg += out.decode("utf-8", errors="replace")
|
||||
msg += "\nCommand line: " + " ".join(cmd)
|
||||
raise RuntimeError(msg)
|
||||
|
||||
|
||||
def test_extern_object():
|
||||
with tempfile.TemporaryDirectory() as temp_dir_str:
|
||||
path = Path(temp_dir_str) / "main.o"
|
||||
_compile_cc(
|
||||
src=Path(__file__).parent / "frontend_nn_extern_module.cc",
|
||||
dst=path,
|
||||
)
|
||||
|
||||
class TestModule(nn.Module):
|
||||
def __init__(self):
|
||||
self.ext_mod = None
|
||||
|
||||
def _get_ext_mod(self):
|
||||
if self.ext_mod is None:
|
||||
self.ext_mod = nn.ObjectModule(
|
||||
{
|
||||
"ext_scalar_add": _infer_scalar_add,
|
||||
"ext_test_sym": _infer_test_sym,
|
||||
},
|
||||
path,
|
||||
)
|
||||
nn.add_extern(self.ext_mod)
|
||||
return self.ext_mod
|
||||
|
||||
def scalar_add(self, a: nn.Tensor, b: nn.Tensor): # pylint: disable=invalid-name
|
||||
return self._get_ext_mod()["ext_scalar_add"](a, b)
|
||||
|
||||
def test_sym(self, a: nn.Tensor, b: nn.Tensor): # pylint: disable=invalid-name
|
||||
return self._get_ext_mod()["ext_test_sym"](a, b)
|
||||
|
||||
mod, _, ext_mods = TestModule().export_tvm(
|
||||
spec={
|
||||
"scalar_add": {
|
||||
"a": spec.Tensor((), "float32"),
|
||||
"b": spec.Tensor((), "float32"),
|
||||
},
|
||||
"test_sym": {
|
||||
"a": spec.Tensor(("x", "y", 1), "float32"),
|
||||
"b": spec.Tensor(("y", "z", 5), "float32"),
|
||||
},
|
||||
},
|
||||
allow_extern=True,
|
||||
)
|
||||
_check_ir_equality(mod)
|
||||
mod = AttachExternModules(ext_mods)(mod) # pylint: disable=not-callable
|
||||
compiled = tvm.runtime.vm.VirtualMachine(
|
||||
tvm.compile(mod, target="llvm"),
|
||||
device=tvm.cpu(),
|
||||
)
|
||||
_test_scalar_add(compiled["scalar_add"])
|
||||
_test_infer_sym(compiled["test_sym"], x=3, y=4, z=2)
|
||||
|
||||
|
||||
def test_extern_source():
|
||||
source = Path(__file__).parent / "frontend_nn_extern_module.cc"
|
||||
|
||||
class TestModule(nn.Module):
|
||||
def __init__(self):
|
||||
self.ext_mod = None
|
||||
|
||||
def _get_ext_mod(self):
|
||||
if self.ext_mod is None:
|
||||
self.ext_mod = nn.SourceModule(
|
||||
{
|
||||
"ext_scalar_add": _infer_scalar_add,
|
||||
"ext_test_sym": _infer_test_sym,
|
||||
},
|
||||
source_code=source,
|
||||
source_format="cpp",
|
||||
)
|
||||
nn.add_extern(self.ext_mod)
|
||||
return self.ext_mod
|
||||
|
||||
def scalar_add(self, a: nn.Tensor, b: nn.Tensor): # pylint: disable=invalid-name
|
||||
return self._get_ext_mod()["ext_scalar_add"](a, b)
|
||||
|
||||
def test_sym(self, a: nn.Tensor, b: nn.Tensor): # pylint: disable=invalid-name
|
||||
return self._get_ext_mod()["ext_test_sym"](a, b)
|
||||
|
||||
mod, _, ext_mods = TestModule().export_tvm(
|
||||
spec={
|
||||
"scalar_add": {
|
||||
"a": spec.Tensor((), "float32"),
|
||||
"b": spec.Tensor((), "float32"),
|
||||
},
|
||||
"test_sym": {
|
||||
"a": spec.Tensor(("x", "y", 1), "float32"),
|
||||
"b": spec.Tensor(("y", "z", 5), "float32"),
|
||||
},
|
||||
},
|
||||
allow_extern=True,
|
||||
)
|
||||
_check_ir_equality(mod)
|
||||
mod = AttachExternModules(ext_mods)(mod) # pylint: disable=not-callable
|
||||
compiled = tvm.runtime.vm.VirtualMachine(
|
||||
tvm.compile(mod, target="llvm"),
|
||||
device=tvm.cpu(),
|
||||
)
|
||||
_test_scalar_add(compiled["scalar_add"])
|
||||
_test_infer_sym(compiled["test_sym"], x=3, y=4, z=2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_extern_object()
|
||||
test_extern_source()
|
||||
Reference in New Issue
Block a user