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 trait

AOperandLoader

One BK-strip A contribution for the shared Apple GEMM body.

accumulate_strip is the single seam between the plain-GEMM and fused-conv kernels: for strip index k_strip, it loads this strip's A side (slab tile vs online im2col gather) and runs the MMA into accum. Both impls consume the same pre-tiled b_sub and write the same accum; only the A side differs. Any per-simdgroup slab is built once by the caller before the K-loop and held in the impl (see DenseALoader), so accumulate_strip only adds this strip's K offset (k_strip * BK).

The loader is parametrized on the GEMM in_type, so the seam method names the concrete _BodyMma[Self.in_type] / its AccumType. The shared body infers the loader's CONCRETE type (infer-only L) and keys its own mma_op/accum/B on L.in_type so both sides of the erased dispatch agree (KB patterns/trait-type-erasure-and-stride-layout-workaround); at the run/run_conv call sites L.in_type == Self.in_type is proven, so there is no runtime cost.

Implemented traits

AnyType

comptime members

b_type

comptime b_type

in_type

comptime in_type

num_m_mmas

comptime num_m_mmas

num_n_mmas

comptime num_n_mmas

Required methods

accumulate_strip

def accumulate_strip[bounded: Bool, b_layout: TensorLayout](mut self, mut mma_op: MmaOpApple[DType.float32, Self.in_type, _Self.num_m_mmas, _Self.num_n_mmas, b_type=Self.b_type], mut accum: Array[SIMD[DType.float32, SIMDLength(8)], MmaOpApple[DType.float32, Self.in_type, _Self.num_m_mmas, _Self.num_n_mmas, b_type=Self.b_type].num_accum], b_sub: TileTensor[Self.b_type, b_layout, ImmutAnyOrigin], conv: ConvIm2colParams, k_strip: Int32, *, a_valid_rows: Int32, b_valid_cols: Int, k_valid: Int)