[Fix][Relax] Lower non-contiguous WebGPU cumsum

python/tvm/relax/backend/dispatch_sort_scan.py

@@ -146,40 +146,58 @@ def visit_call_(self, call: relax.Call) -> relax.Expr:
             # TODO(tvm-team): Support fully dynamic case with `shape=None`
             if shape is None:
                 raise ValueError("non-symbolic shape is not supported for now")
+            shape_values = [shape[i] for i in range(len(shape))]
             kwargs = {}
+            normalized_axis = axis
+            if normalized_axis is not None and normalized_axis < 0:
+                normalized_axis += len(shape)
             if (
-                shape is not None
-                and (axis == -1 or axis == len(shape) - 1)
+                normalized_axis is not None
+                and (normalized_axis == len(shape) - 1 or tgt.kind.name == "webgpu")
                 and self.is_gpu_target(tgt)
                 and not can_use_thrust(tgt, "tvm.contrib.thrust.sum_scan")
                 and call.op.name == "relax.cumsum"
                 and call.attrs.exclusive == 0
             ):
                 from tvm.relax.backend.gpu_generic import (  # pylint: disable=import-outside-toplevel
                     gpu_2d_continuous_cumsum,
+                    gpu_3d_axis_1_cumsum,
                 )
 
-                dim = 1
-                for i in range(len(shape) - 1):
-                    dim *= shape[i]
+                input_tensor = call.args[0]
                 in_dtype = call.args[0].ty.dtype
                 out_dtype = call.attrs.dtype
                 out_dtype = out_dtype or in_dtype
-                cumsum_2d_shape = relax.ShapeExpr([dim, shape[-1]])
+
+                if normalized_axis == len(shape) - 1:
+                    outer = reduce(mul, shape_values[:-1], 1)
+                    kernel_shape = relax.ShapeExpr([outer, shape[-1]])
+                    kernel = gpu_2d_continuous_cumsum(
+                        in_dtype=in_dtype,
+                        out_dtype=out_dtype,
+                    )
+                    kernel_name = "gpu_2d_continuous_cumsum"
+                else:
+                    outer = reduce(mul, shape_values[:normalized_axis], 1)
+                    inner = reduce(mul, shape_values[normalized_axis + 1 :], 1)
+                    kernel_shape = relax.ShapeExpr([outer, shape[normalized_axis], inner])
+                    kernel = gpu_3d_axis_1_cumsum(
+                        in_dtype=in_dtype,
+                        out_dtype=out_dtype,
+                    )
+                    kernel_name = "gpu_3d_axis_1_cumsum"
+
                 reshape = relax.call_pure_packed(
                     "vm.builtin.reshape",
-                    call.args[0],
-                    cumsum_2d_shape,
-                    ty_args=relax.TensorType(cumsum_2d_shape, out_dtype),
-                )
-                gv = self.builder_.add_func(
-                    gpu_2d_continuous_cumsum(in_dtype=in_dtype, out_dtype=out_dtype),
-                    "gpu_2d_continuous_cumsum",
+                    input_tensor,
+                    kernel_shape,
+                    ty_args=relax.TensorType(kernel_shape, in_dtype, vdevice=call.ty.vdevice),
                 )
+                gv = self.builder_.add_func(kernel, kernel_name)
                 cumsum = relax.call_tir(
                     gv,
                     reshape,
-                    out_ty=relax.TensorType(cumsum_2d_shape, out_dtype),
+                    out_ty=relax.TensorType(kernel_shape, out_dtype, vdevice=call.ty.vdevice),
                 )
                 return relax.call_pure_packed(
                     "vm.builtin.reshape",

python/tvm/relax/backend/gpu_generic/__init__.py

@@ -17,7 +17,7 @@
 # under the License.
 """The Relax Metal backend compilation pipeline and other passes."""
 
-from .cumsum import gpu_2d_continuous_cumsum
+from .cumsum import gpu_2d_continuous_cumsum, gpu_3d_axis_1_cumsum
 from .pipeline import (
     dataflow_lower_passes,
     finalize_passes,

python/tvm/relax/backend/gpu_generic/cumsum.py

@@ -193,3 +193,43 @@ def cumsum(var_a: T.handle, var_out: T.handle):
         update_cross_block(m, n, Tmp, Out, src_offset=0, out_offset=0)
 
     return cumsum
+
+
+def gpu_3d_axis_1_cumsum(
+    tx_len: int = 128,
+    in_dtype: str = "int32",
+    out_dtype: str | None = None,
+) -> PrimFunc:
+    """Generate a correctness fallback that scans axis 1 of a contiguous 3D tensor.
+
+    Each thread handles one pair of outer and inner indices and scans the
+    middle axis sequentially.  The dispatcher collapses arbitrary-rank inputs
+    around the scan axis into this 3D representation.  This fallback avoids a
+    transposed scan on targets where that lowering is unavailable; it is not a
+    parallel scan optimization for rank-3 tensors.
+    """
+
+    out_dtype = out_dtype or in_dtype
+    TX = T.int64(tx_len)
+
+    @T.prim_func(private=True, s_tir=True)
+    def cumsum(var_a: T.handle, var_out: T.handle):
+        T.func_attr({"tirx.is_scheduled": True})
+        outer, scan, inner = T.int64(), T.int64(), T.int64()
+        A = T.match_buffer(var_a, [outer, scan, inner], dtype=in_dtype)
+        Out = T.match_buffer(var_out, [outer, scan, inner], dtype=out_dtype)
+
+        for bx in T.thread_binding(T.ceildiv(outer * inner, TX), thread="blockIdx.x"):
+            for tx in T.thread_binding(TX, thread="threadIdx.x"):
+                row: T.let[T.int64] = bx * TX + tx
+                with T.sblock():
+                    accumulator = T.sblock_alloc_buffer((), out_dtype, scope="local")
+                    if row < outer * inner:
+                        outer_idx: T.let[T.int64] = row // inner
+                        inner_idx: T.let[T.int64] = row % inner
+                        accumulator[()] = T.Cast(out_dtype, 0)
+                        for k in T.serial(scan):
+                            accumulator[()] += T.Cast(out_dtype, A[outer_idx, k, inner_idx])
+                            Out[outer_idx, k, inner_idx] = accumulator[()]
+
+    return cumsum

tests/python/relax/test_backend_dispatch_sort_scan.py

@@ -448,5 +448,128 @@ def run_and_check():
     tvm.testing.run_with_gpu_lock(run_and_check)
 
 
+@pytest.mark.parametrize(
+    "shape, axis, in_dtype, out_dtype, expected_kernel, expected_kernel_rank",
+    [
+        ((3, 5), 0, "float32", None, "gpu_3d_axis_1_cumsum", 3),
+        ((2, 3, 4, 5), 1, "float32", None, "gpu_3d_axis_1_cumsum", 3),
+        ((2, 3, 4, 5), -2, "int32", "float32", "gpu_3d_axis_1_cumsum", 3),
+        ((2, 3, 4, 5), -1, "float32", None, "gpu_2d_continuous_cumsum", 2),
+    ],
+)
+def test_dispatch_cumsum_webgpu_axes_and_dtypes(
+    shape, axis, in_dtype, out_dtype, expected_kernel, expected_kernel_rank
+):
+    """WebGPU dispatch collapses arbitrary-rank scans to the appropriate kernel."""
+
+    vdevice = tvm.ir.VDevice("webgpu", 0)
+    x = relax.Var("x", relax.TensorType(shape, in_dtype, vdevice=vdevice))
+    bb = relax.BlockBuilder()
+    with bb.function("main", (x,)):
+        out = bb.emit(relax.op.cumsum(x, axis=axis, dtype=out_dtype))
+        bb.emit_func_output(out)
+    before = bb.finalize()
+    before.update_global_info("vdevice", [vdevice])
+
+    target = tvm.target.Target("webgpu", host="llvm")
+    with target:
+        mod = DispatchSortScan()(before)
+
+    called_kernels = []
+    permute_count = 0
+
+    def collect_calls(expr):
+        nonlocal permute_count
+        if isinstance(expr, relax.Call) and getattr(expr.op, "name", None) == (
+            "relax.permute_dims"
+        ):
+            permute_count += 1
+        if isinstance(expr, relax.Call) and getattr(expr.op, "name", None) == "relax.call_tir":
+            called_kernels.append(expr.args[0].name_hint)
+
+    relax.analysis.post_order_visit(mod["main"], collect_calls)
+    assert permute_count == 0
+    assert called_kernels == [expected_kernel]
+
+    cumsum = mod[expected_kernel]
+    buffers = [param for param in cumsum.params if tvm.tirx.is_buffer_var(param)]
+    assert len(buffers) == 2
+    assert all(len(buffer.shape) == expected_kernel_rank for buffer in buffers)
+    assert str(buffers[0].dtype) == in_dtype
+    assert str(buffers[1].dtype) == (out_dtype or in_dtype)
+
+    with target:
+        tvm.compile(mod, target)
+
+
+def test_dispatch_cumsum_webgpu_symbolic_non_contiguous_axis():
+    """The serial WebGPU fallback accepts a symbolic scan extent."""
+
+    @I.ir_module
+    class Symbolic:
+        I.module_global_infos({"vdevice": [I.vdevice("webgpu", 0)]})
+
+        @R.function
+        def main(x: R.Tensor((1, "n", 9), "float32", "webgpu")):
+            return R.cumsum(x, axis=1)
+
+    target = tvm.target.Target("webgpu", host="llvm")
+    with target:
+        mod = DispatchSortScan()(Symbolic)
+        tvm.compile(mod, target)
+
+    called_kernels = []
+
+    def collect_calls(expr):
+        if isinstance(expr, relax.Call) and getattr(expr.op, "name", None) == "relax.call_tir":
+            called_kernels.append(expr.args[0].name_hint)
+
+    relax.analysis.post_order_visit(mod["main"], collect_calls)
+    assert called_kernels == ["gpu_3d_axis_1_cumsum"]
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        pytest.param("cuda", marks=pytest.mark.gpu),
+        pytest.param({"kind": "vulkan", "supports_int64": True}, marks=pytest.mark.gpu),
+        pytest.param("metal", marks=pytest.mark.gpu),
+    ],
+)
+@pytest.mark.parametrize(
+    "in_dtype, out_dtype",
+    [("float32", "float32"), ("int32", "int32"), ("int32", "float32")],
+)
+def test_gpu_axis_1_cumsum_numerical(target, in_dtype, out_dtype):
+    """The fallback matches a sequential cumsum for supported WebGPU dtypes."""
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    from tvm.relax.backend.gpu_generic import (  # pylint: disable=import-outside-toplevel
+        gpu_3d_axis_1_cumsum,
+    )
+
+    shape = (2, 5, 7)
+    if in_dtype == "int32":
+        np_data = np.random.randint(-4, 5, shape).astype(in_dtype)
+    else:
+        np_data = np.random.uniform(-2, 2, shape).astype(in_dtype)
+    expected = np.cumsum(np_data, axis=1, dtype=out_dtype)
+
+    func = gpu_3d_axis_1_cumsum(in_dtype=in_dtype, out_dtype=out_dtype).with_attr(
+        "global_symbol", "main"
+    )
+    compiled = tvm.compile(func, target=target)
+
+    def run_and_check():
+        dev = tvm.device_from_target(target)
+        input_tensor = tvm.runtime.tensor(np_data, dev)
+        output_tensor = tvm.runtime.empty(shape, out_dtype, dev)
+        compiled(input_tensor, output_tensor)
+        tvm.testing.assert_allclose(output_tensor.numpy(), expected)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
 if __name__ == "__main__":
     tvm.testing.main()