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,80 @@
+# 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.
+import tvm
+from tvm.s_tir.meta_schedule.testing import te_workload
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,missing-class-docstring,missing-function-docstring
+# fmt: off
+
+
+@I.ir_module(s_tir=True)
+class Module:
+    @T.prim_func(s_tir=True)
+    def main(
+        A: T.Buffer((729, 729), "float32"),
+        B: T.Buffer((729, 729), "float32"),
+        C: T.Buffer((729, 729), "float32"),
+    ):
+        T.func_attr(
+            {
+                "global_symbol": "test",
+                "target": tvm.target.Target("llvm", host="llvm"),
+                "tirx.noalias": True,
+            }
+        )
+        # with T.sblock("root"):
+        for i, j, k in T.grid(729, 729, 729):
+            with T.sblock("C"):
+                v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
+                T.reads(A[v_i, v_k], B[v_k, v_j])
+                T.writes(C[v_i, v_j])
+                with T.init():
+                    C[v_i, v_j] = T.float32(0)
+                C[v_i, v_j] = C[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
+
+# fmt: on
+# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,missing-class-docstring,missing-function-docstring
+
+
+def test_host_func():
+    """Test that host functions are not split."""
+    # te schedule copied from test_tir_transform_split_host_device.py
+
+    func = tvm.te.create_prim_func(
+        te_workload.matmul(729, 729, 729, in_dtype="float32", out_dtype="float32")
+    )
+    mod = tvm.ir.IRModule({"main": func})
+    target = tvm.target.Target("cuda")
+    mod = tvm.tirx.transform.Apply(
+        lambda f: f.with_attr(
+            {
+                "global_symbol": "test",
+                "tirx.is_host_func": True,
+            }
+        )
+    )(mod)
+    mod = tvm.tirx.transform.BindTarget(target)(mod)
+    tvm.ir.assert_structural_equal(mod, Module)
+    assert "tirx.is_host_func" not in mod["main"].attrs, (
+        """Target and is_host_func attributes should be mutually exclusive"""
+    )
+
+
+if __name__ == "__main__":
+    test_host_func()