chore: import upstream snapshot with attribution

tests/python/codegen/test_codegen_assert.py

@@ -0,0 +1,124 @@
+# 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.
+"""AssertStmt codegen tests: verify kind and message_parts produce correct exceptions."""
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+
+codegen_target = tvm.testing.parameter("llvm", "c")
+
+
+def test_assert_runtime_error(codegen_target):
+    """AssertStmt with RuntimeError kind produces RuntimeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("RuntimeError", ["Expected non-null input"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(RuntimeError, match="Expected non-null input"):
+        lib(0)
+
+
+def test_assert_value_error(codegen_target):
+    """AssertStmt with ValueError kind produces ValueError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("ValueError", ["Shape mismatch: expected 4 got 8"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(ValueError, match="Shape mismatch"):
+        lib(0)
+
+
+def test_assert_type_error(codegen_target):
+    """AssertStmt with TypeError kind produces TypeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("TypeError", ["Expected Tensor but got int"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(TypeError, match="Expected Tensor but got int"):
+        lib(0)
+
+
+def test_assert_multi_part_message(codegen_target):
+    """Multi-part messages are correctly concatenated at runtime."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("ValueError", ["Expected shape ", "4", " but got ", "8"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(ValueError, match="Expected shape 4 but got 8"):
+        lib(0)
+
+
+def test_assert_passing_condition(codegen_target):
+    """Passing assertion does not raise."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("RuntimeError", ["This should not be raised"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    lib(1)  # should pass without error
+
+
+def test_assert_many_parts(codegen_target):
+    """Assertion with 8 parts concatenated correctly."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("RuntimeError", ["p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(RuntimeError, match="p0p1p2p3p4p5p6p7"):
+        lib(0)
+
+
+def test_tvmscript_assert_preserves_kind(codegen_target):
+    """Regression: TVMScript structured assert preserves kind at runtime."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("ValueError", ["x must be positive"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(ValueError, match="x must be positive"):
+        lib(0)
+
+
+def test_tvmscript_assert_preserves_parts(codegen_target):
+    """Regression: TVMScript structured assert with separate parts."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32):
+        assert x > 0, ("ValueError", ["x must be ", "positive"])
+
+    lib = tvm.compile(func, target=codegen_target)
+    with pytest.raises(ValueError, match="x must be positive"):
+        lib(0)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_codegen_error_handling.py

@@ -0,0 +1,471 @@
+# 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.
+"""Runtime error message tests for MakePackedAPI + TVMFFIABIBuilder.
+
+All tests compile TVMScript functions and verify the correct Python exception
+type and exact error message at runtime.
+"""
+
+import re
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+# Parameterize over both LLVM and C backends
+codegen_target = tvm.testing.parameter("llvm", "c")
+
+
+# ── Argument count errors ────────────────────────────────────
+
+
+def test_wrong_argument_count_error(codegen_target):
+    """Wrong argument count produces TypeError with function signature."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        n0 = T.int64()
+        A = T.match_buffer(a, (n0,), "float32")
+        B = T.match_buffer(b, (n0,), "float32")
+        for i in range(n0):
+            B[i] = A[i] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a, b)  # correct input should pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Expected 2 arguments when calling:\n"
+            "  `func(A: Tensor([n0], float32), B: Tensor([n0], float32))`"
+        ),
+    ):
+        lib()
+
+
+# ── Type mismatch errors (tensor parameters) ────────────────
+
+
+def test_type_mismatch_non_tensor(codegen_target):
+    """Passing a non-tensor where a tensor is expected raises TypeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        n0 = T.int64()
+        A = T.match_buffer(a, (n0,), "float32")
+        B = T.match_buffer(b, (n0,), "float32")
+        for i in range(n0):
+            B[i] = A[i] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a, b)  # correct input should pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched type on argument #1 when calling:\n"
+            "  `func(A: Tensor([n0], float32), B: Tensor([n0], float32))`,\n"
+            "  expected Tensor"
+        ),
+    ):
+        lib(a, 1)
+
+
+# ── Shape mismatch errors ───────────────────────────────────
+
+
+def test_shape_mismatch_shared_variable(codegen_target):
+    """b has different shape than a when they share symbolic variable n0."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        n0 = T.int64()
+        A = T.match_buffer(a, (n0,), "float32")
+        B = T.match_buffer(b, (n0,), "float32")
+        for i in range(n0):
+            B[i] = A[i] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a, b)  # correct input should pass
+
+    b_short = tvm.runtime.tensor(np.zeros(126, dtype="float32"))
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Mismatched B.shape[0] on argument #1 when calling:\n"
+            "  `func(A: Tensor([n0], float32), B: Tensor([n0], float32))`,\n"
+            "  expected to match A.shape[0]"
+        ),
+    ):
+        lib(a, b_short)
+
+
+def test_invalid_shape_fixed(codegen_target):
+    """Passing wrong shape for a fixed buffer dimension raises ValueError."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((128,), "float32"), b: T.Buffer((128,), "float32")):
+        for i in range(128):
+            b[i] = a[i] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a, b)  # correct input should pass
+
+    a_wrong = tvm.runtime.tensor(np.zeros(256, dtype="float32"))
+    b_wrong = tvm.runtime.tensor(np.zeros(256, dtype="float32"))
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Invalid a.shape[0] on argument #0 when calling:\n"
+            "  `func(a: Tensor([128], float32), b: Tensor([128], float32))`,\n"
+            "  expected 128"
+        ),
+    ):
+        lib(a_wrong, b_wrong)
+
+
+# ── ndim mismatch errors ────────────────────────────────────
+
+
+def test_ndim_mismatch_error(codegen_target):
+    """ndim mismatch produces ValueError with function signature."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((4, 8), "float32"), b: T.Buffer((4, 8), "float32")):
+        for i, j in T.grid(4, 8):
+            b[i, j] = a[i, j]
+
+    lib = tvm.compile(func, target=codegen_target)
+    a_ok = tvm.runtime.tensor(np.zeros((4, 8), dtype="float32"))
+    b_ok = tvm.runtime.tensor(np.zeros((4, 8), dtype="float32"))
+    lib(a_ok, b_ok)  # correct input should pass
+
+    a = tvm.runtime.tensor(np.zeros(4, dtype="float32"))
+    b = tvm.runtime.tensor(np.zeros(4, dtype="float32"))
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Mismatched a.ndim on argument #0 when calling:\n"
+            "  `func(a: Tensor([4, 8], float32), b: Tensor([4, 8], float32))`,\n"
+            "  expected 2"
+        ),
+    ):
+        lib(a, b)
+
+
+# ── dtype mismatch errors ───────────────────────────────────
+
+
+def test_dtype_mismatch_error(codegen_target):
+    """dtype mismatch produces TypeError with function signature."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((8,), "float32"), b: T.Buffer((8,), "float32")):
+        for i in range(8):
+            b[i] = a[i]
+
+    lib = tvm.compile(func, target=codegen_target)
+    a_ok = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    b_ok = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    lib(a_ok, b_ok)  # correct input should pass
+
+    a = tvm.runtime.tensor(np.zeros(8, dtype="int32"))
+    b = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched a.dtype on argument #0 when calling:\n"
+            "  `func(a: Tensor([8], float32), b: Tensor([8], float32))`,\n"
+            "  expected float32"
+        ),
+    ):
+        lib(a, b)
+
+
+# ── Data alignment errors ──────────────────────────────────
+
+
+@pytest.mark.skip(reason="alignment check disabled for now, revisit after merge")
+def test_data_alignment_error(codegen_target):
+    """Misaligned buffer data pointer raises ValueError."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((128,), "float32"), b: T.Buffer((128,), "float32")):
+        for i in range(128):
+            b[i] = a[i] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a_ok, b_ok)  # correct input should pass
+
+    # Slice off first element of a 129-element array to create misaligned data pointer
+    np_arr = np.zeros(129, dtype="float32")
+    a_misaligned = tvm_ffi.from_dlpack(np_arr[1:])
+    b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Misaligned Tensor data on argument #0 when calling:\n"
+            "  `func(a: Tensor([128], float32), b: Tensor([128], float32))`,\n"
+            "  expected data alignment=64 bytes"
+        ),
+    ):
+        lib(a_misaligned, b)
+
+
+# ── Compact strides mismatch errors ────────────────────────
+
+
+def test_strides_mismatch_transposed(codegen_target):
+    """Transposed (non-compact) strides raise ValueError."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((128, 128), "float32"), b: T.Buffer((128, 128), "float32")):
+        for i, j in T.grid(128, 128):
+            b[i, j] = a[i, j] + T.float32(1)
+
+    lib = tvm.compile(func, target=codegen_target)
+    a_ok = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+    b_ok = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+    lib(a_ok, b_ok)  # correct input should pass
+
+    # Use Fortran-order array to get non-compact (non-C-contiguous) strides
+    np_arr = np.asfortranarray(np.zeros((128, 128), dtype="float32"))
+    a_transposed = tvm_ffi.from_dlpack(np_arr)
+    b = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Mismatched a.strides on argument #0 when calling:\n"
+            "  `func(a: Tensor([128, 128], float32), b: Tensor([128, 128], float32))`,\n"
+            "  expected to be compact array"
+        ),
+    ):
+        lib(a_transposed, b)
+
+
+# ── Device mismatch errors ─────────────────────────────────
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_device_mismatch_error():
+    """Passing GPU tensor to CPU function raises ValueError."""
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.Buffer((128,), "float32"), b: T.Buffer((128,), "float32")):
+        for i in range(128):
+            b[i] = a[i] + T.float32(1)
+
+    lib = tvm.compile(func, target="llvm")
+    a_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    b_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+    lib(a_ok, b_ok)  # correct input should pass
+
+    def run_and_check():
+        a_gpu = tvm.runtime.tensor(np.zeros(128, dtype="float32"), device=tvm.cuda(0))
+        b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "Mismatched a.device_type on argument #0 when calling:\n"
+                "  `func(a: Tensor([128], float32), b: Tensor([128], float32))`,\n"
+                "  expected cpu"
+            ),
+        ):
+            lib(a_gpu, b)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ── Scalar type mismatch errors ─────────────────────────────
+
+
+def test_type_mismatch_int_parameter(codegen_target):
+    """Passing a tensor where an int is expected raises TypeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.int32) -> T.int32:
+        if x > 0:
+            return 10
+        else:
+            return 20
+
+    lib = tvm.compile(func, target=codegen_target)
+    assert lib(5) == 10  # correct input should pass
+
+    a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched type on argument #0 when calling:\n  `func(x: int32)`,\n  expected int"
+        ),
+    ):
+        lib(a)
+
+
+def test_type_mismatch_float_parameter(codegen_target):
+    """Passing a tensor where a float is expected raises TypeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.float32) -> T.int32:
+        if x > T.float32(0):
+            return 1
+        else:
+            return 0
+
+    lib = tvm.compile(func, target=codegen_target)
+    assert lib(1.0) == 1  # correct input should pass
+
+    a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched type on argument #0 when calling:\n  `func(x: float32)`,\n  expected float"
+        ),
+    ):
+        lib(a)
+
+
+def test_type_mismatch_bool_parameter(codegen_target):
+    """Passing a tensor where a bool is expected raises TypeError."""
+
+    @T.prim_func(s_tir=True)
+    def func(x: T.bool) -> T.int32:
+        if x:
+            return 1
+        else:
+            return 0
+
+    lib = tvm.compile(func, target=codegen_target)
+    assert lib(True) == 1  # correct input should pass
+
+    a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched type on argument #0 when calling:\n  `func(x: bool)`,\n  expected boolean"
+        ),
+    ):
+        lib(a)
+
+
+# ── Forward-reference symbolic variable ────────────────────
+
+
+def test_forward_reference_symbolic_shape(codegen_target):
+    """Buffers sharing a symbolic var with forward reference compile and run correctly.
+
+    When buffer A has shape (batch_size+1,) and buffer B has shape (batch_size,),
+    batch_size is referenced in A's shape assertion before it is defined from B.
+    The three-sequence separation ensures this works. Also verifies the error
+    message uses rendered access paths (e.g. "B.shape[0] + 1") for shape checks.
+    """
+
+    @T.prim_func(s_tir=True)
+    def func(a: T.handle, b: T.handle):
+        batch_size = T.int64()
+        A = T.match_buffer(a, (batch_size + 1,), "int32")
+        B = T.match_buffer(b, (batch_size,), "int32")
+        for i in range(batch_size):
+            B[i] = A[i] + A[i + 1]
+
+    lib = tvm.compile(func, target=codegen_target)
+    # Correct inputs: A has shape (5,), B has shape (4,)
+    a = tvm.runtime.tensor(np.array([1, 2, 3, 4, 5], dtype="int32"))
+    b = tvm.runtime.tensor(np.zeros(4, dtype="int32"))
+    lib(a, b)
+    np.testing.assert_array_equal(b.numpy(), [3, 5, 7, 9])
+
+    # Wrong shape: A has shape (10,) but B has shape (4,), so batch_size=4 but A needs 5
+    a_wrong = tvm.runtime.tensor(np.zeros(10, dtype="int32"))
+    b_ok = tvm.runtime.tensor(np.zeros(4, dtype="int32"))
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Invalid A.shape[0] on argument #0 when calling:\n"
+            "  `func(A: Tensor([batch_size + T.int64(1)], int32),"
+            " B: Tensor([batch_size], int32))`,\n"
+            "  expected B.shape[0] + 1"
+        ),
+    ):
+        lib(a_wrong, b_ok)
+
+
+# ── Mixed parameter type errors ────────────────────────────
+
+
+def test_invalid_arguments_mixed_params(codegen_target):
+    """Mixed bool + tensor function: type, dtype, and shape errors."""
+
+    @T.prim_func(s_tir=True)
+    def func(a0: T.bool, a1: T.Buffer([10], "float32")) -> T.int32:
+        return 0
+
+    lib = tvm.compile(func, target=codegen_target)
+    lib(True, tvm.runtime.tensor(np.zeros(10, dtype="float32")))  # correct input should pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched type on argument #1 when calling:\n"
+            "  `func(a0: bool, a1: Tensor([10], float32))`,\n"
+            "  expected Tensor"
+        ),
+    ):
+        lib(1, 1)
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            "Mismatched a1.dtype on argument #1 when calling:\n"
+            "  `func(a0: bool, a1: Tensor([10], float32))`,\n"
+            "  expected float32"
+        ),
+    ):
+        lib(1, tvm.runtime.empty([10], "int32"))
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Invalid a1.shape[0] on argument #1 when calling:\n"
+            "  `func(a0: bool, a1: Tensor([10], float32))`,\n"
+            "  expected 10"
+        ),
+    ):
+        lib(False, tvm.runtime.empty([11], "float32"))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_gpu_codegen_allreduce.py

@@ -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: E741
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def _reduce_sum_module(d1, d2, d3):
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1, d1, d2, d3), "float32"), B: T.Buffer((1, d1, d2), "float32")):
+            for i in T.thread_binding(1, thread="blockIdx.x"):
+                for j in T.thread_binding(d1, thread="threadIdx.z"):
+                    for k in T.thread_binding(d2, thread="threadIdx.y"):
+                        for l in T.thread_binding(d3, thread="threadIdx.x"):
+                            with T.sblock("reduce"):
+                                vi, vj, vk, vl = T.axis.remap("SSSR", [i, j, k, l])
+                                T.reads(A[vi, vj, vk, vl])
+                                T.writes(B[vi, vj, vk])
+                                with T.init():
+                                    B[vi, vj, vk] = T.float32(0.0)
+                                B[vi, vj, vk] = B[vi, vj, vk] + A[vi, vj, vk, vl]
+
+    return Module
+
+
+def _reduce_max_module(d1, d2, d3):
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1, d1, d2, d3), "float32"), B: T.Buffer((1, d1, d2), "float32")):
+            for i in T.thread_binding(1, thread="blockIdx.x"):
+                for j in T.thread_binding(d1, thread="threadIdx.z"):
+                    for k in T.thread_binding(d2, thread="threadIdx.y"):
+                        for l in T.thread_binding(d3, thread="threadIdx.x"):
+                            with T.sblock("reduce"):
+                                vi, vj, vk, vl = T.axis.remap("SSSR", [i, j, k, l])
+                                T.reads(A[vi, vj, vk, vl])
+                                T.writes(B[vi, vj, vk])
+                                with T.init():
+                                    B[vi, vj, vk] = T.float32(-3.4028234663852886e38)
+                                B[vi, vj, vk] = T.max(B[vi, vj, vk], A[vi, vj, vk, vl])
+
+    return Module
+
+
+def generate_param_sets():
+    for d1 in range(1, 5):
+        for d2 in range(1, 5):
+            for d3 in [2, 4, 8, 12, 16, 32, 48, 64, 100, 128, 201, 256, 512, 1024]:
+                if d1 * d2 * d3 < 1024:
+                    yield (d1, d2, d3)
+
+
+dims = tvm.testing.parameter(*generate_param_sets())
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+    ],
+)
+def test_allreduce_sum(dims, target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    d1, d2, d3 = dims
+    mod = _reduce_sum_module(d1, d2, d3)
+    f = tvm.compile(mod, target=target)
+
+    # prepare input and output array
+    a_np = np.random.rand(1, d1, d2, d3).astype("float32")
+    b_np = a_np.sum(axis=-1).astype("float32")
+
+    def run_and_check():
+        dev = tvm.device(target)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(np.zeros_like(b_np), dev)
+        f(a, b)
+        tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-6, atol=1e-6)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+define_metal_compile_callback = tvm.testing.parameter(True, False)
+
+
+@pytest.fixture
+def optional_metal_compile_callback(define_metal_compile_callback):
+    name = "tvm_callback_metal_compile"
+    cached = tvm.get_global_func(name, allow_missing=True)
+
+    if define_metal_compile_callback:
+
+        @tvm.register_global_func(name, override=True)
+        def compile_metal(src, target):
+            from tvm.support.xcode import compile_metal  # pylint: disable=import-outside-toplevel
+
+            return compile_metal(src, sdk="macosx")
+
+    yield
+
+    if define_metal_compile_callback:
+        if cached is None:
+            tvm_ffi.registry.remove_global_func(name)
+        else:
+            tvm.register_global_func(name, cached, override=True)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_allreduce_sum_compile(optional_metal_compile_callback):
+    # Disable the parametrization over dims, at least for now
+    dims = (1, 1, 2)
+    target = "metal"
+
+    d1, d2, d3 = dims
+    mod = _reduce_sum_module(d1, d2, d3)
+    tvm.compile(mod, target=target)
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+    ],
+)
+def test_allreduce_max(dims, target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    d1, d2, d3 = dims
+    mod = _reduce_max_module(d1, d2, d3)
+    f = tvm.compile(mod, target=target)
+
+    # prepare input and output array
+    a_np = -np.random.rand(1, d1, d2, d3).astype("float32")
+    b_np = a_np.max(axis=-1).astype("float32")
+
+    def run_and_check():
+        dev = tvm.device(target)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(np.zeros_like(b_np), dev)
+        f(a, b)
+        tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-6, atol=1e-6)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_inject_ptx_ldg32.py

@@ -0,0 +1,74 @@
+# 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 numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@T.prim_func(s_tir=True)
+def vector_add(A: T.Buffer((16), "float32"), B: T.Buffer((32), "float32")) -> None:
+    T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
+    bx = T.env_thread("blockIdx.x")
+    tx = T.env_thread("threadIdx.x")
+    T.launch_thread(bx, 1)
+    T.launch_thread(tx, 32)
+    with T.sblock():
+        A_local = T.sblock_alloc_buffer((32), "float32", scope="local")
+
+        with T.sblock():
+            T.reads(A[0:16])
+            T.writes(A_local[0:32])
+            A_local[tx] = T.if_then_else(tx % 2 == 0, A[tx // 2], T.float32(0), dtype="float32")
+            B[tx] = A_local[tx] + 1.0
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_inject_ptx_intrin():
+    f = vector_add
+    arch = tvm.support.nvcc.get_target_compute_version()
+    major, _ = tvm.support.nvcc.parse_compute_version(arch)
+    if major < 8:
+        # Require at least SM80
+        return
+    with tvm.transform.PassContext(config={"tirx.ptx.ldg32": True}):
+        mod = tvm.compile(f, target="cuda")
+    A_np = np.random.rand(16).astype("float32")
+    B_np = np.zeros(32).astype("float32")
+    C_np = np.zeros(32).astype("float32")
+
+    for i in range(32):
+        if i % 2 == 0:
+            C_np[i] = A_np[i // 2]
+        C_np[i] += 1.0
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_nd = tvm.runtime.tensor(A_np, device=dev)
+        B_nd = tvm.runtime.tensor(B_np, device=dev)
+        mod(A_nd, B_nd)
+        tvm.testing.assert_allclose(B_nd.numpy(), C_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    test_inject_ptx_intrin()

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()

tests/python/codegen/test_target_codegen_aarch64.py

@@ -0,0 +1,647 @@
+# 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.
+
+"""
+Codegen tests for AArch64
+"""
+
+import re
+
+import pytest
+
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.target.codegen import llvm_version_major
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_mul(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] * B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"mul\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_add(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] + B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"add\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_sub(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] - B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"sub\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_muladd(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle, var_D: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            D = T.match_buffer(var_D, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("D"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i], C[v_i])
+                    T.writes(D[v_i])
+                    D[v_i] = A[v_i] * B[v_i] + C[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        # Group the mad|mla alternation: a top-level alternation would let a bare
+        # "mad" match anywhere in the assembly (e.g. inside scalar "fmadd").
+        r"(?:mad|mla)\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]",
+        assembly,
+    )
+
+    if llvm_version_major() >= 18 and dtype in ("float", "float16"):
+        # Newer LLVM cost models (observed with LLVM 20) prefer a fixed-width
+        # NEON main loop over a scalable SVE loop for floating-point fmuladd
+        # on generic AArch64 targets, so also accept the NEON form.
+        loads += re.findall(r"ld[rp]\tq[0-9]", assembly)
+        matches += re.findall(r"fml[as]\tv[0-9]+\.[0-9]+[hs]", assembly)
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_max(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = T.max(A[v_i], B[v_i])
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    compare = re.findall(
+        r"cmgt\tp[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+    select = re.findall("sel\tz[0-9].[shdb], p[0-9], z[0-9].[shdb], z[0-9].[shdb]", assembly)
+    max_instr = re.findall(
+        r"max\tz[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert (len(compare) > 1 and len(select) == len(compare)) or len(max_instr) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_min(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = T.min(A[v_i], B[v_i])
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    compare = re.findall(
+        r"cmgt\tp[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+    select = re.findall("sel\tz[0-9].[shdb], p[0-9], z[0-9].[shdb], z[0-9].[shdb]", assembly)
+    min_instr = re.findall(
+        r"min\tz[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert (len(compare) > 1 and len(select) == len(compare)) or len(min_instr) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_div(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = tvm.tirx.div(A[v_i], B[v_i])
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"div\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) >= 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"]
+)
+def test_mod(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = T.floormod(A[v_i], B[v_i])
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"mls\tz[0-9].[shdb],( p[0-9]/[m],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 0
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_eq(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), "bool")
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] == B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"cm(p)?eq\tp[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    # The number of SVE compares depends on the LLVM cost model: LLVM <= 17
+    # interleaves the scalable loop by two, while LLVM 20 emits a fixed-width
+    # NEON main loop with a single predicated SVE epilogue.
+    assert len(matches) > 0
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["float", "float16", "uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"],
+)
+def test_neq(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), "bool")
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] != B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"cm(p)?(gt|ne)\tp[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    # The number of SVE compares depends on the LLVM cost model: LLVM <= 17
+    # interleaves the scalable loop by two, while LLVM 20 emits a fixed-width
+    # NEON main loop with a single predicated SVE epilogue.
+    assert len(matches) > 0
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"]
+)
+def test_or(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] | B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"orr\tz[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"]
+)
+def test_and(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i], B[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] & B[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"and\tz[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.parametrize(
+    "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"]
+)
+def test_not(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = ~A[v_i]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+    matches = re.findall(
+        r"eor\tz[0-9].[shdb],( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]", assembly
+    )
+
+    assert len(loads) > 1
+    assert len(matches) > 1
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
+)
+@pytest.mark.xfail(
+    reason="Awaiting llvm support for gathered loads",
+    strict=True,
+)
+@pytest.mark.parametrize(
+    "dtype", ["uint8", "uint16", "uint32", "uint64", "int8", "int16", "int32", "int64"]
+)
+def test_memcpy(dtype):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,), dtype=dtype)
+            B = T.match_buffer(var_B, (m,), "int32")
+            C = T.match_buffer(var_C, (m,), dtype=dtype)
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(B[v_i], A[B[v_i]])
+                    T.writes(C[v_i])
+                    C[v_i] = A[B[v_i]]
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    assembly = f.inspect_source("asm")
+    loads = re.findall("ld1[whdb]	{ z", assembly)
+
+    assert len(loads) > 0
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 13,
+    reason="Function attribute vscale_range() is not supported in earlier versions of LLVM",
+)
+@pytest.mark.parametrize(
+    "mattr,expect_attr",
+    [
+        ("+neon", False),
+        ("+sve", True),
+        # Since LLVM 19, SVE/SVE2 are optional extensions of Armv9.0-A, so
+        # "+v9a" no longer implies "+sve" and no vscale_range() is added:
+        # https://releases.llvm.org/19.1.0/docs/ReleaseNotes.html#changes-to-the-aarch64-backend
+        ("+v9a", llvm_version_major() < 19),
+        ("+sme", True),
+    ],
+)
+def test_vscale_range_function_attribute(mattr, expect_attr):
+    target = {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": [mattr]}
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_C: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            m = T.int32()
+            A = T.match_buffer(var_A, (m,))
+            C = T.match_buffer(var_C, (m,))
+            for i in range(m):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(m, i)
+                    T.reads(A[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] + T.float32(1)
+
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(Module)
+
+    # Check if the vscale_range() attribute exists
+    ll = f.inspect_source("ll")
+    attr = re.findall(r".*vscale_range\(\d+,\d+\)*.", ll)
+
+    if expect_attr:
+        assert len(attr) > 0, "Function attribute vscale_range() was not found in generated LLVM IR"
+    else:
+        assert len(attr) == 0, (
+            "Unexpected function attribute vscale_range() was found in generated LLVM IR"
+        )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_arm.py

@@ -0,0 +1,137 @@
+# 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 re
+
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_popcount():
+    target = {
+        "kind": "llvm",
+        "mtriple": "armv7l-none-linux-gnueabihf",
+        "mcpu": "cortex-a53",
+        "mattr": ["+neon"],
+    }
+
+    def check_correct_assembly(type, elements, counts):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((elements,), type), B: T.Buffer((elements,), type)):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.vectorized(elements):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(elements, i)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = T.popcount(A[v_i])
+
+        f = tvm.tirx.build(Module, target=target)
+        # Verify we see the correct number of vpaddl and vcnt instructions in the assembly
+        assembly = f.inspect_source("asm")
+        matches = re.findall("vpaddl", assembly)
+        assert len(matches) == counts
+        matches = re.findall("vcnt", assembly)
+        assert len(matches) == 1
+
+    check_correct_assembly("uint16", 8, 1)
+    check_correct_assembly("uint16", 4, 1)
+    check_correct_assembly("uint32", 4, 2)
+    check_correct_assembly("uint32", 2, 2)
+    check_correct_assembly("uint64", 2, 3)
+
+
+def test_vmlal_s16():
+    target = {
+        "kind": "llvm",
+        "mtriple": "armv7l-none-linux-gnueabihf",
+        "mcpu": "cortex-a53",
+        "mattr": ["+neon"],
+    }
+
+    def check_correct_assembly(N):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(var_A: T.handle, var_B: T.handle, C: T.Buffer((N,), "int32")):
+                T.func_attr({"tirx.noalias": True})
+                K = T.int32()
+                A = T.match_buffer(var_A, (K, N), "int8")
+                B = T.match_buffer(var_B, (K, N), "int8")
+                for n in T.vectorized(N):
+                    for rv in range(K):
+                        with T.sblock("C"):
+                            v_n, v_rv = T.axis.remap("SR", [n, rv])
+                            T.reads(A[v_rv, v_n], B[v_rv, v_n])
+                            T.writes(C[v_n])
+                            with T.init():
+                                C[v_n] = 0
+                            C[v_n] = C[v_n] + T.Cast("int32", A[v_rv, v_n]) * T.Cast(
+                                "int32", B[v_rv, v_n]
+                            )
+
+        f = tvm.tirx.build(Module, target=target)
+
+        # Verify we see the correct number of vmlal.s16 instructions
+        assembly = f.inspect_source("asm")
+        matches = re.findall("vmlal.s16", assembly)
+        assert len(matches) == N // 4
+
+    check_correct_assembly(8)
+    check_correct_assembly(16)
+    check_correct_assembly(32)
+    check_correct_assembly(64)
+
+    def check_broadcast_correct_assembly(N):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(var_A: T.handle, var_B: T.handle, C: T.Buffer((N,), "int32")):
+                T.func_attr({"tirx.noalias": True})
+                K = T.int32()
+                A = T.match_buffer(var_A, (K, N), "int8")
+                B = T.match_buffer(var_B, (K,), "int8")
+                for n in T.vectorized(N):
+                    for rv in range(K):
+                        with T.sblock("C"):
+                            v_n, v_rv = T.axis.remap("SR", [n, rv])
+                            T.reads(A[v_rv, v_n], B[v_rv])
+                            T.writes(C[v_n])
+                            with T.init():
+                                C[v_n] = 0
+                            C[v_n] = C[v_n] + T.Cast("int32", A[v_rv, v_n]) * T.Cast(
+                                "int32", B[v_rv]
+                            )
+
+        f = tvm.tirx.build(Module, target=target)
+
+        # Verify we see the correct number of vmlal.s16 instructions
+        assembly = f.inspect_source("asm")
+        matches = re.findall("vmlal.s16", assembly)
+        assert len(matches) == N // 4
+
+    check_broadcast_correct_assembly(8)
+    check_broadcast_correct_assembly(16)
+    check_broadcast_correct_assembly(32)
+    check_broadcast_correct_assembly(64)
+
+
+if __name__ == "__main__":
+    test_popcount()
+    test_vmlal_s16()

tests/python/codegen/test_target_codegen_blob.py

@@ -0,0 +1,108 @@
+# 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: F821
+
+import ctypes
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.support import cc, popen_pool, tar, utils
+
+
+@pytest.mark.gpu
+def test_cuda_multi_lib():
+    pytest.importorskip("cloudpickle")
+
+    # test combining two system lib together
+    # each contains a fatbin component in cuda
+    for device in ["llvm", "cuda"]:
+        if not tvm.testing.device_enabled(device):
+            print(f"skip because {device} is not enabled...")
+            return
+
+    @tvm.script.ir_module
+    class ModA:
+        I.module_attrs({"system_lib_prefix": "modA_"})
+
+        @T.prim_func(s_tir=True)
+        def my_inplace_update(x: T.Buffer((12), "float32")) -> None:
+            T.func_attr({"global_symbol": "modA_my_inplace_update"})
+            for bx in T.thread_binding(T.int64(1), thread="blockIdx.x"):
+                for tx in T.thread_binding(T.int64(12), thread="threadIdx.x"):
+                    x[tx] = x[tx] + 1
+
+    @tvm.script.ir_module
+    class ModB:
+        I.module_attrs({"system_lib_prefix": "modB_"})
+
+        @T.prim_func(s_tir=True)
+        def my_inplace_update(x: T.Buffer((12), "float32")) -> None:
+            T.func_attr({"global_symbol": "modB_my_inplace_update"})
+            for bx in T.thread_binding(T.int64(1), thread="blockIdx.x"):
+                for tx in T.thread_binding(T.int64(12), thread="threadIdx.x"):
+                    x[tx] = x[tx] + 2
+
+    temp = utils.tempdir()
+    target = tvm.target.Target("cuda", host="llvm")
+    libA = tvm.compile(ModA, target=target)
+    libB = tvm.compile(ModB, target=target)
+
+    pathA = temp.relpath("libA.tar")
+    pathB = temp.relpath("libB.tar")
+    pathAll = temp.relpath("libAll.a")
+
+    path_dso = temp.relpath("mylib.so")
+    libA.export_library(pathA, fcompile=tar.tar)
+    libB.export_library(pathB, fcompile=tar.tar)
+    cc.create_staticlib(pathAll, [pathA, pathB])
+    # package two static libs together
+    cc.create_shared(path_dso, ["-Wl,--whole-archive", pathAll, "-Wl,--no-whole-archive"])
+
+    def popen_check():
+        def run_and_check():
+            # Load dll, will trigger system library registration
+            ctypes.CDLL(path_dso)
+            # Load the system wide library
+            dev = tvm.cuda()
+            a_np = np.random.uniform(size=12).astype("float32")
+            a_nd = tvm.runtime.tensor(a_np, dev)
+            b_nd = tvm.runtime.tensor(a_np, dev)
+            syslibA = tvm.runtime.system_lib("modA_")
+            syslibB = tvm.runtime.system_lib("modB_")
+            # reload same lib twice
+            syslibA = tvm.runtime.system_lib("modA_")
+            syslibA["my_inplace_update"](a_nd)
+            syslibB["my_inplace_update"](b_nd)
+            np.testing.assert_equal(a_nd.numpy(), a_np + 1)
+            np.testing.assert_equal(b_nd.numpy(), a_np + 2)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    # system lib should be loaded in different process
+    worker = popen_pool.PopenWorker()
+    worker.send(popen_check)
+    worker.recv()
+
+
+if __name__ == "__main__":
+    test_synthetic()
+    test_cuda_multilib()

tests/python/codegen/test_target_codegen_bool.py

@@ -0,0 +1,109 @@
+# 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.
+"""codegen related to bool types"""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("target", ["llvm", "cuda", "rocm", "vulkan", "metal", "opencl"])
+def test_cmp_load_store(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @I.ir_module(s_tir=True)
+    class GPUModule:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((32,), "float32"),
+            B: T.Buffer((32,), "float32"),
+            D: T.Buffer((32,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            C = T.sblock_alloc_buffer((32,), "bool")
+            for i0_0 in T.thread_binding(8, thread="blockIdx.x"):
+                for i0_1 in T.thread_binding(4, thread="blockIdx.x"):
+                    with T.sblock("C"):
+                        v_i0 = T.axis.spatial(32, i0_0 * 4 + i0_1)
+                        T.reads(B[v_i0], A[v_i0])
+                        T.writes(C[v_i0])
+                        C[v_i0] = B[v_i0] < A[v_i0]
+            for i0_0 in T.thread_binding(8, thread="blockIdx.x"):
+                for i0_1 in T.thread_binding(4, thread="blockIdx.x"):
+                    with T.sblock("D"):
+                        v_i0 = T.axis.spatial(32, i0_0 * 4 + i0_1)
+                        T.reads(C[v_i0], A[v_i0])
+                        T.writes(D[v_i0])
+                        D[v_i0] = T.Cast("float32", C[v_i0] and T.float32(1.0) < A[v_i0])
+
+    @I.ir_module(s_tir=True)
+    class CPUModule:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((32,), "float32"),
+            B: T.Buffer((32,), "float32"),
+            D: T.Buffer((32,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            C = T.sblock_alloc_buffer((32,), "bool")
+            for i0 in range(32):
+                with T.sblock("C"):
+                    v_i0 = T.axis.spatial(32, i0)
+                    T.reads(B[v_i0], A[v_i0])
+                    T.writes(C[v_i0])
+                    C[v_i0] = B[v_i0] < A[v_i0]
+            for i0 in range(32):
+                with T.sblock("D"):
+                    v_i0 = T.axis.spatial(32, i0)
+                    T.reads(C[v_i0], A[v_i0])
+                    T.writes(D[v_i0])
+                    D[v_i0] = T.Cast("float32", C[v_i0] and T.float32(1.0) < A[v_i0])
+
+    arr_size = 32
+    is_gpu = tvm.target.Target(target).kind.name != "llvm"
+    mod = GPUModule if is_gpu else CPUModule
+
+    f = tvm.compile(mod, target=target)
+
+    a_np = np.random.uniform(size=arr_size).astype("float32")
+    b_np = np.random.uniform(size=arr_size).astype("float32")
+
+    def run_and_check():
+        dev = tvm.device(target)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        d = tvm.runtime.tensor(np.zeros(arr_size, dtype="float32"), dev)
+        f(a, b, d)
+        np.testing.assert_equal(
+            d.numpy(),
+            np.logical_and(a_np > b_np, a_np > 1).astype("float32"),
+        )
+
+    if is_gpu:
+        tvm.testing.run_with_gpu_lock(run_and_check)
+    else:
+        run_and_check()
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_c_host.py

@@ -0,0 +1,249 @@
+# 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 numpy as np
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.support import utils
+
+
+def test_add():
+    nn = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def test_fadd(
+            A: T.Buffer((1024,), "float32"),
+            B: T.Buffer((1024,), "float32"),
+            C: T.Buffer((1024,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in range(1024):
+                with T.sblock("C"):
+                    v_i0 = T.axis.spatial(1024, i0)
+                    T.reads(A[v_i0], B[v_i0])
+                    T.writes(C[v_i0])
+                    C[v_i0] = A[v_i0] + B[v_i0]
+
+    def check_c():
+        mhost = tvm.compile(Module, target="c")
+        temp = utils.tempdir()
+        path_dso = temp.relpath("temp.so")
+        mhost.export_library(path_dso)
+        m = tvm.runtime.load_module(path_dso)
+        fadd = m["test_fadd"]
+        dev = tvm.cpu(0)
+        n = nn
+        a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+        b = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+        c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+        fadd(a, b, c)
+        tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
+
+    check_c()
+
+
+def test_reinterpret():
+    nn = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def test_reinterpret(
+            A: T.Buffer((1024,), "int32"),
+            B: T.Buffer((1024,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in range(1024):
+                with T.sblock("B"):
+                    v_i0 = T.axis.spatial(1024, i0)
+                    T.reads(A[v_i0])
+                    T.writes(B[v_i0])
+                    B[v_i0] = T.reinterpret("float32", A[v_i0] + 2)
+
+    def check_c():
+        mhost = tvm.compile(Module, target="c")
+        temp = utils.tempdir()
+        path_dso = temp.relpath("temp.so")
+        mhost.export_library(path_dso)
+        m = tvm.runtime.load_module(path_dso)
+        fadd = m["test_reinterpret"]
+        dev = tvm.cpu(0)
+        n = nn
+        a = tvm.runtime.tensor(np.random.randint(-(2**30), 2**30, size=n).astype("int32"), dev)
+        b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+        fadd(a, b)
+        tvm.testing.assert_allclose(b.numpy(), (2 + a.numpy()).view("float32"))
+
+    check_c()
+
+
+def test_ceil():
+    nn = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def test_ceil(
+            A: T.Buffer((1024,), "float32"),
+            B: T.Buffer((1024,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in range(1024):
+                with T.sblock("B"):
+                    v_i0 = T.axis.spatial(1024, i0)
+                    T.reads(A[v_i0])
+                    T.writes(B[v_i0])
+                    B[v_i0] = T.ceil(A[v_i0])
+
+    def check_c():
+        mhost = tvm.compile(Module, target="c")
+        temp = utils.tempdir()
+        path_dso = temp.relpath("temp.so")
+        mhost.export_library(path_dso)
+        m = tvm.runtime.load_module(path_dso)
+        fceil = m["test_ceil"]
+        dev = tvm.cpu(0)
+        n = nn
+        a = tvm.runtime.tensor(np.random.rand(n).astype("float32"), dev)
+        b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+        fceil(a, b)
+        tvm.testing.assert_allclose(b.numpy(), (np.ceil(a.numpy()).view("float32")))
+
+    check_c()
+
+
+def test_floor():
+    nn = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def test_floor(
+            A: T.Buffer((1024,), "float32"),
+            B: T.Buffer((1024,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in range(1024):
+                with T.sblock("B"):
+                    v_i0 = T.axis.spatial(1024, i0)
+                    T.reads(A[v_i0])
+                    T.writes(B[v_i0])
+                    B[v_i0] = T.floor(A[v_i0])
+
+    def check_c():
+        mhost = tvm.compile(Module, target="c")
+        temp = utils.tempdir()
+        path_dso = temp.relpath("temp.so")
+        mhost.export_library(path_dso)
+        m = tvm.runtime.load_module(path_dso)
+        ffloor = m["test_floor"]
+        dev = tvm.cpu(0)
+        n = nn
+        a = tvm.runtime.tensor(np.random.rand(n).astype("float32"), dev)
+        b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+        ffloor(a, b)
+        tvm.testing.assert_allclose(b.numpy(), (np.floor(a.numpy()).view("float32")))
+
+    check_c()
+
+
+def test_round():
+    nn = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def test_round(
+            A: T.Buffer((1024,), "float32"),
+            B: T.Buffer((1024,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in range(1024):
+                with T.sblock("B"):
+                    v_i0 = T.axis.spatial(1024, i0)
+                    T.reads(A[v_i0])
+                    T.writes(B[v_i0])
+                    B[v_i0] = T.round(A[v_i0])
+
+    def check_c():
+        mhost = tvm.compile(Module, target="c")
+        temp = utils.tempdir()
+        path_dso = temp.relpath("temp.so")
+        mhost.export_library(path_dso)
+        m = tvm.runtime.load_module(path_dso)
+        fround = m["test_round"]
+        dev = tvm.cpu(0)
+        n = nn
+        a = tvm.runtime.tensor(np.random.rand(n).astype("float32"), dev)
+        b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+        fround(a, b)
+        tvm.testing.assert_allclose(b.numpy(), (np.round(a.numpy()).view("float32")))
+
+    check_c()
+
+
+def test_subroutine_call():
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer(1, dtype="float32")):
+            Module.subroutine(A.data)
+
+        @T.prim_func(private=True, s_tir=True)
+        def subroutine(A_data: T.handle("float32")):
+            A = T.decl_buffer(1, dtype="float32", data=A_data)
+            A[0] = 42.0
+
+    built = tvm.tirx.build(Module, target="c")
+
+    source = built.inspect_source()
+    assert source.count("__tvm_ffi_main(void*") == 2, (
+        "Expected two occurrences, for forward-declaration and definition"
+    )
+    assert source.count("subroutine(float*") == 2, (
+        "Expected two occurrences, for forward-declaration and definition"
+    )
+    assert source.count("subroutine(") == 3, (
+        "Expected three occurrences, for forward-declaration, definition, and call from main."
+    )
+
+
+def test_workspace_allocation_cast():
+    @I.ir_module
+    class Module:
+        @T.prim_func
+        def main(A: T.Buffer((256,), "float32")):
+            workspace = T.alloc_buffer((256,), "float32", scope="global")
+            for i in range(256):
+                workspace[i] = A[i]
+            for i in range(256):
+                A[i] = workspace[i]
+
+    built = tvm.tirx.build(Module, target="c")
+    assert "((float*)TVMBackendAllocWorkspace(" in built.inspect_source()
+
+    temp = utils.tempdir()
+    built.export_library(temp.relpath("workspace.so"))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_cross_llvm.py

@@ -0,0 +1,103 @@
+# 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: F401, F841
+"""Test cross compilation"""
+
+import os
+import struct
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import rpc
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.support import cc, utils
+from tvm.testing import env
+
+
+@I.ir_module(s_tir=True)
+class AddModule:
+    @T.prim_func(s_tir=True)
+    def main(
+        A: T.Buffer((1024,), "float32"),
+        B: T.Buffer((1024,), "float32"),
+        C: T.Buffer((1024,), "float32"),
+    ):
+        T.func_attr({"tirx.noalias": True})
+        for i0_0 in T.parallel(256):
+            for i0_1 in T.vectorized(4):
+                with T.sblock("C"):
+                    v_i0 = T.axis.spatial(1024, i0_0 * 4 + i0_1)
+                    T.reads(A[v_i0], B[v_i0])
+                    T.writes(C[v_i0])
+                    C[v_i0] = A[v_i0] + B[v_i0]
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_llvm_add_pipeline():
+    nn = 1024
+
+    def verify_elf(path, e_machine):
+        with open(path, "rb") as fi:
+            arr = fi.read(20)
+            assert struct.unpack("ccc", arr[1:4]) == (b"E", b"L", b"F")
+            endian = struct.unpack("b", arr[0x5:0x6])[0]
+            endian = "<" if endian == 1 else ">"
+            assert struct.unpack(endian + "h", arr[0x12:0x14])[0] == e_machine
+
+    def build_arm():
+        target = {"kind": "llvm", "mtriple": "armv7-none-linux-gnueabihf"}
+        if not tvm.runtime.enabled("llvm"):
+            print(f"Skip because {target} is not enabled..")
+            return
+        temp = utils.tempdir()
+        f = tvm.tirx.build(AddModule, target=target)
+        path = temp.relpath("myadd.o")
+        f.write_to_file(path)
+        verify_elf(path, 0x28)
+        asm_path = temp.relpath("myadd.asm")
+        f.write_to_file(asm_path)
+        # Do a RPC verification, launch kernel on Arm Board if available.
+        host = os.environ.get("TVM_RPC_ARM_HOST", None)
+        remote = None
+        if host:
+            port = int(os.environ["TVM_RPC_ARM_PORT"])
+            try:
+                remote = rpc.connect(host, port)
+            except RuntimeError as e:
+                pass
+
+        if remote:
+            remote.upload(path)
+            farm = remote.load_module("myadd.o")
+            dev = remote.cpu(0)
+            n = nn
+            a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+            b = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+            c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+            farm(a, b, c)
+            tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
+            print("Verification finish on remote..")
+
+    build_arm()
+
+
+if __name__ == "__main__":
+    test_llvm_add_pipeline()

tests/python/codegen/test_target_codegen_cuda_fastmath.py

@@ -0,0 +1,306 @@
+# 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 re
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+import tvm.tirx as tirx
+from tvm.ir.module import IRModule
+from tvm.runtime.executable import Executable
+from tvm.script import tirx as T
+from tvm.support.nvcc import have_fp16
+from tvm.testing import env
+
+VECTOR_N_INPUTS = 8
+
+
+def make_prim_func(
+    name: str,
+    dtype: str,
+    num_inputs: int,
+    op: Callable[[tirx.Expr, ...], tirx.Expr],
+) -> tirx.PrimFunc:
+    """Make a primitive function that applies the given operation to the input buffer."""
+    if num_inputs == 1:
+
+        @T.prim_func
+        def kernel(
+            A: T.Buffer((VECTOR_N_INPUTS,), dtype),
+            B: T.Buffer((VECTOR_N_INPUTS,), dtype),
+        ):
+            T.func_attr({"global_symbol": name + "_kernel", "tirx.noalias": True})
+            for i in T.thread_binding(VECTOR_N_INPUTS, thread="threadIdx.x"):
+                B[i] = op(A[i])
+
+        return kernel
+    elif num_inputs == 2:
+
+        @T.prim_func
+        def kernel(
+            A: T.Buffer((VECTOR_N_INPUTS,), dtype),
+            E: T.Buffer((VECTOR_N_INPUTS,), dtype),
+            B: T.Buffer((VECTOR_N_INPUTS,), dtype),
+        ):
+            T.func_attr({"global_symbol": name + "_kernel", "tirx.noalias": True})
+            for i in T.thread_binding(VECTOR_N_INPUTS, thread="threadIdx.x"):
+                B[i] = op(A[i], E[i])
+
+        return kernel
+    else:
+        raise ValueError(f"Unsupported number of inputs: {num_inputs}")
+
+
+@dataclass(frozen=True)
+class MathCase:
+    name: str
+    op: Callable[[tirx.Expr, ...], tirx.Expr]
+    num_inputs: int
+    default_intrinsic_f16: str
+    default_intrinsic_bf16: str
+    default_intrinsic_f32: str
+    default_intrinsic_f64: str
+    fast_math_intrinsic_f32: str
+    np_ref: object
+    rtol: float = 1e-5
+    atol: float = 1e-6
+
+
+MATH_CASES = [
+    MathCase(
+        "exp_case",
+        T.exp,
+        1,
+        "hexp",
+        "hexp",
+        "expf",
+        "exp",
+        "__expf",
+        lambda x: np.exp(x),
+    ),
+    MathCase(
+        "exp10_case",
+        T.exp10,
+        1,
+        "hexp10",
+        "hexp10",
+        "exp10f",
+        "exp10",
+        "__exp10f",
+        lambda x: np.power(10.0, x),
+    ),
+    MathCase(
+        "log_case",
+        T.log,
+        1,
+        "hlog",
+        "hlog",
+        "logf",
+        "log",
+        "__logf",
+        lambda x: np.log(x),
+    ),
+    MathCase(
+        "log2_case",
+        T.log2,
+        1,
+        "hlog2",
+        "hlog2",
+        "log2f",
+        "log2",
+        "__log2f",
+        lambda x: np.log2(x),
+    ),
+    MathCase(
+        "log10_case",
+        T.log10,
+        1,
+        "hlog10",
+        "hlog10",
+        "log10f",
+        "log10",
+        "__log10f",
+        lambda x: np.log10(x),
+    ),
+    MathCase(
+        "tan_case",
+        T.tan,
+        1,
+        "htan",
+        "htan",
+        "tanf",
+        "tan",
+        "tanf",
+        lambda x: np.tan(x),
+    ),
+    MathCase(
+        "cos_case",
+        T.cos,
+        1,
+        "hcos",
+        "hcos",
+        "cosf",
+        "cos",
+        "__cosf",
+        lambda x: np.cos(x),
+    ),
+    MathCase(
+        "sin_case",
+        T.sin,
+        1,
+        "hsin",
+        "hsin",
+        "sinf",
+        "sin",
+        "__sinf",
+        lambda x: np.sin(x),
+    ),
+    MathCase(
+        "tanh_case",
+        T.tanh,
+        1,
+        "htanh",
+        "htanh",
+        "tanhf",
+        "tanh",
+        "__tanhf",
+        lambda x: np.tanh(x),
+    ),
+    MathCase(
+        "pow_case",
+        T.pow,
+        2,
+        "hpow",
+        "hpow",
+        "powf",
+        "pow",
+        "__powf",
+        lambda x, y: np.power(x, y),
+    ),
+]
+
+
+def make_mod(
+    dtype: str, case: MathCase, enable_fast_math: bool
+) -> tuple[tvm.target.Target, tvm.IRModule]:
+    """Make a module for the given dtype and case."""
+    target = tvm.target.Target("cuda")
+    prim_func = make_prim_func(case.name, dtype, case.num_inputs, case.op)
+    return target, tvm.IRModule.from_expr(prim_func.with_attr("target", target))
+
+
+def expected_intrinsic(dtype: str, case: MathCase, enable_fast_math: bool) -> str:
+    """Get the expected intrinsic for the given dtype and case."""
+    if dtype == "float16":
+        return case.default_intrinsic_f16
+    elif dtype == "bfloat16":
+        return case.default_intrinsic_bf16
+    elif dtype == "float32":
+        return case.fast_math_intrinsic_f32 if enable_fast_math else case.default_intrinsic_f32
+    elif dtype == "float64":
+        return case.default_intrinsic_f64
+    else:
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+
+def check_lowered_ir(
+    dtype: str, case: MathCase, enable_fast_math: bool
+) -> tuple[tvm.target.Target, IRModule]:
+    """Check the lowered IR for the given dtype and case."""
+    target, mod = make_mod(dtype, case, enable_fast_math)
+    with tvm.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
+        lowered_mod = tvm.tirx.transform.LowerIntrin()(mod)
+    script = lowered_mod.script(show_meta=False)
+    expected = expected_intrinsic(dtype, case, enable_fast_math)
+    assert re.search(rf"""["']{re.escape(expected)}["']""", script)
+    return target, lowered_mod
+
+
+def check_cuda_source(
+    target: tvm.target.Target,
+    mod: IRModule,
+    dtype: str,
+    case: MathCase,
+    enable_fast_math: bool,
+) -> Executable:
+    """Check the CUDA source for the given dtype and case."""
+    with tvm.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
+        executable = tvm.compile(mod, target=target)
+    source = executable.mod.imports[0].inspect_source()
+    expected = expected_intrinsic(dtype, case, enable_fast_math)
+    assert re.search(rf"(?<!_)\b{re.escape(expected)}\s*\(", source)
+    return executable
+
+
+def make_numpy_inputs(dtype: str, case: MathCase):
+    """Make the numpy inputs for the given dtype and case."""
+    lhs = np.array([0.25, 0.5, 1.0, 2.0, 4.0, 9.0, 16.0, 10.0], dtype=dtype)
+    if case.num_inputs == 1:
+        return [lhs]
+    elif case.num_inputs == 2:
+        rhs = np.array([2.0, 3.0, 0.5, 1.5, 0.25, 0.5, 2.0, 1.0], dtype=dtype)
+        return [lhs, rhs]
+    else:
+        raise ValueError(f"Unsupported number of inputs: {case.num_inputs}")
+
+
+def check_runtime(dtype: str, case: MathCase, executable: Executable):
+    """Check the runtime for the given dtype and case."""
+    np_inputs = make_numpy_inputs(dtype, case)
+    expected = case.np_ref(*[arr.astype(dtype) for arr in np_inputs]).astype(dtype)
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        tvm_inputs = [tvm.runtime.tensor(arr, device=dev) for arr in np_inputs]
+        output = tvm.runtime.empty((VECTOR_N_INPUTS,), dtype, dev)
+
+        executable(*tvm_inputs, output)
+
+        actual = output.numpy()
+        np.testing.assert_allclose(actual, expected, rtol=case.rtol, atol=case.atol)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize("enable_fast_math", [False, True], ids=["default", "fast_math"])
+def test_cuda_math_intrinsic_lowering_pass_context(enable_fast_math):
+    check_lowered_ir("float32", MATH_CASES[0], enable_fast_math)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize(
+    "dtype",
+    ["float16", "bfloat16", "float32", "float64"],
+)
+@pytest.mark.parametrize("case", MATH_CASES, ids=lambda case: f"{case.name}")
+@pytest.mark.parametrize("enable_fast_math", [False, True], ids=["default", "fast_math"])
+def test_cuda_math_intrinsic_lowering_source_and_runtime(dtype, case, enable_fast_math):
+    if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version):
+        pytest.skip("GPU does not support float16")
+    if dtype == "bfloat16" and case.name.startswith("pow_"):
+        pytest.skip("pow_argnames=case is only supported for float")
+
+    target, lowered_mod = check_lowered_ir(dtype, case, enable_fast_math)
+    executable = check_cuda_source(target, lowered_mod, dtype, case, enable_fast_math)
+    check_runtime(dtype, case, executable)

tests/python/codegen/test_target_codegen_cuda_fp4.py

@@ -0,0 +1,278 @@
+# 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.
+
+from itertools import product
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+try:
+    from ml_dtypes import float4_e2m1fn
+
+    ML_DTYPES_AVAILABLE = True
+except ImportError:
+    ML_DTYPES_AVAILABLE = False
+
+
+@pytest.mark.parametrize("promoted_dtype", ["float32x2", "float16x2"])
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_e2m1_vector_conversions(promoted_dtype):
+    native_dtype = "float4_e2m1fnx2"
+    vector_length = 64
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((vector_length,), native_dtype),
+            B: T.Buffer((vector_length,), native_dtype),
+            C: T.Buffer((vector_length,), native_dtype),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(vector_length // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(vector_length, i_0 * 32 + i_1)
+                        T.reads(A[v_i], B[v_i])
+                        T.writes(C[v_i])
+                        C[v_i] = T.Cast(
+                            native_dtype,
+                            T.Cast(promoted_dtype, A[v_i]) + T.Cast(promoted_dtype, B[v_i]),
+                        )
+
+    target = "cuda"
+    fadd = tvm.compile(Module, target=target)
+
+    if "x" in native_dtype:
+        lanes = int(native_dtype.split("x")[-1])
+    else:
+        lanes = 1
+
+    if "x" in promoted_dtype:
+        promoted_base_dtype = promoted_dtype.split("x")[0]
+    else:
+        promoted_base_dtype = promoted_dtype
+
+    np_shape = (vector_length, lanes) if lanes > 1 else (vector_length,)
+
+    # Create test data - either using ml_dtypes if available, or using int8 with valid FP4 values
+    if ML_DTYPES_AVAILABLE:
+        a_np = np.random.uniform(low=0, high=5, size=np_shape).astype(float4_e2m1fn)
+        b_np = np.random.uniform(low=0, high=5, size=np_shape).astype(float4_e2m1fn)
+    else:
+        # float4_e2m1fn possible values: [0, 0.5, 1, 1.5, 2, 3, 4, 6]
+        # We will create int8 arrays with valid FP4 bit patterns
+        valid_fp4_values = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]  # 4-bit values
+        a_np = np.random.choice(valid_fp4_values, size=np_shape).astype(np.int8)
+        b_np = np.random.choice(valid_fp4_values, size=np_shape).astype(np.int8)
+
+    def run_and_check():
+        dev = tvm.device(target, 0)
+        a = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev)
+        a.copyfrom(a_np)
+        b = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev)
+        b.copyfrom(b_np)
+        c = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype, device=dev)
+        fadd(a, b, c)
+        # For the comparison, we will convert result to the promoted dtype and compare
+        # Note: When ml_dtypes is not available, we skip the numpy-level computation comparison
+        # and just verify that the CUDA kernel compiles and executes without error
+        c_result = c.numpy().astype(promoted_base_dtype)
+        if ML_DTYPES_AVAILABLE:
+            # Full comparison when ml_dtypes is available
+            expected = (a_np + b_np).astype(promoted_base_dtype)
+            tvm.testing.assert_allclose(c_result, expected)
+        else:
+            # When ml_dtypes is not available, we just verify the comparison ran successfully
+            # by checking that we got a result with the expected shape and dtype
+            assert c_result.shape == np_shape
+            assert c_result.dtype == promoted_base_dtype
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+def _shuffle_reinterpret_module(n, num_blocks, vector_length, num_elem_per_storage):
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((n // num_elem_per_storage,), "uint32"),
+            B: T.Buffer((n,), "float16"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    for i_2 in T.vectorized(vector_length):
+                        with T.sblock("C"):
+                            v_i = T.axis.spatial(
+                                n, i_0 * 32 * vector_length + i_1 * vector_length + i_2
+                            )
+                            T.reads(A[v_i])
+                            T.writes(B[v_i])
+                            B[v_i] = T.Shuffle(
+                                [
+                                    T.reinterpret(
+                                        "float4_e2m1fnx2",
+                                        T.bitwise_and(
+                                            T.shift_right(
+                                                A[v_i // num_elem_per_storage],
+                                                ((v_i % num_elem_per_storage) // 2 * 4 * 2).astype(
+                                                    "uint32"
+                                                ),
+                                            ),
+                                            T.uint32((1 << (4 * 2)) - 1),
+                                        ).astype("uint8"),
+                                    ).astype("float16x2")
+                                ],
+                                indices=[v_i % 2],
+                            )
+
+    return Module
+
+
+def _scalar_reinterpret_module(n, num_blocks, vector_length, num_elem_per_storage):
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((n // num_elem_per_storage,), "uint32"),
+            B: T.Buffer((n,), "float16"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    for i_2 in T.vectorized(vector_length):
+                        with T.sblock("C"):
+                            v_i = T.axis.spatial(
+                                n, i_0 * 32 * vector_length + i_1 * vector_length + i_2
+                            )
+                            T.reads(A[v_i])
+                            T.writes(B[v_i])
+                            B[v_i] = T.reinterpret(
+                                "float4_e2m1fn",
+                                T.bitwise_and(
+                                    T.shift_right(
+                                        A[v_i // num_elem_per_storage],
+                                        (v_i % num_elem_per_storage * 4).astype("uint32"),
+                                    ),
+                                    T.uint32((1 << 4) - 1),
+                                ).astype("uint8"),
+                            ).astype("float16")
+
+    return Module
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_e2m1_dequantize():
+    n = 128
+
+    dev = tvm.device("cuda", 0)
+    target = tvm.target.Target.from_device(dev)
+    num_elem_per_storage = 32 // 4
+
+    # We only test the whether the code can be compiled.
+    for func_type, vector_length in product(["shuffle", "scalar"], [1, 2, 4]):
+        if func_type == "shuffle" and vector_length == 1:
+            # Vectorize is necessary for shuffle.
+            continue
+
+        num_blocks = n // (32 * vector_length)
+
+        if func_type == "shuffle":
+            mod = _shuffle_reinterpret_module(n, num_blocks, vector_length, num_elem_per_storage)
+        else:
+            mod = _scalar_reinterpret_module(n, num_blocks, vector_length, num_elem_per_storage)
+
+        tvm.compile(mod, target=target)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_e2m1_scalar_buffer_offset():
+    """Regression test: float4_e2m1fn scalar buffer access uses correct byte offset.
+
+    In CUDA sizeof(__nv_fp4_e2m1) = 1 byte, but fp4 data packs 2 elements per
+    byte.  GetBufferRef must emit ``index / 2`` so that the element index is
+    converted to the correct byte offset.  Without the fix the index was used
+    as-is, producing addresses 2x too large — reading garbage from out-of-bounds
+    memory instead of the correct fp4 value.
+
+    We verify by writing known fp4 values, casting each element to float16 on
+    the GPU, and checking the results match the expected fp4->fp16 conversion.
+    """
+    n = 128
+
+    @T.prim_func(s_tir=True)
+    def func(A_raw: T.Buffer((n // 2,), "uint8"), B: T.Buffer((n,), "float16")):
+        T.func_attr({"tir.noalias": True})
+        A = T.decl_buffer((n,), "float4_e2m1fn", data=A_raw.data)
+        for i in range(n):
+            with T.sblock("B"):
+                vi = T.axis.spatial(n, i)
+                T.reads(A[vi])
+                T.writes(B[vi])
+                B[vi] = T.Cast("float16", A[vi])
+
+    sch = tvm.s_tir.Schedule(func)
+    block = sch.get_sblock("B")
+    loops = sch.get_loops(block)
+    bx, tx = sch.split(loops[0], factors=[None, 32])
+    sch.bind(bx, "blockIdx.x")
+    sch.bind(tx, "threadIdx.x")
+
+    target = "cuda"
+    fadd = tvm.compile(sch.mod, target=target)
+
+    # float4_e2m1fn: 4-bit values 0..15, two packed per byte.
+    # Encoding (sign | exp1 | man1 man0):
+    #   0→0.0  1→0.5  2→1.0  3→1.5  4→2.0  5→3.0  6→4.0  7→6.0
+    #   8→-0.0 9→-0.5 10→-1.0 … 15→-6.0
+    fp4_to_fp16 = np.array(
+        [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
+        dtype=np.float16,
+    )
+
+    # Pack DIFFERENT fp4 values in low/high nibbles so the test verifies
+    # both byte offset (/2) AND correct nibble extraction (% 2 shift).
+    fp4_elements = np.array([i % 16 for i in range(n)], dtype=np.uint8)
+    packed = np.zeros(n // 2, dtype=np.uint8)
+    for i in range(0, n, 2):
+        packed[i // 2] = fp4_elements[i] | (fp4_elements[i + 1] << 4)
+
+    expected = fp4_to_fp16[fp4_elements]
+
+    def run_and_check():
+        dev = tvm.device(target, 0)
+        a = tvm.runtime.empty(shape=(n // 2,), dtype="uint8", device=dev)
+        a.copyfrom(packed)
+        b = tvm.runtime.empty(shape=(n,), dtype="float16", device=dev)
+        fadd(a, b)
+        tvm.testing.assert_allclose(b.numpy(), expected)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_device.py

@@ -0,0 +1,119 @@
+# 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 numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
+def test_large_uint_imm():
+    value = (1 << 63) + 123
+    value_const = tvm.tirx.const(value, "uint64")
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((12,), "uint64")):
+            T.func_attr({"tirx.noalias": True})
+            for i0_0 in T.thread_binding(6, thread="blockIdx.x"):
+                for i0_1 in T.thread_binding(2, thread="threadIdx.x"):
+                    with T.sblock("A"):
+                        v_i0 = T.axis.spatial(12, i0_0 * 2 + i0_1)
+                        T.reads()
+                        T.writes(A[v_i0])
+                        A[v_i0] = value_const + T.uint64(3)
+
+    def check_target(target):
+        target_kind = target["kind"] if isinstance(target, dict) else target
+        if not tvm.testing.device_enabled(target_kind):
+            return
+        f = tvm.compile(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target_kind, 0)
+            a = tvm.runtime.empty((12,), dtype="uint64", device=dev)
+            f(a)
+            assert a.numpy()[0] == value + 3
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_target("cuda")
+    check_target({"kind": "vulkan", "from_device": 0})
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
+def test_add_pipeline():
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, B: T.Buffer((), "float32"), var_D: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            n = T.int32()
+            A = T.match_buffer(var_A, (n,))
+            D = T.match_buffer(var_D, (n,))
+            C = T.sblock_alloc_buffer((n,))
+            for i0_0 in T.thread_binding((n + 255) // 256, thread="blockIdx.x"):
+                for i0_1 in T.thread_binding(256, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i0 = T.axis.spatial(n, i0_0 * 256 + i0_1)
+                        T.where(i0_0 * 256 + i0_1 < n)
+                        T.reads(A[v_i0], B[()])
+                        T.writes(C[v_i0])
+                        C[v_i0] = A[v_i0] + B[()]
+            for i0_0 in T.thread_binding((n + 255) // 256, thread="blockIdx.x"):
+                for i0_1 in T.thread_binding(256, thread="threadIdx.x"):
+                    with T.sblock("D"):
+                        v_i0 = T.axis.spatial(n, i0_0 * 256 + i0_1)
+                        T.where(i0_0 * 256 + i0_1 < n)
+                        T.reads(C[v_i0])
+                        T.writes(D[v_i0])
+                        D[v_i0] = C[v_i0] + T.float32(1.0)
+
+    def check_target(device, host):
+        if not tvm.testing.device_enabled(device) or not tvm.testing.device_enabled(host):
+            return
+        target = tvm.target.Target(device, host)
+        mhost = tvm.tirx.build(Module, target=target)
+        f = mhost.main
+        n = 1027
+
+        def run_and_check():
+            dev = tvm.device(device, 0)
+            a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+            b = tvm.runtime.tensor(np.random.uniform(size=()).astype("float32"), dev)
+            d = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+            f(a, b, d)
+            tvm.testing.assert_allclose(d.numpy(), a.numpy() + b.numpy() + 1)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_target("cuda", host="llvm")
+    # check_target("nvptx", host="llvm")  # nvptx kernel entry-point lookup not wired here
+    check_target("vulkan", host="llvm")
+    check_target("rocm", host="llvm")
+
+
+if __name__ == "__main__":
+    test_large_uint_imm()
+    test_add_pipeline()

tests/python/codegen/test_target_codegen_extern.py

@@ -0,0 +1,109 @@
+# 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 ir as I
+from tvm.script import tirx as T
+
+
+@pytest.mark.gpu
+def test_add_pipeline():
+    """Test extern-style add pipeline with vectorized operations."""
+    nn = 64
+    max_threads = 4
+
+    # CPU version: serial loop with vectorized operations
+    @I.ir_module
+    class ModuleCPU:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((64,), "float32"), C: T.Buffer((64,), "float32")):
+            for i in T.serial((64 + 1) // 2):
+                C[T.Ramp(i * 2, 1, 2)] = A[T.Ramp(i * 2, 1, 2)] + T.Broadcast(T.float32(1), 2)
+
+    # GPU version: thread bindings with vectorized operations
+    @I.ir_module
+    class ModuleGPU:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((64,), "float32"), C: T.Buffer((64,), "float32")):
+            bx = T.launch_thread("blockIdx.x", (64 + 4 - 1) // 4)
+            tx = T.launch_thread("threadIdx.x", 4)
+            idx = bx * 4 + tx
+            if T.likely(idx < 64):
+                C[T.Ramp(idx * 2, 1, 2)] = A[T.Ramp(idx * 2, 1, 2)] + T.Broadcast(T.float32(1), 2)
+
+    def check_target(target):
+        if not tvm.testing.device_enabled(target):
+            return
+        mod = ModuleGPU if target in ["opencl", "cuda"] else ModuleCPU
+        # build and invoke the kernel.
+        f = tvm.compile(mod, target=target)
+        n = nn
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+            c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+            f(a, c)
+            tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)
+
+        if target == "llvm":
+            run_and_check()
+        else:
+            tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_target("llvm")
+    check_target("opencl")
+    check_target("cuda")
+
+
+def test_pack_buffer_simple():
+    """Test call_packed with buffer arguments."""
+    nn = 1024
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1024,), "float32"), C: T.Buffer((1024,), "float32")):
+            T.evaluate(T.call_packed("my_extern_array_func1", A, C))
+
+    @tvm.register_global_func
+    def my_extern_array_func1(aa, bb):
+        aa.copyto(bb)
+
+    def check_target(target):
+        if not tvm.testing.device_enabled(target):
+            return
+        # build and invoke the kernel.
+        f = tvm.compile(Module, target=target)
+        dev = tvm.cpu(0)
+        # launch the kernel.
+        n = nn
+        a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
+        c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+
+        f(a, c)
+        tvm.testing.assert_allclose(c.numpy(), a.numpy())
+
+    check_target("llvm")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_gpu_common.py

@@ -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.
+from functools import partial
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+        pytest.param({"kind": "vulkan", "supports_int64": True}, marks=pytest.mark.gpu),
+        pytest.param("opencl", marks=pytest.mark.gpu),
+    ],
+)
+@pytest.mark.parametrize("dtype", ["int32", "uint32", "int64", "uint64"])
+def test_int_intrin(target, dtype):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    test_funcs = [
+        (T.clz, lambda x, dtype: int(dtype[-2:]) - (len(bin(x)) - 2)),
+    ]
+
+    for tvm_intrin, np_func in test_funcs:
+        n = 128
+
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(
+                A: T.Buffer((n,), dtype),
+                B: T.Buffer((n,), dtype),
+            ):
+                T.func_attr({"tirx.noalias": True})
+                for i0 in T.thread_binding(n, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i0 = T.axis.spatial(n, i0)
+                        T.reads(A[v_i0])
+                        T.writes(B[v_i0])
+                        B[v_i0] = tvm_intrin(A[v_i0])
+
+        f = tvm.compile(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target["kind"] if isinstance(target, dict) else target)
+            a = tvm.runtime.tensor(np.random.randint(0, 100000, size=n).astype(dtype), dev)
+            b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev)
+            f(a, b)
+            ref = np.vectorize(partial(np_func, dtype=dtype))(a.numpy())
+            tvm.testing.assert_allclose(b.numpy(), ref)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_hexagon.py

@@ -0,0 +1,122 @@
+# 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 re
+
+import pytest
+
+import tvm
+import tvm.contrib.hexagon as hexagon
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.fixture(autouse=True)
+def register_linker():
+    original_linker = hexagon.hexagon_link()
+    # Register a phony linker, so that we can test codegen without a Hexagon toolchain.
+    hexagon.register_linker(lambda: "/bin/true")
+    yield None
+    # Restore registration.
+    hexagon.register_linker(original_linker)
+
+
+@pytest.mark.skipif(not env.has_hexagon(), reason="need hexagon")
+def test_basic():
+    target = tvm.target.Target("qcom/hexagon-v66")
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            C: T.Buffer((128,), "uint8"),
+            A: T.Buffer((128,), "uint8"),
+            A_1: T.Buffer((128,), "uint8"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i in range(128):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(128, i)
+                    T.reads(A[v_i], A_1[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = A[v_i] + A_1[v_i]
+
+    hexm = tvm.compile(Module, target=tvm.target.Target(target, target))
+    asm = hexm.inspect_source("s")
+    vadds = re.findall(r"v[0-9]+.b = vadd\(v[0-9]+.b,v[0-9]+.b\)", asm)
+    assert vadds  # Check that it's non-empty
+
+
+@pytest.mark.skipif(not env.has_hexagon(), reason="need hexagon")
+def test_llvm_target_features():
+    target = tvm.target.Target("qcom/hexagon-v66")
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def add_one(C: T.Buffer((128,), "int32"), A: T.Buffer((128,), "uint8")):
+            T.func_attr({"tirx.noalias": True})
+            for i in range(128):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(128, i)
+                    T.reads(A[v_i])
+                    T.writes(C[v_i])
+                    C[v_i] = T.Cast("int32", A[v_i]) + 1
+
+    m = tvm.compile(Module, target=tvm.target.Target(target, target))
+    llvm_ir = m.inspect_source("ll")
+    # Make sure we find +hvx-length128b in "attributes".
+    fs = re.findall(r"attributes.*\+hvx-length128b", llvm_ir)
+    assert fs  # Check that it's non-empty
+
+
+@pytest.mark.skipif(not env.has_hexagon(), reason="need hexagon")
+def test_llvm_options():
+    target = tvm.target.Target(
+        {
+            "kind": "hexagon",
+            "mtriple": "hexagon",
+            "mcpu": "hexagonv66",
+            "mattr": ["+hvxv66", "+hvx-length128b"],
+            "num-cores": 4,
+            "vtcm-capacity": 262144,
+            "llvm-options": ["-hexagon-noopt"],
+        }
+    )
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(compute: T.Buffer((10,), "int32")):
+            T.func_attr({"tirx.noalias": True})
+            for _ in range(10):
+                with T.sblock("compute"):
+                    v__ = T.axis.spatial(10, _)
+                    T.reads()
+                    T.writes(compute[v__])
+                    compute[v__] = 0
+
+    # Check that BuildHexagon hasn't crashed because of target attribute
+    # type mismatch.
+    tvm.compile(Module, target=tvm.target.Target(target, target))
+    assert re.search("-hexagon-noopt", str(target))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_llvm_vla.py

@@ -0,0 +1,212 @@
+# 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.
+
+"""
+Codegen tests for VLA extensions
+"""
+
+import re
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.target.codegen import llvm_version_major
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_codegen_vscale(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    vscale = tvm.tirx.vscale()
+
+    @T.prim_func(s_tir=True)
+    def main(A: T.Buffer((5,), "int32")):
+        for i in range(5):
+            A[i] = 2 * vscale
+
+    with tvm.target.Target(target):
+        build_mod = tvm.tirx.build(main)
+
+    llvm = build_mod.inspect_source()
+    assert re.findall(r"llvm.vscale.i32", llvm), "No vscale in generated LLVM."
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_scalable_buffer_load_store(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def my_func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (128,), "float32")
+        B = T.match_buffer(b, (128,), "float32")
+        T.func_attr({"global_symbol": "my_module", "tirx.noalias": True})
+        B[T.ramp(0, 1, 4 * T.vscale())] = A[T.ramp(0, 1, 4 * T.vscale())]
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.build(my_func)
+
+    llvm = mod.inspect_source("ll")
+    assert re.findall(r"load <vscale x 4 x float>", llvm), "No scalable load in generated LLVM."
+    assert re.findall(r" store <vscale x 4 x float>", llvm), "No scalable store in generated LLVM."
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_scalable_broadcast(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def my_func(a: T.handle):
+        A = T.match_buffer(a, (128,), "float32")
+        T.func_attr({"global_symbol": "my_module", "tirx.noalias": True})
+        A[T.ramp(0, 1, 4 * T.vscale())] = T.broadcast(1, 4 * T.vscale())
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.build(my_func)
+
+    llvm = mod.inspect_source("ll")
+    # Older LLVM versions print the broadcast as a shufflevector of an insertelement,
+    # newer ones print it as a splat constant.
+    assert (
+        "shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float>" in llvm
+        or "store <vscale x 4 x float> splat (float 1.000000e+00)" in llvm
+    ), "No scalable broadcast in generated LLVM."
+    assert re.findall(r" store <vscale x 4 x float>", llvm), "No scalable store in generated LLVM."
+
+
+@pytest.mark.skip(
+    reason="Vscale and get.active.lane.mask are not supported in earlier versions of LLVM",
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_get_active_lane_mask(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle):
+        A = T.match_buffer(a, (30,), "int1")
+        for i in range(T.ceildiv(30, T.vscale() * 4)):
+            A[i : i + T.vscale() * 4] = T.get_active_lane_mask("uint1xvscalex4", i, 30)
+
+    with tvm.target.Target(target):
+        out = tvm.tirx.build(before)
+
+    ll = out.inspect_source("ll")
+    assert "get.active.lane.mask" in ll
+
+
+@pytest.mark.skip(
+    reason="Vscale and get.active.lane.mask are not supported in earlier versions of LLVM",
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_predicated_scalable_buffer(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in T.serial(T.ceildiv(16, 4 * T.vscale())):
+            for i_1 in T.vectorized(4 * T.vscale()):
+                if i_0 * 4 * T.vscale() + i_1 < 14:
+                    B[i_0 * 4 * T.vscale() + i_1] = A[i_0 * 4 * T.vscale() + i_1] + 1.0
+
+    with tvm.target.Target(target):
+        out = tvm.tirx.build(before)
+
+    ll = out.inspect_source("ll")
+    assert "get.active.lane.mask" in ll
+    assert "llvm.masked.load" in ll
+    assert "llvm.masked.store" in ll
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_metal.py

@@ -0,0 +1,353 @@
+# 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 numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_metal_inf_nan():
+    target = "metal"
+
+    def check_inf_nan(n, value, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(
+                A: T.Buffer((1,), dtype),
+                C: T.Buffer((1,), dtype),
+            ):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(1, i)
+                        T.reads()
+                        T.writes(C[v_i])
+                        C[v_i] = T.Cast(dtype, value)
+
+        fun = tvm.compile(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_inf_nan(1, -float("inf"), "float32")
+    check_inf_nan(1, -float("inf"), "float16")
+    check_inf_nan(1, float("inf"), "float32")
+    check_inf_nan(1, float("inf"), "float16")
+    check_inf_nan(1, float("nan"), "float32")
+    check_inf_nan(1, float("nan"), "float16")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_unaligned_vectorize():
+    @tvm.script.ir_module
+    class IRModule:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((2, 3), "float32"), B: T.Buffer((6,), "float32")):
+            T.func_attr({"global_symbol": "main"})
+            for i0_1 in T.thread_binding(3, thread="threadIdx.x"):
+                for i0_0 in T.vectorized(2):
+                    with T.sblock("block"):
+                        vi0 = T.axis.spatial(6, i0_0 * 3 + i0_1)
+                        B[vi0] = A[vi0 // 3, vi0 % 3]
+
+    target = "metal"
+    a = (np.arange(6).reshape(2, 3)).astype("float32")
+    f = tvm.compile(IRModule, target=target)
+
+    def run_and_check():
+        dev = tvm.metal()
+        a_nd = tvm.runtime.tensor(a, dev)
+        b_nd = tvm.runtime.empty((6,), "float32", dev)
+        f(a_nd, b_nd)
+        tvm.testing.assert_allclose(b_nd.numpy(), a.reshape(6), atol=1e-5, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_metal_erf():
+    target = "metal"
+
+    def check_erf(n, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(
+                A: T.Buffer((1,), dtype),
+                C: T.Buffer((1,), dtype),
+            ):
+                T.func_attr({"tirx.noalias": True})
+                for i0 in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i0 = T.axis.spatial(1, i0)
+                        T.reads(A[v_i0])
+                        T.writes(C[v_i0])
+                        C[v_i0] = T.erf(A[v_i0])
+
+        fun = tvm.compile(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_erf(1, "float32")
+    check_erf(1, "float16")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_ramp():
+    target = "metal"
+
+    @tvm.script.ir_module
+    class IRModule:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1, 2), "int32")):
+            T.func_attr({"global_symbol": "main"})
+            for i in T.thread_binding(1, thread="threadIdx.x"):
+                with T.sblock("block"):
+                    tx = T.axis.spatial(1, i)
+                    r = T.ramp(tx, 3, 2)
+                    A[0, T.ramp(0, 1, 2)] = r
+
+    f = tvm.compile(IRModule, target=target)
+
+    def run_and_check():
+        dev = tvm.metal()
+        a_nd = tvm.runtime.empty((1, 2), "int32", dev)
+        f(a_nd)
+        assert tuple(a_nd.numpy()[0, :]) == (0, 3)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_select_vectorize():
+    @tvm.script.ir_module
+    class IRModule:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((6), "float32"), B: T.Buffer((6,), "float32")):
+            T.func_attr({"global_symbol": "main"})
+            for i0_1 in T.thread_binding(3, thread="threadIdx.x"):
+                for i0_0 in T.vectorized(2):
+                    with T.sblock("block"):
+                        vi0 = T.axis.spatial(6, i0_0 * 3 + i0_1)
+                        B[vi0] = T.Select((vi0 % 2) == 0, A[vi0], T.float32(0))
+
+    target = "metal"
+    a = np.arange(6).astype("float32")
+    f = tvm.compile(IRModule, target=target)
+    a.reshape(3, 2)[:, 1] = 0
+
+    def run_and_check():
+        dev = tvm.metal()
+        a_nd = tvm.runtime.tensor(a, dev)
+        b_nd = tvm.runtime.empty((6,), "float32", dev)
+        f(a_nd, b_nd)
+        tvm.testing.assert_allclose(b_nd.numpy(), a, atol=1e-5, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_vectorized_uint8():
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((16), "uint8"), B: T.Buffer((16), "float32")):
+        for i in T.thread_binding(4, thread="threadIdx.x"):
+            for j in T.vectorized(4):
+                with T.sblock("block"):
+                    vi = T.axis.spatial(16, i * 4 + j)
+                    B[vi] = T.Cast("float32", A[vi])
+
+    a = np.arange(16).astype("uint8")
+    f = tvm.compile(func, target="metal")
+
+    def run_and_check():
+        dev = tvm.metal()
+        a_nd = tvm.runtime.tensor(a, dev)
+        b_nd = tvm.runtime.empty((16,), "float32", dev)
+        f(a_nd, b_nd)
+        tvm.testing.assert_allclose(b_nd.numpy(), a.astype("float32"), atol=1e-5, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_func_with_trailing_pod_params():
+    from tvm.support import xcode  # pylint: disable=import-outside-toplevel
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((16), "float32"), B: T.Buffer((16), "float32"), x: T.float32):
+        for i in T.thread_binding(16, thread="threadIdx.x"):
+            with T.sblock("block"):
+                vi = T.axis.spatial(16, i)
+                B[vi] = A[vi] + x
+
+    @tvm.register_global_func("tvm_callback_metal_compile")
+    def compile_metal(src, target):
+        return xcode.compile_metal(src)
+
+    mod = tvm.IRModule({"main": func})
+
+    f = tvm.tirx.build(mod, target="metal")
+    src: str = f.imports[0].inspect_source()
+    occurrences = src.count("struct func_kernel_args_t")
+    assert occurrences == 1, occurrences
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_metal_compile_callback_source_passthrough():
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+
+    seen = {}
+
+    def inspect_callback(src, target):
+        # Pure inspection callback: capture the source, return it untouched and
+        # declare it is still textual MSL so it is compiled at load time.
+        seen["src"] = src
+        return (src, "metal")
+
+    tvm.register_global_func("tvm_callback_metal_compile", inspect_callback, override=True)
+    try:
+        f = tvm.compile(Module, target="metal")
+        dev = tvm.metal()
+        a = np.random.rand(n).astype("float32")
+        a_nd = tvm.runtime.tensor(a, dev)
+        b_nd = tvm.runtime.empty((n,), "float32", dev)
+        f(a_nd, b_nd)
+        dev.sync()
+    finally:
+        tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile")
+
+    assert "src" in seen and len(seen["src"]) > 0
+    tvm.testing.assert_allclose(b_nd.numpy(), a + 1.0, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_metal_compile_callback_mixed_formats_rejected():
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((n,), "float32"),
+            B: T.Buffer((n,), "float32"),
+            C: T.Buffer((n,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            # Two independent thread-bound regions -> two device kernels, so the
+            # compile callback is invoked twice within one module.
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+            for j_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for j_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_j = T.axis.spatial(n, j_0 * 32 + j_1)
+                        T.reads(A[v_j])
+                        T.writes(C[v_j])
+                        C[v_j] = A[v_j] + 2.0
+
+    calls = {"n": 0}
+
+    def mixed_callback(src, target):
+        calls["n"] += 1
+        if calls["n"] == 1:
+            # Treated as a compiled metallib payload.
+            return src
+        # Second kernel declares textual MSL, contradicting the metallib above.
+        return (src, "metal")
+
+    tvm.register_global_func("tvm_callback_metal_compile", mixed_callback, override=True)
+    try:
+        with pytest.raises(Exception, match="inconsistent formats"):
+            tvm.compile(Module, target="metal")
+    finally:
+        tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_metal(), reason="need metal")
+def test_export_load_with_fallback(monkeypatch, tmp_path):
+    """Force the codegen wrapper into the fallback branch, then export."""
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+
+    monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+    host_lib = tvm.compile(Module, target="metal")
+    monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
+
+    lib_path = str(tmp_path / "lib.so")
+    host_lib.export_library(lib_path)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_opencl.py

@@ -0,0 +1,334 @@
+# 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: E501
+import re
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+target = "opencl"
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+def test_opencl_ternary_expression():
+    def check_if_then_else(n, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(1, i)
+                        T.reads(A[0])
+                        T.writes(C[v_i])
+                        C[v_i] = T.max(
+                            T.Cast(dtype, 2),
+                            T.if_then_else(
+                                0 < T.Cast("int32", A[0]),
+                                T.Cast(dtype, 1),
+                                T.Cast(dtype, 3),
+                            ),
+                        )
+
+        fun = tvm.tirx.build(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    def check_select(n, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(1, i)
+                        T.reads(A[0])
+                        T.writes(C[v_i])
+                        C[v_i] = T.max(
+                            T.Cast(dtype, 2),
+                            T.Select(
+                                0 < T.Cast("int32", A[0]),
+                                T.Cast(dtype, 1),
+                                T.Cast(dtype, 3),
+                            ),
+                        )
+
+        fun = tvm.tirx.build(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_if_then_else(1, "int8")
+    check_if_then_else(1, "uint8")
+    check_if_then_else(1, "int16")
+    check_if_then_else(1, "uint16")
+    check_select(1, "int8")
+    check_select(1, "uint8")
+    check_select(1, "int16")
+    check_select(1, "uint16")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+def test_opencl_inf_nan():
+    def check_inf_nan(n, value, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(1, i)
+                        T.reads()
+                        T.writes(C[v_i])
+                        C[v_i] = T.Cast(dtype, value)
+
+        fun = tvm.tirx.build(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_inf_nan(1, -float("inf"), "float32")
+    check_inf_nan(1, -float("inf"), "float64")
+    check_inf_nan(1, float("inf"), "float32")
+    check_inf_nan(1, float("inf"), "float64")
+    check_inf_nan(1, float("nan"), "float32")
+    check_inf_nan(1, float("nan"), "float64")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+def test_opencl_max():
+    def check_max(n, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(1, i)
+                        T.reads(A[0])
+                        T.writes(C[v_i])
+                        C[v_i] = T.max(A[0] + T.Cast(dtype, 1), T.Cast(dtype, 0))
+
+        fun = tvm.tirx.build(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_max(1, "int8")
+    check_max(1, "uint8")
+    check_max(1, "int16")
+    check_max(1, "uint16")
+    check_max(1, "float32")
+    check_max(1, "float64")
+
+
+def test_opencl_erf():
+    def check_erf(n, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i0 in T.thread_binding(1, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i0 = T.axis.spatial(1, i0)
+                        T.reads(A[v_i0])
+                        T.writes(C[v_i0])
+                        C[v_i0] = T.erf(A[v_i0])
+
+        fun = tvm.tirx.build(Module, target=target)
+
+        source_str = fun.imports[0].inspect_source()
+        matches = re.findall("erf", source_str)
+        error_matches = re.findall("erff", source_str)
+        assert len(matches) == 1 and len(error_matches) == 0
+
+    check_erf(1, "float32")
+    check_erf(1, "float64")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+def test_opencl_type_casting():
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(C: T.Buffer((32,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(8, thread="threadIdx.x"):
+                for i_1 in T.vectorized(4):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(32, i_0 * 4 + i_1)
+                        T.reads()
+                        T.writes(C[v_i])
+                        C[v_i] = T.Select(
+                            v_i // 4 == 3 and v_i % 3 == 1, T.float32(1.0), T.float32(0.0)
+                        )
+
+    def check_type_casting(n, dtype):
+        fun = tvm.tirx.build(Module, target=target)
+        assembly = fun.imports[0].inspect_source()
+        lcond = "convert_int4(((convert_uint4(((uint4)(((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3)))))"
+        rcond = "(convert_uint4(((((int4)(((convert_int(get_local_id(0))))+(1*0), ((convert_int(get_local_id(0))))+(1*1), ((convert_int(get_local_id(0))))+(1*2), ((convert_int(get_local_id(0))))+(1*3))) % ((int4)(3, 3, 3, 3))) == ((int4)(1, 1, 1, 1))))))))"
+        pattern_cond = f"({lcond} && {rcond})"
+        assert assembly.count(pattern_cond) != 0
+
+        def run_and_check():
+            dev = tvm.device(target, 0)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_type_casting(32, "float32")
+    # fp16 is not yet supported in ci
+    # check_type_casting(dev, 16, "float16")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("opencl", marks=pytest.mark.gpu),
+        pytest.param({"kind": "opencl", "device": "adreno"}, marks=pytest.mark.gpu),
+    ],
+)
+def test_opencl_ceil_log2(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    def _check(target, n, dtype):
+        target_obj = tvm.target.Target(target)
+        is_adreno = "adreno" in target_obj.attrs.get("device", "")
+        inter_dtype = "float32" if is_adreno else "float64"
+
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(C: T.Buffer((n,), "int32")):
+                T.func_attr({"tirx.noalias": True})
+                for i in T.thread_binding(n, thread="threadIdx.x"):
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(n, i)
+                        T.reads()
+                        T.writes(C[v_i])
+                        C[v_i] = T.Cast("int32", T.ceil(T.log2(T.Cast(inter_dtype, v_i))))
+
+        fun = tvm.tirx.build(Module, target=target)
+        assembly = fun.imports[0].inspect_source()
+        if is_adreno:
+            pattern = "convert_float"
+        else:
+            pattern = "convert_double"
+        assert assembly.count(pattern) != 0
+
+    _check(target, 32, "float32")
+
+
+def _get_maximum_kernel_args(source):
+    def get_kernel_args(source):
+        import re
+
+        p = re.tirx.build(r"__kernel void .+\((.*)\)")
+        args = p.findall(source)
+        return args
+
+    args = get_kernel_args(source)
+    max_args = len(args[0].split(","))
+    for arg_line in args:
+        max_args = max(max_args, len(arg_line.split(",")))
+    return max_args
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_opencl(), reason="need opencl")
+def test_export_load_with_fallback(monkeypatch, tmp_path):
+    """Force the codegen wrapper into the fallback branch, then export+load+run."""
+    import numpy as np
+
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+
+    monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+    host_lib = tvm.compile(Module, target=target)
+    monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
+
+    lib_path = str(tmp_path / "lib.so")
+    host_lib.export_library(lib_path)
+    reloaded = tvm.runtime.load_module(lib_path)
+
+    a_np = np.random.uniform(size=(n,)).astype("float32")
+    b_np = np.zeros((n,), dtype="float32")
+
+    def run_and_check():
+        dev = tvm.device(target, 0)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        reloaded["main"](a, b)
+        np.testing.assert_allclose(b.numpy(), a_np + 1.0, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_riscv.py

@@ -0,0 +1,199 @@
+# 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: E501, F841
+
+import re
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.target.codegen import target_has_features
+from tvm.testing import env
+
+
+@pytest.mark.skipif(not env.has_llvm_min_version(14), reason="need llvm >= 14")
+@pytest.mark.parametrize(
+    "target",
+    [
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv32-linux-gnu",
+            "mcpu": "generic-rv32",
+            "mattr": ["+i", "+m"],
+        },
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv32-linux-gnu",
+            "mcpu": "generic-rv32",
+            "mattr": ["+i", "+m", "+v"],
+        },
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m"],
+        },
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_rvv(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    def check_rvv_presence(N, extent):
+        @T.prim_func(s_tir=True)
+        def load_vec(A: T.Buffer((N,), "int8")):
+            for j in T.vectorized(0, extent):
+                A[j] = 1
+
+        f = tvm.tirx.build(load_vec, target)
+        # Check RVV `vsetvli` prensence
+        assembly = f.inspect_source("asm")
+        if target_has_features("v"):
+            assert "vsetvli" in assembly
+        else:
+            assert "vsetvli" not in assembly
+
+    with tvm.target.Target(target):
+        check_rvv_presence(16, 32)
+
+
+@pytest.mark.skipif(not env.has_llvm_min_version(14), reason="need llvm >= 14")
+@pytest.mark.parametrize(
+    "target",
+    [
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv32-linux-gnu",
+            "mcpu": "generic-rv32",
+            "mattr": ["+i", "+m", "+v"],
+        },
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_rvv_vscale_llvm_dbginfo(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def rvv_with_vscale(A_handle: T.handle, B_handle: T.handle, C_handle: T.handle):
+        A = T.match_buffer(A_handle, (8,), dtype="float32", align=4, offset_factor=1)
+        B = T.match_buffer(B_handle, (4, 8), dtype="float32", align=4, offset_factor=1, strides=[8, 1])
+        C = T.match_buffer(C_handle, (4,), dtype="float32", align=4, offset_factor=1)
+        with T.sblock("root"):
+            T.reads(A[0:8], B[0:4, 0:8])
+            zero = T.call_llvm_intrin("float32xvscalex2", "llvm.riscv.vfmv.v.f", T.Broadcast(T.float32(0.0), T.vscale() * 2), C[0], T.uint64(1))
+            vec_A = T.call_llvm_intrin("float32xvscalex4", "llvm.riscv.vle", T.Broadcast(T.float32(0.0), T.vscale() * 4), T.tvm_access_ptr(T.type_annotation("float32"), A.data, 0, 8, 1), T.int64(8))
+            vec_B = T.call_llvm_intrin("float32xvscalex4", "llvm.riscv.vle", T.Broadcast(T.float32(0.0), T.vscale() * 4), T.tvm_access_ptr(T.type_annotation("float32"), B.data, 0 * 8, 8, 1), T.int64(8))
+            prod = T.call_llvm_intrin("float32xvscalex4", "llvm.riscv.vfmul", T.Broadcast(T.float32(0.0), T.vscale() * 4), vec_A, vec_B, T.uint64(7), T.uint64(8))
+            redsum = T.call_llvm_intrin("float32xvscalex2", "llvm.riscv.vfredusum", T.Broadcast(T.float32(0.0), T.vscale() * 2), prod, zero, T.uint64(7), T.uint64(8))
+    # fmt: on
+
+    # tvm.error.InternalError: Can't fetch the lanes of a scalable vector at a compile time.
+    with tvm.target.Target(target):
+        f = tvm.tirx.build(rvv_with_vscale, target)
+
+
+@pytest.mark.skipif(not env.has_llvm_min_version(14), reason="need llvm >= 14")
+def test_rvv_fixed_width_vectorized_loop_uses_scalable_chunks():
+    @T.prim_func(s_tir=True)
+    def fixed16_negative(
+        A: T.Buffer((14, 23, 67, 99), "float32"),
+        B: T.Buffer((14, 23, 67, 99), "float32"),
+    ):
+        for n, c, h, wo in T.grid(14, 23, 67, 7):
+            for wi in T.vectorized(0, 16):
+                if wo * 16 + wi < 99:
+                    B[n, c, h, wo * 16 + wi] = T.float32(0) - A[n, c, h, wo * 16 + wi]
+
+    @T.prim_func(s_tir=True)
+    def fixed16_negative_int64(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
+        for wi in T.vectorized(T.int64(0), T.int64(16)):
+            B[wi] = T.float32(0) - A[wi]
+
+    target = tvm.target.Target(
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        }
+    )
+
+    def check_codegen(func):
+        with target:
+            f = tvm.tirx.build(func, target)
+
+        assembly = f.inspect_source("asm")
+        assert "vle32.v" in assembly
+        assert "vse32.v" in assembly
+        assert not re.search(r"\bflw\b", assembly)
+        assert not re.search(r"\bfsub\.s\b", assembly)
+        assert not re.search(r"\bfsw\b", assembly)
+
+    check_codegen(fixed16_negative)
+    check_codegen(fixed16_negative_int64)
+
+
+@pytest.mark.skipif(not env.has_llvm_min_version(14), reason="need llvm >= 14")
+def test_rvv_scalable_ramp_expression():
+    @T.prim_func(s_tir=True)
+    def ramp_compare(B: T.Buffer((16,), "int32")):
+        for i in T.vectorized(16):
+            B[i] = T.Select(i * 3 + 5 < 29, i * 3 + 5, -1)
+
+    target = tvm.target.Target(
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        }
+    )
+
+    with target:
+        f = tvm.tirx.build(ramp_compare, target)
+
+    assembly = f.inspect_source("asm")
+    assert "vid.v" in assembly
+    assert re.search(r"\bvmul\.v", assembly)
+    assert re.search(r"\bvadd\.v", assembly)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_rocm.py

@@ -0,0 +1,221 @@
+# 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 ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_rocm_inf_nan():
+    def check_inf_nan(n, value, dtype):
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((1,), dtype), C: T.Buffer((1,), dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i_0 in T.thread_binding(1, thread="blockIdx.x"):
+                    for i_1 in T.thread_binding(128, thread="threadIdx.x"):
+                        with T.sblock("C"):
+                            v_i = T.axis.spatial(1, i_0 * 128 + i_1)
+                            T.where(i_0 * 128 + i_1 < 1)
+                            T.reads()
+                            T.writes(C[v_i])
+                            C[v_i] = T.Cast(dtype, value)
+
+        fun = tvm.compile(Module, "rocm")
+
+        def run_and_check():
+            dev = tvm.rocm(0)
+            a = tvm.runtime.empty((n,), dtype, dev)
+            c = tvm.runtime.empty((n,), dtype, dev)
+            fun(a, c)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_inf_nan(1, -float("inf"), "float32")
+    check_inf_nan(1, -float("inf"), "float64")
+    check_inf_nan(1, float("inf"), "float32")
+    check_inf_nan(1, float("inf"), "float64")
+    check_inf_nan(1, float("nan"), "float32")
+    check_inf_nan(1, float("nan"), "float64")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_rocm_copy():
+    def check_rocm(dtype, n):
+        def run_and_check():
+            dev = tvm.rocm(0)
+            a_np = np.random.uniform(size=(n,)).astype(dtype)
+            a = tvm.runtime.empty((n,), dtype, dev).copyfrom(a_np)
+            b_np = a.numpy()
+            tvm.testing.assert_allclose(a_np, b_np)
+            tvm.testing.assert_allclose(a_np, a.numpy())
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    for _ in range(100):
+        dtype = np.random.choice(["float32", "float16", "int8", "int32"])
+        logN = np.random.randint(1, 15)
+        peturb = np.random.uniform(low=0.5, high=1.5)
+        check_rocm(dtype, int(peturb * (2**logN)))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_rocm_vectorize_add():
+    def check_rocm(dtype, n, lanes):
+        vec_dtype = f"{dtype}x{lanes}"
+        num_blocks = n // 4
+
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(A: T.Buffer((n,), vec_dtype), B: T.Buffer((n,), vec_dtype)):
+                T.func_attr({"tirx.noalias": True})
+                for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
+                    for i_1 in T.thread_binding(4, thread="threadIdx.x"):
+                        with T.sblock("B"):
+                            v_i = T.axis.spatial(n, i_0 * 4 + i_1)
+                            T.reads(A[v_i])
+                            T.writes(B[v_i])
+                            B[v_i] = A[v_i] + T.Broadcast(T.Cast(dtype, 1), lanes)
+
+        fun = tvm.compile(Module, target="rocm")
+
+        def run_and_check():
+            dev = tvm.rocm(0)
+            a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np.random.uniform(size=(n, lanes)))
+            c = tvm.runtime.empty((n,), vec_dtype, dev)
+            fun(a, c)
+            tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    check_rocm("float32", 64, 2)
+    check_rocm("float16", 64, 2)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_rocm_warp_shuffle():
+    @T.prim_func(s_tir=True)
+    def func(
+        A_handle: T.handle,
+    ):
+        A = T.match_buffer(A_handle, (32,), dtype="float32")
+
+        for bx in T.thread_binding(1, thread="blockIdx.x"):
+            for tx in T.thread_binding(32, thread="threadIdx.x"):
+                with T.sblock("test"):
+                    A_local = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                    mask = T.sblock_alloc_buffer((1,), "uint32", scope="local")
+                    t0 = T.sblock_alloc_buffer((1,), "float32", scope="local")
+
+                    A_local[0] = A[tx]
+                    A_local[0] = T.tvm_warp_shuffle(mask[0], A_local[0], 0, 32, 32)
+                    A[tx] = A_local[0]
+
+    mod = tvm.compile(func, target="rocm")
+
+    def run_and_check():
+        dev = tvm.rocm(0)
+        a = tvm.runtime.tensor(np.random.uniform(size=(32,)).astype("float32"), dev)
+        mod(a)
+        tvm.testing.assert_allclose(a.numpy(), np.ones((32,)) * a.numpy()[0])
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_rocm_vectorized_exp():
+    @T.prim_func(s_tir=True)
+    def func(
+        A_handle: T.handle,
+        B_handle: T.handle,
+    ):
+        A = T.match_buffer(A_handle, (4,), dtype="float32")
+        B = T.match_buffer(B_handle, (4,), dtype="float32")
+
+        for bx in T.thread_binding(1, thread="blockIdx.x"):
+            for tx in T.thread_binding(1, thread="threadIdx.x"):
+                with T.sblock("test"):
+                    for i in T.vectorized(0, 4):
+                        B[i] = T.exp2(A[i])
+
+    mod = tvm.compile(func, target="rocm")
+
+    def run_and_check():
+        dev = tvm.rocm(0)
+        a = tvm.runtime.tensor(np.ones((4,)).astype("float32"), dev)
+        b = tvm.runtime.tensor(np.zeros((4,)).astype("float32"), dev)
+        mod(a, b)
+        tvm.testing.assert_allclose(b.numpy(), np.exp2(a.numpy()))
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_export_load_with_fallback(monkeypatch, tmp_path):
+    """Force the codegen wrapper into the fallback branch, then export+load+run."""
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+
+    monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+    host_lib = tvm.compile(Module, target="rocm")
+    monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
+
+    lib_path = str(tmp_path / "lib.so")
+    host_lib.export_library(lib_path)
+    reloaded = tvm.runtime.load_module(lib_path)
+
+    a_np = np.random.uniform(size=(n,)).astype("float32")
+    b_np = np.zeros((n,), dtype="float32")
+
+    def run_and_check():
+        dev = tvm.rocm(0)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        reloaded["main"](a, b)
+        np.testing.assert_allclose(b.numpy(), a_np + 1.0, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_static_init.py

@@ -0,0 +1,52 @@
+# 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 ctypes
+
+import numpy as np
+
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_static_init():
+    @tvm.register_global_func("test_static_callback")
+    def test_cb(sh, A):
+        assert isinstance(sh, ctypes.c_void_p)
+        return sh
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def ramp(A: T.handle):
+            T.func_attr({"global_symbol": "ramp"})
+            n = T.int64()
+            Ab = T.match_buffer(A, (n,), "int64")
+            T.call_packed(
+                "test_static_callback",
+                T.call_intrin("handle", "tirx.tvm_static_handle"),
+                Ab.data,
+            )
+
+    mod = Module
+    f = tvm.driver.build(mod, target="llvm")
+    a = tvm.runtime.tensor(np.zeros(10, dtype="int64"))
+    f(a)
+
+
+if __name__ == "__main__":
+    test_static_init()

tests/python/codegen/test_target_codegen_vulkan.py

@@ -0,0 +1,699 @@
+# 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: E501, F841
+
+import re
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import ir as I_builder
+from tvm.script.ir_builder import tirx as T_builder
+from tvm.testing import env
+
+dtype = tvm.testing.parameter("float32", "int32", "float16", "int8")
+fuzz_seed = tvm.testing.parameter(range(25))
+
+
+# Explicitly specify a target, as this test is looking at the
+# generated shader code, and is not running on an actual device.
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled(
+        {
+            "kind": "vulkan",
+            "supports_int8": 1,
+            "supports_8bit_buffer": 1,
+            "supports_storage_buffer_storage_class": 1,
+            "supports_float16": 1,
+            "supports_16bit_buffer": 1,
+        }
+    ),
+    reason="vulkan not enabled",
+)
+def test_vector_comparison(dtype):
+    target = {
+        "kind": "vulkan",
+        "supports_int8": 1,
+        "supports_8bit_buffer": 1,
+        "supports_storage_buffer_storage_class": 1,
+        "supports_float16": 1,
+        "supports_16bit_buffer": 1,
+    }
+    target = tvm.target.Target(target)
+    zero = tvm.tirx.const(0, dtype)
+    one = tvm.tirx.const(1, dtype)
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1024,), dtype), B: T.Buffer((1024,), dtype)):
+            for i_0 in T.thread_binding(8, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    for i_2 in T.vectorized(4):
+                        with T.sblock("B"):
+                            v_i = T.axis.spatial(1024, i_0 * 128 + i_1 * 4 + i_2)
+                            B[v_i] = T.Select(A[v_i] >= zero, A[v_i] + one, zero)
+
+    # Build
+    f = tvm.tirx.build(Module, target=target)
+
+    # Verify we generate the boolx4 type declaration and the OpSelect
+    # v4{float,half,int} instruction
+    assembly = f.imports[0].inspect_source()
+    matches = re.findall("%v4bool = OpTypeVector %bool 4", assembly)
+    assert len(matches) == 1
+    matches = re.findall("OpSelect %v4.*", assembly)
+    assert len(matches) == 1
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        "llvm",
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param("rocm", marks=pytest.mark.gpu),
+        pytest.param("vulkan", marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+        pytest.param("opencl", marks=pytest.mark.gpu),
+        pytest.param("nvptx", marks=pytest.mark.gpu),
+    ],
+)
+def test_array_copy(target, dtype, fuzz_seed):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    np.random.seed(fuzz_seed)
+
+    log_arr_size = np.random.uniform(low=np.log(1), high=np.log(32768))
+    arr_size = np.exp(log_arr_size).astype(int)
+    a_np = np.random.uniform(size=(arr_size,)).astype(dtype)
+
+    def run_and_check():
+        dev = tvm.device(target)
+        a = tvm.runtime.empty((arr_size,), dtype, dev).copyfrom(a_np)
+        tvm.testing.assert_allclose(a_np, a.numpy())
+
+    if target == "llvm":
+        run_and_check()
+    else:
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_array_vectorize_add(dtype):
+    target = {"kind": "vulkan", "from_device": 0}
+    target = tvm.target.Target(target)
+    arr_size = 64
+    lanes = 2
+
+    vec_dtype = f"{dtype}x{lanes}"
+    one = tvm.tirx.const(1, vec_dtype)
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((64,), vec_dtype), B: T.Buffer((64,), vec_dtype)):
+            for i_0 in T.thread_binding(16, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(4, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(64, i_0 * 4 + i_1)
+                        B[v_i] = A[v_i] + one
+
+    f = tvm.compile(Module, target=target)
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        a = tvm.runtime.empty((arr_size,), vec_dtype, dev).copyfrom(
+            np.random.uniform(size=(arr_size, lanes))
+        )
+        c = tvm.runtime.empty((arr_size,), vec_dtype, dev)
+        f(a, c)
+        tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vulkan_bool_load():
+    target = {"kind": "vulkan", "from_device": 0}
+    target = tvm.target.Target(target)
+    arr_size = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((1024,), "bool"), B: T.Buffer((1024,), "int32")):
+            for i_0 in T.thread_binding(8, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(128, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(1024, i_0 * 128 + i_1)
+                        B[v_i] = T.Cast("int32", A[v_i])
+
+    f = tvm.compile(Module, target=target)
+
+    a_np = np.random.uniform(size=arr_size) > 0.5
+    b_np = np.zeros((arr_size,), dtype="int32")
+    ref = a_np.astype(np.int32)
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        f(a, b)
+        tvm.testing.assert_allclose(b.numpy(), ref)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+vulkan_parameter_impl = tvm.testing.parameter("push_constants", "ubo")
+vulkan_parameter_dtype = tvm.testing.parameter("int32", "float32", "int64")
+
+
+# Only run on vulkan because extremely large numbers of input
+# parameters can crash cuda/llvm compiler.
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vulkan_constant_passing(vulkan_parameter_impl, vulkan_parameter_dtype):
+    target = {"kind": "vulkan", "from_device": 0}
+    target = tvm.target.Target(target)
+    dtype = vulkan_parameter_dtype
+
+    if not target.attrs.get("supports_int64", False):
+        pytest.xfail("Vulkan target does not support Int64 variables")
+
+    # f_add has 3+num_int_params scalar parameters.  The other three
+    # are length_n, stride1, and stride2.
+    if vulkan_parameter_impl == "push_constants":
+        # 4 params, 32 bytes.  Within 128-byte spec-guaranteed size of
+        # push constants.  Uses push constants.
+        num_int_params = 1
+    else:
+        # 24 params, 192 bytes.  May be above spec-guaranteed size of 128
+        # bytes for push constants.  Uses either push constants or UBO,
+        # depending on the device.
+        max_push_constants_size = int(target.attrs.get("max_push_constants_size", 128))
+        max_int_params_in_push = max_push_constants_size // 8 - 3
+        num_int_params = max_int_params_in_push + 1
+
+    # Build IRModule programmatically since num_int_params is dynamic
+    with IRBuilder() as ib:
+        with I_builder.ir_module():
+            with T_builder.prim_func():
+                T_builder.func_name("main")
+                scalar_vars = []
+                for i in range(num_int_params):
+                    v = T_builder.arg(f"scale{i}", tvm.tirx.Var("", dtype))
+                    scalar_vars.append(v)
+                var_A = T_builder.arg("var_A", T_builder.handle())
+                var_B = T_builder.arg("var_B", T_builder.handle())
+                T_builder.func_attr({"tirx.noalias": True})
+                n_var = T_builder.int32()
+                A = T_builder.match_buffer(var_A, (n_var,), dtype)
+                B = T_builder.match_buffer(var_B, (n_var,), dtype)
+                scalar_sum = scalar_vars[0]
+                for s in scalar_vars[1:]:
+                    scalar_sum = scalar_sum + s
+                with T_builder.thread_binding(
+                    tvm.tirx.ceildiv(n_var, 64), thread="blockIdx.x"
+                ) as i_0:
+                    with T_builder.thread_binding(64, thread="threadIdx.x") as i_1:
+                        with T_builder.sblock("B"):
+                            v_i = T_builder.axis.spatial(n_var, i_0 * 64 + i_1)
+                            T_builder.where(i_0 * 64 + i_1 < n_var)
+                            T_builder.reads(A[v_i])
+                            T_builder.writes(B[v_i])
+                            T_builder.buffer_store(B, scalar_sum + A[v_i], [v_i])
+    mod = ib.get()
+    f_add = tvm.compile(mod, target=target)
+
+    n = 1024
+    scalars = np.array([1 for _ in range(num_int_params)]).astype(dtype)
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        a = tvm.runtime.tensor(np.random.uniform(size=n).astype(dtype), dev)
+        b = tvm.runtime.tensor(np.zeros(n, dtype=dtype), dev)
+        f_add(*scalars, a, b)
+        tvm.testing.assert_allclose(a.numpy() + sum(scalars), b.numpy())
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vulkan_while_if():
+    target = {"kind": "vulkan", "from_device": 0}
+    target = tvm.target.Target(target)
+    n = 1
+    dtype = "int32"
+
+    @T.prim_func(s_tir=True)
+    def while_if_gpu(A: T.Buffer((1,), "int32"), B: T.Buffer((1,), "int32")):
+        for bx in T.thread_binding(1, thread="blockIdx.x"):
+            iterations = T.decl_buffer((1,), "int32", scope="local")
+            iterations[0] = 0
+            B[0] = 0
+            while iterations[0] < T.if_then_else(A[0] > 0, 10, 20):
+                iterations[0] = iterations[0] + 1
+                B[0] = B[0] + iterations[0]
+
+    mod = tvm.IRModule.from_expr(while_if_gpu.with_attr("target", target))
+    compiled_func = tvm.compile(mod, target=target)
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        for input_value, expected in [(5, [55]), (-5, [210])]:
+            a = tvm.runtime.tensor(np.array([input_value], dtype=dtype), dev)
+            b = tvm.runtime.tensor(np.zeros(n, dtype=dtype), dev)
+            compiled_func(a, b)
+            tvm.testing.assert_allclose(b.numpy(), expected)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vulkan_local_threadidx():
+    target = {"kind": "vulkan", "from_device": 0}
+    target = tvm.target.Target(target)
+    n = 32
+
+    @T.prim_func(s_tir=True)
+    def local_threadidx_func(A: T.Buffer((32,), "int32"), B: T.Buffer((32,), "int32")):
+        # First block with thread extent 16
+        for _ in range(1):
+            for tx in T.thread_binding(16, thread="threadIdx.x"):
+                B[tx + 0] = A[tx + 0]
+        # Second block with thread extent 16
+        for _ in range(1):
+            for tx in T.thread_binding(16, thread="threadIdx.x"):
+                B[tx + 16] = A[tx + 16]
+
+    mod = tvm.IRModule.from_expr(local_threadidx_func)
+    func = tvm.compile(mod, target=target)
+
+    a_np = np.arange(n).astype(dtype="int32")
+    b_np = np.zeros((n,), dtype="int32")
+
+    def run_and_check():
+        dev = tvm.device(target.kind.name)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        func(a, b)
+        tvm.testing.assert_allclose(b.numpy(), a_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vectorized_index_ramp():
+    """Test vectorized copy with ramp indices (load N values, write to N locations)"""
+    target = {"kind": "vulkan", "from_device": 0}
+    n = 4
+    ramp_index = tvm.tirx.Ramp(0, 1, 4)
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            A = T.match_buffer(var_A, (n,), "int32", offset_factor=1)
+            B = T.match_buffer(var_B, (n,), "int32", offset_factor=1)
+            with T.sblock("compute"):
+                T.reads()
+                T.writes()
+                bx = T.launch_thread("blockIdx.x", 1)
+                B[ramp_index] = A[ramp_index]
+
+    f = tvm.compile(Module, target=target)
+
+    a_np = np.random.randint(np.iinfo("int32").max, size=n).astype("int32")
+    b_np = np.zeros(n, dtype="int32")
+
+    def run_and_check():
+        dev = tvm.device(target["kind"])
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        f(a, b)
+        tvm.testing.assert_allclose(b.numpy(), a_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_vectorized_index_broadcast():
+    """Test broadcast index (load 1 value, write to N locations)"""
+    target = {"kind": "vulkan", "from_device": 0}
+    n = 4
+    broadcast_index = tvm.tirx.Broadcast(0, 4)
+    ramp_index = tvm.tirx.Ramp(0, 1, 4)
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(var_A: T.handle, var_B: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            A = T.match_buffer(var_A, (n,), "int32", offset_factor=1)
+            B = T.match_buffer(var_B, (n,), "int32", offset_factor=1)
+            with T.sblock("compute"):
+                T.reads()
+                T.writes()
+                bx = T.launch_thread("blockIdx.x", 1)
+                # Load from broadcast index (single element), store to ramp index
+                B[ramp_index] = A[broadcast_index]
+
+    f = tvm.compile(Module, target=target)
+
+    a_np = np.random.randint(np.iinfo("int32").max, size=n).astype("int32")
+    b_np = np.zeros(n, dtype="int32")
+
+    def run_and_check():
+        dev = tvm.device(target["kind"])
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        f(a, b)
+        # All elements of b should be a[0] (broadcast load)
+        tvm.testing.assert_allclose(b.numpy(), np.full(n, a_np[0]))
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
+    reason="vulkan not enabled",
+)
+def test_negative_operand_divmod():
+    """Test handling of negative offsets to floormod/floordiv
+
+    Even though the SPIR-V spec states that OpSRem and OpSMod can give
+    the signed modulo, the Vulkan spec states that any use of negative
+    operands is undefined behavior.  This test starts with negative
+    operands to floordiv, validating that they are simplified into the
+    corresponding positive operands, such that the final TIR can be
+    expressed using only positive operands.
+
+    SPIR-V: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpSRem
+    Vulkan: https://registry.khronos.org/vulkan/specs/1.3/html/chap37.html#spirvenv-op-prec
+    """
+    target = {"kind": "vulkan", "from_device": 0}
+
+    N = 32
+    offset = 16
+    divisor = 5
+
+    @T.prim_func(s_tir=True)
+    def func(A: T.Buffer((N, 2), "int32")):
+        for i in T.thread_binding(N, thread="threadIdx.x"):
+            with T.sblock("A"):
+                v_i = T.axis.spatial(N, i)
+                A[v_i, 0] = T.floordiv(v_i - offset, divisor)
+                A[v_i, 1] = T.floormod(v_i - offset, divisor)
+
+    built = tvm.compile(func, target=target)
+
+    def run_and_check():
+        dev = tvm.device(target["kind"])
+        a_dev = tvm.runtime.empty([N, 2], "int32", dev)
+        built(a_dev)
+        a = a_dev.numpy()
+        np.testing.assert_array_equal(a[:, 0], (np.arange(N) - offset) // divisor)
+        np.testing.assert_array_equal(a[:, 1], (np.arange(N) - offset) % divisor)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize("out_dtype", ["float32", "float16"])
+def test_cooperative_matrix(out_dtype):
+    M, N, K = 16, 16, 32
+
+    # fmt: off
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(X: T.Buffer((16, 32), "float16"), W: T.Buffer((32, 16), "float16"), compute: T.Buffer((16, 16), out_dtype)):
+            T.func_attr({"tirx.noalias": True})
+            X_shared = T.sblock_alloc_buffer((16, 32), "float16", scope="shared")
+            W_shared = T.sblock_alloc_buffer((32, 16), "float16", scope="shared")
+            X_shared_wmma_matrix_a = T.sblock_alloc_buffer((16, 32), "float16", scope="wmma.matrix_a")
+            W_shared_wmma_matrix_b = T.sblock_alloc_buffer((32, 16), "float16", scope="wmma.matrix_b")
+            compute_wmma_accumulator = T.sblock_alloc_buffer((16, 16), out_dtype, scope="wmma.accumulator")
+            for i_0_j_0_fused in T.thread_binding(1, thread="blockIdx.x"):
+                with T.sblock("compute_init_o"):
+                    v_i_o = T.axis.spatial(1, 0)
+                    v_j_o = T.axis.spatial(1, 0)
+                    T.reads()
+                    T.writes(compute_wmma_accumulator[0:16, 0:16])
+                    C = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
+                    T.tvm_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0.0))
+                for k_0 in range(2):
+                    for ax0_ax1_fused_0 in range(2):
+                        for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
+                            for ax0_ax1_fused_2 in T.vectorized(4):
+                                with T.sblock("X_shared"):
+                                    v0 = T.axis.spatial(16, (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 16)
+                                    v1 = T.axis.spatial(32, k_0 * 16 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 16)
+                                    T.reads(X[v0, v1])
+                                    T.writes(X_shared[v0, v1])
+                                    X_shared[v0, v1] = X[v0, v1]
+                    for ax0_ax1_fused_0 in range(2):
+                        for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
+                            for ax0_ax1_fused_2 in T.vectorized(4):
+                                with T.sblock("W_shared"):
+                                    v0 = T.axis.spatial(32, k_0 * 16 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 16)
+                                    v1 = T.axis.spatial(16, (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 16)
+                                    T.reads(W[v0, v1])
+                                    T.writes(W_shared[v0, v1])
+                                    W_shared[v0, v1] = W[v0, v1]
+                    for ax0_0 in T.unroll(1):
+                        for ax1_0 in T.unroll(1):
+                            with T.sblock("X_shared_wmma.matrix_a_o"):
+                                v0_o = T.axis.spatial(1, ax0_0)
+                                v1_o = T.axis.spatial(2, k_0 + ax1_0)
+                                T.reads(X_shared[0:16, v1_o * 16:v1_o * 16 + 16])
+                                T.writes(X_shared_wmma_matrix_a[0:16, v1_o * 16:v1_o * 16 + 16])
+                                A = T.match_buffer(X_shared[0:16, v1_o * 16:v1_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=16)
+                                C = T.match_buffer(X_shared_wmma_matrix_a[0:16, v1_o * 16:v1_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_a", offset_factor=16)
+                                T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
+                    for ax0_0 in T.unroll(1):
+                        for ax1_0 in T.unroll(1):
+                            with T.sblock("W_shared_wmma.matrix_b_o"):
+                                v0_o = T.axis.spatial(2, k_0 + ax0_0)
+                                v1_o = T.axis.spatial(1, ax1_0)
+                                T.reads(W_shared[v0_o * 16:v0_o * 16 + 16, 0:16])
+                                T.writes(W_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, 0:16])
+                                A = T.match_buffer(W_shared[v0_o * 16:v0_o * 16 + 16, 0:16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=16)
+                                C = T.match_buffer(W_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, 0:16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_b", offset_factor=16)
+                                T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
+                    with T.sblock("compute_update_o"):
+                        v_i_o = T.axis.spatial(1, 0)
+                        v_j_o = T.axis.spatial(1, 0)
+                        v_k_o = T.axis.reduce(2, k_0)
+                        T.reads(compute_wmma_accumulator[0:16, 0:16], X_shared_wmma_matrix_a[0:16, v_k_o * 16:v_k_o * 16 + 16], W_shared_wmma_matrix_b[v_k_o * 16:v_k_o * 16 + 16, 0:16])
+                        T.writes(compute_wmma_accumulator[0:16, 0:16])
+                        A = T.match_buffer(X_shared_wmma_matrix_a[0:16, v_k_o * 16:v_k_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="wmma.matrix_a", offset_factor=16)
+                        B = T.match_buffer(W_shared_wmma_matrix_b[v_k_o * 16:v_k_o * 16 + 16, 0:16], (16, 16), "float16", strides=("B_s0", "B_s1"), scope="wmma.matrix_b", offset_factor=16)
+                        C = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
+                        T.tvm_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
+                with T.sblock("compute_wmma.accumulator_o"):
+                    v0_o = T.axis.spatial(1, 0)
+                    v1_o = T.axis.spatial(1, 0)
+                    T.reads(compute_wmma_accumulator[0:16, 0:16])
+                    T.writes(compute[0:16, 0:16])
+                    A = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("A_s0", "A_s1"), scope="wmma.accumulator", offset_factor=16)
+                    C = T.match_buffer(compute[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), offset_factor=16)
+                    T.tvm_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation(out_dtype), C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
+    # fmt: on
+
+    target = {"kind": "vulkan", "from_device": 0}
+    tgt_attrs = tvm.target.Target(target).attrs
+
+    if tgt_attrs.get("supports_cooperative_matrix"):
+        f = tvm.compile(Module, target=target)
+
+        def run_and_check():
+            dev = tvm.device("vulkan", 0)
+            A = tvm.runtime.tensor(np.random.randn(M, K).astype("float16"), dev)
+            B = tvm.runtime.tensor(np.random.randn(K, N).astype("float16"), dev)
+            C = tvm.runtime.tensor(np.random.randn(M, N).astype(out_dtype), dev)
+            f(A, B, C)
+            A_np = A.numpy()
+            B_np = B.numpy()
+            ref = np.dot(A_np.astype("float32"), B_np.astype("float32"))
+            tvm.testing.assert_allclose(C.numpy(), ref, rtol=1e-2, atol=1e-2)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
+def test_codegen_decl_buffer():
+    """The codegen should accept DeclBuffer nodes in its input"""
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def kernel():
+            T.func_attr({"calling_conv": 2, "global_symbol": "kernel", "tirx.noalias": True})
+            A = T.alloc_buffer((256,), dtype="float32", scope="local")
+            A_buf = T.decl_buffer([256], dtype="float32", scope="local", data=A.data)
+
+    target = tvm.target.Target("vulkan")
+    vulkan_codegen = tvm.get_global_func("target.build.vulkan")
+    vulkan_codegen(Module, target)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
+def test_codegen_static_shared_memory():
+    """The codegen should accept static shared/workgroup allocations."""
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
+            A_shared = T.alloc_buffer((128,), dtype="float32", scope="shared")
+
+            for bx in T.thread_binding(1, thread="blockIdx.x"):
+                for tx in T.thread_binding(128, thread="threadIdx.x"):
+                    A_shared[tx] = A[tx]
+                    B[tx] = A_shared[tx]
+
+    tvm.compile(Module, target="vulkan")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
+def test_unary():
+    test_funcs = [
+        (tvm.tirx.sin, lambda x: np.sin(x)),
+        (tvm.tirx.cos, lambda x: np.cos(x)),
+        (tvm.tirx.tan, lambda x: np.tan(x)),
+        (tvm.tirx.sinh, lambda x: np.sinh(x)),
+        (tvm.tirx.cosh, lambda x: np.cosh(x)),
+        (tvm.tirx.tanh, lambda x: np.tanh(x)),
+        (tvm.tirx.asin, lambda x: np.arcsin(x)),
+        (tvm.tirx.acos, lambda x: np.arccos(x)),
+        (tvm.tirx.atan, lambda x: np.arctan(x)),
+        (tvm.tirx.asinh, lambda x: np.arcsinh(x)),
+        (tvm.tirx.acosh, lambda x: np.arccosh(x)),
+        (tvm.tirx.atanh, lambda x: np.arctanh(x)),
+    ]
+
+    def run_test(tvm_intrin, np_func):
+        n = 16
+
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(var_A: T.handle, var_B: T.handle):
+                m = T.int32()
+                A = T.match_buffer(var_A, (m,), "float32")
+                B = T.match_buffer(var_B, (m,), "float32")
+                for i_0 in T.thread_binding((m + 63) // 64, thread="blockIdx.x"):
+                    for i_1 in T.thread_binding(64, thread="threadIdx.x"):
+                        with T.sblock("B"):
+                            v_i = T.axis.spatial(m, i_0 * 64 + i_1)
+                            T.where(i_0 * 64 + i_1 < m)
+                            T.reads(A[v_i])
+                            T.writes(B[v_i])
+                            B[v_i] = tvm_intrin(A[v_i])
+
+        target = tvm.target.Target("vulkan")
+        func = tvm.compile(Module, target=target)
+
+        if tvm_intrin in [tvm.tirx.asin, tvm.tirx.acos]:
+            data = np.random.uniform(-1.0, 1.0, size=n)
+        elif tvm_intrin == tvm.tirx.atanh:
+            data = np.random.uniform(-0.999, 0.999, size=n)
+        elif tvm_intrin == tvm.tirx.acosh:
+            data = np.random.uniform(1.0, 5.0, size=n)
+        else:
+            data = np.random.uniform(0.1, 0.9, size=n)
+
+        def run_and_check():
+            dev = tvm.device(target.kind.name, 0)
+            a = tvm.runtime.tensor(data.astype("float32"), dev)
+            b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+            func(a, b)
+            tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=1e-3, rtol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    for func in test_funcs:
+        run_test(*func)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
+def test_export_load_with_fallback(monkeypatch, tmp_path):
+    """Force the codegen wrapper into the fallback branch, then export."""
+    n = 1024
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
+            T.func_attr({"tirx.noalias": True})
+            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
+                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
+                    with T.sblock("B"):
+                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
+                        T.reads(A[v_i])
+                        T.writes(B[v_i])
+                        B[v_i] = A[v_i] + 1.0
+
+    monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+    host_lib = tvm.compile(Module, target="vulkan")
+    monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
+
+    lib_path = str(tmp_path / "lib.so")
+    host_lib.export_library(lib_path)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/codegen/test_target_codegen_x86.py

@@ -0,0 +1,88 @@
+# 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: E741
+import platform
+import re
+
+import pytest
+
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+from tvm.testing import env
+
+llvm_version = tvm.target.codegen.llvm_version_major()
+machine = platform.machine()
+
+if machine not in ["x86_64", "AMD64", "amd64"]:
+    pytest.skip(f"Requires x86_64, but machine is {machine}", allow_module_level=True)
+
+
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+@pytest.mark.skipif(llvm_version < 6, reason=f"Requires LLVM 6+, got {llvm_version}")
+def test_fp16_to_fp32():
+    def fp16_to_fp32(target, width, match=None, not_match=None):
+        elements = 64
+
+        @I.ir_module(s_tir=True)
+        class Module:
+            @T.prim_func(s_tir=True)
+            def main(
+                A: T.Buffer((elements, width), "float16"),
+                B: T.Buffer((elements, width), "float32"),
+            ):
+                T.func_attr({"tirx.noalias": True})
+                for i0 in range(elements):
+                    for i1 in T.vectorized(width):
+                        with T.sblock("B"):
+                            v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+                            T.reads(A[v_i0, v_i1])
+                            T.writes(B[v_i0, v_i1])
+                            B[v_i0, v_i1] = T.Cast("float32", A[v_i0, v_i1])
+
+        f = tvm.tirx.build(Module, target=target)
+
+        assembly = f.inspect_source("asm").splitlines()
+        if match:
+            matches = [l for l in assembly if re.search(match, l)]
+            assert matches
+        if not_match:
+            not_matches = [l for l in assembly if re.search(not_match, l)]
+            assert not not_matches
+
+    fp16_to_fp32({"kind": "llvm", "mcpu": "skylake-avx512"}, 15, match="vcvtph2ps.*mm")
+    fp16_to_fp32({"kind": "llvm", "mcpu": "skylake-avx512"}, 16, match="vcvtph2ps.*mm")
+    fp16_to_fp32({"kind": "llvm", "mcpu": "skylake-avx512"}, 17, match="vcvtph2ps.*mm")
+    fp16_to_fp32({"kind": "llvm", "mcpu": "skylake-avx512"}, 49, match="vcvtph2ps.*mm")
+    fp16_to_fp32(
+        {"kind": "llvm", "mcpu": "skylake-avx512", "mattr": ["-avx512f"]}, 49, match="vcvtph2ps.*mm"
+    )
+    fp16_to_fp32(
+        {"kind": "llvm", "mcpu": "skylake-avx512", "mattr": ["-f16c", "-avx512f"]},
+        49,
+        not_match="vcvtph2ps",
+    )
+    fp16_to_fp32({"kind": "llvm", "mcpu": "core-avx2"}, 8, match="vcvtph2ps.*mm")
+    fp16_to_fp32({"kind": "llvm", "mcpu": "core-avx2"}, 9, match="vcvtph2ps.*mm")
+    fp16_to_fp32("llvm", 9, not_match="vcvtph2ps")
+
+
+is_32bit = platform.architecture()[0] == "32bit"
+
+
+if __name__ == "__main__":
+    test_fp16_to_fp32()