diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py index 76951718c497..d4f3e1319f12 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,43 @@ 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, + ) + 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..630b7fe63da3 100644 --- a/python/tvm/relax/backend/gpu_generic/cumsum.py +++ b/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 diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py b/tests/python/relax/test_backend_dispatch_sort_scan.py index 2cab14f36a26..479909d941ae 100644 --- a/tests/python/relax/test_backend_dispatch_sort_scan.py +++ b/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()