For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

reduce

def reduce[params: ContextParams, //, TileFn: def[ws: Int, _r: Int](IndexList[_r]) -> None & ImplicitlyCopyable](row_coords: Coord, axis_size: Int, mut ctx: Context[params], tile_fn: TileFn)

Drives the tier-aware iteration over the reduce axis, with no monoid state — pure per-tile iteration for map/emit terminals (see the state-carrying overload below for reduce phases).

tile_fn[ws, _r] is invoked per tile with the scaffolder's SIMD width and the row's coords. It's a value closure (its copy-captured state — input/output closures — rides the value).

reduce runs iteration only; the body calls rowwise.pjoin separately for each state needing a cross-thread combine. Two-pass algorithms (softmax pass 2, layernorm rewrite) use reduce with no pjoin — pure map-style iteration.

Tier semantics, picked from ctx:

• Warp tier (cooperative, emit_tile_width == 1): lane i runs one scalar tile_fn[1] on element i (axis_size <= WARP_SIZE).
• Block tier (cooperative, emit_tile_width == 1): threads grid-stride by BLOCK_SIZE * simd_width along the axis; full SIMD tile_fn[simd_width] on aligned chunks, scalar tile_fn[1] on the tail.
• Tiled tier (emit_tile_width > 1): the thread owns emit_tile_width adjacent rows. Each axis step issues one tile_fn[emit_tile_width]; the body dispatches each lane to its own state. No cross-thread combine — rowwise.pjoin no-ops.
• Serial tier (one thread, one row): scalar tile_fn[1] per axis step.
• Split-K tier (num_rows < sm_count, large row): each block handles a slice of one row; rowwise.pjoin coordinates across blocks.

Parameters:

• ​params (ContextParams): The comptime dispatch parameters.
• ​TileFn (def[ws: Int, _r: Int](IndexList[_r]) -> None & ImplicitlyCopyable): The value-closure type of tile_fn.

Args:

• ​row_coords (Coord): The current row's coords. ctx.axis is the reduce axis; other dims are pinned.
• ​axis_size (Int): Length of the reduce axis.
• ​ctx (Context[params]): The dispatch bundle.
• ​tile_fn (TileFn): Per-tile callback; closes over input/output closures.

def reduce[State: ReduceOp, params: ContextParams, //, TileFn: def[ws: Int, _r: Int](mut State, IndexList[_r]) -> None & ImplicitlyCopyable](row_coords: Coord, axis_size: Int, mut ctx: Context[params], mut state: State, tile_fn: TileFn) where (eq TileFn.State, State)

State-carrying overload of reduce: same tier-aware iteration as above, but tile_fn takes the caller's monoid state as a mut argument (rather than capturing it — a captured accumulator can't be mutated through a value closure) and folds each tile into it.

Parameters:

• ​State (ReduceOp): The monoid type being accumulated.
• ​params (ContextParams): The comptime dispatch parameters.
• ​TileFn (def[ws: Int, _r: Int](mut State, IndexList[_r]) -> None & ImplicitlyCopyable): The value-closure type of tile_fn.

Args:

• ​row_coords (Coord): The current row's coords. ctx.axis is the reduce axis; other dims are pinned.
• ​axis_size (Int): Length of the reduce axis.
• ​ctx (Context[params]): The dispatch bundle.
• ​state (State): The caller's monoid accumulator, threaded through by mut reference and mutated in place by tile_fn.
• ​tile_fn (TileFn): Per-tile callback; closes over input closures and folds each tile into state.