Repository navigation
[TIRx][CUDA] uint32 index dtypes, directive fixes, and PTX ISA coverage - #20110
Merged
Merged
Conversation
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.
tqchen
approved these changes
Aug 9, 2026
tqchen
pushed a commit
that referenced
this pull request
Oct 16, 2026
tqchen
pushed a commit
that referenced
this pull request
Oct 18, 2026
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
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.
Three independent TIRx changes, readable commit by commit.
1.
[TIRx][TVMScript]Support uint32 loop and scope-id var dtypesLoop 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 toint32. Adds an explicit per-sitedtype=accepting"int32"(the default) or"uint32", so kernels can get u32 index registers, feed.u32PTX operands directly, and skip a manualCastat every use site.Forconstructor promotesIntImmbounds across signedness as well as width, with a representability check.minalone. This also fixesfor i in range(<uint32 expr>), which previously tripped the loop_var/extent dtype equality check.int32, so launch-paramIterVars are unchanged and lowering casts at the bind.dtype=for scope ids, whose dtype is not recoverable from the extents. Loops already round-trip via their bound dtypes.T.thread_bindingstays int32-only, since theIterVarconstructor requires an int-coded extent dtype.2.
[FIX][TIRx]Preserve CUDA sync and unroll directives#pragma unrolllost its count: an explicitunroll=Nprinted as a bare#pragma unroll. The codegen now emits the count from either theForKind::kUnrolledextent or apragma_unrollannotation, and the printer round-trips the annotation's value instead of flattening it toTrue.mbarrier.try_wait.parityintrinsic 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.14Registers the instructions the table-driven
T.ptxdialect was missing across nine ISA sections, taking it from 175 to 314 entries and from 162,716 to 198,709 variants.add/sub/mul/mad(incl..wide),mul24,mad24,sad,div,rem,abs,neg,popc,clz,bfind,brev,bfe,bfi,szext,bmsk,dp4a,dp2amad,div,abs,sqrt,rsqrt,sin,cos,lg2,tanh,testp,copysign,rcp's f64 approximationfma(with.relu/.oob),neg,abs,tanh,ex2srctypelines — verified cell by cell and test-pinnedset,setp,selp,slctset,setpover the half type gridand,or,xor,not,cnot,lop3,shf,shl,shrmov,prmt,ldu,prefetchu,applypriority,discard,isspacep,getctarank,cvt.pack,shfl.sync, the fullcvtagrid,st.async,multimem,createpolicy,cp.async.bulk.prefetch,tensormap.replacevote.sync,match.sync,activemask,redux.sync,elect.sync,atom.cas/.exch/half/vector,red.async, the state-returningmbarrierarrives and non-parity waits,bar.red,tensormap.cp_fenceproxyMechanism extensions, each mirroring something the table already had:
OperandSlot.dtypemay 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.u32only if both source types are. Neither is nameable as a slot reference, because neither appears in the instruction text. Same totality contract as the existinglanes/sinkablecallables.OperandSlot.pipejoins adjacent operands with|, for the ISA's paired destinations (setp'sp|q,lop3'sd|p,shfl.sync,elect.sync). It sits beside the existingbracketmechanism 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/b64iname 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.u32would also have dropped the signed spelling ptxas accepts.and,or,not) round-trip asT.ptx.and_. Without it a kernel using them printedT.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.ftzonsetfollows the source precision, not the line it is written onmultimempairs an ordering semantic with a scope — both halves required, not "either may be omitted"bmsk,elect.sync,match.sync,redux's bitwise line andmbarrier's state token take no float register.acc::f32applies only to.add;multimem's.addrow excludes.s64Forms the ISA spells but the CUDA 13.2 toolchain (PTX ISA 9.2) cannot assemble —
clmad,multimem.st.async,multimem.red.async, thembarrierlayout facility,phase_typewaits — are excluded with the reason stated where they belong, as are the deprecated non-.syncshflandvote, which ptxas rejects outright at sm_70 and above.Verification
PTX_CERT=1 pytest -n 16 -k certify tests/python/tirx/codegen/test_ptx_dialect.py)test_ptx_dialect.py43 passed; 280 passed across the adjacent TIRx suitessetp/selppre-commit run --all-filesclean