gemini-code-assist[bot] commented on code in PR #19921:
URL: https://github.com/apache/tvm/pull/19921#discussion_r3503736557
##########
python/tvm/topi/gpu/sort.py:
##########
@@ -30,9 +30,12 @@
def _get_threads(nthread_tx, nthread_bx, nthread_by):
+ # Swap blockIdx.x and blockIdx.y mappings.
+ # The batch dimension (nthread_by) can easily exceed the CUDA limit of
65535 for gridDim.y.
+ # By mapping nthread_by to blockIdx.x (limit ~2 billion), we prevent
CUDA_ERROR_INVALID_VALUE on large inputs.
tx = te.thread_axis("threadIdx.x")
- bx = te.thread_axis("blockIdx.x")
- by = te.thread_axis("blockIdx.y")
+ bx = te.thread_axis("blockIdx.y")
+ by = te.thread_axis("blockIdx.x")
Review Comment:
While swapping `blockIdx.x` and `blockIdx.y` successfully prevents crashes
when the batch dimension (`nthread_by`) is large, it introduces a new crash
risk when the sorted dimension (`size` or `shape[axis]`) is large.\n\nOn CUDA,
`gridDim.y` has a strict hardware limit of 65,535, whereas `gridDim.x` can go
up to 2^31 - 1. By mapping `bx` (the block index for the sorted dimension) to
`blockIdx.y`, we are now bound by the 65,535 limit on `bx`.\n\nThis can easily
be exceeded in practice:\n1. In `_odd_even_sort` (line 99), `nthread_bx` is
computed as `ceil_div(size, block_size)`. Since `block_size = 128`, any input
where the sorted dimension `size` exceeds `65535 * 128 = 8,388,480` (approx
8.4M elements) will result in `nthread_bx > 65535`, causing a
`CUDA_ERROR_INVALID_VALUE` crash during kernel launch.\n2. In `_sort_init`
(line 57), `nthread_bx = ceil_div(shape[axis], max_threads)`. If `max_threads =
1024`, sorting a dimension of size > 67,107,840 will also crash.\n3. In
`_sort_comm
on` (line 502), `nbx = ceil_div(width, max_threads * thread_work)`. If `size >
268,431,360`, it will crash.\n\nTo make this robust for both large batch
dimensions and large sorted dimensions, consider fusing the block dimensions
into a single 1D grid on `blockIdx.x` (with extent `nthread_bx * nthread_by`),
and then reconstructing the virtual `bx` and `by` indices inside the kernel
using division and modulo (e.g., `bx = bx_fused % nthread_bx` and `by =
bx_fused // nthread_bx`).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
