chore: import upstream snapshot with attribution

tests/python/codegen/test_target_codegen.py

@@ -0,0 +1,166 @@
+# 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 numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+
+
+def test_buffer_store_predicate_not_supported():
+    target = "c"
+
+    @T.prim_func(s_tir=True)
+    def func(b: T.handle):
+        B = T.match_buffer(b, (8,), "float32")
+        B.vstore([T.Ramp(0, 2, 4)], T.Broadcast(1.0, 4), predicate=T.Broadcast(T.bool(True), 4))
+
+    err_msg = "Predicated buffer store is not supported."
+    with pytest.raises(RuntimeError, match=err_msg):
+        with tvm.target.Target(target):
+            tvm.compile(func)
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("opencl", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+        pytest.param("rocm", marks=pytest.mark.gpu),
+        pytest.param({"kind": "vulkan", "from_device": 0}, marks=pytest.mark.gpu),
+    ],
+)
+def test_buffer_store_predicate_not_supported_gpu(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (2, 3), "float32")
+        B = T.match_buffer(b, (6,), "float32")
+        T.func_attr({"global_symbol": "main"})
+        for i_0 in T.thread_binding(3, thread="threadIdx.x"):
+            B.vstore(
+                [T.Ramp(i_0, 1, 4)], T.Broadcast(1.0, 4), predicate=T.Broadcast(T.bool(True), 4)
+            )
+
+    err_msg = "Predicated buffer store is not supported."
+    with pytest.raises(RuntimeError, match=err_msg):
+        with tvm.target.Target(target):
+            tvm.compile(func)
+
+
+def test_buffer_load_predicate_not_supported():
+    target = "c"
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (8,), "float32")
+        B = T.match_buffer(b, (8,), "float32")
+        for i_0 in range(4):
+            B.vstore(
+                [T.Ramp(0, 2, 4)],
+                A.vload([T.Ramp(i_0, 1, 4)], predicate=T.Broadcast(T.bool(True), 4)),
+            )
+
+    err_msg = "Predicated buffer load is not supported."
+    with pytest.raises(RuntimeError, match=err_msg):
+        with tvm.target.Target(target):
+            tvm.compile(func)
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("opencl", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+        pytest.param("rocm", marks=pytest.mark.gpu),
+        pytest.param({"kind": "vulkan", "from_device": 0}, marks=pytest.mark.gpu),
+    ],
+)
+def test_buffer_load_predicate_not_supported_gpu(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (8,), "float32")
+        B = T.match_buffer(b, (8,), "float32")
+        for i_0 in T.thread_binding(3, thread="threadIdx.x"):
+            B.vstore(
+                [T.Ramp(0, 2, 4)],
+                A.vload([T.Ramp(i_0, 1, 4)], predicate=T.Broadcast(T.bool(True), 4)),
+            )
+
+    err_msg = "Predicated buffer load is not supported."
+    with pytest.raises(RuntimeError, match=err_msg):
+        with tvm.target.Target(target):
+            tvm.compile(func)
+
+
+@pytest.mark.parametrize("target", ["c", "llvm"])
+def test_codegen_loop_step(target):
+    if target != "c" and not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def test_loop_step(
+        A: T.Buffer((1024,), "float32"),
+        B: T.Buffer((1024,), "float32"),
+        C: T.Buffer((1024,), "float32"),
+    ):
+        for i in T.serial(3, 1024, step=96):
+            C[i] = A[i] + B[i]
+
+    with tvm.transform.PassContext(disabled_pass=["s_tir.CanonicalizeLoop"]):
+        lib = tvm.compile(test_loop_step, target=target)
+
+    src = lib.mod.inspect_source()
+    if target == "c":
+        assert src.find("for (int32_t i = 3; i < 1024; i += 96)") >= 0
+
+    dev = tvm.device(target, 0)
+    a_np = np.random.rand(1024).astype("float32")
+    b_np = np.random.rand(1024).astype("float32")
+    c_np = np.zeros(1024, dtype="float32")
+    a_tvm = tvm.runtime.tensor(a_np, dev)
+    b_tvm = tvm.runtime.tensor(b_np, dev)
+    c_tvm = tvm.runtime.tensor(c_np, dev)
+
+    lib(a_tvm, b_tvm, c_tvm)
+
+    c_result = c_tvm.numpy()
+
+    # Check that the loop executes at positions 3, 99, 195, 291, 387, 483, 579, 675, 771, 867, 963
+    for i in range(3, 1024, 96):
+        tvm.testing.assert_allclose(c_result[i], a_np[i] + b_np[i], rtol=1e-5)
+
+    # Assert non-touched positions remain zero
+    for i in range(0, 3):
+        assert c_result[i] == 0.0
+    for i in range(4, 1024):
+        if (i - 3) % 96 != 0:
+            assert c_result[i] == 0.0
+
+
+if __name__ == "__main__":
+    tvm.testing.main()