chore: import upstream snapshot with attribution

tests/python/relax/test_vm_codegen_tir.py

@@ -0,0 +1,224 @@
+# 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.
+"""Test the TIR codegen path of VM compiled mode.
+
+Restrictions: all shape lowered, explicit allocation.
+"""
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.ir import assert_structural_equal
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+
+def get_tir_mod(mod):
+    builder = relax.ExecBuilder()
+    return relax.vm_build._vmcodegen(builder, mod, exec_mode="compiled")
+
+
+def test_add():
+    @tvm.script.ir_module
+    class Before:
+        @R.function(pure=False)
+        def foo(x: R.Tensor):
+            R.func_attr({"global_symbol": "foo"})
+            z = R.call_packed("test.vm.add", x, x, ty_args=(R.Tensor))
+            return z
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def __vmtir__foo(ctx_ptr: T.handle, r: T.handle, c: T.handle, f: T.handle):
+            T.func_attr({"global_symbol": "__vmtir__foo"})
+            T.anylist_setitem_call_packed(
+                r,
+                T.int32(2),
+                "test.vm.add",
+                T.anylist_getitem(r, T.int32(0)),
+                T.anylist_getitem(r, T.int32(0)),
+            )
+            T.anylist_setitem_call_packed(
+                r, T.int32(1), "vm.builtin.copy", T.anylist_getitem(r, T.int32(2))
+            )
+
+    before = Before
+    expected = Expected
+    after = get_tir_mod(before)
+    assert_structural_equal(expected, after)
+
+
+def test_tir_call():
+    @tvm.script.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def shape_func(H: T.Buffer(T.int64(4), "int64")):
+            T.func_attr({"global_symbol": "shape_func"})
+            # generated compute function
+            H[T.int64(0)] = H[T.int64(0)] + T.int64(1)
+
+        @R.function(pure=False)
+        def foo(x: R.Tensor([4], "int64")):
+            R.func_attr({"global_symbol": "foo"})
+            _ = Before.shape_func(x)
+            return x
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def shape_func(H: T.Buffer(T.int64(4), "int64")):
+            T.func_attr({"global_symbol": "shape_func"})
+            # generated compute function
+            H[T.int64(0)] = H[T.int64(0)] + T.int64(1)
+
+        @T.prim_func(s_tir=True)
+        def __vmtir__foo(ctx_ptr: T.handle, r: T.handle, c: T.handle, f: T.handle):
+            T.func_attr({"global_symbol": "__vmtir__foo"})
+            T.call_cpacked("shape_func", T.anylist_getitem(r, T.int32(0)))
+            T.anylist_setitem_call_packed(
+                r, T.int32(1), "vm.builtin.copy", T.anylist_getitem(r, T.int32(0))
+            )
+
+    before = Before
+    expected = Expected
+    after = get_tir_mod(before)
+    assert_structural_equal(expected, after)
+
+
+def test_if_cond():
+    @tvm.script.ir_module
+    class Before:
+        @R.function(pure=False)
+        def ife(cond: R.Tensor((), "bool"), x: R.Tensor) -> R.Tensor:
+            R.func_attr({"global_symbol": "ife"})
+            if cond:
+                w = R.call_packed("test.vm.add", x, x, ty_args=(R.Tensor))
+            else:
+                w = R.call_packed("test.vm.mul", x, x, ty_args=(R.Tensor))
+            return w
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def __vmtir__ife(ctx_ptr: T.handle, r: T.handle, c: T.handle, f: T.handle):
+            T.func_attr({"global_symbol": "__vmtir__ife"})
+            if T.Call(
+                tvm.ir.Op.get("tirx.tvm_call_packed"),
+                ["vm.builtin.read_if_cond", T.anylist_getitem(r, T.int32(0))],
+                ret_ty="bool",
+            ):
+                T.anylist_setitem_call_packed(
+                    r,
+                    T.int32(4),
+                    "test.vm.add",
+                    T.anylist_getitem(r, T.int32(1)),
+                    T.anylist_getitem(r, T.int32(1)),
+                )
+                T.anylist_setitem_call_packed(
+                    r, T.int32(3), "vm.builtin.copy", T.anylist_getitem(r, T.int32(4))
+                )
+            else:
+                T.anylist_setitem_call_packed(
+                    r,
+                    T.int32(5),
+                    "test.vm.mul",
+                    T.anylist_getitem(r, T.int32(1)),
+                    T.anylist_getitem(r, T.int32(1)),
+                )
+                T.anylist_setitem_call_packed(
+                    r, T.int32(3), "vm.builtin.copy", T.anylist_getitem(r, T.int32(5))
+                )
+            T.anylist_setitem_call_packed(
+                r, T.int32(2), "vm.builtin.copy", T.anylist_getitem(r, T.int32(3))
+            )
+
+    before = Before
+    expected = Expected
+    after = get_tir_mod(before)
+    assert_structural_equal(expected, after)
+
+
+def test_const():
+    @tvm.script.ir_module
+    class Before:
+        @R.function
+        def main(x: R.Tensor):
+            R.func_attr({"global_symbol": "main"})
+            y = R.const([1, 2])
+            z = (y, R.const([3, 4]), x)
+            return z
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def __vmtir__main(ctx_ptr: T.handle, r: T.handle, c: T.handle, f: T.handle):
+            # function attr dict
+            T.func_attr({"global_symbol": "__vmtir__main"})
+            # body
+            T.anylist_setitem_call_packed(
+                r,
+                T.int32(2),
+                "vm.builtin.make_tuple",
+                T.anylist_getitem(c, T.int32(0)),
+                T.anylist_getitem(c, T.int32(1)),
+                T.anylist_getitem(r, T.int32(0)),
+            )
+            T.anylist_setitem_call_packed(
+                r, T.int32(1), "vm.builtin.copy", T.anylist_getitem(r, T.int32(2))
+            )
+
+    before = Before
+    expected = Expected
+    after = get_tir_mod(before)
+    assert_structural_equal(expected, after)
+
+
+def test_const_call():
+    @tvm.script.ir_module
+    class Before:
+        @R.function(pure=False)
+        def main(x: R.Tensor):
+            R.func_attr({"global_symbol": "main"})
+            y = R.const([1, 2])
+            z = R.call_packed("test.vm.add", x, y, ty_args=(R.Tensor))
+            return z
+
+    @tvm.script.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def __vmtir__main(ctx_ptr: T.handle, r: T.handle, c: T.handle, f: T.handle):
+            # function attr dict
+            T.func_attr({"global_symbol": "__vmtir__main"})
+            # body
+            T.anylist_setitem_call_packed(
+                r,
+                2,
+                "test.vm.add",
+                T.anylist_getitem(r, 0),
+                T.anylist_getitem(c, 0),
+            )
+            T.anylist_setitem_call_packed(r, 1, "vm.builtin.copy", T.anylist_getitem(r, 2))
+
+    before = Before
+    expected = Expected
+    after = get_tir_mod(before)
+    assert_structural_equal(expected, after)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()