From ba185119a9e4c67aaab7876cb875e33b8212832c Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Tue, 11 Aug 2026 07:25:26 -0400 Subject: [PATCH] [FIX][RELAX] Lower non-contiguous WebGPU cumsum --- .../tvm/relax/backend/dispatch_sort_scan.py | 47 +++++-- .../tvm/relax/backend/gpu_generic/__init__.py | 2 +- .../tvm/relax/backend/gpu_generic/cumsum.py | 86 ++++++++++-- .../relax/test_backend_dispatch_sort_scan.py | 124 ++++++++++++++++++ 4 files changed, 236 insertions(+), 23 deletions(-) diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py index 76951718c497..df1543761de4 100644 --- a/python/tvm/relax/backend/dispatch_sort_scan.py +++ b/python/tvm/relax/backend/dispatch_sort_scan.py @@ -146,10 +146,14 @@ 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" @@ -157,29 +161,44 @@ def visit_call_(self, call: relax.Call) -> relax.Expr: ): 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, + index_bits=32 if tgt.kind.name == "webgpu" else 64, + ) + 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", diff --git a/python/tvm/relax/backend/gpu_generic/__init__.py b/python/tvm/relax/backend/gpu_generic/__init__.py index 3e8bd0f8f07f..bb67a4721c96 100644 --- a/python/tvm/relax/backend/gpu_generic/__init__.py +++ b/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, diff --git a/python/tvm/relax/backend/gpu_generic/cumsum.py b/python/tvm/relax/backend/gpu_generic/cumsum.py index 9676131f46de..c97aab463aa0 100644 --- a/python/tvm/relax/backend/gpu_generic/cumsum.py +++ b/python/tvm/relax/backend/gpu_generic/cumsum.py @@ -34,6 +34,7 @@ def gpu_2d_continuous_cumsum( thread_elem: int = 4, in_dtype: str = "int32", out_dtype: str | None = None, + index_bits: int = 64, ) -> PrimFunc: """Generate GPU kernel for 2D continuous cumsum, i.e. The cumsum axis is -1 @@ -54,6 +55,9 @@ def gpu_2d_continuous_cumsum( out_dtype : Optional[str] The output data type, if None, it will be the same as in_dtype + index_bits : int + The number of bits available for signed index expressions + Returns ------- cumsum : PrimFunc @@ -69,6 +73,8 @@ def gpu_2d_continuous_cumsum( if not _is_power_of_two(TX) or not _is_power_of_two(TY) or not _is_power_of_two(N): raise ValueError("Configuration of TX, TY, N must be power of 2") + if index_bits not in (32, 64): + raise ValueError("index_bits must be either 32 or 64") # number of elements to be processed by single warp warp_elem = T.int64(tx_len * thread_elem) @@ -77,11 +83,13 @@ def gpu_2d_continuous_cumsum( LOG_TX = T.int64(int(math.log2(tx_len))) LOG_BLOCK_N = T.int64(int(math.log2(tx_len * ty_len * thread_elem))) + MAX_INDEX_SHIFT = T.int64(index_bits - 1) @T.macro def block_inclusive_inside_block( batch: T.int64, cur_len: T.int64, + num_blocks: T.int64, source: T.Buffer, output: T.Buffer, tmp_buf: T.Buffer, @@ -89,7 +97,7 @@ def block_inclusive_inside_block( tmp_offset: T.int64, ): for by in T.thread_binding(batch, thread="blockIdx.y"): - for bx in T.thread_binding(T.ceildiv(cur_len, block_elem), thread="blockIdx.x"): + for bx in T.thread_binding(num_blocks, thread="blockIdx.x"): with T.sblock(): local_buf = T.sblock_alloc_buffer((thread_elem,), out_dtype, scope="local") shared_buf = T.sblock_alloc_buffer((block_elem,), out_dtype, scope="shared") @@ -138,13 +146,14 @@ def block_inclusive_inside_block( def update_cross_block( batch: T.int64, cur_len: T.int64, + num_blocks: T.int64, source: T.Buffer, output: T.Buffer, src_offset: T.int64, out_offset: T.int64, ): for by in T.thread_binding(batch, thread="blockIdx.y"): - for bx in T.thread_binding(T.ceildiv(cur_len, block_elem), thread="blockIdx.x"): + for bx in T.thread_binding(num_blocks, thread="blockIdx.x"): for ty in T.thread_binding(TY, thread="threadIdx.y"): for tx in T.thread_binding(TX, thread="threadIdx.x"): for i in T.serial(N): @@ -166,30 +175,91 @@ def cumsum(var_a: T.handle, var_out: T.handle): ) block_inclusive_inside_block( - m, n, A, Out, Tmp, src_offset=T.int64(0), tmp_offset=T.int64(0) + m, + n, + T.ceildiv(n, block_elem), + A, + Out, + Tmp, + src_offset=T.int64(0), + tmp_offset=T.int64(0), ) for i in range(total_rounds): - cur_len: T.let[T.int64] = T.ceildiv(n, 1 << (LOG_BLOCK_N * (i + 1))) + shift: T.let[T.int64] = T.min(T.max(LOG_BLOCK_N * (i + 1), T.int64(0)), MAX_INDEX_SHIFT) + block_shift: T.let[T.int64] = T.min(shift + LOG_BLOCK_N, MAX_INDEX_SHIFT) + # n is non-negative and WebGPU indices are narrowed to signed int32. + # Spell out positive ceildiv by a power of two so lowering does not + # introduce an int64 sign-bit test (`remainder >> 63`). + cur_len: T.let[T.int64] = ((n - 1) >> shift) + 1 + num_blocks: T.let[T.int64] = ((n - 1) >> block_shift) + 1 block_inclusive_inside_block( m, cur_len, + num_blocks, Tmp, Tmp, Tmp, src_offset=i * T.ceildiv(n, block_elem), tmp_offset=(i + 1) * T.ceildiv(n, block_elem), ) - for i in range(total_rounds - 1): - real_idx: T.let[T.int64] = total_rounds - 1 - i - 1 - cur_len: T.let[T.int64] = T.ceildiv(n, 1 << (LOG_BLOCK_N * (real_idx + 1))) + reverse_rounds: T.let[T.int64] = T.max(total_rounds - 1, 0) + for i in range(reverse_rounds): + real_idx: T.let[T.int64] = reverse_rounds - 1 - i + shift: T.let[T.int64] = T.min( + T.max(LOG_BLOCK_N * (real_idx + 1), T.int64(0)), MAX_INDEX_SHIFT + ) + block_shift: T.let[T.int64] = T.min(shift + LOG_BLOCK_N, MAX_INDEX_SHIFT) + cur_len: T.let[T.int64] = ((n - 1) >> shift) + 1 + num_blocks: T.let[T.int64] = ((n - 1) >> block_shift) + 1 update_cross_block( m, cur_len, + num_blocks, Tmp, Tmp, src_offset=(real_idx + 1) * T.ceildiv(n, block_elem), out_offset=real_idx * T.ceildiv(n, block_elem), ) - update_cross_block(m, n, Tmp, Out, src_offset=0, out_offset=0) + update_cross_block(m, n, T.ceildiv(n, block_elem), 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 diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py b/tests/python/relax/test_backend_dispatch_sort_scan.py index 9965b58c159e..d37eaeae74de 100644 --- a/tests/python/relax/test_backend_dispatch_sort_scan.py +++ b/tests/python/relax/test_backend_dispatch_sort_scan.py @@ -484,5 +484,129 @@ 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), + # A short scan keeps total_rounds at zero in gpu_2d_continuous_cumsum. + ((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()