spectrometerHBH opened a new pull request, #20110: URL: https://github.com/apache/tvm/pull/20110
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 `IterVar`s 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 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
