Repository navigation
[TIRx][CUDA] Replace source helpers with typed PTX forms - #20140
Merged
spectrometerHBH merged 3 commits intoAug 17, 2026
Merged
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. |
wodylm
pushed a commit
to wodylm/mcTVM
that referenced
this pull request
Sep 15, 2026
## 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_entry` region. This PR adds those capabilities without kernel-specific hooks. ## Change-by-change rationale ### 1. One `preserve_dst` flag for predicated destinations Files: `python/tvm/backend/cuda/ptx/engine.py`, `render.py`, and `table.py`. - Change: PTX calls use one Boolean choice, `preserve_dst=False` by default and `preserve_dst=True` when the old destination must remain live. - Why one flag is sufficient: a false PTX predicate always suppresses the instruction write. The only compiler-level choice is whether the inline-asm destination also needs its previous C++ value as an input. `False` means it does not; `True` means 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. - Why not two flags: `undefined_dst` was exactly the complement of `preserve_dst`, so exposing both created contradictory and invalid Boolean combinations. The renderer now derives the default undefined policy from `entry.has_dst` plus `preserve_dst=False`. - Validation: `preserve_dst=True` requires both `pred=...` and an ordinary written destination. Accumulators already have `rw="rw"` and need no extra flag; stores and other destination-free instructions reject it. - IR/codegen identity: the marker is now only `pred` or `pred,keep`. The renderer still emits distinct `_pred_undef` and `_pred_keep` helper 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 predicated `ex2` uses `preserve_dst=True` to retain the zero selected for inactive lanes. ### 2. Explicit MXF4 block32 `tcgen05.mma` form File: `python/tvm/backend/cuda/ptx/table.py`. - Change: add SS and TS entries for `tcgen05.mma.*.block_scale.block16/block32` with a dedicated `block_size` modifier. - Why: the paged MQA kernel needs `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. - ISA constraint: a validator encodes the legal matrix instead of accepting every token combination: `mxf8f6f4` and `mxf4` allow only block32; `mxf4nvf4` allows block16 and block32. Keeping this as a separate table form preserves exact instruction spelling and rejects invalid `mxf4.block16` at IR construction time. ### 3. Typed SMEM descriptor address updates File: `python/tvm/backend/cuda/tile_primitive/common.py`. - Change: introduce `smem_desc_replace_lo` and implement `smem_desc_add_16B_offset` by reinterpreting the descriptor as `uint32x2`, updating only lane 0, and repacking it. - Why: the low 32 bits are the address lane and the high 32 bits contain descriptor control fields. A generic 64-bit add could carry into those control fields. The previous implementation used a CUDA union inside `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.py` and `python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py`. - Change: replace `smem_desc_make_lo_uniform` CUDA helpers with `mov.b64` unpacking, `shfl.sync.idx.b32` of the low lane, and `mov.b64` repacking while preserving the high lane. - Why: only the address lane needs to be broadcast from lane 0; the descriptor metadata must remain unchanged. Expressing the sequence in typed PTX removes an arbitrary source boundary and lets normal operand/type validation see the operation. - Lowering detail: the tile-primitive implementation emits explicit local buffers and a `SeqStmt` because it constructs IR nodes below the script layer. The script-level `SmemDescriptor` method emits the same typed operation sequence at its abstraction layer. ### 5. `T.device_entry` as the device-side target boundary File: `src/tirx/transform/bind_target.cc`. - Change: both the function classifier and call substitutor treat `attr::kDeviceEntry` as entering GPU scope, alongside thread extents and virtual-thread attributes. - Why: a TIRx PrimFunc may contain host preparation followed by a device body marked by `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. - Scope: this does not globally classify the whole function as device code; it changes classification only while visiting the body dominated by the existing canonical device-entry marker. 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.py` and `python/tvm/script/tirx.pyi`. - Change: generated PTX chain signatures expose `pred` and, for families with an ordinary written destination, the single `preserve_dst: bool = False` flag. - Why: the static typing surface must match the runtime dialect API. The stub remains generated from the instruction table rather than becoming a second hand-maintained source of truth. - Generator fix: line wrapping now checks the length of the final indented stub line directly. Removing the redundant flag exposed an off-by-one boundary in the previous indirect threshold. ### 7. Regression tests Files: `tests/python/tirx/codegen/test_ptx_dialect.py`, `tests/python/tirx-transform/test_tir_transform_helpers.py`, and `tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py`. - Predicated destinations: verify default write-only versus opted-in read-write constraints, `_pred_undef` versus `_pred_keep` helper identity, IR markers and round trips, invalid flag placement, and ptxas acceptance. - MXF4 block size: verify the exact block32 opcode reaches generated CUDA/PTX and ptxas, and verify invalid `mxf4.block16` is rejected. - BindTarget: use one private helper from both host code and the region after `T.device_entry`, then verify separate LLVM and CUDA-targeted functions are produced. - Descriptor lowering: verify all descriptor modes still build and that warp-scoped uniformization appears as typed `T.ptx.shfl_sync` before the kernel replacement point instead of checking for source-helper names. ## Downstream evidence The downstream migration is mlc-ai/TIRx-kernels#62. - GDN predication produced byte-identical SASS; fixed-config baseline/candidate ratios were 0.9998x, 1.0001x, and 1.0098x. After simplifying to one flag, all three configs passed again against the frozen FlashInfer oracle. - The typed private weighted-ReLU helper retained 168 registers, zero stack/local memory, and the intended packed arithmetic; ratios were 1.0011x and 1.0116x. - The typed acquire polling helper retained the acquire load/control structure and completed TP1 and TP4 correctness; ratios were about 1.0003x and 0.9955x. - The affected downstream correctness matrix passed, and the low-level IR registry audit no longer needs source helpers for any of these cases. ## Validation - `python -m pytest -q tests/python/tirx/codegen/test_ptx_dialect.py`: 47 passed, 32 skipped - BindTarget regression selection: 1 passed - descriptor lowering regression selection: 6 passed - downstream GDN frozen-oracle correctness: 3 configs passed - downstream low-level IR checker: 2 passed - changed-file pre-commit: passed in both repositories - generated `tirx.pyi`: byte-identical to `python -m tvm.backend.cuda.ptx.gen_stubs` output ## Related Downstream kernel migration: https://github.com/mlc-ai/TIRx-kernels/pull/62
tqchen
pushed a commit
that referenced
this pull request
Oct 16, 2026
## 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_entry` region. This PR adds those capabilities without kernel-specific hooks. ## Change-by-change rationale ### 1. One `preserve_dst` flag for predicated destinations Files: `python/tvm/backend/cuda/ptx/engine.py`, `render.py`, and `table.py`. - Change: PTX calls use one Boolean choice, `preserve_dst=False` by default and `preserve_dst=True` when the old destination must remain live. - Why one flag is sufficient: a false PTX predicate always suppresses the instruction write. The only compiler-level choice is whether the inline-asm destination also needs its previous C++ value as an input. `False` means it does not; `True` means 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. - Why not two flags: `undefined_dst` was exactly the complement of `preserve_dst`, so exposing both created contradictory and invalid Boolean combinations. The renderer now derives the default undefined policy from `entry.has_dst` plus `preserve_dst=False`. - Validation: `preserve_dst=True` requires both `pred=...` and an ordinary written destination. Accumulators already have `rw="rw"` and need no extra flag; stores and other destination-free instructions reject it. - IR/codegen identity: the marker is now only `pred` or `pred,keep`. The renderer still emits distinct `_pred_undef` and `_pred_keep` helper 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 predicated `ex2` uses `preserve_dst=True` to retain the zero selected for inactive lanes. ### 2. Explicit MXF4 block32 `tcgen05.mma` form File: `python/tvm/backend/cuda/ptx/table.py`. - Change: add SS and TS entries for `tcgen05.mma.*.block_scale.block16/block32` with a dedicated `block_size` modifier. - Why: the paged MQA kernel needs `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. - ISA constraint: a validator encodes the legal matrix instead of accepting every token combination: `mxf8f6f4` and `mxf4` allow only block32; `mxf4nvf4` allows block16 and block32. Keeping this as a separate table form preserves exact instruction spelling and rejects invalid `mxf4.block16` at IR construction time. ### 3. Typed SMEM descriptor address updates File: `python/tvm/backend/cuda/tile_primitive/common.py`. - Change: introduce `smem_desc_replace_lo` and implement `smem_desc_add_16B_offset` by reinterpreting the descriptor as `uint32x2`, updating only lane 0, and repacking it. - Why: the low 32 bits are the address lane and the high 32 bits contain descriptor control fields. A generic 64-bit add could carry into those control fields. The previous implementation used a CUDA union inside `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.py` and `python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py`. - Change: replace `smem_desc_make_lo_uniform` CUDA helpers with `mov.b64` unpacking, `shfl.sync.idx.b32` of the low lane, and `mov.b64` repacking while preserving the high lane. - Why: only the address lane needs to be broadcast from lane 0; the descriptor metadata must remain unchanged. Expressing the sequence in typed PTX removes an arbitrary source boundary and lets normal operand/type validation see the operation. - Lowering detail: the tile-primitive implementation emits explicit local buffers and a `SeqStmt` because it constructs IR nodes below the script layer. The script-level `SmemDescriptor` method emits the same typed operation sequence at its abstraction layer. ### 5. `T.device_entry` as the device-side target boundary File: `src/tirx/transform/bind_target.cc`. - Change: both the function classifier and call substitutor treat `attr::kDeviceEntry` as entering GPU scope, alongside thread extents and virtual-thread attributes. - Why: a TIRx PrimFunc may contain host preparation followed by a device body marked by `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. - Scope: this does not globally classify the whole function as device code; it changes classification only while visiting the body dominated by the existing canonical device-entry marker. 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.py` and `python/tvm/script/tirx.pyi`. - Change: generated PTX chain signatures expose `pred` and, for families with an ordinary written destination, the single `preserve_dst: bool = False` flag. - Why: the static typing surface must match the runtime dialect API. The stub remains generated from the instruction table rather than becoming a second hand-maintained source of truth. - Generator fix: line wrapping now checks the length of the final indented stub line directly. Removing the redundant flag exposed an off-by-one boundary in the previous indirect threshold. ### 7. Regression tests Files: `tests/python/tirx/codegen/test_ptx_dialect.py`, `tests/python/tirx-transform/test_tir_transform_helpers.py`, and `tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py`. - Predicated destinations: verify default write-only versus opted-in read-write constraints, `_pred_undef` versus `_pred_keep` helper identity, IR markers and round trips, invalid flag placement, and ptxas acceptance. - MXF4 block size: verify the exact block32 opcode reaches generated CUDA/PTX and ptxas, and verify invalid `mxf4.block16` is rejected. - BindTarget: use one private helper from both host code and the region after `T.device_entry`, then verify separate LLVM and CUDA-targeted functions are produced. - Descriptor lowering: verify all descriptor modes still build and that warp-scoped uniformization appears as typed `T.ptx.shfl_sync` before the kernel replacement point instead of checking for source-helper names. ## Downstream evidence The downstream migration is mlc-ai/TIRx-kernels#62. - GDN predication produced byte-identical SASS; fixed-config baseline/candidate ratios were 0.9998x, 1.0001x, and 1.0098x. After simplifying to one flag, all three configs passed again against the frozen FlashInfer oracle. - The typed private weighted-ReLU helper retained 168 registers, zero stack/local memory, and the intended packed arithmetic; ratios were 1.0011x and 1.0116x. - The typed acquire polling helper retained the acquire load/control structure and completed TP1 and TP4 correctness; ratios were about 1.0003x and 0.9955x. - The affected downstream correctness matrix passed, and the low-level IR registry audit no longer needs source helpers for any of these cases. ## Validation - `python -m pytest -q tests/python/tirx/codegen/test_ptx_dialect.py`: 47 passed, 32 skipped - BindTarget regression selection: 1 passed - descriptor lowering regression selection: 6 passed - downstream GDN frozen-oracle correctness: 3 configs passed - downstream low-level IR checker: 2 passed - changed-file pre-commit: passed in both repositories - generated `tirx.pyi`: byte-identical to `python -m tvm.backend.cuda.ptx.gen_stubs` output ## Related Downstream kernel migration: https://github.com/mlc-ai/TIRx-kernels/pull/62
tqchen
pushed a commit
that referenced
this pull request
Oct 18, 2026
## 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_entry` region. This PR adds those capabilities without kernel-specific hooks. ## Change-by-change rationale ### 1. One `preserve_dst` flag for predicated destinations Files: `python/tvm/backend/cuda/ptx/engine.py`, `render.py`, and `table.py`. - Change: PTX calls use one Boolean choice, `preserve_dst=False` by default and `preserve_dst=True` when the old destination must remain live. - Why one flag is sufficient: a false PTX predicate always suppresses the instruction write. The only compiler-level choice is whether the inline-asm destination also needs its previous C++ value as an input. `False` means it does not; `True` means 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. - Why not two flags: `undefined_dst` was exactly the complement of `preserve_dst`, so exposing both created contradictory and invalid Boolean combinations. The renderer now derives the default undefined policy from `entry.has_dst` plus `preserve_dst=False`. - Validation: `preserve_dst=True` requires both `pred=...` and an ordinary written destination. Accumulators already have `rw="rw"` and need no extra flag; stores and other destination-free instructions reject it. - IR/codegen identity: the marker is now only `pred` or `pred,keep`. The renderer still emits distinct `_pred_undef` and `_pred_keep` helper 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 predicated `ex2` uses `preserve_dst=True` to retain the zero selected for inactive lanes. ### 2. Explicit MXF4 block32 `tcgen05.mma` form File: `python/tvm/backend/cuda/ptx/table.py`. - Change: add SS and TS entries for `tcgen05.mma.*.block_scale.block16/block32` with a dedicated `block_size` modifier. - Why: the paged MQA kernel needs `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. - ISA constraint: a validator encodes the legal matrix instead of accepting every token combination: `mxf8f6f4` and `mxf4` allow only block32; `mxf4nvf4` allows block16 and block32. Keeping this as a separate table form preserves exact instruction spelling and rejects invalid `mxf4.block16` at IR construction time. ### 3. Typed SMEM descriptor address updates File: `python/tvm/backend/cuda/tile_primitive/common.py`. - Change: introduce `smem_desc_replace_lo` and implement `smem_desc_add_16B_offset` by reinterpreting the descriptor as `uint32x2`, updating only lane 0, and repacking it. - Why: the low 32 bits are the address lane and the high 32 bits contain descriptor control fields. A generic 64-bit add could carry into those control fields. The previous implementation used a CUDA union inside `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.py` and `python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py`. - Change: replace `smem_desc_make_lo_uniform` CUDA helpers with `mov.b64` unpacking, `shfl.sync.idx.b32` of the low lane, and `mov.b64` repacking while preserving the high lane. - Why: only the address lane needs to be broadcast from lane 0; the descriptor metadata must remain unchanged. Expressing the sequence in typed PTX removes an arbitrary source boundary and lets normal operand/type validation see the operation. - Lowering detail: the tile-primitive implementation emits explicit local buffers and a `SeqStmt` because it constructs IR nodes below the script layer. The script-level `SmemDescriptor` method emits the same typed operation sequence at its abstraction layer. ### 5. `T.device_entry` as the device-side target boundary File: `src/tirx/transform/bind_target.cc`. - Change: both the function classifier and call substitutor treat `attr::kDeviceEntry` as entering GPU scope, alongside thread extents and virtual-thread attributes. - Why: a TIRx PrimFunc may contain host preparation followed by a device body marked by `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. - Scope: this does not globally classify the whole function as device code; it changes classification only while visiting the body dominated by the existing canonical device-entry marker. 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.py` and `python/tvm/script/tirx.pyi`. - Change: generated PTX chain signatures expose `pred` and, for families with an ordinary written destination, the single `preserve_dst: bool = False` flag. - Why: the static typing surface must match the runtime dialect API. The stub remains generated from the instruction table rather than becoming a second hand-maintained source of truth. - Generator fix: line wrapping now checks the length of the final indented stub line directly. Removing the redundant flag exposed an off-by-one boundary in the previous indirect threshold. ### 7. Regression tests Files: `tests/python/tirx/codegen/test_ptx_dialect.py`, `tests/python/tirx-transform/test_tir_transform_helpers.py`, and `tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py`. - Predicated destinations: verify default write-only versus opted-in read-write constraints, `_pred_undef` versus `_pred_keep` helper identity, IR markers and round trips, invalid flag placement, and ptxas acceptance. - MXF4 block size: verify the exact block32 opcode reaches generated CUDA/PTX and ptxas, and verify invalid `mxf4.block16` is rejected. - BindTarget: use one private helper from both host code and the region after `T.device_entry`, then verify separate LLVM and CUDA-targeted functions are produced. - Descriptor lowering: verify all descriptor modes still build and that warp-scoped uniformization appears as typed `T.ptx.shfl_sync` before the kernel replacement point instead of checking for source-helper names. ## Downstream evidence The downstream migration is mlc-ai/TIRx-kernels#62. - GDN predication produced byte-identical SASS; fixed-config baseline/candidate ratios were 0.9998x, 1.0001x, and 1.0098x. After simplifying to one flag, all three configs passed again against the frozen FlashInfer oracle. - The typed private weighted-ReLU helper retained 168 registers, zero stack/local memory, and the intended packed arithmetic; ratios were 1.0011x and 1.0116x. - The typed acquire polling helper retained the acquire load/control structure and completed TP1 and TP4 correctness; ratios were about 1.0003x and 0.9955x. - The affected downstream correctness matrix passed, and the low-level IR registry audit no longer needs source helpers for any of these cases. ## Validation - `python -m pytest -q tests/python/tirx/codegen/test_ptx_dialect.py`: 47 passed, 32 skipped - BindTarget regression selection: 1 passed - descriptor lowering regression selection: 6 passed - downstream GDN frozen-oracle correctness: 3 configs passed - downstream low-level IR checker: 2 passed - changed-file pre-commit: passed in both repositories - generated `tirx.pyi`: byte-identical to `python -m tvm.backend.cuda.ptx.gen_stubs` output ## Related Downstream kernel migration: https://github.com/mlc-ai/TIRx-kernels/pull/62
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: https://github.com/mlc-ai/TIRx-kernels/pull/62