[Lang] Reducer v2: first-class deferred reduction epochs with planned physical lowering (#2940)
* [Lang] Reducer v2: first-class deferred reduction epochs with planned physical lowering
Fixes #2408 at its root. v1 implemented reducers as fragment buffers with
a special layout, so no pass owned the contribution-multiplicity
semantics: loop-layout replication silently multiplied contributions
(128 threads summing 8 elements returned 576 instead of 36), the
finalize width had to be guessed from the storage layout's
ReplicateExtent, and every compiler stage carried buffer-identity
special cases to compensate (see #2897, #2881).
Reducer v2 makes the reducer a first-class deferred reduction epoch.
Core invariant: `T.reducer_update` is the single source of contribution
multiplicity; physical layout can never change how many times a logical
contribution is combined.
New IR contract (`local.reducer` virtual scope + three first-class ops):
acc = T.alloc_reducer(shape, dtype, op="sum", seed=None)
T.reducer_init(acc)
T.reducer_update(acc[indices], contribution) # inside T.Parallel
T.finalize_reducer(acc, dst) # out-of-place
Pipeline:
* VerifyReducerEpoch (early): lifecycle state machine + opaque-access
rules; ordinary reads/writes, clear/fill, aliasing, updates outside
T.Parallel, double init/finalize are compile errors.
* ReducerPlanAndMaterialize (after LayoutInference; loop layouts are
read-only inputs): decides storage, communication, and execution
multiplicity per epoch, then materializes the first-class ops into
ordinary IR. Two physical plans:
- Narrow plan: project the reduction axes (parallel vars absent from
the update indices, plus loop replication) out of the update site's
loop layout to obtain compact per-thread partial storage; updates
run unguarded on every replica; the collective covers only the
thread-expression splits sourced from reduction axes (reusing the
ComputeReducerLayout / CollectThreadReduceSteps machinery T.reduce
is built on). Zero splits means zero collectives (LocalComplete).
Proof obligations: direct var-to-dim indices, replica-safe
contribution values, full-block coverage, a single power-of-two
step, structurally equal plans across sites, destination-layout
containment (with a controlled override for provably unconstrained
single-slot destinations).
- Wide plan (FullParticipant baseline, always available): one full
logical-shape partial per participant, canonical-replica-guarded
updates, participant-wide AllReduce, seed combined exactly once
post-collective. Any narrow-plan proof failure falls back here
deterministically - never to a compile error or wrong code.
* Generic multiplicity marker `tl.parallel_multiplicity`: PartitionLoop
lowers it to a REP == 0 guard knowing nothing about reducers; the
fully_replicated_reducer_buffers compatibility hook is retired.
* FinalizeReducerOp takes explicit plan operands (reducing_threads,
scale), closing the derive-width-from-ReplicateExtent channel.
* VerifyReducerConsumed (after LowerTileOp): surviving v2 constructs
are compile errors, never silently-wrong code.
* Pass config `tl.reducer_force_baseline` forces the wide plan for
differential testing.
Legacy compatibility and v1 removal: CanonicalizeLegacyReducer rewrites
v1 syntax (clear + read-modify-write stores + in-place finalize) into
v2 ops on a strict whitelist; non-identity fills on idempotent combines
are forwarded as one-time seeds, non-zero sum fills are rejected, batch
annotations are forwarded, and post-finalize reads are redirected to a
synthesized destination. With the shim in place, the entire v1 lowering
path is deleted: LayoutReducer, the PartitionLoop/LowerParallelLoop
reducer-buffer plumbing, the ParallelOp rep=ALL special cases, and the
vectorizer's reducer parameter chain.
Testing (H100, kernel cache disabled): 128 passed / 2 skipped across
the legacy reduce suite (93, unchanged, through the shim), the new
reducer v2 suite (22, incl. the 36-vs-576 regression, narrow-plan
codegen assertions, lifecycle diagnostics, forced-baseline
differentials), and the copy suite; the pipelined GEMV example (legacy
syntax) remains numerically correct.
* [Refactor] Consolidate reducer ops into src/op/reducer.{h,cc}
FinalizeReducerOp is no longer a standalone legacy op — it is the
collective that ReducerPlanAndMaterialize emits — so it moves next to
the first-class reducer ops. Its combine-op field now uses the same
ReducerV2OpType enum as the planner, removing the implicit
static_cast coupling between two parallel enums.
The former LayoutReducer pass files are gone entirely: the
ReducerInfo/ReducerRepType annotation types had no consumers left, and
the legacy `reducer_info` annotation key (still read by the
canonicalization shim and the data-race verifier) moves to
op/reducer.h. Stale includes in layout_inference / lower_tile_op are
dropped.
No functional change; 115 reducer tests pass.
* [Doc] Clarify the intent of tl.reducer_force_baseline
The switch is not a workaround for expected narrow-plan bugs; it exposes
the reducer design's canonical reference lowering for differential
testing, field escape-hatch use, and plan-choice A/B measurement.
* [Layout] Opaque tile-op loops must not self-plan their layout
A parallel loop whose body contains an opaque tile-op call (e.g. a
reducer update) has no visible write to anchor layout inference. In
free mode such a loop used to self-plan a partition, and the
register-count-based attempt selection then preferred it: the plan
minimizes this loop's own footprint (e.g. one x-vector element per
thread) while silently forcing a column-major layout onto the fragment
operands it reads, degrading their producing copies from vectorized to
scalar accesses. On a 4096x4096 fp16 GEMV this cost 40% end-to-end
versus the legacy reducer lowering.
Fix (generic, not reducer-aware): when a loop carries an opaque
tile-op call, prefer propagating its layout from a fragment operand
that is addressed purely by the parallel loop vars, deferring while
such an operand is still unsolved so the producing ops fix the
operands first. Self-planning remains the fallback when no operand can
drive the loop (e.g. operands indexed by inner serial reduction vars,
which require the loop to own the operand layout instead) or when
operand propagation cannot express the loop's iteration space.
With the fix the GEMV benchmark matches the legacy lowering (0.075 ms
vs 0.074 ms) and narrow-plan row reductions stay ahead (AllReduce<16>
with one partial slot per thread, ~6% faster than legacy); 129 tests
pass across the reducer v1/v2 and copy suites.
* [Lang] Reducer v2 stage 4: narrow-plan coverage, packed accumulation, bitwise ops
Narrow-plan coverage and observability:
- tl.enable_reducer_plan_verbose pass config: INFO-level per-epoch plan
decisions (narrow/wide, collective widths, rejection reasons).
- Multi-step collectives: interleaved thread splits lower to one AllReduce
per step instead of falling back to the wide plan.
- Destination override generalized: multi-slot destinations and transitive
copy chains (e.g. fp32->fp16 staging hops) adopt the induced layout when
the whole chain is unconstrained; ParallelOp candidate selection falls
back to the operand-compatible layout when the register-preferred
candidate conflicts with an overridden buffer.
- Update-index generalization: permuted loop-var indices and constant unit
dims (acc[j, i], acc[0, i]) rebuild the induced fragment over the
reducer's own dim order; all-constant scalar indices keep the synthetic
unit input ComputeReducerLayout leaves behind.
- Seed and batch now work under narrow plans: the finalize lowering
combines the seed once per slot after the collective, and explicit-width
steps compose with batched AllReduce.
Packed partial accumulation (16-bit floats):
- Single-site narrow epochs on fp16/bf16 split each per-thread partial
into two lanes keyed by the parity of an on-thread reduction source (the
innermost enclosing serial loop, or the innermost parallel reduction dim
when its low bit stays on one thread), halving the serial combine
dependence chain.
- A per-thread fold recombines the lanes into the plain induced storage
before the (unchanged) collective, so the plan's communication and
destination proofs are untouched. Lane parity never changes which thread
a contribution lands on, so the plan stays correct under any layout.
Bitwise combine ops:
- alloc_reducer(op="bitand"/"bitor"/"bitxor") for integer dtypes under
both plans (identities: all-ones / 0 / 0). Float dtypes are rejected at
allocation time; the legacy v1 path stays sum/max/min only.
Tests: reducer v2 suite grows to 38 (packed structural + numeric +
differential, bitwise under both plans, permuted/unit-dim/scalar indices,
narrow seed); the v1 codegen batch tests pin tl.reducer_force_baseline
since narrow plans legitimately skip run_batch. v1 suite stays green (93).
* [Lang] Move the reducer seed to T.reducer_init(acc, init)
The logical starting value used to be declared at allocation time
(alloc_reducer(seed=...)) and smuggled to the planner through the
reducer_info_v2 block annotation. Declaring it where the epoch opens
reads naturally ("the reduction starts from this value") and makes it a
first-class operand of the init op:
acc = T.alloc_reducer((M,), dtype, op="sum")
T.reducer_init(acc, 100.0) # optional; default = combine identity
- ReducerInitOp accepts an optional second argument carrying the value;
the annotation now carries only the combine op.
- VerifyReducerEpoch checks the value's dtype at the init site (Python
numbers are auto-converted by the frontend).
- CanonicalizeLegacyReducer forwards v1 non-identity idempotent fills as
the synthesized reducer_init's argument instead of via the annotation.
- Semantics unchanged: physical partials always initialize to the
identity; the value is combined exactly once per logical output at
finalize time, so physical replication can never multiply it.
Tests: seed tests migrated to the new API plus a dtype-mismatch
diagnostic (v2 suite now 39); v1 suite stays green (93).
* [Layout] Name the operand-driven-loop contract: kTLPerIterationOp
The opaque-loop rule (6749631b) worked but read poorly: detection
sniffed the tile-op registry (minus a tl.region special case), the
policy lived in comments spread over three code blocks, and "can this
operand drive the loop?" was answered twice with different machinery
(a try/catch around the generic source-buffer choice, plus a separate
drivability scan for the defer decision).
Restructure without changing the architecture:
- New op attribute kTLPerIterationOp declares the contract at the op
definition site: the op executes once per iteration inside T.Parallel
and its region-mediated effects are invisible to statement-level
scans. tl.reducer_update is its only member today. Both consumers now
read the attribute instead of keeping their own lists: ParallelOp's
detection (previously registry sniffing) and the nested-loop
checker's _PARALLEL_SAFE_TILE_OPS whitelist (deleted).
- Free-mode inference for such loops is now a single decision:
FindOperandDriver returns the best solved drivable operand (indices
pure in the loop vars, most index dims first) or reports a pending
one; the loop adopts, defers, or falls through to self-planning. The
exception-as-control-flow try/catch is gone from the policy path (a
small guard remains for non-affine driver indices, which drivability
does not test). One behavioral refinement: when the generic
most-indices read buffer is not drivable but another solved drivable
operand exists, the loop now adopts that operand instead of falling
back to self-planning.
Verified: reducer v2 (39) + v1 (93) + copy suites pass; the reducer
benchmark (rowsum fp32/fp16, GEMV fp16) reproduces identical plans and
timings to the previous commit.
* [Lang] reducer_update is a per-iteration builtin intrinsic, not a tile op
reducer_update never was a tile op in substance: it executes once per
iteration INSIDE T.Parallel (no other tile op does), owns no layout (the
enclosing loop and the planner decide physics), and never lowers on its
own. Registering it as one forced two categories of friction: consumers
that had to exempt it (the nested-loop checker whitelist, ParallelOp's
registry sniffing before kTLPerIterationOp), and an argument convention
(tl.region wrapping) whose only real payload was the multi-dim indices.
Re-register it as the builtin `tl.reducer_update`:
- args are (acc[indices], value) with a plain BufferLoad target — an
update descriptor keeping structured indices for the planner. Every
access scan now sees the acc read for free; the write stays
undeclared, which is sound for reducers specifically: updates commute
(combine ops are associative/commutative) and VerifyReducerEpoch pins
init/finalize to straight-line code, so no cross-statement write
ordering exists to lose inside pipelined/parallel regions.
- ParseReducerUpdate replaces the TileOperator ctor; verifier, planner
and the legacy shim consume it. The verifier visits the target's
indices but not the target load itself (a reducer load anywhere else
is still the usual compile error).
- Pipeline statement classification counts kTLPerIterationOp calls as
compute (BufferRegionCollector, IsPureCopyStmt/GetSinglePureCopyTileOp
would otherwise classify an update-only loop as copy-like now that the
call no longer parses as a tile op); the shared IsPerIterationOpCall
helper also replaces ParallelOp's inline attribute lookup.
- The nested-loop checker exemption is gone entirely: a builtin has no
TLOpBuilder, so it passes the tile-op ban by construction.
Verified: reducer v2 (39) + v1 (93) + copy suites pass; the reducer
benchmark (all cases T.Pipelined, exercising the pipeline dependency and
classification paths) reproduces identical plans and timings.
* [Cleanup] Drop the TLPerIterationOp attribute for a direct predicate
The attribute named a category that does not carve reality at a joint:
nearly every builtin intrinsic "executes per iteration" inside
T.Parallel, so the name suggested a structural property while actually
being an opt-in marker with exactly one member. What its two consumers
really key on is reducer_update itself — the hidden write to a
deferred-layout buffer (ParallelOp inference) and compute-vs-copy
statement classification (pipeline planning).
Replace it with IsReducerUpdateCall next to the op declaration; both
consumers and ParallelOpNode::has_reducer_update_ now say plainly what
they check. A future sibling op (e.g. a rescale intrinsic) extends the
one predicate.
* [Doc] Record the operand-driven defer protocol's known limitation
Deferring assumes the pending drivable operand will be solved by its
producer. An operand with no producer at all (an update reading a
fragment nothing ever writes) defers on every free-level round and
surfaces as a compile-time "layout can not be inferred" ICHECK — never
wrong code, and only on degenerate input. A sound fix needs the
inference engine to order solving (or guarantee a final round); that
redesign is deliberately deferred to a follow-up rather than patched
into the protocol here.
* [Layout] Drop the operand-driven defer protocol from ParallelOp
The three-piece workaround (has_reducer_update_ detection, the
FindOperandDriver three-state ladder, free-mode defer) patched solve
ordering for anchor-less loops from inside a single op. The real defect
is the engine's register-count attempt scoring: the bias row-broadcast
counterexample reproduces the same mis-selection with no anchor-less
loop involved, so the protocol treated a symptom. Direction for the
proper fix (IO-aware cost model + declared solve ordering) is recorded
in the layout RFC.
Reducer-update loops return to the generic two-candidate free
inference. The candidate-validation fallback still bounds the damage:
fp16 GEMV 0.0615 ms vs 0.0736 ms v1 baseline (down from 0.0346 ms with
the protocol), rowsum unchanged across all sizes. reducer v2 (39),
v1 reduce (93), atomic, and copy suites pass.
* [Cleanup] Unwrap the free-inference branch's vacuous layout guard
The inner !loop_layout_.defined() check was the seam where the deleted
operand-driven block could adopt a driver layout before the generic
two-candidate path ran. The enclosing branch already requires an
undefined layout and nothing assigns one in between, so the guard is
always true now.
* [Cleanup] Drop the unreachable post-selection candidate fallback
ChooseBestCandidate already validates both candidates with the same
flags and returns a valid one whenever either validates (single-valid
returns it directly, both-valid picks among valid, both-invalid leaves
nothing to fall back to). The post-hoc re-validation could therefore
never flip the choice; an instrumented run across the reducer suites
and the rowsum/GEMV benchmarks confirms it never fires. It was added
alongside the destination-chain override work and mistakenly credited
for making multi-slot overrides safe — the chain registration is what
did that. Comments referencing the fallback updated.
* [Cleanup] Tighten loop_partition after the multiplicity-marker rework
- MultiplicityMarkerLowerer only rewrites statement-level AttrStmt
markers; derive from StmtMutator instead of StmtExprMutator so
expression subtrees are not traversed.
- Drop a duplicated example paragraph in the partition guard comment.
- Discard InverseWithLevel's IterMapLevel explicitly instead of
keeping an unused pair binding.
* [Test] Cover reducer accumulation inside a T.Pipelined body
Pipeline planning must classify the per-iteration update intrinsic as
compute while the shared staging copy forms the only copy stage; the
epoch spans every pipeline iteration with a single finalize. This path
previously had no coverage at all.
Warp specialization is disabled in the test: a WS-split epoch is out of
scope for reducer v2 and currently surfaces as the finalize lowering's
thread-bounds ICHECK rather than a user-facing diagnostic.
* [Cleanup] Revert the pipeline-planning reducer classification marks
The three IsReducerUpdateCall sites marked update intrinsics as compute
during pipeline statement classification. Empirically they change
nothing: with the marks compiled out, both a plain pipelined
accumulation and an adversarial mixed statement (copy-shaped shared
store plus an update in one parallel loop) produce byte-identical
kernels — a bare intrinsic never classifies as a copy on any consumer
path that matters while warp-specialized reducer epochs are rejected
outright. pipeline_planning.cc returns to its upstream state, and the
now-unused IsReducerUpdateCall predicate is dropped with it.
* [Lint] Apply pre-commit clang-format L
Lei Wang committed
024fc72896910a9459825dc25ffa7bf81b73758e
Parent: 75cf763
Committed by GitHub <noreply@github.com>
on 8/13/2026, 5:53:24 AM