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 50c9ca7ee1 [Fix][TOPI] Fuse GPU scan blocks to avoid CUDA gridDim.y
overflow (#20108)
50c9ca7ee1 is described below
commit 50c9ca7ee16a6e65cd47a6e55685181663859bef
Author: MiaoMing Chen <[email protected]>
AuthorDate: Thu Aug 20 08:12:48 2026 +0800
[Fix][TOPI] Fuse GPU scan blocks to avoid CUDA gridDim.y overflow (#20108)
## Motivation
Fixes #20106.
The fallback TOPI GPU scan maps the per-row scan block to `blockIdx.x`
and the
batch to `blockIdx.y`. CUDA limits `gridDim.y` to 65,535, so a valid
Relax
`cumprod` with shape `(65536, 1)` fails at execution time with
`CUDA_ERROR_INVALID_VALUE` and `grid=(1,65536,1)`.
## Changes
Fuse the virtual scan-block and batch dimensions into a single
`blockIdx.x`
launch dimension, then recover both indices with integer division and
modulo.
Apply this mapping to the initial copy, up-sweep, and down-sweep stages.
Fusing
the dimensions fixes large batches without moving the same 65,535 limit
onto
long scan axes.
Add a dynamic-shape CUDA `cumprod` regression with an input of shape
`(65536, 3)`. The three-element scan axis exercises the up-sweep and
down-sweep
stages while retaining the batch size that previously caused the invalid
launch.
## Testing
- `python -m pytest
tests/python/relax/test_backend_dispatch_sort_scan.py -xvs`
- 10 passed, 1 skipped, 1 pre-existing xpassed
- Manual CUDA checks for the issue shape `(65536, 1)` and a multi-block
scan
shape `(8, 2000)`
- `pre-commit run --files python/tvm/topi/gpu/scan.py
tests/python/relax/test_backend_dispatch_sort_scan.py`
---
python/tvm/topi/gpu/scan.py | 48 +++++++++++-----------
.../relax/test_backend_dispatch_sort_scan.py | 36 ++++++++++++++++
2 files changed, 60 insertions(+), 24 deletions(-)
diff --git a/python/tvm/topi/gpu/scan.py b/python/tvm/topi/gpu/scan.py
index a0b35bcb19..e576e7d3e8 100644
--- a/python/tvm/topi/gpu/scan.py
+++ b/python/tvm/topi/gpu/scan.py
@@ -102,25 +102,25 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=operator.add, identity
reduction[bx] = cast(identity_value, out_dtype)
with T.Else():
nthread_tx = max_threads
- nthread_bx = ceil_div(scan_axis_size, max_threads)
- nthread_by = batch_size
+ blocks_per_batch = ceil_div(scan_axis_size, max_threads)
+ # Flatten the batch and scan-block axes into blockIdx.x. On
CUDA,
+ # blockIdx.y is limited to 65535 even when blockIdx.x can be
much larger.
# Copy data to output
tx = te.thread_axis("threadIdx.x")
bx = te.thread_axis("blockIdx.x")
- by = te.thread_axis("blockIdx.y")
with T.frame_scope(
[
T.attr(tx, "thread_extent", nthread_tx),
- T.attr(bx, "thread_extent", nthread_bx),
- T.attr(by, "thread_extent", nthread_by),
+ T.attr(bx, "thread_extent", blocks_per_batch *
batch_size),
]
):
- tid = bx * nthread_tx + tx
+ batch = tvm.tirx.indexdiv(bx, blocks_per_batch)
+ tid = tvm.tirx.indexmod(bx, blocks_per_batch) * nthread_tx
+ tx
with T.If(tid < scan_axis_size):
with T.Then():
- output[by * scan_axis_size + tid] = cast(
- data[by * scan_axis_size + tid], out_dtype
+ output[batch * scan_axis_size + tid] = cast(
+ data[batch * scan_axis_size + tid], out_dtype
)
# The following algorithm performs parallel exclusive scan
@@ -132,7 +132,7 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=operator.add, identity
tx = te.thread_axis("threadIdx.x")
bx = te.thread_axis("blockIdx.x")
- by = te.thread_axis("blockIdx.y")
+ blocks_per_batch = cast(ceil_div(scan_axis_size,
max_threads * width), "int32")
start_buf = T.decl_buffer([1], "int32", scope="local")
middle_buf = T.decl_buffer([1], "int32", scope="local")
end_buf = T.decl_buffer([1], "int32", scope="local")
@@ -142,12 +142,12 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=operator.add, identity
T.attr(
bx,
"thread_extent",
- cast(ceil_div(scan_axis_size, max_threads *
width), "int32"),
+ blocks_per_batch * batch_size,
),
- T.attr(by, "thread_extent", nthread_by),
]
):
- tid = bx * nthread_tx + tx
+ batch = tvm.tirx.indexdiv(bx, blocks_per_batch)
+ tid = tvm.tirx.indexmod(bx, blocks_per_batch) *
nthread_tx + tx
start = T.buffer_proxy(start_buf)
middle = T.buffer_proxy(middle_buf)
end = T.buffer_proxy(end_buf)
@@ -158,9 +158,9 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=operator.add, identity
end[0] = tvm.te.min(start[0] + width,
scan_axis_size)
with T.If(middle[0] < scan_axis_size):
with T.Then():
- output[by * scan_axis_size + end[0] -
1] = binop(
- output[by * scan_axis_size +
end[0] - 1],
- output[by * scan_axis_size +
middle[0] - 1],
+ output[batch * scan_axis_size + end[0]
- 1] = binop(
+ output[batch * scan_axis_size +
end[0] - 1],
+ output[batch * scan_axis_size +
middle[0] - 1],
)
# Down Sweep of exclusive scan
@@ -177,7 +177,7 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=operator.add, identity
tx = te.thread_axis("threadIdx.x")
bx = te.thread_axis("blockIdx.x")
- by = te.thread_axis("blockIdx.y")
+ blocks_per_batch = cast(ceil_div(scan_axis_size,
max_threads * width), "int32")
start_buf = T.decl_buffer([1], "int32", scope="local")
middle_buf = T.decl_buffer([1], "int32", scope="local")
end_buf = T.decl_buffer([1], "int32", scope="local")
@@ -188,12 +188,12 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=operator.add, identity
T.attr(
bx,
"thread_extent",
- cast(ceil_div(scan_axis_size, max_threads *
width), "int32"),
+ blocks_per_batch * batch_size,
),
- T.attr(by, "thread_extent", nthread_by),
]
):
- tid = bx * nthread_tx + tx
+ batch = tvm.tirx.indexdiv(bx, blocks_per_batch)
+ tid = tvm.tirx.indexmod(bx, blocks_per_batch) *
nthread_tx + tx
start = T.buffer_proxy(start_buf)
middle = T.buffer_proxy(middle_buf)
end = T.buffer_proxy(end_buf)
@@ -205,12 +205,12 @@ def exclusive_scan_ir(data, output, reduction=None,
binop=operator.add, identity
end[0] = tvm.tirx.min(start[0] + width,
scan_axis_size)
with T.If(middle[0] < scan_axis_size):
with T.Then():
- tmp[0] = output[by * scan_axis_size +
middle[0] - 1]
- output[by * scan_axis_size + middle[0]
- 1] = output[
- by * scan_axis_size + end[0] - 1
+ tmp[0] = output[batch * scan_axis_size
+ middle[0] - 1]
+ output[batch * scan_axis_size +
middle[0] - 1] = output[
+ batch * scan_axis_size + end[0] - 1
]
- output[by * scan_axis_size + end[0] -
1] = binop(
- output[by * scan_axis_size +
end[0] - 1], tmp[0]
+ output[batch * scan_axis_size + end[0]
- 1] = binop(
+ output[batch * scan_axis_size +
end[0] - 1], tmp[0]
)
return ib.get()
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
index 2cab14f36a..9965b58c15 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -448,5 +448,41 @@ def test_dispatch_cumsum_gpu(target):
tvm.testing.run_with_gpu_lock(run_and_check)
[email protected]
+def test_dispatch_cumprod_cuda_large_batch():
+ """Test that GPU scan supports more batches than CUDA's grid-y limit."""
+ target = "cuda"
+ if not tvm.testing.device_enabled(target):
+ pytest.skip(f"{target} not enabled")
+
+ @I.ir_module
+ class Module:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")):
+ with R.dataflow():
+ gv = R.cumprod(x, axis=1)
+ R.output(gv)
+ return gv
+
+ np_data = np.ones((65536, 3), dtype="float32")
+ np_data[:, 0] = np.arange(65536) % 7 + 1
+ np_data[:, 1] = 2
+ np_data[:, 2] = 3
+ np_cumprod = np.cumprod(np_data, axis=1)
+
+ with tvm.target.Target(target):
+ mod = DispatchSortScan()(Module)
+ ex = tvm.compile(mod, target)
+
+ def run_and_check():
+ dev = tvm.device_from_target(target)
+ vm = tvm.relax.VirtualMachine(ex, dev)
+ tvm_data = tvm.runtime.tensor(np_data, dev)
+ cumprod = vm["main"](tvm_data)
+ tvm.testing.assert_allclose(cumprod.numpy(), np_cumprod)
+
+ tvm.testing.run_with_gpu_lock(run_and_check)
+
+
if __name__ == "__main__":
tvm.testing.main()
