[TOPI][CUDA] Fix topk/sort gridDim overflow by remapping grid axes in sort_ir

[cchung100m]

Hi Committers,

This PR addresses issue #19549. Any suggestions would be appreciated if you are available.

[gemini-code-assist]

Code Review

This pull request refactors the GPU merge sort implementation in python/tvm/topi/gpu/sort.py by introducing a separate blockIdx.z thread axis (bz) instead of packing dimensions into blockIdx.y and manually unpacking them using division and modulo. Feedback on this change highlights a potential CUDA overflow issue: since blockIdx.z has a limit of 65,535, sorting arrays larger than ~8.38 million elements will cause kernel launch failures. To resolve this, the reviewer suggested swapping the thread axis mappings so that bz is bound to blockIdx.x (which has a much higher limit) and bx is bound to blockIdx.z, safely extending the maximum supported size to approximately 268 million elements.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

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!

Suggested change

	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),
	]
	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),
	]
	):