IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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 struct

Row

struct Row[params: ContextParams, accum: DType, dtype: DType, axis: Int, rank: Int, AxisSize: CoordLike, is_cached: Bool]

The body's handle to one logical row. See the module-level note.

Carries no input_fn field: it's threaded as a fresh value-closure ARGUMENT to every method that needs it (a "forwarding bridge"), rather than stored as a typed field. A field whose type is a trait-bound closure that references a sibling struct parameter (dtype) hits a Mojo limitation the moment the struct itself is named generically elsewhere (e.g. as another method's parameter or return type) — "lacking evidence to prove correctness" even when the constraint holds by construction. Method-level trait-bound params (fresh per call, not part of Self's persistent identity) don't hit this; every other closure here (Contribute, G, Write, ...) already uses that shape.

Parameters​

  • ​params (ContextParams): Comptime tier parameters (target, tier, widths).
  • ​accum (DType): Accumulator dtype the reduce closures compute in.
  • ​dtype (DType): The primary (axis-walked) input dtype.
  • ​axis (Int): Reduced axis (statically known).
  • ​rank (Int): Tensor rank.
  • ​AxisSize (CoordLike): The reduced-axis length's CoordLike type — ComptimeInt[N] when known at comptime (enables the register-resident _fuse path) or a dynamic Scalar type otherwise. One type carries both the comptime-known-ness and (via .value()) the runtime length, so the caller supplies a single AxisSize-typed value instead of a separate comptime "is it static" signal.
  • ​is_cached (Bool): Whether this view may cache the row — true for normalizing multi-phase bodies; false for single-pass true reductions where caching buys nothing.

Fields​

  • ​row_il (IndexList[rank]):
  • ​axis_size (AxisSize):
  • ​ctx (Context[params]):

Implemented traits​

AnyType, Copyable, Deinitable, Movable

Methods​

__init__​

def __init__[InputFn: def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable](out self, row_coords: Coord, axis_size: AxisSize, ctx: Context[params], input_fn: InputFn) where (eq InputFn.dtype, dtype)

Builds the view, staging the row strip into registers on cache-eligible tiers.

Parameters:

  • ​InputFn (def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable): The value-closure type of input_fn.

Args:

  • ​row_coords (Coord): The row's coords.
  • ​axis_size (AxisSize): Length of the reduce axis, as an AxisSize-typed value (ComptimeInt[N]() when statically known, else a dynamic Scalar).
  • ​ctx (Context[params]): The dispatch bundle.
  • ​input_fn (InputFn): Loads a tile of the primary input at a coord. Used only here (fuse-path staging); not retained as a field — later non-fuse-path methods take it again as an argument.

reduce​

def reduce[M: ReduceOp, Contribute: def[w: Int](SIMD[dtype, w], IndexList[rank]) -> SIMD[accum, w] & RegisterPassable & ImplicitlyCopyable, InputFn: def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable](mut self, contribute: Contribute, input_fn: InputFn) -> M where (eq Contribute.dtype, dtype) where (eq Contribute.rank, rank) where (eq Contribute.accum, accum) where (eq InputFn.dtype, dtype)

Runs one reduction phase: folds contribute(tile, idx) into M across the row, cross-thread-joins, and returns the joined state. The body reads M's result fields (e.g. ReduceSum.acc[0], Welford.mean[0], ArgMin.best_idx).

contribute is a value closure (its copy-captured state rides the value); the tier loop's reduce_tile captures it directly. input_fn (the primary-input loader) is likewise a fresh value argument, used only on the non-fuse (streaming) path — see the module note above Row on why it isn't a stored field.

Parameters:

  • ​M (ReduceOp): The monoid for this phase (instantiated at the view's SIMD width, e.g. ReduceSum[accum, params.simd_width]).
  • ​Contribute (def[w: Int](SIMD[dtype, w], IndexList[rank]) -> SIMD[accum, w] & RegisterPassable & ImplicitlyCopyable): The value-closure type of contribute.
  • ​InputFn (def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable): The value-closure type of input_fn.

Args:

  • ​contribute (Contribute): Per-element math; maps the loaded primary tile + its coord to the accumulator-dtype contribution.
  • ​input_fn (InputFn): Loads a tile of the primary input at a coord (non-fuse path only).

Returns:

M: The joined M, broadcast to every participant.

def reduce[T: DType, shared: Bool, //, M: ReduceOp, Contribute: def[w: Int](SIMD[T, w], IndexList[rank]) -> SIMD[accum, w] & RegisterPassable & ImplicitlyCopyable, InputFn: def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable, Compute: def[w: Int](SIMD[dtype, w], IndexList[rank]) -> SIMD[T, w] & RegisterPassable & ImplicitlyCopyable](mut self, over: RowCache[params, T, dtype, axis, rank, (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))) if is_cached and (xor is_cpu[params.target](), True) and (params != ReduceTier(Int(3))) and (eq (add rank, -1), axis) and (xor (eq (xor (lt AxisSize.static_value if AxisSize.is_static_value else Int(0), 1), True), False), True) and (eq (AxisSize.static_value if AxisSize.is_static_value else Int(0) % params), 0) and (xor (eq (lt (mul (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))), size_of[SIMD[dtype, SIMDLength(params.simd_width)]]()), 257), False), True) else Int(1), params.simd_width, AxisSize.static_value if AxisSize.is_static_value else Int(0), is_cached and not is_cpu[params.target]().__bool__() and (params != ReduceTier.Splitk) and (axis == (rank - Int(1))) and (AxisSize.static_value if AxisSize.is_static_value else Int(0) > Int(0)) and ((AxisSize.static_value if AxisSize.is_static_value else Int(0) % params) == Int(0)) and (Int((mul (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))), size_of[SIMD[dtype, SIMDLength(params.simd_width)]]())) <= Int(256)), shared], contribute: Contribute, input_fn: InputFn, compute: Compute) -> M where (eq Contribute.T, T) where (eq Contribute.rank, rank) where (eq Contribute.accum, accum) where (eq InputFn.dtype, dtype) where (eq Compute.dtype, dtype) where (eq Compute.rank, rank) where (eq Compute.T, T)

Cached overload of reduce: folds contribute over a previously cached row instead of the primary input — reads the staged value at each column (registers on the fast path, recompute fallback otherwise). Used for a phase that reduces a computed cross-phase value (e.g. sum of the fused intermediate squared).

Resolved at compile time from plain Mojo function overloading — whether the call site passes over — not a runtime-checked Optional, so the cached-vs-plain choice costs nothing at runtime.

contribute is a value closure (its copy-captured state rides the value); the tier loop's reduce_tile captures it directly. input_fn / compute (the same producer passed to the cache call that built over) are likewise fresh value arguments, used only on the non-fuse (recompute) path — RowCache doesn't carry them as fields (see its module note). axis_size / ctx are read from self rather than threaded as separate arguments — the owning view already holds both.

Parameters:

  • ​T (DType): The staged value's dtype (inferred from over).
  • ​shared (Bool): The staged handle's shmem flag (inferred).
  • ​M (ReduceOp): The monoid for this phase.
  • ​Contribute (def[w: Int](SIMD[T, w], IndexList[rank]) -> SIMD[accum, w] & RegisterPassable & ImplicitlyCopyable): The value-closure type of contribute.
  • ​InputFn (def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable): The value-closure type of input_fn.
  • ​Compute (def[w: Int](SIMD[dtype, w], IndexList[rank]) -> SIMD[T, w] & RegisterPassable & ImplicitlyCopyable): The value-closure type of compute.

Args:

Returns:

M: The joined M, broadcast to every participant.

elementwise​

def elementwise[G: def[w: Int](SIMD[dtype, w], IndexList[rank]) -> None & RegisterPassable & ImplicitlyCopyable, InputFn: def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable](mut self, g: G, input_fn: InputFn) where (eq G.dtype, dtype) where (eq G.rank, rank) where (eq InputFn.dtype, dtype)

Per-element terminal (normalizing reductions): runs g(tile, idx) over the row; g computes and stores each output element. Reuses the register cache when present.

g is a value closure (its copy-captured state rides the value); the tier loop's emit_tile captures it directly. input_fn is likewise a fresh value argument (non-fuse path only) — see the module note above Row.

Parameters:

Args:

  • ​g (G): Per-element write; closes over the joined results and side inputs, and stores via the body's output closure.
  • ​input_fn (InputFn): Loads a tile of the primary input at a coord (non-fuse path only).

def elementwise[T: DType, shared: Bool, //, G: def[w: Int](SIMD[T, w], IndexList[rank]) -> None & ImplicitlyCopyable, InputFn: def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable, Compute: def[w: Int](SIMD[dtype, w], IndexList[rank]) -> SIMD[T, w] & RegisterPassable & ImplicitlyCopyable](mut self, over: RowCache[params, T, dtype, axis, rank, (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))) if is_cached and (xor is_cpu[params.target](), True) and (params != ReduceTier(Int(3))) and (eq (add rank, -1), axis) and (xor (eq (xor (lt AxisSize.static_value if AxisSize.is_static_value else Int(0), 1), True), False), True) and (eq (AxisSize.static_value if AxisSize.is_static_value else Int(0) % params), 0) and (xor (eq (lt (mul (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))), size_of[SIMD[dtype, SIMDLength(params.simd_width)]]()), 257), False), True) else Int(1), params.simd_width, AxisSize.static_value if AxisSize.is_static_value else Int(0), is_cached and not is_cpu[params.target]().__bool__() and (params != ReduceTier.Splitk) and (axis == (rank - Int(1))) and (AxisSize.static_value if AxisSize.is_static_value else Int(0) > Int(0)) and ((AxisSize.static_value if AxisSize.is_static_value else Int(0) % params) == Int(0)) and (Int((mul (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))), size_of[SIMD[dtype, SIMDLength(params.simd_width)]]())) <= Int(256)), shared], g: G, input_fn: InputFn, compute: Compute) where (eq G.T, T) where (eq G.rank, rank) where (eq InputFn.dtype, dtype) where (eq Compute.dtype, dtype) where (eq Compute.rank, rank) where (eq Compute.T, T)

Cached overload of elementwise: the per-element terminal g receives the staged value at each column (from over) instead of the primary input — registers on the fast path, recompute fallback otherwise. g may additionally read other columns of over via over.load (e.g. rope's rotate-half partner, staged in shmem).

Resolved at compile time from plain Mojo function overloading — whether the call site passes over — not a runtime-checked Optional, so the cached-vs-plain choice costs nothing at runtime.

g is a value closure (its copy-captured state rides the value); the tier loop's emit_tile captures it directly. input_fn / compute (the same producer passed to the cache call that built over) are likewise fresh value arguments, used only on the non-fuse (recompute) path — RowCache doesn't carry them as fields (see its module note). axis_size / ctx are read from self rather than threaded as separate arguments — the owning view already holds both.

The bound on G is ImplicitlyCopyable (not RegisterPassable like the other value overloads) because g may copy-capture non-register-passable state (e.g. over itself, for a terminal that also reads a different column via over.load); this overload runs inside the launched kernel (not across a launch boundary), so register-passability is not required.

Parameters:

  • ​T (DType): The staged value's dtype (inferred from over).
  • ​shared (Bool): The staged handle's shmem flag (inferred).
  • ​G (def[w: Int](SIMD[T, w], IndexList[rank]) -> None & ImplicitlyCopyable): The value-closure type of g.
  • ​InputFn (def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable): The value-closure type of input_fn.
  • ​Compute (def[w: Int](SIMD[dtype, w], IndexList[rank]) -> SIMD[T, w] & RegisterPassable & ImplicitlyCopyable): The value-closure type of compute.

Args:

emit​

def emit[Write: def(IndexList[rank]) -> None & RegisterPassable & ImplicitlyCopyable](mut self, write: Write) where (eq Write.rank, rank)

Per-row terminal (true reductions): runs write(out_coord) once on the canonical writer for this row. write closes over the joined results and stores them, and is taken as a value arg (its copy-captured state rides the value) rather than a comptime capturing parameter.

Unlike the per-element terminal (elementwise, plain or cached), emit needs no comptime-form + value-shim pair: once is an in-kernel thread predicate (not a launch boundary and not a per-element loop), so the emit_once wrapper (which pins the reduced axis and calls write) is itself a value closure passed straight to once, with no per-element env-hoist risk.

Parameters:

Args:

  • ​write (Write): Per-row terminal store, passed as a value. Stores the row's output(s) at the collapsed coordinate (reduced axis pinned to 0).

cache​

def cache[T: DType, Compute: def[w: Int](SIMD[dtype, w], IndexList[rank]) -> SIMD[T, w] & RegisterPassable & ImplicitlyCopyable, shared: Bool = False](mut self, compute: Compute) -> RowCache[params, T, dtype, axis, rank, (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))) if is_cached and (xor is_cpu[params.target](), True) and (params != ReduceTier(Int(3))) and (eq (add rank, -1), axis) and (xor (eq (xor (lt AxisSize.static_value if AxisSize.is_static_value else Int(0), 1), True), False), True) and (eq (AxisSize.static_value if AxisSize.is_static_value else Int(0) % params), 0) and (xor (eq (lt (mul (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))), size_of[SIMD[dtype, SIMDLength(params.simd_width)]]()), 257), False), True) else Int(1), params.simd_width, AxisSize.static_value if AxisSize.is_static_value else Int(0), is_cached and not is_cpu[params.target]().__bool__() and (params != ReduceTier.Splitk) and (axis == (rank - Int(1))) and (AxisSize.static_value if AxisSize.is_static_value else Int(0) > Int(0)) and ((AxisSize.static_value if AxisSize.is_static_value else Int(0) % params) == Int(0)) and (Int((mul (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))), size_of[SIMD[dtype, SIMDLength(params.simd_width)]]())) <= Int(256)), shared] where (eq Compute.dtype, dtype) where (eq Compute.rank, rank) where (eq Compute.T, T)

Caches a per-row value the later phases reuse: runs compute(tile, idx) once over the row and stages the result so the cached overloads of reduce / elementwise (and, with shared=True, load at any column) read it back instead of recomputing.

On a cache-eligible tier this writes each participant's tiles into registers (and, when shared, publishes the row to a shmem strip with one row-local sync); otherwise the returned handle recomputes on access (always correct) — the caller re-supplies input_fn / compute to the cached reduce / elementwise overloads on that path, since RowCache doesn't carry them as fields (see its module note). compute runs exactly once per element on the fast path but may re-run on the fallback, so any side effect it performs (e.g. a residual emit) must be idempotent.

Parameters:

  • ​T (DType): The cached value's dtype.
  • ​Compute (def[w: Int](SIMD[dtype, w], IndexList[rank]) -> SIMD[T, w] & RegisterPassable & ImplicitlyCopyable): The value-closure type of compute.
  • ​shared (Bool): Publish the row to shmem for cross-participant load.

Args:

  • ​compute (Compute): Per-element producer from the primary tile + coord.

Returns:

RowCache[params, T, dtype, axis, rank, (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))) if is_cached and (xor is_cpu[params.target](), True) and (params != ReduceTier(Int(3))) and (eq (add rank, -1), axis) and (xor (eq (xor (lt AxisSize.static_value if AxisSize.is_static_value else Int(0), 1), True), False), True) and (eq (AxisSize.static_value if AxisSize.is_static_value else Int(0) % params), 0) and (xor (eq (lt (mul (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))), size_of[SIMD[dtype, SIMDLength(params.simd_width)]]()), 257), False), True) else Int(1), params.simd_width, AxisSize.static_value if AxisSize.is_static_value else Int(0), is_cached and not is_cpu[params.target]().__bool__() and (params != ReduceTier.Splitk) and (axis == (rank - Int(1))) and (AxisSize.static_value if AxisSize.is_static_value else Int(0) > Int(0)) and ((AxisSize.static_value if AxisSize.is_static_value else Int(0) % params) == Int(0)) and (Int((mul (Int((add (mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width), AxisSize.static_value if AxisSize.is_static_value else Int(0), -1)) // Int((mul _resolve_warp_size() if (params == ReduceTier(Int(1))) else params.BLOCK_SIZE, params.simd_width))), size_of[SIMD[dtype, SIMDLength(params.simd_width)]]())) <= Int(256)), shared]: A RowCache handle over the cached row.

Models

Cookbook

Puzzles

Blog

Community

Contact