This is an automated email from the ASF dual-hosted git repository.
yaxingcai 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 70d23100f6 [Unity] Clear slots before each shape lowering (#15712)
70d23100f6 is described below
commit 70d23100f6afbda8c1d9bb5149da5fe6c2e1a7bf
Author: Wuwei Lin <[email protected]>
AuthorDate: Fri Sep 8 16:27:38 2023 -0700
[Unity] Clear slots before each shape lowering (#15712)
---
src/relax/backend/vm/vm_shape_lower.cc | 2 +
.../relax/test_backend_transform_shape_lower.py | 103 +++++++++++++++++++++
2 files changed, 105 insertions(+)
diff --git a/src/relax/backend/vm/vm_shape_lower.cc
b/src/relax/backend/vm/vm_shape_lower.cc
index 97e20a6b86..a5252be50b 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -222,6 +222,8 @@ class VMShapeLowerMutator
// Unit rewrite function per function.
Function Rewrite(GlobalVar gvar, Function func) {
// prepare mapping and heap var
+ slot_vec_.clear();
+ slot_map_.clear();
PrimExprSlotCollector::Collect(func, &slot_vec_, &slot_map_);
heap_size_ = IntImm(ShapeDType(), static_cast<int64_t>(slot_vec_.size()));
VarBinding shape_heap_binding = this->AllocShapeHeapBinding(heap_size_);
diff --git a/tests/python/relax/test_backend_transform_shape_lower.py
b/tests/python/relax/test_backend_transform_shape_lower.py
index 859df1c9ea..a5d4395e3c 100644
--- a/tests/python/relax/test_backend_transform_shape_lower.py
+++ b/tests/python/relax/test_backend_transform_shape_lower.py
@@ -22,6 +22,7 @@ from tvm.ir import assert_structural_equal
from tvm.relax.testing.runtime_builtin import MakeShapeCode, MatchShapeCode
from tvm.script import relax as R
from tvm.script import tir as T
+from tvm.script import ir as I
# note: we expected RemovePurityChecking to be run first, so we force purity
in most test cases
@@ -451,5 +452,107 @@ def test_return_match_check():
assert_structural_equal(after, expected)
+def test_symbolic_shape_multiple_function():
+ MS = MatchShapeCode
+ MK = MakeShapeCode
+
+ @I.ir_module
+ class Before:
+ @R.function
+ def fn1(A: R.Tensor(("m", "n"), dtype="float32")):
+ R.func_attr({"relax.force_pure": True})
+ m = T.int64()
+ n = T.int64()
+ return A
+
+ @R.function
+ def fn2(A: R.Tensor(("n", "m"), dtype="float32")):
+ R.func_attr({"relax.force_pure": True})
+ n = T.int64()
+ m = T.int64()
+ return A
+
+ # slot assignment:
+ sindex_fn1 = {
+ "m": 0,
+ "n": 1,
+ }
+ sindex_fn2 = {
+ "n": 0,
+ "m": 1,
+ }
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def fn1(A: R.Tensor(("m", "n"), dtype="float32")) -> R.Tensor(("m",
"n"), dtype="float32"):
+ R.func_attr({"relax.force_pure": True})
+ m = T.int64()
+ n = T.int64()
+ shape_heap: R.Tensor(dtype="int64", ndim=1) =
R.call_builtin_with_ctx(
+ "vm.builtin.alloc_shape_heap",
+ (R.prim_value(2),),
+ sinfo_args=(R.Tensor(dtype="int64", ndim=1),),
+ )
+ _: R.Tuple = R.call_packed(
+ "vm.builtin.check_tensor_info",
+ A,
+ R.prim_value(2),
+ R.dtype("float32"),
+ R.str(""),
+ sinfo_args=(R.Tuple,),
+ )
+ _1: R.Tuple = R.call_packed(
+ "vm.builtin.match_shape",
+ A,
+ shape_heap,
+ R.prim_value(2),
+ MS.STORE_TO_HEAP,
+ sindex_fn1["m"],
+ MS.STORE_TO_HEAP,
+ sindex_fn1["n"],
+ R.str(""),
+ sinfo_args=(R.Tuple,),
+ )
+ return A
+
+ @R.function
+ def fn2(A: R.Tensor(("n", "m"), dtype="float32")) -> R.Tensor(("n",
"m"), dtype="float32"):
+ R.func_attr({"relax.force_pure": True})
+ n = T.int64()
+ m = T.int64()
+ shape_heap: R.Tensor(dtype="int64", ndim=1) =
R.call_builtin_with_ctx(
+ "vm.builtin.alloc_shape_heap",
+ (R.prim_value(2),),
+ sinfo_args=(R.Tensor(dtype="int64", ndim=1),),
+ )
+ _2: R.Tuple = R.call_packed(
+ "vm.builtin.check_tensor_info",
+ A,
+ R.prim_value(2),
+ R.dtype("float32"),
+ R.str(""),
+ sinfo_args=(R.Tuple,),
+ )
+ _3: R.Tuple = R.call_packed(
+ "vm.builtin.match_shape",
+ A,
+ shape_heap,
+ R.prim_value(2),
+ MS.STORE_TO_HEAP,
+ sindex_fn2["n"],
+ MS.STORE_TO_HEAP,
+ sindex_fn2["m"],
+ R.str(""),
+ sinfo_args=(R.Tuple,),
+ )
+ return A
+
+ before = Before
+ expected = Expected
+ after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
+ assert_structural_equal(after, expected)
+
+
if __name__ == "__main__":
tvm.testing.main()
