chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,870 @@
+# 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 numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.testing import env
+from tvm.tirx.layout import S, TileLayout, wg_local_layout
+
+
+@pytest.mark.parametrize(
+    "input",
+    [
+        ######### basic test #########
+        (
+            (32, 32),  # g_shape
+            (0, 0),  # st_a
+            (0, 0),  # st_b
+            (0, 0),  # st_res
+            (32, 32),  # extent_a
+            (32, 32),  # extent_b
+            (32, 32),  # extent_res
+            64,  # thread_cnt
+            0,  # device_id
+        ),
+        ######### offset test #########
+        (
+            (32, 8, 12),  # g_shape
+            (10, 0, 3),  # st_a
+            (14, 1, 4),  # st_b
+            (20, 0, 2),  # st_res
+            (5, 6, 7),  # extent_a
+            (5, 6, 7),  # extent_b
+            (5, 6, 7),  # extent_res
+            64,  # thread_cnt
+            0,  # device_id
+        ),
+        ######### broadcast test #########
+        (
+            (32, 8, 12),  # g_shape
+            (10, 0, 3),  # st_a
+            (14, 1, 4),  # st_b
+            (20, 0, 2),  # st_res
+            (5, 6, 7),  # extent_a
+            (1, 6, 1),  # extent_b
+            (5, 6, 7),  # extent_res
+            64,  # thread_cnt
+            0,  # device_id
+        ),
+    ],
+)
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
+@pytest.mark.parametrize("operands_type", ["region_region", "region_const", "const_region"])
+@pytest.mark.parametrize("dtype", ["float16"])
+def test_binary_op_shared(input, op_type, operands_type, dtype):
+    # skip test
+    if op_type in ["sub", "fdiv"] and operands_type == "const_region":
+        return
+
+    g_shape, st_a, st_b, st_res, ext_a, ext_b, ext_res, thread_cnt, device_id = input
+    g_layout = s_layout = TileLayout(S[g_shape])
+
+    copy_slice = list(slice(None) for i in range(len(g_shape)))
+    map_slice_a = list(slice(st_a[i], st_a[i] + ext_a[i]) for i in range(len(g_shape)))
+    map_slice_b = list(slice(st_b[i], st_b[i] + ext_b[i]) for i in range(len(g_shape)))
+    map_slice_res = list(slice(st_res[i], st_res[i] + ext_res[i]) for i in range(len(g_shape)))
+
+    const = T.float16(3.0) if dtype == "float16" else T.float32(3.0)
+
+    # fmt: off
+    @T.prim_func
+    def binary_op_region_region(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, g_shape, dtype, layout=g_layout)
+        B = T.match_buffer(B_ptr, g_shape, dtype, layout=g_layout)
+
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([thread_cnt])
+        A_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=s_layout)
+        B_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=s_layout)
+
+        Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
+        Tx.cta.copy(B_smem[tuple(copy_slice)], B[tuple(copy_slice)])
+        T.cuda.cta_sync()
+        if op_type == "add":
+            Tx.cta.add(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)])  # noqa: E501
+        elif op_type == "sub":
+            Tx.cta.sub(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)])  # noqa: E501
+        elif op_type == "mul":
+            Tx.cta.mul(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)])  # noqa: E501
+        elif op_type == "fdiv":
+            Tx.cta.fdiv(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)])  # noqa: E501
+        T.cuda.cta_sync()
+        Tx.cta.copy(A[tuple(copy_slice)], A_smem[tuple(copy_slice)])
+
+    @T.prim_func
+    def binary_op_const_region_or_region_const(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, g_shape, dtype, layout=g_layout)
+        _B = T.match_buffer(B_ptr, g_shape, dtype, layout=g_layout)
+
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([thread_cnt])
+        A_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=s_layout)
+
+        Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
+        T.cuda.cta_sync()
+        if op_type == "add":
+            if operands_type == "const_region":
+                Tx.cta.add(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
+            elif operands_type == "region_const":
+                Tx.cta.add(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
+        elif op_type == "sub":
+            if operands_type == "const_region":
+                Tx.cta.sub(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
+            elif operands_type == "region_const":
+                Tx.cta.sub(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
+        elif op_type == "mul":
+            if operands_type == "const_region":
+                Tx.cta.mul(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
+            elif operands_type == "region_const":
+                Tx.cta.mul(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
+        elif op_type == "fdiv":
+            if operands_type == "const_region":
+                Tx.cta.fdiv(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
+            elif operands_type == "region_const":
+                Tx.cta.fdiv(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
+        T.cuda.cta_sync()
+        Tx.cta.copy(A[tuple(copy_slice)], A_smem[tuple(copy_slice)])
+        # fmt: on
+
+    def get_prim_func(operands_type):
+        if operands_type == "region_region":
+            return binary_op_region_region
+        elif operands_type in ["const_region", "region_const"]:
+            return binary_op_const_region_or_region_const
+        raise ValueError(f"operands_type={operands_type} is not supported")
+
+    def get_ref(A_np, B_np):
+        A_ref = A_np.copy()
+        if op_type == "add":
+            if operands_type == "region_region":
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] + B_np[tuple(map_slice_b)]
+            elif operands_type in ["const_region", "region_const"]:
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] + 3.0
+        elif op_type == "sub":
+            if operands_type == "region_region":
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] - B_np[tuple(map_slice_b)]
+            elif operands_type in ["const_region", "region_const"]:
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] - 3.0
+        elif op_type == "mul":
+            if operands_type == "region_region":
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] * B_np[tuple(map_slice_b)]
+            elif operands_type in ["const_region", "region_const"]:
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] * 3.0
+        elif op_type == "fdiv":
+            if operands_type == "region_region":
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] / B_np[tuple(map_slice_b)]
+            elif operands_type in ["const_region", "region_const"]:
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] / 3.0
+
+        return A_ref
+
+    target = tvm.target.Target("cuda")
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(*g_shape).astype(dtype)
+        B_np = np.random.rand(*g_shape).astype(dtype)
+        mod = tvm.IRModule({"main": get_prim_func(operands_type)})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        print(f"compiled source code: {mod.mod.imports[0].inspect_source()}")
+        A_ref = get_ref(A_np, B_np)
+        atol = 1e-3
+
+        def run_and_check():
+            dev = tvm.cuda(device_id)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            mod(A, B)
+            tvm.testing.assert_allclose(A_ref, A.numpy(), atol=atol)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize("op_type", ["sub", "fdiv"])
+def test_binary_non_commutative_const_lhs_rejected(op_type):
+    dtype = "float16"
+    shape = (16, 16)
+    layout = TileLayout(S[shape])
+    const = T.float16(3.0)
+
+    with pytest.raises(Exception):
+
+        @T.prim_func
+        def bad_kernel() -> None:
+            T.device_entry()
+            _bx = T.cta_id([1])
+            _tid = T.thread_id([64])
+            A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=layout)
+            if op_type == "sub":
+                Tx.cta.sub(A_smem, const, A_smem)
+            elif op_type == "fdiv":
+                Tx.cta.fdiv(A_smem, const, A_smem)
+
+        target = tvm.target.Target("cuda")
+        with target:
+            mod = tvm.IRModule({"main": bad_kernel})
+            tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("exec_scope", ["warp", "warpgroup"])
+@pytest.mark.parametrize("op_type", ["add", "mul"])
+def test_binary_op_shared_subcta_scope(exec_scope, op_type):
+    """Test binary ops in warp/warpgroup scope with shared memory."""
+    dtype = "float16"
+    n_warps = 4 if exec_scope == "warpgroup" else 1
+    g_shape = (n_warps * 32, 8)
+    tx_op = {
+        ("warp", "add"): Tx.warp.add,
+        ("warp", "mul"): Tx.warp.mul,
+        ("warpgroup", "add"): Tx.wg.add,
+        ("warpgroup", "mul"): Tx.wg.mul,
+    }[(exec_scope, op_type)]
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, g_shape, dtype, layout=TileLayout(S[g_shape]))
+        B = T.match_buffer(B_ptr, g_shape, dtype, layout=TileLayout(S[g_shape]))
+        T.device_entry()
+        warp_id = T.warp_id([(256) // 32])
+        wg_id = T.warpgroup_id([(256) // 128])
+        _bx = T.cta_id([1])
+        _tid = T.thread_id([256])
+        A_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=TileLayout(S[g_shape]))
+        B_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=TileLayout(S[g_shape]))
+        Tx.cta.copy(A_smem, A)
+        Tx.cta.copy(B_smem, B)
+        T.cuda.cta_sync()
+        if exec_scope == "warp":
+            if warp_id == 5:
+                tx_op(A_smem, A_smem, B_smem)
+        elif exec_scope == "warpgroup":
+            if wg_id == 1:
+                tx_op(A_smem, A_smem, B_smem)
+        T.cuda.cta_sync()
+        Tx.cta.copy(A, A_smem)
+
+    target = tvm.target.Target("cuda")
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(*g_shape).astype(dtype)
+        B_np = np.random.rand(*g_shape).astype(dtype)
+        mod = tvm.IRModule({"main": kernel})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        np_op = {"add": np.add, "mul": np.multiply}[op_type]
+        A_ref = np_op(A_np, B_np).astype(dtype)
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            mod(A, B)
+            tvm.testing.assert_allclose(A_ref, A.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("exec_scope", ["cta", "warpgroup", "warp"])
+@pytest.mark.parametrize("rhs_kind", ["region", "broadcast", "const"])
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
+def test_binary_op_local_subcta_trivial(exec_scope, rhs_kind, op_type):
+    dtype = "float16"
+    m, n = 4, 8
+    n_threads = 256 if exec_scope == "cta" else (128 if exec_scope == "warpgroup" else 32)
+    # in this test, use warp3/warpgroup1 to test
+    thr_str = 0 if exec_scope == "cta" else (128 if exec_scope == "warpgroup" else 32 * 3)
+    a_shape = (n_threads, m, n)
+    b_shape = (n_threads, m, n if rhs_kind == "region" else 1)
+    c_shape = a_shape
+    const = T.float16(1.25)
+    tx_op = {"add": Tx.add, "sub": Tx.sub, "mul": Tx.mul, "fdiv": Tx.fdiv}[op_type]
+    tid_in_scope_fn = {"cta": T.thread_id, "warpgroup": T.thread_id_in_wg, "warp": T.lane_id}[
+        exec_scope
+    ]
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
+        B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
+        C = T.match_buffer(C_ptr, c_shape, dtype, layout=TileLayout(S[c_shape]))
+
+        T.device_entry()
+        wg_id = T.warpgroup_id([(256) // 128])
+        warp_id = T.warp_id([(256) // 32])
+        _bx = T.cta_id([1])
+        _tid = T.thread_id([256])
+        tid_in_scope = tid_in_scope_fn([n_threads])
+        b_n = T.meta_var(n if rhs_kind == "region" else 1)
+        A_local = T.alloc_buffer((m, n), dtype, scope="local", layout=TileLayout(S[(m, n)]))
+        C_local = T.alloc_buffer((m, n), dtype, scope="local", layout=TileLayout(S[(m, n)]))
+        B_local = T.alloc_buffer((m, b_n), dtype, scope="local", layout=TileLayout(S[(m, b_n)]))
+
+        if thr_str <= _tid and _tid < thr_str + n_threads:
+            for i in T.serial(m):
+                for j in T.serial(n):
+                    A_local[i, j] = A[tid_in_scope, i, j]
+            if rhs_kind != "const":
+                for i in T.serial(m):
+                    for j in T.serial(b_n):
+                        B_local[i, j] = B[tid_in_scope, i, j]
+
+        if exec_scope == "cta":
+            if rhs_kind == "const":
+                tx_op(C_local, A_local, const)
+            else:
+                tx_op(C_local, A_local, B_local)
+        elif exec_scope == "warpgroup":
+            if wg_id == 1:
+                if rhs_kind == "const":
+                    tx_op(C_local, A_local, const)
+                else:
+                    tx_op(C_local, A_local, B_local)
+        else:
+            if warp_id == 3:
+                if rhs_kind == "const":
+                    tx_op(C_local, A_local, const)
+                else:
+                    tx_op(C_local, A_local, B_local)
+                # T.cuda.cta_sync()
+
+        if thr_str <= _tid and _tid < thr_str + n_threads:
+            for i in T.serial(m):
+                for j in T.serial(n):
+                    C[tid_in_scope, i, j] = C_local[i, j]
+
+    target = tvm.target.Target("cuda")
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(*a_shape).astype(dtype)
+        B_np = np.random.rand(*b_shape).astype(dtype)
+        C_np = np.zeros(c_shape, dtype=dtype)
+        mod = tvm.IRModule({"main": kernel})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        print(f"compiled source code: {mod.mod.imports[0].inspect_source()}")
+        np_op = {"add": np.add, "sub": np.subtract, "mul": np.multiply, "fdiv": np.divide}[op_type]
+        if rhs_kind == "region":
+            C_ref = np_op(A_np, B_np)
+        elif rhs_kind == "broadcast":
+            C_ref = np_op(A_np, np.repeat(B_np, n, axis=2))
+        else:
+            C_ref = np_op(A_np, const.value)
+        atol = 1e-2 if op_type == "fdiv" else 1e-3
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            C = tvm.runtime.tensor(C_np, dev)
+            mod(A, B, C)
+            tvm.testing.assert_allclose(C_ref, C.numpy(), atol=atol)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize(
+    "input",
+    [
+        ######### basic test #########
+        (
+            (64, 32),  # a_shape
+            (64, 32),  # b_shape
+            (64, 32),  # res_shape
+            64,  # thread_cnt
+            0,  # device_id
+        ),
+        ######### broadcast test #########
+        (
+            (32, 5, 4),  # a_shape
+            (32, 1, 4),  # b_shape
+            (32, 5, 4),  # res_shape
+            32,  # thread_cnt (≥ warp size so sctx.intra at cta scope models cleanly)
+            0,  # device_id
+        ),
+    ],
+)
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("storage_scope", ["shared", "local"])
+@pytest.mark.parametrize("exec_scope", ["cta", "thread"])
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
+@pytest.mark.parametrize("dtype", ["float16"])
+def test_binary_op_vectorized(input, storage_scope, exec_scope, op_type, dtype):
+    a_shape, b_shape, res_shape, thread_cnt, device_id = input
+    tx_op = {"add": Tx.add, "sub": Tx.sub, "mul": Tx.mul, "fdiv": Tx.fdiv}[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def test_binary_cta(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
+        B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([thread_cnt])
+        if storage_scope == "shared":
+            A_smem = T.alloc_buffer(
+                a_shape, dtype, scope="shared", layout=TileLayout(S[a_shape])
+            )
+            B_smem = T.alloc_buffer(
+                b_shape, dtype, scope="shared", layout=TileLayout(S[b_shape])
+            )
+            Tx.cta.copy(A_smem, A)
+            Tx.cta.copy(B_smem, B)
+            T.cuda.cta_sync()
+            tx_op(A_smem, A_smem, B_smem)
+            T.cuda.cta_sync()
+            Tx.cta.copy(A, A_smem)
+        if storage_scope == "local":
+            A_local = T.alloc_buffer(
+                a_shape[1:], dtype, scope="local", layout=TileLayout(S[a_shape[1:]])
+            )
+            B_local = T.alloc_buffer(
+                b_shape[1:], dtype, scope="local", layout=TileLayout(S[b_shape[1:]])
+            )
+            Tx.copy(A_local, A[tx])
+            Tx.copy(B_local, B[tx])
+            tx_op(A_local, A_local, B_local)
+            Tx.copy(A[tx], A_local)
+
+    @T.prim_func
+    def test_binary_thread(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
+        B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([thread_cnt])
+        if storage_scope == "shared":
+            A_smem = T.alloc_buffer(
+                a_shape, dtype, scope="shared", layout=TileLayout(S[a_shape])
+            )
+            B_smem = T.alloc_buffer(
+                b_shape, dtype, scope="shared", layout=TileLayout(S[b_shape])
+            )
+            Tx.copy(A_smem, A)
+            Tx.copy(B_smem, B)
+            T.cuda.cta_sync()
+            tx_op(A_smem, A_smem, B_smem)
+            T.cuda.cta_sync()
+            Tx.copy(A, A_smem)
+        elif storage_scope == "local":
+            A_local = T.alloc_buffer(
+                a_shape[1:], dtype, scope="local", layout=TileLayout(S[a_shape[1:]])
+            )
+            B_local = T.alloc_buffer(
+                b_shape[1:], dtype, scope="local", layout=TileLayout(S[b_shape[1:]])
+            )
+            Tx.copy(A_local, A[tx])
+            Tx.copy(B_local, B[tx])
+            tx_op(A_local, A_local, B_local)
+            Tx.copy(A[tx], A_local)
+        # fmt: on
+
+    def get_prim_func():
+        if exec_scope == "cta":
+            return test_binary_cta
+        elif exec_scope == "thread":
+            return test_binary_thread
+        else:
+            raise ValueError(f"exec_scope={exec_scope} is not supported")
+
+    target = tvm.target.Target("cuda")
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(*a_shape).astype(dtype)
+        B_np = np.random.rand(*b_shape).astype(dtype)
+        mod = tvm.IRModule({"main": get_prim_func()})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        print(f"compiled source code: {mod.mod.imports[0].inspect_source()}")
+        np_op = {"add": np.add, "sub": np.subtract, "mul": np.multiply, "fdiv": np.divide}[op_type]
+        A_ref = np_op(A_np, B_np)
+        atol = 1e-2 if op_type == "fdiv" else 1e-3
+
+        def run_and_check():
+            dev = tvm.cuda(device_id)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            mod(A, B)
+            tvm.testing.assert_allclose(A_ref, A.numpy(), atol=atol)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul"])
+def test_binary_op_packed_f32x2_auto_dispatch(op_type):
+    target = tvm.target.Target("cuda")
+    arch = target.arch if hasattr(target, "arch") else ""
+    if not arch.startswith("sm_"):
+        pytest.skip(f"unknown target arch: {arch}")
+    sm_digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
+    if not sm_digits:
+        pytest.skip(f"cannot parse target arch: {arch}")
+    sm_version = int(sm_digits)
+    if sm_version < 100:
+        pytest.skip(f"packed_f32x2 auto-dispatch requires sm_100+, got {arch}")
+
+    a_shape, b_shape = (64, 32), (64, 32)
+    dtype = "float32"
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
+        B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([64])
+        A_local = T.alloc_buffer(
+            a_shape[1:], dtype, scope="local", layout=TileLayout(S[a_shape[1:]])
+        )
+        B_local = T.alloc_buffer(
+            b_shape[1:], dtype, scope="local", layout=TileLayout(S[b_shape[1:]])
+        )
+        Tx.copy(A_local, A[tx])
+        Tx.copy(B_local, B[tx])
+        if op_type == "add":
+            Tx.add(A_local, A_local, B_local)
+        elif op_type == "sub":
+            Tx.sub(A_local, A_local, B_local)
+        elif op_type == "mul":
+            Tx.mul(A_local, A_local, B_local)
+        Tx.copy(A[tx], A_local)
+
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(*a_shape).astype(dtype)
+        B_np = np.random.rand(*b_shape).astype(dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = mod.mod.imports[0].inspect_source()
+        ptx_pat = {
+            "add": r"add\.[a-z]+\.ftz\.f32x2",
+            "sub": r"sub\.[a-z]+\.ftz\.f32x2",
+            "mul": r"mul\.[a-z]+\.ftz\.f32x2",
+        }[op_type]
+        builtin_pat = {
+            "add": r"tvm_builtin_ptx_add_packed_",
+            "sub": r"tvm_builtin_ptx_sub_packed_",
+            "mul": r"tvm_builtin_ptx_mul_packed_",
+        }[op_type]
+        assert re.search(ptx_pat, src) or re.search(builtin_pat, src), src
+        if op_type == "add":
+            A_ref = A_np + B_np
+        elif op_type == "sub":
+            A_ref = A_np - B_np
+        elif op_type == "mul":
+            A_ref = A_np * B_np
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            mod(A, B)
+            tvm.testing.assert_allclose(A_ref, A.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("op_name", ["add", "sub", "mul"])
+def test_binary_op_warpgroup_wg_local_layout(op_name):
+    dtype = "float32"
+    rows, cols = 128, 16
+    target = tvm.target.Target("cuda")
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        B = T.match_buffer(B_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        C = T.match_buffer(C_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        wg_id = T.warpgroup_id([1])
+        tid = T.thread_id_in_wg([rows])
+
+        lhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        rhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        out = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        lhs_row = lhs.local(cols)
+        rhs_row = rhs.local(cols)
+        out_row = out.local(cols)
+        for i in T.serial(cols):
+            lhs_row[i] = A[tid, i]
+            rhs_row[i] = B[tid, i]
+            out_row[i] = T.float32(0)
+        if op_name == "add":
+            Tx.wg.add(out, lhs, rhs)
+        elif op_name == "sub":
+            Tx.wg.sub(out, lhs, rhs)
+        elif op_name == "mul":
+            Tx.wg.mul(out, lhs, rhs)
+        out_row_1 = out.local(cols)
+        for i in T.serial(cols):
+            C[tid, i] = out_row_1[i]
+
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(rows, cols).astype(dtype)
+        B_np = np.random.rand(rows, cols).astype(dtype)
+        C_np = np.zeros((rows, cols), dtype=dtype)
+
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        if op_name == "add":
+            C_ref = A_np + B_np
+        elif op_name == "sub":
+            C_ref = A_np - B_np
+        else:
+            C_ref = A_np * B_np
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            C = tvm.runtime.tensor(C_np, dev)
+            mod(A, B, C)
+            tvm.testing.assert_allclose(C_ref, C.numpy(), atol=1e-5)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize("op_name,ptx_op", [("add", "add"), ("sub", "sub"), ("mul", "mul")])
+def test_binary_op_warpgroup_wg_local_emits_packed_f32x2(op_name, ptx_op):
+    """Warpgroup-scope binary on a wg-local fp32 view must lower to packed
+    f32x2 PTX on SM100+, mirroring the thread-scope packed dispatch.
+
+    Regression test for the fa4 perf path: rescale-style ``T.{add,sub,mul}``
+    calls in warpgroup scope used to fall through to scalar codegen because
+    ``_emit_binary_local_view`` only emitted ``op_func(...)`` per element.
+    """
+    target = tvm.target.Target("cuda")
+    arch = target.arch if hasattr(target, "arch") else ""
+    if not arch.startswith("sm_"):
+        pytest.skip(f"unknown target arch: {arch}")
+    sm_digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
+    if not sm_digits or int(sm_digits) < 100:
+        pytest.skip(f"packed_f32x2 wg-local path requires sm_100+, got {arch}")
+
+    dtype = "float32"
+    rows, cols = 128, 16
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        B = T.match_buffer(B_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        C = T.match_buffer(C_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _wg_id = T.warpgroup_id([1])
+        tid = T.thread_id_in_wg([rows])
+
+        lhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        rhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        out = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        lhs_row = lhs.local(cols)
+        rhs_row = rhs.local(cols)
+        out_row = out.local(cols)
+        for i in T.serial(cols):
+            lhs_row[i] = A[tid, i]
+            rhs_row[i] = B[tid, i]
+            out_row[i] = T.float32(0)
+        if op_name == "add":
+            Tx.wg.add(out, lhs, rhs)
+        elif op_name == "sub":
+            Tx.wg.sub(out, lhs, rhs)
+        else:
+            Tx.wg.mul(out, lhs, rhs)
+        out_row_1 = out.local(cols)
+        for i in T.serial(cols):
+            C[tid, i] = out_row_1[i]
+
+    with target:
+        mod = tvm.IRModule({"main": test_func})
+        ex = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = ex.mod.imports[0].inspect_source()
+
+    # Codegen must use the packed f32x2 path, not scalar fallback.
+    assert re.search(rf"{ptx_op}\.[a-z]+\.ftz\.f32x2", src) or re.search(
+        rf"tvm_builtin_ptx_{ptx_op}_packed_[a-z]+_f32x2", src
+    ), f"expected packed f32x2 PTX for op={op_name}, source preview:\n{src[:2000]}"
+
+
+def test_fma_warpgroup_wg_local_emits_packed_f32x2():
+    """Same regression coverage as the binary case but for ``T.fma``."""
+    target = tvm.target.Target("cuda")
+    arch = target.arch if hasattr(target, "arch") else ""
+    if not arch.startswith("sm_"):
+        pytest.skip(f"unknown target arch: {arch}")
+    sm_digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
+    if not sm_digits or int(sm_digits) < 100:
+        pytest.skip(f"packed_f32x2 wg-local path requires sm_100+, got {arch}")
+
+    dtype = "float32"
+    rows, cols = 128, 16
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, C_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        C = T.match_buffer(C_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _wg_id = T.warpgroup_id([1])
+        tid = T.thread_id_in_wg([rows])
+
+        buf = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        buf_row = buf.local(cols)
+        for i in T.serial(cols):
+            buf_row[i] = A[tid, i]
+        Tx.wg.fma(buf, buf, T.float32(2.0), T.float32(0.5))
+        buf_row_1 = buf.local(cols)
+        for i in T.serial(cols):
+            C[tid, i] = buf_row_1[i]
+
+    with target:
+        mod = tvm.IRModule({"main": test_func})
+        ex = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = ex.mod.imports[0].inspect_source()
+
+    assert re.search(r"fma\.[a-z]+\.ftz\.f32x2", src) or re.search(
+        r"tvm_builtin_ptx_fma_packed_[a-z]+_f32x2", src
+    ), f"expected packed f32x2 fma PTX, source preview:\n{src[:2000]}"
+
+
+# -----------------------------------------------------------------------------
+# Dispatch codegen checks (no GPU runtime — explicit target arch).
+# These complement the existing `*_warpgroup_wg_local_layout` / `*_auto_dispatch`
+# variants by forcing the arch in the Target dict, so the codegen path runs
+# even on hosts where ``Target("cuda")`` cannot detect the GPU.
+# -----------------------------------------------------------------------------
+def test_binary_add_f32_sm100_packed_f32x2_dispatch():
+    """add f32 + all-local → reg.py + add_f32x2 packed (no T.vectorized)."""
+    shape = (64, 32)
+    lay = TileLayout(S[shape])
+
+    @T.prim_func
+    def k(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, "float32", layout=lay)
+        B = T.match_buffer(B_ptr, shape, "float32", layout=lay)
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([64])
+        ra = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        rb = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        Tx.copy(ra, A[tx])
+        Tx.copy(rb, B[tx])
+        Tx.add(ra, ra, rb)
+        Tx.copy(A[tx], ra)
+
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
+    with target:
+        mod = tvm.IRModule({"main": k})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = mod.mod.imports[0].inspect_source()
+    assert re.search(r"add\.[a-z]+\.ftz\.f32x2", src) or re.search(
+        r"tvm_builtin_ptx_add_packed_", src
+    ), f"expected packed add_f32x2; got:\n{src[:2000]}"
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_binary_maximum_reg():
+    N = 128
+
+    @T.prim_func
+    def relu_max(a_ptr: T.handle, b_ptr: T.handle, c_ptr: T.handle) -> None:
+        A = T.match_buffer(a_ptr, (N,), "float32")
+        B = T.match_buffer(b_ptr, (N,), "float32")
+        C = T.match_buffer(c_ptr, (N,), "float32")
+        T.device_entry()
+        T.warp_id([4])
+        T.cta_id([1])
+        T.warpgroup_id([1])
+        tid = T.thread_id_in_wg([N])
+        a_reg = T.alloc_local((1,), "float32")
+        b_reg = T.alloc_local((1,), "float32")
+        Tx.copy(a_reg[:], A[tid : tid + 1])
+        Tx.copy(b_reg[:], B[tid : tid + 1])
+        Tx.maximum(a_reg[:], a_reg[:], b_reg[:])
+        Tx.copy(C[tid : tid + 1], a_reg[:])
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.compile(tvm.IRModule({"main": relu_max}), target=target, tir_pipeline="tirx")
+    np.random.seed(0)
+    a = np.random.randn(N).astype("float32")
+    b = np.random.randn(N).astype("float32")
+    c = np.zeros(N, dtype="float32")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        a_t, b_t, c_t = (tvm.runtime.tensor(x, dev) for x in (a, b, c))
+        mod["main"](a_t, b_t, c_t)
+        np.testing.assert_allclose(c_t.numpy(), np.maximum(a, b), atol=0, rtol=0)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+def test_binary_add_f16_scalar_fallback_dispatch():
+    """add f16 has no packed VecImpl → reg.py scalar fallback (T.vectorized)."""
+    shape = (64, 32)
+    lay = TileLayout(S[shape])
+
+    @T.prim_func
+    def k(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, "float16", layout=lay)
+        B = T.match_buffer(B_ptr, shape, "float16", layout=lay)
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([64])
+        ra = T.alloc_buffer(shape[1:], "float16", scope="local", layout=TileLayout(S[shape[1:]]))
+        rb = T.alloc_buffer(shape[1:], "float16", scope="local", layout=TileLayout(S[shape[1:]]))
+        Tx.copy(ra, A[tx])
+        Tx.copy(rb, B[tx])
+        Tx.add(ra, ra, rb)
+        Tx.copy(A[tx], ra)
+
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
+    with target:
+        mod = tvm.IRModule({"main": k})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = mod.mod.imports[0].inspect_source()
+    assert "half" in src or "__half" in src, f"expected scalar half add; got:\n{src[:2000]}"
+    assert not re.search(r"add\.[a-z]+\.ftz\.f(32|16)x2", src), (
+        f"unexpected packed f32x2/f16x2 add in scalar-fallback path; got:\n{src[:2000]}"
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()
@@ -0,0 +1,403 @@
+# 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.
+"""Tests for FMA op dispatch, layout=None local dispatch, scalar broadcast,
+and rounding mode support."""
+
+import re
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.testing import env
+from tvm.tirx.layout import S, TileLayout, wg_local_layout
+
+
+def _get_sm_version():
+    target = tvm.target.Target("cuda")
+    arch = target.arch if hasattr(target, "arch") else ""
+    if not arch.startswith("sm_"):
+        return 0
+    digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
+    return int(digits) if digits else 0
+
+
+# ---------------------------------------------------------------------------
+# FMA op: scalar scale + scalar bias
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_fma_scalar_scalar():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed fma requires sm_100+, got sm_{sm}")
+
+    N = 128
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    scale_val = 0.5
+    bias_val = -1.0
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([N])
+        buf = T.alloc_buffer((1,), dtype, scope="local", layout=TileLayout(S[1]))
+        Tx.copy(buf, A[tx : tx + 1])
+        Tx.fma(buf, buf, T.float32(scale_val), T.float32(bias_val))
+        Tx.copy(A[tx : tx + 1], buf)
+
+    with target:
+        A_np = np.random.rand(N).astype(dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            mod(A)
+            expected = A_np * scale_val + bias_val
+            tvm.testing.assert_allclose(expected, A.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# FMA op: buffer scale + scalar bias (Horner pattern)
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_fma_buffer_scale_scalar_bias():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed fma requires sm_100+, got sm_{sm}")
+
+    N = 2
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    coeff = 0.695
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        B = T.match_buffer(B_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _tx = T.thread_id([1])
+        acc = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        frac = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        Tx.copy(acc, A[0:N])
+        Tx.copy(frac, B[0:N])
+        Tx.fma(acc, acc, frac, T.float32(coeff))
+        Tx.copy(A[0:N], acc)
+
+    with target:
+        A_np = np.random.rand(N).astype(dtype)
+        B_np = np.random.rand(N).astype(dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            mod(A, B)
+            expected = A_np * B_np + coeff
+            tvm.testing.assert_allclose(expected, A.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# Binary op with scalar broadcast (Expr scalar, e.g. BufferLoad)
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_mul_scalar_broadcast():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed mul requires sm_100+, got sm_{sm}")
+
+    N = 16
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, S_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        Scale = T.match_buffer(S_ptr, (1,), dtype, layout=TileLayout(S[1]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _tx = T.thread_id([1])
+        a_local = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        s_local = T.alloc_buffer((1,), dtype, scope="local", layout=TileLayout(S[1]))
+        Tx.copy(a_local, A[0:N])
+        Tx.copy(s_local, Scale[0:1])
+        Tx.mul(a_local, a_local, s_local[0])
+        Tx.copy(A[0:N], a_local)
+
+    with target:
+        A_np = np.random.rand(N).astype(dtype)
+        S_np = np.array([2.5], dtype=dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_dev = tvm.runtime.tensor(A_np, dev)
+            S_dev = tvm.runtime.tensor(S_np, dev)
+            mod(A_dev, S_dev)
+            expected = A_np * S_np[0]
+            tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# Binary add with rounding mode
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_add_rounding_mode():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed add with rounding requires sm_100+, got sm_{sm}")
+
+    N = 2
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    round_const = float(2**23 + 2**22)
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _tx = T.thread_id([1])
+        buf = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        Tx.copy(buf, A[0:N])
+        Tx.add(buf, buf, T.float32(round_const), rounding_mode="rm")
+        Tx.copy(A[0:N], buf)
+
+    with target:
+        A_np = np.array([1.3, 2.7], dtype=dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        # Check that the PTX uses the rounding mode
+        src = mod.mod.imports[0].inspect_source()
+        assert re.search(r"add\.rm\.ftz\.f32x2", src) or re.search(
+            r"tvm_builtin_ptx_add_packed_", src
+        ), f"Expected packed add with rm rounding in PTX:\n{src}"
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_dev = tvm.runtime.tensor(A_np, dev)
+            mod(A_dev)
+            expected = A_np + round_const
+            tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1.0)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# FMA op: layout=None local buffer (no TileLayout)
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_fma_no_layout():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed fma requires sm_100+, got sm_{sm}")
+
+    N = 4
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    scale_val = 2.0
+    bias_val = 1.0
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _tx = T.thread_id([1])
+        buf = T.alloc_local([N], dtype)
+        for i in T.serial(N):
+            buf[i] = A[i]
+        Tx.fma(buf[0:N], buf[0:N], T.float32(scale_val), T.float32(bias_val))
+        for i in T.serial(N):
+            A[i] = buf[i]
+
+    with target:
+        A_np = np.array([1.0, 2.0, 3.0, 4.0], dtype=dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_dev = tvm.runtime.tensor(A_np, dev)
+            mod(A_dev)
+            expected = A_np * scale_val + bias_val
+            tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# Binary sub with rounding mode (buffer-buffer)
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_sub_buffer_buffer_rounding():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed sub with rounding requires sm_100+, got sm_{sm}")
+
+    N = 2
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        B = T.match_buffer(B_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _tx = T.thread_id([1])
+        a_buf = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        b_buf = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        Tx.copy(a_buf, A[0:N])
+        Tx.copy(b_buf, B[0:N])
+        Tx.sub(a_buf, a_buf, b_buf, rounding_mode="rn")
+        Tx.copy(A[0:N], a_buf)
+
+    with target:
+        A_np = np.array([3.14, 2.71], dtype=dtype)
+        B_np = np.array([1.41, 0.57], dtype=dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = mod.mod.imports[0].inspect_source()
+        assert re.search(r"sub\.rn\.ftz\.f32x2", src) or re.search(
+            r"tvm_builtin_ptx_sub_packed_", src
+        ), f"Expected packed sub with rn rounding in PTX:\n{src}"
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_dev = tvm.runtime.tensor(A_np, dev)
+            B_dev = tvm.runtime.tensor(B_np, dev)
+            mod(A_dev, B_dev)
+            expected = A_np - B_np
+            tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1e-6)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_fma_warpgroup_wg_local_layout():
+    rows, cols = 128, 8
+    dtype = "float32"
+    scale_val = 1.5
+    bias_val = -0.25
+    target = tvm.target.Target("cuda")
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        B = T.match_buffer(B_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        wg_id = T.warpgroup_id([1])
+        tid = T.thread_id_in_wg([rows])
+
+        reg = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        reg_row = reg.local(cols)
+        for i in T.serial(cols):
+            reg_row[i] = A[tid, i]
+        Tx.wg.fma(reg, reg, T.float32(scale_val), T.float32(bias_val))
+        reg_row_1 = reg.local(cols)
+        for i in T.serial(cols):
+            B[tid, i] = reg_row_1[i]
+
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(rows, cols).astype(dtype)
+        B_np = np.zeros((rows, cols), dtype=dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_dev = tvm.runtime.tensor(A_np, dev)
+            B_dev = tvm.runtime.tensor(B_np, dev)
+            mod(A_dev, B_dev)
+            expected = A_np * scale_val + bias_val
+            tvm.testing.assert_allclose(expected, B_dev.numpy(), atol=1e-5)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# -----------------------------------------------------------------------------
+# Dispatch codegen check (no GPU runtime — explicit target arch).
+# Complements ``test_fma_warpgroup_wg_local_emits_packed_f32x2`` (which uses
+# the host-detected ``Target("cuda")`` and skips when arch < sm_100).
+# -----------------------------------------------------------------------------
+def test_fma_f32_sm100_packed_f32x2_dispatch():
+    """fma f32 + all-local → reg.py + fma_f32x2 packed (no T.vectorized)."""
+    shape = (64, 32)
+    lay = TileLayout(S[shape])
+
+    @T.prim_func
+    def k(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, D_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, "float32", layout=lay)
+        B = T.match_buffer(B_ptr, shape, "float32", layout=lay)
+        C = T.match_buffer(C_ptr, shape, "float32", layout=lay)
+        D = T.match_buffer(D_ptr, shape, "float32", layout=lay)
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([64])
+        ra = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        rb = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        rc = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        rd = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        Tx.copy(ra, A[tx])
+        Tx.copy(rb, B[tx])
+        Tx.copy(rc, C[tx])
+        Tx.fma(rd, ra, rb, rc)
+        Tx.copy(D[tx], rd)
+
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
+    with target:
+        mod = tvm.IRModule({"main": k})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = mod.mod.imports[0].inspect_source()
+    assert re.search(r"fma\.[a-z]+\.ftz\.f32x2", src) or re.search(
+        r"tvm_builtin_ptx_fma_packed_", src
+    ), f"expected packed fma_f32x2; got:\n{src[:2000]}"
+
+
+if __name__ == "__main__":
+    tvm.testing.main()