Skip to content

[TIRx][CUDA] uint32 index dtypes, directive fixes, and PTX ISA coverage - #20110

Merged
spectrometerHBH merged 4 commits into
apache:mainfrom
spectrometerHBH:main
Aug 9, 2026
Merged

spectrometerHBH merged 4 commits into
apache:mainfrom
spectrometerHBH:main

Conversation

@spectrometerHBH

Copy link
Copy Markdown
Contributor

Three independent TIRx changes, readable commit by commit.


1. [TIRx][TVMScript] Support uint32 loop and scope-id var dtypes

Loop vars (T.serial/parallel/vectorized/unroll/grid) and scope-id vars (T.thread_id, T.lane_id, T.warp_id, T.cta_id, …) were pinned to int32. Adds an explicit per-site dtype= accepting "int32" (the default) or "uint32", so kernels can get u32 index registers, feed .u32 PTX operands directly, and skip a manual Cast at every use site.

  • For constructor promotes IntImm bounds across signedness as well as width, with a representability check.
  • Frames infer the loop var dtype by C-style signedness promotion rather than taking the code from min alone. This also fixes for i in range(<uint32 expr>), which previously tripped the loop_var/extent dtype equality check.
  • For scope ids the dtype applies to the def vars only; extents stay int32, so launch-param IterVars are unchanged and lowering casts at the bind.
  • The printer emits dtype= for scope ids, whose dtype is not recoverable from the extents. Loops already round-trip via their bound dtypes.

T.thread_binding stays int32-only, since the IterVar constructor requires an int-coded extent dtype.

2. [FIX][TIRx] Preserve CUDA sync and unroll directives

  • #pragma unroll lost its count: an explicit unroll=N printed as a bare #pragma unroll. The codegen now emits the count from either the ForKind::kUnrolled extent or a pragma_unroll annotation, and the printer round-trips the annotation's value instead of flattening it to True.
  • The mbarrier.try_wait.parity intrinsic now spells .acquire.cta, matching the ordering semantics the wait actually needs.

3. [TIRx][CUDA] Extend the PTX dialect across ISA 9.7.1-9.7.14

Registers the instructions the table-driven T.ptx dialect was missing across nine ISA sections, taking it from 175 to 314 entries and from 162,716 to 198,709 variants.

ISA Added
9.7.1 integer arithmetic add/sub/mul/mad (incl. .wide), mul24, mad24, sad, div, rem, abs, neg, popc, clz, bfind, brev, bfe, bfi, szext, bmsk, dp4a, dp2a
9.7.3 floating point mad, div, abs, sqrt, rsqrt, sin, cos, lg2, tanh, testp, copysign, rcp's f64 approximation
9.7.4 half precision fma (with .relu/.oob), neg, abs, tanh, ex2
9.7.5 mixed precision already covered by the srctype lines — verified cell by cell and test-pinned
9.7.6 comparison/selection set, setp, selp, slct
9.7.7 half comparison set, setp over the half type grid
9.7.8 logic and shift and, or, xor, not, cnot, lop3, shf, shl, shr
9.7.9 data movement scalar mov, prmt, ldu, prefetchu, applypriority, discard, isspacep, getctarank, cvt.pack, shfl.sync, the full cvta grid, st.async, multimem, createpolicy, cp.async.bulk.prefetch, tensormap.replace
9.7.14 parallel sync vote.sync, match.sync, activemask, redux.sync, elect.sync, atom.cas/.exch/half/vector, red.async, the state-returning mbarrier arrives and non-parity waits, bar.red, tensormap.cp_fenceproxy

Mechanism extensions, each mirroring something the table already had:

  • OperandSlot.dtype may be a function of the modifier map, for operands the ISA types by formula rather than by a written token — mul.wide's result is "twice as wide as a and b", dp4a's accumulator is .u32 only if both source types are. Neither is nameable as a slot reference, because neither appears in the instruction text. Same totality contract as the existing lanes/sinkable callables.
  • OperandSlot.pipe joins adjacent operands with |, for the ISA's paired destinations (setp's p|q, lop3's d|p, shfl.sync, elect.sync). It sits beside the existing bracket mechanism and composes with it — lop3's pair holds two different register classes and renders %0|pd0, each half keeping its own constraint and bridge.
  • b32i/b64i name a bit-size type with the float carrier removed. Several instructions refuse a float register in one position while accepting either integer signedness; pinning those operands to .u32 would also have dropped the signed spelling ptxas accepts.
  • The script printer escapes the family name, so the three keyword mnemonics (and, or, not) round-trip as T.ptx.and_. Without it a kernel using them printed T.ptx.and(...), a syntax error on reparse.

Method. Every check function was written after probing the cells it rejects, so the table encodes measured behaviour rather than a reading of the spec. That caught several places where ptxas and the ISA text disagree, each recorded at the entry:

  • st.async's release line takes no 8-bit type, though the ISA lists .b8/.u8/.s8
  • .ftz on set follows the source precision, not the line it is written on
  • multimem pairs an ordering semantic with a scope — both halves required, not "either may be omitted"
  • bmsk, elect.sync, match.sync, redux's bitwise line and mbarrier's state token take no float register
  • .acc::f32 applies only to .add; multimem's .add row excludes .s64

Forms the ISA spells but the CUDA 13.2 toolchain (PTX ISA 9.2) cannot assemble — clmad, multimem.st.async, multimem.red.async, the mbarrier layout facility, phase_type waits — are excluded with the reason stated where they belong, as are the deprecated non-.sync shfl and vote, which ptxas rejects outright at sm_70 and above.


Verification

  • Full ptxas certification green on all 32 shards, covering all 198,709 variants at sm_90, sm_90a, sm_100 and sm_100a (PTX_CERT=1 pytest -n 16 -k certify tests/python/tirx/codegen/test_ptx_dialect.py)
  • test_ptx_dialect.py 43 passed; 280 passed across the adjacent TIRx suites
  • New goldens and one dispatch trace test per ISA section, each exercising shared-mnemonic resolution, the script round trip, and the error cases the checks encode
  • Three existing test guards generalized from table-key/prefix matching to the canonical mnemonic, now that several instructions share a name or collide with the bridge machinery's own setp/selp
  • pre-commit run --all-files clean

Loop vars (T.serial/parallel/vectorized/unroll/grid) and scope-id vars
(T.thread_id, T.lane_id, T.warp_id, T.cta_id, ...) were pinned to int32.
Add an explicit per-site `dtype=` accepting "int32" (the default) or
"uint32", so kernels can get u32 index registers, feed .u32 PTX operands
directly, and avoid a manual Cast at every use site.

- For constructor: promote IntImm bounds across signedness as well as
  width, with a representability check.
- For frames: infer the loop var dtype using C-style signedness promotion
  instead of taking the code from `min` alone. This also fixes
  `for i in range(<uint32 expr>)`, which previously tripped the
  loop_var/extent dtype equality check.
- Scope ids: the dtype applies to the def vars only. Extents stay int32,
  so launch-param IterVars are unchanged and lowering casts at the bind.
- Printer: emit `dtype=` for scope ids, whose dtype is not recoverable
  from the extents. Loops already round-trip via their bound dtypes.

T.thread_binding stays int32-only, since the IterVar constructor requires
an int-coded extent dtype.
Registers the instructions the table-driven `T.ptx` dialect was missing
across nine ISA sections, taking it from 175 to 314 entries and from
162,716 to 198,709 certified variants:

  9.7.1  integer arithmetic     add/sub/mul/mad (+.wide), mul24, mad24,
                                sad, div, rem, abs, neg, popc, clz,
                                bfind, brev, bfe, bfi, szext, bmsk,
                                dp4a, dp2a
  9.7.3  floating point         mad, div, abs, sqrt, rsqrt, sin, cos,
                                lg2, tanh, testp, copysign, rcp's f64
                                approximation
  9.7.4  half precision         fma (with .relu/.oob), neg, abs, tanh, ex2
  9.7.5  mixed precision        already covered by the srctype lines;
                                verified cell by cell and test-pinned
  9.7.6  comparison/selection   set, setp, selp, slct
  9.7.7  half comparison        set, setp over the half type grid
  9.7.8  logic and shift        and, or, xor, not, cnot, lop3, shf, shl, shr
  9.7.9  data movement          scalar mov, prmt, ldu, prefetchu,
                                applypriority, discard, isspacep,
                                getctarank, cvt.pack, shfl.sync, the full
                                cvta grid, st.async, multimem, createpolicy,
                                cp.async.bulk.prefetch, tensormap.replace
  9.7.14 parallel sync          vote.sync, match.sync, activemask,
                                redux.sync, elect.sync, atom.cas/.exch/
                                half/vector, red.async, the state-returning
                                mbarrier arrives and non-parity waits,
                                bar.red, tensormap.cp_fenceproxy

Four mechanism extensions were needed, each mirroring an existing one:

  - `OperandSlot.dtype` may be a function of the modifier map, for the
    operands the ISA types by formula rather than by a written token
    (mul.wide's double-width result, dp4a's derived accumulator). Same
    contract as the existing `lanes`/`sinkable` callables.
  - `OperandSlot.pipe` joins adjacent operands with `|`, for the ISA's
    paired destinations (setp's p|q, lop3's d|p, shfl.sync, elect.sync).
    Sits beside the `bracket` mechanism and composes with it.
  - `b32i`/`b64i` name a bit-size type with the float carrier removed.
    A handful of instructions refuse a float register in one position
    while taking either integer signedness; pinning those operands to
    `.u32` would have dropped the signed spelling ptxas accepts.
  - The script printer now escapes the family name, so the three
    keyword mnemonics (and, or, not) round-trip as `T.ptx.and_`.

Every registered variant is proven against ptxas rather than read off
the spec: each check function was written after probing the cells it
rejects, and several encode measured facts where the toolchain and the
ISA text disagree -- st.async's release line takes no 8-bit type,
`.ftz` on `set` follows the source precision, multimem pairs an
ordering semantic with a scope, and bmsk takes no float register.
Forms the ISA spells but this toolchain cannot assemble (clmad,
multimem.st.async, multimem.red.async, the mbarrier layout facility)
are excluded with the reason stated at the entry.
test_uint32_loop_var_runs_correctly launches the kernel on the device,
but carried neither of the guards its siblings in this file use, so on a
CPU-only node it failed with "Device API cuda is not enabled" instead of
skipping.

The neighbouring test_uint32_loop_var_and_scope_id_emit_unsigned only
inspects the generated source, so it correctly stays unguarded.
@spectrometerHBH
spectrometerHBH merged commit 0468e13 into apache:main Aug 9, 2026
8 checks passed
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