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

Mojo struct

Im2colALoader

struct Im2colALoader[dtype: DType, BK: Int = Int(16), num_m: Int = Int(2), num_n: Int = Int(2), c_aligned: Bool = False]

Fused-conv A loader: input_ptr + this simdgroup's prebaked pixel anchors and carried K-state. The gather reads the A MMA-fragment from NHWC via MmaOpApple.mma_im2col (the im2col matrix is non-affine, so it is not a distribute-expressible TileTensor -- KB exceptions/apple-mma-fragment-is-not-distribute-expressible). The anchor prebake + K-state strength-reduction design: KB kernels/apple-conv2d-im2col.

conv is deliberately NOT a struct field: an 11-Int32 ConvIm2colParams held in a by-value loader spilled to a GENERIC addrspace alloca that Metal's AIR backend mishandles (the FLUX-concat addrspace-loss class) and crashed MTLCompilerService -- so it is threaded as an accumulate_strip arg. The prebaked anchors are plain Int32 arrays (safe, like DenseALoader's slab).

Lifetime: input_ptr uses UntrackedOrigin (struct fields cannot expose AnyOrigin); the NHWC input kernel arg outlives the K-loop gather.

Parameters

dtype (DType): Element type of the NHWC input operand (bf16 for now).
BK (Int): K-strip depth per accumulate step; must match the body's BK tiling (defaults to 16).
num_m (Int): M MMA-fragment count per simdgroup, SG_M / 16 (defaults to 2).
num_n (Int): N MMA-fragment count per simdgroup, SG_N / 16 (defaults to 2).
c_aligned (Bool): When True, assume conv.C is a multiple of 8 so the channel run is contiguous and a single width-8 load suffices (defaults to False).

Fields

input_ptr (Pointer[Scalar[dtype], ImmUntrackedOrigin, _safe=False]):
h_base (Array[Int32, Im2colALoader[dtype, BK, num_m, num_n, c_aligned].NUM_ROWS]):
w_base (Array[Int32, Im2colALoader[dtype, BK, num_m, num_n, c_aligned].NUM_ROWS]):
batch_base (Array[Int32, Im2colALoader[dtype, BK, num_m, num_n, c_aligned].NUM_ROWS]):
c0 (Int32):
r (Int32):
s (Int32):
k_total (Int32):

Implemented traits

AOperandLoader, AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable

`comptime` members

`b_type`

comptime b_type = dtype

`in_type`

comptime in_type = dtype

`num_m_mmas`

comptime num_m_mmas = num_m

`num_n_mmas`

comptime num_n_mmas = num_n

`NUM_ROWS`

comptime NUM_ROWS = (num_m * Int(2))

Methods

`init`

def __init__(out self, input_ptr: Pointer[Scalar[dtype], ImmUntrackedOrigin, _safe=False], row_base: Int32, rb: Int32, cb: Int32, conv: ConvIm2colParams)

Prebake this lane's per-row im2col anchors + seed the K-state.

row_base is the simdgroup's absolute M base (_sg_row_base); rb/cb are this lane's MMA-fragment row/col (matching MmaOpApple.rb/.cb). Row ri covers M-fragment ri // 2, row-half ri % 2, i.e. output pixel row_base + (ri//2)*16 + (ri%2)*8 + rb. The K-state is decomposed once here for the lane's first strip (k0base = 2*cb); every later strip advances it incrementally in accumulate_strip.

Args:

input_ptr (Pointer[Scalar[dtype], ImmUntrackedOrigin, _safe=False]): Pointer to the 4-D NHWC input tensor; the im2col gather reads A MMA-fragments from it.
row_base (Int32): Simdgroup's absolute M-row base (_sg_row_base); anchors this lane's output pixels.
rb (Int32): This lane's MMA-fragment row offset within the simdgroup (matches MmaOpApple.rb).
cb (Int32): This lane's MMA-fragment column offset within the simdgroup (matches MmaOpApple.cb); seeds the K-state at k0base = 2*cb.
conv (ConvIm2colParams): Conv geometry for the im2col gather; threaded as a value arg, not a struct field (see the struct docstring for the Metal addrspace reason).

`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], mut accum: Array[SIMD[DType.float32, SIMDLength(8)], MmaOpApple[DType.float32, Self.in_type, Self.num_m_mmas, Self.num_n_mmas].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)

Parameters​

Fields​

Implemented traits​

comptime members​

b_type​

in_type​

num_m_mmas​

num_n_mmas​

NUM_ROWS​

Methods​

__init__​

accumulate_strip​