wehub-resource-sync
253 lines
8.6 KiB
Python
253 lines
8.6 KiB
Python
# 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()
|