# 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()