Skip to content

[TIRx][CUDA] Add PTX address expressions with immediate byte offsets - #20153

Open
spectrometerHBH wants to merge 2 commits into
apache:mainfrom
spectrometerHBH:ptx-addr-expressions
Open

[TIRx][CUDA] Add PTX address expressions with immediate byte offsets#20153
spectrometerHBH wants to merge 2 commits into
apache:mainfrom
spectrometerHBH:ptx-addr-expressions

Conversation

@spectrometerHBH

Copy link
Copy Markdown
Contributor

This PR adds T.ptx.addr(base, byte_offset) to the TIRx PTX dialect: a pure expression that folds a compile-time signed byte displacement into the PTX address operand, rendering as [%N+imm] instead of requiring a separate address computation before the instruction.

  • tirx.ptx.addr is an expression only in the outer PTX call's IR: the helper still receives the coerced base register, while the displacement becomes renderer metadata baked into the instruction text and the helper name (_addr<slot>_p<imm> / _m<imm>).
  • Table-level allow_imm_offset classification of address slots, with validation that rejects immediate offsets on operand classes that cannot take them (e.g. tmem addresses).
  • Immediate operands are now validated to be compile-time IntImm at CUDA codegen, with an actionable error pointing at explicitly-unrolled loops.
  • Displacements are range-checked to int32; zero offsets normalize to the bare form so existing helper names are untouched.

Also includes a small test fix: test_tirx_kernels_registry_correctness.py accepts both the old and new MegaMoE kernel registry names (deepgemm_fp8_fp4_mega_moe / sm100_fp8_fp4_mega_moe), so the test works against tirx-kernels checkouts from either side of the rename.

Tested with tests/python/tirx/codegen/test_ptx_addr.py (new, 12 cases), plus the full tests/python/tirx/ suite on sm100.

)

tirx-kernels renamed deepgemm_fp8_fp4_mega_moe to sm100_fp8_fp4_mega_moe.
List both in _DISTRIBUTED_KERNELS so the test works against checkouts from
either side of the rename, matching how _load_workloads already tolerates
older and newer bench-suite layouts.
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