This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 1f352bf9c1 [Fix][Relax] Lower non-contiguous WebGPU cumsum (#20133)
1f352bf9c1 is described below
commit 1f352bf9c1af19b9f43584a5d27e9a14c30b468b
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Mon Aug 24 19:43:23 2026 -0700
[Fix][Relax] Lower non-contiguous WebGPU cumsum (#20133)
This PR adds direct WebGPU lowering for inclusive `cumsum` along
non-innermost axes. Previously, these operations used TOPI's
transpose-based GPU scan, which materializes a full-size transposed
tensor and requires multiple WebGPU dispatches. This is particularly
problematic when the scan extent is symbolic.
The new lowering:
- Normalizes the scan axis and reshapes known-rank inputs to `[outer,
scan, inner]` without transposing data
- Uses a correctness fallback where each GPU thread scans one `(outer,
inner)` row serially
- Preserves input/output dtypes and virtual-device information
- Retains the existing parallel 2D kernel for innermost-axis scans
- Clamps the existing kernel's reverse-round count so short scans do not
generate negative shifts during WebGPU index narrowing
---
python/tvm/relax/backend/dispatch_sort_scan.py | 47 ++++---
python/tvm/relax/backend/gpu_generic/__init__.py | 2 +-
python/tvm/relax/backend/gpu_generic/cumsum.py | 104 +++++++++++++--
.../relax/test_backend_dispatch_sort_scan.py | 140 +++++++++++++++++++++
4 files changed, 266 insertions(+), 27 deletions(-)
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py
b/python/tvm/relax/backend/dispatch_sort_scan.py
index 76951718c4..df1543761d 100644
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -146,10 +146,14 @@ class SortScanDispatcher(BackendDispatcher):
# 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 @@ class SortScanDispatcher(BackendDispatcher):
):
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 3e8bd0f8f0..bb67a4721c 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 9676131f46..9d79cf7079 100644
--- a/python/tvm/relax/backend/gpu_generic/cumsum.py
+++ b/python/tvm/relax/backend/gpu_generic/cumsum.py
@@ -28,12 +28,22 @@ def _is_power_of_two(n: int):
return n > 0 and (n & (n - 1)) == 0
+def _get_total_rounds(n, log_block_n: int, index_bits: int):
+ """Count the hierarchy levels without introducing integer division."""
+ total_rounds = T.int64(0)
+ for round_index in range(1, (index_bits - 1) // log_block_n + 1):
+ threshold = T.int64(1 << (round_index * log_block_n - 1))
+ total_rounds += T.Cast("int64", n > threshold)
+ return total_rounds
+
+
def gpu_2d_continuous_cumsum(
ty_len: int = 4,
tx_len: int = 32,
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 +64,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 +82,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)
@@ -76,12 +91,15 @@ def gpu_2d_continuous_cumsum(
block_elem = T.int64(tx_len * ty_len * thread_elem)
LOG_TX = T.int64(int(math.log2(tx_len)))
- LOG_BLOCK_N = T.int64(int(math.log2(tx_len * ty_len * thread_elem)))
+ log_block_n = int(math.log2(tx_len * ty_len * thread_elem))
+ LOG_BLOCK_N = T.int64(log_block_n)
+ 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 +107,7 @@ def gpu_2d_continuous_cumsum(
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 +156,14 @@ def gpu_2d_continuous_cumsum(
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):
@@ -161,35 +180,96 @@ def gpu_2d_continuous_cumsum(
A = T.match_buffer(var_a, [m, n], dtype=in_dtype)
Out = T.match_buffer(var_out, [m, n], dtype=out_dtype)
Tmp = T.alloc_buffer([m, n], dtype=out_dtype)
- total_rounds: T.let[T.int64] = (
- T.Cast("int64", T.ceil(T.log2(T.Cast("float32", n)))) //
LOG_BLOCK_N
- )
+ # LowerIntrin may implement signed FloorDiv using a sign-bit shift.
Keep
+ # hierarchy counting division-free so WebGPU can narrow indices to
int32.
+ total_rounds: T.let[T.int64] = _get_total_rounds(n, log_block_n,
index_bits)
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 9965b58c15..7fcd264917 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -484,5 +484,145 @@ def test_dispatch_cumprod_cuda_large_batch():
tvm.testing.run_with_gpu_lock(run_and_check)
[email protected](
+ "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)
+
+ if expected_kernel == "gpu_2d_continuous_cumsum":
+ floor_divisors = []
+
+ def collect_floor_divisors(node):
+ if isinstance(node, tirx.FloorDiv):
+ floor_divisors.append(node.b)
+
+ tirx.stmt_functor.post_order_visit(cumsum.body, collect_floor_divisors)
+ assert floor_divisors
+ assert all(
+ isinstance(divisor, tirx.IntImm)
+ and divisor.value > 0
+ and divisor.value & (divisor.value - 1) == 0
+ for divisor in floor_divisors
+ )
+
+ 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"]
+
+
[email protected](
+ "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),
+ ],
+)
[email protected](
+ "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()
