[pallas] compose narrowing views on resident VMEM refs - #3273
Open
ethche wants to merge 8 commits into
Open
Conversation
added 8 commits
July 30, 2026 20:37
Snapshot of the resident-Ref subview implementation as reviewed. Defers
jax array materialization of a tile load so inner-loop sub-selections read
directly out of VMEM instead of dynamic-slicing a materialized block.
Known issues recorded in review, addressed by the follow-up refactor:
- read-after-write miscompile: the lazy Ref observes stores to the
underlying tensor issued after the load
- tile.count regresses to a silent wrong value inside fori_loop /
emit_pipeline (end_var_name set without begin_var_name)
- test_resident_ref_static_unroll and
test_resident_ref_does_not_miscompile_cross_dim_read are invalid
kernels (host variable mutated on device) and never run
- test_resident_ref_flatten_worklist_grouping_two is rejected by
compact_worklist, leaving the grouping-2 variant machinery untested
Replaces the resident-Ref transform DAG with a two-node plan. A tile load
whose consumers all narrow the tiled dimension to a contiguous run is emitted
as a Pallas Ref, and each consumer reads a slice out of it. The load node
stays in the outer loop body, so the block is still staged and DMA
double-buffered once per outer tile; only the read of an individual slice
moves inside the inner loop.
Layered so that folding is an optimization and never a correctness
precondition:
- `hl.subscript`'s fake takes an explicit allowlist of narrowing forms on
Pallas (constant index, constant slice, `tile.begin`, `tile.index`) rather
than deferring to `SubscriptIndexing.compute_shape`, which also concretized
backed symints and changed `[:, None]` shapes for every Pallas kernel.
- Each of those forms has a materializing lowering via
`jax.lax.dynamic_slice_in_dim`, so declining to fold always leaves a
working kernel. Codegen records whether a load really became a Ref and
the consumers pick their path from that, replacing three codegen-time
`InvalidConfig` raises that could leave a kernel uncompilable.
- `pallas_resident_subviews` turns folding off, for A/B measurement and as
an escape hatch. It is a user knob, not an autotuner field, so the search
space is unchanged.
Fixes from the review of the previous implementation:
- A store or atomic anywhere on the loaded tensor now blocks folding. The
lazy Ref previously observed writes issued after the load, so a kernel
that read a block and then overwrote it returned the post-write data.
- `_codegen_fori_loop` / `_codegen_emit_pipeline` record `begin_var_name`
alongside `end_var_name`. Setting only the end made `tile.count` compute
`cdiv(end - 0, block)` and silently return a wrong count for a loop with a
non-zero begin.
- A data-dependent run start is clamped into the block instead of being
matched against the exact `if live >= N` predicate that guarded it. This
is what `jax.lax.dynamic_slice` does with an out-of-range start, so the
folded and materializing lowerings agree on every input.
- `_capture_edges` matched per-config graph copies against the originals by
node signature. Placeholders and the parent's capture argument correspond
positionally by construction, so the copies alone are enough.
- Transform nodes and value transports no longer carry annotations: a Ref
survives `_new_var` and capture placeholders as an ordinary SSA value.
Scope cuts, each covered by a test that pins the fallback:
- Composing through a reshape or view between load and subview.
- Narrowing a dimension the load did not tile.
- The worklist-grouping-2 block-size variants, which had no kernel that
reached them.
Verified in Pallas interpret mode: 207 passed across test_pallas.py,
test_pallas_worklist.py and test_pallas_load_store.py; the GDN and short-conv
benchmark kernels still fold to the same reads with DMA prefetch intact.
Review found two ways the materializing lowering could disagree with the
shape the fake reported, which broke the invariant that declining to fold
always leaves a working kernel.
A bounded slice reaching past the dimension it narrows now fails with
InvalidConfig naming the extent. Python clips such a run while the traced
shape does not, so `x_block[1:10]` on a four-row block reported nine rows and
held three; a kernel dividing by `part.size(0)` returned 0.667 instead of 2.
Whether a run fits depends on the configured block size, so this is reported
where the size is known rather than at trace time.
An index that crossed a scope boundary lost its provenance. The fake accepts
any rank-one index, but the lowering only recognized a `tile_index` or `iota`
node written in place, so binding `tile.index` to a name before a branch --
or using `tile.index + 1` -- raised InvalidIndexingType. Both lowerings now
resolve an index through `_new_var` renames and capture placeholders once, in
the planning pass, and any index that still cannot be proven contiguous
lowers as a `jnp.take` gather instead of failing. A transported tile run
keeps the contiguous `dynamic_slice_in_dim` form rather than degrading to a
gather.
Also:
- `pallas_resident_subviews` rejects non-boolean values. It read as a plain
`config.get`, so None silently disabled an optimization that defaults on.
- `hl.subscript`'s docstring described the pre-existing None/`:` restriction
as the whole surface.
- Adds the load -> view -> subview test for the scoped-out composed-view
case; the existing test covered load -> subview -> view, which folds.
Pallas interpret mode: 211 passed across the three pallas suites. The GDN and
short-conv kernels still fold to the same five Ref loads and five subview
reads with DMA prefetch intact.
All three were reported in review and reproduced before fixing. Negative bounds are refused at trace time. The fake recorded `stop - start` while codegen emitted a Python slice, which clips a negative bound against the dimension: `x_block[-3:2]` on a four-row block reported five rows and produced one. The dimension's size is not visible at trace time, so the bound cannot be normalized there and is rejected instead. A tile run that reaches the subscript as arithmetic keeps its mask. `tile.index + 0` traces to an `add`, which the lowering did not recognize as a tile run and treated as an arbitrary gather -- dropping the tail mask, so a three-element loop tiled by two read a padding row. Runs are now recognized from the `tile_with_offset` provenance the device IR already records, via `subscript_tile_info`, falling back to a bare `tile_index` node whose block id lives on its argument rather than its value. A folded dynamic run reads its start through the resolved index. `_dynamic_begin_expr` took `kwargs["start"]` off the subscript's immediate index node, so a run bound to a name before a branch arrived as a placeholder and emitted `jnp.clip(None, ...)`. The unqualified `jnp.take` fallback is removed. Lowering a real gather needs the backend's indirect-load machinery along with its own bounds and mask semantics, which is a larger surface than this change should carry, so a shifted tile run or any other computed row vector is now refused with InvalidIndexingType rather than approximated. This also matches the hardware: Mosaic cannot lower `dynamic_slice`, so the materializing path for a dynamic index is interpret-only today and the un-folded fallback is not the safety net the earlier commit message claimed. Four regression tests replace the gather test. Pallas interpret mode: 214 passed across the three pallas suites. Re-measured on tpu7x after the change: short conv 84.67us vs 445.35us (5.26x), GDN 2749.00us vs 3248.04us (1.18x), both unchanged.
Narrowing a device value is now lowered exactly one way -- reading a run out
of the resident Pallas Ref the block was loaded into -- and everything else is
rejected with the reason. This replaces a second, materializing lowering that
produced a P1 in each of three review rounds and that Mosaic cannot execute
anyway: it has neither `dynamic_slice` nor a gather, so a run selected out of
a value already in vector registers has no lowering at all.
The three defects that motivated it, each reproduced first:
- A tail run was clamped backward before masking. `dynamic_slice_in_dim`
pulls the start back to keep the run in bounds, after which the mask zeroed
the wrong lanes: rows [1,2,4] tiled by two summed to 5 instead of 7.
- A strided `iota` was treated as contiguous, because the value path lacked
the `step == 1` check the planner had. `hl.arange(0, 4, 2)` read rows 0,1.
- A mask recorded the input axis while `_mask_expr` works on the output
shape, so a `None` ahead of the run masked the wrong dimension.
All three are divergences between two sites that recognized the same four index
forms. Rather than re-synchronize them, `_normalize` is now the only place an
index is recognized; it produces one `_Selector`, planning attaches it to the
subscript, and codegen consumes it without re-deriving anything. A `DYNAMIC`
run carries its start as a value so codegen never returns to the node for it.
Also:
- Planning runs per config and rejects up front, so a subscript that cannot
be lowered fails with the cause rather than reaching codegen. Structural
causes raise InvalidIndexingType; config-dependent ones (run wider than the
block, memory space) raise InvalidConfig so autotuning can skip them.
- A block consumed whole elsewhere cannot stay resident, and the error names
that consumer's source location -- the user's mistake is on a different
line from the one that fails.
- A tile loop that encloses the subscript indirectly, through an `if`, now
counts as the driving loop, so those runs fold instead of being refused.
- Constants read out of the Ref too, so there is one rule rather than two.
- `pallas_resident_subviews` and `DeviceFunction.pallas_ref_loads` are gone.
The knob could not express an A/B once narrowing required residency, and a
fallback that Mosaic cannot run is not a fallback -- addressability may be
checked late but now only ever raises.
- Plain `None`/`:` subscripts are delegated back to the common lowering
untouched.
- The fake rejects only what is undecidable from shapes. A negative slice
bound is one of those; a negative scalar is not, and is refused as a
deliberate language restriction, which the comments now distinguish.
Known residual: the outer load's range mask is still judged twice, by
`whole_dim`/`_bounded_by_block` in planning and by `_load_mask_expr` at
codegen. They agree today and nonuniform-tail tests cover the behaviour, but
it is the same two-opinions shape as the defects above.
Pallas interpret mode: 212 passed across the three pallas suites. Re-measured
on tpu7x: short conv 84.67us vs 445.33us (5.26x), GDN 2750.38us vs 3248.22us
(1.18x), both unchanged by the rewrite.
Plan address-preserving view chains per configuration and retain VMEM references until their value boundary. Lower narrowing through composed Ref transforms, preserve worklist variants, and cover the path with CPU and standalone JAX TPU tests.
Use one typed node plan, keep rejection diagnostics local to the planner, and make pallas_ref registration the address-preserving view contract. Collapse redundant selector state and reject mutations by storage identity so writes through aliases remain safe.
Document strict capture arity validation and detect impossible cycles while resolving resident Ref index provenance.
meta-cla
Bot
added
the
CLA Signed
ethche
requested review from
AmesingFlank,
cota,
norx1991 and
thcmbs
and removed request for
AmesingFlank
August 4, 2026 00:46
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.
Keep Pallas block loads as VMEM
Refs through registered, address-preserving view operations and nested control-flow captures. Contiguous inner indexing can then lower directly topl.dswithout materializing and repeatedly slicing the entire outer block.This introduces no new Helion language API. Existing source such as:
now lowers to the equivalent of:
The outer load remains compatible with asynchronous DMA and double buffering.
Benchmarks
This enabled ragged_causal_conv1d. The Helion kernel uses the newly enabled resident inner-row access:
TPU7x device results for H=4, D=128, and kernel width 4:
This is a kernel-body comparison. Helion receives the selected convolution state directly, while the public wrapper also processes state-index metadata.
The same lowering enables a recurrent Helion GDN implementation comparable to TPU Inference v1 fused_recurrent_gdn. H=4, K=V=128: