[TIRX] Represent buffers as typed variables - #20079
Conversation
Make buffer identity an ordinary ir.Var carrying BufferType. The previous parallel Buffer object and pointer-typed data Var could diverge under generic IR substitution and rebuilding; BufferType is now the single source of truth, and buffer_data(buffer) is the explicit physical-pointer projection. This lets standard Var scoping, substitution, undefined-variable analysis, and identity maps handle buffers without a second object hierarchy. It also removes duplicate constructor state and representation-specific remapping from the common path. C++ migration: use BufferVar as a checked Var view, read logical metadata through buffer->..., use buffer.var() for identity, and use buffer.data() or builtin::buffer_data for the pointer projection. Metadata changes rebuild the BufferType and rebind the new Var through AllocBuffer or DeclBuffer. Python migration: constructors return tvm.ir.Var, value.ty carries BufferType, and is_buffer_var(value) is the runtime discriminator. Buffer remains a source-compatibility alias for imports and annotations, so isinstance(value, Buffer) must not be used to distinguish buffer variables.
Keep physical-buffer ownership explicit after Buffer becomes a typed Var. Storage planning, vector type rewriting, cleanup, and all-reduce lowering now maintain pass-local alias-to-root maps only where liveness or allocation remapping requires them; simpler passes decode buffer_data projections locally. Replace superseded DeclBuffer bindings instead of retaining stale declarations, reject aliases whose physical source has not been registered, and repair merged tagged-memory children after their parent allocation is rebuilt. Restore the alias-lifetime and original-output regressions that exercise these invariants. Preserve the Python compatibility surface deliberately: Buffer remains an annotation/import alias, is_buffer_var is the runtime discriminator, legacy metadata properties reject non-buffer Vars with AttributeError, and dtype continues to return DataType. Validated with a clean C++ build, 239 focused passing tests (3 skipped), and the complete TIRx/S-TIR Python suite; remaining broad-suite failures require unavailable LLVM, CUDA, NKI, Torch, or target-feature support.
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
Rebase validation is complete on exact base |
…ering, and typed-buffer migration fixes (#20080) ## Summary This PR upstreams our TIRx development branch, rebased onto current main (36cf270). It contains two commits. **feat(tirx): consolidated fork delta over apache main** — the development stack: - **op-dispatch**: dense fp8/tf32 tcgen05 `gemm_async` paths, per-MMA SMEM descriptor with hoist/recompute, TMA planning split and hardening, Layout F sub-slab selection in tcgen05 ld/st - **layout**: buffer dim-surgery views (`unflatten`/`flatten`/`select`/`narrow`/`sub`/`rearrange`), `SwizzleLayout` folded into `ComposeLayout`, physical offset and rearrange axis-name preservation - **lower-tirx**: FlashMLA CUDA intrinsics and sparse decode lowering, PrimType dtype handling, unsigned swizzle iter patterns, IKET profiling, dynamic while loops kept unrolled-by-1 in CUDA codegen - **tvmscript**: PTX `cvt` instruction forms - **infra**: Triton-standard bench harness (proton timer, cooldown, timer alignment), `TVM_CUDA_NVRTC_EXTRA_OPTS`, nvcc arch-suffix and fast-math opt-out fixes, GemmComm distributed benchmarks **fix(lower-tirx): keep buffer identity coherent across buffer rebuilds** — fixes for regressions from the typed-buffer-variable migration (#20079). These only manifest on sm_100 execution paths, which CI skips without such a GPU: - `FlattenBuffer` is restructured around its invariant: each n-d buffer flattens to a 1-d storage husk and every access `buf[x]` rewrites to `buf'[f(x)]` with `f(x) = layout.apply(x, shape) + elem_offset`. The pair {rewritten geometry, husk} is derived exactly once at each buffer definition point; use sites only look it up, and a use before its definition is a hard error. Previously, loads embedded in view shapes, strides, and folded elem_offsets kept referencing pre-rebuild buffer identities, which surfaced as `MakePackedAPI` "used but not passed as API arguments" failures. - `buffer_data` projections of device-local views carried into host-side tensormap-init statements are now resolved onto their storage root (a PrimFunc parameter) by `TilePrimitiveDispatch` at the moment the statements are hoisted to host scope; `LowerTIRxCleanup`'s alias lookup no longer tolerates forward references. - `LowerTIRxCleanup` rewrites buffer-type shape/stride fields when rebuilding buffers; `LowerTIRxOpaque`'s unit-loop Var visitor no longer shadows the base visitor's buffer remapping. - Dynamic shared memory size is declared as kernel-level metadata (`tirx.dyn_smem_bytes`, emitted by `SMEMPool.commit()`) and read by `SplitHostDevice`, instead of being patched into the shared.dyn allocation's extent, which kept buffer-referencing metadata alive across every buffer-rebuilding pass. The allocation extent is an extern placeholder and is no longer consulted. ## Testing - Full `tests/python/tirx` suite on B200 (sm_100a): 2620 passed, 0 failed (includes the gpu-gated tcgen05/TMA tests that CI skips) - Kernel benchmark suite: 113/113 workloads; MQA logits kernels measure at or above the DeepGEMM reference - New `tests/python/tirx/transform/test_transform_flatten_buffer.py` pins the FlattenBuffer invariant; all three tests fail on the pre-fix pass
This refactors TIRx buffers into ordinary
ir.Varidentities carrying animmutable
BufferType.Rationale
The old representation kept a standalone Buffer object, its logical metadata,
and a separate pointer-typed data Var. Generic Var substitution and
buffer-specific rebuilding could therefore let identity, metadata, and pointer
type drift apart. In this change,
BufferTypeis the single source of thelogical contract and
buffer_data(buffer)is the explicit physical-pointerprojection;
DeclBufferrecords a view's binding.Simplification
maps now apply directly to buffers.
BufferType, so the physical pointer andbuffer metadata cannot disagree.
constructor state, or a parallel buffer-only substitution mechanism.
at the replacing
AllocBufferorDeclBuffer.Physical-root tracking remains deliberately local to the passes whose
correctness depends on allocation ownership:
allocation.
Each is a plain pass-local
Map<Var, Var>flattened when aDeclBufferisvisited. Parameters,
buffer_map, andAllocBufferseed roots; an alias to anunseen typed source is rejected rather than silently treated as a new root.
Simpler passes decode the one-step
buffer_dataprojection locally.Migration guide
C++
BufferVaras a checked view over an ordinaryVar.buffer->...; usebuffer.var()for identity.buffer.data()orbuiltin::buffer_data()only when a physical pointerexpression is required.
BufferType, create a fresh Var, and bind it atthe replacement definition site.
DeclBufferreplaces its superseded declaration consistently;it must not leave a stale declaration beside the new binding.
Python
tvm.ir.Var;.tycontainsBufferType.tvm.tirx.is_buffer_var(value)for runtime discrimination.tvm.tirx.Bufferremains an import and annotation compatibility alias fortvm.ir.Var, soisinstance(value, Buffer)matches every Var.tvm.tirxinstalls the legacy metadata properties on Var.Non-buffer Vars raise
AttributeError; buffer.dtypekeeps its historicalDataTyperesult.Behavior preservation
The behavior-preserving follow-up restores established LowerTIRx,
FlattenBuffer, BF16, alias-lifetime, and declaration-replacement behavior while
keeping ownership logic pass-local.
The final focused correction completes typed-buffer integration at structural
comparison, retyping, specialization, annotation, external-output, access
detection, projection lowering, printer, and target-codegen boundaries. It
also preserves the physical source handle type for declared aliases and
retains allocation-backed versus parameter-backed identity where lowering
requires that distinction.
Validation
pointer, and rank-0/rank-2 access-pointer regressions: 11 passed.
PrimTypeconversionboundaries; the corrected runtime-trace and minimal-LLVM files pass
8 tests.
commit-3 regression suite: 467 passed, 8 xfailed, 1 xpassed.
nvrtcVersion() == 13.1in the testprocess and loaded
libnvrtc.so.13.1.115. After the representation-safeTensorMap
address_ofmigration, the two complete CUDA-codegen files passed376 tests with 159 skipped.
git diff --check:passed.