This is an automated email from the ASF dual-hosted git repository.
sanirudh 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 71cdd460b4 [Unity][Frontend][NN] Op print_ (#15604)
71cdd460b4 is described below
commit 71cdd460b4c212cf75f55e746e997cedb35994bf
Author: Lesheng Jin <[email protected]>
AuthorDate: Thu Aug 24 03:36:18 2023 -0700
[Unity][Frontend][NN] Op print_ (#15604)
* print done
* fix
---
python/tvm/relax/frontend/nn/modules.py | 10 +++++++-
python/tvm/relax/frontend/nn/op.py | 5 ++++
tests/python/relax/test_frontend_nn_op.py | 38 +++++++++++++++++++++++++++++++
3 files changed, 52 insertions(+), 1 deletion(-)
diff --git a/python/tvm/relax/frontend/nn/modules.py
b/python/tvm/relax/frontend/nn/modules.py
index 9dfa36a990..cd94c23115 100644
--- a/python/tvm/relax/frontend/nn/modules.py
+++ b/python/tvm/relax/frontend/nn/modules.py
@@ -53,7 +53,15 @@ class IOEffect(Effect):
def print_(self, tensor: Tensor) -> None:
"""Encloses the side effect of NDArray printing"""
- raise NotImplementedError
+ self.effect = rx.BlockBuilder.current().emit(
+ rx.call_pure_packed(
+ rx.extern("effect.print"),
+ self.effect,
+ tensor._expr, # pylint: disable=protected-access
+ sinfo_args=[rx.ObjectStructInfo()],
+ ),
+ name_hint=self.effect.name_hint,
+ )
@register_func("effect.print")
diff --git a/python/tvm/relax/frontend/nn/op.py
b/python/tvm/relax/frontend/nn/op.py
index ed2b32b7ea..e5485d54c7 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -25,6 +25,7 @@ from ... import op as _op
from ...block_builder import BlockBuilder
from ...struct_info import TensorStructInfo, TupleStructInfo
from .core import Tensor
+from .spec import SpecBuilder
IntExpr = Union[int, _tir.PrimExpr]
@@ -938,3 +939,7 @@ def tensor_expr_op(
),
name=name_hint,
)
+
+
+def print_(array: Tensor):
+ SpecBuilder.current().io_effect.print_(array)
diff --git a/tests/python/relax/test_frontend_nn_op.py
b/tests/python/relax/test_frontend_nn_op.py
index 27d7e6d2ff..048a671101 100644
--- a/tests/python/relax/test_frontend_nn_op.py
+++ b/tests/python/relax/test_frontend_nn_op.py
@@ -15,6 +15,9 @@
# specific language governing permissions and limitations
# under the License.
import pytest
+import torch
+import sys
+import io
import tvm
import tvm.testing
@@ -304,5 +307,40 @@ def test_tensor_expr_op():
tvm.ir.assert_structural_equal(irmodule, Expected)
+def test_print():
+ class Model(Module):
+ def test(self, x: Tensor):
+ z = op.add(x, x)
+ op.print_(z)
+ return x
+
+ # fmt: off
+ @I.ir_module
+ class Expected:
+ @R.function
+ def _initialize_effect() -> R.Tuple(R.Object):
+ with R.dataflow():
+ _io: R.Object = R.null_value()
+ lv: R.Tuple(R.Object) = (_io,)
+ gv: R.Tuple(R.Object) = lv
+ R.output(gv)
+ return gv
+
+ @R.function
+ def test(x: R.Tensor((10, 10), dtype="float32"), _io: R.Object) ->
R.Tuple(R.Tensor((10, 10), dtype="float32"), R.Tuple(R.Object)):
+ with R.dataflow():
+ add: R.Tensor((10, 10), dtype="float32") = R.add(x, x)
+ _io1: R.Object = R.call_pure_packed("effect.print", _io, add,
sinfo_args=(R.Object(),))
+ gv1: R.Tuple(R.Tensor((10, 10), dtype="float32"),
R.Tuple(R.Object)) = x, (_io1,)
+ R.output(gv1)
+ return gv1
+ # fmt: on
+
+ m = Model()
+ irmodule, params = m.export_tvm(spec={"test": {"x": spec.Tensor([10, 10],
"float32")}})
+
+ tvm.ir.assert_structural_equal(irmodule["test"], Expected["test"])
+
+
if __name__ == "__main__":
tvm.testing.main()
