Skip to content

[TIRx][CUDA] Support exact required block dimensions - #20223

Merged
jinhongyii merged 1 commit into
apache:mainfrom
spectrometerHBH:feat/tirx-required-block-size
Aug 29, 2026
Merged

jinhongyii merged 1 commit into
apache:mainfrom
spectrometerHBH:feat/tirx-required-block-size

Conversation

@spectrometerHBH

Copy link
Copy Markdown
Contributor

This adds an exact CUDA launch contract for TIRx kernels that need PTX .reqntid rather than advisory launch bounds.

  • preserve tirx.required_block_size through host/device splitting
  • emit CUDA 13 __block_size__ with static thread and cluster dimensions
  • pass the required-block flag through runtime metadata without adding a packed operand
  • launch with CU_LAUNCH_KERNEL_REQUIRED_BLOCK_DIM and avoid duplicate cluster attributes
  • reject incompatible launch-bound and max-register controls

Validation:

  • cmake --build build --parallel
  • python -m pytest -q tests/python/tirx-transform/test_tir_transform_split_host_device.py tests/python/tirx/codegen/test_codegen_cuda.py (225 passed)
  • pre-commit run --files include/tvm/tirx/function.h src/backend/cuda/codegen/codegen_cuda.cc src/backend/cuda/runtime/cuda_module.cc src/runtime/metadata.h src/runtime/thread_storage_scope.h src/tirx/transform/split_host_device.cc tests/python/tirx-transform/test_tir_transform_split_host_device.py tests/python/tirx/codegen/test_codegen_cuda.py
  • CUDA 13.2 SM100 runtime launch through a TIRx kernel with generated PTX .reqntid 128,1,1

@jinhongyii
jinhongyii merged commit 0d9c938 into apache:main Aug 29, 2026
8 checks passed
tqchen pushed a commit that referenced this pull request Oct 16, 2026
This adds an exact CUDA launch contract for TIRx kernels that need PTX
`.reqntid` rather than advisory launch bounds.

- preserve `tirx.required_block_size` through host/device splitting
- emit CUDA 13 `__block_size__` with static thread and cluster
dimensions
- pass the required-block flag through runtime metadata without adding a
packed operand
- launch with `CU_LAUNCH_KERNEL_REQUIRED_BLOCK_DIM` and avoid duplicate
cluster attributes
- reject incompatible launch-bound and max-register controls

Validation:
- `cmake --build build --parallel`
- `python -m pytest -q
tests/python/tirx-transform/test_tir_transform_split_host_device.py
tests/python/tirx/codegen/test_codegen_cuda.py` (225 passed)
- `pre-commit run --files include/tvm/tirx/function.h
src/backend/cuda/codegen/codegen_cuda.cc
src/backend/cuda/runtime/cuda_module.cc src/runtime/metadata.h
src/runtime/thread_storage_scope.h
src/tirx/transform/split_host_device.cc
tests/python/tirx-transform/test_tir_transform_split_host_device.py
tests/python/tirx/codegen/test_codegen_cuda.py`
- CUDA 13.2 SM100 runtime launch through a TIRx kernel with generated
PTX `.reqntid 128,1,1`
tqchen pushed a commit that referenced this pull request Oct 18, 2026
This adds an exact CUDA launch contract for TIRx kernels that need PTX
`.reqntid` rather than advisory launch bounds.

- preserve `tirx.required_block_size` through host/device splitting
- emit CUDA 13 `__block_size__` with static thread and cluster
dimensions
- pass the required-block flag through runtime metadata without adding a
packed operand
- launch with `CU_LAUNCH_KERNEL_REQUIRED_BLOCK_DIM` and avoid duplicate
cluster attributes
- reject incompatible launch-bound and max-register controls

Validation:
- `cmake --build build --parallel`
- `python -m pytest -q
tests/python/tirx-transform/test_tir_transform_split_host_device.py
tests/python/tirx/codegen/test_codegen_cuda.py` (225 passed)
- `pre-commit run --files include/tvm/tirx/function.h
src/backend/cuda/codegen/codegen_cuda.cc
src/backend/cuda/runtime/cuda_module.cc src/runtime/metadata.h
src/runtime/thread_storage_scope.h
src/tirx/transform/split_host_device.cc
tests/python/tirx-transform/test_tir_transform_split_host_device.py
tests/python/tirx/codegen/test_codegen_cuda.py`
- CUDA 13.2 SM100 runtime launch through a TIRx kernel with generated
PTX `.reqntid 128,1,1`
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants