[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.

[cchung100m]

cc @mshr-h @tlopex