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
Inner_matmul_neon
struct Inner_matmul_neon
ARM NEON lane-FMA microkernel for CPU matmul.
Implements InnerMatmulKernel using vectorized lane-broadcast FMA to
accumulate partial products. Loads simd_size elements of A, broadcasts
each lane across all kernel rows, and computes dot products against packed
B columns. Handles tail elements in the K dimension with scalar cleanup.
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
InnerMatmulKernel,
Movable
Methods
__inner_matmul__
def __inner_matmul__[kernel_rows: Int, kernel_cols: Int, simd_size: Int](self, c: TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b_packed: TileTensor[Storage=b_packed.Storage, address_space=b_packed.address_space, linear_idx_type=b_packed.linear_idx_type], global_offset: GemmShape, global_bound: GemmShape, tile_n_k: IndexList[Int(2)], skip_boundary_check: Bool)
Utility function on the inner loop. Run the inner kernel on the whole (kernel_rows, TileN, TileK) tile.
Parameters:
- kernel_rows (
Int): Number of rows in the M-dimension tile processed per invocation. - kernel_cols (
Int): Number of columns in the N-dimension tile processed per invocation; must be a multiple ofsimd_size. - simd_size (
Int): SIMD vector width used for lane-broadcast FMA on packed B elements.
Args:
- c (
TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output C matrix tile accumulating the partial products. - a (
TileTensor[Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Input A matrix tile read in row-major order. - b_packed (
TileTensor[Storage=b_packed.Storage, address_space=b_packed.address_space, linear_idx_type=b_packed.linear_idx_type]): Packed B matrix tile in cache-friendly layout; must be rank 3. - global_offset (
GemmShape): Global (M, N, K) coordinate offset of this tile within the whole matmul problem space. - global_bound (
GemmShape): Global (M, N, K) extent of the full matmul problem space, used for boundary clipping on the C store. - tile_n_k (
IndexList[Int(2)]): (n, k) extents of the tile to process in the N and K dimensions. - skip_boundary_check (
Bool): Whether the tile lies fully withinglobal_boundso boundary clipping can be skipped.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!