Skip to main content

allreduce_section

Function allreduce_section 

Source
pub fn allreduce_section(
    col: &TpCollective,
    shape: (usize, usize),
    floats: bool,
    target: &SpatialTarget,
    certification: SectionCertification,
    narrow: &mut dyn FnMut(SpatialProgram) -> SpatialProgram,
) -> Result<SpatialProgram, TpReject>
Expand description

ONE all-reduce as a section: out = Σ_c partials[c] over a rows × cols tensor, tile by tile — chip c’s partial Loaded on chip c (where it lives) and handed across the chip boundary, all of them folded by an AllReduce on chip 0, the sum stored as the replicated out.

★ BUILT AS A SUM, THEN SPREAD. The fold is lowered first as what it is on one chip — a chips-ary elementwise sum, tiled, placed and scheduled by the pinned passes — and then each operand’s Load moves to its chip, its hand-off becomes one that crosses a chip boundary, and the fold becomes the AllReduce whose combine it already is. The passes run again from the placement, so the collective is sized and choreographed like any other hand-off, and certified as certification says: on its own, or as a section of the composed program. The one-chip sum is only scheduled — its certificate would be dropped with the placement it was bound to.

★ AN INTEGER SUM IS EXACT. Float partials (floats) are summed in the target’s element precision, an approximation whose order the claim names. Integer partials — an int8 model’s accumulators — are summed in the widest integer the elementwise datapath holds, exactly, in any order, and their storage is narrowed by narrow as every integer section’s is before it is tiled.