[TIRx][CUDA] Replace source helpers with typed PTX forms - #20140
Merged
spectrometerHBH merged 3 commits intoAug 17, 2026
Conversation
spectrometerHBH
approved these changes
Aug 16, 2026
Member
|
@tvm-bot run |
Contributor
|
Failed to re-run CI in https://github.com/apache/tvm/actions/runs/31983859770 Detailswith response |
Contributor
Author
|
Fixed the three CPU CI failures in af44f71:
Local CPU-condition reproduction of the original failing selection: 1 passed, 2 skipped. Changed-file pre-commit and diff checks pass. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
TIRx low-level kernels need to reject
tirx.cuda.func_call: it embeds arbitrary CUDA source outside the typed IR and PTX dialect, so the compiler cannot validate operand types, destination liveness, instruction legality, or target placement.The downstream audit in mlc-ai/TIRx-kernels#62 found four missing general capabilities: inactive-path liveness for predicated destinations, the MXF4 block32 MMA spelling, typed SMEM descriptor operations, and correct target binding for private PrimFunc helpers called from a
T.device_entryregion. This PR adds those capabilities without kernel-specific hooks.Change-by-change rationale
1. One
preserve_dstflag for predicated destinationsFiles:
python/tvm/backend/cuda/ptx/engine.py,render.py, andtable.py.preserve_dst=Falseby default andpreserve_dst=Truewhen the old destination must remain live.Falsemeans it does not;Truemeans it does.preserve_dst=False: use the normal write-only"="constraint. The inactive value is undefined to the caller and must be merged or guarded before use.preserve_dst=True: use a read-write"+"constraint and seed any carrier bridge from the old destination. The inactive path therefore retains that value.undefined_dstwas exactly the complement ofpreserve_dst, so exposing both created contradictory and invalid Boolean combinations. The renderer now derives the default undefined policy fromentry.has_dstpluspreserve_dst=False.preserve_dst=Truerequires bothpred=...and an ordinary written destination. Accumulators already haverw="rw"and need no extra flag; stores and other destination-free instructions reject it.predorpred,keep. The renderer still emits distinct_pred_undefand_pred_keephelper names, preserving helper identity and the exact"="versus"+"ABI.Hardware behavior is unchanged in both cases: predicate false means the PTX instruction does not execute. The flag only states whether the surrounding compiler must keep the previous destination value live.
The downstream GDN path needs both choices in one expression chain: its predicated shared loads use the default undefined destination until a
selp, while the later predicatedex2usespreserve_dst=Trueto retain the zero selected for inactive lanes.2. Explicit MXF4 block32
tcgen05.mmaformFile:
python/tvm/backend/cuda/ptx/table.py.tcgen05.mma.*.block_scale.block16/block32with a dedicatedblock_sizemodifier.tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32. The table previously covered the scale-vector spelling but not the explicit scale-block-size spelling, forcing the kernel to hide a valid instruction in a CUDA source helper.mxf8f6f4andmxf4allow only block32;mxf4nvf4allows block16 and block32. Keeping this as a separate table form preserves exact instruction spelling and rejects invalidmxf4.block16at IR construction time.3. Typed SMEM descriptor address updates
File:
python/tvm/backend/cuda/tile_primitive/common.py.smem_desc_replace_loand implementsmem_desc_add_16B_offsetby reinterpreting the descriptor asuint32x2, updating only lane 0, and repacking it.func_call; the typed implementation makes the no-carry uint32 behavior explicit and gives descriptor address replacement one reusable authority.4. Typed descriptor warp uniformization
Files:
python/tvm/backend/cuda/lang/smem_desc.pyandpython/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py.smem_desc_make_lo_uniformCUDA helpers withmov.b64unpacking,shfl.sync.idx.b32of the low lane, andmov.b64repacking while preserving the high lane.SeqStmtbecause it constructs IR nodes below the script layer. The script-levelSmemDescriptormethod emits the same typed operation sequence at its abstraction layer.5.
T.device_entryas the device-side target boundaryFile:
src/tirx/transform/bind_target.cc.attr::kDeviceEntryas entering GPU scope, alongside thread extents and virtual-thread attributes.T.device_entry. A private typed helper called from that device body must receive the CUDA target. Without recognizing the marker, BindTarget can classify or clone the helper as host-side, producing a cross-target call instead of a device helper.This enables the downstream weighted-ReLU reduction and acquire polling loop to remain typed private PrimFuncs rather than CUDA source helpers.
6. Generated PTX typing surface
Files:
python/tvm/backend/cuda/ptx/gen_stubs.pyandpython/tvm/script/tirx.pyi.predand, for families with an ordinary written destination, the singlepreserve_dst: bool = Falseflag.7. Regression tests
Files:
tests/python/tirx/codegen/test_ptx_dialect.py,tests/python/tirx-transform/test_tir_transform_helpers.py, andtests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py._pred_undefversus_pred_keephelper identity, IR markers and round trips, invalid flag placement, and ptxas acceptance.mxf4.block16is rejected.T.device_entry, then verify separate LLVM and CUDA-targeted functions are produced.T.ptx.shfl_syncbefore the kernel replacement point instead of checking for source-helper names.Downstream evidence
The downstream migration is mlc-ai/TIRx-kernels#62.
Validation
python -m pytest -q tests/python/tirx/codegen/test_ptx_dialect.py: 47 passed, 32 skippedtirx.pyi: byte-identical topython -m tvm.backend.cuda.ptx.gen_stubsoutputRelated
Downstream kernel migration: mlc-ai/TIRx-kernels#62