[TIRx] tcgen05 dispatch paths, buffer dim-surgery views, FlashMLA lowering, and typed-buffer migration fixes - #20080
Open
spectrometerHBH wants to merge 2 commits into
Open
[TIRx] tcgen05 dispatch paths, buffer dim-surgery views, FlashMLA lowering, and typed-buffer migration fixes#20080spectrometerHBH wants to merge 2 commits into
spectrometerHBH wants to merge 2 commits into
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
spectrometerHBH
changed the title
Contributor
Author
|
@tvm-bot rerun |
1 similar comment
Contributor
Author
|
@tvm-bot rerun |
Contributor
Author
|
@tvm-bot rerun |
1 similar comment
Contributor
Author
|
@tvm-bot rerun |
Squash of the fork's development stack. Major pieces: - op-dispatch: dense fp8/tf32 tcgen05 gemm_async paths, per-MMA SMEM descriptor hoist/recompute, TMA planning split/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/rearrange axis-name preservation - lower-tirx: FlashMLA CUDA intrinsics + 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 - tests: MegaKernelMoE scheduler coverage, TMA host-init dtype coverage, remote mbarrier arrive API alignment
After the typed-buffer-variable migration, a buffer's identity is the
variable itself. A pass that rewrites the program while rebuilding
buffers is applying a substitution over buffer identities, and that
substitution must reach every expression position — including the ones
tucked inside buffer types and pass-internal folds. Several spots
missed this and broke kernels using shared-memory pools, views
sized/offset by local scalars, or host-side tensormap encodes of
device-declared views.
FlattenBuffer is restructured around its invariant instead of patched:
each n-d buffer flattens to a 1-d storage husk buf' (same data origin,
dtype, alignment and scope; no layout, no elem_offset), and every
access buf[x] rewrites to buf'[f(x)] with
f(x) = layout.apply(x, shape) + elem_offset. Since the folded indices
live in the rewritten program, f's coefficients (runtime elem_offset,
symbolic shapes/strides, layout iters) must be rewritten too. The pass
now walks the AST top-down and derives, exactly once at each buffer's
definition point (AllocBuffer/DeclBuffer/params), the pair
{fold view = rewritten geometry, buf'}; use sites only look the pair
up. The substitution is applied exactly once per subexpression by
construction, and a use before its definition is a hard error at the
fault instead of a stale reference that surfaces later in
MakePackedAPI.
A buffer_data projection is only resolvable where the buffer's
definition is visible, and TMA host-init statements violated that:
they are emitted inside the kernel region but hoisted to host scope,
carrying projections of device-local views. TilePrimitiveDispatch now
tracks each buffer's storage root as it walks definitions and rewrites
those projections onto the root — a PrimFunc parameter, visible on the
host — at the moment the statements are hoisted. With host statements
self-contained, LowerTIRxCleanup's alias lookup no longer has to
tolerate forward references, so a projection with no visible
definition is an error there instead of silently resolving to itself.
Two more substitution obligations were unmet:
- LowerTIRxCleanup rebuilt buffers without rewriting the buffer-type
shape/stride fields, so a view sized by a local scalar kept pointing
at the pre-rebuild variable.
- LowerTIRxOpaque's unit-loop Var visitor returned non-unit-loop
variables verbatim, shadowing the base visitor's buffer remapping.
The shared-memory pool no longer needs buffer-referencing metadata at
all. Its allocation is an extern placeholder whose extent never
reaches the generated code, so the dynamic-shared-memory launch size
is declared directly instead of being patched into the allocation:
- SMEMPool.commit() appends a leaf AttrStmt
(node=0, key="tirx.dyn_smem_bytes", IntImm value) to the kernel
region; hand-written kernels declare theirs with the existing
T.attr({"tirx.dyn_smem_bytes": n}) sugar.
- SplitHostDevice reads the declaration as the dynamic-shared-memory
launch argument and strips it when finalizing the kernel; allocation
extents are no longer consulted, and a shared.dyn allocation without
a declaration is a hard error.
- LowerTIRxOpaque's pool-size collection and extent patching are
removed.
Fork tests updated for the typed-buffer surface;
test_tcgen05_mma_ss_no_tma builds its smem descriptors with
first-class ptr_to() instead of a reinterpreting access_ptr over the
raw arena, satisfying tvm_access_ptr's element-type check; and
test_transform_flatten_buffer.py pins the FlattenBuffer invariant
(rewritten references in view shapes and folded elem_offsets; identity
preservation for already-flat buffers); all three of its tests fail on
the pre-fix pass.
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
1 participant
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.
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:
gemm_asyncpaths, per-MMA SMEM descriptor with hoist/recompute, TMA planning split and hardening, Layout F sub-slab selection in tcgen05 ld/stunflatten/flatten/select/narrow/sub/rearrange),SwizzleLayoutfolded intoComposeLayout, physical offset and rearrange axis-name preservationcvtinstruction formsTVM_CUDA_NVRTC_EXTRA_OPTS, nvcc arch-suffix and fast-math opt-out fixes, GemmComm distributed benchmarksfix(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:
FlattenBufferis restructured around its invariant: each n-d buffer flattens to a 1-d storage husk and every accessbuf[x]rewrites tobuf'[f(x)]withf(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 asMakePackedAPI"used but not passed as API arguments" failures.buffer_dataprojections of device-local views carried into host-side tensormap-init statements are now resolved onto their storage root (a PrimFunc parameter) byTilePrimitiveDispatchat the moment the statements are hoisted to host scope;LowerTIRxCleanup's alias lookup no longer tolerates forward references.LowerTIRxCleanuprewrites buffer-type shape/stride fields when rebuilding buffers;LowerTIRxOpaque's unit-loop Var visitor no longer shadows the base visitor's buffer remapping.tirx.dyn_smem_bytes, emitted bySMEMPool.commit()) and read bySplitHostDevice, 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
tests/python/tirxsuite on B200 (sm_100a): 2620 passed, 0 failed (includes the gpu-gated tcgen05/TMA tests that CI skips)tests/python/tirx/transform/test_transform_flatten_buffer.pypins the FlattenBuffer invariant; all three tests fail on the pre-fix pass