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 module

rowwise

Unified row-wise reduction scaffolder.

Author surface for kernels targeting both CPU and GPU with one body. Re-exports ContextParams and Context from algorithm.rowwise_types, and provides comptime-dispatched reduce / pjoin / once / simd / launch that pick the CPU or GPU backend from params.target (a comptime StaticString field on ContextParams).

Author contract β€” one body, two targets:

def reduce_sum[...](shape, target="cpu", ctx=None) raises:
    var axis_size = ...

    @always_inline
    def body[
        params: rowwise.ContextParams
    ](row_coords: Coord, mut ctx: rowwise.Context[params]) {var axis_size}:
        comptime W = params.emit_tile_width
        var state = ReduceSum[accum_type, W]()

        @always_inline
        def tile_fn[ws: Int, _r: Int](mut state: ReduceSum[accum_type, W], coords: IndexList[_r]):
            ...
            state.accumulate[dtype, ws](...)

        rowwise.reduce(row_coords, axis_size, ctx, state, tile_fn)
        rowwise.pjoin(state, ctx)

        @always_inline
        def emit() {...}:
            ...

        rowwise.once(emit, ctx)

    # `num_phases=1` if the output collapses the axis (one value per
    # row); `num_phases>1` if the output keeps the input shape
    # (per-element emission along axis).
    rowwise.launch[body, axis=..., simd_width=..., target=target, num_phases=1, ...](shape, ctx)

The body never branches on target directly. rowwise.reduce / pjoin / once / simd each comptime-dispatch via comptime if is_cpu[params.target]():, so GPU primitives never appear in CPU codegen and vice versa. tile_fn and emit are value closures (their copy-captured state rides the value); a reduce phase's state is threaded through reduce as a mut argument rather than captured, since a captured accumulator can't be mutated through a value closure.

computationally_expensive is a GPU-only author hint, forwarded to the GPU backend and ignored by the CPU backend. BLOCK_SIZE / TILED_BLOCK_SIZE / COOPERATIVE_BLOCK_SIZE / WARP_BLOCK_WARPS are GPU launch-geometry constants private to launch β€” no caller has ever needed to override them, so they are not part of its public signature (see the tier-dispatch notes above launch for what they do). The CPU backend runs sync_parallelize over output rows / slices (no DeviceContext); the GPU backend launches kernels through the provided DeviceContext.

Structs​

  • ​Row: The body's handle to one logical row. See the module-level note.
  • ​RowCache: A row value cached once by Row.cache and read by later phases. See the module note above for the two backings.

Functions​

  • ​launch: Top-level scaffolder. num_phases picks reduce-shaped (num_phases == 1, output collapses the reduced axis, terminal emit) vs. normalize-shaped (num_phases > 1, output keeps the input shape, terminal elementwise) β€” see the module note above for the full picture, including the GPU tier-dispatch decision tree.
  • ​once: Runs emit exactly once per logical output row.
  • ​pick_simd_width: Picks the SIMD width a rowwise body should use throughout.
  • ​pjoin: Cross-thread join for a single monoid state.
  • ​reduce: Drives tile_fn over the reduce axis on the target backend, with no monoid state β€” pure per-tile iteration. See the state-carrying overload below for reduce phases.
  • ​strided_load: Loads simd_width values from addr with a compile-time stride.