spectrometerHBH opened a new pull request, #20080:
URL: https://github.com/apache/tvm/pull/20080
## Summary
This PR upstreams our TIRx development branch, rebased onto current main
(36cf270e91). 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 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]
