chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,133 @@
+# 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.
+# pylint: disable=invalid-name, missing-docstring
+
+import numpy as np
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+
+def test_transform_fuse_transpose_matmul():
+    @I.ir_module(s_tir=True)
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((128, 256), "float32"),
+            w: R.Tensor((128, 256), "float32"),
+        ) -> R.Tensor((128, 128), "float32"):
+            with R.dataflow():
+                wT = R.permute_dims(w, [1, 0])
+                o = R.matmul(x, wT)
+                R.output(o)
+            return o
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(private=True, s_tir=True)
+        def NT_matmul(
+            x: T.Buffer((T.int64(128), T.int64(256)), "float32"),
+            w: T.Buffer((T.int64(128), T.int64(256)), "float32"),
+            NT_matmul: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            # with T.sblock("root"):
+            for i0, i1, k in T.grid(T.int64(128), T.int64(128), T.int64(256)):
+                with T.sblock("NT_matmul"):
+                    v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
+                    T.reads(x[v_i0, v_k], w[v_i1, v_k])
+                    T.writes(NT_matmul[v_i0, v_i1])
+                    with T.init():
+                        NT_matmul[v_i0, v_i1] = T.float32(0)
+                    NT_matmul[v_i0, v_i1] = NT_matmul[v_i0, v_i1] + x[v_i0, v_k] * w[v_i1, v_k]
+
+        @R.function
+        def main(
+            x: R.Tensor((128, 256), dtype="float32"), w: R.Tensor((128, 256), dtype="float32")
+        ) -> R.Tensor((128, 128), dtype="float32"):
+            cls = Expected
+            with R.dataflow():
+                gv = R.call_tir(cls.NT_matmul, (x, w), out_ty=R.Tensor((128, 128), dtype="float32"))
+                R.output(gv)
+            return gv
+
+    after = tvm.ir.transform.Sequential(
+        [
+            relax.transform.FuseTransposeMatmul(),
+            relax.transform.FuseTIR(),  # Only used for remove unused primitive function
+        ]
+    )(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+def test_transform_fuse_transpose_matmul_const():
+    w = relax.const(np.random.uniform(-1e-3, 1e-3, (128, 256)), "float32")
+
+    @I.ir_module(s_tir=True)
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((128, 256), "float32"),
+        ) -> R.Tensor((128, 128), "float32"):
+            with R.dataflow():
+                wT = R.permute_dims(w, [1, 0])
+                o = R.matmul(x, wT)
+                R.output(o)
+            return o
+
+    @I.ir_module(s_tir=True)
+    class Expected:
+        @T.prim_func(private=True, s_tir=True)
+        def NT_matmul(
+            x: T.Buffer((T.int64(128), T.int64(256)), "float32"),
+            w: T.Buffer((T.int64(128), T.int64(256)), "float32"),
+            NT_matmul: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            # with T.sblock("root"):
+            for i0, i1, k in T.grid(T.int64(128), T.int64(128), T.int64(256)):
+                with T.sblock("NT_matmul"):
+                    v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
+                    T.reads(x[v_i0, v_k], w[v_i1, v_k])
+                    T.writes(NT_matmul[v_i0, v_i1])
+                    with T.init():
+                        NT_matmul[v_i0, v_i1] = T.float32(0)
+                    NT_matmul[v_i0, v_i1] = NT_matmul[v_i0, v_i1] + x[v_i0, v_k] * w[v_i1, v_k]
+
+        @R.function
+        def main(x: R.Tensor((128, 256), dtype="float32")) -> R.Tensor((128, 128), dtype="float32"):
+            cls = Expected
+            with R.dataflow():
+                gv = R.call_tir(cls.NT_matmul, (x, w), out_ty=R.Tensor((128, 128), dtype="float32"))
+                R.output(gv)
+            return gv
+
+    after = tvm.ir.transform.Sequential(
+        [
+            relax.transform.FuseTransposeMatmul(),
+            relax.transform.FuseTIR(),  # Only used for remove unused primitive function
+        ]
+    )(Before)
+    tvm.ir.assert_structural_equal(after, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()