[Refactor] Give reducers a first-class PartialFragment layout solved by layout inference (#3093)
* Steer reducer destinations during layout inference * Remove extra reducer steering attempt filtering * handle reducer problem3 * bug fix * [Refactor] Give reducers a first-class PartialFragment layout solved by layout inference The reducer's physical plan (narrow vs wide partials) used to be decided after inference by ReducerPlanAndMaterialize, with dst-steering emulating the missing propagation through side channels (wake tables, hint plumbing, ownership gates); the same induced layout was computed twice and kept consistent by convention only, and the reducer split the use-def graph into disconnected components. This makes the reducer's layout a first-class solver value that fully determines the physical plan: * New layout kind `tl.PartialFragment` (PartialFragmentNode : FragmentNode): Fragment algebra whose replication coordinate carries an explicit combine decomposition — under a low-bits convention enforced at construction, `_rep % combine_size` enumerates addend lanes the finalize collective must reduce and `_rep / combine_size` enumerates equal-value copy groups. (storage algebra, combine width) is the complete plan: the collective steps are derivable from the node alone (CombineSteps), the FullParticipant object (combine == replicate == participants, thread == _rep) is the wide fallback, and combine width is part of layout identity (same-storage different-decomposition proposals widen at the commit point instead of unifying silently). A boundary-preserving condensation (CondenseReplicateVarKeepingBoundary) keeps the convention true through projection; kind-aware IsEqual in both dispatch directions prevents silent unification with plain Fragments. * Update nests (ParallelOp) record their reducer_update sites and, once their loop layout solves, propose the induced partial (per-site AnalyzeReducerUpdateSite proofs) or the participant-wide plan. The commit point resolves multi-site disagreement with a monotone widen-on-conflict rule; zero-update epochs are seeded wide before any inference level runs; tl.reducer_force_baseline pins the proposal. * finalize_reducer reads the solved partial and steers an unconstrained destination to its post-collective reading. The wake table, the LayoutInferArgs.reducer_update_sites hint plumbing and the reducer_use_list_ component glue are deleted: reducers travel ordinary use_list_ buffer edges. * ReducerPlanAndMaterialize reads the plan off the node: storage = AsPostCollective(), steps = CombineSteps(), wide lowering iff the node is the FullParticipant object. TryNarrowPlan is retired; what remains (DeriveNarrowLowering) derives only the lowering details that live outside the layout — the packed-lane heuristic and the destination-containment gate (which still demotes gracefully to the wide plan, preserving the unsatisfiable-steering fallback contract). The destination override chains are kept as a verified-no-op fallback with retirement telemetry. * T.annotate_layout accepts a PartialFragment on reducer buffers, pinning plans the automatic solution cannot reach (e.g. a middle plan between narrow and wide). Annotated partials are strict: a disagreeing update site raises with the narrow-rejection reason, and ComputeReducerPinnedCandidate back-propagates the annotation through the update loop into its input fragments. * Layout cost models price register-array spills: a fragment access whose physical slot index depends on the thread variable (detected by the shared FragmentThreadIndexProbe, extracted from HasThreadDependentFragmentIndex) demotes the whole array to local memory, charged as byte-denominated traffic in both models' `mem` channel so it competes with — rather than vetoes — global-bandwidth estimates. ChooseBestCandidate's anti-spill preference now applies regardless of where the loop stores (a fragment->fragment publication copy spills exactly like a fragment->global one). With generation and pricing combined, inference produces the wide publication chain's replicated layouts on its own: the reducer_v2 suite and the mhc_post_bwd production kernels pass with the destination override fully disabled. Testing (all with the kernel cache disabled): reducer_v2 suite 70 passed (new coverage: pinned middle plans, annotation steering input fragments, conflict diagnostics, combine-step roundtrip, same-storage different-combine widening at commit); language suite 1261 passed; layout-inference cost-model tests and maint/layout_inference goldens unchanged; generated CUDA for existing reducer kernels is byte-identical; mhc_post_bwd correctness passes and its six benchmark sizes show no latency change with zero STL/LDL in SASS. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
L
Lei Wang committed
6279e735f991a78d860fc242915c702d00ec8434
Parent: 905c75a
Committed by GitHub <noreply@github.com>
on 8/27/2026, 6:20:55 AM