You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ju...@apache.org on 2023/11/22 07:55:37 UTC
(tvm) branch unity updated: [Unity] Support constant args in `nn.ExternModule` (#16130)
This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new aae1112a65 [Unity] Support constant args in `nn.ExternModule` (#16130)
aae1112a65 is described below
commit aae1112a6565816dd60e95f6cc4e3bd4fd9c965a
Author: Yaxing Cai <ca...@gmail.com>
AuthorDate: Tue Nov 21 23:55:29 2023 -0800
[Unity] Support constant args in `nn.ExternModule` (#16130)
This PR introduces `spec.ConstInt`, `spec.ConstFloat` and `spec.ConstString` for the `nn.ExternModule` args. So we will enable the constant input for extern module functions.
---
python/tvm/relax/frontend/nn/core.py | 27 ++++++++-
python/tvm/relax/frontend/nn/spec.py | 42 +++++++++++++-
.../python/relax/test_frontend_nn_extern_module.py | 64 +++++++++++++++++++++-
3 files changed, 127 insertions(+), 6 deletions(-)
diff --git a/python/tvm/relax/frontend/nn/core.py b/python/tvm/relax/frontend/nn/core.py
index ffda5c7901..e6c9a64f92 100644
--- a/python/tvm/relax/frontend/nn/core.py
+++ b/python/tvm/relax/frontend/nn/core.py
@@ -474,9 +474,15 @@ class ExternModule(Module):
from . import spec as _spec
from .op import _wrap_nested
- def extern_func(*args: Tensor) -> Tensor:
+ def extern_func(
+ *args: List[
+ Union[_spec.Tensor, _spec.ConstInt, _spec.ConstFloat, _spec.ConstString]
+ ]
+ ) -> Tensor:
spec2var = {}
for arg, arg_spec in zip(args, function_spec.args):
+ if not isinstance(arg_spec, _spec.Tensor):
+ continue
for value, value_spec in zip(arg.shape, arg_spec.shape):
if isinstance(value_spec, str):
if not value_spec in spec2var:
@@ -503,10 +509,27 @@ class ExternModule(Module):
out_shape, # type: ignore[arg-type]
func_spec_ret.dtype,
)
+ relax_args = []
+ for arg, arg_spec in zip(args, function_spec.args):
+ if isinstance(arg_spec, _spec.Tensor):
+ relax_args.append(arg._expr)
+ elif isinstance(arg_spec, _spec.ConstInt):
+ if arg_spec.dtype is None:
+ relax_args.append(rx.PrimValue(int(arg)))
+ else:
+ relax_args.append(rx.PrimValue(tir.IntImm(arg_spec.dtype, arg)))
+ elif isinstance(arg_spec, _spec.ConstFloat):
+ if arg_spec.dtype is None:
+ relax_args.append(rx.PrimValue(float(arg)))
+ else:
+ relax_args.append(rx.PrimValue(tir.FloatImm(arg_spec.dtype, arg)))
+ elif isinstance(arg_spec, _spec.ConstString):
+ relax_args.append(rx.StringImm(arg))
+
ret_tensor = _wrap_nested(
call_dps_packed(
func_name,
- args=RxTuple([tensor._expr for tensor in args]),
+ args=RxTuple(relax_args),
out_sinfo=out_sinfo,
),
func_name,
diff --git a/python/tvm/relax/frontend/nn/spec.py b/python/tvm/relax/frontend/nn/spec.py
index fb7f2588fa..07d21cff1b 100644
--- a/python/tvm/relax/frontend/nn/spec.py
+++ b/python/tvm/relax/frontend/nn/spec.py
@@ -78,6 +78,44 @@ class Tuple: # pylint: disable=too-few-public-methods
return self.elements.__repr__()
+class ConstInt: # pylint: disable=too-few-public-methods
+ """An integer constant"""
+
+ dtype: typing.Optional[str]
+
+ def __init__(self, dtype: str = None) -> None:
+ self.dtype = dtype
+
+ def __repr__(self) -> str:
+ if self.dtype is None:
+ return "const.int"
+ return f"const.int({self.dtype})"
+
+
+class ConstFloat: # pylint: disable=too-few-public-methods
+ """A float constant"""
+
+ dtype: typing.Optional[str]
+
+ def __init__(self, dtype: str = None) -> None:
+ self.dtype = dtype
+
+ def __repr__(self) -> str:
+ if self.dtype is None:
+ return "const.float"
+ return f"const.float({self.dtype})"
+
+
+class ConstString: # pylint: disable=too-few-public-methods
+ """A string constant"""
+
+ def __init__(self) -> None:
+ pass
+
+ def __repr__(self) -> str:
+ return "const.string"
+
+
class MethodSpec:
"""A spec for a compiled method"""
@@ -298,13 +336,13 @@ class ModuleSpec:
class ExternFunctionSpec: # pylint: disable=too-few-public-methods
"""A spec for a compiled external function."""
- args: typing.List[Tensor]
+ args: typing.List[typing.Union[Tensor, ConstInt, ConstFloat, ConstString]]
ret: typing.Union[Tensor, typing.List[Tensor]]
symbol: typing.Optional[str]
def __init__(
self,
- args: typing.List[Tensor],
+ args: typing.List[typing.Union[Tensor, ConstInt, ConstFloat, ConstString]],
ret: typing.Union[Tensor, typing.List[Tensor]],
symbol: typing.Optional[str] = None,
) -> None:
diff --git a/tests/python/relax/test_frontend_nn_extern_module.py b/tests/python/relax/test_frontend_nn_extern_module.py
index 0e5dc707b8..2f00284a91 100644
--- a/tests/python/relax/test_frontend_nn_extern_module.py
+++ b/tests/python/relax/test_frontend_nn_extern_module.py
@@ -191,6 +191,66 @@ def test_extern_module():
scalar_c = compiled["scalar_add"](scalar_a, scalar_b)
+def test_extern_spec():
+ class TestModule(nn.Module):
+ def __init__(self) -> None:
+ self.ext_mod = nn.ExternModule(
+ spec.ExternModuleSpec(
+ library=tvm.runtime.Module(None),
+ functions=[
+ spec.ExternFunctionSpec(
+ args=[
+ spec.Tensor((2, 4), "float16"),
+ spec.ConstInt(),
+ spec.ConstInt("int32"),
+ spec.ConstFloat(),
+ spec.ConstFloat("float16"),
+ spec.ConstString(),
+ ],
+ ret=spec.Tensor((2, 4), "float16"),
+ symbol="test",
+ )
+ ],
+ )
+ )
+
+ def forward(self, x: nn.Tensor):
+ return self.ext_mod.get_extern_func("test")(x, 1, 2, 3.0, 4.0, "123")
+
+ @I.ir_module
+ class ExpectedModule:
+ I.module_attrs({"external_mods": [None]})
+
+ @R.function
+ def forward(x: R.Tensor((2, 4), dtype="float16")) -> R.Tensor((2, 4), dtype="float16"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ test = R.call_dps_packed(
+ "test",
+ (
+ x,
+ R.prim_value(1),
+ R.prim_value(T.int32(2)),
+ R.prim_value(T.float32(3)),
+ R.prim_value(T.float16(4)),
+ R.str("123"),
+ ),
+ out_sinfo=R.Tensor((2, 4), dtype="float16"),
+ )
+ gv: R.Tensor((2, 4), dtype="float16") = test
+ R.output(gv)
+ return gv
+
+ model = TestModule()
+ ir_module, _ = model.export_tvm(
+ spec={
+ "forward": {
+ "x": spec.Tensor((2, 4), "float16"),
+ },
+ }
+ )
+ tvm.ir.assert_structural_equal(ir_module, ExpectedModule)
+
+
if __name__ == "__main__":
- test_extern_module()
- # tvm.testing.main()
+ tvm.testing.main()