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_vnni
struct Inner_matmul_vnni[saturated_vnni: Bool]
Integer CPU matmul microkernel using x86 VNNI or ARM NEON dot-product.
Implements InnerMatmulKernel using 4-element 8-bit integer dot-product
instructions to accumulate int32 partial products. On x86 it accumulates
from a uint8 A and an int8 B, emitting the VNNI vpdpbusd
instruction when the target has VNNI (AVX-512 VNNI or AVX-VNNI) and falling
back to an AVX2 emulation otherwise (AVX2 is the minimum, not AVX-512). On
ARM it accumulates from int8/int8 or uint8/uint8 operands using the
NEON dot-product (_neon_dotprod). Handles K tails (remainder < 4) by
packing A into a local buffer or using AVX-512 masked loads.
Parameters
- saturated_vnni (
Bool): WhenTrue, uses the saturating x86 variant (dot_i8_to_i32_saturated_x86), which requires theAoperand to lie in[0, 127]. This saves instructions on the AVX2 emulation path used when the target lacks VNNI; on VNNI hardware it emits the same instruction as the default.
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 processed per inner kernel iteration along the M dimension. - kernel_cols (
Int): Number of columns processed per inner kernel iteration along the N dimension. - simd_size (
Int): SIMD vector width used for the int8 dot-product instructions.
Args:
- c (
TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output matrix tile accumulating the matmul partial sums. - a (
TileTensor[Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Input A matrix tile being processed. - 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. - global_offset (
GemmShape): Global (M, N, K) coordinate offset for this tile. - global_bound (
GemmShape): Global (M, N, K) bound of the matrices, used for boundary checks. - tile_n_k (
IndexList[Int(2)]): Dynamic tile sizes along the N and K dimensions. - skip_boundary_check (
Bool): Whether to skip boundary checks when storing results.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!