chore: import upstream snapshot with attribution

tests/python/tirx-transform/test_tir_transform_helpers.py

@@ -0,0 +1,345 @@
+# 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
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_annotate_entry_func_single_primfunc():
+    @tvm.script.ir_module
+    class MockModule:
+        @T.prim_func(private=True, s_tir=True)
+        def func1(A: T.Buffer((16,), "float32")):
+            for i in T.serial(16):
+                if i == 5:
+                    if i == 5:
+                        A[i] = 0.0
+
+    mod = MockModule
+    assert mod
+    assert "tirx.is_entry_func" not in (mod["func1"].attrs or {})
+    after = tvm.tirx.transform.AnnotateEntryFunc()(mod)
+    assert (
+        after["func1"].attrs
+        and "tirx.is_entry_func" in after["func1"].attrs
+        and after["func1"].attrs["tirx.is_entry_func"]
+    )
+
+
+# Test module
+@tvm.script.ir_module
+class MockModule:
+    @T.prim_func(private=True, s_tir=True)
+    def func1(A: T.Buffer((16,), "float32")):
+        for i in T.serial(16):
+            if i == 5:
+                if i == 5:
+                    A[i] = 0.0
+
+    @T.prim_func(private=True, s_tir=True)
+    def func2(A: T.Buffer((32,), "float32")):
+        for i in T.serial(32):
+            if i == 15:
+                if i == 15:
+                    A[i] = 0.0
+
+
+@pytest.mark.xfail
+def test_annotate_entry_func_multiple_primfunc():
+    mod = MockModule
+    assert mod
+    assert "target" not in (mod["func1"].attrs or {})
+    assert "target" not in (mod["func2"].attrs or {})
+    # This should fail
+    after = tvm.tirx.transform.AnnotateEntryFunc()(mod)
+
+
+def test_bind_target():
+    mod = MockModule
+    assert mod
+
+    target = tvm.target.Target("cuda")
+    assert "target" not in (mod["func1"].attrs or {})
+    assert "target" not in (mod["func2"].attrs or {})
+    after = tvm.tirx.transform.BindTarget(target)(mod)
+
+    assert "target" in after["func1"].attrs
+    assert after["func1"].attrs["target"] == target
+    assert "target" in after["func2"].attrs
+    assert after["func2"].attrs["target"] == target
+
+
+def test_bind_target_adds_attribute():
+    """BindTarget adds the "target" attribute"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("cuda")})
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda"))(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_with_host_to_exposed_function():
+    """BindTarget adds the host target to externally-exposed functions"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"global_symbol": "main"})
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"global_symbol": "main", "target": T.target("cuda", host="llvm")})
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda", host="llvm"))(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_with_host_to_internal_function():
+    """Internal functions have a target annotation, but without the host
+
+    The host portion of the target annotation provides host
+    parameters, and is used to expose a function externally as part of
+    `MakePackedAPI` and `MakeUnpackedAPI`.  For internal functions, no
+    external exposure is required, so the host attribute should not be
+    used.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def main():
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(private=True, s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("cuda")})
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda", host="llvm"))(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_ignores_existing():
+    """BindTarget should not replace existing annotations"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"target": T.target("nvptx")})
+            T.evaluate(0)
+
+    Expected = Before
+
+    After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda"))(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_updates_host():
+    """BindTarget should update host for existing annotations"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr({"global_symbol": "func", "target": T.target("nvptx")})
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def main():
+            T.func_attr(
+                {
+                    "global_symbol": "func",
+                    "target": T.target("nvptx", host={"kind": "llvm", "opt-level": 0}),
+                }
+            )
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.BindTarget(
+        tvm.target.Target("cuda", host={"kind": "llvm", "opt-level": 0})
+    )(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_multiple_functions():
+    """BindTarget may apply to multiple functions in a module"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def func1():
+            T.evaluate(0)
+
+        @T.prim_func(s_tir=True)
+        def func2():
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(s_tir=True)
+        def func1():
+            T.func_attr({"target": T.target("cuda")})
+            T.evaluate(0)
+
+        @T.prim_func(s_tir=True)
+        def func2():
+            T.func_attr({"target": T.target("cuda")})
+            T.evaluate(0)
+
+    After = tvm.tirx.transform.BindTarget(tvm.target.Target("cuda"))(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_bind_target_with_device_host_call_same_func():
+    """BindTarget should bind the device target to the function if it is called from device"""
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(private=True, s_tir=True)
+        def add(a: T.int32, b: T.int32) -> T.int32:
+            return a + b
+
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((128, 128), "int32"),
+            B: T.Buffer((128, 128), "int32"),
+            C: T.Buffer((128, 128), "int32"),
+        ):
+            T.func_attr({"global_symbol": "main"})
+            length: T.let[T.int32] = Before.add(64, 64)  # Call from host
+            for bx in T.thread_binding(length, "blockIdx.x"):
+                for tx in T.thread_binding(length, "threadIdx.x"):
+                    C[bx, tx] = Before.add(A[bx, tx], B[bx, tx])  # Call from device
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func(private=True, s_tir=True)
+        def add(a: T.int32, b: T.int32) -> T.int32:
+            T.func_attr({"target": T.target("cuda")})
+            return a + b
+
+        @T.prim_func(private=True, s_tir=True)
+        def add_host(a: T.int32, b: T.int32) -> T.int32:
+            T.func_attr({"target": T.target({"kind": "llvm", "opt-level": 0})})
+            return a + b
+
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((128, 128), "int32"),
+            B: T.Buffer((128, 128), "int32"),
+            C: T.Buffer((128, 128), "int32"),
+        ):
+            T.func_attr(
+                {
+                    "global_symbol": "main",
+                    "target": T.target("cuda", host={"kind": "llvm", "opt-level": 0}),
+                }
+            )
+            length: T.let[T.int32] = Expected.add_host(64, 64)  # Call from host
+            for bx in T.thread_binding(length, "blockIdx.x"):
+                for tx in T.thread_binding(length, "threadIdx.x"):
+                    C[bx, tx] = Expected.add(A[bx, tx], B[bx, tx])  # Call from device
+
+    After = tvm.tirx.transform.BindTarget(
+        tvm.target.Target("cuda", host={"kind": "llvm", "opt-level": 0})
+    )(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_filter_primfunc():
+    mod = MockModule
+    assert mod
+    # Annotate each function for testing
+    mod["func1"] = mod["func1"].with_attr("temp", "test1")
+    mod["func2"] = mod["func2"].with_attr("temp", "test2")
+
+    # Test condition that does not filter out anything
+    def checker_filter_out_none(func: tvm.tirx.PrimFunc):
+        return "temp" in func.attrs
+
+    after = tvm.tirx.transform.Filter(checker_filter_out_none)(mod)
+    assert len(after.functions) == 2
+    # Filtered functions should satisfy the given condition.
+    assert checker_filter_out_none(after["func1"])
+    assert checker_filter_out_none(after["func2"])
+
+    # Test condition that selectively filters out primfuncs
+    def checker_filter_out_one(func: tvm.tirx.PrimFunc):
+        return ("temp" in func.attrs) and func.attrs["temp"] == "test1"
+
+    after = tvm.tirx.transform.Filter(checker_filter_out_one)(mod)
+    assert len(after.functions) == 1
+    # Filtered functions should satisfy the given condition.
+    assert checker_filter_out_one(after["func1"])
+
+    # Test condition that filters out everything
+    def checker_filter_out_both(func: tvm.tirx.PrimFunc):
+        return "invalid_attr" in func.attrs
+
+    after = tvm.tirx.transform.Filter(checker_filter_out_both)(mod)
+    assert len(after.functions) == 0
+
+
+def test_filter_removes_global_var_map():
+    """Filtering out a function should be identical to never adding it
+
+    This test is to guard against hidden state in the IRModule that
+    remains after filtering.  Previously, this was observed in the
+    `IRModuleNode::global_var_map_`, which retained entries of
+    filtered-out functions.
+    """
+
+    @I.ir_module
+    class Before:
+        @T.prim_func(s_tir=True)
+        def func():
+            T.evaluate(0)
+
+    @I.ir_module
+    class Expected:
+        pass
+
+    After = tvm.tirx.transform.Filter(lambda prim_func: False)(Before)
+    tvm.ir.assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()