gemini-code-assist[bot] commented on code in PR #19900:
URL: https://github.com/apache/tvm/pull/19900#discussion_r3488150416
##########
python/tvm/topi/gpu/sort.py:
##########
@@ -499,17 +499,17 @@ def dual_mergepath(
nbx = tvm.tirx.generic.cast(ceil_div(width, max_threads *
thread_work), "int32")
nbz = tvm.tirx.generic.cast(ceil_div(size, width), "int32")
- tx, bx, by, _, _, _ = _get_threads(ntx, nbx, nthread_by * nbz)
+ tx, bx, by, _, _, _ = _get_threads(ntx, nbx, nthread_by)
+ bz = te.thread_axis("blockIdx.z")
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),
]
Review Comment:
In CUDA, the maximum limit for `blockIdx.y` and `blockIdx.z` is 65,535,
while `blockIdx.x` can go up to $2^{31} - 1$ (2,147,483,647).
With the current changes in this PR:
- `bx` (extent `nbx`) is bound to `blockIdx.x`
- `by` (extent `nthread_by`) is bound to `blockIdx.y`
- `bz` (extent `nbz`) is bound to `blockIdx.z`
Since `nbz = ceil_div(size, width)` and `width` starts at `block_size`
(128), at the beginning of the merge sort, `nbz` can be up to `ceil_div(size,
128)`. If `size` exceeds `128 * 65,535 = 8,388,480` (approx 8.38 million
elements), `nbz` will exceed 65,535, causing a `blockIdx.z` overflow and kernel
launch failure.
To support much larger sizes without overflow, we can swap the mapping of
`bx` and `bz`:
- Bind `bz` (extent `nbz`) to `blockIdx.x` (limit 2,147,483,647)
- Bind `bx` (extent `nbx`) to `blockIdx.z` (limit 65,535)
Since `nbx = ceil_div(width, max_threads * thread_work) <= ceil_div(size,
4096)`, `nbx` will only exceed 65,535 if `size` exceeds `4096 * 65,535 =
268,435,456` (approx 268 million elements). This simple swap increases the
maximum supported sort size by 32x (from 8.38M to 268M elements) without adding
any division/modulo overhead!
```suggestion
tx = te.thread_axis("threadIdx.x")
bx = te.thread_axis("blockIdx.z")
by = te.thread_axis("blockIdx.y")
bz = te.thread_axis("blockIdx.x")
with T.frame_scope(
[
T.attr(tx, "thread_extent", ntx),
T.attr(bx, "thread_extent", nbx),
T.attr(by, "thread_extent", nthread_by),
T.attr(bz, "thread_extent", nbz),
]
):
```
--
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]
