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,174 @@
+# 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.
+# ruff: noqa: F841
+
+import pytest
+from tvm_ffi.access_path import AccessPath
+
+import tvm
+from tvm.ir import assert_structural_equal
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def _error_message(exception):
+    return str(exception)
+
+
+def _expected_result(func1, func2, objpath1, objpath2):
+    return f"""StructuralEqual check failed, caused by lhs at {objpath1}:
+{func1.script(path_to_underline=[objpath1], syntax_sugar=False)}
+and rhs at {objpath2}:
+{func2.script(path_to_underline=[objpath2], syntax_sugar=False)}"""
+
+
+def test_prim_type_hidden_path_exact_message():
+    with pytest.raises(ValueError) as exc_info:
+        assert_structural_equal(tvm.ir.PrimType("int32"), tvm.ir.PrimType("float32"))
+
+    assert str(exc_info.value) == (
+        "StructuralEqual check failed, caused by lhs at <root>.dtype:\n"
+        "Access path: <root>.dtype\n"
+        "Note: The underlined object is the nearest visible parent of this path.\n\n"
+        "T.int32\n"
+        "^^^^^^^\n"
+        "and rhs at <root>.dtype:\n"
+        "Access path: <root>.dtype\n"
+        "Note: The underlined object is the nearest visible parent of this path.\n\n"
+        "T.float32\n"
+        "^^^^^^^^^"
+    )
+
+
+def test_prim_func_buffer_map():
+    @T.prim_func(s_tir=True)
+    def func1(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (128, 128))
+        B = T.match_buffer(b, (128, 128))
+
+    @T.prim_func(s_tir=True)
+    def func2(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (128, 128))
+        B = T.match_buffer(b, (128, 256))
+
+    func1 = func1.with_attr("global_symbol", "main")
+    func2 = func2.with_attr("global_symbol", "main")
+
+    with pytest.raises(ValueError) as ve:
+        assert_structural_equal(func1, func2)
+    assert _error_message(ve.value) == _expected_result(
+        func1,
+        func2,
+        AccessPath.root()
+        .attr("buffer_map")
+        .map_item(func1.params[1])
+        .attr("shape")
+        .array_item(1)
+        .attr("value"),
+        AccessPath.root()
+        .attr("buffer_map")
+        .map_item(func2.params[1])
+        .attr("shape")
+        .array_item(1)
+        .attr("value"),
+    )
+
+
+def test_evaluate():
+    @I.ir_module(s_tir=True)
+    class module1:
+        @T.prim_func(s_tir=True)
+        def func():
+            T.evaluate(0)
+
+    @I.ir_module(s_tir=True)
+    class module2:
+        @T.prim_func(s_tir=True)
+        def func():
+            T.evaluate(1)
+
+    with pytest.raises(ValueError) as ve:
+        assert_structural_equal(module1, module2)
+    assert _error_message(ve.value) == _expected_result(
+        module1,
+        module2,
+        AccessPath.root()
+        .attr("functions")
+        .map_item(module1.get_global_var("func"))
+        .attr("body")
+        .attr("value")
+        .attr("value"),
+        AccessPath.root()
+        .attr("functions")
+        .map_item(module2.get_global_var("func"))
+        .attr("body")
+        .attr("value")
+        .attr("value"),
+    )
+
+
+def test_allocate():
+    @T.prim_func(s_tir=True)
+    def func1():
+        a = T.alloc_buffer((128, 128), dtype="float32")
+
+    @T.prim_func(s_tir=True)
+    def func2():
+        a = T.alloc_buffer((256, 128), dtype="float32")
+
+    func1 = func1.with_attr("global_symbol", "main")
+    func2 = func2.with_attr("global_symbol", "main")
+
+    with pytest.raises(ValueError) as ve:
+        assert_structural_equal(func1, func2)
+
+    assert _error_message(ve.value) == _expected_result(
+        func1,
+        func2,
+        AccessPath.root().attr("body").attr("buffer").attr("shape").array_item(0).attr("value"),
+        AccessPath.root().attr("body").attr("buffer").attr("shape").array_item(0).attr("value"),
+    )
+
+
+def test_for():
+    @T.prim_func(s_tir=True)
+    def func1():
+        for i, j in T.grid(128, 128):
+            with T.sblock():
+                pass
+
+    @T.prim_func(s_tir=True)
+    def func2():
+        for i, j, k in T.grid(128, 128, 128):
+            with T.sblock():
+                pass
+
+    func1 = func1.with_attr("global_symbol", "main")
+    func2 = func2.with_attr("global_symbol", "main")
+
+    with pytest.raises(ValueError) as ve:
+        assert_structural_equal(func1, func2)
+    assert _error_message(ve.value) == _expected_result(
+        func1,
+        func2,
+        AccessPath.root().attr("body").attr("block").attr("body").attr("body").attr("body"),
+        AccessPath.root().attr("body").attr("block").attr("body").attr("body").attr("body"),
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()