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 7924acaef4 [TOPI][CUDA] Fix topk/sort gridDim overflow by remapping
grid axes in sort_ir (#19900)
7924acaef4 is described below
commit 7924acaef4277042b6ba981d5b9d44e0611f9610
Author: Neo Chien <[email protected]>
AuthorDate: Mon Sep 14 11:54:58 2026 +0800
[TOPI][CUDA] Fix topk/sort gridDim overflow by remapping grid axes in
sort_ir (#19900)
Hi Committers,
This PR addresses issue https://github.com/apache/tvm/issues/19549. Any
suggestions would be appreciated if you are available.
---------
Co-authored-by: cchung100m <[email protected]>
---
python/tvm/topi/gpu/sort.py | 27 ++++++---
.../relax/test_backend_dispatch_sort_scan.py | 67 ++++++++++++++++++++++
2 files changed, 87 insertions(+), 7 deletions(-)
diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py
index 2ca4056438..eabed85a9b 100644
--- a/python/tvm/topi/gpu/sort.py
+++ b/python/tvm/topi/gpu/sort.py
@@ -30,9 +30,15 @@ from ..utils import ceil_div, prod, swap
def _get_threads(nthread_tx, nthread_bx, nthread_by):
+ target = tvm.target.Target.current(allow_none=True)
+ is_cuda = target is not None and target.kind.name == "cuda"
tx = te.thread_axis("threadIdx.x")
- bx = te.thread_axis("blockIdx.x")
- by = te.thread_axis("blockIdx.y")
+ if is_cuda:
+ bx = te.thread_axis("blockIdx.y")
+ by = te.thread_axis("blockIdx.x")
+ else:
+ bx = te.thread_axis("blockIdx.x")
+ by = te.thread_axis("blockIdx.y")
return tx, bx, by, nthread_tx, nthread_bx, nthread_by
@@ -501,17 +507,24 @@ def _sort_common(
nbx = cast(ceil_div(width, max_threads * thread_work), "int32")
nbz = cast(ceil_div(size, width), "int32")
- tx, bx, by, _, _, _ = _get_threads(ntx, nbx, nthread_by * nbz)
+ is_cuda = target.kind.name == "cuda"
+ tx = te.thread_axis("threadIdx.x")
+ bx = te.thread_axis("blockIdx.z") # nbx
+ if is_cuda:
+ by = te.thread_axis("blockIdx.x") # batch
+ bz = te.thread_axis("blockIdx.y") # nbz
+ else:
+ by = te.thread_axis("blockIdx.y") # batch
+ bz = te.thread_axis("blockIdx.x") # nbz
with T.frame_scope(
[
T.attr(tx, "thread_extent", ntx),
T.attr(bx, "thread_extent", nbx),
- T.attr(by, "thread_extent", nthread_by * nbz),
+ T.attr(by, "thread_extent", nthread_by),
+ T.attr(bz, "thread_extent", nbz),
]
):
- by_val = by % nthread_by
- bz = by // nthread_by
- base_idx = by_val * size
+ base_idx = by * size
# calculate the start, mid, and end points of this section
start_pos = width * bz
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
index df4aca033a..c2dc7ef66a 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -411,6 +411,73 @@ def test_dispatch_topk_gpu():
assert_structural_equal(mod, expected_mod)
[email protected]
[email protected]("size", [300, 600])
+def test_dispatch_sort_cuda_large_batch(size):
+ 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.sort(x, axis=-1, descending=False)
+ R.output(gv)
+ return gv
+
+ batch = 65600 # > CUDA's 65535 gridDim.y/z limit
+ np_data = np.random.uniform(size=(batch, size)).astype("float32")
+ np_sorted = np.sort(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)
+ sorted_data = vm["main"](tvm_data)
+ tvm.testing.assert_allclose(sorted_data.numpy(), np_sorted)
+
+ tvm.testing.run_with_gpu_lock(run_and_check)
+
+
[email protected]
+def test_dispatch_topk_cuda_large_batch():
+ 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.topk(x, k=1, axis=-1, ret_type="values", largest=True)
+ R.output(gv)
+ return gv
+
+ batch, size = 65600, 300
+ np_data = np.random.uniform(size=(batch, size)).astype("float32")
+ np_values = np.sort(np_data, axis=-1)[:, -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)
+ values = vm["main"](tvm_data)
+ tvm.testing.assert_allclose(values.numpy(), np_values)
+
+ tvm.testing.run_with_gpu_lock(run_and_check)
+
+
@pytest.mark.parametrize(
"target",
[
