[tvm] branch unity updated: [Unity] FuseOps skipping PrimValues (#14687)

[syfeng]

This is an automated email from the ASF dual-hosted git repository.

syfeng 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 4b59ee6e9d [Unity] FuseOps skipping PrimValues (#14687)
4b59ee6e9d is described below

commit 4b59ee6e9db8269c09096569b9c3ceb81a789ec5
Author: Ruihang Lai <ruiha...@cs.cmu.edu>
AuthorDate: Fri Apr 21 02:49:43 2023 -0400

    [Unity] FuseOps skipping PrimValues (#14687)
    
    Previously, FuseOps does not take PrimValue into account and will error
    on PrimValues for a sanity check inside. This PR supports FuseOps with
    PrimValues so that FuseOps will skip the PrimValues - no need to handle
    them in a special way.
    
    One regression test is added.
---
 src/relax/transform/fuse_ops.cc               |  4 +++-
 tests/python/relax/test_transform_fuse_ops.py | 24 ++++++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc
index c9c36bfcd8..f50a578954 100644
--- a/src/relax/transform/fuse_ops.cc
+++ b/src/relax/transform/fuse_ops.cc
@@ -253,7 +253,9 @@ class GraphCreator : public ExprVisitor {
     IndexedForwardGraph::Node* leaf_node = nullptr;
     if (it != graph_.node_map.end()) {
       leaf_node = it->second;
-    } else if (leaf_expr->IsInstance<ConstantNode>() || 
leaf_expr->IsInstance<ShapeExprNode>()) {
+    } else if (leaf_expr->IsInstance<ConstantNode>() || 
leaf_expr->IsInstance<ShapeExprNode>() ||
+               leaf_expr->IsInstance<PrimValueNode>() || 
leaf_expr->IsInstance<StringImmNode>() ||
+               leaf_expr->IsInstance<DataTypeImmNode>()) {
       leaf_node = CreateNode(leaf_expr.get());
       // Since we never fuse constants, the pattern of the constant is set to 
`kOpaque`.
       SetNodePattern(leaf_node, OpPatternKind::kOpaque);
diff --git a/tests/python/relax/test_transform_fuse_ops.py 
b/tests/python/relax/test_transform_fuse_ops.py
index 285a78a30e..1a4af26bd8 100644
--- a/tests/python/relax/test_transform_fuse_ops.py
+++ b/tests/python/relax/test_transform_fuse_ops.py
@@ -1387,5 +1387,29 @@ def test_shape_expr_arg():
     _check(Before, Expected)
 
 
+def test_skipping_primvalue():
+    @I.ir_module
+    class Module:
+        @R.function
+        def main(inp: R.Tensor((2, 2), dtype="float32")) -> R.Tensor((2, 2), 
dtype="float32"):
+            with R.dataflow():
+                lv = R.call_packed(
+                    "my_func1", inp, R.prim_value(0), sinfo_args=[R.Tensor((2, 
2), dtype="float32")]
+                )
+                lv1 = R.call_packed(
+                    "my_func2", lv, R.str("str"), sinfo_args=[R.Tensor((2, 2), 
dtype="float32")]
+                )
+                gv = R.call_packed(
+                    "my_func3",
+                    lv1,
+                    R.dtype("float32"),
+                    sinfo_args=[R.Tensor((2, 2), dtype="float32")],
+                )
+                R.output(gv)
+            return gv
+
+    _check(Module, Module)
+
+
 if __name__ == "__main__":
     tvm.testing.main()