Python Kernel Library

Pre-built AIE kernel wrappers for common operations. These provide ready-to-use Worker-compatible callables backed by optimized native AIE code. For the C++ kernel sources these wrap, see C++ AIE kernels.

Element-wise operations

Element-wise kernel factories: passthrough, scale, add, mul, relu.

add_ref

add_ref(a, b)

Numpy reference for add: element-wise a + b in float32.

Source code in python/iron/kernels/eltwise.py

def add_ref(a, b):
    """Numpy reference for [`add`][iron.kernels.eltwise.add]: element-wise ``a + b`` in float32."""
    return a.astype(np.float32) + b.astype(np.float32)

mul_ref

mul_ref(a, b)

Numpy reference for mul: element-wise a * b in float32.

Source code in python/iron/kernels/eltwise.py

def mul_ref(a, b):
    """Numpy reference for [`mul`][iron.kernels.eltwise.mul]: element-wise ``a * b`` in float32."""
    return a.astype(np.float32) * b.astype(np.float32)

scale_ref

scale_ref(x, factor)

Numpy reference for scale: x * factor[0] in int64.

factor is the 1-element int32 buffer the kernel reads the multiplier from; the harness casts the int64 product back to x.dtype, wrapping on overflow like the C++ store does.

Source code in python/iron/kernels/eltwise.py

def scale_ref(x, factor):
    """Numpy reference for [`scale`][iron.kernels.eltwise.scale]: ``x * factor[0]`` in int64.

    ``factor`` is the 1-element int32 buffer the kernel reads the multiplier
    from; the harness casts the int64 product back to ``x.dtype``, wrapping
    on overflow like the C++ store does.
    """
    return x.astype(np.int64) * np.int64(np.asarray(factor).reshape(-1)[0])

passthrough

passthrough(
    tile_size: int = 4096, dtype: type = int32
) -> ExternalFunction

Element-wise passthrough kernel: copies input tile to output tile.

Parameters:

Name	Type	Description	Default
tile_size	int	Number of elements per tile, a positive whole number of 64-byte vectors. Compiled into the kernel's loop bound.	4096
dtype	type	Element data type (np.uint8, np.int16, or np.int32).	int32

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for passThroughLine.

Raises:

Type	Description
ValueError	When dtype is not np.uint8, np.int16, or np.int32, or when tile_size is not a whole number of 64-byte vectors, or is nonpositive.

Source code in python/iron/kernels/eltwise.py

def passthrough(tile_size: int = 4096, dtype: type = np.int32) -> ExternalFunction:
    """Element-wise passthrough kernel: copies input tile to output tile.

    Args:
        tile_size: Number of elements per tile, a positive whole number of
            64-byte vectors. Compiled into the kernel's loop bound.
        dtype: Element data type (``np.uint8``, ``np.int16``, or ``np.int32``).

    Returns:
        ExternalFunction configured for ``passThroughLine``.

    Raises:
        ValueError: When ``dtype`` is not ``np.uint8``, ``np.int16``, or
            ``np.int32``, or when ``tile_size`` is not a whole number of
            64-byte vectors, or is nonpositive.
    """
    bit_width = _dtype_to_bit_width(dtype, factory_name="passthrough")
    _require_vector_alignment(
        "passthrough",
        tile_size,
        _PASSTHROUGH_VEC_BYTES // np.dtype(dtype).itemsize,
    )
    tile_ty = np.ndarray[(tile_size,), np.dtype[dtype]]
    return _make_extern(
        "passThroughLine",
        _kernel_source("eltwise/passThrough.cc"),
        [tile_ty, tile_ty, np.int32],
        compile_flags=[
            f"-DBIT_WIDTH={bit_width}",
            f"-DPASSTHROUGH_ELEMS={tile_size}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            reference=lambda x: x,
            tolerance=Tolerance.exact(note="lossless copy"),
        ),
    )

scale

scale(
    tile_size: int = 1024,
    dtype: type = int32,
    vectorized: bool = True,
    use_chess: bool = False,
) -> ExternalFunction

Scalar-multiply kernel: multiplies each element of an input tile by a factor.

Parameters:

Name	Type	Description	Default
tile_size	int	Number of elements per tile.	1024
dtype	type	Element data type. Must be np.int16 or np.int32.	int32
vectorized	bool	If True use the vectorized path; False selects scalar.	True
use_chess	bool	When True, build the .o with xchesscc_wrapper instead of Peano.	False

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the scale kernel.

Raises:

Type	Description
ValueError	When dtype is not np.int16 or np.int32.

Source code in python/iron/kernels/eltwise.py

@dtypes(({"dtype": np.int16}, {"dtype": np.int32}))
def scale(
    tile_size: int = 1024,
    dtype: type = np.int32,
    vectorized: bool = True,
    use_chess: bool = False,
) -> ExternalFunction:
    """Scalar-multiply kernel: multiplies each element of an input tile by a factor.

    Args:
        tile_size: Number of elements per tile.
        dtype: Element data type. Must be ``np.int16`` or ``np.int32``.
        vectorized: If ``True`` use the vectorized path; ``False`` selects scalar.
        use_chess: When ``True``, build the .o with ``xchesscc_wrapper``
            instead of Peano.

    Returns:
        ExternalFunction configured for the scale kernel.

    Raises:
        ValueError: When ``dtype`` is not ``np.int16`` or ``np.int32``.
    """
    if dtype not in (np.int16, np.int32):
        raise ValueError(f"scale() dtype must be np.int16 or np.int32, got {dtype}")
    if vectorized:
        _require_vector_alignment("scale", tile_size, 64 // np.dtype(dtype).itemsize)

    tile_ty = np.ndarray[(tile_size,), np.dtype[dtype]]
    scalar_ty = np.ndarray[(1,), np.dtype[np.int32]]
    func_variant = "vector" if vectorized else "scalar"
    bit_width = 16 if dtype == np.int16 else 32
    return _make_extern(
        f"vector_scalar_mul_{func_variant}",
        _kernel_source("eltwise/scale.cc"),
        [tile_ty, tile_ty, scalar_ty, np.int32],
        compile_flags=[f"-DBIT_WIDTH={bit_width}", f"-DSCALE_ELEMS={tile_size}"],
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param, Param),
            parameter_bindings=((3, tile_size),),
            reference=scale_ref,
            acc_dtype=np.int32 if dtype == np.int16 else np.int64,  # acc32 / acc64
            reduction=1,
            tolerance=Tolerance.exact(
                note="integer multiply; overflow wraps like the C++ store"
            ),
        ),
    )

add

add(
    tile_size: int = 1024,
    dtype: type = bfloat16,
    vectorized: bool = True,
) -> ExternalFunction

Element-wise bf16 addition (tile_size must be 1024, hard-coded in C++).

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per tile (must be 1024).	1024
dtype	type	Element data type (only bfloat16 supported).	bfloat16
vectorized	bool	If True use vectorized path; False selects scalar.	True

Returns:

Type	Description
ExternalFunction	ExternalFunction for eltwise_add_bf16.

Raises:

Type	Description
ValueError	When dtype is not bfloat16.

Source code in python/iron/kernels/eltwise.py

def add(
    tile_size: int = 1024, dtype: type = bfloat16, vectorized: bool = True
) -> ExternalFunction:
    """Element-wise bf16 addition (tile_size must be 1024, hard-coded in C++).

    Args:
        tile_size: Elements per tile (must be 1024).
        dtype: Element data type (only ``bfloat16`` supported).
        vectorized: If ``True`` use vectorized path; ``False`` selects scalar.

    Returns:
        ExternalFunction for eltwise_add_bf16.

    Raises:
        ValueError: When ``dtype`` is not ``bfloat16``.
    """
    return _eltwise_bf16_kernel("add", tile_size, dtype, vectorized)

mul

mul(
    tile_size: int = 1024,
    dtype: type = bfloat16,
    vectorized: bool = True,
) -> ExternalFunction

Element-wise bf16 multiplication (tile_size must be 1024, hard-coded in C++).

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per tile (must be 1024).	1024
dtype	type	Element data type (only bfloat16 supported).	bfloat16
vectorized	bool	If True use vectorized path; False selects scalar.	True

Returns:

Type	Description
ExternalFunction	ExternalFunction for eltwise_mul_bf16.

Raises:

Type	Description
ValueError	When dtype is not bfloat16.

Source code in python/iron/kernels/eltwise.py

def mul(
    tile_size: int = 1024, dtype: type = bfloat16, vectorized: bool = True
) -> ExternalFunction:
    """Element-wise bf16 multiplication (tile_size must be 1024, hard-coded in C++).

    Args:
        tile_size: Elements per tile (must be 1024).
        dtype: Element data type (only ``bfloat16`` supported).
        vectorized: If ``True`` use vectorized path; ``False`` selects scalar.

    Returns:
        ExternalFunction for eltwise_mul_bf16.

    Raises:
        ValueError: When ``dtype`` is not ``bfloat16``.
    """
    return _eltwise_bf16_kernel("mul", tile_size, dtype, vectorized)

mul_add

mul_add(tile_size: int = 1024) -> ExternalFunction

c = a * b or c = a + b on bf16 tiles, chosen per call by is_mul.

One kernel for a two-phase runtime-parameter design: programming_examples/ml/scale_shift computes A * B with is_mul = 1 and then + C with is_mul = 0 on the same workers. aie_kernels/eltwise/scale_shift.cc fixes the tile at 1024 elements.

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per tile (must be 1024).	1024

Raises:

Type	Description
ValueError	When tile_size is not 1024.

Source code in python/iron/kernels/eltwise.py

def mul_add(tile_size: int = 1024) -> ExternalFunction:
    """``c = a * b`` or ``c = a + b`` on bf16 tiles, chosen per call by ``is_mul``.

    One kernel for a two-phase runtime-parameter design:
    programming_examples/ml/scale_shift computes ``A * B`` with ``is_mul = 1``
    and then ``+ C`` with ``is_mul = 0`` on the same workers.
    ``aie_kernels/eltwise/scale_shift.cc`` fixes the tile at 1024 elements.

    Args:
        tile_size: Elements per tile (must be 1024).

    Raises:
        ValueError: When ``tile_size`` is not 1024.
    """
    _require_fixed_tile_size("mul_add", tile_size, 1024)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _make_extern(
        "eltwise_mul_add_bf16_vector",
        _kernel_source("eltwise/scale_shift.cc"),
        [tile_ty, tile_ty, tile_ty, np.int32],
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, In, Out, Param),
            reference=mul_add_ref,
            acc_dtype=np.float32,
            tolerance=_BF16_ROUNDTRIP,
            ops_per_call=tile_size,
        ),
    )

mul_add_ref

mul_add_ref(a, b, is_mul)

Numpy reference for mul_add: a * b if is_mul else a + b.

Source code in python/iron/kernels/eltwise.py

def mul_add_ref(a, b, is_mul):
    """Numpy reference for [`mul_add`][iron.kernels.eltwise.mul_add]: ``a * b`` if ``is_mul`` else ``a + b``."""
    a32, b32 = a.astype(np.float32), b.astype(np.float32)
    return (a32 * b32 if int(is_mul) else a32 + b32).astype(a.dtype)

relu

relu(tile_size: int = 1024) -> ExternalFunction

Element-wise bf16 ReLU (tile_size must be 1024, hard-coded in C++).

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per tile (must be 1024).	1024

Returns:

Type	Description
ExternalFunction	ExternalFunction for bf16_relu.

Raises:

Type	Description
ValueError	When tile_size is not 1024.

Source code in python/iron/kernels/eltwise.py

def relu(tile_size: int = 1024) -> ExternalFunction:
    """Element-wise bf16 ReLU (tile_size must be 1024, hard-coded in C++).

    Args:
        tile_size: Elements per tile (must be 1024).

    Returns:
        ExternalFunction for bf16_relu.

    Raises:
        ValueError: When ``tile_size`` is not 1024.
    """
    _require_fixed_tile_size("relu", tile_size, _RELU_FIXED_TILE)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _make_extern(
        "bf16_relu",
        _kernel_source("eltwise/relu.cc"),
        [tile_ty, tile_ty],
        compile_flags=[f"-DRELU_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=lambda x: np.maximum(x.astype(np.float32), 0.0),
            tolerance=Tolerance.exact(note="selection: max(x, 0) is exact in bf16"),
        ),
    )

add_sized

add_sized(tile_size: int = 1024) -> ExternalFunction

Element-wise bf16 addition, with a compiled-in element count.

Runtime-size sibling of add; design passes (a, b, c, size) for ABI compatibility. Scalar tails are supported.

Source code in python/iron/kernels/eltwise.py

def add_sized(tile_size: int = 1024) -> ExternalFunction:
    """Element-wise bf16 addition, with a compiled-in element count.

    Runtime-size sibling of [`add`][iron.kernels.eltwise.add]; design passes
    ``(a, b, c, size)`` for ABI compatibility. Scalar tails are supported.
    """
    if tile_size <= 0:
        raise ValueError(f"add_sized() tile_size must be positive, got {tile_size}")
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _make_extern(
        "eltwise_add_bf16_vector_size",
        _kernel_source("eltwise/add.cc"),
        [tile_ty, tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DADD_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, In, Out, Param),
            parameter_bindings=((3, tile_size),),
            reference=add_ref,
            acc_dtype=np.float32,
            tolerance=_BF16_ROUNDTRIP,
            ops_per_call=tile_size,
        ),
    )

mul_sized

mul_sized(tile_size: int = 1024) -> ExternalFunction

Element-wise bf16 multiplication, with a compiled-in element count.

Runtime-size sibling of mul; design passes (a, b, c, size) for ABI compatibility. Scalar tails are supported.

Source code in python/iron/kernels/eltwise.py

def mul_sized(tile_size: int = 1024) -> ExternalFunction:
    """Element-wise bf16 multiplication, with a compiled-in element count.

    Runtime-size sibling of [`mul`][iron.kernels.eltwise.mul]; design passes
    ``(a, b, c, size)`` for ABI compatibility. Scalar tails are supported.
    """
    if tile_size <= 0:
        raise ValueError(f"mul_sized() tile_size must be positive, got {tile_size}")
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _make_extern(
        "eltwise_mul_bf16_vector_size",
        _kernel_source("eltwise/mul.cc"),
        [tile_ty, tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DMUL_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, In, Out, Param),
            parameter_bindings=((3, tile_size),),
            reference=mul_ref,
            acc_dtype=np.float32,
            tolerance=_BF16_ROUNDTRIP,
            ops_per_call=tile_size,
        ),
    )

relu_sized

relu_sized(tile_size: int = 1024) -> ExternalFunction

Element-wise bf16 ReLU, with a compiled-in element count.

Runtime-size sibling of relu; design passes (in, out, size) for ABI compatibility. Not LUT-based. Positive multiples of 32 elements are supported.

Source code in python/iron/kernels/eltwise.py

def relu_sized(tile_size: int = 1024) -> ExternalFunction:
    """Element-wise bf16 ReLU, with a compiled-in element count.

    Runtime-size sibling of [`relu`][iron.kernels.eltwise.relu]; design passes
    ``(in, out, size)`` for ABI compatibility. Not LUT-based. Positive
    multiples of 32 elements are supported.
    """
    _require_vector_alignment("relu_sized", tile_size, 32)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _make_extern(
        "relu_bf16_size",
        _kernel_source("eltwise/relu.cc"),
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DRELU_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            reference=lambda x: np.maximum(x.astype(np.float32), 0.0),
            tolerance=Tolerance.exact(note="selection: max(x, 0) is exact in bf16"),
            ops_per_call=tile_size,
        ),
    )

Data movement

Data-movement / conversion kernel factories: affine_cast, axpy, convert_copy, expand, transpose.

Each wraps one source under aie_kernels/datamovement/ — plain aie_api vector code with no LUT dependency. convert_copy binds cast_f32_bf16.cc, the f32->bf16 cast with host-matching conv_even rounding, and affine_cast applies a per-column gamma/beta ahead of the same cast.

axpy_ref

axpy_ref(x, y, a)

Numpy reference for axpy: a * x + y in float32.

The kernel broadcasts a as bf16, so it is rounded to bf16 here.

Source code in python/iron/kernels/datamovement.py

def axpy_ref(x, y, a):
    """Numpy reference for [`axpy`][iron.kernels.datamovement.axpy]: ``a * x + y`` in float32.

    The kernel broadcasts ``a`` as bf16, so it is rounded to bf16 here.
    """
    a = np.float32(bfloat16(a))
    return a * x.astype(np.float32) + y.astype(np.float32)

convert_copy_ref

convert_copy_ref(x)

Numpy reference for convert_copy.

ml_dtypes rounds f32 -> bf16 half-to-even, exactly as the kernel's conv_even does, so the cast is the reference and the match is bit-for-bit.

Source code in python/iron/kernels/datamovement.py

def convert_copy_ref(x):
    """Numpy reference for [`convert_copy`][iron.kernels.datamovement.convert_copy].

    ``ml_dtypes`` rounds f32 -> bf16 half-to-even, exactly as the kernel's
    ``conv_even`` does, so the cast is the reference and the match is
    bit-for-bit.
    """
    return x.astype(bfloat16)

affine_cast_ref

affine_cast_ref(x, gamma_beta)

Numpy reference for affine_cast.

gamma_beta is gamma then beta, cols float32 values each; x is row-major with cols per row. The multiply and add round in float32 and the bf16 cast half-to-even, as the kernel does.

Source code in python/iron/kernels/datamovement.py

def affine_cast_ref(x, gamma_beta):
    """Numpy reference for [`affine_cast`][iron.kernels.datamovement.affine_cast].

    ``gamma_beta`` is ``gamma`` then ``beta``, ``cols`` float32 values each;
    ``x`` is row-major with ``cols`` per row. The multiply and add round in
    float32 and the bf16 cast half-to-even, as the kernel does.
    """
    gb = np.asarray(gamma_beta, dtype=np.float32).reshape(-1)
    cols = gb.size // 2
    x32 = np.asarray(x, dtype=np.float32)
    y = x32.reshape(-1, cols) * gb[:cols] + gb[cols:]
    return y.astype(bfloat16).reshape(x32.shape)

expand_ref

expand_ref(payload, *, tile_size: int, group_size: int)

Numpy reference for expand.

payload holds, per tile, tile_size packed uint4 values (tile_size // 2 bytes, low nibble first) followed by one bf16 scale per group_size elements; the result is nibble * scale-of-its-group.

Source code in python/iron/kernels/datamovement.py

def expand_ref(payload, *, tile_size: int, group_size: int):
    """Numpy reference for [`expand`][iron.kernels.datamovement.expand].

    ``payload`` holds, per tile, ``tile_size`` packed uint4 values
    (``tile_size // 2`` bytes, low nibble first) followed by one bf16 scale per
    ``group_size`` elements; the result is ``nibble * scale-of-its-group``.
    """
    payload = np.asarray(payload, dtype=np.uint8)
    payload = payload.reshape(-1, payload.shape[-1])
    n_scales = tile_size // group_size
    packed = payload[:, : tile_size // 2]
    scales = np.ascontiguousarray(payload[:, tile_size // 2 :]).view(bfloat16)
    scales = scales.reshape(-1, n_scales).astype(np.float32)
    nibbles = np.empty((payload.shape[0], tile_size), np.float32)
    nibbles[:, 0::2] = packed & 0x0F
    nibbles[:, 1::2] = packed >> 4
    return nibbles * np.repeat(scales, group_size, axis=1)

expand_sample

expand_sample(
    rng, calls: int, *, tile_size: int, group_size: int
) -> list

Random expand payloads: packed uint4 values then bf16 scales in [0.1, 1).

Source code in python/iron/kernels/datamovement.py

def expand_sample(rng, calls: int, *, tile_size: int, group_size: int) -> list:
    """Random ``expand`` payloads: packed uint4 values then bf16 scales in [0.1, 1)."""
    n_scales = tile_size // group_size
    nibbles = rng.integers(0, 16, size=(calls, tile_size), dtype=np.uint8)
    packed = (nibbles[:, 0::2] | (nibbles[:, 1::2] << 4)).astype(np.uint8)
    scales = rng.uniform(0.1, 1.0, size=(calls, n_scales)).astype(bfloat16)
    return [np.concatenate([packed, scales.view(np.uint8)], axis=1)]

transpose_ref

transpose_ref(x, *, dim_m: int, dim_n: int, subtile: int)

Numpy reference for transpose.

Transposes each subtile x subtile block of the dim_n x dim_m matrix in place -- the blocks move, the matrix does not.

Source code in python/iron/kernels/datamovement.py

def transpose_ref(x, *, dim_m: int, dim_n: int, subtile: int):
    """Numpy reference for [`transpose`][iron.kernels.datamovement.transpose].

    Transposes each ``subtile`` x ``subtile`` block of the ``dim_n`` x ``dim_m``
    matrix in place -- the blocks move, the matrix does not.
    """
    x = np.asarray(x)
    mats = x.reshape(-1, dim_n, dim_m)
    out = mats.copy()
    for r in range(0, dim_n, subtile):
        for c in range(0, dim_m, subtile):
            out[:, r : r + subtile, c : c + subtile] = np.swapaxes(
                mats[:, r : r + subtile, c : c + subtile], 1, 2
            )
    return out.reshape(x.shape)

axpy

axpy(
    tile_size: int = 1024, vectorized: bool = True
) -> ExternalFunction

SAXPY kernel: z = a * x + y over bf16 tiles.

The scalar a and element count are passed to the kernel at runtime, so a design supplies (x, y, a, z, size). The vectorized path processes 64 elements per iteration; tile_size must therefore be a multiple of 64.

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per tile (multiple of 64 for the vectorized path).	1024
vectorized	bool	If True bind saxpy; False binds saxpy_scalar.	True

Returns:

Type	Description
ExternalFunction	ExternalFunction for the saxpy kernel.

Raises:

Type	Description
ValueError	When vectorized and tile_size is not a multiple of 64.

Source code in python/iron/kernels/datamovement.py

def axpy(tile_size: int = 1024, vectorized: bool = True) -> ExternalFunction:
    """SAXPY kernel: ``z = a * x + y`` over bf16 tiles.

    The scalar ``a`` and element count are passed to the kernel at runtime, so a
    design supplies ``(x, y, a, z, size)``.  The vectorized path processes 64
    elements per iteration; ``tile_size`` must therefore be a multiple of 64.

    Args:
        tile_size: Elements per tile (multiple of 64 for the vectorized path).
        vectorized: If ``True`` bind ``saxpy``; ``False`` binds ``saxpy_scalar``.

    Returns:
        ExternalFunction for the saxpy kernel.

    Raises:
        ValueError: When ``vectorized`` and ``tile_size`` is not a multiple of 64.
    """
    if vectorized and tile_size % _AXPY_VEC != 0:
        raise ValueError(
            f"axpy() vectorized tile_size must be a multiple of {_AXPY_VEC}, "
            f"got {tile_size}."
        )
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    # saxpy takes float a; saxpy_scalar takes bfloat16 a.
    a_ty = np.float32 if vectorized else bfloat16
    func = "saxpy" if vectorized else "saxpy_scalar"
    return _make_extern(
        func,
        _kernel_source("datamovement/axpy.cc"),
        [tile_ty, tile_ty, a_ty, tile_ty, np.int32],
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, In, Param, Out, Param),
            parameter_bindings=((4, tile_size),),
            reference=axpy_ref,
            acc_dtype=np.float32,
            reduction=1,
            tolerance=_BF16_ROUNDTRIP,
            ops_per_call=2 * tile_size,
        ),
    )

convert_copy

convert_copy(tile_size: int = 1024) -> ExternalFunction

Convert-copy kernel: element-preserving float32 -> bfloat16.

Reads a length-tile_size f32 tile and writes the same number of bf16 elements (halving the byte footprint). Element count is a runtime arg; the kernel processes 16 elements per iteration, so tile_size must be a multiple of 16.

Backed by aie_kernels/datamovement/cast_f32_bf16.cc (symbol cast_f32_bf16_row), which rounds with conv_even — bit-for-bit agreeing with a host AVX512-BF16 pack — and restores the core's rounding mode on exit. The same source builds for aie2.

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per tile (multiple of 16).	1024

Returns:

Type	Description
ExternalFunction	ExternalFunction for cast_f32_bf16_row.

Raises:

Type	Description
ValueError	When tile_size is not a multiple of 16.

Source code in python/iron/kernels/datamovement.py

def convert_copy(tile_size: int = 1024) -> ExternalFunction:
    """Convert-copy kernel: element-preserving ``float32`` -> ``bfloat16``.

    Reads a length-``tile_size`` f32 tile and writes the same number of bf16
    elements (halving the byte footprint).  Element count is a runtime arg; the
    kernel processes 16 elements per iteration, so ``tile_size`` must be a
    multiple of 16.

    Backed by ``aie_kernels/datamovement/cast_f32_bf16.cc`` (symbol
    ``cast_f32_bf16_row``), which rounds with ``conv_even`` — bit-for-bit
    agreeing with a host AVX512-BF16 pack — and restores the core's rounding
    mode on exit.  The same source builds for aie2.

    Args:
        tile_size: Elements per tile (multiple of 16).

    Returns:
        ExternalFunction for ``cast_f32_bf16_row``.

    Raises:
        ValueError: When ``tile_size`` is not a multiple of 16.
    """
    if tile_size % 16 != 0:
        raise ValueError(
            f"convert_copy() tile_size must be a multiple of 16, got {tile_size}."
        )
    in_ty = np.ndarray[(tile_size,), np.dtype[np.float32]]
    out_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _make_extern(
        "cast_f32_bf16_row",
        _kernel_source("datamovement/cast_f32_bf16.cc"),
        [in_ty, out_ty, np.int32],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            reference=convert_copy_ref,
            tolerance=Tolerance.exact(
                note="conv_even rounding matches ml_dtypes bit-for-bit (test_kernels_e2e)"
            ),
        ),
    )

affine_cast

affine_cast(
    rows: int = 96, cols: int = 32
) -> ExternalFunction

Per-column affine transform narrowed to bf16: out = bfloat16(in * gamma + beta).

Works on a row-major rows x cols float32 tile; gamma and beta hold one float32 value per column and arrive packed in one 2 * cols buffer, gamma first. The bf16 store rounds with conv_even, as convert_copy does, and the kernel restores the core's rounding mode on exit.

Parameters:

Name	Type	Description	Default
rows	int	Rows per tile.	96
cols	int	Columns per tile (multiple of 16).	32

Returns:

Type	Description
ExternalFunction	ExternalFunction for affine_cast_f32_bf16.

Raises:

Type	Description
ValueError	When rows is below 1 or cols is not a positive multiple of 16.

Source code in python/iron/kernels/datamovement.py

def affine_cast(rows: int = 96, cols: int = 32) -> ExternalFunction:
    """Per-column affine transform narrowed to bf16: ``out = bfloat16(in * gamma + beta)``.

    Works on a row-major ``rows`` x ``cols`` float32 tile; ``gamma`` and
    ``beta`` hold one float32 value per column and arrive packed in one
    ``2 * cols`` buffer, ``gamma`` first. The bf16 store rounds with
    ``conv_even``, as [`convert_copy`][iron.kernels.datamovement.convert_copy]
    does, and the kernel restores the core's rounding mode on exit.

    Args:
        rows: Rows per tile.
        cols: Columns per tile (multiple of 16).

    Returns:
        ExternalFunction for ``affine_cast_f32_bf16``.

    Raises:
        ValueError: When ``rows`` is below 1 or ``cols`` is not a positive
            multiple of 16.
    """
    if rows < 1 or cols < 16 or cols % 16 != 0:
        raise ValueError(
            "affine_cast() needs rows >= 1 and cols a positive multiple of 16, "
            f"got rows={rows}, cols={cols}."
        )
    in_ty = np.ndarray[(rows * cols,), np.dtype[np.float32]]
    gb_ty = np.ndarray[(2 * cols,), np.dtype[np.float32]]
    out_ty = np.ndarray[(rows * cols,), np.dtype[bfloat16]]
    return _make_extern(
        "affine_cast_f32_bf16",
        _kernel_source("datamovement/affine_cast_f32_bf16.cc"),
        [in_ty, gb_ty, out_ty, np.int32, np.int32],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Param, Out, Param, Param),
            parameter_bindings=((3, rows), (4, cols)),
            reference=affine_cast_ref,
            acc_dtype=np.float32,
            reduction=1,
            tolerance=Tolerance.bf16_ulps(
                1,
                atol=2.0**-126,
                note="aie::mul emulates the float32 product in bf16 terms and "
                "can land one float32 ulp off, which tips a bf16 tie: 3 of "
                "1769472 outputs one ulp off on npu2. atol is the smallest "
                "normal bf16, for the device's subnormal flush to zero",
            ),
            ops_per_call=2 * rows * cols,
        ),
    )

expand

expand(
    tile_size: int = 1024, group_size: int = 32
) -> ExternalFunction

Dequantize kernel: uint4 -> bfloat16 with per-group scale factors.

Each tile holds tile_size packed unsigned int4 values followed by one bf16 scale factor per group_size-element group; the kernel zero-extends and scales into tile_size bf16 outputs (no zero point). tile_size and group_size are baked in at compile time via -DTILE_SIZE / -DGROUP_SIZE (group_size must be a multiple of 32, matching the C++ static_assert).

Parameters:

Name	Type	Description	Default
tile_size	int	Number of uint4 elements per tile.	1024
group_size	int	Elements sharing one scale factor (multiple of 32).	32

Returns:

Type	Description
ExternalFunction	ExternalFunction for expand_uint4_to_bfloat16.

Raises:

Type	Description
ValueError	When group_size is not a multiple of 32.

Source code in python/iron/kernels/datamovement.py

def expand(tile_size: int = 1024, group_size: int = 32) -> ExternalFunction:
    """Dequantize kernel: ``uint4`` -> ``bfloat16`` with per-group scale factors.

    Each tile holds ``tile_size`` packed unsigned int4 values followed by one
    bf16 scale factor per ``group_size``-element group; the kernel zero-extends
    and scales into ``tile_size`` bf16 outputs (no zero point).  ``tile_size``
    and ``group_size`` are baked in at compile time via ``-DTILE_SIZE`` /
    ``-DGROUP_SIZE`` (group_size must be a multiple of 32, matching the C++
    ``static_assert``).

    Args:
        tile_size: Number of uint4 elements per tile.
        group_size: Elements sharing one scale factor (multiple of 32).

    Returns:
        ExternalFunction for ``expand_uint4_to_bfloat16``.

    Raises:
        ValueError: When ``group_size`` is not a multiple of 32.
    """
    if group_size % 32 != 0:
        raise ValueError(
            f"expand() group_size must be a multiple of 32, got {group_size}."
        )
    # Input tile layout the kernel expects: tile_size packed int4s
    # (= tile_size//2 bytes) IMMEDIATELY followed by one bf16 scale factor per
    # group (the kernel reads them from ``in + N/2``).  So the buffer is larger
    # than just the int4 payload; model the whole thing as raw uint8 or the
    # func.call operand type won't match the design's ObjectFifo.
    n_scales = tile_size // group_size
    in_bytes = tile_size // 2 + n_scales * 2  # int4 payload + bf16 scales
    in_ty = np.ndarray[(in_bytes,), np.dtype[np.uint8]]
    out_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _make_extern(
        "expand_uint4_to_bfloat16",
        _kernel_source("datamovement/expand.cc"),
        [in_ty, out_ty],
        compile_flags=[f"-DTILE_SIZE={tile_size}", f"-DGROUP_SIZE={group_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, Out),
            reference=lambda p: expand_ref(
                p, tile_size=tile_size, group_size=group_size
            ),
            tolerance=_BF16_ROUNDTRIP,
            ops_per_call=tile_size,
            sample=lambda rng, calls: expand_sample(
                rng, calls, tile_size=tile_size, group_size=group_size
            ),
        ),
    )

rope

rope(
    tile_size: int = 1024,
    two_halves: bool = False,
    *,
    cols: int | None = None
) -> ExternalFunction

RoPE positional rotation over bf16 tiles; dims read at runtime.

Design passes (in, lut, out, dims). two_halves selects the HuggingFace-style rope_two_halves over the Llama-paper interleave rope. cols aliases tile_size. Both architectures use the generic source; rows must be positive multiples of 16 (interleaved) or 32 (two halves, keeping each half 32-byte aligned). Each input row has its own streamed (cos, sin) LUT.

Source code in python/iron/kernels/datamovement.py

def rope(
    tile_size: int = 1024, two_halves: bool = False, *, cols: int | None = None
) -> ExternalFunction:
    """RoPE positional rotation over bf16 tiles; ``dims`` read at runtime.

    Design passes ``(in, lut, out, dims)``.  ``two_halves`` selects the
    HuggingFace-style ``rope_two_halves`` over the Llama-paper interleave
    ``rope``. ``cols`` aliases ``tile_size``. Both architectures use the generic
    source; rows must be positive multiples of 16 (interleaved) or 32
    (two halves, keeping each half 32-byte aligned). Each input row has its own
    streamed (cos, sin) LUT.
    """
    tile_size = _row_size("rope", tile_size, cols, 32 if two_halves else 16)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    func = "rope_two_halves" if two_halves else "rope"
    return _make_extern(
        func,
        _kernel_source("datamovement/rope.cc"),
        [tile_ty, tile_ty, tile_ty, np.int32],
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, In, Out, Param),
            parameter_bindings=((3, tile_size),),
            reference=lambda x, lut: rope_ref(x, lut, two_halves=two_halves),
            acc_dtype=np.float32,
            reduction=2,
            tolerance=Tolerance.relative(
                0.128, note="programming_examples/ml/rope: default bf16 rtol"
            ),
            ops_per_call=3 * tile_size,
        ),
    )

rope_ref

rope_ref(x, lut, *, two_halves: bool = False)

Rotate bf16 pairs by an interleaved (cos, sin) LUT, in either RoPE layout.

Source code in python/iron/kernels/datamovement.py

def rope_ref(x, lut, *, two_halves: bool = False):
    """Rotate bf16 pairs by an interleaved (cos, sin) LUT, in either RoPE layout."""
    x32, l32 = x.astype(np.float32), lut.astype(np.float32)
    cos_v, sin_v = l32[..., 0::2], l32[..., 1::2]
    if two_halves:
        half = x32.shape[-1] // 2
        x1, x2 = x32[..., :half], x32[..., half:]
        return np.concatenate(
            (x1 * cos_v - x2 * sin_v, x2 * cos_v + x1 * sin_v), axis=-1
        ).astype(x.dtype)
    x_even, x_odd = x32[..., 0::2], x32[..., 1::2]
    out = np.empty_like(x32)
    out[..., 0::2] = x_even * cos_v - x_odd * sin_v
    out[..., 1::2] = x_even * sin_v + x_odd * cos_v
    return out.astype(x.dtype)

transpose

transpose(
    dim_m: int = 32,
    dim_n: int = 32,
    subtile: int = 4,
    dtype: type = bfloat16,
) -> ExternalFunction

Blocked transpose through aie::transpose.

Transposes each subtile x subtile block of a dim_n x dim_m matrix in place: the blocks stay put, the elements inside them move. dim_m / dim_n are compile-time (-DDIM_m / -DDIM_n). The kernel only moves bytes, so any 1-, 2- or 4-byte dtype works and selects -DBIT_WIDTH as the other generic kernels do; bf16 is the default. programming_examples/basic/transposes uses this kernel for its combined strategy.

Parameters:

Name	Type	Description	Default
dim_m	int	Inner (contiguous) dimension.	32
dim_n	int	Outer dimension.	32
subtile	int	Block size to transpose, 4 (transpose_4x4) or 8 (transpose_8x8).	4
dtype	type	Element type, 1, 2 or 4 bytes wide.	bfloat16

Returns:

Type	Description
ExternalFunction	ExternalFunction for the selected transpose variant.

Raises:

Type	Description
ValueError	When subtile is not 4 or 8, when dtype is not 1, 2 or 4 bytes, or when the shape does not divide into the strips the kernel walks (dim_n a multiple of subtile; dim_m a multiple of the strip width, at least 16 bytes long).

Source code in python/iron/kernels/datamovement.py

@dtypes(
    (
        {"dtype": bfloat16},
        {"dtype": np.uint8},
        {"dtype": np.uint16},
        {"dtype": np.uint32},
    )
)
def transpose(
    dim_m: int = 32, dim_n: int = 32, subtile: int = 4, dtype: type = bfloat16
) -> ExternalFunction:
    """Blocked transpose through ``aie::transpose``.

    Transposes each ``subtile`` x ``subtile`` block of a ``dim_n`` x ``dim_m``
    matrix in place: the blocks stay put, the elements inside them move.
    ``dim_m`` / ``dim_n`` are compile-time (``-DDIM_m`` / ``-DDIM_n``). The
    kernel only moves bytes, so any 1-, 2- or 4-byte ``dtype`` works and
    selects ``-DBIT_WIDTH`` as the other generic kernels do; bf16 is the
    default. programming_examples/basic/transposes uses this kernel for its
    ``combined`` strategy.

    Args:
        dim_m: Inner (contiguous) dimension.
        dim_n: Outer dimension.
        subtile: Block size to transpose, 4 (``transpose_4x4``) or 8
            (``transpose_8x8``).
        dtype: Element type, 1, 2 or 4 bytes wide.

    Returns:
        ExternalFunction for the selected transpose variant.

    Raises:
        ValueError: When ``subtile`` is not 4 or 8, when ``dtype`` is not 1, 2
            or 4 bytes, or when the shape does not divide into the strips the
            kernel walks (``dim_n`` a multiple of ``subtile``; ``dim_m`` a
            multiple of the strip width, at least 16 bytes long).
    """
    if subtile not in (4, 8):
        raise ValueError(f"transpose() subtile must be 4 or 8, got {subtile}.")
    if dim_m <= 0 or dim_n <= 0:
        raise ValueError(
            f"transpose() dim_m and dim_n must be positive, got {dim_m}x{dim_n}."
        )
    width = np.dtype(dtype).itemsize
    if width not in (1, 2, 4):
        raise ValueError(
            f"transpose() dtype must be 1, 2 or 4 bytes wide, got {dtype}."
        )
    bits = 8 * width
    strip_w, _ = _transpose_strip(dim_m, subtile, bits)
    if dim_n % subtile or dim_m % strip_w or strip_w * width < 16:
        raise ValueError(
            f"transpose() {dim_m}x{dim_n} with {subtile}x{subtile} blocks of "
            f"{width}-byte elements: dim_n must be a multiple of {subtile} and "
            f"dim_m a multiple of {strip_w} (the kernel's strip width), with "
            "dim_m at least 16 bytes long."
        )
    flags = [f"-DDIM_m={dim_m}", f"-DDIM_n={dim_n}", f"-DBIT_WIDTH={bits}"]
    tile_ty = np.ndarray[(dim_m * dim_n,), np.dtype[dtype]]
    return _make_extern(
        f"transpose_{subtile}x{subtile}",
        _kernel_source("datamovement/transpose.cc"),
        [tile_ty, tile_ty],
        compile_flags=flags,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=lambda x: transpose_ref(
                x, dim_m=dim_m, dim_n=dim_n, subtile=subtile
            ),
            tolerance=Tolerance.exact(note="data movement only; lossless"),
            ops_per_call=0,
        ),
    )

Independent, target-native zero-fill kernel.

zero

zero(
    tile_size: int | tuple[int, ...] = 1024,
    dtype: type | dtype = int32,
    *,
    vectorized: bool = True,
    use_chess: bool = False
) -> ExternalFunction

Fill one tile with zeros, independently of any compute kernel.

tile_size is an element count or shape. For v8bfp16ebs8 it counts eight-value blocks, matching the ndarray ABI; all nine bytes of every block (exponent and mantissas) are cleared. Vector stores use the target's native width, with a scalar tail for smaller tiles.

Source code in python/iron/kernels/zero.py

@dtypes(tuple({"dtype": dtype} for dtype in _TYPES) + ({"dtype": v8bfp16ebs8},))
def zero(
    tile_size: int | tuple[int, ...] = 1024,
    dtype: type | np.dtype = np.int32,
    *,
    vectorized: bool = True,
    use_chess: bool = False,
) -> ExternalFunction:
    """Fill one tile with zeros, independently of any compute kernel.

    ``tile_size`` is an element count or shape. For ``v8bfp16ebs8`` it
    counts eight-value blocks, matching the ndarray ABI; all nine bytes
    of every block (exponent and mantissas) are cleared. Vector stores
    use the target's native width, with a scalar tail for smaller tiles.
    """
    try:
        shape = (
            (tile_size,)
            if isinstance(tile_size, (int, np.integer))
            else tuple(tile_size)
        )
        shape = tuple(operator.index(n) for n in shape)
    except TypeError as exc:
        raise ValueError("zero: tile_size must be a positive integer or shape") from exc
    if not shape or any(n <= 0 for n in shape):
        raise ValueError("zero: tile_size must be a positive integer or shape")
    size = math.prod(shape)
    block = dtype is v8bfp16ebs8
    if block:
        if not _arch_traits().bfp16:
            raise NotImplementedError("zero: bfp16ebs8 requires an NPU2 device")
        from aie.utils import bfp

        layout = TensorLayout(
            (size * bfp.BLOCK,),
            pack=lambda x: bfp.encode(x).reshape(len(x), -1),
            unpack=lambda x: bfp.decode(x).reshape(len(x), -1),
        )
        reference_dtype = np.float32
        ctype, count = "uint8_t", size * bfp.BLOCK_BYTES
    else:
        dtype = np.dtype(dtype).type
        if dtype not in _TYPES:
            raise ValueError(f"zero: unsupported dtype {dtype}")
        layout = TensorLayout(shape)
        reference_dtype = dtype
        ctype, count = _TYPES[dtype], size
    flags = [f"-DZERO_TYPE={ctype}", f"-DTILE_SIZE={count}"]
    if not vectorized:
        flags.append("-DZERO_SCALAR")
    return _make_extern(
        "zero",
        _kernel_source("zero/zero.cc"),
        [np.ndarray[shape, np.dtype[dtype]]],
        compile_flags=flags,
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(Out,),
            layouts=(layout,),
            reference=lambda: np.zeros((1, *layout.shape), dtype=reference_dtype),
            tolerance=Tolerance.exact(note="zero fill"),
            ops_per_call=0,
        ),
    )

Quantization

Packed quantization kernel factories and byte-exact host references.

q4nx_dequant_ref

q4nx_dequant_ref(
    payload,
    *,
    m_tile=32,
    k_tile=256,
    group=32,
    ct_k=128,
    s=8,
    t=8
)

Dequantize packed q4nx to the kernel's bfp16ebs8 output bytes.

Input has shape (..., m_tile*k_tile//2 + 4*m_tile*k_tile//group). It contains little-endian bf16 scales, then bf16 minima, both indexed [k_group, n], followed by unsigned nibbles (low nibble first) indexed [n//16, k, n%16]. Values are min + scale * nibble.

Accumulator results narrow to bf16 with floor rounding before BFP conversion. Output is uint8 with the same leading dimensions, indexed [k//ct_k, n//8, (k%ct_k)//8, n%8, 9]: each nine-byte block is an exponent followed by eight signed mantissas for consecutive k values. Finite scales/minima and finite dequantized bf16 results are required.

Source code in python/iron/kernels/quant.py

def q4nx_dequant_ref(payload, *, m_tile=32, k_tile=256, group=32, ct_k=128, s=8, t=8):
    """Dequantize packed q4nx to the kernel's bfp16ebs8 output bytes.

    Input has shape ``(..., m_tile*k_tile//2 + 4*m_tile*k_tile//group)``.
    It contains little-endian bf16 scales, then bf16 minima, both indexed
    ``[k_group, n]``, followed by unsigned nibbles (low nibble first) indexed
    ``[n//16, k, n%16]``. Values are ``min + scale * nibble``.

    Accumulator results narrow to bf16 with floor rounding before BFP
    conversion. Output is uint8 with the same leading dimensions, indexed
    ``[k//ct_k, n//8, (k%ct_k)//8, n%8, 9]``: each nine-byte block is an
    exponent followed by eight signed mantissas for consecutive k values.
    Finite scales/minima and finite dequantized bf16 results are required.
    """
    geometry = _geometry(m_tile, k_tile, group, ct_k, s, t)
    m_tile, k_tile, group, ct_k, s, t = geometry.values()
    scales_count = m_tile * k_tile // group
    input_bytes = m_tile * k_tile // 2 + 4 * scales_count
    payload = np.asarray(payload, dtype=np.uint8)
    if payload.ndim == 0 or payload.shape[-1] != input_bytes:
        raise ValueError(f"q4nx_dequant_ref: expected {input_bytes} bytes per tile")
    lead = payload.shape[:-1]
    data = payload.reshape(-1, input_bytes)
    if not len(data):
        return np.empty((*lead, m_tile * k_tile * 9 // 8), dtype=np.uint8)
    params = np.ascontiguousarray(data[:, : 4 * scales_count]).view("<u2")
    params = (params.astype(np.uint32) << 16).view(np.float32)
    if not np.isfinite(params).all():
        raise ValueError("q4nx_dequant_ref: scales and minima must be finite")
    scales, mins = params.reshape(-1, 2, k_tile // group, m_tile).transpose(1, 0, 2, 3)
    packed = data[:, 4 * scales_count :]
    q = np.empty((len(data), m_tile * k_tile), dtype=np.uint8)
    q[:, 0::2], q[:, 1::2] = packed & 15, packed >> 4
    q = q.reshape(-1, m_tile // 16, k_tile, 16).transpose(0, 2, 1, 3)
    q = q.reshape(-1, k_tile, m_tile)
    # A bf16 scale times a four-bit integer is exact in float32. Float64
    # models the fused addition before its single float32 accumulator rounding.
    with np.errstate(over="ignore", invalid="ignore"):
        values = (
            np.repeat(scales, group, axis=1).astype(np.float64) * q
            + np.repeat(mins, group, axis=1).astype(np.float64)
        ).astype(np.float32)
        values = _bf16_floor(values)
    if not np.isfinite(values).all():
        raise ValueError("q4nx_dequant_ref: dequantized bf16 values must be finite")
    values = values.reshape(-1, k_tile // ct_k, ct_k // s, s, m_tile // t, t)
    values = values.transpose(0, 1, 4, 2, 5, 3)
    encoded = bfp.encode(values.reshape(len(data), -1))
    return encoded.reshape(*lead, m_tile * k_tile * 9 // 8)

q4nx_dequant

q4nx_dequant(
    *, m_tile=32, k_tile=256, group=32, ct_k=128, s=8, t=8
) -> ExternalFunction

AIE2P-only q4nx dequantization into GEMM B-operand BFP storage.

m_tile counts n rows (a multiple of 16); k_tile must be divisible by both group and ct_k. The latter two are positive multiples of eight; s and t must be eight. Groups need not divide k slices. See q4nx_dequant_ref for the packed input and output layouts.

Both arguments are uint8 byte buffers, including the bfp16ebs8 output, so the generic harness compares the complete encoded result byte for byte rather than decoding or quantizing it again. The kernel saves, selects and restores floor rounding itself; no setup kernel is needed.

Source code in python/iron/kernels/quant.py

def q4nx_dequant(
    *, m_tile=32, k_tile=256, group=32, ct_k=128, s=8, t=8
) -> ExternalFunction:
    """AIE2P-only q4nx dequantization into GEMM B-operand BFP storage.

    ``m_tile`` counts n rows (a multiple of 16); ``k_tile`` must be divisible
    by both ``group`` and ``ct_k``. The latter two are positive multiples of
    eight; ``s`` and ``t`` must be eight. Groups need not divide k slices.
    See ``q4nx_dequant_ref`` for the packed input and output layouts.

    Both arguments are uint8 byte buffers, including the bfp16ebs8 output,
    so the generic harness compares the complete encoded result byte for
    byte rather than decoding or quantizing it again. The kernel saves,
    selects and restores floor rounding itself; no setup kernel is needed.
    """
    geometry = _geometry(m_tile, k_tile, group, ct_k, s, t)
    m_tile, k_tile, group, ct_k, s, t = geometry.values()
    if not _arch_traits().bfp16:
        raise NotImplementedError("q4nx_dequant() is only available on aie2p.")
    input_bytes = m_tile * k_tile // 2 + 4 * m_tile * k_tile // group
    output_bytes = m_tile * k_tile * 9 // 8
    return _make_extern(
        "q4nx_dequant_bfp",
        _kernel_source("quant/q4nx_dequant.cc"),
        [
            np.ndarray[(input_bytes,), np.dtype[np.uint8]],
            np.ndarray[(output_bytes,), np.dtype[np.uint8]],
        ],
        # The inner loop is one long latency chain, so it is allowed five
        # pipeline stages instead of the default three.
        compile_flags=[
            f"-DQ4NX_{name.upper()}={value}" for name, value in geometry.items()
        ]
        + ["-mllvm", "--aie-pipeliner-max-stagecount=5"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=partial(q4nx_dequant_ref, **geometry),
            sample=partial(_q4nx_sample, **geometry),
            tolerance=Tolerance.exact(
                note="floor bf16 narrowing and bfp16ebs8 conversion, byte for byte"
            ),
            ops_per_call=2 * m_tile * k_tile,
            acc_dtype=np.float32,
            reduction=1,
        ),
    )

Core state

set_rounding: the core's rounding-mode register.

The AIE core narrows accumulators (an srs shift, a bf16 store) in the rounding mode its mode register holds, and a fresh core boots in floor. A kernel that needs another mode names a setter as its contract's setup, and a design calls it once before the kernel's first call; the mode persists on that core until something changes it.

RoundingMode

Bases: str, Enum

An aie::rounding_mode, named as the C++ enumerator is.

set_rounding

set_rounding(
    mode: RoundingMode = CONV_EVEN,
) -> ExternalFunction

Set the core's rounding mode using always-inline, merge-linked LLVM IR.

Parameters:

Name	Type	Description	Default
mode	RoundingMode	The RoundingMode to set.	CONV_EVEN

Returns:

Type	Description
ExternalFunction	ExternalFunction set_rounding_<mode>, which takes no arguments.

Source code in python/iron/kernels/core.py

def set_rounding(mode: RoundingMode = RoundingMode.CONV_EVEN) -> ExternalFunction:
    """Set the core's rounding mode using always-inline, merge-linked LLVM IR.

    Args:
        mode: The [`RoundingMode`][iron.kernels.core.RoundingMode] to set.

    Returns:
        ExternalFunction ``set_rounding_<mode>``, which takes no arguments.
    """
    mode = RoundingMode(mode)
    return _make_extern(
        f"set_rounding_{mode}",
        _kernel_source("core/set_rounding.cc"),
        [],
        compile_flags=[f"-DROUNDING_MODE={mode}"],
        inline=True,
        # Sets core state and has no data arguments.
        contract=KernelContract(
            roles=(), trace=Trace.none("runs once before the calls; nothing to time")
        ),
    )

Reduction

Reduction kernel factories: reduce_add, reduce_min, reduce_max, compute_max.

reduce_add_ref

reduce_add_ref(x)

Numpy reference for reduce_add: per-tile sum.

In int64 for integer tiles, in float64 for bfloat16 ones.

Source code in python/iron/kernels/reduce.py

def reduce_add_ref(x):
    """Numpy reference for [`reduce_add`][iron.kernels.reduce.reduce_add]: per-tile sum.

    In int64 for integer tiles, in float64 for bfloat16 ones.
    """
    x = np.asarray(x)
    wide = np.int64 if np.issubdtype(x.dtype, np.integer) else np.float64
    return x.astype(wide).sum(axis=-1, keepdims=True)

reduce_min_ref

reduce_min_ref(x)

Numpy reference for reduce_min: per-tile minimum.

Source code in python/iron/kernels/reduce.py

def reduce_min_ref(x):
    """Numpy reference for [`reduce_min`][iron.kernels.reduce.reduce_min]: per-tile minimum."""
    return x.min(axis=-1, keepdims=True)

reduce_max_ref

reduce_max_ref(x)

Numpy reference for reduce_max: per-tile maximum.

Source code in python/iron/kernels/reduce.py

def reduce_max_ref(x):
    """Numpy reference for [`reduce_max`][iron.kernels.reduce.reduce_max]: per-tile maximum."""
    return x.max(axis=-1, keepdims=True)

reduce_add

reduce_add(
    tile_size: int = 1024,
    dtype: type = int32,
    vectorized: bool = True,
) -> ExternalFunction

Reduction kernel: sums all elements of a tile to a scalar (int32 or bfloat16).

A bfloat16 tile is summed in fp32 and the result rounded once, to the nearest bfloat16.

Parameters:

Name	Type	Description	Default
tile_size	int	Number of elements in the input tile.	1024
dtype	type	Element data type (np.int32 or bfloat16).	int32
vectorized	bool	If True use vectorized path; False selects scalar.	True

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the reduce_add kernel.

Raises:

Type	Description
ValueError	When dtype is not np.int32 or bfloat16.

Source code in python/iron/kernels/reduce.py

@dtypes(({"dtype": np.int32}, {"dtype": bfloat16}))
def reduce_add(
    tile_size: int = 1024, dtype: type = np.int32, vectorized: bool = True
) -> ExternalFunction:
    """Reduction kernel: sums all elements of a tile to a scalar (int32 or bfloat16).

    A bfloat16 tile is summed in fp32 and the result rounded once, to the
    nearest bfloat16.

    Args:
        tile_size: Number of elements in the input tile.
        dtype: Element data type (``np.int32`` or ``bfloat16``).
        vectorized: If ``True`` use vectorized path; ``False`` selects scalar.

    Returns:
        ExternalFunction configured for the reduce_add kernel.

    Raises:
        ValueError: When ``dtype`` is not ``np.int32`` or ``bfloat16``.
    """
    return _reduce_kernel("add", tile_size, dtype, vectorized)

reduce_min

reduce_min(
    tile_size: int = 1024,
    dtype: type = int32,
    vectorized: bool = True,
) -> ExternalFunction

Reduction kernel: finds the minimum element of a tile (int32 or bfloat16).

Parameters:

Name	Type	Description	Default
tile_size	int	Number of elements in the input tile.	1024
dtype	type	Element data type (np.int32 or bfloat16).	int32
vectorized	bool	If True use vectorized path; False selects scalar.	True

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the reduce_min kernel.

Raises:

Type	Description
ValueError	When dtype is not np.int32 or bfloat16.

Source code in python/iron/kernels/reduce.py

@dtypes(({"dtype": np.int32}, {"dtype": bfloat16}))
def reduce_min(
    tile_size: int = 1024, dtype: type = np.int32, vectorized: bool = True
) -> ExternalFunction:
    """Reduction kernel: finds the minimum element of a tile (int32 or bfloat16).

    Args:
        tile_size: Number of elements in the input tile.
        dtype: Element data type (``np.int32`` or ``bfloat16``).
        vectorized: If ``True`` use vectorized path; ``False`` selects scalar.

    Returns:
        ExternalFunction configured for the reduce_min kernel.

    Raises:
        ValueError: When ``dtype`` is not ``np.int32`` or ``bfloat16``.
    """
    return _reduce_kernel("min", tile_size, dtype, vectorized)

reduce_max

reduce_max(
    tile_size: int = 1024,
    dtype: type = int32,
    vectorized: bool = True,
) -> ExternalFunction

Reduction kernel: finds the maximum element of a tile (int32 or bfloat16).

Parameters:

Name	Type	Description	Default
tile_size	int	Number of elements in the input tile.	1024
dtype	type	Element data type (np.int32 or bfloat16).	int32
vectorized	bool	If True use vectorized path; False selects scalar.	True

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the reduce_max kernel.

Raises:

Type	Description
ValueError	When dtype is not np.int32 or bfloat16.

Source code in python/iron/kernels/reduce.py

@dtypes(({"dtype": np.int32}, {"dtype": bfloat16}))
def reduce_max(
    tile_size: int = 1024, dtype: type = np.int32, vectorized: bool = True
) -> ExternalFunction:
    """Reduction kernel: finds the maximum element of a tile (int32 or bfloat16).

    Args:
        tile_size: Number of elements in the input tile.
        dtype: Element data type (``np.int32`` or ``bfloat16``).
        vectorized: If ``True`` use vectorized path; ``False`` selects scalar.

    Returns:
        ExternalFunction configured for the reduce_max kernel.

    Raises:
        ValueError: When ``dtype`` is not ``np.int32`` or ``bfloat16``.
    """
    is_bf16 = np.dtype(dtype) == np.dtype(bfloat16)
    is_int32 = np.dtype(dtype) == np.dtype(np.int32)
    if not is_bf16 and not is_int32:
        raise ValueError(
            f"reduce_max() dtype must be np.int32 or bfloat16, got {dtype}"
        )

    actual_dtype = bfloat16 if is_bf16 else np.int32
    if vectorized:
        _require_vector_alignment(
            "reduce_max",
            tile_size,
            _REDUCE_VEC_ELEMS_BF16 if is_bf16 else _REDUCE_VEC_ELEMS,
        )
    in_ty = np.ndarray[(tile_size,), np.dtype[actual_dtype]]
    # The C++ kernel writes one scalar; the output tile must still be at least
    # 4 bytes for shim-DMA alignment, so bfloat16 callers get out_size=2 even
    # though they only read the first element.
    out_ty = np.ndarray[(_min_dma_aligned_elems(actual_dtype),), np.dtype[actual_dtype]]

    func_variant = "vector" if vectorized else "scalar"
    suffix = "_bfloat16" if is_bf16 else ""
    return _make_extern(
        f"reduce_max_{func_variant}{suffix}",
        _kernel_source("reduce/reduce_max.cc"),
        [in_ty, out_ty, np.int32],
        compile_flags=[f"-DREDUCE_MAX_ELEMS={tile_size}"],
        contract=_reduce_contract("max", tile_size),
    )

compute_max

compute_max(dtype: type = int32) -> ExternalFunction

Pairwise scalar max — companion to reduce_max.

Used for multi-core reductions where each core produces a partial max and a final tree reduces them pairwise.

Lives in the same reduce_max.cc as reduce_max, but uses an unspecialized object independent of reduction tile sizes.

Parameters:

Name	Type	Description	Default
dtype	type	Element data type (np.int32 or bfloat16).	int32

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the compute_max kernel; signature
ExternalFunction	is (out_ty, out_ty, out_ty) where out_ty is a one-element
ExternalFunction	(DMA-aligned) tile of dtype.

Raises:

Type	Description
ValueError	When dtype is not np.int32 or bfloat16.

Source code in python/iron/kernels/reduce.py

@dtypes(({"dtype": np.int32}, {"dtype": bfloat16}))
def compute_max(dtype: type = np.int32) -> ExternalFunction:
    """Pairwise scalar max — companion to [`reduce_max`][iron.kernels.reduce.reduce_max].

    Used for multi-core reductions where each core produces a partial max and a
    final tree reduces them pairwise.

    Lives in the same ``reduce_max.cc`` as [`reduce_max`][iron.kernels.reduce.reduce_max],
    but uses an unspecialized object independent of reduction tile sizes.

    Args:
        dtype: Element data type (``np.int32`` or ``bfloat16``).

    Returns:
        ExternalFunction configured for the ``compute_max`` kernel; signature
        is ``(out_ty, out_ty, out_ty)`` where ``out_ty`` is a one-element
        (DMA-aligned) tile of ``dtype``.

    Raises:
        ValueError: When ``dtype`` is not ``np.int32`` or ``bfloat16``.
    """
    is_bf16 = np.dtype(dtype) == np.dtype(bfloat16)
    is_int32 = np.dtype(dtype) == np.dtype(np.int32)
    if not is_bf16 and not is_int32:
        raise ValueError(
            f"compute_max() dtype must be np.int32 or bfloat16, got {dtype}"
        )
    actual_dtype = bfloat16 if is_bf16 else np.int32
    out_ty = np.ndarray[(_min_dma_aligned_elems(actual_dtype),), np.dtype[actual_dtype]]

    suffix = "_bfloat16" if is_bf16 else ""
    return _make_extern(
        f"compute_max{suffix}",
        _kernel_source("reduce/reduce_max.cc"),
        [out_ty, out_ty, out_ty],
        object_file_name=_REDUCE_MAX_OBJ,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            reference=compute_max_ref,
            tolerance=Tolerance.exact(note="selection"),
            ops_per_call=1,
            out_valid=1,
        ),
    )

compute_max_ref

compute_max_ref(a, b)

Numpy reference for compute_max.

The kernel compares only element 0 of each (DMA-padded) input tile and writes element 0 of the output; the reference does the same, returning max(a[..., 0], b[..., 0]) with a trailing axis of length 1.

Source code in python/iron/kernels/reduce.py

def compute_max_ref(a, b):
    """Numpy reference for [`compute_max`][iron.kernels.reduce.compute_max].

    The kernel compares only element 0 of each (DMA-padded) input tile and
    writes element 0 of the output; the reference does the same, returning
    ``max(a[..., 0], b[..., 0])`` with a trailing axis of length 1.
    """
    return np.maximum(np.asarray(a)[..., :1], np.asarray(b)[..., :1])

Linear algebra

Linear algebra kernel factories: mm, mv, cascade_mm.

StreamDimsABC

Bases: NamedTuple

The three to_stream a matmul design needs, one per operand.

None for an operand a build streams untransformed.

MatrixKernel

Bases: _ZeroInitializedKernel

A kernel whose first three operands are the A, B and C of a product.

The blocking and the DMA transforms are declared once, on the contract's operand layouts (TensorLayout.block and .stream); these properties read them back in the form the matrix designs consume.

mac_dims property

mac_dims: tuple[int, int, int]

(r, s, t): the MMUL micro-tile of A (r, s), B (s, t) and C (r, t).

stream_dims property

stream_dims: StreamDimsABC

The to_stream a design applies to A, B and C; None streams as stored.

mm_acc_dtype

mm_acc_dtype(input_dtype)

Return what mm.cc accumulates in for input_dtype.

accauto is acc32 for 8-bit and acc64 for 16-bit integer inputs, and float32 for bf16.

Source code in python/iron/kernels/linalg.py

def mm_acc_dtype(input_dtype):
    """Return what ``mm.cc`` accumulates in for ``input_dtype``.

    ``accauto`` is acc32 for 8-bit and acc64 for 16-bit integer inputs, and
    float32 for bf16.
    """
    dt = np.dtype(input_dtype)
    if dt == np.dtype(np.int8):
        return np.int32
    if dt == np.dtype(np.int16):
        return np.int64
    return np.float32

mm_ref

mm_ref(a, b)

Numpy reference for mm: a @ b accumulated in int64 or float64.

Source code in python/iron/kernels/linalg.py

def mm_ref(a, b):
    """Numpy reference for [`mm`][iron.kernels.linalg.mm]: ``a @ b`` accumulated in int64 or float64."""
    acc = np.int64 if np.issubdtype(a.dtype, np.integer) else np.float64
    return a.astype(acc) @ b.astype(acc)

mv_ref

mv_ref(a, b)

Numpy reference for mv: a @ b accumulated in int64.

Source code in python/iron/kernels/linalg.py

def mv_ref(a, b):
    """Numpy reference for [`mv`][iron.kernels.linalg.mv]: ``a @ b`` accumulated in int64."""
    return a.astype(np.int64) @ b.astype(np.int64)

mm_bfp_ref

mm_bfp_ref(a, b)

Numpy reference for mm_bfp: a @ b on bfp16ebs8-quantized operands.

a is (M, K) and b (K, N) float; each is quantized the way the host encodes it for the kernel (blocks of 8 along K, see aie.utils.bfp) and the product is accumulated in float64. The kernel's own output is bfp16ebs8 too, which the tolerance covers.

Source code in python/iron/kernels/linalg.py

def mm_bfp_ref(a, b):
    """Numpy reference for [`mm_bfp`][iron.kernels.linalg.mm_bfp]: ``a @ b`` on bfp16ebs8-quantized operands.

    ``a`` is ``(M, K)`` and ``b`` ``(K, N)`` float; each is quantized the way
    the host encodes it for the kernel (blocks of 8 along ``K``, see
    ``aie.utils.bfp``) and the product is accumulated in float64. The
    kernel's own output is bfp16ebs8 too, which the tolerance covers.
    """
    from aie.utils import bfp

    aq = bfp.quantize(np.asarray(a, dtype=np.float32)).astype(np.float64)
    bq = bfp.quantize(np.ascontiguousarray(np.asarray(b, dtype=np.float32).T)).T
    return aq @ bq.astype(np.float64)

mm_bfp_mixed_ref

mm_bfp_mixed_ref(a, b)

Numpy reference for mm_bfp with mixed=True.

a arrives bf16 and the core converts it, under the rounding mode mm_bfp_mixed.cc pins, so it is quantized conv_even here. b is encoded by the host, which truncates, so it keeps the default as in mm_bfp_ref.

Pairing A with the wrong mode is the difference between 10 mismatching outputs and 2791, on a 64x64x64 tile of large inputs; leaving A unquantized altogether gives 1369.

Source code in python/iron/kernels/linalg.py

def mm_bfp_mixed_ref(a, b):
    """Numpy reference for [`mm_bfp`][iron.kernels.linalg.mm_bfp] with ``mixed=True``.

    ``a`` arrives bf16 and the core converts it, under the rounding mode
    ``mm_bfp_mixed.cc`` pins, so it is quantized ``conv_even`` here. ``b`` is
    encoded by the host, which truncates, so it keeps the default as in
    [`mm_bfp_ref`][iron.kernels.linalg.mm_bfp_ref].

    Pairing A with the wrong mode is the difference between 10 mismatching
    outputs and 2791, on a 64x64x64 tile of large inputs; leaving A
    unquantized altogether gives 1369.
    """
    from aie.utils import bfp

    aq = bfp.quantize(np.asarray(a, dtype=np.float32), rounding="conv_even").astype(
        np.float64
    )
    bq = bfp.quantize(np.ascontiguousarray(np.asarray(b, dtype=np.float32).T)).T
    return aq @ bq.astype(np.float64)

mm_tile_ref

mm_tile_ref(a, b, *, dim_m: int, dim_k: int, dim_n: int)

One mm call: a (dim_m, dim_k) tile times a (dim_k, dim_n) one.

Tiles arrive flattened as (calls, ...), one row per call, and one (dim_m * dim_n,) row comes back per call.

Source code in python/iron/kernels/linalg.py

def mm_tile_ref(a, b, *, dim_m: int, dim_k: int, dim_n: int):
    """One [`mm`][iron.kernels.linalg.mm] call: a ``(dim_m, dim_k)`` tile times a ``(dim_k, dim_n)`` one.

    Tiles arrive flattened as ``(calls, ...)``, one row per call, and one
    ``(dim_m * dim_n,)`` row comes back per call.
    """
    a = np.asarray(a)
    acc = np.int64 if np.issubdtype(a.dtype, np.integer) else np.float64
    a = a.reshape(-1, dim_m, dim_k).astype(acc)
    b = np.asarray(b).reshape(-1, dim_k, dim_n).astype(acc)
    return (a @ b).reshape(len(a), dim_m * dim_n)

mv_tile_ref

mv_tile_ref(a, b, *, dim_m: int, dim_k: int)

One mv call: a (dim_m, dim_k) tile times a (dim_k,) vector.

Source code in python/iron/kernels/linalg.py

def mv_tile_ref(a, b, *, dim_m: int, dim_k: int):
    """One [`mv`][iron.kernels.linalg.mv] call: a ``(dim_m, dim_k)`` tile times a ``(dim_k,)`` vector."""
    a = np.asarray(a).reshape(-1, dim_m, dim_k).astype(np.int64)
    b = np.asarray(b).reshape(-1, dim_k).astype(np.int64)
    return np.einsum("cmk,ck->cm", a, b)

mm_bfp_tile_ref

mm_bfp_tile_ref(
    a,
    b,
    *,
    dim_m: int,
    dim_k: int,
    dim_n: int,
    mixed: bool = False
)

One mm_bfp call, on operands quantized as the host encodes them.

Blocks of 8 run along K for both operands, so B is quantized transposed. With mixed the A tile stays bf16 and the core converts it itself, with a rounding this does not model -- which is what the wider mixed tolerance covers. See mm_bfp_mixed_ref for why modelling it as a host-side quantize is worse, not better.

Source code in python/iron/kernels/linalg.py

def mm_bfp_tile_ref(a, b, *, dim_m: int, dim_k: int, dim_n: int, mixed: bool = False):
    """One [`mm_bfp`][iron.kernels.linalg.mm_bfp] call, on operands quantized as the host encodes them.

    Blocks of 8 run along K for both operands, so B is quantized transposed.
    With ``mixed`` the A tile stays bf16 and the core converts it itself, with
    a rounding this does not model -- which is what the wider mixed tolerance
    covers. See [`mm_bfp_mixed_ref`][iron.kernels.linalg.mm_bfp_mixed_ref] for
    why modelling it as a host-side quantize is worse, not better.
    """
    from aie.utils import bfp

    a = np.asarray(a).reshape(-1, dim_m, dim_k)
    b = np.asarray(b).reshape(-1, dim_k, dim_n)
    aq = (
        bfp.quantize(a.astype(np.float32), rounding="conv_even").astype(np.float64)
        if mixed
        else bfp.quantize(a.astype(np.float32)).astype(np.float64)
    )
    bt = np.ascontiguousarray(b.astype(np.float32).swapaxes(-1, -2))
    bq = bfp.quantize(bt).swapaxes(-1, -2).astype(np.float64)
    return (aq @ bq).reshape(len(a), dim_m * dim_n)

mm_stream_dims

mm_stream_dims(
    dim_m: int,
    dim_k: int,
    dim_n: int,
    mac_dims,
    *,
    b_col_maj: bool = False,
    c_col_maj: bool = False
) -> StreamDimsABC

DMA to_stream that feed mm.cc its (r, s, t) micro-tiles.

mm.cc consumes A, B and produces C in the micro-tile blocking given by mac_dims; a plain row-major stream yields wrong numbers, not an error. Every matmul design (single_core, whole_array, cascade, ...) derives these same three transforms from mac_dims; kernels.mm(...).stream_dims carries them so designs do not re-derive them. Keys "A", "B", "C".

b_col_maj describes a B tile stored as (n, k) (the transpose) and c_col_maj a C tile emitted as (n, m), matching the kernel's -DB_COL_MAJ / -DC_COL_MAJ builds.

Source code in python/iron/kernels/linalg.py

def mm_stream_dims(
    dim_m: int,
    dim_k: int,
    dim_n: int,
    mac_dims,
    *,
    b_col_maj: bool = False,
    c_col_maj: bool = False,
) -> StreamDimsABC:
    """DMA ``to_stream`` that feed ``mm.cc`` its (r, s, t) micro-tiles.

    ``mm.cc`` consumes A, B and produces C in the micro-tile blocking given by
    ``mac_dims``; a plain row-major stream yields wrong numbers, not an error.
    Every matmul design (single_core, whole_array, cascade, ...) derives these
    same three transforms from ``mac_dims``; ``kernels.mm(...).stream_dims``
    carries them so designs do not re-derive them. Keys ``"A"``, ``"B"``, ``"C"``.

    ``b_col_maj`` describes a B tile stored as ``(n, k)`` (the transpose) and
    ``c_col_maj`` a C tile emitted as ``(n, m)``, matching the kernel's
    ``-DB_COL_MAJ`` / ``-DC_COL_MAJ`` builds.
    """
    r, s, t = mac_dims
    m, k, n = dim_m, dim_k, dim_n
    # A and B are read row-major and emitted as (r x s) / (s x t) blocks.
    a = TensorAccessPattern.full((m, k)).tile((r, s))
    b_shape, b_tile = ((n, k), (t, s)) if b_col_maj else ((k, n), (s, t))
    b = TensorAccessPattern.full(b_shape).tile(b_tile)
    # C goes the other way: the DMA reads the core's block-ordered buffer and
    # emits it row-major, so the intra-tile row term sits outside the tile
    # index -- (r, t) before (n//t, r*t). That is the inverse of tiling.
    c_shape, c_tile = ((n, m), (t, r)) if c_col_maj else ((m, n), (r, t))
    c = TensorAccessPattern.full(c_shape).tile(c_tile).inverse()
    return StreamDimsABC(A=a, B=b, C=c)

mm

mm(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    input_dtype: type = int16,
    output_dtype: type = int16,
    vectorized: bool = True,
    b_col_maj: bool = False,
    c_col_maj: bool = False,
    use_chess: bool = False,
    emulate_bf16_mmul_with_bfp16: bool = False,
    round_conv_even: bool = False,
) -> MatrixKernel

Matrix-multiply kernel: C += A * B.

.zero initializes the accumulator using the independent, reusable kernels.zero(dim_m * dim_n, output_dtype) kernel. The contract declares the same initializer for the generic harness.

Parameters:

Name	Type	Description	Default
dim_m	int	Number of rows of A / C.	64
dim_k	int	Number of columns of A / rows of B.	64
dim_n	int	Number of columns of B / C.	64
input_dtype	type	Input element type (np.int8, np.int16, or bfloat16).	int16
output_dtype	type	Output element type.	int16
vectorized	bool	If True use the vectorized variant.	True
b_col_maj	bool	If True compile with -DB_COL_MAJ so the kernel consumes B laid out column-major. Must agree with the design's B to_stream.	False
c_col_maj	bool	If True compile with -DC_COL_MAJ so the kernel writes C laid out column-major. Must agree with the design's C output to_stream.	False
use_chess	bool	If True build with xchesscc_wrapper instead of Peano's clang++. All ExternalFunctions in a single @iron.jit design must share the same toolchain.	False
emulate_bf16_mmul_with_bfp16	bool	AIE2P only, bf16 inputs only. When True compile with -DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16 so the kernel uses BFP16-based emulation of the bf16 MMUL. Changes the micro-kernel dims to (8, 8, 8); designs reading .mac_dims will see the new geometry automatically. Ignored for non-bf16 inputs and on AIE2.	False
round_conv_even	bool	AIE2 only, bf16 inputs only. When True compile with -DROUND_CONV_EVEN so the kernel itself selects round-to-nearest-even for each call and restores the caller's mode, in place of the contract's conv_even setup. AIE2P's kernel always does this, so it is ignored there.	False

Returns:

Type	Description
MatrixKernel	ExternalFunction configured for the matmul kernel.

Raises:

Type	Description
ValueError	When (input_dtype, output_dtype) is not a supported combination.

Source code in python/iron/kernels/linalg.py

@dtypes(
    tuple({"input_dtype": i, "output_dtype": o} for (i, o) in _MM_MAC_DIMS["aie2p"])
)
def mm(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    input_dtype: type = np.int16,
    output_dtype: type = np.int16,
    vectorized: bool = True,
    b_col_maj: bool = False,
    c_col_maj: bool = False,
    use_chess: bool = False,
    emulate_bf16_mmul_with_bfp16: bool = False,
    round_conv_even: bool = False,
) -> MatrixKernel:
    """Matrix-multiply kernel: C += A * B.

    ``.zero`` initializes the accumulator using the independent, reusable
    ``kernels.zero(dim_m * dim_n, output_dtype)`` kernel. The contract declares
    the same initializer for the generic harness.

    Args:
        dim_m: Number of rows of A / C.
        dim_k: Number of columns of A / rows of B.
        dim_n: Number of columns of B / C.
        input_dtype: Input element type (``np.int8``, ``np.int16``, or ``bfloat16``).
        output_dtype: Output element type.
        vectorized: If ``True`` use the vectorized variant.
        b_col_maj: If ``True`` compile with ``-DB_COL_MAJ`` so the kernel
            consumes B laid out column-major.  Must agree with the
            design's B ``to_stream``.
        c_col_maj: If ``True`` compile with ``-DC_COL_MAJ`` so the kernel
            writes C laid out column-major.  Must agree with the design's
            C output ``to_stream``.
        use_chess: If ``True`` build with ``xchesscc_wrapper`` instead of
            Peano's ``clang++``.  All ExternalFunctions in a single
            ``@iron.jit`` design must share the same toolchain.
        emulate_bf16_mmul_with_bfp16: AIE2P only, bf16 inputs only.  When
            ``True`` compile with ``-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16``
            so the kernel uses BFP16-based emulation of the bf16 MMUL.
            Changes the micro-kernel dims to (8, 8, 8); designs reading
            ``.mac_dims`` will see the new geometry automatically.  Ignored
            for non-bf16 inputs and on AIE2.
        round_conv_even: AIE2 only, bf16 inputs only.  When ``True``
            compile with ``-DROUND_CONV_EVEN`` so the kernel itself selects
            round-to-nearest-even for each call and restores the caller's
            mode, in place of the contract's ``conv_even`` setup.  AIE2P's
            kernel always does this, so it is ignored there.

    Returns:
        ExternalFunction configured for the matmul kernel.

    Raises:
        ValueError: When ``(input_dtype, output_dtype)`` is not a supported combination.
    """
    key = (input_dtype, output_dtype)
    if key not in _MM_COMBOS:
        raise ValueError(
            f"mm(): unsupported (input_dtype, output_dtype) = {key}. Supported: {list(_MM_COMBOS.keys())}"
        )

    suffix, only_flag = _MM_COMBOS[key]
    prefix = "matmul" if vectorized else "matmul_scalar"
    a_ty = np.ndarray[(dim_m * dim_k,), np.dtype[input_dtype]]
    b_ty = np.ndarray[(dim_k * dim_n,), np.dtype[input_dtype]]
    c_ty = np.ndarray[(dim_m * dim_n,), np.dtype[output_dtype]]
    compile_flags = [
        f"-DDIM_M={dim_m}",
        f"-DDIM_K={dim_k}",
        f"-DDIM_N={dim_n}",
        f"-D{only_flag}",
    ]
    if b_col_maj:
        compile_flags.append("-DB_COL_MAJ")
    if c_col_maj:
        compile_flags.append("-DC_COL_MAJ")
    arch = _detect_arch()
    bf16_emulated = (
        emulate_bf16_mmul_with_bfp16
        and ARCH_TRAITS[arch].bfp16
        and input_dtype is bfloat16
    )
    if bf16_emulated:
        compile_flags.append("-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16")
    self_rounding = round_conv_even and arch != "aie2p" and input_dtype is bfloat16
    if self_rounding:
        compile_flags.append("-DROUND_CONV_EVEN")
    # The scalar kernel walks its operands element by element in row-major
    # order: its micro-tile is 1x1x1 and nothing is streamed transformed.
    r, s, t = _MatMulFactory.mac_dims(
        input_dtype,
        output_dtype,
        arch=arch,
        emulate_bf16_mmul_with_bfp16=emulate_bf16_mmul_with_bfp16,
        vectorized=vectorized,
    )
    streams = mm_stream_dims(
        dim_m, dim_k, dim_n, (r, s, t), b_col_maj=b_col_maj, c_col_maj=c_col_maj
    )
    # Host-side layout the streams assume: with b_col_maj, B is given as
    # (n, k) tiles of B^T; with c_col_maj, C is emitted as (n, m) tiles of
    # C^T. The layouts carry that, so the builder transposes the host
    # operands without knowing which kernel it is building.
    layouts = (
        _tile_layout((dim_m, dim_k), streams.A if vectorized else None, block=(r, s)),
        _tile_layout(
            (dim_k, dim_n),
            streams.B if vectorized else None,
            axes=(1, 0) if b_col_maj else None,
            block=(s, t),
        ),
        _tile_layout(
            (dim_m, dim_n),
            streams.C if vectorized else None,
            axes=(1, 0) if c_col_maj else None,
            inverse=True,
            block=(r, t),
        ),
    )
    return _make_extern(
        f"{prefix}_{suffix}",
        _kernel_source("linalg/mm.cc"),
        [a_ty, b_ty, c_ty],
        compile_flags=compile_flags,
        use_chess=use_chess,
        cls=MatrixKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            layouts=layouts,
            # Tuned or not: at most 512 B on aie2
            # and 192 B on aie2p, except aie2p's int8 -> int32 kernel, whose
            # fully unrolled K loop spills about 16 B per unit of K (c_col_maj:
            # 1088 B at K = 56, 6208 B at K = 384, at most 1024 B up to K =
            # 48). mm_aie2p.h rolls K from K = 416 (192 B), and a 16x16 tile
            # does not spill (640 B). The fit covers only that kernel over 48 <
            # dim_k < 416 (787 M/N/K/layout combos, see 0c60e6be3f2); the rest
            # fit the 1024 B default. Too small a value fails the aiecc build
            # loudly, as checkStackSizeRequirements reads the real
            # .stack_sizes. A chess core cannot be measured, so use_chess keeps
            # the matrix_multiplication examples' constant.
            stack_bytes=(
                0xD00
                if use_chess
                else (
                    16 * dim_k + 256
                    if arch == "aie2p"
                    and vectorized
                    and key == (np.int8, np.int32)
                    and 48 < dim_k < 416
                    else None
                )
            ),
            # mm_aie2p.h sets conv_even itself and restores it; mm_aie2.h
            # does so only under round_conv_even, and otherwise stores bf16
            # in whatever mode the core is in.
            setup=(
                conv_even
                if arch != "aie2p" and output_dtype is bfloat16 and not self_rounding
                else None
            ),
            roles=(In, In, InOut),
            reference=partial(mm_tile_ref, dim_m=dim_m, dim_k=dim_k, dim_n=dim_n),
            initializers=((2, _zero_output),),
            acc_dtype=mm_acc_dtype(input_dtype),
            reduction=dim_k,
            tolerance=_linalg_tolerance(input_dtype),
            ops_per_call=2 * dim_m * dim_k * dim_n,
        ),
    )

mv

mv(
    dim_m: int = 32,
    dim_k: int = 32,
    input_dtype: type = int16,
    output_dtype: type = int32,
    vectorized: bool = True,
    use_chess: bool = False,
    vec_size: int = 64,
    output_rows: int | None = None,
    a_col_maj: bool = False,
) -> ExternalFunction

Matrix-vector multiply kernel: c += A * b.

(np.int16, np.int32) builds aie_kernels/linalg/mv_i16.cc; its vectorized path reads A word-transposed, which A's layout carries (contract.layouts[0].stream). Its .zero companion initializes C with the independent kernels.zero(dim_m, output_dtype). (bfloat16, bfloat16) builds aie_kernels/linalg/mv_bf16.cc, IRON's GEMV kernel, whose signature is (m, row_offset, A, b, c): row_offset shifts the write into c so one core can fill several output blocks; A is row-major.

a_col_maj=True builds the same file with -DA_COL_MAJ, signature (flags, A, b, acc, c): A is dim_k stored rows of a (K, M) matrix, dim_m elements each (the transpose), and b their dim_k elements of the vector. acc holds vec_size * dim_m float32 partial sums that carry from call to call, so a whole K is a MV_COL_MAJ_FIRST call, calls with no flags, and a MV_COL_MAJ_LAST call, which writes the dim_m outputs to c (or one call with both flags). The sums are the row-major kernel's at the same vec_size over that K, in its order, so over A.T it returns the same bits, on rows the row-major kernel computes in whole groups of four (the source spells out the order). Choose vec_size as the row-major kernel would for the whole K.

Parameters:

Name	Type	Description	Default
dim_m	int	Number of rows of A (output vector length). With a_col_maj: 16, 32 or a multiple of 64.	32
dim_k	int	Number of columns of A (input vector length). With a_col_maj: the stored rows of A one call takes.	32
input_dtype	type	Input element type: np.int16 or bfloat16.	int16
output_dtype	type	Output element type: np.int32 for np.int16 inputs, bfloat16 for bfloat16 inputs.	int32
vectorized	bool	If True use the vectorized variant.	True
use_chess	bool	If True build the .o with xchesscc_wrapper instead of Peano. See mm for the design-level constraint (all EFs in one design must agree).	False
vec_size	int	bf16 only: the kernel's VEC_SIZE accumulation width; 16, 32 or 64 with a_col_maj.	64
output_rows	int | None	bf16 only: the rows of a C tile that successive calls fill dim_m rows at a time through row_offset, with b held for all of them, as amd/IRON's GEMV core does (its tile_size_output). None: every call writes its own dim_m rows.	None
a_col_maj	bool	bf16 and vectorized only: A is stored (dim_k, dim_m), the transpose, with partial sums carried across calls (above).	False

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the matvec kernel.

Raises:

Type	Description
ValueError	When the dtype combination is not supported.

Source code in python/iron/kernels/linalg.py

@dtypes(
    (
        {"input_dtype": np.int16, "output_dtype": np.int32},
        # dim_k: the bf16 kernel accumulates VEC_SIZE (64) elements at a
        # time, so the default dim_k of 32 is too short for it.
        {"input_dtype": bfloat16, "output_dtype": bfloat16, "dim_k": 256},
    )
)
def mv(
    dim_m: int = 32,
    dim_k: int = 32,
    input_dtype: type = np.int16,
    output_dtype: type = np.int32,
    vectorized: bool = True,
    use_chess: bool = False,
    vec_size: int = 64,
    output_rows: int | None = None,
    a_col_maj: bool = False,
) -> ExternalFunction:
    """Matrix-vector multiply kernel: c += A * b.

    ``(np.int16, np.int32)`` builds ``aie_kernels/linalg/mv_i16.cc``; its
    vectorized path reads A word-transposed, which A's layout carries
    (``contract.layouts[0].stream``). Its ``.zero`` companion initializes C
    with the independent ``kernels.zero(dim_m, output_dtype)``.
    ``(bfloat16, bfloat16)`` builds
    ``aie_kernels/linalg/mv_bf16.cc``, IRON's ``GEMV`` kernel, whose signature
    is ``(m, row_offset, A, b, c)``: ``row_offset`` shifts the write into
    ``c`` so one core can fill several output blocks; A is row-major.

    ``a_col_maj=True`` builds the same file with ``-DA_COL_MAJ``, signature
    ``(flags, A, b, acc, c)``: ``A`` is ``dim_k`` stored rows of a ``(K, M)``
    matrix, ``dim_m`` elements each (the transpose), and ``b`` their ``dim_k``
    elements of the vector. ``acc`` holds ``vec_size * dim_m`` float32 partial
    sums that carry from call to call, so a whole ``K`` is a
    ``MV_COL_MAJ_FIRST`` call, calls with no flags, and a ``MV_COL_MAJ_LAST``
    call, which writes the ``dim_m`` outputs to ``c`` (or one call with both
    flags). The sums are the row-major kernel's at the same ``vec_size`` over
    that ``K``, in its order, so over ``A.T`` it returns the same bits, on
    rows the row-major kernel computes in whole groups of four (the source
    spells out the order). Choose ``vec_size`` as the row-major kernel would
    for the whole ``K``.

    Args:
        dim_m: Number of rows of A (output vector length). With
            ``a_col_maj``: 16, 32 or a multiple of 64.
        dim_k: Number of columns of A (input vector length). With
            ``a_col_maj``: the stored rows of A one call takes.
        input_dtype: Input element type: ``np.int16`` or ``bfloat16``.
        output_dtype: Output element type: ``np.int32`` for ``np.int16``
            inputs, ``bfloat16`` for ``bfloat16`` inputs.
        vectorized: If ``True`` use the vectorized variant.
        use_chess: If ``True`` build the .o with ``xchesscc_wrapper``
            instead of Peano.  See [`mm`][iron.kernels.linalg.mm] for the design-level
            constraint (all EFs in one design must agree).
        vec_size: bf16 only: the kernel's ``VEC_SIZE`` accumulation width;
            16, 32 or 64 with ``a_col_maj``.
        output_rows: bf16 only: the rows of a C tile that successive calls
            fill ``dim_m`` rows at a time through ``row_offset``, with b
            held for all of them, as amd/IRON's GEMV core does (its
            ``tile_size_output``). ``None``: every call writes its own
            ``dim_m`` rows.
        a_col_maj: bf16 and vectorized only: A is stored ``(dim_k, dim_m)``,
            the transpose, with partial sums carried across calls (above).

    Returns:
        ExternalFunction configured for the matvec kernel.

    Raises:
        ValueError: When the dtype combination is not supported.
    """
    if (input_dtype, output_dtype) == (bfloat16, bfloat16):
        return _mv_bf16(
            dim_m, dim_k, vectorized, use_chess, vec_size, output_rows, a_col_maj
        )
    if a_col_maj:
        raise ValueError("mv(): a_col_maj is a bf16 layout")
    if output_rows is not None:
        raise ValueError("mv(): output_rows needs the bf16 kernel's row_offset")
    if input_dtype != np.int16 or output_dtype != np.int32:
        raise ValueError(
            f"mv(): only (np.int16, np.int32) and (bfloat16, bfloat16) are supported, got ({input_dtype}, {output_dtype})"
        )

    prefix = "matvec_vectorized" if vectorized else "matvec_scalar"
    a_ty = np.ndarray[(dim_m * dim_k,), np.dtype[np.int16]]
    b_ty = np.ndarray[(dim_k,), np.dtype[np.int16]]
    c_ty = np.ndarray[(dim_m,), np.dtype[np.int32]]
    # The vectorized kernel reads A in a "32-bit-word transposed" layout (see
    # aie_kernels/linalg/mv_i16.cc): 2-byte elements are packed two per word, rows
    # of each 2-column word slowly, m rows then the next 2-col word. A design
    # applies this as from_stream on the hop into the core, reading it
    # from the layout (programming_examples/basic/matrix_multiplication/
    # matrix_vector does).
    a_from_stream = (
        TensorAccessPattern.full((dim_k // 2, dim_m, 2)).permute((1, 0, 2))
        if vectorized
        else None
    )
    return _make_extern(
        f"{prefix}_i16_i32",
        _kernel_source("linalg/mv_i16.cc"),
        [a_ty, b_ty, c_ty],
        compile_flags=[f"-DDIM_M={dim_m}", f"-DDIM_K={dim_k}"],
        use_chess=use_chess,
        cls=_ZeroInitializedKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, InOut),
            layouts=(
                _tile_layout((dim_m, dim_k), a_from_stream, inverse=True),
                TensorLayout((dim_k,)),
                TensorLayout((dim_m,)),
            ),
            reference=partial(mv_tile_ref, dim_m=dim_m, dim_k=dim_k),
            initializers=((2, _zero_output),),
            acc_dtype=np.int32,  # acc32
            reduction=dim_k,
            tolerance=Tolerance.exact(note="int16 x int16 accumulated in int32"),
            ops_per_call=2 * dim_m * dim_k,
        ),
    )

mv_bf16_ref

mv_bf16_ref(m, row_offset, a, b)

Numpy reference for the bf16 mv: a @ b over m rows, accumulated in float32.

a is row-major (rows, K); only the first m rows are computed, and the kernel writes them at row_offset into c.

Source code in python/iron/kernels/linalg.py

def mv_bf16_ref(m, row_offset, a, b):
    """Numpy reference for the bf16 [`mv`][iron.kernels.linalg.mv]: ``a @ b`` over ``m`` rows, accumulated in float32.

    ``a`` is row-major ``(rows, K)``; only the first ``m`` rows are computed,
    and the kernel writes them at ``row_offset`` into ``c``.
    """
    a = np.asarray(a).reshape(-1, np.asarray(b).shape[-1])
    out = a[: int(m)].astype(np.float32) @ np.asarray(b).astype(np.float32)
    return np.concatenate([np.zeros(int(row_offset), np.float32), out])

mm_bfp

mm_bfp(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    mixed: bool = False,
) -> MatrixKernel

Block-floating-point matmul C += A @ B on bfp16ebs8 blocks (aie2p only).

mixed=False (aie_kernels/linalg/mm_bfp.cc): A, B and C are v8bfp16ebs8 blocks, all pre-shuffled into the mmul layout, so no DMA transform applies (stream_dims is None for every operand). mixed=True (mm_bfp_mixed.cc): A is bf16 in the (r, s, t) micro-tile layout, B is bfp16ebs8, C is bf16; stream_dims.A and .C carry the transforms and .B is None.

Initialize C with the independent kernels.zero factory. The host holds B transposed (b_col_maj), and every bfp16ebs8 operand is encoded and shuffled into the mmul tile layout on the host with aie.utils.bfp, which is what the generic harness does; the contract's reference multiplies the quantized operands. These are the kernels programming_examples/ml/block_datatypes/matrix_multiplication build.

Parameters:

Name	Type	Description	Default
dim_m	int	Tile rows of A and C (multiple of 8).	64
dim_k	int	Tile columns of A / rows of B (multiple of 8).	64
dim_n	int	Tile columns of B and C (multiple of 8).	64
mixed	bool	bf16 A and C with bfp16 B.	False

Source code in python/iron/kernels/linalg.py

@dtypes(({"mixed": False}, {"mixed": True}))
def mm_bfp(
    dim_m: int = 64, dim_k: int = 64, dim_n: int = 64, mixed: bool = False
) -> MatrixKernel:
    """Block-floating-point matmul ``C += A @ B`` on bfp16ebs8 blocks (aie2p only).

    ``mixed=False`` (``aie_kernels/linalg/mm_bfp.cc``): A, B and C are
    ``v8bfp16ebs8`` blocks, all pre-shuffled into the mmul layout, so no
    DMA transform applies (``stream_dims`` is ``None`` for every operand).
    ``mixed=True`` (``mm_bfp_mixed.cc``): A is bf16 in the (r, s, t)
    micro-tile layout, B is bfp16ebs8, C is bf16; ``stream_dims.A`` and
    ``.C`` carry the transforms and ``.B`` is ``None``.

    Initialize C with the independent ``kernels.zero`` factory.
    The host holds B transposed (``b_col_maj``), and
    every bfp16ebs8 operand is encoded and shuffled into the mmul tile
    layout on the host with ``aie.utils.bfp``, which is what the generic
    harness does; the contract's reference multiplies the quantized
    operands. These are the kernels
    programming_examples/ml/block_datatypes/matrix_multiplication build.

    Args:
        dim_m: Tile rows of A and C (multiple of 8).
        dim_k: Tile columns of A / rows of B (multiple of 8).
        dim_n: Tile columns of B and C (multiple of 8).
        mixed: bf16 A and C with bfp16 B.
    """
    if not _arch_traits().bfp16:
        raise NotImplementedError(
            "mm_bfp: bfp16ebs8 is an AIE2P type; select an NPU2 device"
        )
    for name, v in (("dim_m", dim_m), ("dim_k", dim_k), ("dim_n", dim_n)):
        if v <= 0 or v % 8:
            raise ValueError(
                f"mm_bfp: {name} must be a positive multiple of 8, got {v}"
            )
    flags = [f"-DDIM_M={dim_m}", f"-DDIM_K={dim_k}", f"-DDIM_N={dim_n}"]
    b_ty = np.ndarray[(dim_k * dim_n // 8,), np.dtype[v8bfp16ebs8]]
    if mixed:
        source = _kernel_source("linalg/mm_bfp_mixed.cc")
        a_ty = np.ndarray[(dim_m * dim_k,), np.dtype[bfloat16]]
        c_ty = np.ndarray[(dim_m * dim_n,), np.dtype[bfloat16]]
        symbol = "matmul_vectorized_different_datatypes"
    else:
        source = _kernel_source("linalg/mm_bfp.cc")
        a_ty = np.ndarray[(dim_m * dim_k // 8,), np.dtype[v8bfp16ebs8]]
        c_ty = np.ndarray[(dim_m * dim_n // 8,), np.dtype[v8bfp16ebs8]]
        symbol = "matmul_vectorized_bfp16"
    streams = mm_stream_dims(dim_m, dim_k, dim_n, _BFP_MAC_DIMS)
    # The kernel reads B transposed (8x8 sub-tiles of B^T), so the host B
    # buffer is B^T (N, K), as the block_datatypes examples tile it; the
    # bfp16ebs8 operands are pre-shuffled on the host and stream as stored.
    layouts = (
        (
            _tile_layout((dim_m, dim_k), streams.A, block=_BFP_MAC_DIMS[:2])
            if mixed
            else _block_layout((dim_m, dim_k))
        ),
        _block_layout((dim_k, dim_n), axes=(1, 0)),
        (
            _tile_layout(
                (dim_m, dim_n), streams.C, inverse=True, block=_BFP_MAC_DIMS[::2]
            )
            if mixed
            else _block_layout((dim_m, dim_n))
        ),
    )
    return _make_extern(
        symbol,
        source,
        [a_ty, b_ty, c_ty],
        compile_flags=flags + ["-DMATMUL_ONLY"],
        cls=MatrixKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            layouts=layouts,
            setup=conv_even,
            roles=(In, In, InOut),
            reference=partial(
                mm_bfp_tile_ref, dim_m=dim_m, dim_k=dim_k, dim_n=dim_n, mixed=mixed
            ),
            initializers=((2, _zero_output),),
            acc_dtype=np.float32,
            reduction=dim_k,
            tolerance=_BFP_MIXED_TOLERANCE if mixed else _BFP_TOLERANCE,
            ops_per_call=2 * dim_m * dim_k * dim_n,
        ),
    )

mm_bfp_shuffle

mm_bfp_shuffle(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    *,
    in_shape: tuple | None = None,
    out_shape: tuple | None = None,
    unshuffle: bool = False
) -> ExternalFunction

Scalar shuffle of a bfp16ebs8 tile into (or out of) the mmul block layout (aie2p).

scalar_shuffle(in, out, tile_width, tile_height, unshuffle) from mm_bfp.cc; the in-core-shuffle block-datatype examples run it before mm_bfp. By default the input tile is sized like mm_bfp's A and the output like its C; in_shape / out_shape (in v8bfp16ebs8 blocks) override that, e.g. (dim_m, dim_k // 8) twice to shuffle an A tile in place, matching the ObjectFifo types a design already uses.

Parameters:

Name	Type	Description	Default
dim_m	int	Tile rows (multiple of 8).	64
dim_k	int	A's tile columns (multiple of 8).	64
dim_n	int	C's tile columns (multiple of 8).	64
in_shape	tuple | None	Input tile shape in blocks; default (dim_m * dim_k // 8,).	None
out_shape	tuple | None	Output tile shape in blocks; default (dim_m * dim_n // 8,).	None
unshuffle	bool	Validate the other direction, out of the block layout back to row-major; the harness passes it as the call's last argument.	False

Source code in python/iron/kernels/linalg.py

def mm_bfp_shuffle(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    *,
    in_shape: tuple | None = None,
    out_shape: tuple | None = None,
    unshuffle: bool = False,
) -> ExternalFunction:
    """Scalar shuffle of a bfp16ebs8 tile into (or out of) the mmul block layout (aie2p).

    ``scalar_shuffle(in, out, tile_width, tile_height, unshuffle)`` from
    ``mm_bfp.cc``; the in-core-shuffle block-datatype examples run it before
    ``mm_bfp``. By default the input tile is sized like ``mm_bfp``'s A and
    the output like its C; ``in_shape`` / ``out_shape`` (in ``v8bfp16ebs8``
    blocks) override that, e.g. ``(dim_m, dim_k // 8)`` twice to shuffle an
    A tile in place, matching the ObjectFifo types a design already uses.

    Args:
        dim_m: Tile rows (multiple of 8).
        dim_k: A's tile columns (multiple of 8).
        dim_n: C's tile columns (multiple of 8).
        in_shape: Input tile shape in blocks; default ``(dim_m * dim_k // 8,)``.
        out_shape: Output tile shape in blocks; default ``(dim_m * dim_n // 8,)``.
        unshuffle: Validate the other direction, out of the block layout
            back to row-major; the harness passes it as the call's last
            argument.
    """
    if not _arch_traits().bfp16:
        raise NotImplementedError(
            "mm_bfp_shuffle: bfp16ebs8 is an AIE2P type; select an NPU2 device"
        )
    flags = [f"-DDIM_M={dim_m}", f"-DDIM_K={dim_k}", f"-DDIM_N={dim_n}"]
    in_shape = tuple(in_shape) if in_shape else (dim_m * dim_k // 8,)
    out_shape = tuple(out_shape) if out_shape else (dim_m * dim_n // 8,)
    in_ty = np.ndarray[in_shape, np.dtype[v8bfp16ebs8]]
    out_ty = np.ndarray[out_shape, np.dtype[v8bfp16ebs8]]
    from aie.utils import bfp

    logical_shape = (dim_m, dim_k)
    plain = TensorLayout(
        logical_shape,
        pack=lambda x: bfp.encode(x).reshape(len(x), -1),
        unpack=lambda x: bfp.decode(x).reshape(len(x), *logical_shape),
    )
    blocked = _block_layout(logical_shape)
    extern = _make_extern(
        "scalar_shuffle",
        _kernel_source("linalg/mm_bfp.cc"),
        [in_ty, out_ty, np.int16, np.int16, np.int16],
        compile_flags=flags + ["-DSHUFFLE_ONLY"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param, Param, Param),
            parameter_bindings=((2, dim_k), (3, dim_m), (4, int(unshuffle))),
            layouts=(
                blocked if unshuffle else plain,
                plain if unshuffle else blocked,
                None,
                None,
                None,
            ),
            reference=lambda x: bfp.quantize(x),
            tolerance=Tolerance.exact(note="byte permutation preserves encoded values"),
            ops_per_call=0,
            unsupported=(
                None
                if np.prod(in_shape) == np.prod(out_shape) == dim_m * dim_k // 8
                else "independent shuffle validation requires both buffers to hold "
                "exactly the dim_m by dim_k tile; custom unequal buffers need "
                "their enclosing design's runtime dimensions"
            ),
        ),
    )
    return extern

mha

mha(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    pv: bool = False,
    b_col_maj: bool = False,
    emulate_bf16_mmul_with_bfp16: bool = False,
) -> MatrixKernel

Flash-attention toolkit from aie_kernels/linalg/mha.cc.

One translation unit that includes softmax_aie2p.h and mm_aie2p.h and exports the symbols an attention dataflow composes over one micro-tile. The returned kernel is one of the toolkit's two matmuls, both accumulating into C, so it is a MatrixKernel judged like mm. By default that is the QK^T product matmul_bf16_bf16_wrapper, mm_aie2p.h's bf16 product on its 4x8x8 micro-tile behind an idx_buffer gate (the call runs when idx[0] <= idx[1]) bound here to [0, 0]. With pv it is the P*V product matmul_bf16_bf16_rowmaj, mha.cc's own expansion of the native 8x8x8 micro-tile and ungated. matmul_PV is that same product preceded by a row rescale, which needs the online softmax's running state and so is exercised by test_mha_e2e.py instead. mha.mac_dims(pv=...) answers either product's micro-tile without building a kernel, as mm.mac_dims does.

Bind the others from the same object with fn.object_file.bind(symbol, arg_types): matmul_bf16_bf16_wrapper_scalar, partial_softmax, matmul_PV, rescale_O, init_scale_buffer. It declares but does not define passThroughLine: take that from passthrough(dtype=np.int32), as IRON's MHA operator does.

Parameters:

Name	Type	Description	Default
dim_m	int	Rows of the micro-tile (multiple of 16).	64
dim_k	int	Depth of the micro-tile (multiple of 8).	64
dim_n	int	Columns of the micro-tile (multiple of 16).	64
pv	bool	If True return the P*V product instead of QK^T.	False
b_col_maj	bool	If True compile with -DB_COL_MAJ so the QK^T product consumes K as stored, (n, k), rather than transposed. matmul_bf16_bf16_rowmaj is row-major either way.	False
emulate_bf16_mmul_with_bfp16	bool	As for mm: both products use BFP16-based emulation, and QK^T's micro-tile becomes (8, 8, 8). Ignored on AIE2.	False

Source code in python/iron/kernels/linalg.py

def mha(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    pv: bool = False,
    b_col_maj: bool = False,
    emulate_bf16_mmul_with_bfp16: bool = False,
) -> MatrixKernel:
    """Flash-attention toolkit from ``aie_kernels/linalg/mha.cc``.

    One translation unit that includes ``softmax_aie2p.h`` and ``mm_aie2p.h`` and
    exports the symbols an attention dataflow composes over one micro-tile.
    The returned kernel is one of the toolkit's two matmuls, both
    accumulating into ``C``, so it is a
    [`MatrixKernel`][iron.kernels.linalg.MatrixKernel] judged like
    [`mm`][iron.kernels.linalg.mm]. By default that is the ``QK^T`` product
    ``matmul_bf16_bf16_wrapper``, ``mm_aie2p.h``'s bf16 product on its 4x8x8
    micro-tile behind an ``idx_buffer`` gate (the call runs when
    ``idx[0] <= idx[1]``) bound here to ``[0, 0]``. With ``pv`` it is the
    ``P*V`` product ``matmul_bf16_bf16_rowmaj``, mha.cc's own expansion of
    the native 8x8x8 micro-tile and ungated. ``matmul_PV`` is that same
    product preceded by a row rescale, which needs the online softmax's
    running state and so is exercised by ``test_mha_e2e.py`` instead.
    ``mha.mac_dims(pv=...)`` answers either product's micro-tile without
    building a kernel, as ``mm.mac_dims`` does.

    Bind the others from the same object with
    ``fn.object_file.bind(symbol, arg_types)``:
    ``matmul_bf16_bf16_wrapper_scalar``, ``partial_softmax``,
    ``matmul_PV``, ``rescale_O``, ``init_scale_buffer``. It declares but
    does not define ``passThroughLine``: take that from
    ``passthrough(dtype=np.int32)``, as IRON's MHA operator does.

    Args:
        dim_m: Rows of the micro-tile (multiple of 16).
        dim_k: Depth of the micro-tile (multiple of 8).
        dim_n: Columns of the micro-tile (multiple of 16).
        pv: If ``True`` return the ``P*V`` product instead of ``QK^T``.
        b_col_maj: If ``True`` compile with ``-DB_COL_MAJ`` so the ``QK^T``
            product consumes ``K`` as stored, ``(n, k)``, rather than
            transposed.  ``matmul_bf16_bf16_rowmaj`` is row-major either way.
        emulate_bf16_mmul_with_bfp16: As for [`mm`][iron.kernels.linalg.mm]:
            both products use BFP16-based emulation, and ``QK^T``'s
            micro-tile becomes (8, 8, 8).  Ignored on AIE2.
    """
    for name, v, mult in (
        ("dim_m", dim_m, 16),
        ("dim_k", dim_k, 8),
        ("dim_n", dim_n, 16),
    ):
        if v <= 0 or v % mult:
            raise ValueError(
                f"mha: {name} must be a positive multiple of {mult}, got {v}"
            )
    tile = np.ndarray[(dim_m * dim_n,), np.dtype[bfloat16]]
    a_ty = np.ndarray[(dim_m * dim_k,), np.dtype[bfloat16]]
    b_ty = np.ndarray[(dim_k * dim_n,), np.dtype[bfloat16]]
    idx = np.ndarray[(2,), np.dtype[np.int32]]
    # mha.cc only calls mm.cc's bf16 products, so build none of the others.
    flags = [
        f"-DDIM_M={dim_m}",
        f"-DDIM_K={dim_k}",
        f"-DDIM_N={dim_n}",
        "-Dbf16_bf16_ONLY",
    ]
    if b_col_maj:
        flags.append("-DB_COL_MAJ")
    emulate_bf16_mmul_with_bfp16 = emulate_bf16_mmul_with_bfp16 and _arch_traits().bfp16
    if emulate_bf16_mmul_with_bfp16:
        flags.append("-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16")
    # mha.cc includes mm_aie2p.h without C_COL_MAJ, and without B_COL_MAJ unless
    # b_col_maj. matmul_bf16_bf16_rowmaj is always row-major.
    b_col_maj = b_col_maj and not pv
    r, s, t = _MhaFactory.mac_dims(
        pv=pv, emulate_bf16_mmul_with_bfp16=emulate_bf16_mmul_with_bfp16
    )
    streams = mm_stream_dims(dim_m, dim_k, dim_n, (r, s, t), b_col_maj=b_col_maj)
    return _make_extern(
        "matmul_bf16_bf16_rowmaj" if pv else "matmul_bf16_bf16_wrapper",
        _kernel_source("linalg/mha.cc"),
        [a_ty, b_ty, tile] if pv else [a_ty, b_ty, tile, idx],
        compile_flags=flags,
        cls=MatrixKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            layouts=(
                _tile_layout((dim_m, dim_k), streams.A, block=(r, s)),
                _tile_layout(
                    (dim_k, dim_n),
                    streams.B,
                    axes=(1, 0) if b_col_maj else None,
                    block=(s, t),
                ),
                _tile_layout((dim_m, dim_n), streams.C, inverse=True, block=(r, t)),
                *(() if pv else (None,)),
            ),
            roles=(In, In, InOut) if pv else (In, In, InOut, Param),
            parameter_bindings=(() if pv else ((3, np.array([0, 0], np.int32)),)),
            reference=partial(mm_tile_ref, dim_m=dim_m, dim_k=dim_k, dim_n=dim_n),
            initializers=((2, _zero_output),),
            acc_dtype=np.float32,
            reduction=dim_k,
            tolerance=_linalg_tolerance(bfloat16),
            ops_per_call=2 * dim_m * dim_k * dim_n,
        ),
    )

mha_softmax

mha_softmax() -> ExternalFunction

One 64x64 block of mha.cc's online softmax, partial_softmax.

Writes the block's unnormalized weights P = exp2(A * s - m), with s = log2(e) / 8 and m each query row's running maximum, and updates the running state scale_buffer: [m, m, l, exp2(m_prev - m)], 64 rows each. The state is zeroed before every call, so each call is a first key block against a running maximum of 0; the carry across blocks is test_mha_e2e.py's. The kernel masks by overwriting A's masked entries in place.

Which block it is stays a runtime operand, as in the kernel: idx is (key block, query block) (equal on the causal diagonal, key block past query block skipped), and S_q_eff/S_kv_eff are the sequence lengths whose tails pad the block.

Source code in python/iron/kernels/linalg.py

def mha_softmax() -> ExternalFunction:
    """One 64x64 block of ``mha.cc``'s online softmax, ``partial_softmax``.

    Writes the block's unnormalized weights ``P = exp2(A * s - m)``, with
    ``s = log2(e) / 8`` and ``m`` each query row's running maximum, and
    updates the running state ``scale_buffer``: ``[m, m, l, exp2(m_prev -
    m)]``, 64 rows each. The state is zeroed before every call, so each call
    is a first key block against a running maximum of 0; the carry across
    blocks is ``test_mha_e2e.py``'s. The kernel masks by overwriting
    ``A``'s masked entries in place.

    Which block it is stays a runtime operand, as in the kernel: ``idx`` is
    ``(key block, query block)`` (equal on the causal diagonal, key block
    past query block skipped), and ``S_q_eff``/``S_kv_eff`` are the
    sequence lengths whose tails pad the block.
    """
    b = _MHA_BLOCK
    tile = np.ndarray[(b * b,), np.dtype[bfloat16]]
    state = np.ndarray[(4 * b,), np.dtype[bfloat16]]
    idx = np.ndarray[(2,), np.dtype[np.int32]]
    scale = float(bfloat16(np.log2(np.e) / np.sqrt(b)))
    return _make_extern(
        "partial_softmax",
        _kernel_source("linalg/mha.cc"),
        [tile, tile, state, idx, bfloat16, *([np.int32] * 4)],
        compile_flags=[f"-DDIM_M={b}", f"-DDIM_K={b}", f"-DDIM_N={b}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, InOut, *([Param] * 6)),
            parameter_bindings=((4, scale), (5, b), (6, b)),
            initializers=((2, _zero_output),),
            reference=partial(mha_softmax_ref, scale=scale),
            # The untuned loop on aie2, where a whole 64-lane row spills; the
            # tuned one fits the default.
            stack_bytes=(
                1376
                if _detect_arch() == "aie2" and _tuned_arch() == "portable"
                else None
            ),
            # aie::exp2<bfloat16> interpolates 2**frac linearly, overshooting
            # by up to 6.15%, and two bf16 roundings bring it to 6.98%
            # (test_mha_e2e.py's _RTOL_EXP2). This form's rtol multiplies
            # |a| + |b|, so half of that; the floor is test_mha_e2e's
            # 4 bf16 steps at 1, P's top. aie2 has no aie::exp2 and takes
            # exp2_bf16.h's cubic instead, 0.74% from exp2 at worst with the
            # bf16 store on npu1 (0.37% of |a| + |b|), so its rtol is 0.4% of
            # |a| + |b| and its floor only admits the underflow to 0.
            tolerance=(
                Tolerance.relative(
                    0.004,
                    2.0**-120,
                    note="exp2_bf16.h cubic, 0.74% worst element on npu1",
                )
                if not _arch_traits().native_exp2
                else Tolerance.relative(
                    0.035,
                    4 * 2.0**-7,
                    note="aie::exp2 interpolant envelope, 6.98% (test_mha_e2e.py)",
                )
            ),
        ),
    )

mha_softmax_ref

mha_softmax_ref(a, idx, s_q_eff, s_kv_eff, *, scale)

Numpy reference for mha_softmax: (P, scale_buffer).

True exp2 rather than the device's interpolant, from a zeroed running state. m is rounded to bf16 before P uses it, as the kernel stores it. A padded row keeps m = l = 0; a skipped or wholly padded block leaves the state untouched.

Source code in python/iron/kernels/linalg.py

def mha_softmax_ref(a, idx, s_q_eff, s_kv_eff, *, scale):
    """Numpy reference for [`mha_softmax`][iron.kernels.linalg.mha_softmax]: ``(P, scale_buffer)``.

    True ``exp2`` rather than the device's interpolant, from a zeroed
    running state. ``m`` is rounded to bf16 before ``P`` uses it, as the
    kernel stores it. A padded row keeps ``m = l = 0``; a skipped or wholly
    padded block leaves the state untouched.
    """
    b = _MHA_BLOCK
    kv, q = (int(i) for i in np.asarray(idx).ravel())
    rows, cols = np.indices((b, b))
    keep = (rows < s_q_eff - q * b) & (cols < s_kv_eff - kv * b)
    if kv == q:
        keep &= cols <= rows
    if kv > q or not keep.any():
        return np.zeros((len(a), b * b)), np.zeros((len(a), 4 * b))
    scaled = a.astype(np.float32).reshape(-1, b, b) * np.float32(scale)
    m = np.where(keep, scaled, -np.inf).max(axis=2)
    m = np.maximum(m.astype(bfloat16).astype(np.float32), 0)
    p = np.exp2(np.where(keep, scaled - m[..., None], -np.inf))
    state = np.concatenate([m, m, p.sum(axis=2), np.exp2(-m)], axis=1)
    return p.reshape(len(p), -1), state

prefill_fv

prefill_fv(head_dim: int = 512) -> ExternalFunction

Flash-attention prefill toolkit from aie_kernels/linalg/flash_attn_prefill.cc.

One translation unit per geometry, exporting the five steps an attention prefill dataflow composes over one query chunk. The returned kernel is the y += S*V step, mm.cc-style bf16 products on the native 8x8x8 micro-tile accumulating into a float32 y, so it is sampled and judged like mm. Its key-chunk index j is bound to 0.

Bind the others from the same object with fn.object_file.bind(symbol, arg_types): prefill_round_begin, prefill_qk_step, prefill_block_mid, prefill_epilogue.

One reference serves both geometries even though they decompose differently, because PrefillGeom<256>::reorder_s deinterleaves its 2x2 attn_qk output into the same block order the 1x1 geometry produces directly, and block_mid runs it before this step sees S. The geometries differ only in operand storage order, which the layouts carry.

head_dim picks the geometry, and two instantiations coexist in one design: differing -DPREFILL_HEAD_DIM gives each its own object and its own symbol prefix.

Parameters:

Name	Type	Description	Default
head_dim	int	512 for global attention, 256 for sliding-window.	512

Source code in python/iron/kernels/linalg.py

def prefill_fv(head_dim: int = 512) -> ExternalFunction:
    """Flash-attention prefill toolkit from ``aie_kernels/linalg/flash_attn_prefill.cc``.

    One translation unit per geometry, exporting the five steps an attention
    prefill dataflow composes over one query chunk. The returned kernel is the
    ``y += S*V`` step, ``mm.cc``-style bf16 products on the native 8x8x8
    micro-tile accumulating into a float32 y, so it is sampled and judged like
    [`mm`][iron.kernels.linalg.mm]. Its key-chunk index ``j`` is bound to 0.

    Bind the others from the same object with
    ``fn.object_file.bind(symbol, arg_types)``: ``prefill_round_begin``,
    ``prefill_qk_step``, ``prefill_block_mid``, ``prefill_epilogue``.

    One reference serves both geometries even though they decompose
    differently, because ``PrefillGeom<256>::reorder_s`` deinterleaves its 2x2
    ``attn_qk`` output into the same block order the 1x1 geometry produces
    directly, and ``block_mid`` runs it before this step sees S. The
    geometries differ only in operand *storage* order, which the layouts carry.

    ``head_dim`` picks the geometry, and two instantiations coexist in one
    design: differing ``-DPREFILL_HEAD_DIM`` gives each its own object and its
    own symbol prefix.

    Args:
        head_dim: 512 for global attention, 256 for sliding-window.
    """
    if head_dim not in _PREFILL_GEOM:
        raise ValueError(f"prefill_fv: head_dim must be 512 or 256, got {head_dim}")
    lq, lk, stack_bytes = _PREFILL_GEOM[head_dim]
    if _detect_arch() == "aie2":
        stack_bytes = _PREFILL_STACK_AIE2
    # flash_attn_prefill.h's MMUL is aie::mmul<8, 8, 8, bf16, bf16>, the native
    # bf16 micro-tile, not the (4, 8, 8) _MM_MAC_DIMS records for mm.cc.
    r = s = t = 8
    y_ty = np.ndarray[(lq * head_dim,), np.dtype[np.float32]]
    s_ty = np.ndarray[(lq * lk,), np.dtype[bfloat16]]
    v_ty = np.ndarray[(lk * head_dim,), np.dtype[bfloat16]]
    streams = mm_stream_dims(lq, lk, head_dim, (r, s, t))
    # attn_fv walks V n-block-outer, k-block-inner -- the transpose of mm.cc's
    # B block order, so this is built here rather than taken from streams.B. At
    # head_dim 512 LK is 8, making the k term degenerate, so only the 256
    # geometry can tell the two orders apart.
    v_dims = TensorAccessPattern.full((lk, head_dim)).tile((s, t)).permute((1, 0, 2, 3))
    return _make_extern(
        "prefill_fv_step",
        _kernel_source("linalg/flash_attn_prefill.cc"),
        [y_ty, s_ty, v_ty, np.int32],
        compile_flags=[f"-DPREFILL_HEAD_DIM={head_dim}"],
        cls=_ZeroInitializedKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            layouts=(
                _tile_layout((lq, head_dim), streams.C, inverse=True, block=(r, t)),
                _tile_layout((lq, lk), streams.A, block=(r, s)),
                _tile_layout((lk, head_dim), v_dims, block=(s, t)),
                None,
            ),
            roles=(InOut, In, In, Param),
            parameter_bindings=((3, 0),),
            initializers=((0, partial(_zero_output, index=0)),),
            reference=partial(prefill_fv_ref, dim_m=lq, dim_k=lk, dim_n=head_dim),
            acc_dtype=np.float32,
            reduction=lk,
            stack_bytes=stack_bytes,
            # Derived, not inherited: _linalg_tolerance(bfloat16)'s 0.05/0.5 is
            # for a kernel that narrows C back to bf16, and y here is float32.
            # bf16 mantissas are 8 bits, so every product is exact in f32
            # (8 + 8 < 24) and the only error against a float64 reference is
            # f32 summation order over lk terms, bounded by lk * 2**-24 ~ 1e-6.
            # The margin below that is ~10x, not the ~50000x inheriting would
            # have given, which would hide nearly any real bug. Outside aie2p's
            # tuned branch the mmul sums in a different order, and on large inputs whose
            # products cancel that order error exceeds any relative bound, so
            # it is judged against the order bound itself.
            tolerance=(
                Tolerance.bounded(
                    partial(_prefill_fv_bound, dim_m=lq, dim_k=lk, dim_n=head_dim),
                    note="f32 summation order over lk exact bf16 products, in "
                    "the kernel and the reference; measured worst 0.12 of it "
                    "on npu1",
                )
                if _tuned_arch() != "aie2p"
                else Tolerance.relative(
                    1e-5,
                    1e-5,
                    note="f32 accumulation order over lk<=16 exact bf16 "
                    "products: lk * 2**-24 ~ 1e-6; verified on npu2 at this bound",
                )
            ),
            ops_per_call=2 * lq * lk * head_dim,
        ),
    )

prefill_fv_ref

prefill_fv_ref(s, v, *, dim_m: int, dim_k: int, dim_n: int)

Numpy reference for prefill_fv: one y += S @ V key chunk.

y is not an argument: it is InOut, so the contract's initializer zeroes it before each independent call and the reference computes the whole product.

Accumulated in float32, not the float64 of mm_tile_ref. The kernel's y is float32 and its MMUL accumulates there, so a float64 reference is more precise than the kernel is defined to be and the comparison measures that gap rather than correctness. The difference is invisible when the product narrows back to bf16, as every other matmul here does, and decisive when it does not: on inputs near 1e4 the eight products reach 1e8 and cancel, leaving a result near zero whose float64 and float32 sums differ by a whole float32 ULP of the intermediate. Products themselves are exact either way -- bf16 carries 8 mantissa bits and 8 + 8 < 24.

Source code in python/iron/kernels/linalg.py

def prefill_fv_ref(s, v, *, dim_m: int, dim_k: int, dim_n: int):
    """Numpy reference for [`prefill_fv`][iron.kernels.linalg.prefill_fv]: one ``y += S @ V`` key chunk.

    ``y`` is not an argument: it is ``InOut``, so the contract's initializer
    zeroes it before each independent call and the reference computes the
    whole product.

    Accumulated in **float32**, not the float64 of
    [`mm_tile_ref`][iron.kernels.linalg.mm_tile_ref]. The kernel's y is
    float32 and its MMUL accumulates there, so a float64 reference is more
    precise than the kernel is defined to be and the comparison measures that
    gap rather than correctness. The difference is invisible when the product
    narrows back to bf16, as every other matmul here does, and decisive when
    it does not: on inputs near 1e4 the eight products reach 1e8 and cancel,
    leaving a result near zero whose float64 and float32 sums differ by a
    whole float32 ULP of the intermediate. Products themselves are exact
    either way -- bf16 carries 8 mantissa bits and 8 + 8 < 24.
    """
    s = np.asarray(s).reshape(-1, dim_m, dim_k).astype(np.float32)
    v = np.asarray(v).reshape(-1, dim_k, dim_n).astype(np.float32)
    return (s @ v).reshape(len(s), dim_m * dim_n)

mm_bfp_shuffle_ref

mm_bfp_shuffle_ref(
    tile, tile_width, tile_height, unshuffle
)

Numpy reference for mm_bfp_shuffle: aie.utils.bfp.shuffle of one tile's bytes.

tile is the encoded (tile_height, tile_width) tile as bytes (tile_width in values); returns the reordered bytes.

Source code in python/iron/kernels/linalg.py

def mm_bfp_shuffle_ref(tile, tile_width, tile_height, unshuffle):
    """Numpy reference for [`mm_bfp_shuffle`][iron.kernels.linalg.mm_bfp_shuffle]: ``aie.utils.bfp.shuffle`` of one tile's bytes.

    ``tile`` is the encoded ``(tile_height, tile_width)`` tile as bytes
    (``tile_width`` in values); returns the reordered bytes.
    """
    from aie.utils import bfp

    w, h = int(tile_width), int(tile_height)
    return bfp.shuffle(tile, w, h, w, h, unshuffle=bool(unshuffle)).ravel()

cascade_mm

cascade_mm(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    input_dtype: type = int16,
    output_dtype: type = int16,
    use_chess: bool = False,
) -> _CascadeMatrixKernel

Build the GET half of a cascade matrix multiply: C += A * B + cascade.

cascade_mm.cc emits all three cascade variants (get_only, put_only, put_get) in one object. This binds get_only (also available as .get_only), with .put_only and .put_get siblings; cascade_mm_put is the PUT half that feeds it, and put_get serves longer chains: fn.object_file.bind("matmul_scalar_cascade_put_get_<dtype>", fn.arg_types()). .zero initializes accumulators using independent kernels.zero. The pair is a two-tile design, which the generic builder does not run; the device test builds and judges it (test/python/npu/test_kernels_e2e.py). On AIE2 each output's partial sum crosses the cascade as one 32-bit lane, a float chain as the float's bits. On AIE2P, when dim_m and dim_k are multiples of 8 and dim_n of 16, the whole accumulator crosses the cascade; other shapes run the AIE2 kernel. Either way a float chain sums in float and rounds once, to nearest even, at the GET half.

Parameters:

Name	Type	Description	Default
dim_m	int	Number of rows of A / C.	64
dim_k	int	Number of columns of A / rows of B.	64
dim_n	int	Number of columns of B / C.	64
input_dtype	type	Input element type.	int16
output_dtype	type	Output element type.	int16
use_chess	bool	If True build the .o with xchesscc_wrapper instead of Peano.	False

Raises:

Type	Description
ValueError	When the dtype combination is not supported.

Source code in python/iron/kernels/linalg.py

def cascade_mm(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    input_dtype: type = np.int16,
    output_dtype: type = np.int16,
    use_chess: bool = False,
) -> _CascadeMatrixKernel:
    r"""Build the GET half of a cascade matrix multiply: ``C += A * B + cascade``.

    cascade_mm.cc emits all three cascade variants (``get_only``,
    ``put_only``, ``put_get``) in one object. This binds ``get_only`` (also
    available as ``.get_only``), with ``.put_only`` and ``.put_get`` siblings;
    [`cascade_mm_put`][iron.kernels.linalg.cascade_mm_put] is the PUT half
    that feeds it, and ``put_get`` serves longer chains:
    ``fn.object_file.bind("matmul_scalar_cascade_put_get_<dtype>", fn.arg_types())``.
    ``.zero`` initializes accumulators using independent ``kernels.zero``.
    The pair is a two-tile
    design, which the generic builder does not run; the device test builds
    and judges it (``test/python/npu/test_kernels_e2e.py``). On AIE2 each
    output's partial sum crosses the cascade as one 32-bit lane, a float
    chain as the float's bits. On AIE2P, when ``dim_m`` and ``dim_k`` are
    multiples of 8 and ``dim_n`` of 16, the whole accumulator crosses the
    cascade; other shapes run the AIE2 kernel. Either way a float chain
    sums in float and rounds once, to nearest even, at the GET half.

    Args:
        dim_m: Number of rows of A / C.
        dim_k: Number of columns of A / rows of B.
        dim_n: Number of columns of B / C.
        input_dtype: Input element type.
        output_dtype: Output element type.
        use_chess: If ``True`` build the .o with ``xchesscc_wrapper``
            instead of Peano.

    Raises:
        ValueError: When the dtype combination is not supported.
    """
    key = (input_dtype, output_dtype)
    if key not in _CASCADE_COMBOS:
        raise ValueError(
            f"cascade_mm(): unsupported (input_dtype, output_dtype) = {key}. Supported: {list(_CASCADE_COMBOS.keys())}"
        )

    suffix = _CASCADE_COMBOS[key]
    a_ty = np.ndarray[(dim_m * dim_k,), np.dtype[input_dtype]]
    b_ty = np.ndarray[(dim_k * dim_n,), np.dtype[input_dtype]]
    c_ty = np.ndarray[(dim_m * dim_n,), np.dtype[output_dtype]]
    r, s, t = _CascadeMatMulFactory.mac_dims(input_dtype, output_dtype)
    extern = _make_extern(
        f"matmul_scalar_cascade_get_only_{suffix}",
        _kernel_source("linalg/cascade_mm.cc"),
        [a_ty, b_ty, c_ty],
        compile_flags=[
            f"-DDIM_M={dim_m}",
            f"-DDIM_K={dim_k}",
            f"-DDIM_N={dim_n}",
        ],
        use_chess=use_chess,
        cls=_CascadeMatrixKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, InOut),
            # Row-major operands on both targets, nothing streamed
            # transformed, so the layouts carry the 1x1x1 blocking and no
            # stream (see _CASCADE_MM_SCALAR_DIMS).
            layouts=(
                _tile_layout((dim_m, dim_k), block=(r, s)),
                _tile_layout((dim_k, dim_n), block=(s, t)),
                _tile_layout((dim_m, dim_n), block=(r, t)),
            ),
            unsupported=(
                "the GET half of a cascade pair: one input arrives on the "
                "cascade stream; the device test builds and judges the pair"
            ),
            initializers=((2, _zero_output),),
            acc_dtype=mm_acc_dtype(input_dtype),
            reduction=2 * dim_k,
            tolerance=_linalg_tolerance(input_dtype),
            ops_per_call=2 * dim_m * (2 * dim_k) * dim_n,
        ),
    )
    extern.get_only = extern
    extern.put_only = extern.object_file.bind(
        f"matmul_scalar_cascade_put_only_{suffix}", [a_ty, b_ty, c_ty]
    )
    extern.put_get = extern.object_file.bind(
        f"matmul_scalar_cascade_put_get_{suffix}", [a_ty, b_ty, c_ty]
    )
    return extern

cascade_mm_put

cascade_mm_put(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    input_dtype: type = int16,
    output_dtype: type = int16,
    use_chess: bool = False,
) -> MatrixKernel

Build the PUT half of cascade_mm: A * B onto the cascade stream.

Same object and arguments as the GET half. put_only never touches its third argument (the ABI just mirrors get_only), so the contract binds it to zeros. Its result leaves on the cascade stream, so it is judged with its GET half by the device test, not by the generic builder.

Source code in python/iron/kernels/linalg.py

def cascade_mm_put(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    input_dtype: type = np.int16,
    output_dtype: type = np.int16,
    use_chess: bool = False,
) -> MatrixKernel:
    """Build the PUT half of [`cascade_mm`][iron.kernels.linalg.cascade_mm]: ``A * B`` onto the cascade stream.

    Same object and arguments as the GET half. ``put_only`` never touches
    its third argument (the ABI just mirrors ``get_only``), so the contract
    binds it to zeros. Its result leaves on the cascade stream, so it is
    judged with its GET half by the device test, not by the generic builder.
    """
    key = (input_dtype, output_dtype)
    if key not in _CASCADE_COMBOS:
        raise ValueError(
            f"cascade_mm_put(): unsupported (input_dtype, output_dtype) = {key}. Supported: {list(_CASCADE_COMBOS.keys())}"
        )
    suffix = _CASCADE_COMBOS[key]
    a_ty = np.ndarray[(dim_m * dim_k,), np.dtype[input_dtype]]
    b_ty = np.ndarray[(dim_k * dim_n,), np.dtype[input_dtype]]
    c_ty = np.ndarray[(dim_m * dim_n,), np.dtype[output_dtype]]
    r, s, t = _CascadeMatMulFactory.mac_dims(input_dtype, output_dtype)
    return _make_extern(
        f"matmul_scalar_cascade_put_only_{suffix}",
        _kernel_source("linalg/cascade_mm.cc"),
        [a_ty, b_ty, c_ty],
        compile_flags=[
            f"-DDIM_M={dim_m}",
            f"-DDIM_K={dim_k}",
            f"-DDIM_N={dim_n}",
        ],
        use_chess=use_chess,
        cls=MatrixKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Param),
            parameter_bindings=((2, np.zeros(dim_m * dim_n, dtype=output_dtype)),),
            layouts=(
                _tile_layout((dim_m, dim_k), block=(r, s)),
                _tile_layout((dim_k, dim_n), block=(s, t)),
                None,
            ),
            unsupported="the PUT half of a cascade pair: its result leaves on the cascade stream",
            acc_dtype=mm_acc_dtype(input_dtype),
            reduction=dim_k,
            ops_per_call=2 * dim_m * dim_k * dim_n,
        ),
    )

fused_mm computes one bf16 matrix-multiply tile on a single core, with an f32 reduction and an optional activation and clamp fused into the drain.

Single-tile composition of the fused GEMM init, reduction and drain ABI.

fused_mm

fused_mm(
    *,
    dim_m=32,
    dim_k=32,
    dim_n=16,
    band_m=16,
    chunk_k=16,
    out_chunk=64,
    epilogue="none",
    clamp=None,
    bfp16_b=False,
    epilogue_modes=None,
    rounding="conv_even",
    gelu="fp32",
    mmul_shape=None,
    c_depth=2,
    step_markers=False
) -> _FusedMMKernel

Compute one bf16 A @ B tile, with an f32 reduction and fused epilogue.

This bounded composition holds its operands and accumulator on one core; it is not the streaming whole-matrix operator. A and B use bf16 storage on both architectures unless bfp16_b is set. band_m and chunk_k subdivide the reduction; out_chunk subdivides the drain with a depth of two. The contract supplies the blocked storage layouts, including B's column-major ordering of row-major microblocks.

epilogue is none, gelu (the kernel's sigmoid approximation, not the tanh GELU curve), silu or sigmoid. An optional finite (min, max) clamp follows the activation, before the bf16 conversion.

bfp16_b (aie2p only) selects the prepacked-B form amd/IRON's flm GEMM builds: B arrives as bfp16ebs8 blocks the host packs once, the mmul is 8x8x8, and the core converts A to bfp16 itself. The host rounds B in the rounding mode, as IRON's pack_b does, and the reference multiplies both operands as the core sees them.

epilogue_modes lists the activations compiled in, for a caller that selects one at run time through mm_fused_epilogue_chunk's mode argument. It defaults to (epilogue,). epilogue is the mode fused_mm_tile and the reference use, and must be in the list.

rounding is the core's rounding mode for every f32 -> bf16 and bf16 -> bfp16 conversion: conv_even or floor, the mode a core powers up in.

gelu selects how the gelu mode computes. fp32 applies it to the f32 accumulator and rounds once. bf16_steps rounds the accumulator to bf16 and rounds again after each step of x * sigmoid(1.702x), as FastFlowLM's shipped mm overlay does. With rounding="floor" it reproduces that overlay bit for bit on AIE2P.

mmul_shape overrides the (r, s, t) of the mmul. c_depth is the depth of the C fifo, which a caller's drain loop unrolls by. step_markers brackets each init, k step and drain chunk with event markers, for a caller that calls those entry points itself. Otherwise one pair brackets the whole fused_mm_tile call.

The returned function carries every entry point of its object as an attribute named by the C symbol: fn.mm_fused_acc_init, fn.mm_fused_k_step, fn.mm_fused_epilogue_chunk and fn.fused_mm_tile (fn itself).

Source code in python/iron/kernels/fused.py

def fused_mm(
    *,
    dim_m=32,
    dim_k=32,
    dim_n=16,
    band_m=16,
    chunk_k=16,
    out_chunk=64,
    epilogue="none",
    clamp=None,
    bfp16_b=False,
    epilogue_modes=None,
    rounding="conv_even",
    gelu="fp32",
    mmul_shape=None,
    c_depth=2,
    step_markers=False,
) -> _FusedMMKernel:
    """Compute one bf16 ``A @ B`` tile, with an f32 reduction and fused epilogue.

    This bounded composition holds its operands and accumulator on one core;
    it is not the streaming whole-matrix operator. A and B use bf16 storage
    on both architectures unless ``bfp16_b`` is set.
    ``band_m`` and ``chunk_k`` subdivide the reduction; ``out_chunk`` subdivides
    the drain with a depth of two. The contract supplies the blocked storage
    layouts, including B's column-major ordering of row-major microblocks.

    ``epilogue`` is ``none``, ``gelu`` (the kernel's sigmoid approximation,
    not the tanh GELU curve), ``silu`` or ``sigmoid``. An optional finite
    ``(min, max)`` clamp follows the activation, before the bf16 conversion.

    ``bfp16_b`` (aie2p only) selects the prepacked-B form amd/IRON's flm GEMM
    builds: B arrives as bfp16ebs8 blocks the host packs once, the mmul is
    8x8x8, and the core converts A to bfp16 itself. The host rounds B in the
    ``rounding`` mode, as IRON's ``pack_b`` does, and the reference
    multiplies both operands as the core sees them.

    ``epilogue_modes`` lists the activations compiled in, for a caller that
    selects one at run time through ``mm_fused_epilogue_chunk``'s mode
    argument. It defaults to ``(epilogue,)``. ``epilogue`` is the mode
    ``fused_mm_tile`` and the reference use, and must be in the list.

    ``rounding`` is the core's rounding mode for every f32 -> bf16 and
    bf16 -> bfp16 conversion: ``conv_even`` or ``floor``, the mode a core
    powers up in.

    ``gelu`` selects how the gelu mode computes. ``fp32`` applies it to the
    f32 accumulator and rounds once. ``bf16_steps`` rounds the accumulator to
    bf16 and rounds again after each step of ``x * sigmoid(1.702x)``, as
    FastFlowLM's shipped mm overlay does. With ``rounding="floor"`` it
    reproduces that overlay bit for bit on AIE2P.

    ``mmul_shape`` overrides the ``(r, s, t)`` of the mmul. ``c_depth`` is
    the depth of the C fifo, which a caller's drain loop unrolls by.
    ``step_markers`` brackets each init, k step and drain chunk with event
    markers, for a caller that calls those entry points itself. Otherwise one
    pair brackets the whole ``fused_mm_tile`` call.

    The returned function carries every entry point of its object as an
    attribute named by the C symbol: ``fn.mm_fused_acc_init``,
    ``fn.mm_fused_k_step``, ``fn.mm_fused_epilogue_chunk`` and
    ``fn.fused_mm_tile`` (``fn`` itself).
    """
    arch = _detect_arch()
    device = _device()
    if bfp16_b and not ARCH_TRAITS[arch].bfp16:
        raise ValueError("fused_mm: bfp16_b needs aie2p; bfp16ebs8 is an AIE2P type")
    if bfp16_b:
        r, s, t = 8, 8, 8
    else:
        r, s, t = (4, 8, 4) if arch == "aie2" else (4, 8, 8)
    if mmul_shape is not None:
        if bfp16_b and tuple(mmul_shape) != (8, 8, 8):
            raise ValueError("fused_mm: bfp16_b needs mmul_shape (8, 8, 8)")
        r, s, t = mmul_shape
    dims = (dim_m, dim_k, dim_n, band_m, chunk_k, out_chunk, c_depth, r, s, t)
    if any(not isinstance(d, int) or isinstance(d, bool) or d <= 0 for d in dims):
        raise ValueError("fused_mm dimensions must be positive integers")
    if (
        dim_m % band_m
        or band_m % (2 * r)
        or dim_n % (2 * t)
        or dim_k % chunk_k
        or chunk_k % s
        or out_chunk % 16
        or (dim_m * dim_n) % (c_depth * out_chunk)
    ):
        raise ValueError("fused_mm dimensions violate band, mmul or drain divisibility")
    modes = _MODES
    if epilogue not in modes:
        raise ValueError(f"unknown fused_mm epilogue: {epilogue}")
    if epilogue_modes is None:
        epilogue_modes = (epilogue,)
    if any(m not in modes for m in epilogue_modes):
        raise ValueError(f"unknown fused_mm epilogue in {epilogue_modes}")
    if epilogue not in epilogue_modes:
        raise ValueError(f"epilogue {epilogue} is not in epilogue_modes")
    if rounding not in ("conv_even", "floor"):
        raise ValueError(f"fused_mm: rounding must be conv_even or floor: {rounding}")
    if gelu not in ("fp32", "bf16_steps"):
        raise ValueError(f"fused_mm: gelu must be fp32 or bf16_steps: {gelu}")
    steps = gelu == "bf16_steps" and epilogue == "gelu"
    if clamp is not None:
        if len(clamp) != 2 or not np.isfinite(clamp).all() or clamp[0] > clamp[1]:
            raise ValueError("clamp must be a finite (min, max) pair with min <= max")
        clamp = tuple(float(v) for v in clamp)

    def pack_a(a):
        return (
            a.reshape(-1, dim_m // r, r, dim_k // chunk_k, chunk_k // s, s)
            .transpose(0, 3, 1, 4, 2, 5)
            .reshape(len(a), -1)
        )

    def unpack_a(a):
        return (
            a.reshape(-1, dim_k // chunk_k, dim_m // r, chunk_k // s, r, s)
            .transpose(0, 2, 4, 1, 3, 5)
            .reshape(-1, dim_m, dim_k)
        )

    def pack_b(b):
        return (
            b.reshape(-1, dim_k // chunk_k, chunk_k // s, s, dim_n // t, t)
            .transpose(0, 1, 4, 2, 3, 5)
            .reshape(len(b), -1)
        )

    def unpack_b(b):
        return (
            b.reshape(-1, dim_k // chunk_k, dim_n // t, chunk_k // s, s, t)
            .transpose(0, 1, 3, 4, 2, 5)
            .reshape(-1, dim_k, dim_n)
        )

    # The bfp16 blocks are t-major (block (i, j) holds B^T), because
    # mac_8x8_8x8T takes B transposed; that also makes each shared exponent
    # span 8 consecutive k of one column, the grouping the mac expects. The
    # blocks round as amd/IRON's weight packer does.
    def pack_b_bfp(b):
        blocks = (
            b.reshape(-1, dim_k // chunk_k, chunk_k // s, s, dim_n // t, t)
            .transpose(0, 1, 4, 2, 5, 3)
            .reshape(len(b), -1)
        )
        return bfp.encode(blocks, rounding=rounding)

    def unpack_b_bfp(b):
        return (
            bfp.decode(b)
            .reshape(-1, dim_k // chunk_k, dim_n // t, chunk_k // s, t, s)
            .transpose(0, 1, 3, 5, 2, 4)
            .reshape(-1, dim_k, dim_n)
        )

    def pack_c(c):
        return (
            c.reshape(-1, dim_m // r, r, dim_n // t, t)
            .transpose(0, 1, 3, 2, 4)
            .reshape(len(c), -1)
        )

    def unpack_c(c):
        return (
            c.reshape(-1, dim_m // r, dim_n // t, r, t)
            .transpose(0, 1, 3, 2, 4)
            .reshape(-1, dim_m, dim_n)
        )

    def operands(a, b):
        a = a.reshape(-1, dim_m, dim_k).astype(np.float32)
        b = b.reshape(-1, dim_k, dim_n).astype(np.float32)
        if bfp16_b:
            # The core converts A itself; B is whatever the host packed.
            a = bfp.quantize(a, rounding=rounding)
            b = unpack_b_bfp(pack_b_bfp(b))
        return a.astype(np.float64), b.astype(np.float64)

    # float64 throughout: in float32, 1 + tanh cancels for large negative
    # inputs and the product's rounding depends on numpy's summation order.
    # gelu_bf16_steps_vec with an exact tanh.
    def gelu_steps(c):
        x = _bf16_round(c, rounding)
        y = _bf16_round(x * _GELU_SCALE_BF16, rounding)
        sig = _bf16_round(np.tanh(y * 0.5) + 1, rounding) * 0.5
        return _bf16_round(x * sig, rounding)

    def activate(c):
        if steps:
            c = gelu_steps(c)
        elif epilogue != "none":
            x = c * 1.702 if epilogue == "gelu" else c
            sigmoid = (np.tanh(x * 0.5) + 1) * 0.5
            c = sigmoid if epilogue == "sigmoid" else c * sigmoid
        if clamp is not None:
            c = np.clip(c, clamp[0], clamp[1])
        return c

    def reference(a, b):
        a, b = operands(a, b)
        c = activate(a @ b)
        return c.reshape(len(c), -1)

    # Each step of gelu_steps rounds in the core and in the reference. A
    # rounding moves its result by under one ulp of the larger operand, so
    # the step's input error plus that ulp bounds its output error. tanh is
    # 1-Lipschitz and vtanh adds its own error.
    def gelu_steps_bound(c, acc_err):
        def ulp(v, dv):
            return _bf16_ulp(np.abs(v) + dv)

        x = _bf16_round(c, rounding)
        dx = acc_err + ulp(x, acc_err)
        y = _bf16_round(x * _GELU_SCALE_BF16, rounding)
        dy = _GELU_SCALE_BF16 * dx + ulp(y, _GELU_SCALE_BF16 * dx)
        # vtanh's error is not monotone, so take it at both ends of the
        # argument's interval and at its centre.
        vtanh = np.maximum.reduce(
            [_vtanh_error(0.5 * (np.abs(y) + d)) for d in (-dy, 0, dy)]
        )
        dt = 0.5 * dy + vtanh
        sig = _bf16_round(np.tanh(y * 0.5) + 1, rounding)
        dsig = 0.5 * (dt + ulp(sig, dt))
        out = x * sig * 0.5
        dout = 0.5 * sig * dx + np.abs(x) * dsig + dx * dsig
        return dout + ulp(out, dout) + 2.0**-126

    # Floor rounding errs by up to one f32 ulp per accumulation, not half.
    acc_ulps = 2.0**-23 if rounding == "floor" else 2.0**-24

    # The core's f32 sums are off by at most dim_k f32 ulps of sum |a*b|,
    # and the activation's slope (under 1.1) carries that through. On
    # AIE2P tanh is vtanh (see activation._vtanh_error): sigmoid scales
    # its error by 1/2 and x * sigmoid(u) by |x|/2. Every other step is an
    # exact bf16 product or an f32 add, and one output ulp covers the store.
    def error_bound(a, b):
        a, b = operands(a, b)
        c = a @ b
        if steps:
            acc_err = dim_k * acc_ulps * (np.abs(a) @ np.abs(b))
            return gelu_steps_bound(c, acc_err).reshape(len(c), -1)
        err = dim_k * acc_ulps * (np.abs(a) @ np.abs(b)) * _SLOPE[epilogue]
        if epilogue == "sigmoid":
            err = err + 0.5 * _vtanh_error(0.5 * c)
        elif epilogue != "none":
            u = 0.851 * c if epilogue == "gelu" else 0.5 * c
            err = err + 0.5 * np.abs(c) * _vtanh_error(u)
        err = err + _bf16_ulp(activate(c)) + 2.0**-126
        return err.reshape(len(c), -1)

    flags = {
        "TILE_M": dim_m,
        "TILE_MA": band_m,
        "TILE_K": dim_k,
        "TILE_N": dim_n,
        "CT_K": chunk_k,
        "R": r,
        "S": s,
        "T": t,
        "OUT_CHUNK": out_chunk,
        "C_DEPTH": c_depth,
        # The kernel selects the activation at runtime; the mask only decides
        # which bodies are compiled in, so each mode costs program memory.
        "EPILOGUE_MODE_MASK": sum(1 << modes[m] for m in set(epilogue_modes)),
    }
    compile_flags = [
        *(["-DROUND_CONV_EVEN"] if rounding == "conv_even" else []),
        *(f"-DMM_FUSED_{name}={value}" for name, value in flags.items()),
        *_portable_flags(),
    ]
    if bfp16_b:
        # amd/IRON's pair: the first selects aie_api's bfp16-emulated bf16
        # mmul, the second the prepacked B storage.
        compile_flags += [
            "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16",
            "-DMM_FUSED_BFP16_B",
        ]
    if gelu == "bf16_steps":
        compile_flags.append("-DMM_FUSED_GELU_BF16_STEPS")
    if step_markers:
        compile_flags.append("-DMM_FUSED_STEP_MARKERS")
    # An absent clamp is (-inf, +inf), which leaves every finite value
    # untouched, so there is no unclamped path to select between.
    bounds = clamp if clamp is not None else (-np.inf, np.inf)
    clamp_bits = tuple(int(np.float32(v).view(np.int32)) for v in bounds)
    source = _kernel_source("fused/fused_mm_tile.cc")
    include_dirs = _include_dirs()
    native_tanh = ARCH_TRAITS[arch].native_tanh
    if not native_tanh:
        from aie.utils import config

        runtime = Path(config.aie_runtime_lib_dir()) / arch.upper()
        include_dirs.append(str(runtime))
    # Include the complete recipe, not just geometry: architecture, source
    # location and runtime includes can change without changing the operands.
    key = (
        "fused_mm_tile",
        source,
        tuple(include_dirs),
        tuple(compile_flags),
        False,
        arch,
        device.default_core_stack_bytes,
        # The clamp is no longer a compile flag, so two clamps of the same
        # kernel share compile_flags. They still need their own bindings,
        # reference and tolerance, so the bounds belong in the key.
        clamp_bits,
    )
    # With one mode compiled in, the mask names the bound mode. With several,
    # the bound mode needs its own place in the key, for the clamp's reason.
    if len(set(epilogue_modes)) > 1:
        key += (modes[epilogue],)
    prefix = hashlib.sha256(repr(key).encode()).hexdigest()[:16]
    contract = KernelContract(
        trace=(
            Trace.partial("markers bracket each init, k step and drain chunk")
            if step_markers
            else Trace.whole_call()
        ),
        roles=(In, In, Out, Param, Param, Param),
        # The operands are held on the core in this blocking; nothing is
        # streamed transformed, the host packs them (block, no stream).
        layouts=(
            TensorLayout((dim_m, dim_k), pack_a, unpack_a, block=(r, s)),
            (
                TensorLayout((dim_k, dim_n), pack_b_bfp, unpack_b_bfp, block=(s, t))
                if bfp16_b
                else TensorLayout((dim_k, dim_n), pack_b, unpack_b, block=(s, t))
            ),
            TensorLayout((dim_m, dim_n), pack_c, unpack_c, block=(r, t)),
            None,
            None,
            None,
        ),
        # Bound here rather than left to the caller: `epilogue` and `clamp`
        # stay factory arguments, so `reference` below closes over the same
        # values the core is given.
        parameter_bindings=(
            (3, modes[epilogue]),
            (4, clamp_bits[0]),
            (5, clamp_bits[1]),
        ),
        reference=reference,
        # Without a tanh instruction the epilogue reads getTanhBf16's table,
        # which this box cannot measure.
        tolerance=(
            Tolerance.relative(
                0.02 if epilogue == "none" else 0.04,
                0.01 if epilogue == "none" else 0.04,
                note="bf16 store; activated path additionally narrows tanh to bf16",
            )
            if not native_tanh
            else Tolerance.bounded(
                error_bound,
                note=(
                    "f32 accumulation, vtanh's error measured on npu2 and one "
                    "bf16 ulp per rounding step, carried through each step"
                    if steps
                    else "f32 accumulation, vtanh's error measured on npu2 and "
                    "one bf16 store ulp, per output; fails a 1.5% change to "
                    "gelu's 1.702 and a 3% change to silu's or sigmoid's 1/2"
                ),
            )
        ),
        acc_dtype=np.float32,
        reduction=dim_k,
        ops_per_call=2 * dim_m * dim_k * dim_n,
        # The aie2 epilogue reaches getTanhBf16; aie2p has no table.
        uses_lut=True,
        # Reserve the f32 accumulator plus call frames and epilogue spills:
        # AIE2P SiLU with clamp needs 1600 bytes beyond the accumulator.
        # aiecc still checks the measured linked stack.
        stack_bytes=np.dtype(np.float32).itemsize * dim_m * dim_n
        + max(device.default_core_stack_bytes, 2048),
    )
    fn = _FusedMMKernel(
        "fused_mm_tile",
        source_file=str(source),
        arg_types=[
            np.ndarray[(dim_m * dim_k,), np.dtype[bfloat16]],
            (
                np.ndarray[(dim_k * dim_n // 8,), np.dtype[v8bfp16ebs8]]
                if bfp16_b
                else np.ndarray[(dim_k * dim_n,), np.dtype[bfloat16]]
            ),
            np.ndarray[(dim_m * dim_n,), np.dtype[bfloat16]],
            np.int32,
            np.int32,
            np.int32,
        ],
        include_dirs=include_dirs,
        compile_flags=compile_flags,
        # Object compilation renames every defined symbol, including the
        # included init/k_step/epilogue, zero kernels and AIE2 LUT exports.
        symbol_prefix=prefix,
        contract=contract,
    )
    acc = np.ndarray[(dim_m * dim_n,), np.dtype[np.float32]]
    b_chunk = (
        np.ndarray[(chunk_k * dim_n // 8,), np.dtype[v8bfp16ebs8]]
        if bfp16_b
        else np.ndarray[(chunk_k * dim_n,), np.dtype[bfloat16]]
    )
    bind = fn.object_file.bind
    fn.fused_mm_tile = fn
    fn.mm_fused_acc_init = bind("mm_fused_acc_init", [acc])
    # The trailing int32 is the A band.
    fn.mm_fused_k_step = bind(
        "mm_fused_k_step",
        [np.ndarray[(band_m * chunk_k,), np.dtype[bfloat16]], b_chunk, acc, np.int32],
    )
    # outer, half, mode, clamp_min_bits, clamp_max_bits
    fn.mm_fused_epilogue_chunk = bind(
        "mm_fused_epilogue_chunk",
        [np.ndarray[(out_chunk,), np.dtype[bfloat16]], acc] + [np.int32] * 5,
    )
    return fn

Convolution

Convolution kernel factories: conv2dk1/3/14, bottleneck (bn_*) variants.

The conv2dk1, conv2dk1_i8, conv2dk1_skip, conv2dk3 and conv2dk14 factories specialize their dimensions at compile time. Their runtime dimension arguments remain in the ABI and must match the factory dimensions; scales, region checks and channel offsets remain runtime values.

conv2dk1_ref

conv2dk1_ref(
    x,
    weights,
    input_width,
    input_channels,
    output_channels,
    scale,
)

Numpy reference for conv2dk1: 1x1 conv, requantized.

Layouts are the kernel's: activations [C/8][W][8] (x is one line of input_width * input_channels values, or (calls, ...) of them), weights [OC/8][IC/8][ic8][oc8], output [OC/8][W][8] as uint8. out = sat_u8((sum_ic x * w + 2**(scale-1)) >> scale), i.e. a fused ReLU. Exact for the scalar path of conv2dk1.cc.

Source code in python/iron/kernels/conv.py

def conv2dk1_ref(x, weights, input_width, input_channels, output_channels, scale):
    """Numpy reference for [`conv2dk1`][iron.kernels.conv.conv2dk1]: 1x1 conv, requantized.

    Layouts are the kernel's: activations ``[C/8][W][8]`` (``x`` is one
    line of ``input_width * input_channels`` values, or ``(calls, ...)`` of
    them), weights ``[OC/8][IC/8][ic8][oc8]``, output ``[OC/8][W][8]`` as
    ``uint8``. ``out = sat_u8((sum_ic x * w + 2**(scale-1)) >> scale)``,
    i.e. a fused ReLU. Exact for the scalar path of ``conv2dk1.cc``.
    """
    W, IC, OC = int(input_width), int(input_channels), int(output_channels)
    acc, lead = _conv1x1_acc(x, weights, W, IC, OC)
    return _requant(acc, scale).reshape(*lead, W * OC)

conv2dk1_i8_ref

conv2dk1_i8_ref(
    x,
    weights,
    input_width,
    input_channels,
    output_channels,
    scale,
)

Numpy reference for conv2dk1_i8: 1x1 conv to int8.

The layouts of conv2dk1_ref with an int8 output and no ReLU: out = sat_i8((sum_ic x * w + 2**(scale-1)) >> scale). Exact for the scalar path of conv2dk1_i8.cc.

Source code in python/iron/kernels/conv.py

def conv2dk1_i8_ref(x, weights, input_width, input_channels, output_channels, scale):
    """Numpy reference for [`conv2dk1_i8`][iron.kernels.conv.conv2dk1_i8]: 1x1 conv to ``int8``.

    The layouts of [`conv2dk1_ref`][iron.kernels.conv.conv2dk1_ref] with an
    ``int8`` output and no ReLU:
    ``out = sat_i8((sum_ic x * w + 2**(scale-1)) >> scale)``. Exact for the
    scalar path of ``conv2dk1_i8.cc``.
    """
    W, IC, OC = int(input_width), int(input_channels), int(output_channels)
    acc, lead = _conv1x1_acc(x, weights, W, IC, OC)
    return _requant(acc, scale, -128, 127, np.int8).reshape(*lead, W * OC)

conv2dk1_skip_ref

conv2dk1_skip_ref(
    x0,
    x1,
    weights,
    skip,
    input_width,
    input_channels,
    output_channels,
    scale,
    skip_scale,
)

Numpy reference for conv2dk1_skip: 1x1 conv plus residual.

x0 and x1 each hold half the input channels ([IC/16][W][8] lines, x1 the upper half); weights are [OC/8][IC/8][ic8][oc8] over all of them and skip is an [OC/8][W][8] line. The conv sum is requantized and saturated to int8 first, then the residual is added and the total requantized to uint8:

conv = sat_i8((sum_ic x * w + 2**(scale-1)) >> scale)
out  = sat_u8((conv + skip + 2**(skip_scale-1)) >> skip_scale)

Exact for the scalar path of conv2dk1_skip.cc. A shift of 0 is no shift (the scalar path's 1 << -1 is not defined for it; the vector path handles 0).

Source code in python/iron/kernels/conv.py

def conv2dk1_skip_ref(
    x0,
    x1,
    weights,
    skip,
    input_width,
    input_channels,
    output_channels,
    scale,
    skip_scale,
):
    """Numpy reference for [`conv2dk1_skip`][iron.kernels.conv.conv2dk1_skip]: 1x1 conv plus residual.

    ``x0`` and ``x1`` each hold half the input channels (``[IC/16][W][8]``
    lines, ``x1`` the upper half); weights are ``[OC/8][IC/8][ic8][oc8]``
    over all of them and ``skip`` is an ``[OC/8][W][8]`` line. The conv sum
    is requantized and saturated to ``int8`` first, then the residual is
    added and the total requantized to ``uint8``:

    ```text
    conv = sat_i8((sum_ic x * w + 2**(scale-1)) >> scale)
    out  = sat_u8((conv + skip + 2**(skip_scale-1)) >> skip_scale)
    ```

    Exact for the scalar path of ``conv2dk1_skip.cc``. A shift of 0 is no
    shift (the scalar path's ``1 << -1`` is not defined for it; the vector
    path handles 0).
    """
    W, IC, OC = int(input_width), int(input_channels), int(output_channels)
    x0, x1 = np.asarray(x0), np.asarray(x1)
    lead = x0.shape[:-1]
    halves = (x0.reshape(*lead, IC // 16, W, 8), x1.reshape(*lead, IC // 16, W, 8))
    x = np.concatenate(halves, axis=-3).reshape(*lead, W * IC)
    acc, _ = _conv1x1_acc(x, weights, W, IC, OC)
    conv = _requant(acc, scale, -128, 127, np.int64)
    total = conv + np.asarray(skip).reshape(*lead, OC // 8, W, 8).astype(np.int64)
    return _requant(total, skip_scale).reshape(*lead, W * OC)

conv2dk3_ref

conv2dk3_ref(
    line0,
    line1,
    line2,
    weights,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
)

Numpy reference for conv2dk3: 3x3 conv over three lines.

Produces the output line for line1. Activations are [C/8][W][8] lines; weights [WOC/8][IC/8][3 rows][kernel_width][ic8][oc8] where WOC may exceed output_channels (channel_offset selects this call's slice, in units of 8 channels). The spatial border is zero padded; check is 0 (top: line0 ignored), 1 (middle) or 2 (bottom: line2 ignored), matching the kernel's region enum. out = sat_u8((sum + 2**(scale-1)) >> scale). Exact for the scalar path of conv2dk3.cc; kernel_height is accepted for the signature and is 3.

Source code in python/iron/kernels/conv.py

def conv2dk3_ref(
    line0,
    line1,
    line2,
    weights,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
):
    """Numpy reference for [`conv2dk3`][iron.kernels.conv.conv2dk3]: 3x3 conv over three lines.

    Produces the output line for ``line1``. Activations are ``[C/8][W][8]``
    lines; weights ``[WOC/8][IC/8][3 rows][kernel_width][ic8][oc8]`` where
    ``WOC`` may exceed ``output_channels`` (``channel_offset`` selects this
    call's slice, in units of 8 channels). The spatial border is zero
    padded; ``check`` is 0 (top: ``line0`` ignored), 1 (middle) or 2
    (bottom: ``line2`` ignored), matching the kernel's ``region`` enum.
    ``out = sat_u8((sum + 2**(scale-1)) >> scale)``. Exact for the scalar
    path of ``conv2dk3.cc``; ``kernel_height`` is accepted for the
    signature and is 3.
    """
    acc, lead, W, OC = _conv3x3_acc(
        (line0, line1, line2),
        weights,
        input_width,
        input_channels,
        output_channels,
        kernel_width,
        check,
        channel_offset,
    )
    del kernel_height
    return _requant(acc, scale).reshape(*lead, W * OC)

conv2dk1_skip_init_ref

conv2dk1_skip_init_ref(
    x0,
    x1,
    weights,
    skip,
    input_width,
    input_channels,
    output_channels,
    skip_input_channels,
    scale,
    skip_scale,
    scale_skip_conv,
)

Numpy reference for conv2dk1_skip_init: 1x1 conv plus a projected residual.

Like conv2dk1_skip_ref, but the residual is itself a 1x1 conv of skip ([ICs/8][W][8]) with the weights stored after the main ones ([OC/8][ICs/8][ic8][oc8] at offset OC * IC):

conv = sat_i8((sum_ic x * w + 2**(scale-1)) >> scale)
proj = sat_i8((sum_ics skip * ws + 2**(scale_skip_conv-1)) >> scale_skip_conv)
out  = sat_u8((conv + proj + 2**(skip_scale-1)) >> skip_scale)

Exact for the scalar path of conv2dk1_skip_init.cc; a shift of 0 is no shift.

Source code in python/iron/kernels/conv.py

def conv2dk1_skip_init_ref(
    x0,
    x1,
    weights,
    skip,
    input_width,
    input_channels,
    output_channels,
    skip_input_channels,
    scale,
    skip_scale,
    scale_skip_conv,
):
    """Numpy reference for [`conv2dk1_skip_init`][iron.kernels.conv.conv2dk1_skip_init]: 1x1 conv plus a projected residual.

    Like [`conv2dk1_skip_ref`][iron.kernels.conv.conv2dk1_skip_ref], but the
    residual is itself a 1x1 conv of ``skip`` (``[ICs/8][W][8]``) with the
    weights stored after the main ones (``[OC/8][ICs/8][ic8][oc8]`` at
    offset ``OC * IC``):

    ```text
    conv = sat_i8((sum_ic x * w + 2**(scale-1)) >> scale)
    proj = sat_i8((sum_ics skip * ws + 2**(scale_skip_conv-1)) >> scale_skip_conv)
    out  = sat_u8((conv + proj + 2**(skip_scale-1)) >> skip_scale)
    ```

    Exact for the scalar path of ``conv2dk1_skip_init.cc``; a shift of 0
    is no shift.
    """
    W, IC, OC = int(input_width), int(input_channels), int(output_channels)
    ICs = int(skip_input_channels)
    x0, x1 = np.asarray(x0), np.asarray(x1)
    lead = x0.shape[:-1]
    halves = (x0.reshape(*lead, IC // 16, W, 8), x1.reshape(*lead, IC // 16, W, 8))
    x = np.concatenate(halves, axis=-3).reshape(*lead, W * IC)
    w = np.asarray(weights, dtype=np.int8)
    acc, _ = _conv1x1_acc(x, w[: OC * IC], W, IC, OC)
    conv = _requant(acc, scale, -128, 127, np.int64)
    acc_s, _ = _conv1x1_acc(skip, w[OC * IC : OC * IC + OC * ICs], W, ICs, OC)
    proj = _requant(acc_s, scale_skip_conv, -128, 127, np.int64)
    return _requant(conv + proj, skip_scale).reshape(*lead, W * OC)

conv2dk14_ref

conv2dk14_ref(
    x,
    weights,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    scale,
)

Numpy reference for conv2dk14: a KxK patch conv (stride K) to int8.

One call covers T = input_width / kernel_width patches of K*K RGBA pixels. Layouts are the kernel's: input [T/8][P/2][t8][p2][4] uint8 (P = K*K pixels of 4 channels), weights [OC/8][P/2][p2][4][oc8] int8, output [OC/8][T][oc8] int8: out = sat_i8((sum_{p,c} x * w + 2**(scale-1)) >> scale). input_channels is accepted for the signature (the pixel is RGBA). Exact for the scalar path of conv2dk14.cc.

Source code in python/iron/kernels/conv.py

def conv2dk14_ref(
    x, weights, input_width, input_channels, output_channels, kernel_width, scale
):
    """Numpy reference for [`conv2dk14`][iron.kernels.conv.conv2dk14]: a KxK patch conv (stride K) to ``int8``.

    One call covers ``T = input_width / kernel_width`` patches of ``K*K``
    RGBA pixels. Layouts are the kernel's: input ``[T/8][P/2][t8][p2][4]``
    ``uint8`` (``P = K*K`` pixels of 4 channels), weights
    ``[OC/8][P/2][p2][4][oc8]`` ``int8``, output ``[OC/8][T][oc8]``
    ``int8``: ``out = sat_i8((sum_{p,c} x * w + 2**(scale-1)) >> scale)``.
    ``input_channels`` is accepted for the signature (the pixel is RGBA).
    Exact for the scalar path of ``conv2dk14.cc``.
    """
    del input_channels
    W, OC, K = int(input_width), int(output_channels), int(kernel_width)
    T, P = W // K, K * K
    x = np.asarray(x)
    lead = x.shape[:-1]
    xi = x.reshape(*lead, T // 8, P // 2, 8, 2, 4).astype(np.int64)
    xi = np.moveaxis(xi, -3, -4).reshape(*lead, T, P, 4)  # [T][P][c]
    w = np.asarray(weights, dtype=np.int8).reshape(OC // 8, P, 4, 8).astype(np.int64)
    acc = np.einsum("...tpc,opcq->...otq", xi, w)  # [OC/8][T][oc8]
    return _requant(acc, scale, -128, 127, np.int8).reshape(*lead, OC * T)

dwconv1d_channels_first

dwconv1d_channels_first(
    seq_len: int = 1024,
    kernel_size: int = 9,
    bias: bool = True,
) -> ExternalFunction

Depthwise 1-D cross-correlation on one bf16 channel.

out[t] = bias + sum_p w[p] * x_pad[t + p] for t < seq_len: a 'same' convolution when x_pad is the channel zero-padded by (kernel_size - 1) // 2 on each side plus DWCONV1D_TAIL don't-care elements (programming_examples/ml/dwconv1d builds it that way). The weight row holds kernel_size taps followed by the bias, whether or not bias is enabled.

One channel per call with time contiguous, vectorized along time with scalar taps. See dwconv1d_channels_last for the transposed layout, and the "Choosing a depthwise conv1d" section of aie_kernels/README.md for which to reach for.

Parameters:

Name	Type	Description	Default
seq_len	int	Outputs per call (multiple of 16).	1024
kernel_size	int	Taps, 1 to 17.	9
bias	bool	Add the trailing weight as a bias.	True

Source code in python/iron/kernels/conv.py

def dwconv1d_channels_first(
    seq_len: int = 1024, kernel_size: int = 9, bias: bool = True
) -> ExternalFunction:
    """Depthwise 1-D cross-correlation on one bf16 channel.

    ``out[t] = bias + sum_p w[p] * x_pad[t + p]`` for ``t < seq_len``: a
    'same' convolution when ``x_pad`` is the channel zero-padded by
    ``(kernel_size - 1) // 2`` on each side plus ``DWCONV1D_TAIL`` don't-care
    elements (programming_examples/ml/dwconv1d builds it that way). The
    weight row holds ``kernel_size`` taps followed by the bias, whether or
    not ``bias`` is enabled.

    One channel per call with time contiguous, vectorized along time with
    scalar taps. See
    [`dwconv1d_channels_last`][iron.kernels.conv.dwconv1d_channels_last] for
    the transposed layout, and the "Choosing a depthwise conv1d" section of
    ``aie_kernels/README.md`` for which to reach for.

    Args:
        seq_len: Outputs per call (multiple of 16).
        kernel_size: Taps, 1 to 17.
        bias: Add the trailing weight as a bias.
    """
    if not 1 <= kernel_size <= 17:
        raise ValueError(
            f"dwconv1d_channels_first: kernel_size must be 1..17, got {kernel_size}"
        )
    if seq_len <= 0 or seq_len % 16:
        raise ValueError(
            "dwconv1d_channels_first: seq_len must be a positive multiple of 16, "
            f"got {seq_len}"
        )
    in_ty = np.ndarray[(seq_len + DWCONV1D_TAIL,), np.dtype[bfloat16]]
    w_ty = np.ndarray[(kernel_size + 1,), np.dtype[bfloat16]]
    out_ty = np.ndarray[(seq_len,), np.dtype[bfloat16]]
    return _make_extern(
        "dwconv1d_channels_first_bf16",
        _kernel_source("conv/dwconv1d_channels_first.cc"),
        [in_ty, w_ty, out_ty, np.int32],
        compile_flags=[
            f"-DDWCONV1D_CF_K={kernel_size}",
            f"-DDWCONV1D_CF_BIAS={int(bias)}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out, Param),
            reference=lambda x, w, n: dwconv1d_channels_first_ref(
                x, w, n, kernel_size=kernel_size, bias=bias
            ),
            acc_dtype=np.float32,
            reduction=kernel_size,
            tolerance=Tolerance.relative(
                0.128, 0.05, note="programming_examples/ml/dwconv1d: atol 0.05"
            ),
            ops_per_call=2 * kernel_size * seq_len,
            stack_bytes=(
                # At 17 taps
                1888
                if _detect_arch() == "aie2" and _tuned_arch() == "portable"
                else None
            ),
        ),
    )

dwconv1d_channels_first_ref

dwconv1d_channels_first_ref(
    x_pad, w, seq_len, *, kernel_size: int, bias: bool
)

Numpy reference for dwconv1d_channels_first on the padded row(s).

x_pad is (..., seq_len + DWCONV1D_TAIL); w is (..., kernel_size + 1).

Source code in python/iron/kernels/conv.py

def dwconv1d_channels_first_ref(x_pad, w, seq_len, *, kernel_size: int, bias: bool):
    """Numpy reference for [`dwconv1d_channels_first`][iron.kernels.conv.dwconv1d_channels_first] on the padded row(s).

    ``x_pad`` is ``(..., seq_len + DWCONV1D_TAIL)``; ``w`` is ``(..., kernel_size + 1)``.
    """
    n = int(seq_len)
    x32 = np.asarray(x_pad).astype(np.float32)
    w32 = np.asarray(w).astype(np.float32)
    out = np.zeros(x32.shape[:-1] + (n,), dtype=np.float32)
    for p in range(kernel_size):
        out += w32[..., p : p + 1] * x32[..., p : p + n]
    if bias:
        out += w32[..., kernel_size : kernel_size + 1]
    return out.astype(np.asarray(x_pad).dtype)

dwconv1d

dwconv1d(
    seq_len: int = 1024,
    kernel_size: int = 9,
    bias: bool = True,
) -> ExternalFunction

Compatibility alias for dwconv1d_channels_first.

Source code in python/iron/kernels/conv.py

def dwconv1d(
    seq_len: int = 1024, kernel_size: int = 9, bias: bool = True
) -> ExternalFunction:
    """Compatibility alias for [`dwconv1d_channels_first`][iron.kernels.conv.dwconv1d_channels_first]."""
    return dwconv1d_channels_first(seq_len, kernel_size, bias)

dwconv1d_ref

dwconv1d_ref(
    x_pad, w, seq_len, *, kernel_size: int, bias: bool
)

Compatibility alias for dwconv1d_channels_first_ref.

Source code in python/iron/kernels/conv.py

def dwconv1d_ref(x_pad, w, seq_len, *, kernel_size: int, bias: bool):
    """Compatibility alias for [`dwconv1d_channels_first_ref`][iron.kernels.conv.dwconv1d_channels_first_ref]."""
    return dwconv1d_channels_first_ref(
        x_pad, w, seq_len, kernel_size=kernel_size, bias=bias
    )

dwconv1d_channels_last

dwconv1d_channels_last(
    channels: int = 256, clamp: bool = True
) -> ExternalFunction

Depthwise 1-D conv over a channels-last layout, 5 taps.

y[c] = clamp(sum_{t=0..4} w_t[c] * x_t[c], lo, hi) for c < channels: one output timestep across every channel, with per-channel taps. The five taps arrive as five separate base pointers, oldest first, so a depth-5 ObjectFifo is itself the sliding window.

Counterpart to dwconv1d_channels_first; layout picks the vectorization axis, so neither subsumes the other. See the "Choosing a depthwise conv1d" section of aie_kernels/README.md.

The five weight planes are five independent arguments, like the taps, so where they live is the design's business: they need not be one buffer, or evenly spaced.

lo/hi are runtime buffers the design writes, bound here to +/-_CLAMP_LIMIT so the kernel is judged against a reference clamping to the same pair.

Parameters:

Name	Type	Description	Default
channels	int	Channels per call (multiple of 32).	256
clamp	bool	Clamp the result to the runtime lo/hi buffers.	True

Source code in python/iron/kernels/conv.py

def dwconv1d_channels_last(channels: int = 256, clamp: bool = True) -> ExternalFunction:
    """Depthwise 1-D conv over a channels-last layout, 5 taps.

    ``y[c] = clamp(sum_{t=0..4} w_t[c] * x_t[c], lo, hi)`` for ``c < channels``:
    one output timestep across every channel, with per-channel taps. The five
    taps arrive as five separate base pointers, oldest first, so a depth-5
    ObjectFifo is itself the sliding window.

    Counterpart to
    [`dwconv1d_channels_first`][iron.kernels.conv.dwconv1d_channels_first];
    layout picks the vectorization axis, so neither subsumes the other. See
    the "Choosing a depthwise conv1d" section of ``aie_kernels/README.md``.

    The five weight planes are five independent arguments, like the taps, so
    where they live is the design's business: they need not be one buffer, or
    evenly spaced.

    ``lo``/``hi`` are runtime buffers the design writes, bound here to
    ``+/-_CLAMP_LIMIT`` so the kernel is judged against a reference clamping to
    the same pair.

    Args:
        channels: Channels per call (multiple of 32).
        clamp: Clamp the result to the runtime ``lo``/``hi`` buffers.
    """
    if channels <= 0 or channels % 32:
        raise ValueError(
            "dwconv1d_channels_last: channels must be a positive multiple of the "
            f"32-lane store, got {channels}"
        )
    _TAPS = 5
    plane_ty = np.ndarray[(channels,), np.dtype[bfloat16]]
    lim_ty = np.ndarray[(1,), np.dtype[np.float32]]
    return _make_extern(
        "dwconv1d_channels_last_k5_bf16",
        _kernel_source("conv/dwconv1d_channels_last.cc"),
        [*([plane_ty] * 2 * _TAPS), plane_ty, lim_ty, lim_ty],
        compile_flags=[
            f"-DDWCONV1D_CL_C={channels}",
            f"-DDWCONV1D_CL_CLAMP={int(clamp)}",
        ],
        contract=KernelContract(
            alignments=_vector_args(*range(2 * _TAPS + 1)),
            trace=Trace.whole_call(),
            # Over 32 to 1280 channels: tuned for aie2p, 1280 B on the
            # 64-lane path and 1024 B on the generic one; portable, 1280 B on
            # aie2p and 416 B on aie2. The tuned aie2 build unrolls every
            # channel: 672 B up to 384 channels, then up to 9 B per channel
            # (8608 B at 960).
            stack_bytes={
                "aie2": None if channels <= 384 else 9 * channels - 32,
                "aie2p": 1280 if channels % 64 == 0 else None,
            }.get(_tuned_arch(), 1280),
            setup=conv_even,
            # lo/hi are buffers the design writes, so they are Param like
            # mha's idx gate: bound here rather than sampled, which also keeps
            # lo <= hi (aie::clamp does not define the inverted pair).
            roles=(*((In,) * 2 * _TAPS), Out, Param, Param),
            parameter_bindings=(
                (11, np.array([-_CLAMP_LIMIT], np.float32)),
                (12, np.array([_CLAMP_LIMIT], np.float32)),
            ),
            reference=partial(
                dwconv1d_channels_last_ref,
                lo=-_CLAMP_LIMIT,
                hi=_CLAMP_LIMIT,
                clamp=clamp,
            ),
            acc_dtype=np.float32,
            reduction=_TAPS,
            # Derived: the five products are exact in f32 (bf16 carries 8
            # mantissa bits, 8 + 8 < 24) and the sum rounds at most 5 * 2**-24
            # before one bf16 narrowing on store, which is 2**-9 relative. The
            # narrowing dominates by four orders of magnitude, so the bound is
            # one bf16 ulp with no room to spare for a real error.
            tolerance=Tolerance.relative(
                2**-8,
                0.0,
                note="one bf16 ulp from the single narrowing store; the f32 "
                "5-term sum contributes 5 * 2**-24. Verified on npu2",
            ),
            ops_per_call=2 * _TAPS * channels,
        ),
    )

dwconv1d_channels_last_ref

dwconv1d_channels_last_ref(
    w_0,
    w_1,
    w_2,
    w_3,
    w_4,
    x_0,
    x_1,
    x_2,
    x_3,
    x_4,
    *,
    lo,
    hi,
    clamp: bool
)

Numpy reference for dwconv1d_channels_last: one timestep over all channels.

Each tap is an independent plane, so this is a plain sum_t w_t * x_t. Accumulated in float32, which is what the kernel's accfloat is, then narrowed once on store; clamp applies after the narrowing, as the kernel's aie::clamp does on the already-bf16 vector.

Source code in python/iron/kernels/conv.py

def dwconv1d_channels_last_ref(
    w_0, w_1, w_2, w_3, w_4, x_0, x_1, x_2, x_3, x_4, *, lo, hi, clamp: bool
):
    """Numpy reference for [`dwconv1d_channels_last`][iron.kernels.conv.dwconv1d_channels_last]: one timestep over all channels.

    Each tap is an independent plane, so this is a plain ``sum_t w_t * x_t``.
    Accumulated in float32, which is what the kernel's ``accfloat`` is, then
    narrowed once on store; ``clamp`` applies after the narrowing, as the
    kernel's ``aie::clamp`` does on the already-bf16 vector.
    """
    ws = [np.asarray(v).astype(np.float32) for v in (w_0, w_1, w_2, w_3, w_4)]
    xs = [np.asarray(v).astype(np.float32) for v in (x_0, x_1, x_2, x_3, x_4)]
    acc = np.zeros(xs[0].shape, dtype=np.float32)
    for wt, xt in zip(ws, xs):
        acc += wt * xt
    out = acc.astype(np.asarray(x_0).dtype)
    if clamp:
        out = np.clip(out, np.asarray(lo, out.dtype), np.asarray(hi, out.dtype))
    return out

conv2dk1

conv2dk1(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    act_dtype: type = int8,
) -> ExternalFunction

1x1 convolution kernel.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels.	64
act_dtype	type	Activation data type (np.int8 or np.uint8).	int8

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk1 kernel.

Raises:

Type	Description
ValueError	When act_dtype is not np.int8 or np.uint8.

Source code in python/iron/kernels/conv.py

@dtypes(({"act_dtype": np.int8}, {"act_dtype": np.uint8}))
def conv2dk1(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    act_dtype: type = np.int8,
) -> ExternalFunction:
    """1x1 convolution kernel.

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels.
        act_dtype: Activation data type (``np.int8`` or ``np.uint8``).

    Returns:
        ExternalFunction configured for the conv2dk1 kernel.

    Raises:
        ValueError: When ``act_dtype`` is not ``np.int8`` or ``np.uint8``.
    """
    func_name, flags = _conv_act_dtype_info(
        "conv2dk1", act_dtype, factory_name="conv2dk1"
    )
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[act_dtype]]
    wt_ty = np.ndarray[(input_channels * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    return _make_extern(
        func_name,
        _kernel_source("conv/conv2dk1.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(4)],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            alignments=_vector_args(0, 1, 2),
            trace=Trace.whole_call(),
            # 1088 B tuned for aie2p, which keeps the oc-invariant input block
            # on the stack; 2752 B untuned on aie2p (1504 B on aie2); 288 B
            # tuned for aie2
            stack_bytes={"aie2": None, "aie2p": 1088}.get(_tuned_arch(), 2752),
            roles=(In, Param, Out, Param, Param, Param, Param),
            reference=conv2dk1_ref,
            acc_dtype=np.int32,
            reduction=input_channels,
            tolerance=Tolerance.exact(
                note="measured bit-exact against the reference over every data case"
            ),
            ops_per_call=2 * input_width * input_channels * output_channels,
        ),
    )

conv2dk3

conv2dk3(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    act_dtype: type = int8,
    weight_output_channels: int | None = None,
) -> ExternalFunction

3x3 convolution kernel.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels produced by this call.	64
act_dtype	type	Activation data type (np.int8 or np.uint8).	int8
weight_output_channels	int | None	Total number of output channels stored in the weights buffer. Defaults to output_channels. Set higher than output_channels when the weights buffer is shared across multiple workers that each produce a slice of the output (the channel_offset runtime arg selects a worker's slice).	None

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk3 kernel.

Raises:

Type	Description
ValueError	When act_dtype is not np.int8 or np.uint8, or when input_width is not 32.

Source code in python/iron/kernels/conv.py

@dtypes(({"act_dtype": np.int8}, {"act_dtype": np.uint8}))
def conv2dk3(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    act_dtype: type = np.int8,
    weight_output_channels: int | None = None,
) -> ExternalFunction:
    """3x3 convolution kernel.

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels produced by this call.
        act_dtype: Activation data type (``np.int8`` or ``np.uint8``).
        weight_output_channels: Total number of output channels stored in the
            weights buffer. Defaults to ``output_channels``. Set higher than
            ``output_channels`` when the weights buffer is shared across
            multiple workers that each produce a slice of the output (the
            ``channel_offset`` runtime arg selects a worker's slice).

    Returns:
        ExternalFunction configured for the conv2dk3 kernel.

    Raises:
        ValueError: When ``act_dtype`` is not ``np.int8`` or ``np.uint8``, or
            when ``input_width`` is not 32.

    """
    if input_width != 32:
        # Every compiled variant of the vector kernel (aie2, aie2p, and the
        # portable branch) hard-codes a 32-pixel row inside conv2dk3.cc and
        # ignores the runtime_input_width arg; a different width would
        # silently compute over the wrong pixels instead of raising.
        raise ValueError(f"conv2dk3: input_width must be 32, got {input_width}")
    func_name, flags = _conv_act_dtype_info(
        "conv2dk3", act_dtype, factory_name="conv2dk3"
    )
    if weight_output_channels is None:
        weight_output_channels = output_channels
    line_size = input_width * input_channels
    line_ty = np.ndarray[(line_size,), np.dtype[act_dtype]]
    wt_ty = np.ndarray[
        (3 * 3 * input_channels * weight_output_channels,), np.dtype[np.int8]
    ]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    return _make_extern(
        func_name,
        _kernel_source("conv/conv2dk3.cc"),
        [line_ty, line_ty, line_ty, wt_ty, out_ty, *_i32s(8)],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels)
        + ["-DCONV_KERNEL_WIDTH=3", "-DCONV_KERNEL_HEIGHT=3"],
        contract=KernelContract(
            alignments=_vector_args(0, 1, 2, 3, 4),
            trace=Trace.whole_call(),
            # 384 B tuned for aie2p, 4736 B untuned; 0 B tuned for aie2
            stack_bytes={"aie2": None, "aie2p": 384}.get(_tuned_arch(), 4736),
            roles=(In, In, In, Param, Out, *((Param,) * 8)),
            reference=conv2dk3_ref,
            acc_dtype=np.int32,
            reduction=9 * input_channels,
            tolerance=Tolerance.exact(
                note="measured bit-exact against the reference over every data case"
            ),
            ops_per_call=2 * 9 * input_width * input_channels * output_channels,
        ),
    )

conv2dk1_skip

conv2dk1_skip(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    act_dtype: type = int8,
) -> ExternalFunction

1x1 convolution kernel with skip (residual) connection.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels.	64
act_dtype	type	Activation data type (np.int8 or np.uint8).	int8

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk1_skip kernel.

Raises:

Type	Description
ValueError	When act_dtype is not np.int8 or np.uint8.

Note

The activations are uint8 in two half-channel tensors whatever act_dtype, which types the residual (skip) only. The generic harness packs tensors of one type into one fifo, so an int8 residual streams beside the uint8 activations in a second fifo.

Source code in python/iron/kernels/conv.py

@dtypes(({"act_dtype": np.int8}, {"act_dtype": np.uint8}))
def conv2dk1_skip(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    act_dtype: type = np.int8,
) -> ExternalFunction:
    """1x1 convolution kernel with skip (residual) connection.

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels.
        act_dtype: Activation data type (``np.int8`` or ``np.uint8``).

    Returns:
        ExternalFunction configured for the conv2dk1_skip kernel.

    Raises:
        ValueError: When ``act_dtype`` is not ``np.int8`` or ``np.uint8``.

    Note:
        The activations are ``uint8`` in two half-channel tensors whatever
        ``act_dtype``, which types the residual (``skip``) only. The generic
        harness packs tensors of one type into one fifo, so an ``int8``
        residual streams beside the ``uint8`` activations in a second fifo.
    """
    func_name, flags = _conv_act_dtype_info(
        "conv2dk1_skip", act_dtype, factory_name="conv2dk1_skip"
    )
    half_ch = input_channels // 2
    in0_ty = np.ndarray[(input_width * half_ch,), np.dtype[np.uint8]]
    in1_ty = np.ndarray[(input_width * half_ch,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(input_channels * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    skip_ty = np.ndarray[(input_width * output_channels,), np.dtype[act_dtype]]
    return _make_extern(
        func_name,
        _kernel_source("conv/conv2dk1_skip.cc"),
        [in0_ty, in1_ty, wt_ty, out_ty, skip_ty, *_i32s(5)],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            alignments=_vector_args(0, 1, 2, 3, 4),
            trace=Trace.whole_call(),
            # With an int8 skip: 512 B tuned for aie2p, 2816 B untuned; 32 B
            # tuned for aie2
            stack_bytes={"aie2": None, "aie2p": 512}.get(_tuned_arch(), 2816),
            roles=(In, In, Param, Out, In, *((Param,) * 5)),
            reference=conv2dk1_skip_ref,
            acc_dtype=np.int32,
            reduction=input_channels,
            tolerance=Tolerance.exact(
                note="measured bit-exact against the reference over every data case"
            ),
            ops_per_call=2 * input_width * input_channels * output_channels
            + input_width * output_channels,
        ),
    )

conv2dk1_i8

conv2dk1_i8(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
) -> ExternalFunction

1x1 convolution kernel with int8 activations/weights/output.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels.	64

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk1_i8 kernel.

Source code in python/iron/kernels/conv.py

def conv2dk1_i8(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
) -> ExternalFunction:
    """1x1 convolution kernel with int8 activations/weights/output.

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels.

    Returns:
        ExternalFunction configured for the conv2dk1_i8 kernel.
    """
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[np.int8]]
    wt_ty = np.ndarray[(input_channels * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.int8]]
    return _make_extern(
        "conv2dk1_i8",
        _kernel_source("conv/conv2dk1_i8.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(4)],
        compile_flags=["-DINT8_ACT"]
        + _conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            trace=Trace.whole_call(),
            # At most 480 B tuned for aie2, 128 B tuned for aie2p and 256 B
            # portable, all under the 1024 B default
            stack_bytes=None,
            roles=(In, Param, Out, Param, Param, Param, Param),
            reference=conv2dk1_i8_ref,
            acc_dtype=np.int32,
            reduction=input_channels,
            # The only conv kernel that is not bit-exact: its vector path
            # ends in a symmetric_inf srs the scalar reference does not model,
            # measured at 2 of 98304 values, each one LSB out.
            tolerance=Tolerance.lsb(
                1, note="vector path srs rounding; measured within one LSB"
            ),
            ops_per_call=2 * input_width * input_channels * output_channels,
        ),
    )

conv2dk14

conv2dk14(
    input_width: int = 224,
    input_channels: int = 16,
    output_channels: int = 16,
    kernel_width: int = 14,
) -> ExternalFunction

14x14 convolution kernel.

The source lives under aie_kernels/conv/ and builds for aie2 as well, where the vector path has its own AIE2 variant.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input, a multiple of 16 patches.	224
input_channels	int	Number of input channels.	16
output_channels	int	Number of output channels.	16
kernel_width	int	Width (and height) of the convolution kernel, even.	14

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk14 kernel.

Source code in python/iron/kernels/conv.py

def conv2dk14(
    input_width: int = 224,
    input_channels: int = 16,
    output_channels: int = 16,
    kernel_width: int = 14,
) -> ExternalFunction:
    """14x14 convolution kernel.

    The source lives under ``aie_kernels/conv/`` and builds for aie2 as
    well, where the vector path has its own AIE2 variant.

    Args:
        input_width: Spatial width of the input, a multiple of 16 patches.
        input_channels: Number of input channels.
        output_channels: Number of output channels.
        kernel_width: Width (and height) of the convolution kernel, even.

    Returns:
        ExternalFunction configured for the conv2dk14 kernel.
    """
    tiles = input_width // kernel_width
    if kernel_width % 2 or tiles * kernel_width != input_width or tiles % 16:
        raise ValueError(
            "conv2dk14: kernel_width must be even and input_width a multiple "
            f"of 16 * kernel_width, got {input_width} and {kernel_width}"
        )
    pixels = kernel_width * kernel_width
    _RGBA = 4
    in_ty = np.ndarray[(tiles * pixels * _RGBA,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(output_channels * pixels * _RGBA,), np.dtype[np.int8]]
    # One int8 per (output channel, tile): conv2dk14.cc writes
    # output[oc * tiles + tile] and nothing beyond it.
    out_ty = np.ndarray[(output_channels * tiles,), np.dtype[np.int8]]
    return _make_extern(
        "conv2dk14_i8",
        _kernel_source("conv/conv2dk14.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(5)],
        compile_flags=_conv_dimensions(input_width, input_channels, output_channels)
        + [f"-DCONV_KERNEL_WIDTH={kernel_width}"],
        contract=KernelContract(
            alignments=_vector_args(0, 1, 2),
            trace=Trace.whole_call(),
            roles=(In, Param, Out, *((Param,) * 5)),
            reference=conv2dk14_ref,
            acc_dtype=np.int32,
            reduction=pixels * _RGBA,
            tolerance=Tolerance.exact(
                note="measured bit-exact against the reference over every data case"
            ),
            ops_per_call=2 * tiles * pixels * _RGBA * output_channels,
        ),
    )

conv2dk1_skip_init

conv2dk1_skip_init(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    act_dtype: type = int8,
    skip_input_channels: int | None = None,
) -> ExternalFunction

1x1 convolution kernel with skip-init connection.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels.	64
act_dtype	type	Activation data type (np.int8 or np.uint8).	int8
skip_input_channels	int | None	Number of input channels for the skip-projection 1x1 conv whose weights are concatenated after the main conv weights in the same buffer. Defaults to input_channels.	None

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk1_skip_init kernel.

Raises:

Type	Description
ValueError	When act_dtype is not np.int8 or np.uint8, when input_width is not a positive multiple of 32, or when a channel count is not a positive multiple of what the source steps by (16 input channels, 8 output and skip channels).

Source code in python/iron/kernels/conv.py

@dtypes(({"act_dtype": np.int8}, {"act_dtype": np.uint8}))
def conv2dk1_skip_init(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    act_dtype: type = np.int8,
    skip_input_channels: int | None = None,
) -> ExternalFunction:
    """1x1 convolution kernel with skip-init connection.

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels.
        act_dtype: Activation data type (``np.int8`` or ``np.uint8``).
        skip_input_channels: Number of input channels for the skip-projection
            1x1 conv whose weights are concatenated after the main conv
            weights in the same buffer. Defaults to ``input_channels``.

    Returns:
        ExternalFunction configured for the conv2dk1_skip_init kernel.

    Raises:
        ValueError: When ``act_dtype`` is not ``np.int8`` or ``np.uint8``,
            when ``input_width`` is not a positive multiple of 32, or when a
            channel count is not a positive multiple of what the source
            steps by (16 input channels, 8 output and skip channels).
    """
    func_name, flags = _conv_act_dtype_info(
        "conv2dk1_skip_init", act_dtype, factory_name="conv2dk1_skip_init"
    )
    if input_width <= 0 or input_width % 32:
        # The kernel computes whole 32-wide blocks and its tail path was never
        # implemented, so a width that is not a multiple of 32 silently leaves
        # its last columns unwritten. Refuse it rather than return part of an
        # answer.
        raise ValueError(
            f"conv2dk1_skip_init: input_width must be a positive multiple of 32, "
            f"got {input_width}"
        )
    if skip_input_channels is None:
        skip_input_channels = input_channels
    # Both paths in the source count channels in whole steps: the two input
    # halves in 16s (an 8x8 weight block per half), the output and the skip
    # projection in 8s. A count that is not a whole number of steps is
    # silently truncated by the integer division in the loop bounds.
    for label, count, step in (
        ("input_channels", input_channels, 16),
        ("output_channels", output_channels, 8),
        ("skip_input_channels", skip_input_channels, 8),
    ):
        if count <= 0 or count % step:
            raise ValueError(
                f"conv2dk1_skip_init: {label} must be a positive multiple of "
                f"{step}, got {count}"
            )
    half_ch = input_channels // 2
    total_in_ch = input_channels + skip_input_channels
    in0_ty = np.ndarray[(input_width * half_ch,), np.dtype[np.uint8]]
    in1_ty = np.ndarray[(input_width * half_ch,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(total_in_ch * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    skip_ty = np.ndarray[(input_width * skip_input_channels,), np.dtype[act_dtype]]
    return _make_extern(
        func_name,
        _kernel_source("conv/conv2dk1_skip_init.cc"),
        [in0_ty, in1_ty, wt_ty, out_ty, skip_ty, *_i32s(7)],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels)
        + [f"-DCONV_SKIP_INPUT_CHANNELS={skip_input_channels}"],
        contract=KernelContract(
            alignments=_vector_args(0, 1, 2, 3, 4),
            trace=Trace.whole_call(),
            # 1728 B tuned for aie2p, the largest over input_channels 16..256;
            # untuned >=2144 plus __modsi3, which has no .stack_sizes; 288 B
            # tuned for aie2
            stack_bytes={"aie2": None, "aie2p": 1728}.get(_tuned_arch(), 0x2000),
            roles=(In, In, Param, Out, In, *((Param,) * 7)),
            reference=conv2dk1_skip_init_ref,
            acc_dtype=np.int32,
            reduction=max(input_channels, skip_input_channels),
            # Measured bit-exact over every data case at three seeds. The uint8
            # entry point was an empty function and the one-LSB slack these
            # conv kernels used to share was not what hid it, but an exact
            # contract states what this kernel actually owes.
            tolerance=Tolerance.exact(
                note="both paths match the reference bit-for-bit"
            ),
            ops_per_call=2
            * input_width
            * output_channels
            * (input_channels + skip_input_channels),
        ),
    )

bn_conv2dk1_relu_ref

bn_conv2dk1_relu_ref(
    x,
    weights,
    input_width,
    input_channels,
    output_channels,
    scale,
)

Numpy reference for bn_conv2dk1_relu.

The layouts of conv2dk1_ref, rounding half to even: out = sat_u8(round_even(sum_ic x * w, scale)).

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_relu_ref(
    x, weights, input_width, input_channels, output_channels, scale
):
    """Numpy reference for [`bn_conv2dk1_relu`][iron.kernels.conv.bn_conv2dk1_relu].

    The layouts of [`conv2dk1_ref`][iron.kernels.conv.conv2dk1_ref], rounding
    half to even: ``out = sat_u8(round_even(sum_ic x * w, scale))``.
    """
    W, IC, OC = int(input_width), int(input_channels), int(output_channels)
    acc, lead = _conv1x1_acc(x, weights, W, IC, OC)
    return _requant_even(acc, scale).reshape(*lead, W * OC)

bn_conv2dk1_i8_ref

bn_conv2dk1_i8_ref(
    x,
    weights,
    input_width,
    input_channels,
    output_channels,
    scale,
)

Numpy reference for bn_conv2dk1_i8.

uint8 activations, int8 output: out = sat_i8(round_even(sum_ic x * w, scale)).

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_i8_ref(x, weights, input_width, input_channels, output_channels, scale):
    """Numpy reference for [`bn_conv2dk1_i8`][iron.kernels.conv.bn_conv2dk1_i8].

    ``uint8`` activations, ``int8`` output:
    ``out = sat_i8(round_even(sum_ic x * w, scale))``.
    """
    W, IC, OC = int(input_width), int(input_channels), int(output_channels)
    acc, lead = _conv1x1_acc(x, weights, W, IC, OC)
    return _requant_even(acc, scale, -128, 127, np.int8).reshape(*lead, W * OC)

bn_conv2dk1_skip_ref

bn_conv2dk1_skip_ref(
    x,
    weights,
    skip,
    input_width,
    input_channels,
    output_channels,
    scale,
    skip_scale,
)

Numpy reference for bn_conv2dk1_skip.

skip is an [OC/8][W][8] line of either signedness:

conv = sat_i8(round_even(sum_ic x * w, scale))
out  = sat_i8(round_even(conv + skip, skip_scale))

scale must be at least 1; skip_scale 0 is a saturating add.

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_skip_ref(
    x,
    weights,
    skip,
    input_width,
    input_channels,
    output_channels,
    scale,
    skip_scale,
):
    """Numpy reference for [`bn_conv2dk1_skip`][iron.kernels.conv.bn_conv2dk1_skip].

    ``skip`` is an ``[OC/8][W][8]`` line of either signedness:

    ```text
    conv = sat_i8(round_even(sum_ic x * w, scale))
    out  = sat_i8(round_even(conv + skip, skip_scale))
    ```

    ``scale`` must be at least 1; ``skip_scale`` 0 is a saturating add.
    """
    W, IC, OC = int(input_width), int(input_channels), int(output_channels)
    acc, lead = _conv1x1_acc(x, weights, W, IC, OC)
    conv = _requant_even(acc, scale, -128, 127, np.int64)
    total = conv + np.asarray(skip).reshape(*lead, OC // 8, W, 8).astype(np.int64)
    return _requant_even(total, skip_scale, -128, 127, np.int8).reshape(*lead, W * OC)

bn_conv2dk3_ref

bn_conv2dk3_ref(
    line0,
    line1,
    line2,
    weights,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
)

Numpy reference for bn_conv2dk3: 3x3 stride-2 conv.

The layouts and check of conv2dk3_ref; output x reads input pixels 2x-1 .. 2x+1 and the output line is input_width / 2 wide: out = sat_u8(round_even(sum, scale)).

Source code in python/iron/kernels/conv.py

def bn_conv2dk3_ref(
    line0,
    line1,
    line2,
    weights,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
):
    """Numpy reference for [`bn_conv2dk3`][iron.kernels.conv.bn_conv2dk3]: 3x3 stride-2 conv.

    The layouts and ``check`` of [`conv2dk3_ref`][iron.kernels.conv.conv2dk3_ref];
    output ``x`` reads input pixels ``2x-1 .. 2x+1`` and the output line is
    ``input_width / 2`` wide: ``out = sat_u8(round_even(sum, scale))``.
    """
    acc, lead, Wo, OC = _conv3x3_acc(
        (line0, line1, line2),
        weights,
        input_width,
        input_channels,
        output_channels,
        kernel_width,
        check,
        channel_offset,
        stride=2,
    )
    del kernel_height
    return _requant_even(acc, scale).reshape(*lead, Wo * OC)

bn_conv2dk3_dw_ref

bn_conv2dk3_dw_ref(
    line0,
    line1,
    line2,
    weights,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
    *,
    stride: int = 1
)

Numpy reference for bn_conv2dk3_dw: depthwise 3x3 + ReLU.

Lines are [C/8][W][8] uint8, weights [C/8][3 rows][3][c8], check as in conv2dk3_ref: out = sat_u8(round_even(sum, scale)) over output_channels channels and input_width / stride pixels. kernel_width, kernel_height and channel_offset are accepted for the signature.

Source code in python/iron/kernels/conv.py

def bn_conv2dk3_dw_ref(
    line0,
    line1,
    line2,
    weights,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
    *,
    stride: int = 1,
):
    """Numpy reference for [`bn_conv2dk3_dw`][iron.kernels.conv.bn_conv2dk3_dw]: depthwise 3x3 + ReLU.

    Lines are ``[C/8][W][8]`` ``uint8``, weights ``[C/8][3 rows][3][c8]``,
    ``check`` as in [`conv2dk3_ref`][iron.kernels.conv.conv2dk3_ref]:
    ``out = sat_u8(round_even(sum, scale))`` over ``output_channels``
    channels and ``input_width / stride`` pixels. ``kernel_width``,
    ``kernel_height`` and ``channel_offset`` are accepted for the signature.
    """
    del input_channels, kernel_width, kernel_height, channel_offset
    acc, lead, Wo = _dw3x3_acc(
        (line0, line1, line2), weights, input_width, output_channels, check, stride
    )
    return _requant_even(acc, scale).reshape(*lead, Wo * int(output_channels))

bn_conv2dk3_dw_out_split_ref

bn_conv2dk3_dw_out_split_ref(
    line0,
    line1,
    line2,
    weights,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
)

Numpy reference for bn_conv2dk3_dw_out_split.

The stride-1 bn_conv2dk3_dw_ref, its output channels split in two halves, one per output.

Source code in python/iron/kernels/conv.py

def bn_conv2dk3_dw_out_split_ref(
    line0,
    line1,
    line2,
    weights,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
):
    """Numpy reference for [`bn_conv2dk3_dw_out_split`][iron.kernels.conv.bn_conv2dk3_dw_out_split].

    The stride-1 [`bn_conv2dk3_dw_ref`][iron.kernels.conv.bn_conv2dk3_dw_ref],
    its output channels split in two halves, one per output.
    """
    out = bn_conv2dk3_dw_ref(
        line0,
        line1,
        line2,
        weights,
        input_width,
        input_channels,
        output_channels,
        kernel_width,
        kernel_height,
        check,
        scale,
        channel_offset,
    )
    half = out.shape[-1] // 2
    return out[..., :half], out[..., half:]

bn_conv2dk1_relu_xy_pool_padded_ref

bn_conv2dk1_relu_xy_pool_padded_ref(
    x,
    weights,
    input_width,
    input_channels,
    output_channels,
    output_channels_padd,
    scale,
    y_index,
    output_split,
    weight_index,
)

Numpy reference for bn_conv2dk1_relu_xy_pool_padded into a zeroed output.

Each output channel of this call's slice (output_channels / output_split of them, at slice weight_index) sums sat_u8(round_even(sum_ic x * w, scale)) over the input_width pixels. On the last row (y_index == input_width - 1) the sum is averaged over 49 pixels in float32 as the kernel does: an average whose first decimal is 5 rounds to even, any other rounds half up. Channels outside the slice stay 0, as do the pad channels from output_channels up to output_channels_padd.

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_relu_xy_pool_padded_ref(
    x,
    weights,
    input_width,
    input_channels,
    output_channels,
    output_channels_padd,
    scale,
    y_index,
    output_split,
    weight_index,
):
    """Numpy reference for [`bn_conv2dk1_relu_xy_pool_padded`][iron.kernels.conv.bn_conv2dk1_relu_xy_pool_padded] into a zeroed output.

    Each output channel of this call's slice (``output_channels /
    output_split`` of them, at slice ``weight_index``) sums
    ``sat_u8(round_even(sum_ic x * w, scale))`` over the ``input_width``
    pixels. On the last row (``y_index == input_width - 1``) the sum is
    averaged over 49 pixels in ``float32`` as the kernel does: an average
    whose first decimal is 5 rounds to even, any other rounds half up.
    Channels outside the slice stay 0, as do the pad channels from
    ``output_channels`` up to ``output_channels_padd``.
    """
    W, IC, OC = int(input_width), int(input_channels), int(output_channels)
    tile = OC // int(output_split)
    acc, lead = _conv1x1_acc(x, np.asarray(weights)[: IC * tile], W, IC, tile)
    total = _requant_even(acc, scale, dtype=np.int64).sum(axis=-2)  # [tile/8][8]
    total = total.reshape(*lead, tile)
    if int(y_index) == W - 1:
        avg = total.astype(np.float32) / np.float32(49.0)
        whole = avg.astype(np.int32)
        tie = (avg * np.float32(10)).astype(np.int32) % 10 == 5
        even = np.where(whole % 2 == 0, whole, whole + 1)
        total = np.where(tie, even, (avg + np.float32(0.5)).astype(np.int32))
    out = np.zeros((*lead, max(OC, int(output_channels_padd))), dtype=np.uint16)
    start = tile * int(weight_index)
    out[..., start : start + tile] = total.astype(np.uint16)
    return out

bn_fc_relu_ui16_pad_ref

bn_fc_relu_ui16_pad_ref(
    x,
    weights,
    input_width,
    input_channels,
    input_channels_pad,
    output_channels,
    scale,
)

Numpy reference for bn_fc_relu_ui16_pad.

A 1x1 conv of uint16 activations whose weights are laid out for input_channels_pad input channels ([OC/8][ICP/8][ic8][oc8], the first input_channels used); the output is uint16 holding sat_u8(round_even(sum_ic x * w, scale)).

Source code in python/iron/kernels/conv.py

def bn_fc_relu_ui16_pad_ref(
    x, weights, input_width, input_channels, input_channels_pad, output_channels, scale
):
    """Numpy reference for [`bn_fc_relu_ui16_pad`][iron.kernels.conv.bn_fc_relu_ui16_pad].

    A 1x1 conv of ``uint16`` activations whose weights are laid out for
    ``input_channels_pad`` input channels (``[OC/8][ICP/8][ic8][oc8]``, the
    first ``input_channels`` used); the output is ``uint16`` holding
    ``sat_u8(round_even(sum_ic x * w, scale))``.
    """
    W, IC, ICP = int(input_width), int(input_channels), int(input_channels_pad)
    OC = int(output_channels)
    x = np.asarray(x)
    lead = x.shape[:-1]
    xi = x[..., : IC * W].reshape(*lead, IC // 8, W, 8).astype(np.int64)
    w = np.asarray(weights, dtype=np.int8)[: OC * ICP]
    w = w.reshape(OC // 8, ICP // 8, 8, 8)[:, : IC // 8].astype(np.int64)
    acc = np.einsum("...iwc,oicp->...owp", xi, w)
    return _requant_even(acc, scale, dtype=np.uint16).reshape(*lead, W * OC)

bn_conv2dk1_relu

bn_conv2dk1_relu(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
) -> ExternalFunction

Bottleneck 1x1 conv + ReLU kernel (int8 in, uint8 out).

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels.	64

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk1_relu kernel.

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_relu(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
) -> ExternalFunction:
    """Bottleneck 1x1 conv + ReLU kernel (int8 in, uint8 out).

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels.

    Returns:
        ExternalFunction configured for the bn_conv2dk1_relu kernel.
    """
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[np.int8]]
    wt_ty = np.ndarray[(input_channels * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    return _make_extern(
        "conv2dk1_relu_i8_ui8",
        _kernel_source("conv/bn_conv2dk1_relu.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(4)],
        compile_flags=["-DREGULAR", "-DINT8_ACT"]
        + _conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            alignments=_vector_loads(1),
            trace=Trace.whole_call(),
            roles=(In, Param, Out, Param, Param, Param, Param),
            reference=bn_conv2dk1_relu_ref,
            acc_dtype=np.int32,
            reduction=input_channels,
            tolerance=_BN_EXACT,
            ops_per_call=2 * input_width * input_channels * output_channels,
        ),
    )

bn_conv2dk3

bn_conv2dk3(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    weight_output_channels: int | None = None,
) -> ExternalFunction

Bottleneck 3x3 conv with stride-2 kernel (int8 in, uint8 out).

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels produced by this call.	64
weight_output_channels	int | None	Total number of output channels stored in the weights buffer, as for conv2dk3. Defaults to output_channels.	None

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk3 kernel.

Source code in python/iron/kernels/conv.py

def bn_conv2dk3(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    weight_output_channels: int | None = None,
) -> ExternalFunction:
    """Bottleneck 3x3 conv with stride-2 kernel (int8 in, uint8 out).

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels produced by this call.
        weight_output_channels: Total number of output channels stored in the
            weights buffer, as for [`conv2dk3`][iron.kernels.conv.conv2dk3].
            Defaults to ``output_channels``.

    Returns:
        ExternalFunction configured for the bn_conv2dk3 kernel.
    """
    if weight_output_channels is None:
        weight_output_channels = output_channels
    line_size = input_width * input_channels
    line_ty = np.ndarray[(line_size,), np.dtype[np.int8]]
    wt_ty = np.ndarray[
        (3 * 3 * input_channels * weight_output_channels,), np.dtype[np.int8]
    ]
    # Output is half-resolution because the kernel is stride-2.
    out_ty = np.ndarray[((input_width // 2) * output_channels,), np.dtype[np.uint8]]
    return _make_extern(
        "conv2dk3_stride2_i8",
        _kernel_source("conv/bn_conv2dk3.cc"),
        [line_ty, line_ty, line_ty, wt_ty, out_ty, *_i32s(8)],
        compile_flags=_conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Param, Out, *((Param,) * 8)),
            alignments=_vector_loads(0, 1, 2, 3) + _vector_32b(4),
            reference=bn_conv2dk3_ref,
            acc_dtype=np.int32,
            reduction=9 * input_channels,
            tolerance=_BN_EXACT,
            ops_per_call=2 * 9 * (input_width // 2) * input_channels * output_channels,
        ),
    )

bn_conv2dk1_i8

bn_conv2dk1_i8(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
) -> ExternalFunction

Bottleneck 1x1 conv kernel (uint8 in, int8 out).

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels.	64

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk1_i8 kernel.

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_i8(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
) -> ExternalFunction:
    """Bottleneck 1x1 conv kernel (uint8 in, int8 out).

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels.

    Returns:
        ExternalFunction configured for the bn_conv2dk1_i8 kernel.
    """
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(input_channels * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.int8]]
    return _make_extern(
        "conv2dk1_ui8_i8",
        _kernel_source("conv/bn_conv2dk1_i8.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(4)],
        compile_flags=["-DREGULAR", "-DSCALAR"]
        + _conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            alignments=_vector_loads(1),
            trace=Trace.whole_call(),
            roles=(In, Param, Out, Param, Param, Param, Param),
            reference=bn_conv2dk1_i8_ref,
            acc_dtype=np.int32,
            reduction=input_channels,
            tolerance=_BN_EXACT,
            ops_per_call=2 * input_width * input_channels * output_channels,
        ),
    )

bn_conv2dk1_skip

bn_conv2dk1_skip(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    skip_dtype: type = uint8,
) -> ExternalFunction

Bottleneck 1x1 conv with skip connection (uint8 in).

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels.	64
skip_dtype	type	Skip connection data type (np.uint8 or np.int8).	uint8

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk1_skip kernel.

Raises:

Type	Description
ValueError	When skip_dtype is not np.uint8 or np.int8.

Source code in python/iron/kernels/conv.py

@dtypes(({"skip_dtype": np.uint8}, {"skip_dtype": np.int8}))
def bn_conv2dk1_skip(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    skip_dtype: type = np.uint8,
) -> ExternalFunction:
    """Bottleneck 1x1 conv with skip connection (uint8 in).

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels.
        skip_dtype: Skip connection data type (``np.uint8`` or ``np.int8``).

    Returns:
        ExternalFunction configured for the bn_conv2dk1_skip kernel.

    Raises:
        ValueError: When ``skip_dtype`` is not ``np.uint8`` or ``np.int8``.
    """
    if skip_dtype == np.uint8:
        func_name = "conv2dk1_skip_ui8_ui8_i8"
        flags = ["-DREGULAR", "-DSCALAR", "-DUNSIGNED_SKIP"]
    elif skip_dtype == np.int8:
        func_name = "conv2dk1_skip_ui8_i8_i8"
        flags = ["-DREGULAR", "-DSCALAR"]
    else:
        raise ValueError(
            f"bn_conv2dk1_skip(): skip_dtype must be np.uint8 or np.int8, "
            f"got {skip_dtype}"
        )

    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(input_channels * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.int8]]
    skip_ty = np.ndarray[(input_width * output_channels,), np.dtype[skip_dtype]]
    return _make_extern(
        func_name,
        _kernel_source("conv/bn_conv2dk1_skip.cc"),
        [in_ty, wt_ty, out_ty, skip_ty, *_i32s(5)],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            alignments=_vector_loads(1),
            trace=Trace.whole_call(),
            roles=(In, Param, Out, In, *((Param,) * 5)),
            reference=bn_conv2dk1_skip_ref,
            acc_dtype=np.int32,
            reduction=input_channels,
            tolerance=_BN_EXACT,
            ops_per_call=2 * input_width * input_channels * output_channels
            + input_width * output_channels,
        ),
    )

bn_conv2dk3_dw

bn_conv2dk3_dw(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    stride: int = 1,
) -> ExternalFunction

Bottleneck depthwise 3x3 conv + ReLU kernel (uint8 in/out).

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels.	64
stride	int	Convolution stride (1 or 2).	1

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk3_dw kernel.

Raises:

Type	Description
ValueError	When stride is not 1 or 2.

Source code in python/iron/kernels/conv.py

def bn_conv2dk3_dw(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    stride: int = 1,
) -> ExternalFunction:
    """Bottleneck depthwise 3x3 conv + ReLU kernel (uint8 in/out).

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels.
        stride: Convolution stride (1 or 2).

    Returns:
        ExternalFunction configured for the bn_conv2dk3_dw kernel.

    Raises:
        ValueError: When ``stride`` is not 1 or 2.
    """
    if stride not in (1, 2):
        raise ValueError(f"bn_conv2dk3_dw(): stride must be 1 or 2, got {stride}")

    func_name = f"conv2dk3_dw_stride{stride}_relu_ui8_ui8"

    line_size = input_width * input_channels
    line_ty = np.ndarray[(line_size,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(3 * 3 * input_channels,), np.dtype[np.int8]]
    out_size = (input_width // stride) * output_channels
    out_ty = np.ndarray[(out_size,), np.dtype[np.uint8]]

    return _make_extern(
        func_name,
        _kernel_source("conv/bn_conv2dk3_dw.cc"),
        [line_ty, line_ty, line_ty, wt_ty, out_ty, *_i32s(8)],
        compile_flags=["-DREGULAR", "-DSCALAR", f"-DSTRIDE{stride}"]
        + _conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Param, Out, *((Param,) * 8)),
            # The weights are read unaligned.
            alignments=(
                _vector_32b(0, 1, 2, 4) if stride == 1 else _vector_loads(0, 1, 2, 4)
            ),
            reference=partial(bn_conv2dk3_dw_ref, stride=stride),
            acc_dtype=np.int32,
            reduction=9,
            tolerance=_BN_EXACT,
            ops_per_call=2 * 9 * out_size,
        ),
    )

bn_conv2dk1_relu_xy_pool_padded

bn_conv2dk1_relu_xy_pool_padded(
    input_width: int = 7,
    input_channels: int = 80,
    output_channels: int = 1280,
    weight_chunk_count: int | None = None,
) -> ExternalFunction

Fused 1x1 conv + ReLU + xy-pool with channel padding (int8 in, uint16 out).

A post-stage kernel that fuses a pointwise (1x1) convolution, ReLU activation, and global xy avg-pool into a single pass, with output channels padded to a DMA-friendly multiple. Sized for MobileNet V3's post-bottleneck stage where the final 1x1 expand-conv collapses the 7x7 feature map into a 1x1 vector.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	7
input_channels	int	Number of input channels.	80
output_channels	int	Logical output channels (e.g. 1280). Sets both the output buffer length AND, when weight_chunk_count is None, the weight buffer length (input_channels * output_channels).	1280
weight_chunk_count	int | None	Override the weight buffer's element count when the design streams weights in chunks (cascade/output-split). None means use the full input_channels * output_channels tile.	None

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the fused conv+relu+xy_pool kernel.

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_relu_xy_pool_padded(
    input_width: int = 7,
    input_channels: int = 80,
    output_channels: int = 1280,
    weight_chunk_count: int | None = None,
) -> ExternalFunction:
    """Fused 1x1 conv + ReLU + xy-pool with channel padding (int8 in, uint16 out).

    A post-stage kernel that fuses a pointwise (1x1) convolution, ReLU
    activation, and global xy avg-pool into a single pass, with output
    channels padded to a DMA-friendly multiple.  Sized for MobileNet V3's
    post-bottleneck stage where the final 1x1 expand-conv collapses the
    7x7 feature map into a 1x1 vector.

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Logical output channels (e.g. 1280).  Sets both
            the output buffer length AND, when ``weight_chunk_count`` is
            None, the weight buffer length (``input_channels * output_channels``).
        weight_chunk_count: Override the weight buffer's element count when
            the design streams weights in chunks (cascade/output-split).
            ``None`` means use the full ``input_channels * output_channels``
            tile.

    Returns:
        ExternalFunction configured for the fused conv+relu+xy_pool kernel.
    """
    wts_count = (
        weight_chunk_count
        if weight_chunk_count is not None
        else input_channels * output_channels
    )
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[np.int8]]
    wt_ty = np.ndarray[(wts_count,), np.dtype[np.int8]]
    out_ty = np.ndarray[(output_channels,), np.dtype[np.uint16]]
    return _make_extern(
        "conv2dk1_xy_pool_fused_relu_large_padded_i8_ui8",
        _kernel_source("conv/bn_conv2dk1_relu.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(8)],
        compile_flags=["-DSCALAR", "-DCONV_XYPOOL_FUSED_LARGE_PADDED", "-DINT8_ACT"],
        contract=KernelContract(
            alignments=_vector_loads(1),
            trace=Trace.whole_call(),
            roles=(In, Param, InOut, *((Param,) * 8)),
            reference=bn_conv2dk1_relu_xy_pool_padded_ref,
            initializers=((2, _zero_output),),
            acc_dtype=np.int32,
            reduction=input_channels,
            tolerance=_BN_EXACT,
            ops_per_call=2 * input_width * input_channels * output_channels,
        ),
    )

bn_conv2dk1_partial_put_i8

bn_conv2dk1_partial_put_i8(
    input_width: int = 7,
    input_channels: int = 80,
    weight_count: int = 4800,
    *,
    block_index: int = 13,
    oc_blocks: int = 1
) -> ExternalFunction

Cascade-PUT half of a width-split 1x1 conv on int8 activations.

The PUT tile of a two-tile cascade-split pointwise conv: consumes a width slice of the activation, multiplies against its weight half, and emits the partial sum onto the cascade stream (no separate output buffer — cascade-only). Sister of bn_conv2dk1_partial_get_relu_i8.

Currently defined in the .cc only for MobileNet V3's bn13 / bn14 (one wrapper symbol per block); block_index selects which. Generalising this to arbitrary block names would require adding a non-prefixed wrapper to bn_conv2dk1_i8.cc.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input slice.	7
input_channels	int	Number of input channels.	80
weight_count	int	Per-call weight chunk size in elements (the design streams weights in chunks; full weight tensor is shared across multiple kernel invocations).	4800
block_index	int	13 or 14; selects the per-block C++ wrapper.	13
oc_blocks	int	Consecutive 8-channel output blocks one call covers; the oc argument then indexes groups of oc_blocks.	1

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the PUT tile.

Raises:

Type	Description
ValueError	When block_index is not 13 or 14, or oc_blocks < 1.

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_partial_put_i8(
    input_width: int = 7,
    input_channels: int = 80,
    weight_count: int = 4800,
    *,
    block_index: int = 13,
    oc_blocks: int = 1,
) -> ExternalFunction:
    """Cascade-PUT half of a width-split 1x1 conv on int8 activations.

    The PUT tile of a two-tile cascade-split pointwise conv: consumes a
    width slice of the activation, multiplies against its weight half,
    and emits the partial sum onto the cascade stream (no separate
    output buffer — cascade-only).  Sister of
    [`bn_conv2dk1_partial_get_relu_i8`][iron.kernels.conv.bn_conv2dk1_partial_get_relu_i8].

    Currently defined in the .cc only for MobileNet V3's bn13 / bn14
    (one wrapper symbol per block); ``block_index`` selects which.
    Generalising this to arbitrary block names would require adding a
    non-prefixed wrapper to ``bn_conv2dk1_i8.cc``.

    Args:
        input_width: Spatial width of the input slice.
        input_channels: Number of input channels.
        weight_count: Per-call weight chunk size in elements (the design
            streams weights in chunks; full weight tensor is shared
            across multiple kernel invocations).
        block_index: ``13`` or ``14``; selects the per-block C++ wrapper.
        oc_blocks: Consecutive 8-channel output blocks one call covers;
            the ``oc`` argument then indexes groups of ``oc_blocks``.

    Returns:
        ExternalFunction configured for the PUT tile.

    Raises:
        ValueError: When ``block_index`` is not 13 or 14, or
            ``oc_blocks`` < 1.
    """
    _validate_bn_block_index(block_index, "bn_conv2dk1_partial_put_i8")
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[np.int8]]
    wt_ty = np.ndarray[(weight_count,), np.dtype[np.int8]]
    return _make_extern(
        f"bn{block_index}_1_conv2dk1_i8_ui8_partial_width_put_new",
        _kernel_source("conv/bn_conv2dk1_i8.cc"),
        [in_ty, wt_ty, *_i32s(7)],
        compile_flags=[f"-DBN{block_index}_1_PARTIAL_PUT_I8_CAS_WIDTH_NEW"]
        + _cas_oc_blocks_flags(oc_blocks),
    )

bn_conv2dk1_partial_get_relu_i8

bn_conv2dk1_partial_get_relu_i8(
    input_width: int = 7,
    input_channels: int = 80,
    output_channels: int = 480,
    weight_count: int = 4800,
    *,
    block_index: int = 13,
    oc_blocks: int = 1
) -> ExternalFunction

Cascade-GET half of a width-split 1x1 conv + ReLU on int8 activations.

The GET tile of a two-tile cascade-split pointwise conv: consumes the cascade partial sum from its sister PUT tile, finishes the dot product against its weight half, applies ReLU, and writes the full output buffer. Sister of bn_conv2dk1_partial_put_i8.

Currently defined in the .cc only for MobileNet V3's bn13 / bn14 (one wrapper symbol per block); block_index selects which.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input slice.	7
input_channels	int	Number of input channels.	80
output_channels	int	Number of output channels (full L1 output width).	480
weight_count	int	Per-call weight chunk size in elements.	4800
block_index	int	13 or 14; selects the per-block C++ wrapper.	13
oc_blocks	int	Consecutive 8-channel output blocks one call covers; the oc argument then indexes groups of oc_blocks.	1

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the GET tile.

Raises:

Type	Description
ValueError	When block_index is not 13 or 14, or oc_blocks < 1.

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_partial_get_relu_i8(
    input_width: int = 7,
    input_channels: int = 80,
    output_channels: int = 480,
    weight_count: int = 4800,
    *,
    block_index: int = 13,
    oc_blocks: int = 1,
) -> ExternalFunction:
    """Cascade-GET half of a width-split 1x1 conv + ReLU on int8 activations.

    The GET tile of a two-tile cascade-split pointwise conv: consumes
    the cascade partial sum from its sister PUT tile, finishes the dot
    product against its weight half, applies ReLU, and writes the full
    output buffer.  Sister of [`bn_conv2dk1_partial_put_i8`][iron.kernels.conv.bn_conv2dk1_partial_put_i8].

    Currently defined in the .cc only for MobileNet V3's bn13 / bn14
    (one wrapper symbol per block); ``block_index`` selects which.

    Args:
        input_width: Spatial width of the input slice.
        input_channels: Number of input channels.
        output_channels: Number of output channels (full L1 output width).
        weight_count: Per-call weight chunk size in elements.
        block_index: ``13`` or ``14``; selects the per-block C++ wrapper.
        oc_blocks: Consecutive 8-channel output blocks one call covers;
            the ``oc`` argument then indexes groups of ``oc_blocks``.

    Returns:
        ExternalFunction configured for the GET tile.

    Raises:
        ValueError: When ``block_index`` is not 13 or 14, or
            ``oc_blocks`` < 1.
    """
    _validate_bn_block_index(block_index, "bn_conv2dk1_partial_get_relu_i8")
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[np.int8]]
    wt_ty = np.ndarray[(weight_count,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    return _make_extern(
        f"bn{block_index}_1_conv2dk1_i8_ui8_partial_width_get_new",
        _kernel_source("conv/bn_conv2dk1_relu.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(9)],
        compile_flags=[f"-DBN{block_index}_1_PARTIAL_GET_I8_CAS_WIDTH_NEW"]
        + _cas_oc_blocks_flags(oc_blocks),
    )

bn_conv2dk3_dw_out_split

bn_conv2dk3_dw_out_split(
    input_width: int = 7,
    input_channels: int = 480,
    output_split_channels: int = 240,
    *,
    block_index: int = 13
) -> ExternalFunction

Depthwise 3x3 stride-1 conv with split output stream (uint8 in/out).

A variant of bn_conv2dk3_dw (stride=1) that writes its output to TWO separate buffers — the channel dimension is split in half so downstream cascade-PUT tiles can each consume one slice. Used by MobileNet V3's bn13 / bn14 depthwise stage to feed the L3 cascade.

Currently defined in the .cc only via per-block extern wrappers (BN13 or BN14 macro picks the symbol prefix); block_index selects which.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input.	7
input_channels	int	Number of input channels (== output channels — depthwise).	480
output_split_channels	int	Channels per output slice (half of input_channels for the typical 2-way split).	240
block_index	int	13 or 14; selects the per-block C++ wrapper.	13

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the split-output DW kernel.

Raises:

Type	Description
ValueError	When block_index is not 13 or 14.

Source code in python/iron/kernels/conv.py

def bn_conv2dk3_dw_out_split(
    input_width: int = 7,
    input_channels: int = 480,
    output_split_channels: int = 240,
    *,
    block_index: int = 13,
) -> ExternalFunction:
    """Depthwise 3x3 stride-1 conv with split output stream (uint8 in/out).

    A variant of [`bn_conv2dk3_dw`][iron.kernels.conv.bn_conv2dk3_dw] (stride=1) that writes its output
    to TWO separate buffers — the channel dimension is split in half so
    downstream cascade-PUT tiles can each consume one slice.  Used by
    MobileNet V3's bn13 / bn14 depthwise stage to feed the L3 cascade.

    Currently defined in the .cc only via per-block extern wrappers
    (BN13 or BN14 macro picks the symbol prefix); ``block_index`` selects
    which.

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels (== output channels —
            depthwise).
        output_split_channels: Channels per output slice (half of
            ``input_channels`` for the typical 2-way split).
        block_index: ``13`` or ``14``; selects the per-block C++ wrapper.

    Returns:
        ExternalFunction configured for the split-output DW kernel.

    Raises:
        ValueError: When ``block_index`` is not 13 or 14.
    """
    _validate_bn_block_index(block_index, "bn_conv2dk3_dw_out_split")
    if (
        input_width < 2
        or output_split_channels <= 0
        or input_channels != 2 * output_split_channels
        or output_split_channels % 8
    ):
        raise ValueError(
            "bn_conv2dk3_dw_out_split requires width >= 2 and equal channel "
            "halves divisible by 8"
        )
    line_size = input_width * input_channels
    line_ty = np.ndarray[(line_size,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(3 * 3 * input_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_split_channels,), np.dtype[np.uint8]]

    return _make_extern(
        f"bn{block_index}_conv2dk3_ui8_out_split",
        _kernel_source("conv/bn_conv2dk3_dw.cc"),
        [line_ty, line_ty, line_ty, wt_ty, out_ty, out_ty, *_i32s(8)],
        compile_flags=["-DSCALAR", f"-DBN{block_index}", "-DSTRIDE1_OUT_SPLIT"]
        + _conv_dimensions(input_width, input_channels, input_channels),
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Param, Out, Out, *((Param,) * 8)),
            alignments=_vector_32b(0, 1, 2, 4, 5),
            reference=bn_conv2dk3_dw_out_split_ref,
            acc_dtype=np.int32,
            reduction=9,
            tolerance=_BN_EXACT,
            ops_per_call=2 * 9 * input_width * input_channels,
        ),
    )

bn_conv2dk1_input_split_partial_put_ui8

bn_conv2dk1_input_split_partial_put_ui8(
    input_width: int = 7,
    input_channels: int = 240,
    weight_count: int = 9600,
    *,
    block_index: int = 13,
    oc_blocks: int = 1
) -> ExternalFunction

Input-split cascade-PUT half of a 1x1 conv on uint8 activations.

Like bn_conv2dk1_partial_put_i8 but consumes a CHANNEL slice (input-split) of a uint8 activation instead of a width slice of int8. Used by MobileNet V3's bn13 / bn14 L3 stage.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input slice.	7
input_channels	int	Number of input channels (one half of the full input after split).	240
weight_count	int	Per-call weight chunk size in elements.	9600
block_index	int	13 or 14; selects the per-block C++ wrapper.	13
oc_blocks	int	Consecutive 8-channel output blocks one call covers; the oc argument then indexes groups of oc_blocks.	1

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the input-split PUT tile.

Raises:

Type	Description
ValueError	When block_index is not 13 or 14, or oc_blocks < 1.

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_input_split_partial_put_ui8(
    input_width: int = 7,
    input_channels: int = 240,
    weight_count: int = 9600,
    *,
    block_index: int = 13,
    oc_blocks: int = 1,
) -> ExternalFunction:
    """Input-split cascade-PUT half of a 1x1 conv on uint8 activations.

    Like [`bn_conv2dk1_partial_put_i8`][iron.kernels.conv.bn_conv2dk1_partial_put_i8] but consumes a CHANNEL slice
    (input-split) of a uint8 activation instead of a width slice of int8.
    Used by MobileNet V3's bn13 / bn14 L3 stage.

    Args:
        input_width: Spatial width of the input slice.
        input_channels: Number of input channels (one half of the
            full input after split).
        weight_count: Per-call weight chunk size in elements.
        block_index: ``13`` or ``14``; selects the per-block C++ wrapper.
        oc_blocks: Consecutive 8-channel output blocks one call covers;
            the ``oc`` argument then indexes groups of ``oc_blocks``.

    Returns:
        ExternalFunction configured for the input-split PUT tile.

    Raises:
        ValueError: When ``block_index`` is not 13 or 14, or
            ``oc_blocks`` < 1.
    """
    _validate_bn_block_index(block_index, "bn_conv2dk1_input_split_partial_put_ui8")
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(weight_count,), np.dtype[np.int8]]
    return _make_extern(
        f"bn{block_index}_1_conv2dk1_ui8_ui8_input_split_partial_width_put_new",
        _kernel_source("conv/bn_conv2dk1_i8.cc"),
        [in_ty, wt_ty, *_i32s(7)],
        compile_flags=[
            f"-DBN{block_index}_1_INPUT_SPLIT_PARTIAL_PUT_UI8_UI8_CAS_WIDTH_NEW"
        ]
        + _cas_oc_blocks_flags(oc_blocks),
    )

bn_conv2dk1_input_split_partial_skip_get

bn_conv2dk1_input_split_partial_skip_get(
    input_width: int = 7,
    input_channels: int = 240,
    output_channels: int = 80,
    weight_count: int = 9600,
    *,
    block_index: int = 13,
    oc_blocks: int = 1
) -> ExternalFunction

Input-split cascade-GET half of a 1x1 conv + skip-add (uint8 in, int8 out).

The GET tile completes the cascade-split 1x1 + ReLU + residual add pattern: consumes the partial sum from its sister PUT tile, finishes the dot product, adds a skip row of int8 activations, and writes int8 output. Sister of bn_conv2dk1_input_split_partial_put_ui8.

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input slice.	7
input_channels	int	Number of input channels (one half after split).	240
output_channels	int	Final output channels.	80
weight_count	int	Per-call weight chunk size in elements.	9600
block_index	int	13 or 14; selects the per-block C++ wrapper.	13
oc_blocks	int	Consecutive 8-channel output blocks one call covers; the oc argument then indexes groups of oc_blocks.	1

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the input-split skip-GET tile.

Raises:

Type	Description
ValueError	When block_index is not 13 or 14, or oc_blocks < 1.

Source code in python/iron/kernels/conv.py

def bn_conv2dk1_input_split_partial_skip_get(
    input_width: int = 7,
    input_channels: int = 240,
    output_channels: int = 80,
    weight_count: int = 9600,
    *,
    block_index: int = 13,
    oc_blocks: int = 1,
) -> ExternalFunction:
    """Input-split cascade-GET half of a 1x1 conv + skip-add (uint8 in, int8 out).

    The GET tile completes the cascade-split 1x1 + ReLU + residual add
    pattern: consumes the partial sum from its sister PUT tile, finishes
    the dot product, adds a skip row of int8 activations, and writes int8
    output.  Sister of [`bn_conv2dk1_input_split_partial_put_ui8`][iron.kernels.conv.bn_conv2dk1_input_split_partial_put_ui8].

    Args:
        input_width: Spatial width of the input slice.
        input_channels: Number of input channels (one half after split).
        output_channels: Final output channels.
        weight_count: Per-call weight chunk size in elements.
        block_index: ``13`` or ``14``; selects the per-block C++ wrapper.
        oc_blocks: Consecutive 8-channel output blocks one call covers;
            the ``oc`` argument then indexes groups of ``oc_blocks``.

    Returns:
        ExternalFunction configured for the input-split skip-GET tile.

    Raises:
        ValueError: When ``block_index`` is not 13 or 14, or
            ``oc_blocks`` < 1.
    """
    _validate_bn_block_index(block_index, "bn_conv2dk1_input_split_partial_skip_get")
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(weight_count,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.int8]]
    skip_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.int8]]
    return _make_extern(
        f"bn_{block_index}_2_conv2dk1_ui8_i8_i8_scalar_input_split_partial_width_get_new",
        _kernel_source("conv/bn_conv2dk1_skip.cc"),
        [in_ty, wt_ty, out_ty, skip_ty, *_i32s(10)],
        compile_flags=[
            f"-DBN{block_index}_1_INPUT_SPLIT_PARTIAL_GET_UI8_I8_I8_CAS_WIDTH_NEW"
        ]
        + _cas_oc_blocks_flags(oc_blocks),
    )

bn_fc_relu_ui16_pad

bn_fc_relu_ui16_pad(
    input_channels: int = 1280,
    output_channels: int = 16,
    weight_chunk_count: int | None = None,
) -> ExternalFunction

Fully-connected layer (1x1 conv on (1,1,C)) + ReLU, uint16 in/out, with padding.

A post-stage FC kernel used by MobileNet V3's classifier head. Input is a (1,1,input_channels) feature vector held as uint16; output is output_channels uint16 logits. Weights stored in a padded layout (the input_channels_pad runtime arg selects the actual stride).

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (e.g. 1280).	1280
output_channels	int	Number of output channels per call (slice width, since the full FC is split across multiple tiles).	16
weight_chunk_count	int | None	Override the weight buffer's element count when the design streams weights in chunks (cascade/ping-pong). None means use the full input_channels * output_channels tile.	None

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the post-L2 FC kernel.

Source code in python/iron/kernels/conv.py

def bn_fc_relu_ui16_pad(
    input_channels: int = 1280,
    output_channels: int = 16,
    weight_chunk_count: int | None = None,
) -> ExternalFunction:
    """Fully-connected layer (1x1 conv on (1,1,C)) + ReLU, uint16 in/out, with padding.

    A post-stage FC kernel used by MobileNet V3's classifier head.  Input is
    a (1,1,input_channels) feature vector held as uint16; output is
    ``output_channels`` uint16 logits.  Weights stored in a padded layout
    (the ``input_channels_pad`` runtime arg selects the actual stride).

    Args:
        input_channels: Number of input channels (e.g. 1280).
        output_channels: Number of output channels per call (slice width,
            since the full FC is split across multiple tiles).
        weight_chunk_count: Override the weight buffer's element count when
            the design streams weights in chunks (cascade/ping-pong).
            ``None`` means use the full ``input_channels * output_channels``
            tile.

    Returns:
        ExternalFunction configured for the post-L2 FC kernel.
    """
    wts_count = (
        weight_chunk_count
        if weight_chunk_count is not None
        else input_channels * output_channels
    )
    in_ty = np.ndarray[(input_channels,), np.dtype[np.uint16]]
    wt_ty = np.ndarray[(wts_count,), np.dtype[np.int8]]
    out_ty = np.ndarray[(output_channels,), np.dtype[np.uint16]]
    return _make_extern(
        "post_L2_conv2dk1_relu_i16_ui16_pad",
        _kernel_source("conv/bn_conv2dk1_relu.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(5)],
        compile_flags=["-DSCALAR", "-DPOSTL2_PAD", "-DUINT16_ACT"],
        contract=KernelContract(
            alignments=_vector_loads(1),
            trace=Trace.whole_call(),
            roles=(In, Param, Out, *((Param,) * 5)),
            reference=bn_fc_relu_ui16_pad_ref,
            acc_dtype=np.int32,
            reduction=input_channels,
            tolerance=_BN_EXACT,
            ops_per_call=2 * input_channels * output_channels,
        ),
    )

Activation functions

Activation kernel factories and NumPy reference implementations.

softmax

softmax(tile_size: int = 1024) -> ExternalFunction

Softmax activation kernel for bf16 tiles.

Parameters:

Name	Type	Description	Default
tile_size	int	Number of elements per tile, passed at run time; a positive multiple of 32, the kernel's vector step.	1024

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the softmax kernel.

Source code in python/iron/kernels/activation.py

def softmax(tile_size: int = 1024) -> ExternalFunction:
    """Softmax activation kernel for bf16 tiles.

    Args:
        tile_size: Number of elements per tile, passed at run time; a
            positive multiple of 32, the kernel's vector step.

    Returns:
        ExternalFunction configured for the softmax kernel.
    """
    _require_vector_alignment("softmax", tile_size, _RUNTIME_VECTOR_WIDTH)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "softmax_bf16",
        "softmax.cc",
        [tile_ty, tile_ty, np.int32],
        contract=_unary_lut_contract(
            lambda x: softmax_ref(x, tile_size=tile_size),
            count=tile_size,
            # Softmax outputs sum to 1 over the tile, so a typical element is
            # about 1 / tile_size and the generic LUT floor of 0.05 would
            # accept an all-zero output. A tenth of an average element still
            # covers the exp LUT's underflow on the far tail, while an
            # unwritten tile mismatches on most elements.
            tolerance=Tolerance.relative(
                0.04,
                0.1 / tile_size,
                note="AIE2P exp instruction through the softmax normalisation, "
                "measured on npu2 at 2.94e-2 relative; atol = 0.1 / tile_size "
                "so an unwritten (all-zero) tile fails, since every softmax "
                "output is below a generic absolute floor",
            ),
            # softmax_aie2p.h sets conv_even itself; the aie2 LUT path does not.
            setup=conv_even if _tuned_arch() == "aie2" else None,
        ),
    )

gelu

gelu(tile_size: int = 1024) -> ExternalFunction

GELU activation kernel (tanh approximation) for bf16 tiles (must be 1024).

Source code in python/iron/kernels/activation.py

def gelu(tile_size: int = 1024) -> ExternalFunction:
    """GELU activation kernel (tanh approximation) for bf16 tiles (must be 1024)."""
    return _bf16_lut_factory(
        "gelu",
        "gelu_bf16",
        "gelu.cc",
        tile_size,
        arg_arity=2,
        contract=_unary_lut_contract(
            gelu_ref,
            count=False,
            tolerance=_gelu_tolerance(),
            stack_bytes=_GELU_AIE2_STACK_BYTES if _tuned_arch() == "aie2" else None,
        ),
    )

silu

silu(
    tile_size: int = 1024, use_lut: bool = False
) -> ExternalFunction

SiLU (Swish) activation kernel for bf16 tiles (must be 1024).

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call (must be 1024).	1024
use_lut	bool	Compute tanh from the interpolated LUT rather than AIE2P's vtanh instruction, which makes this 7.5x closer to the true function and judged against an exact model of it. Moot on aie2.	False

Source code in python/iron/kernels/activation.py

def silu(tile_size: int = 1024, use_lut: bool = False) -> ExternalFunction:
    """SiLU (Swish) activation kernel for bf16 tiles (must be 1024).

    Args:
        tile_size: Elements per call (must be 1024).
        use_lut: Compute tanh from the interpolated LUT rather than AIE2P's
            vtanh instruction, which makes this 7.5x closer to the true
            function and judged against an exact model of it. Moot on aie2.
    """
    return _bf16_lut_factory(
        "silu",
        "silu_bf16",
        "silu.cc",
        tile_size,
        arg_arity=2,
        contract=_unary_lut_contract(
            silu_ref,
            count=False,
            use_lut=use_lut,
            elementwise=silu_table_ref if _tuned_arch() == "aie2p" else silu_lut_ref,
            tolerance=_vtanh_family_tolerance("silu"),
        ),
        use_lut_tanh=use_lut,
    )

silu_sized

silu_sized(tile_size: int = 1024) -> ExternalFunction

SiLU (Swish) for bf16 tiles, with a compiled-in element count.

Runtime-size sibling of silu; design keeps the (in, out, size) ABI. Positive whole vectors are required (16 on aie2, 32 on aie2p).

Source code in python/iron/kernels/activation.py

def silu_sized(tile_size: int = 1024) -> ExternalFunction:
    """SiLU (Swish) for bf16 tiles, with a compiled-in element count.

    Runtime-size sibling of [`silu`][iron.kernels.activation.silu]; design
    keeps the ``(in, out, size)`` ABI. Positive whole vectors are required
    (16 on aie2, 32 on aie2p).
    """
    width = _arch_traits().bf16_lanes
    _require_vector_alignment("silu_sized", tile_size, width)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "silu_bf16_size",
        "silu.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DSILU_ELEMS={tile_size}"],
        contract=_unary_lut_contract(
            silu_ref, count=tile_size, tolerance=_vtanh_family_tolerance("silu")
        ),
    )

gelu_sized

gelu_sized(tile_size: int = 1024) -> ExternalFunction

GELU (tanh approx) for bf16 tiles, with a compiled-in element count.

Runtime-size sibling of gelu; design keeps the (in, out, size) ABI. Positive whole vectors only: multiples of 16 on aie2 or 32 on aie2p.

Source code in python/iron/kernels/activation.py

def gelu_sized(tile_size: int = 1024) -> ExternalFunction:
    """GELU (tanh approx) for bf16 tiles, with a compiled-in element count.

    Runtime-size sibling of [`gelu`][iron.kernels.activation.gelu]; design
    keeps the ``(in, out, size)`` ABI. Positive whole vectors only: multiples
    of 16 on aie2 or 32 on aie2p.
    """
    _require_vector_alignment("gelu_sized", tile_size, _arch_traits().bf16_lanes)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "gelu_bf16_size",
        "gelu.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DGELU_ELEMS={tile_size}"],
        contract=_unary_lut_contract(
            gelu_ref,
            count=tile_size,
            tolerance=_gelu_tolerance(),
            stack_bytes=_GELU_AIE2_STACK_BYTES if _tuned_arch() == "aie2" else None,
        ),
    )

swiglu

swiglu(
    tile_size: int = 1024, use_lut: bool = False
) -> ExternalFunction

SwiGLU gated activation kernel for bf16 tiles (must be 1024).

out = (x * w1) * silu(x * w2); see swiglu_ref.

Source code in python/iron/kernels/activation.py

def swiglu(tile_size: int = 1024, use_lut: bool = False) -> ExternalFunction:
    """SwiGLU gated activation kernel for bf16 tiles (must be 1024).

    ``out = (x * w1) * silu(x * w2)``; see [`swiglu_ref`][iron.kernels.activation.swiglu_ref].
    """
    use_lut_model = use_lut or not _arch_traits().native_tanh
    return _bf16_lut_factory(
        "swiglu",
        "swiglu_bf16",
        "swiglu.cc",
        tile_size,
        arg_arity=4,
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, In, In, Out),
            reference=(
                swiglu_ref
                if not use_lut_model
                else swiglu_table_ref if _tuned_arch() == "aie2p" else swiglu_lut_ref
            ),
            acc_dtype=bfloat16,
            tolerance=(
                _LUT_MODEL_TOLERANCE
                if use_lut_model
                else _vtanh_family_tolerance("swiglu")
            ),
            ops_per_call=6 * tile_size,
            uses_lut=True,
        ),
        use_lut_tanh=use_lut,
    )

bf16_exp

bf16_exp(tile_size: int = 1024) -> ExternalFunction

Element-wise exponential kernel for bf16 tiles (must be 1024).

Computes exp(clip(x, -88, 88)): the kernel saturates rather than overflowing for real inputs, including infinities. On AIE2P a range-reduced polynomial replaces the table and rounds every result correctly, subnormal outputs included. See bf16_exp_ref for why that clamp matches the AIE2 table's domain.

Source code in python/iron/kernels/activation.py

def bf16_exp(tile_size: int = 1024) -> ExternalFunction:
    """Element-wise exponential kernel for bf16 tiles (must be 1024).

    Computes ``exp(clip(x, -88, 88))``: the kernel saturates rather than
    overflowing for real inputs, including infinities. On AIE2P a
    range-reduced polynomial replaces the table and rounds every result
    correctly, subnormal outputs included. See
    [`bf16_exp_ref`][iron.kernels.activation.bf16_exp_ref] for why that
    clamp matches the AIE2 table's domain.
    """
    return _bf16_lut_factory(
        "bf16_exp",
        "exp_bf16_1024",
        "bf16_exp.cc",
        tile_size,
        arg_arity=2,
        contract=_unary_lut_contract(
            bf16_exp_ref,
            count=False,
            # Only aie2's tuned branch reaches getExpBf16. The other computes a
            # range-reduced polynomial, which this model does not describe,
            # so it keeps the true-function reference and a measured bound,
            # whether or not there is a tanh instruction.
            elementwise=bf16_exp_lut_ref if _tuned_arch() == "aie2" else None,
            lut_tolerance=_EXP_LUT_TOLERANCE,
            use_lut=True,
            tolerance=_EXP_POLY_TOLERANCE,
            # The polynomial sets conv_even itself.
            setup=conv_even if _tuned_arch() == "aie2" else None,
            stack_bytes=(
                None if _tuned_arch() == "aie2" else _BF16_EXP_POLY_STACK_BYTES
            ),
        ),
    )

exp2f_vec

exp2f_vec(
    tile_size: int = 1024, min_x: float = -111.0
) -> ExternalFunction

Software f32 2**x kernel: a minimax poly, not a LUT.

A float32-output alternative to bf16_exp with a separately configurable input domain. See aie_kernels/activation/exp2f_vec.cc for the accuracy rationale: 9.2e-6 relative error on aie2p, 8.9e-5 on aie2.

The same source builds for aie2.

Parameters:

Name	Type	Description	Default
tile_size	int	Number of elements per tile; must be a multiple of 16 (the kernel's vector width).	1024
min_x	float	Input is clamped to this before evaluation. -126 is the hard floor (one f32 exponent field); on aie2p the kernel holds its accuracy down to it.	-111.0

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the exp2f_vec kernel.

Raises:

Type	Description
ValueError	If tile_size is not a multiple of 16, or min_x is below -126.

Source code in python/iron/kernels/activation.py

def exp2f_vec(tile_size: int = 1024, min_x: float = -111.0) -> ExternalFunction:
    """Software f32 ``2**x`` kernel: a minimax poly, not a LUT.

    A float32-output alternative to [`bf16_exp`]
    [iron.kernels.activation.bf16_exp] with a separately configurable
    input domain. See ``aie_kernels/activation/exp2f_vec.cc`` for the
    accuracy rationale: 9.2e-6 relative error on aie2p, 8.9e-5 on aie2.

    The same source builds for aie2.

    Args:
        tile_size: Number of elements per tile; must be a multiple of 16
            (the kernel's vector width).
        min_x: Input is clamped to this before evaluation. -126 is the
            hard floor (one f32 exponent field); on aie2p the kernel holds
            its accuracy down to it.

    Returns:
        ExternalFunction configured for the exp2f_vec kernel.

    Raises:
        ValueError: If tile_size is not a multiple of 16, or min_x is
            below -126.
    """
    if tile_size % 16 != 0:
        raise ValueError(
            f"exp2f_vec: tile_size must be a multiple of 16, got {tile_size}"
        )
    if min_x < -126.0:
        raise ValueError(
            f"exp2f_vec: min_x must be >= -126 (the kernel builds 2**k in the "
            f"f32 exponent field, whose smallest normal exponent is -126), "
            f"got {min_x}"
        )
    source = _kernel_source("activation/exp2f_vec.cc")
    tile_ty = np.ndarray[(tile_size,), np.dtype[np.float32]]
    return _make_extern(
        "exp2f_vec_f32",
        source,
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DEXP2F_VEC_MIN_X={float(min_x)!r}f"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            # The aie2p branch sets conv_even itself.
            setup=None if _tuned_arch() == "aie2p" else conv_even,
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            reference=lambda x: exp2f_vec_ref(x, min_x=min_x),
            acc_dtype=np.float32,
            tolerance=Tolerance.relative(
                1e-3,
                note="measured 9.2e-6 on aie2p and 8.9e-5 on aie2; clamping "
                "[127.999, 128) costs up to 7.8e-4",
            ),
            # aiecc measured 1984 on aie2p (448 portable); remarks gives the
            # kernel 1792 on aie2 (832 portable).
            stack_bytes={"aie2": 2048, "aie2p": 2048}.get(_tuned_arch()),
        ),
    )

tanh

tanh(
    tile_size: int = 1024, use_lut: bool = False
) -> ExternalFunction

Tanh for bf16 tiles of a positive multiple of 32 elements.

The count is compiled in; retain tile_size as a trailing int argument (e.g. via transform_parallel(pass_size_to_kernel=True)).

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call (a positive multiple of 32).	1024
use_lut	bool	Compute tanh from the interpolated LUT rather than AIE2P's vtanh instruction. Moot on aie2, which only has the LUT. See tanh_lut_ref for what the LUT computes and why it is the more accurate of the two.	False

Source code in python/iron/kernels/activation.py

def tanh(tile_size: int = 1024, use_lut: bool = False) -> ExternalFunction:
    """Tanh for bf16 tiles of a positive multiple of 32 elements.

    The count is compiled in; retain
    ``tile_size`` as a trailing ``int`` argument (e.g. via
    ``transform_parallel(pass_size_to_kernel=True)``).

    Args:
        tile_size: Elements per call (a positive multiple of 32).
        use_lut: Compute tanh from the interpolated LUT rather than AIE2P's
            vtanh instruction. Moot on aie2, which only has the LUT. See
            [`tanh_lut_ref`][iron.kernels.activation.tanh_lut_ref] for what
            the LUT computes and why it is the more accurate of the two.
    """
    _require_vector_alignment("tanh", tile_size, _RUNTIME_VECTOR_WIDTH)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "tanh_bf16",
        "tanh.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DTANH_ELEMS={tile_size}"],
        contract=_unary_lut_contract(
            tanh_ref,
            count=tile_size,
            use_lut=use_lut,
            elementwise=tanh_lut_ref,
            # aie2 never reaches this: _unary_lut_contract swaps in the
            # model tolerance there, since the LUT is its only tanh.
            tolerance=_VTANH_TOLERANCE,
        ),
        use_lut_tanh=use_lut,
    )

sigmoid

sigmoid(
    tile_size: int = 1024, use_lut: bool = False
) -> ExternalFunction

Sigmoid for bf16 tiles of a positive multiple of 32 elements.

The count is compiled in; retain tile_size as a trailing ABI argument.

Source code in python/iron/kernels/activation.py

def sigmoid(tile_size: int = 1024, use_lut: bool = False) -> ExternalFunction:
    """Sigmoid for bf16 tiles of a positive multiple of 32 elements.

    The count is compiled in; retain ``tile_size`` as a trailing ABI argument.
    """
    _require_vector_alignment("sigmoid", tile_size, _RUNTIME_VECTOR_WIDTH)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "sigmoid_bf16",
        "sigmoid.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DSIGMOID_ELEMS={tile_size}"],
        contract=_unary_lut_contract(
            sigmoid_ref,
            count=tile_size,
            use_lut=use_lut,
            elementwise=(
                sigmoid_table_ref if _tuned_arch() == "aie2p" else sigmoid_lut_ref
            ),
            tolerance=_vtanh_family_tolerance("sigmoid"),
        ),
        use_lut_tanh=use_lut,
    )

leaky_relu

leaky_relu(tile_size: int = 1024) -> ExternalFunction

Leaky ReLU for bf16 tiles of at least 64 elements, in multiples of 32.

The count is compiled in, but the ABI retains (tile_size, alpha) as trailing int/bfloat16 arguments. The slope remains runtime-valued.

Source code in python/iron/kernels/activation.py

def leaky_relu(tile_size: int = 1024) -> ExternalFunction:
    """Leaky ReLU for bf16 tiles of at least 64 elements, in multiples of 32.

    The count is compiled in, but the ABI retains ``(tile_size, alpha)`` as
    trailing ``int``/``bfloat16`` arguments. The slope remains runtime-valued.
    """
    if tile_size < 64 or tile_size % _RUNTIME_VECTOR_WIDTH:
        raise ValueError(
            "leaky_relu: tile_size must be a multiple of "
            f"{_RUNTIME_VECTOR_WIDTH} and at least 64, got {tile_size}"
        )
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "leaky_relu_bf16",
        "leaky_relu.cc",
        [tile_ty, tile_ty, np.int32, bfloat16],
        compile_flags=[f"-DLEAKY_RELU_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, Out, Param, Param),
            parameter_bindings=((2, tile_size),),
            reference=leaky_relu_ref,
            acc_dtype=bfloat16,
            # max(x, alpha*x) introduces exactly one rounding, on alpha*x.
            # Measured bit-exact on npu2 over 262144 elements at alpha=0.5,
            # which is a power of two and so rounds trivially; one ulp covers
            # an alpha that does not. The 0.03/0.05 with a 2% budget this
            # replaces was three orders of magnitude looser than the
            # arithmetic allows.
            tolerance=Tolerance.bf16_ulps(
                1, note="one rounding, on alpha*x; measured bit-exact at alpha=0.5"
            ),
        ),
    )

relu_ref

relu_ref(x)

Numpy reference for a ReLU kernel — element-wise max(x, 0).

Exact; tolerance comparison is not needed. See aie.utils.verify for the relaxed bf16/LUT-style comparators most kernels here want.

Source code in python/iron/kernels/activation.py

def relu_ref(x):
    """Numpy reference for a ReLU kernel — element-wise `max(x, 0)`.

    Exact; tolerance comparison is not needed.  See `aie.utils.verify`
    for the relaxed bf16/LUT-style comparators most kernels here want.
    """
    return np.maximum(x.astype(np.float32), 0.0).astype(x.dtype)

silu_ref

silu_ref(x)

Numpy reference for silu (Swish) — x * sigmoid(x).

LUT-approximation territory; pair with rtol=0.128 (the default in count_mismatches) when verifying.

Source code in python/iron/kernels/activation.py

def silu_ref(x):
    """Numpy reference for [`silu`][iron.kernels.activation.silu] (Swish) — ``x * sigmoid(x)``.

    LUT-approximation territory; pair with ``rtol=0.128`` (the default
    in `count_mismatches`) when verifying.
    """
    xf = _no_neg_inf(_f64(x))
    with np.errstate(over="ignore"):
        return _rounded(xf / (1.0 + np.exp(-xf)), x)

gelu_ref

gelu_ref(x)

Numpy reference for gelu.

Tanh approximation 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))). Matches the C++ kernel's tanh-GELU formula. It is evaluated in float64: in float32, 1 + tanh cancels for x below about -4.5 and leaves values over ten times too large.

Source code in python/iron/kernels/activation.py

def gelu_ref(x):
    """Numpy reference for [`gelu`][iron.kernels.activation.gelu].

    Tanh approximation ``0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))``.
    Matches the C++ kernel's tanh-GELU formula. It is evaluated in float64:
    in float32, ``1 + tanh`` cancels for x below about -4.5 and leaves
    values over ten times too large.
    """
    xf = _no_neg_inf(x.astype(np.float64))
    with np.errstate(over="ignore"):
        inner = math.sqrt(2.0 / math.pi) * (xf + 0.044715 * xf**3)
    return (0.5 * xf * (1.0 + np.tanh(inner))).astype(x.dtype)

bf16_exp_lut_ref

bf16_exp_lut_ref(x)

Model of getExpBf16, the LUT exponential the bf16_exp kernel uses.

The kernel clamps to +/-_EXP_BF16_CLAMP, converts to Q8 with a floor (bfloat16_to_int(x, 8)), then reads the byte halves of that fixed-point key as two table indices and multiplies: exp(x) = exp(int) * exp(frac). Both tables hold exactly bfloat16(exp(.)), so they are written here as that rule rather than as 512 opaque floats -- the unreachable middle of the integer table (keys the clamp cannot produce) is the only part that is not an exponential, and it is never read.

The product is exact in f32 (two bf16 operands), so the model differs from the device only where the hardware flushes the one subnormal table entry, bfloat16(exp(-88)).

Source code in python/iron/kernels/activation.py

def bf16_exp_lut_ref(x):
    """Model of ``getExpBf16``, the LUT exponential the bf16_exp kernel uses.

    The kernel clamps to ``+/-_EXP_BF16_CLAMP``, converts to Q8 with a floor
    (``bfloat16_to_int(x, 8)``), then reads the byte halves of that fixed-point
    key as two table indices and multiplies: ``exp(x) = exp(int) * exp(frac)``.
    Both tables hold exactly ``bfloat16(exp(.))``, so they are written here as
    that rule rather than as 512 opaque floats -- the unreachable middle of the
    integer table (keys the clamp cannot produce) is the only part that is not
    an exponential, and it is never read.

    The product is exact in f32 (two bf16 operands), so the model differs from
    the device only where the hardware flushes the one subnormal table entry,
    ``bfloat16(exp(-88))``.
    """
    xf = np.clip(np.asarray(x).astype(np.float32), -_EXP_BF16_CLAMP, _EXP_BF16_CLAMP)
    key = np.floor(xf * 256.0).astype(np.int32).astype(np.int16).astype(np.uint16)
    i = np.arange(256)
    ilut = np.where(
        i <= 88,
        np.exp(np.minimum(i, 88.0)),
        np.where(i >= 168, np.exp(i - 256.0), np.exp(88.0)),
    )
    ilut = np.asarray(ilut, np.float32).astype(bfloat16).astype(np.float32)
    flut = np.exp(np.arange(256) / 256.0).astype(np.float32).astype(bfloat16)
    return (ilut[key >> 8] * flut.astype(np.float32)[key & 255]).astype(np.float32)

sigmoid_lut_ref

sigmoid_lut_ref(x)

Model of sigmoid's LUT build on aie2.

AIE2P's reads a table of its own; see sigmoid_table_ref. This follows activation/sigmoid.cc's aie2 branch step for step: x/2 is exact (0.5 is a power of two), the accumulator overload of tanh_bf16_v16 narrows to bf16 before the table, and the +1 and *0.5 stay in the accumulator so there is a single store rounding at the end.

Source code in python/iron/kernels/activation.py

def sigmoid_lut_ref(x):
    """Model of [`sigmoid`][iron.kernels.activation.sigmoid]'s LUT build on aie2.

    AIE2P's reads a table of its own; see
    [`sigmoid_table_ref`][iron.kernels.activation.sigmoid_table_ref]. This
    follows activation/sigmoid.cc's aie2 branch step for step: ``x/2`` is
    exact (0.5 is a power of two), the accumulator overload of ``tanh_bf16_v16`` narrows to bf16
    before the table, and the ``+1`` and ``*0.5`` stay in the accumulator so
    there is a single store rounding at the end.
    """
    xf = np.asarray(x).astype(np.float32)
    t = np.asarray(tanh_lut_ref(_bf16(xf * 0.5)), np.float32)
    return _bf16((t + 1.0) * 0.5).astype(np.asarray(x).dtype)

sigmoid_table_ref

sigmoid_table_ref(x)

Model of AIE2P's sigmoid built with use_lut=True.

activation/sigmoid.cc reads its own table there: getTanhBf16's segments rewritten for 0.5 + 0.5 * tanh(x/2), so segment e has slope slope[e] / 4 and offset 0.5 + 0.5 * offset[e] and covers 0.5 of x. x is clamped to [-8, 8 - 1/32]; the one rounding is the store to bf16, as in tanh_lut_ref. sigmoid_lut_ref also rounds tanh to bf16 before 0.5 * (1 + t), which makes it 0 over [-7.5, -6.9], where this is not.

Source code in python/iron/kernels/activation.py

def sigmoid_table_ref(x):
    """Model of AIE2P's [`sigmoid`][iron.kernels.activation.sigmoid] built with ``use_lut=True``.

    activation/sigmoid.cc reads its own table there: getTanhBf16's segments
    rewritten for ``0.5 + 0.5 * tanh(x/2)``, so segment ``e`` has slope
    ``slope[e] / 4`` and offset ``0.5 + 0.5 * offset[e]`` and covers 0.5 of x.
    x is clamped to ``[-8, 8 - 1/32]``; the one rounding is the store to bf16,
    as in [`tanh_lut_ref`][iron.kernels.activation.tanh_lut_ref].
    [`sigmoid_lut_ref`][iron.kernels.activation.sigmoid_lut_ref] also rounds
    ``tanh`` to bf16 before ``0.5 * (1 + t)``, which makes it 0 over
    ``[-7.5, -6.9]``, where this is not.
    """
    xf = np.clip(np.asarray(x).astype(np.float32), -8.0, 8.0 - 1.0 / 32)
    e = np.clip(np.floor(xf * 2.0).astype(np.int64), -16, 15) + 16
    slope = np.asarray(_TANH_LUT_SLOPE, np.float32)[e] / 4
    offset = 0.5 + 0.5 * np.asarray(_TANH_LUT_OFFSET, np.float32)[e]
    return (slope * xf + offset).astype(bfloat16).astype(np.asarray(x).dtype)

silu_lut_ref

silu_lut_ref(x)

Model of silu's LUT build on aie2.

AIE2P's multiplies by its sigmoid table instead; see silu_table_ref. activation/silu.cc narrows the sigmoid factor to bf16 before the final multiply, so that rounding is modelled too, not folded away. The sigmoid is exactly 0 from x = -8 down, and x is clamped there before the multiply, so -inf gives 0 rather than NaN. A subnormal product is flushed to zero.

Source code in python/iron/kernels/activation.py

def silu_lut_ref(x):
    """Model of [`silu`][iron.kernels.activation.silu]'s LUT build on aie2.

    AIE2P's multiplies by its sigmoid table instead; see
    [`silu_table_ref`][iron.kernels.activation.silu_table_ref].
    activation/silu.cc narrows the sigmoid factor to bf16 before the final
    multiply, so that rounding is modelled too, not folded away. The sigmoid
    is exactly 0 from x = -8 down, and x is clamped there before the multiply,
    so -inf gives 0 rather than NaN. A subnormal product is flushed to zero.
    """
    xf = np.asarray(x).astype(np.float32)
    sig = np.asarray(sigmoid_lut_ref(xf), np.float32)
    return _bf16_ftz(np.maximum(xf, -8.0) * sig).astype(np.asarray(x).dtype)

silu_table_ref

silu_table_ref(x)

Model of AIE2P's silu built with use_lut=True.

silu_lut_ref with AIE2P's sigmoid table, sigmoid_table_ref, as the bf16 factor.

Source code in python/iron/kernels/activation.py

def silu_table_ref(x):
    """Model of AIE2P's [`silu`][iron.kernels.activation.silu] built with ``use_lut=True``.

    [`silu_lut_ref`][iron.kernels.activation.silu_lut_ref] with AIE2P's
    sigmoid table,
    [`sigmoid_table_ref`][iron.kernels.activation.sigmoid_table_ref], as the
    bf16 factor.
    """
    xf = np.asarray(x).astype(np.float32)
    sig = np.asarray(sigmoid_table_ref(xf), np.float32)
    return _bf16_ftz(np.maximum(xf, -8.0) * sig).astype(np.asarray(x).dtype)

swiglu_lut_ref

swiglu_lut_ref(x, w1, w2)

Model of swiglu's LUT build on aie2.

AIE2P's reads its sigmoid table instead; see swiglu_table_ref. activation/swiglu.cc narrows after every multiply -- x*w1, x*w2, the sigmoid factor and the silu product each land in a bf16 register before the next step -- which is what this reproduces. x*w2 is clamped at -8 before its multiply, as in silu_lut_ref, and where the silu product is 0 the output is 0, so an overflowed x*w1 does not make inf * 0. Each subnormal product is flushed to zero, so a subnormal x*w2 zeroes the output through the gate.

Source code in python/iron/kernels/activation.py

def swiglu_lut_ref(x, w1, w2):
    """Model of [`swiglu`][iron.kernels.activation.swiglu]'s LUT build on aie2.

    AIE2P's reads its sigmoid table instead; see
    [`swiglu_table_ref`][iron.kernels.activation.swiglu_table_ref].
    activation/swiglu.cc narrows after every multiply -- ``x*w1``, ``x*w2``, the
    sigmoid factor and the silu product each land in a bf16 register before
    the next step -- which is what this reproduces. ``x*w2`` is clamped at -8
    before its multiply, as in silu_lut_ref, and where the silu product is 0
    the output is 0, so an overflowed ``x*w1`` does not make ``inf * 0``.
    Each subnormal product is flushed to zero, so a subnormal ``x*w2`` zeroes
    the output through the gate.
    """
    return _swiglu_model(x, w1, w2, sigmoid_lut_ref)

swiglu_table_ref

swiglu_table_ref(x, w1, w2)

Model of AIE2P's swiglu built with use_lut=True.

swiglu_lut_ref with AIE2P's sigmoid table, sigmoid_table_ref, as the bf16 factor.

Source code in python/iron/kernels/activation.py

def swiglu_table_ref(x, w1, w2):
    """Model of AIE2P's [`swiglu`][iron.kernels.activation.swiglu] built with ``use_lut=True``.

    [`swiglu_lut_ref`][iron.kernels.activation.swiglu_lut_ref] with AIE2P's
    sigmoid table,
    [`sigmoid_table_ref`][iron.kernels.activation.sigmoid_table_ref], as the
    bf16 factor.
    """
    return _swiglu_model(x, w1, w2, sigmoid_table_ref)

tanh_lut_ref

tanh_lut_ref(x)

Numpy model of getTanhBf16, the interpolated-LUT tanh.

The kernel clamps x to the table's range [-4, 4 - 1/64], then evaluates slope[e] * x + offset[e] for e = floor(4x) + 16: 32 segments of width 0.25 over [-4, 4). The end segments are the constants -1 and +1, so the clamp changes no finite result, and +-inf gives +-1 rather than 0 * inf. The product is exact in f32 (bf16 carries 8 mantissa bits and 8 + 8 < 24), so the only rounding is the accumulator's store back to bf16, which is why the build using this is judged at one ulp rather than a percentage.

This is what tanh computes with use_lut=True, and what it always computes on aie2. The default aie2p build uses the vtanh instruction instead, which is a coarser approximation with no published spec, so it is judged against tanh_ref and a measured bound.

Source code in python/iron/kernels/activation.py

def tanh_lut_ref(x):
    """Numpy model of ``getTanhBf16``, the interpolated-LUT tanh.

    The kernel clamps x to the table's range ``[-4, 4 - 1/64]``, then evaluates
    ``slope[e] * x + offset[e]`` for ``e = floor(4x) + 16``: 32 segments of
    width 0.25 over ``[-4, 4)``. The end segments are the constants -1 and +1,
    so the clamp changes no finite result, and +-inf gives +-1 rather than
    ``0 * inf``. The product is exact in f32 (bf16 carries 8 mantissa bits and
    8 + 8 < 24), so the only rounding is the accumulator's store back to bf16,
    which is why the build using this is judged at one ulp rather than a
    percentage.

    This is what [`tanh`][iron.kernels.activation.tanh] computes with
    ``use_lut=True``, and what it always computes on aie2. The default aie2p
    build uses the ``vtanh`` instruction instead, which is a coarser
    approximation with no published spec, so it is judged against
    [`tanh_ref`][iron.kernels.activation.tanh_ref] and a measured bound.
    """
    xf = np.clip(np.asarray(x).astype(np.float32), -4.0, 4.0 - 1.0 / 64)
    e = np.clip(np.floor(xf * 4.0).astype(np.int64), -16, 15) + 16
    slope = np.asarray(_TANH_LUT_SLOPE, np.float32)[e]
    offset = np.asarray(_TANH_LUT_OFFSET, np.float32)[e]
    return (slope * xf + offset).astype(bfloat16).astype(np.asarray(x).dtype)

tanh_ref

tanh_ref(x)

Numpy reference for tanh — element-wise tanh(x).

LUT/native-approximation territory; pair with rtol=0.128 when verifying.

Source code in python/iron/kernels/activation.py

def tanh_ref(x):
    """Numpy reference for [`tanh`][iron.kernels.activation.tanh] — element-wise ``tanh(x)``.

    LUT/native-approximation territory; pair with ``rtol=0.128`` when verifying.
    """
    return _rounded(np.tanh(_f64(x)), x)

sigmoid_ref

sigmoid_ref(x)

Numpy reference for sigmoid — 1 / (1 + exp(-x)).

LUT-approximation territory; pair with rtol=0.128 when verifying.

Source code in python/iron/kernels/activation.py

def sigmoid_ref(x):
    """Numpy reference for [`sigmoid`][iron.kernels.activation.sigmoid] — ``1 / (1 + exp(-x))``.

    LUT-approximation territory; pair with ``rtol=0.128`` when verifying.
    """
    with np.errstate(over="ignore"):
        return _rounded(1.0 / (1.0 + np.exp(-_f64(x))), x)

leaky_relu_ref

leaky_relu_ref(x, alpha=0.01)

Numpy reference for leaky_relu.

x if x > 0 else alpha * x. alpha must match the slope the design passes to the kernel at runtime; the kernel takes it as bf16, so it is rounded to bf16 here. Exact up to bf16 rounding; pair with a small rtol when verifying.

Source code in python/iron/kernels/activation.py

def leaky_relu_ref(x, alpha=0.01):
    """Numpy reference for [`leaky_relu`][iron.kernels.activation.leaky_relu].

    ``x if x > 0 else alpha * x``.  ``alpha`` must match the slope the design
    passes to the kernel at runtime; the kernel takes it as bf16, so it is
    rounded to bf16 here.  Exact up to bf16 rounding; pair with a small
    ``rtol`` when verifying.
    """
    alpha = np.float32(bfloat16(alpha))
    xf = x.astype(np.float32)
    return np.where(xf > 0.0, xf, alpha * xf).astype(x.dtype)

swiglu_ref

swiglu_ref(x, w1, w2)

Numpy reference for swiglu: (x * w1) * silu(x * w2).

swiglu.cc forms the two products in bf16, then silu of the second through the tanh LUT (0.5 * (1 + tanh(z / 2))). The reference rounds the two products to bf16 as the kernel does and computes the rest in float64; LUT-approximation territory, pair with rtol=0.128.

Where silu is 0 (x * w2 is -inf, or low enough to underflow) the output is 0, as swiglu_lut_ref's is, even where x * w1 overflowed: 0 is the limit as x goes to -inf, and inf * 0 would be NaN.

Source code in python/iron/kernels/activation.py

def swiglu_ref(x, w1, w2):
    """Numpy reference for [`swiglu`][iron.kernels.activation.swiglu]: ``(x * w1) * silu(x * w2)``.

    ``swiglu.cc`` forms the two products in bf16, then ``silu`` of the second
    through the tanh LUT (``0.5 * (1 + tanh(z / 2))``). The reference rounds
    the two products to bf16 as the kernel does and computes the rest in
    float64; LUT-approximation territory, pair with ``rtol=0.128``.

    Where silu is 0 (``x * w2`` is -inf, or low enough to underflow) the
    output is 0, as swiglu_lut_ref's is, even where ``x * w1`` overflowed:
    0 is the limit as x goes to -inf, and ``inf * 0`` would be NaN.
    """
    xf = _f64(x)
    xw1 = _f64(round_to(xf * _f64(w1), bfloat16))
    xw2 = _no_neg_inf(_f64(round_to(xf * _f64(w2), bfloat16)))
    with np.errstate(over="ignore", invalid="ignore"):
        silu = xw2 / (1.0 + np.exp(-xw2))
        out = np.where(silu == 0, np.copysign(0.0, xw1) * silu, xw1 * silu)
    return _rounded(out, x)

bf16_exp_ref

bf16_exp_ref(x)

Numpy reference for bf16_exp — element-wise exp(x).

LUT approximation territory; pair with the canonical 12.8% relative tolerance and stop_at_nonfinite=True (the default in count_mismatches) when verifying.

exp(clip(x, -88, 88)), not plain exp(x): the kernel clamps to EXP_BF16_CLAMP before its Q8 fixed-point table lookup (see aie_runtime_lib/AIE2/lut_based_ops.h), so it saturates rather than overflowing. +88 is the largest value the tables carry. exp(-88) is a nonzero bf16 subnormal (about 6.06e-39), not zero: AIE2P preserves it through integer exponent reconstruction, whereas the AIE2 LUT may flush the tail under the absolute tolerance. The clamp also keeps the result in range: exp(88) = 1.65e+38 fits bf16 and float32 where exp(89) would not.

Source code in python/iron/kernels/activation.py

def bf16_exp_ref(x):
    """Numpy reference for [`bf16_exp`][iron.kernels.activation.bf16_exp] — element-wise ``exp(x)``.

    LUT approximation territory; pair with the canonical 12.8% relative
    tolerance and ``stop_at_nonfinite=True`` (the default in
    `count_mismatches`) when verifying.

    ``exp(clip(x, -88, 88))``, not plain ``exp(x)``: the kernel clamps to
    ``EXP_BF16_CLAMP`` before its Q8 fixed-point table lookup (see
    ``aie_runtime_lib/AIE2/lut_based_ops.h``), so it saturates rather than
    overflowing. ``+88`` is the largest value the tables carry.
    ``exp(-88)`` is a nonzero bf16 subnormal (about ``6.06e-39``), not
    zero: AIE2P preserves it through integer exponent reconstruction,
    whereas the AIE2 LUT may flush the tail under the absolute tolerance.
    The clamp also keeps the result in range: ``exp(88) = 1.65e+38`` fits
    bf16 and float32 where ``exp(89)`` would not.
    """
    xf = np.clip(_f64(x), -_EXP_BF16_CLAMP, _EXP_BF16_CLAMP)
    # The clamp rules out overflow, so that warning stays un-suppressed and
    # would now be a real signal. A NaN input still reaches exp -- the kernel
    # declares nonfinite="unspecified" and callers do feed raw bit patterns
    # (programming_examples/basic/vector_exp sweeps all 65536 of them).
    with np.errstate(invalid="ignore"):
        return _rounded(np.exp(xf), x)

exp2f_vec_ref

exp2f_vec_ref(x, min_x: float = -111.0)

Numpy reference for exp2f_vec: exact 2**x.

Unlike the LUT-based refs above, this is float64 2**x (not a reimplementation of the on-device poly): the kernel holds 8.9e-5 relative error or better, several orders tighter than the LUT-based kernels' 12.8% default, so pair with a correspondingly tight tolerance (e.g. rtol=1e-3) rather than the LUT default.

The kernel clamps its input to min_x before evaluating (see the factory's min_x), so the reference does the same: 2**-5000 is 2**min_x on the device, not zero. Pass the factory's min_x.

Source code in python/iron/kernels/activation.py

def exp2f_vec_ref(x, min_x: float = -111.0):
    """Numpy reference for [`exp2f_vec`][iron.kernels.activation.exp2f_vec]: exact ``2**x``.

    Unlike the LUT-based refs above, this is float64 ``2**x`` (not a
    reimplementation of the on-device poly): the kernel holds 8.9e-5
    relative error or better, several orders tighter than the LUT-based
    kernels' 12.8% default, so pair with a correspondingly tight
    tolerance (e.g. ``rtol=1e-3``) rather than the LUT default.

    The kernel clamps its input to ``min_x`` before evaluating (see the
    factory's ``min_x``), so the reference does the same: ``2**-5000`` is
    ``2**min_x`` on the device, not zero. Pass the factory's ``min_x``.
    """
    xf = np.maximum(x.astype(np.float64), min_x)
    return np.exp2(xf).astype(x.dtype)

softmax_ref

softmax_ref(x, *, tile_size: int = 1024)

Numpy reference for softmax.

The AIE kernel computes softmax independently per tile_size-element tile (no cross-tile reduction), so the reference splits x the same way before applying the softmax, in float64. x.size must be a multiple of tile_size.

Source code in python/iron/kernels/activation.py

def softmax_ref(x, *, tile_size: int = 1024):
    """Numpy reference for [`softmax`][iron.kernels.activation.softmax].

    The AIE kernel computes softmax independently per ``tile_size``-element
    tile (no cross-tile reduction), so the reference splits ``x`` the same
    way before applying the softmax, in float64.  ``x.size`` must be a
    multiple of ``tile_size``.
    """
    xf = _f64(x)
    if xf.size % tile_size != 0:
        raise ValueError(
            f"softmax_ref: x has {xf.size} elements; not a multiple of "
            f"tile_size={tile_size}"
        )
    flat = xf.reshape(-1, tile_size)
    flat = flat - flat.max(axis=1, keepdims=True)
    exp = np.exp(flat)
    out = exp / exp.sum(axis=1, keepdims=True)
    return _rounded(out.reshape(x.shape), x)

Normalization

rms_norm, rms_norm_eps, and layer_norm support both aie2 and aie2p. They accept tile_size (default 1024) or its compatibility alias cols. rope in Data movement has the same size API and supports both interleaved and two_halves=True layouts. The transformer module re-exports the canonical bf16 norm and RoPE factories and references; importing through either module does not select a different implementation. Norm references accept eps.

Normalization kernel factories + numpy references: rms_norm, layer_norm.

rms_norm

rms_norm(
    tile_size: int = 1024, *, cols: int | None = None
) -> ExternalFunction

RMS-norm a bf16 row on aie2/aie2p; (in, out, cols), eps=1e-5.

cols is a compatibility alias for tile_size. Positive row lengths need not be vector-aligned: the kernel handles scalar tails.

Source code in python/iron/kernels/norm.py

def rms_norm(tile_size: int = 1024, *, cols: int | None = None) -> ExternalFunction:
    """RMS-norm a bf16 row on aie2/aie2p; ``(in, out, cols)``, eps=1e-5.

    ``cols`` is a compatibility alias for ``tile_size``. Positive row lengths
    need not be vector-aligned: the kernel handles scalar tails.
    """
    tile_size = _row_size("rms_norm", tile_size, cols)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _norm_extern(
        "rms_norm",
        "rms_norm.cc",
        [tile_ty, tile_ty, np.int32],
        KernelContract(
            trace=Trace.whole_call(),
            setup=None if _tuned_arch() == "aie2" else conv_even,
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            reference=rms_norm_ref,
            acc_dtype=np.float32,
            reduction=tile_size,
            tolerance={
                "aie2": _RMS_NORM_BF16_AIE2,
                "aie2p": _RMS_NORM_BF16_AIE2P,
            }.get(_tuned_arch(), _NORM_BF16),
            ops_per_call=4 * tile_size,
        ),
    )

rms_norm_eps

rms_norm_eps(
    tile_size: int = 1024, *, cols: int | None = None
) -> ExternalFunction

RMS-norm a bf16 row (gamma=1); design passes (in, out, cols, epsilon).

Source code in python/iron/kernels/norm.py

def rms_norm_eps(tile_size: int = 1024, *, cols: int | None = None) -> ExternalFunction:
    """RMS-norm a bf16 row (gamma=1); design passes ``(in, out, cols, epsilon)``."""
    tile_size = _row_size("rms_norm_eps", tile_size, cols)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _norm_extern(
        "rms_norm_eps",
        "rms_norm.cc",
        [tile_ty, tile_ty, np.int32, np.float32],
        KernelContract(
            trace=Trace.whole_call(),
            setup=None if _tuned_arch() == "aie2" else conv_even,
            roles=(In, Out, Param, Param),
            parameter_bindings=((2, tile_size),),
            reference=lambda x, epsilon: rms_norm_ref(x, eps=epsilon),
            acc_dtype=np.float32,
            reduction=tile_size,
            tolerance={
                "aie2": _RMS_NORM_BF16_AIE2,
                "aie2p": _RMS_NORM_BF16_AIE2P,
            }.get(_tuned_arch(), _NORM_BF16),
            ops_per_call=4 * tile_size,
        ),
    )

layer_norm

layer_norm(
    tile_size: int = 1024, *, cols: int | None = None
) -> ExternalFunction

Layer-norm a bf16 row; (in, out, cols), gamma=1, beta=0, eps=1e-5.

cols aliases tile_size, a positive multiple of 16 (the source processes whole 16-lane halves, without a scalar tail).

Source code in python/iron/kernels/norm.py

def layer_norm(tile_size: int = 1024, *, cols: int | None = None) -> ExternalFunction:
    """Layer-norm a bf16 row; ``(in, out, cols)``, gamma=1, beta=0, eps=1e-5.

    ``cols`` aliases ``tile_size``, a positive multiple of 16 (the source
    processes whole 16-lane halves, without a scalar tail).
    """
    tile_size = _row_size("layer_norm", tile_size, cols, 16)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _norm_extern(
        "layer_norm",
        "layer_norm.cc",
        [tile_ty, tile_ty, np.int32],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            reference=layer_norm_ref,
            acc_dtype=np.float32,
            reduction=tile_size,
            tolerance={
                "aie2": _LAYER_NORM_BF16_AIE2,
                "aie2p": _LAYER_NORM_BF16_AIE2P,
            }.get(_tuned_arch(), _NORM_BF16),
            ops_per_call=6 * tile_size,
        ),
    )

rms_norm_ref

rms_norm_ref(x, *, eps: float = 1e-05)

Numpy reference for rms_norm: x / sqrt(mean(x**2) + eps).

Source code in python/iron/kernels/norm.py

def rms_norm_ref(x, *, eps: float = 1e-5):
    """Numpy reference for [`rms_norm`][iron.kernels.norm.rms_norm]: ``x / sqrt(mean(x**2) + eps)``."""
    xf = x.astype(np.float32)
    ms = np.mean(xf * xf, axis=-1, keepdims=True)
    return (xf / np.sqrt(ms + eps)).astype(x.dtype)

layer_norm_ref

layer_norm_ref(x, *, eps: float = 1e-05)

Numpy reference for layer_norm: (x - mean) / sqrt(var + eps).

Source code in python/iron/kernels/norm.py

def layer_norm_ref(x, *, eps: float = 1e-5):
    """Numpy reference for [`layer_norm`][iron.kernels.norm.layer_norm]: ``(x - mean) / sqrt(var + eps)``."""
    xf = x.astype(np.float32)
    mean = np.mean(xf, axis=-1, keepdims=True)
    var = np.mean((xf - mean) ** 2, axis=-1, keepdims=True)
    return ((xf - mean) / np.sqrt(var + eps)).astype(x.dtype)

Transformer blocks

Transformer building blocks: rms_norm, layer_norm (bf16, f32, affine+cast), rope, mm_activation_epilogue.

The bf16 norms and RoPE are re-exported from norm and datamovement; they support aie2 and aie2p, with cols as an alias for tile_size. The activation epilogue and the f32/affine norms take their aie2p source on both generations. Each processes one row (cols elements) per call; the row length is a scalar Param the factory binds to cols. These are the kernels programming_examples/ml/{norm,rope,mm_activation_epilogue} build.

layer_norm_f32

layer_norm_f32(cols: int = 4096) -> ExternalFunction

Row-wise LayerNorm on float32 in and out (gamma = 1, beta = 0, eps 1e-5).

A separate factory rather than a dtype of layer_norm: this one is held to atol 2e-6 (1e-3 in the portable build) instead of the bf16 tolerance, which its reference meets only by computing the variance two-pass in float64. Merging them would put that numerical difference behind a dtype switch.

Parameters:

Name	Type	Description	Default
cols	int	Elements per row (multiple of 16).	4096

Source code in python/iron/kernels/transformer.py

def layer_norm_f32(cols: int = 4096) -> ExternalFunction:
    """Row-wise LayerNorm on float32 in and out (gamma = 1, beta = 0, eps 1e-5).

    A separate factory rather than a dtype of
    [`layer_norm`][iron.kernels.norm.layer_norm]: this one is held to atol 2e-6
    (1e-3 in the portable build) instead of the bf16 tolerance, which its
    reference meets only by computing the variance two-pass in float64.
    Merging them would put that numerical difference behind a dtype switch.

    Args:
        cols: Elements per row (multiple of 16).
    """
    return _row_kernel(
        "layer_norm_f32",
        "layer_norm_f32",
        "transformer/layer_norm_f32.cc",
        cols,
        np.float32,
        np.float32,
        layer_norm_f32_ref,
        _NORM_F32_TOLERANCE.get(_tuned_arch(), _NORM_F32),
        6 * cols,
        # 896 B tuned for aie2p, 160 B tuned for aie2, 832 B untuned on aie2p
        # (672 B on aie2); the 1024 B default covers every build.
    )

layer_norm_affine_cast

layer_norm_affine_cast(
    cols: int = 4096,
) -> ExternalFunction

Row-wise LayerNorm, f32 in, per-column gamma/beta, bf16 out.

The second argument holds gamma (cols values) followed by beta (cols values) as float32: a tensor Param, which the generic builder bakes into a core buffer.

Parameters:

Name	Type	Description	Default
cols	int	Elements per row (multiple of 16).	4096

Source code in python/iron/kernels/transformer.py

def layer_norm_affine_cast(cols: int = 4096) -> ExternalFunction:
    """Row-wise LayerNorm, f32 in, per-column gamma/beta, bf16 out.

    The second argument holds ``gamma`` (``cols`` values) followed by ``beta``
    (``cols`` values) as float32: a tensor ``Param``, which the generic
    builder bakes into a core buffer.

    Args:
        cols: Elements per row (multiple of 16).
    """
    _cols("layer_norm_affine_cast", cols)
    in_ty = np.ndarray[(cols,), np.dtype[np.float32]]
    gb_ty = np.ndarray[(2 * cols,), np.dtype[np.float32]]
    out_ty = np.ndarray[(cols,), np.dtype[bfloat16]]
    return _make_extern(
        "layer_norm_affine_cast",
        _kernel_source("transformer/layer_norm_f32.cc"),
        [in_ty, gb_ty, out_ty, np.int32],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Param, Out, Param),
            parameter_bindings=((3, cols),),
            reference=layer_norm_affine_cast_ref,
            acc_dtype=np.float32,
            reduction=cols,
            tolerance=_AFFINE_TOLERANCE.get(_tuned_arch(), _NORM_BF16),
            ops_per_call=8 * cols,
        ),
    )

mm_activation_epilogue

mm_activation_epilogue(
    tile_size: int = 1024,
) -> ExternalFunction

GEMM epilogue on float32 rows: identity (0), SiLU (1), tanh-GELU (2) or ReLU (3) by mode.

One resident kernel whose mode is a runtime argument, so a design can switch activations without recompiling (programming_examples/ml/mm_activation_epilogue).

AIE2 has no tanh instruction, so there SiLU and GELU read getTanhBf16's table, and its tuned build is judged against a model of that arithmetic instead of the true functions.

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call (multiple of 16).	1024

Source code in python/iron/kernels/transformer.py

def mm_activation_epilogue(tile_size: int = 1024) -> ExternalFunction:
    """GEMM epilogue on float32 rows: identity (0), SiLU (1), tanh-GELU (2) or ReLU (3) by ``mode``.

    One resident kernel whose ``mode`` is a runtime argument, so a design can
    switch activations without recompiling
    (programming_examples/ml/mm_activation_epilogue).

    AIE2 has no tanh instruction, so there SiLU and GELU read getTanhBf16's
    table, and its tuned build is judged against a model of that arithmetic
    instead of the true functions.

    Args:
        tile_size: Elements per call (multiple of 16).
    """
    _cols("mm_activation_epilogue", tile_size)
    tile_ty = np.ndarray[(tile_size,), np.dtype[np.float32]]
    source = _kernel_source("transformer/mm_activation_epilogue.cc")
    lut = not _arch_traits().native_tanh
    lut_model = _tuned_arch() == "aie2"
    flags = None
    if lut:
        # lut_kernel.cc compiles the source next to lut_based_ops.cpp, whose
        # tables getTanhBf16 reads.
        flags = [f'-DAIE_LUT_KERNEL_SOURCE="{source}"', f"-I{_runtime_lib_include()}"]
        source = _kernel_source("common/lut_kernel.cc")
    return _make_extern(
        "mm_activation_epilogue_row",
        source,
        [tile_ty, tile_ty, np.int32, np.int32],
        compile_flags=flags,
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, Out, Param, Param),
            parameter_bindings=((2, tile_size),),
            reference=(
                mm_activation_epilogue_lut_ref
                if lut_model
                else mm_activation_epilogue_ref
            ),
            acc_dtype=np.float32,
            reduction=1,
            tolerance=(
                _EPILOGUE_LUT_TOLERANCE
                if lut_model
                else Tolerance.relative(
                    0.128,
                    0.05,
                    note="programming_examples/ml/mm_activation_epilogue: atol 0.05 "
                    "for the bf16-internal SiLU / GELU, identity and ReLU are exact",
                )
            ),
            ops_per_call=8 * tile_size,
            uses_lut=lut,
        ),
    )

layer_norm_f32_ref

layer_norm_f32_ref(x)

Numpy reference for layer_norm_f32.

Centered two-pass variance in float64 so it stays exact on the non-zero-mean input the f32 kernel is exercised with.

Source code in python/iron/kernels/transformer.py

def layer_norm_f32_ref(x):
    """Numpy reference for [`layer_norm_f32`][iron.kernels.transformer.layer_norm_f32].

    Centered two-pass variance in float64 so it stays exact on the
    non-zero-mean input the f32 kernel is exercised with.
    """
    x64 = x.astype(np.float64)
    mean = x64.mean(axis=-1, keepdims=True)
    var = ((x64 - mean) ** 2).mean(axis=-1, keepdims=True)
    return ((x64 - mean) / np.sqrt(var + _EPS)).astype(np.float32)

layer_norm_affine_cast_ref

layer_norm_affine_cast_ref(x, gamma_beta)

Numpy reference for layer_norm_affine_cast.

gamma_beta is gamma then beta, each cols float32 values.

Source code in python/iron/kernels/transformer.py

def layer_norm_affine_cast_ref(x, gamma_beta):
    """Numpy reference for [`layer_norm_affine_cast`][iron.kernels.transformer.layer_norm_affine_cast].

    ``gamma_beta`` is ``gamma`` then ``beta``, each ``cols`` float32 values.
    """
    x32 = x.astype(np.float32)
    gb = np.asarray(gamma_beta, dtype=np.float32).reshape(-1)
    cols = x32.shape[-1]
    gamma, beta = gb[:cols], gb[cols : 2 * cols]
    mean = x32.mean(axis=-1, keepdims=True)
    var = ((x32 - mean) ** 2).mean(axis=-1, keepdims=True)
    return ((x32 - mean) / np.sqrt(var + _EPS) * gamma + beta).astype(bfloat16)

mm_activation_epilogue_ref

mm_activation_epilogue_ref(x, mode)

Numpy reference for mm_activation_epilogue.

mode 0 identity, 1 x * sigmoid(x), 2 tanh-approximation GELU, 3 max(x, 0).

Source code in python/iron/kernels/transformer.py

def mm_activation_epilogue_ref(x, mode):
    """Numpy reference for [`mm_activation_epilogue`][iron.kernels.transformer.mm_activation_epilogue].

    ``mode`` 0 identity, 1 ``x * sigmoid(x)``, 2 tanh-approximation GELU,
    3 ``max(x, 0)``.
    """
    x32 = x.astype(np.float32)
    mode = int(mode)
    if mode == 0:
        return x32.astype(x.dtype)
    if mode == 1:
        with np.errstate(over="ignore"):
            return (x32 / (1.0 + np.exp(-x32))).astype(x.dtype)
    if mode == 2:
        # gelu's limit at -inf is 0, which the most negative float gives
        # rather than -inf * 0.
        x32 = np.maximum(x32, -np.finfo(np.float32).max)
        inner = 0.7978845608 * (x32 + 0.044715 * x32**3)
        return (0.5 * x32 * (1.0 + np.tanh(inner))).astype(x.dtype)
    if mode == 3:
        return np.maximum(x32, 0.0).astype(x.dtype)
    raise ValueError(f"mm_activation_epilogue mode must be 0, 1, 2 or 3, got {mode}")

mm_activation_epilogue_lut_ref

mm_activation_epilogue_lut_ref(x, mode)

Model of mm_activation_epilogue on aie2.

Follows mm_activation_epilogue.cc's roundings around getTanhBf16 (tanh_lut_ref). SiLU splits x into hi, its top 16 bits, and lo, bf16(x - hi), and multiplies each by the bf16 sigmoid (bf16(t + 1)) / 2, where t is the table's tanh of bf16(x) / 2 narrowed to bf16. hi is finite for any finite x, so huge inputs give about x or 0 rather than NaN; +-inf still gives NaN. GELU runs in bf16: x, x * x and the inner polynomial are each rounded before the next step, and the output is bf16(x / 2) * bf16(t + 1), all reading x clamped at -8 so -inf gives 0. Both return +0 where IEEE arithmetic gives -0, as the accumulator does. Identity and ReLU are exact.

Source code in python/iron/kernels/transformer.py

def mm_activation_epilogue_lut_ref(x, mode):
    """Model of [`mm_activation_epilogue`][iron.kernels.transformer.mm_activation_epilogue] on aie2.

    Follows mm_activation_epilogue.cc's roundings around getTanhBf16
    ([`tanh_lut_ref`][iron.kernels.activation.tanh_lut_ref]). SiLU splits
    ``x`` into ``hi``, its top 16 bits, and ``lo``, ``bf16(x - hi)``, and
    multiplies each by the bf16 sigmoid ``(bf16(t + 1)) / 2``, where ``t`` is
    the table's tanh of ``bf16(x) / 2`` narrowed to bf16. ``hi`` is finite for
    any finite ``x``, so huge inputs give about ``x`` or 0 rather than NaN;
    +-inf still gives NaN. GELU runs in bf16: ``x``, ``x * x`` and the inner
    polynomial are each rounded before the next step, and the output is
    ``bf16(x / 2) * bf16(t + 1)``, all reading ``x`` clamped at -8 so -inf
    gives 0. Both return +0 where IEEE arithmetic gives -0, as the accumulator
    does. Identity and ReLU are exact.
    """
    x32 = np.asarray(x, np.float32)
    mode = int(mode)
    if mode not in (1, 2):
        return mm_activation_epilogue_ref(x, mode)
    with np.errstate(over="ignore", invalid="ignore"):
        if mode == 1:
            hi = (x32.view(np.uint32) & np.uint32(0xFFFF0000)).view(np.float32)
            lo = _bf16(x32 - hi)
            t = tanh_lut_ref(_bf16(x32) * np.float32(0.5))
            sig = _bf16(_bf16(t + np.float32(1.0)) * np.float32(0.5))
            out = hi * sig + lo * sig
        else:
            c0 = _bf16(np.float32(0.7978845608))
            c0c1 = _bf16(np.float32(0.7978845608) * np.float32(0.044715))
            xl = np.maximum(_bf16(x32), np.float32(-8.0))
            poly = _bf16(c0 + c0c1 * _bf16(xl * xl))
            t = tanh_lut_ref(_bf16(xl * poly))
            half_x = _bf16(np.float32(0.5) * xl)
            out = half_x * _bf16(t + np.float32(1.0))
    return (out + np.float32(0.0)).astype(x.dtype)

Vision

Vision kernel factories: color conversion, threshold, filter2d, add_weighted.

rgba2hue

rgba2hue(
    line_width: int = 1920, use_chess: bool = False
) -> ExternalFunction

Convert a line of RGBA pixels to hue values (full-range, 0..255).

Source code in python/iron/kernels/vision.py

def rgba2hue(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Convert a line of RGBA pixels to hue values (full-range, 0..255)."""
    _require_vector_alignment("rgba2hue", line_width, 32, param="line_width")
    # lut_inv.h pins its gather pair with AIE_BANK_A/AIE_BANK_B.
    flags = [f"-I{_runtime_lib_include()}"]
    return _color_convert_kernel(
        "rgba2hueLine",
        "rgba2hue.cc",
        line_width * 4,
        line_width,
        use_chess=use_chess,
        compile_flags=flags,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, line_width),),
            reference=rgba2hue_ref,
            acc_dtype=np.int32,
            reduction=1,
            tolerance=Tolerance.exact(note="integer reciprocal, no rounding slack"),
            uses_lut=True,
        ),
    )

threshold

threshold(
    line_width: int = 1920,
    dtype: type = uint8,
    use_chess: bool = False,
) -> ExternalFunction

Apply a threshold operation to a line of pixels.

Parameters:

Name	Type	Description	Default
line_width	int	Number of elements per line.	1920
dtype	type	Element data type, np.uint8 or np.int16. The source's 32-bit branch multiplies int32 data by int16 coefficients, a MAC AIE2 does not have, so it does not compile and is not offered.	uint8
use_chess	bool	When True, build the .o with xchesscc_wrapper instead of Peano.	False

Raises:

Type	Description
ValueError	When dtype is not np.uint8, np.int16, or np.int32.

Source code in python/iron/kernels/vision.py

@dtypes(({"dtype": np.uint8}, {"dtype": np.int16}, {"dtype": np.int32}))
def threshold(
    line_width: int = 1920, dtype: type = np.uint8, use_chess: bool = False
) -> ExternalFunction:
    """Apply a threshold operation to a line of pixels.

    Args:
        line_width: Number of elements per line.
        dtype: Element data type, ``np.uint8`` or ``np.int16``. The source's
            32-bit branch multiplies int32 data by int16 coefficients, a MAC
            AIE2 does not have, so it does not compile and is not offered.
        use_chess: When ``True``, build the .o with ``xchesscc_wrapper``
            instead of Peano.

    Raises:
        ValueError: When ``dtype`` is not ``np.uint8``, ``np.int16``, or ``np.int32``.
    """
    bit_width = _dtype_to_bit_width(dtype, factory_name="threshold")
    _require_vector_alignment(
        "threshold", line_width, 512 // bit_width, param="line_width"
    )
    scalar_ty = np.int32 if bit_width == 32 else np.int16
    line_ty = np.ndarray[(line_width,), np.dtype[dtype]]
    return _make_extern(
        "thresholdLine",
        _kernel_source("vision/threshold.cc"),
        [line_ty, line_ty, np.int32, scalar_ty, scalar_ty, np.int8],
        compile_flags=[f"-DBIT_WIDTH={bit_width}", f"-DTHRESHOLD_ELEMS={line_width}"],
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param, Param, Param, Param),
            parameter_bindings=((2, line_width),),
            reference=threshold_ref,
            tolerance=Tolerance.exact(note="selection"),
        ),
    )

bitwise_or

bitwise_or(
    line_width: int = 1920,
    dtype: type = uint8,
    use_chess: bool = False,
) -> ExternalFunction

Element-wise bitwise OR of two lines.

Source code in python/iron/kernels/vision.py

@dtypes(({"dtype": np.uint8}, {"dtype": np.int16}, {"dtype": np.int32}))
def bitwise_or(
    line_width: int = 1920, dtype: type = np.uint8, use_chess: bool = False
) -> ExternalFunction:
    """Element-wise bitwise OR of two lines."""
    return _bitwise_kernel("OR", line_width, dtype, use_chess=use_chess)

bitwise_and

bitwise_and(
    line_width: int = 1920,
    dtype: type = uint8,
    use_chess: bool = False,
) -> ExternalFunction

Element-wise bitwise AND of two lines.

Source code in python/iron/kernels/vision.py

@dtypes(({"dtype": np.uint8}, {"dtype": np.int16}, {"dtype": np.int32}))
def bitwise_and(
    line_width: int = 1920, dtype: type = np.uint8, use_chess: bool = False
) -> ExternalFunction:
    """Element-wise bitwise AND of two lines."""
    return _bitwise_kernel("AND", line_width, dtype, use_chess=use_chess)

gray2rgba

gray2rgba(
    line_width: int = 1920, use_chess: bool = False
) -> ExternalFunction

Convert a grayscale line to RGBA.

Source code in python/iron/kernels/vision.py

def gray2rgba(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Convert a grayscale line to RGBA."""
    return _color_convert_kernel(
        "gray2rgbaLine",
        "gray2rgba.cc",
        line_width,
        line_width * 4,
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, line_width),),
            reference=gray2rgba_ref,
            tolerance=Tolerance.exact(note="copy with alpha = 255"),
        ),
    )

rgba2gray

rgba2gray(
    line_width: int = 1920, use_chess: bool = False
) -> ExternalFunction

Convert an RGBA line to grayscale.

Source code in python/iron/kernels/vision.py

def rgba2gray(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Convert an RGBA line to grayscale."""
    flags = []
    if not use_chess and _tuned_arch() == "aie2p":
        # The pre-RA pipeliner's schedule of the 64-pixel loop ends up at II13
        # after register allocation; the postpipeliner finds II11.
        flags += ["-mllvm", "--aie-force-postpipeliner"]
    return _color_convert_kernel(
        "rgba2grayLine",
        "rgba2gray.cc",
        line_width * 4,
        line_width,
        use_chess=use_chess,
        compile_flags=flags,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, line_width),),
            reference=rgba2gray_ref,
            tolerance=Tolerance.exact(
                note="measured bit-exact against the reference over every data case"
            ),
            acc_dtype=np.int32,
            reduction=3,
        ),
    )

filter2d

filter2d(
    line_width: int = 1920, use_chess: bool = False
) -> ExternalFunction

Apply a 3x3 2D convolution filter across three input lines.

Source code in python/iron/kernels/vision.py

def filter2d(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Apply a 3x3 2D convolution filter across three input lines."""
    if line_width % 32 or line_width < 96:
        raise ValueError(
            f"filter2d: line_width must be a multiple of 32 and at least 96 "
            f"(the vector kernel handles the two borders in 32-pixel blocks), "
            f"got {line_width}"
        )
    line_ty = np.ndarray[(line_width,), np.dtype[np.uint8]]
    kernel_ty = np.ndarray[(3, 3), np.dtype[np.int16]]
    return _make_extern(
        "filter2dLine",
        _kernel_source("vision/filter2d.cc"),
        [line_ty, line_ty, line_ty, line_ty, np.int32, kernel_ty],
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Out, Param, Param),
            parameter_bindings=((4, line_width),),
            reference=filter2d_ref,
            acc_dtype=np.int32,
            reduction=9,
            # The one-LSB slack these pixel kernels used to share was
            # absorbing a wrong carry across the 32-pixel boundary here (see
            # filter2d.cc); an exact contract is what would have caught it.
            tolerance=Tolerance.exact(
                note="vector path matches the reference bit-for-bit"
            ),
            ops_per_call=18 * line_width,
        ),
    )

add_weighted

add_weighted(
    line_width: int = 1920,
    dtype: type = uint8,
    use_chess: bool = False,
) -> ExternalFunction

Weighted addition of two lines with a gamma offset.

Parameters:

Name	Type	Description	Default
line_width	int	Number of elements per line.	1920
dtype	type	Element data type, np.uint8 or np.int16. The source's 32-bit branch multiplies int32 data by int16 coefficients, a MAC AIE2 does not have, so it does not compile and is not offered.	uint8
use_chess	bool	When True, build the .o with xchesscc_wrapper instead of Peano.	False

Raises:

Type	Description
ValueError	When dtype is not np.uint8 or np.int16.

Source code in python/iron/kernels/vision.py

@dtypes(({"dtype": np.uint8}, {"dtype": np.int16}))
def add_weighted(
    line_width: int = 1920, dtype: type = np.uint8, use_chess: bool = False
) -> ExternalFunction:
    """Weighted addition of two lines with a gamma offset.

    Args:
        line_width: Number of elements per line.
        dtype: Element data type, ``np.uint8`` or ``np.int16``. The source's
            32-bit branch multiplies int32 data by int16 coefficients, a MAC
            AIE2 does not have, so it does not compile and is not offered.
        use_chess: When ``True``, build the .o with ``xchesscc_wrapper``
            instead of Peano.

    Raises:
        ValueError: When ``dtype`` is not ``np.uint8`` or ``np.int16``.
    """
    bit_width = _dtype_to_bit_width(dtype, factory_name="add_weighted")
    if bit_width == 32:
        raise ValueError(
            "add_weighted: no int32 build; addWeighted.cc has no int32 x int16 MAC. "
            "Use np.uint8 or np.int16."
        )
    gamma_ty = {8: np.int8, 16: np.int16}[bit_width]
    _require_vector_alignment(
        "add_weighted", line_width, 256 // bit_width, param="line_width"
    )
    line_ty = np.ndarray[(line_width,), np.dtype[dtype]]
    return _make_extern(
        "addWeightedLine",
        _kernel_source("vision/addWeighted.cc"),
        [line_ty, line_ty, line_ty, np.int32, np.int16, np.int16, gamma_ty],
        compile_flags=[
            f"-DBIT_WIDTH={bit_width}",
            f"-DADD_WEIGHTED_ELEMS={line_width}",
        ],
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out, Param, Param, Param, Param),
            parameter_bindings=((3, line_width),),
            reference=add_weighted_ref,
            acc_dtype=np.int32,
            reduction=2,
            tolerance=Tolerance.exact(
                note="measured bit-exact against the reference over every data case"
            ),
            ops_per_call=3 * line_width,
        ),
    )

gray2rgba_ref

gray2rgba_ref(y)

Numpy reference for gray2rgba: (y, y, y, 255) per pixel.

Source code in python/iron/kernels/vision.py

def gray2rgba_ref(y):
    """Numpy reference for [`gray2rgba`][iron.kernels.vision.gray2rgba]: ``(y, y, y, 255)`` per pixel."""
    y = np.asarray(y, dtype=np.uint8)
    out = np.empty(y.shape[:-1] + (y.shape[-1] * 4,), dtype=np.uint8)
    px = out.reshape(*y.shape, 4)
    px[..., 0] = px[..., 1] = px[..., 2] = y
    px[..., 3] = 255
    return out

rgba2gray_ref

rgba2gray_ref(rgba)

Numpy reference for rgba2gray: fixed-point BT.470 luma.

Y = (9798 R + 19235 G + 3736 B + 2**14) >> 15, saturated to uint8; alpha is ignored. Within one LSB of the kernel (rounding of the final shift).

Source code in python/iron/kernels/vision.py

def rgba2gray_ref(rgba):
    """Numpy reference for [`rgba2gray`][iron.kernels.vision.rgba2gray]: fixed-point BT.470 luma.

    ``Y = (9798 R + 19235 G + 3736 B + 2**14) >> 15``, saturated to ``uint8``;
    alpha is ignored. Within one LSB of the kernel (rounding of the final
    shift).
    """
    rgba = np.asarray(rgba, dtype=np.uint8)
    px = rgba.reshape(*rgba.shape[:-1], -1, 4).astype(np.int64)
    wr, wg, wb = _GRAY_WEIGHTS
    acc = px[..., 0] * wr + px[..., 1] * wg + px[..., 2] * wb + (1 << 14)
    return np.clip(acc >> 15, 0, 255).astype(np.uint8)

rgba2hue_ref

rgba2hue_ref(rgba)

Numpy reference for rgba2hue: full-range hue.

rgba2hue.cc multiplies by a Q7.9 reciprocal rather than dividing, so with d = max - min of R, G, B and inv = 85 * 512 / d the hue is (offset * 512 + c * inv) >> 10 for whichever channel holds the max: c = G - B at offset 1, B - R at 171, R - G at 341. Each offset carries the + 1 that rounds the final halving, so there is one rounding step rather than two. The cast to uint8 wraps, so a negative hue (R max, G < B) comes out as 256 + h -- the right circular value. Gray pixels (d == 0) are hue 0, and a max held by both G and R goes to G, as the kernel's select order does. inv truncates, which leaves hue up to one LSB below the exact value.

Source code in python/iron/kernels/vision.py

def rgba2hue_ref(rgba):
    """Numpy reference for [`rgba2hue`][iron.kernels.vision.rgba2hue]: full-range hue.

    ``rgba2hue.cc`` multiplies by a Q7.9 reciprocal rather than dividing, so
    with ``d = max - min`` of R, G, B and ``inv = 85 * 512 / d`` the hue is
    ``(offset * 512 + c * inv) >> 10`` for whichever channel holds the max:
    ``c = G - B`` at offset 1, ``B - R`` at 171, ``R - G`` at 341.
    Each offset carries the ``+ 1`` that rounds the final halving, so there is
    one rounding step rather than two. The cast to ``uint8`` wraps, so a
    negative hue (R max, G < B) comes out as ``256 + h`` -- the right circular
    value. Gray pixels (``d == 0``) are hue 0, and a max held by both G and R
    goes to G, as the kernel's select order does. ``inv`` truncates, which
    leaves hue up to one LSB below the exact value.
    """
    rgba = np.asarray(rgba, dtype=np.uint8)
    px = rgba.reshape(*rgba.shape[:-1], -1, 4).astype(np.int64)
    r, g, b = px[..., 0], px[..., 1], px[..., 2]
    mx = np.maximum(np.maximum(r, g), b)
    d = mx - np.minimum(np.minimum(r, g), b)
    inv = (85 * 512) // np.where(d == 0, 1, d)  # avoid /0; masked to 0 below
    h = np.where(
        mx == g,
        (171 * 512 + (b - r) * inv) >> 10,
        np.where(
            mx == r,
            (1 * 512 + (g - b) * inv) >> 10,
            (341 * 512 + (r - g) * inv) >> 10,
        ),
    )
    return (np.where(d == 0, 0, h) & 0xFF).astype(np.uint8)

threshold_ref

threshold_ref(x, thresh, maxval, ttype)

Numpy reference for threshold (OpenCV semantics).

ttype: 0 binary (x > thresh ? maxval : 0), 1 binary inverted, 2 truncate (min(x, thresh)), 3 to-zero (x > thresh ? x : 0), 4 to-zero inverted. Exact.

Source code in python/iron/kernels/vision.py

def threshold_ref(x, thresh, maxval, ttype):
    """Numpy reference for [`threshold`][iron.kernels.vision.threshold] (OpenCV semantics).

    ``ttype``: 0 binary (``x > thresh ? maxval : 0``), 1 binary inverted,
    2 truncate (``min(x, thresh)``), 3 to-zero (``x > thresh ? x : 0``),
    4 to-zero inverted. Exact.
    """
    x = np.asarray(x)
    t, mv, zero = x.dtype.type(thresh), x.dtype.type(maxval), x.dtype.type(0)
    above = x > t
    ttype = int(ttype)
    if ttype == 0:
        return np.where(above, mv, zero)
    if ttype == 1:
        return np.where(above, zero, mv)
    if ttype == 2:
        return np.minimum(x, t)
    if ttype == 3:
        return np.where(above, x, zero)
    if ttype == 4:
        return np.where(above, zero, x)
    raise ValueError(f"threshold type must be 0..4, got {ttype}")

bitwise_or_ref

bitwise_or_ref(a, b)

Numpy reference for bitwise_or. Exact.

Source code in python/iron/kernels/vision.py

def bitwise_or_ref(a, b):
    """Numpy reference for [`bitwise_or`][iron.kernels.vision.bitwise_or]. Exact."""
    return np.bitwise_or(np.asarray(a), np.asarray(b))

bitwise_and_ref

bitwise_and_ref(a, b)

Numpy reference for bitwise_and. Exact.

Source code in python/iron/kernels/vision.py

def bitwise_and_ref(a, b):
    """Numpy reference for [`bitwise_and`][iron.kernels.vision.bitwise_and]. Exact."""
    return np.bitwise_and(np.asarray(a), np.asarray(b))

add_weighted_ref

add_weighted_ref(a, b, alpha, beta, gamma)

Numpy reference for add_weighted: Q2.14 blend.

out = sat(((alpha * a + beta * b) >> 14) + gamma) with alpha and beta as Q2.14 fixed point (8192 is 0.5) and gamma in output units, as OpenCV's addWeighted has it; the kernel reads gamma as the data type, so for uint8 data -56 is 200. The vector path in addWeighted.cc rounds the shift down; its scalar path rounds to nearest, so the two differ by at most one LSB.

Source code in python/iron/kernels/vision.py

def add_weighted_ref(a, b, alpha, beta, gamma):
    """Numpy reference for [`add_weighted`][iron.kernels.vision.add_weighted]: Q2.14 blend.

    ``out = sat(((alpha * a + beta * b) >> 14) + gamma)`` with ``alpha`` and
    ``beta`` as Q2.14 fixed point (``8192`` is 0.5) and ``gamma`` in output
    units, as OpenCV's ``addWeighted`` has it; the kernel reads ``gamma`` as
    the data type, so for ``uint8`` data ``-56`` is ``200``. The vector path in
    ``addWeighted.cc`` rounds the shift down; its scalar path rounds to
    nearest, so the two differ by at most one LSB.
    """
    a = np.asarray(a)
    info = np.iinfo(a.dtype)
    acc = a.astype(np.int64) * int(alpha) + np.asarray(b).astype(np.int64) * int(beta)
    acc = (acc >> 14) + int(np.array(gamma).astype(a.dtype))
    return np.clip(acc, info.min, info.max).astype(a.dtype)

filter2d_ref

filter2d_ref(line0, line1, line2, kernel)

Numpy reference for filter2d: 3x3 correlation of three lines.

Produces the middle line. The vector path keeps only the top byte of each int16 coefficient (k >> 8 as int8) and shifts the sum by 4, so a Q4.12 kernel such as 4096 * [[0, 1, 0], [1, -4, 1], [0, 1, 0]] lands on integer taps. Borders replicate the edge pixel; the result is saturated to uint8. Within one LSB of the kernel.

Source code in python/iron/kernels/vision.py

def filter2d_ref(line0, line1, line2, kernel):
    """Numpy reference for [`filter2d`][iron.kernels.vision.filter2d]: 3x3 correlation of three lines.

    Produces the middle line. The vector path keeps only the top byte of
    each ``int16`` coefficient (``k >> 8`` as ``int8``) and shifts the sum by
    4, so a Q4.12 kernel such as ``4096 * [[0, 1, 0], [1, -4, 1], [0, 1, 0]]``
    lands on integer taps. Borders replicate the edge pixel; the result is
    saturated to ``uint8``. Within one LSB of the kernel.
    """
    k8 = (np.asarray(kernel, dtype=np.int16) >> 8).astype(np.int8).astype(np.int64)
    k8 = k8.reshape(3, 3)
    terms = []
    for r, line in enumerate((line0, line1, line2)):
        v = np.asarray(line, dtype=np.uint8).astype(np.int64)
        left = np.concatenate([v[..., :1], v[..., :-1]], axis=-1)
        right = np.concatenate([v[..., 1:], v[..., -1:]], axis=-1)
        terms.append(left * k8[r, 0] + v * k8[r, 1] + right * k8[r, 2])
    acc = terms[0] + terms[1] + terms[2]
    return np.clip(acc >> 4, 0, 255).astype(np.uint8)

Contracts and the generic design builder

Every factory the generic builder can build carries a KernelContract on the returned ExternalFunction (fn.contract). The field-by-field account is in Kernel Library; the reference below is generated from the dataclass, so it cannot drift from it. Multi-dtype factories publish the combinations they support as factory.dtypes, and the factories above export their references as *_ref functions (add_ref, reduce_max_ref, mm_ref, ...), so host code never reimplements the math.

Shared helpers for the kernels submodules.

KernelContract dataclass

KernelContract(
    roles: tuple[type, ...],
    reference: (
        Callable[..., ndarray | tuple[ndarray, ...]] | None
    ) = None,
    tolerance: Tolerance | None = None,
    ops_per_call: int | None = None,
    out_valid: int | None = None,
    sample: Callable[..., list] | None = None,
    acc_dtype: type | None = None,
    reduction: int | None = None,
    setup: Callable[[], object] | None = None,
    stack_bytes: int | None = None,
    unsupported: str | None = None,
    layouts: tuple[TensorLayout | None, ...] = (),
    parameter_bindings: tuple[tuple[int, object], ...] = (),
    initializers: tuple[tuple[int, Callable], ...] = (),
    out_offset: tuple[int, int] | None = None,
    trace: Trace | None = None,
    uses_lut: bool = False,
    alignments: tuple[tuple[int, int], ...] = (),
)

What a kernel computes, declared next to the factory that builds it.

arg_types fixes each argument's shape and dtype; the contract adds what types cannot say, so aie.iron.algorithms.kernel_design can build, run and judge any factory from this one declaration.

Attributes:

Name	Type	Description
roles	tuple[type, ...]	In, Out, InOut or Param per argument (the first three are the @iron.jit markers). InOut is accumulated into, so it needs an initializer. Param is a scalar or a read-only tensor the generic builder holds fixed across its calls; the argument type decides which. Several outputs are allowed, in argument order.
reference	Callable[..., ndarray | tuple[ndarray, ...]] | None	The host implementation, and the arithmetic model (a saturating kernel's reference clips). Called with every unbound non-output argument in order: In tiles as (calls, n) arrays, Param values as arrays or numbers. Returns the output for all calls, a tuple for several outputs. None builds the kernel but does not judge it.
tolerance	Tolerance | None	How close the device must come; None is Tolerance.default_for the output dtype.
ops_per_call	int | None	Arithmetic operations per call; None means one per output element.
out_valid	int | None	Meaningful leading elements of a DMA-padded output tile; None means the whole tile.
sample	Callable[..., list] | None	sample(rng, calls) -> list[np.ndarray] for inputs with structure a dtype cannot express; None draws random data.
acc_dtype	type | None	The accumulator type, or None when nothing accumulates. With reduction it bounds the inputs so the accumulator cannot overflow.
reduction	int | None	Terms summed into one output element per call; None means one.
setup	Callable[[], object] | None	A kernel to run once on the core first (conv_even sets the rounding mode a bf16 store needs); None when the source sets its own mode or narrows nothing. A Worker handed the kernel calls it before its loop.
stack_bytes	int | None	Core stack a Worker calling this kernel needs, when more than the target's default. Say where the number came from; a note that gives only bytes per build means aiecc's measured_stack_size.
unsupported	str | None	Why the builder cannot run this kernel, or None. A kernel with no output argument (a cascade PUT half) says so here.
layouts	tuple[TensorLayout | None, ...]	A TensorLayout per argument; None is identity.
parameter_bindings	tuple[tuple[int, object], ...]	(index, value) pairs fixing Param operands, counts included; the rest come from the caller.
initializers	tuple[tuple[int, Callable], ...]	(index, factory) pairs for InOut arguments; factory(fn) returns the kernel that initializes the buffer.
out_offset	tuple[int, int] | None	(index, step) for a kernel that writes step elements of its one Out per call, at the offset it reads from bound scalar Param index. The builder then hands every call the same output tile, passes call * step as the offset and drains the tile once, so the calls must fill it exactly. None: each call writes a whole tile of its own.
trace	Trace | None	The Trace shape of the kernel's markers. Every library factory declares one; None (undeclared) is only for ad-hoc kernels, and cycles_per_call refuses it.
uses_lut	bool	Whether the kernel gathers through an aie::lut<4> table pair, so a build should verify the two tables land in different banks. Set it on the contract, not per source file: the LUT often comes in through a header (lut_based_ops.h, lut_inv.h).
alignments	tuple[tuple[int, int], ...]	(index, bytes) pairs for arguments the kernel loads as whole vectors from their start, so they must begin at a multiple of bytes. A call handed a memref.view at a constant offset that breaks one raises.

Overflow, rounding and NaN handling are not declared twice: the reference is the arithmetic model and the tolerance the slack against it.

out_indices property

out_indices: tuple[int, ...]

Output argument positions, in declaration order.

out_index property

out_index: int

Position of the one output, written (Out) or accumulated into (InOut).

Raises for a kernel with several outputs: code that must handle any kernel reads out_indices.

accumulates property

accumulates: bool

Whether the kernel reads its output back (InOut), as C += A * B does.

reference_indices

reference_indices() -> list[int]

Argument positions handed to reference, in order.

An InOut output is excluded like an Out one: the reference computes the result from the declared initializer's state. The builder initializes the buffer before each independent tile call.

Source code in python/iron/kernels/_common.py

def reference_indices(self) -> list[int]:
    """Argument positions handed to ``reference``, in order.

    An ``InOut`` output is excluded like an ``Out`` one: the reference
    computes the result from the declared initializer's state. The
    builder initializes the buffer before each independent tile call.
    """
    bound = dict(self.parameter_bindings)
    return [
        i
        for i, r in enumerate(self.roles)
        if r not in (Out, InOut) and i not in bound
    ]

validate_types

validate_types(arg_types)

Validate contracts against NumPy tensor aliases and scalar dtypes.

Raw MLIR types remain usable by ExternalFunction, but the host contract requires NumPy declarations for sampling, layouts and references.

Source code in python/iron/kernels/_common.py

def validate_types(self, arg_types):
    """Validate contracts against NumPy tensor aliases and scalar dtypes.

    Raw MLIR types remain usable by ExternalFunction, but the host contract
    requires NumPy declarations for sampling, layouts and references.
    """
    if len(arg_types) != len(self.roles):
        raise ValueError("roles must have one entry per argument")
    bound = dict(self.parameter_bindings)
    for i, (role, arg_type) in enumerate(zip(self.roles, arg_types)):
        tensor = _is_tensor_type(arg_type)
        if tensor:
            try:
                np_ndarray_type_get_shape(arg_type)
                dtype = np_ndarray_type_get_dtype(arg_type)
            except (AssertionError, IndexError, TypeError) as exc:
                raise ValueError(
                    f"argument {i}: expected np.ndarray[shape, np.dtype[dtype]]"
                ) from exc
        else:
            dtype = arg_type
        if dtype not in get_args(NpuDType):
            raise ValueError(
                f"argument {i}: kernel contracts require NumPy tensor aliases "
                "or supported NumPy scalar dtypes"
            )
        if not tensor and role is not Param:
            raise ValueError(f"argument {i}: scalar arguments require Param")
        if self.layouts and self.layouts[i] is not None and not tensor:
            raise ValueError(f"argument {i}: layouts require tensor arguments")
        if i not in bound:
            continue
        value = bound[i]
        if tensor:
            value = np.asarray(value)
            layout = self.layouts[i] if self.layouts else None
            shape = (
                layout.shape
                if layout is not None
                else np_ndarray_type_get_shape(arg_type)
            )
            if value.ndim == 0 or value.size != int(np.prod(shape)):
                raise ValueError(
                    f"argument {i}: tensor parameter must contain {shape} elements"
                )
        elif not isinstance(value, (int, float, np.integer, np.floating)):
            raise ValueError(f"argument {i}: expected scalar parameter")

TensorLayout dataclass

TensorLayout(
    shape: tuple[int, ...],
    pack: Callable | None = None,
    unpack: Callable | None = None,
    stream: TensorAccessPattern | None = None,
    block: tuple[int, ...] | None = None,
)

How a kernel wants one tensor operand laid out.

shape is the logical tile. pack and unpack are the reversible host codec between (calls, *shape) and (calls, storage_elements); identity is the default. stream is the DMA transform (to_stream) a design applies on the hop that feeds this operand to the kernel or drains it, None when the operand streams as stored; block is the micro-tile the kernel consumes or produces, (r, s) for an MMUL operand. The codec is built from the same two facts, so the host and the design agree by construction. None of this is an algorithm or a whole-problem iteration schedule.

Param

Read-only test-fixture parameter; its ABI determines scalar or tensor.

The generic harness holds its value fixed across calls. This is not a C++ operand lifetime: direct designs may pass a new value on every kernel call.

Trace dataclass

Trace(shape: str, reason: str | None = None)

How a kernel's event0()/event1() markers bracket one call.

Trace.whole_call(): one pair brackets every call of the entry symbol and nothing it calls emits another, so each trace interval is one call. Trace.none(reason): a call emits no marker. Trace.partial(reason): markers exist but do not bracket each call exactly once (around an inner loop, or skipped on an early return), so intervals cannot be attributed to calls. test_kernel_trace_markers.py checks the declaration against the compiled IR of every library build.

aie.iron.algorithms.kernel_design turns any contract-bearing factory into a design of one Worker, built on the same single-core pipeline as transform, for_each and reduce. What a kernel can answer about itself -- its reference result, its safe input range, which arguments are parameters, how to judge a device output -- lives on ExternalFunction instead, so bringing up a kernel needs no test harness. See Kernel Library for the add-a-kernel procedure and the test tiers built on it.

Import the builder with from aie.iron.algorithms import kernel_design as kd and call kd.design(...), or import design directly from aie.iron.algorithms. The former aie.utils.kernel_harness module has been removed.

Build, sample and check independent kernel calls from their declarations.

Each call consumes one tile per In, reads constant Param values, and writes one tile per output. Layout codecs convert logical tiles to kernel storage on the host. The Worker and the runtime sequence are the shared single-core pipeline's (_pipeline). Whole-problem tiling, reductions across kernel calls and multi-core schedules belong to algorithms, not this kernel-validation harness.

CallCycles dataclass

CallCycles(
    kernel: tuple[int, ...] = (),
    initializers: dict[int, tuple[int, ...]] = dict(),
    setup: tuple[int, ...] = (),
    truncated: bool = False,
    untimed: str | None = None,
)

One traced run's intervals, split by the kernel that emitted them.

kernel holds one interval per call of the measured kernel, in call order; initializers the same for each traced initializer, keyed by the InOut argument it initializes; setup the setup kernel's one interval when it is traced. A trace that fills its buffer keeps a prefix of the stream, which the split still labels correctly, and truncated says the lists are short. untimed is the contract's reason when the kernel's markers do not bracket its calls; then nothing ran.

design

design(
    factory,
    *,
    calls=1,
    scalars=(),
    shape=None,
    params=None,
    aiecc_flags=None,
    guard=False,
    stack_bytes=None,
    arg_byte_offsets=(),
    **factory_kwargs
)

Wrap tile calls; params/scalars supply unbound tensor/scalar Params.

This harness embeds these values for every call; changing them recompiles the design. Direct designs can supply different operands on each call.

stack_bytes replaces the core stack the contract declares, for a kernel built from sources other than the ones the contract was sized for.

With guard=True the core fills each output tile and GUARD_BYTES after it with 0x55 before every call and drains the guard with the tile, so a kernel that skips part of its output or writes past it shows up on the host: size the outputs with output_size(..., guard=True) and split them with strip_guard. bfp outputs carry no guard.

arg_byte_offsets binds tensor Params at a byte offset: ((1, 16),) hands argument 1 a view 16 bytes into a buffer VIEW_PAD larger, poisoned around the data, as a design that packs several weights into one buffer does. The shift is loaded at run time, so the kernel sees an address no design-time check could have folded.

Source code in python/iron/algorithms/kernel_design.py

def design(
    factory,
    *,
    calls=1,
    scalars=(),
    shape=None,
    params=None,
    aiecc_flags=None,
    guard=False,
    stack_bytes=None,
    arg_byte_offsets=(),
    **factory_kwargs,
):
    """Wrap tile calls; ``params``/``scalars`` supply unbound tensor/scalar Params.

    This harness embeds these values for every call; changing them recompiles
    the design. Direct designs can supply different operands on each call.

    ``stack_bytes`` replaces the core stack the contract declares, for a
    kernel built from sources other than the ones the contract was sized for.

    With ``guard=True`` the core fills each output tile and ``GUARD_BYTES``
    after it with ``0x55`` before every call and drains the guard with the
    tile, so a kernel that skips part of its output or writes past it shows
    up on the host:
    size the outputs with ``output_size(..., guard=True)`` and split them
    with ``strip_guard``. bfp outputs carry no guard.

    ``arg_byte_offsets`` binds tensor Params at a byte offset: ``((1, 16),)``
    hands argument 1 a view 16 bytes into a buffer ``VIEW_PAD`` larger,
    poisoned around the data, as a design that packs several weights into
    one buffer does. The shift is loaded at run time, so the kernel sees an
    address no design-time check could have folded.
    """
    calls = _calls(calls, shape)
    _device()
    fn = factory(**factory_kwargs)
    c = _contract(fn)
    if c.unsupported:
        raise ValueError(
            f"{fn.name}: the generic harness cannot build this kernel: {c.unsupported}"
        )
    _fifo_plan(fn)  # the DMA channel budget is checked before anything builds
    offsets = _byte_offsets(fn, arg_byte_offsets)
    flags: list[str] = list(aiecc_flags or ())
    if any(bfp.is_bfp(shape_dtype(t)[1]) for t in fn.arg_types() if _is_tensor_type(t)):
        if "--dynamic-objFifos" not in flags:
            flags.append("--dynamic-objFifos")
    # A gather reads its two tables at once, so they have to sit in different
    # banks. Chess cannot be checked (the flag needs Peano LLVM IR), and an
    # arch whose path has no LUT just finds nothing to check.
    if c.uses_lut and not fn.use_chess:
        if not any(f.startswith("--check-lut-banks") for f in flags):
            flags.append("--check-lut-banks")
    return _stream.specialize(
        factory=factory,
        factory_kwargs=factory_kwargs,
        calls=calls,
        # A key of its own: the contract that sets it can change in a module
        # the cache key never reads, and a stale stack overflows silently.
        stack_bytes=stack_bytes or _stack_bytes(fn),
        scalars=tuple(scalars),
        params=_encode_params(fn, params or ()),
        guard=guard,
        # Only when given, so every other design keeps its cache key.
        **({"arg_byte_offsets": offsets} if offsets else {}),
        **({"aiecc_flags": flags} if flags else {}),
    )

host_args

host_args(fn, *, calls=1, shape=None, guard=False)

Describe physical host buffers, not kernel arguments or constant Params.

Same-type streamed inputs share a buffer; each output has its own. Descriptors expose direction, shape, dtype and n_elements.

Source code in python/iron/algorithms/kernel_design.py

def host_args(fn, *, calls=1, shape=None, guard=False):
    """Describe physical host buffers, not kernel arguments or constant Params.

    Same-type streamed inputs share a buffer; each output has its own.
    Descriptors expose ``direction``, ``shape``, ``dtype`` and ``n_elements``.
    """
    calls = _calls(calls, shape)
    types = fn.arg_types()
    guarded = _guarded(fn, guard)
    result = []
    entries = [(In, g) for g in _fifo_plan(fn)[0]]
    entries += [(Out, [i]) for i in _contract(fn).out_indices]
    for direction, indices in entries:
        i = indices[0]
        dt = shape_dtype(types[i])[1]
        n = elems(types[i])
        if bfp.is_bfp(dt):
            n, dt = n * bfp.BLOCK_BYTES, np.uint8
        elif direction is Out and i in guarded:
            n += _guard_elems(types[i])
        rows = calls if direction is In else _out_tiles(fn, calls)
        s = (rows, n) if len(indices) == 1 else (rows, len(indices), n)
        result.append(_HostBuffer(direction, s, dt))
    return result

sample_inputs

sample_inputs(fn, *, calls=1, shape=None, rng=None)

One array per unbound In/tensor Param; only In has a call dimension.

Source code in python/iron/algorithms/kernel_design.py

def sample_inputs(fn, *, calls=1, shape=None, rng=None):
    """One array per unbound In/tensor Param; only In has a call dimension."""
    calls = _calls(calls, shape)
    rng = np.random.default_rng(0) if rng is None else rng
    c = _contract(fn)
    if c.sample is not None:
        return c.sample(rng, calls)
    result = []
    for i in _tensor_positions(fn)[0]:
        s, dt = shape_dtype(fn.arg_types()[i])
        layout = _layout(c, i)
        s = layout.shape if layout else (int(np.prod(s)),)
        dt = np.float32 if bfp.is_bfp(dt) else dt
        s = ((calls,) if c.roles[i] is In else ()) + s
        if np.issubdtype(np.dtype(dt), np.integer):
            r = fn.input_limit(dt) or _INT_RANGE.get(dt, 1 << 15)
            lo = 0 if np.dtype(dt).kind == "u" else -r
            result.append(rng.integers(lo, r, size=s).astype(dt))
        else:
            result.append(rng.standard_normal(s).astype(np.float32).astype(dt))
    return result

host_layout

host_layout(fn, inputs)

Pack declared layouts and interleave same-type In tiles; omit Param buffers.

Source code in python/iron/algorithms/kernel_design.py

def host_layout(fn, inputs):
    """Pack declared layouts and interleave same-type In tiles; omit Param buffers."""
    c = _contract(fn)
    positions = _tensor_positions(fn)[0]
    if len(inputs) != len(positions):
        raise ValueError(f"{fn.name}: expected {len(positions)} input arrays")
    values = dict(zip(positions, inputs))
    result = []
    for group in _fifo_plan(fn)[0]:
        arrays = []
        for i in group:
            value = np.asarray(values[i])
            layout = _layout(c, i)
            arrays.append(
                layout.encode(value)
                if layout
                else value.reshape(-1, elems(fn.arg_types()[i]))
            )
        result.append(
            np.ascontiguousarray(
                arrays[0] if len(arrays) == 1 else np.stack(arrays, axis=1)
            )
        )
    return result

output_size

output_size(fn, *, calls=1, shape=None, guard=False)

Storage elements per output; a tuple for multiple outputs.

Source code in python/iron/algorithms/kernel_design.py

def output_size(fn, *, calls=1, shape=None, guard=False):
    """Storage elements per output; a tuple for multiple outputs."""
    calls = _calls(calls, shape)
    types = fn.arg_types()
    guarded = _guarded(fn, guard)
    sizes = tuple(
        (elems(types[i]) + (_guard_elems(types[i]) if i in guarded else 0))
        * _out_tiles(fn, calls)
        * (bfp.BLOCK_BYTES if bfp.is_bfp(shape_dtype(types[i])[1]) else 1)
        for i in _contract(fn).out_indices
    )
    return sizes[0] if len(sizes) == 1 else sizes

upload

upload(inputs, out_size, out_dtype, *, fn, poison=False)

Return design inputs and output tensor(s); splat multiple outputs when calling.

Source code in python/iron/algorithms/kernel_design.py

def upload(inputs, out_size, out_dtype, *, fn, poison=False):
    """Return design inputs and output tensor(s); splat multiple outputs when calling."""
    ins = [
        tensor(a.reshape(-1), dtype=a.dtype, device="npu")
        for a in host_layout(fn, inputs)
    ]
    multiple = len(fn.contract.out_indices) > 1
    sizes, dtypes = (out_size, out_dtype) if multiple else ((out_size,), (out_dtype,))
    outs = [
        tensor(
            poisoned(n, dt) if poison else np.zeros(n, dtype=dt), dtype=dt, device="npu"
        )
        for n, dt in zip(sizes, dtypes)
    ]
    return ins, tuple(outs) if multiple else outs[0]

cycles_per_call

cycles_per_call(
    design_,
    inputs,
    out_size,
    out_dtype,
    *,
    fn,
    trace_size,
    workdir,
    calls=1
) -> CallCycles

Trace one run and split its intervals by kernel (CallCycles).

Only a kernel whose contract declares Trace.whole_call() is timed; one declaring none or partial returns its reason without running, and one declaring nothing raises. Traced initializers and a traced setup kernel are split off by position; a partial one raises, because its intervals cannot be labeled. Size trace_size from traced_intervals.

Source code in python/iron/algorithms/kernel_design.py

def cycles_per_call(
    design_, inputs, out_size, out_dtype, *, fn, trace_size, workdir, calls=1
) -> CallCycles:
    """Trace one run and split its intervals by kernel (``CallCycles``).

    Only a kernel whose contract declares ``Trace.whole_call()`` is timed;
    one declaring ``none`` or ``partial`` returns its reason without running,
    and one declaring nothing raises. Traced initializers and a traced setup
    kernel are split off by position; a ``partial`` one raises, because its
    intervals cannot be labeled. Size ``trace_size`` from
    ``traced_intervals``.
    """
    c = _contract(fn)
    if c.trace is None:
        raise ValueError(f"{fn.name}: the contract declares no trace")
    if c.trace.shape != "whole_call":
        return CallCycles(untimed=c.trace.reason)
    calls = _calls(calls)
    setup, inits = _traced(fn)
    workdir = Path(workdir)
    workdir.mkdir(parents=True, exist_ok=True)
    cfg = TraceConfig(trace_size=trace_size, trace_file=str(workdir / "trace.txt"))
    ins, out = upload(inputs, out_size, out_dtype, fn=fn)
    design_(*ins, *(out if isinstance(out, tuple) else (out,)), trace_config=cfg)
    if cfg.physical_mlir_path is None:
        raise RuntimeError("the traced run recorded no physical MLIR path")
    trace_json = workdir / "trace.json"
    cfg.trace_to_json(cfg.physical_mlir_path, str(trace_json))
    durations = [d for p in get_cycles_summary(str(trace_json)) for d in p[1:]]
    head, per_kernel, truncated = split_intervals(
        durations,
        calls=calls,
        per_call=len(inits) + 1,
        setup=int(setup),
        flush=TRACE_FLUSH,
    )
    if not per_kernel[-1]:
        raise RuntimeError(
            f"{fn.name}: the trace holds {len(durations)} intervals and none of "
            "the kernel's; the buffer is too small or the markers are missing"
        )
    return CallCycles(
        kernel=per_kernel[-1],
        initializers=dict(zip(inits, per_kernel)),
        setup=head,
        truncated=truncated,
    )

traced_intervals

traced_intervals(fn, *, calls=1) -> int

How many event0/event1 intervals a traced run of fn emits.

A trace_size that holds fewer truncates the run; 0 when the kernel itself is not timed.

Source code in python/iron/algorithms/kernel_design.py

def traced_intervals(fn, *, calls=1) -> int:
    """How many ``event0``/``event1`` intervals a traced run of ``fn`` emits.

    A ``trace_size`` that holds fewer truncates the run; ``0`` when the
    kernel itself is not timed.
    """
    trace = _contract(fn).trace
    if trace is None or trace.shape != "whole_call":
        return 0
    setup, inits = _traced(fn)
    return int(setup) + _calls(calls) * (len(inits) + 1)

split_intervals

split_intervals(
    durations, *, calls, per_call, setup=0, flush=0
)

Label an interval stream: setup intervals, then per_call per call.

The harness runs the setup kernel once, then each call runs its traced initializers in contract order and the kernel last, so interval j of call n sits at setup + n * per_call + j. Returns (setup intervals, one tuple per per-call kernel, truncated). Up to flush intervals of at most FLUSH_CYCLES may follow, the pairs the core emits to push the trace out. Anything else is some kernel emitting markers it does not declare; one whose extra intervals are that short and that few passes as a flush.

Source code in python/iron/algorithms/kernel_design.py

def split_intervals(durations, *, calls, per_call, setup=0, flush=0):
    """Label an interval stream: ``setup`` intervals, then ``per_call`` per call.

    The harness runs the setup kernel once, then each call runs its traced
    initializers in contract order and the kernel last, so interval ``j`` of
    call ``n`` sits at ``setup + n * per_call + j``. Returns ``(setup
    intervals, one tuple per per-call kernel, truncated)``. Up to ``flush``
    intervals of at most ``FLUSH_CYCLES`` may follow, the pairs the core
    emits to push the trace out. Anything else is some kernel emitting
    markers it does not declare; one whose extra intervals are that short
    and that few passes as a flush.
    """
    durations = [int(d) for d in durations]
    expected = setup + calls * per_call
    tail = durations[expected:]
    if len(tail) > flush or any(d > FLUSH_CYCLES for d in tail):
        raise RuntimeError(
            f"expected {expected} trace intervals and up to {flush} flush pairs, "
            f"got {len(durations)}, the extra {tail[:flush + 1]} cycles; a "
            "kernel on the core emits markers its contract's trace does not "
            "declare"
        )
    durations = durations[:expected]
    head, stream = durations[:setup], durations[setup:]
    return (
        tuple(head),
        [tuple(stream[j::per_call]) for j in range(per_call)],
        len(durations) < expected,
    )

test/python/npu/test_kernels_perf.py times the library: correctness first, then cycles, wall time and build size, gated by a device preflight and a measurement-sanity test. It is an ordinary pytest module, so -k selects cases and the session's exit status decides whether any numbers are written. --baseline-sources DIR measures every selected case a second time with its kernels from DIR and compares the two runs' raw output words and cycles in --perf-meta. The timing helpers live here:

Benchmarking helpers for NPU kernel callables.

Stats keeps the raw sample list and robust statistics (median, MAD, p95, coefficient of variation) beside avg/min/max, so callers can report jitter, not just central tendency. The numbers are numpy's.

Stats dataclass

Stats(
    avg_us: float,
    min_us: float,
    max_us: float,
    median_us: float = 0.0,
    mad_us: float = 0.0,
    p95_us: float = 0.0,
    cov: float = 0.0,
    n: int = 0,
    samples_us: list[float] = list(),
)

as_dict

as_dict(prefix: str = '') -> dict

Flat dict for JSON emission (samples omitted).

Source code in python/utils/benchmark.py

def as_dict(self, prefix: str = "") -> dict:
    """Flat dict for JSON emission (samples omitted)."""
    return {
        f"{prefix}median_us": self.median_us,
        f"{prefix}mad_us": self.mad_us,
        f"{prefix}p95_us": self.p95_us,
        f"{prefix}min_us": self.min_us,
        f"{prefix}max_us": self.max_us,
        f"{prefix}avg_us": self.avg_us,
        f"{prefix}cov": self.cov,
        f"{prefix}n": self.n,
    }

BenchmarkResult dataclass

BenchmarkResult(e2e: Stats, npu: Stats | None)

Preflight dataclass

Preflight(
    npu: str, arch: str, device: str, pmode: str | None
)

What the runtime says about the device a run is about to use.

run_iters

run_iters(
    fn: Callable,
    *args,
    warmup: int = 0,
    iters: int = 1,
    arg_sets: list[tuple] | None = None,
    **kwargs
) -> BenchmarkResult

Invoke fn warmup + iters times, reporting timings.

End-to-end latency is measured around the Python call. If the return value carries an npu_time (nanoseconds, captured by the runtime around kernel.wait()), it is reported separately so callers can see the host-side overhead delta.

arg_sets: an optional list of positional-argument tuples to rotate through, one per iteration, so consecutive runs do not hit the same host buffers (buffer rotation; see the CUTLASS measurement guidelines). When given, *args must be empty.

Source code in python/utils/benchmark.py

def run_iters(
    fn: Callable,
    *args,
    warmup: int = 0,
    iters: int = 1,
    arg_sets: list[tuple] | None = None,
    **kwargs,
) -> BenchmarkResult:
    """Invoke ``fn`` ``warmup + iters`` times, reporting timings.

    End-to-end latency is measured around the Python call. If the return
    value carries an ``npu_time`` (nanoseconds, captured by the runtime
    around ``kernel.wait()``), it is reported separately so callers can see
    the host-side overhead delta.

    ``arg_sets``: an optional list of positional-argument tuples to
    rotate through, one per iteration, so consecutive runs do not hit the
    same host buffers (buffer rotation; see the CUTLASS measurement
    guidelines). When given, ``*args`` must be empty.
    """
    if iters < 1:
        raise ValueError(f"iters must be >= 1 (got {iters})")
    if warmup < 0:
        raise ValueError(f"warmup must be >= 0 (got {warmup})")
    if arg_sets is not None and args:
        raise ValueError("pass either *args or arg_sets, not both")
    if arg_sets is not None and not arg_sets:
        raise ValueError("arg_sets must hold at least one argument tuple")
    rotation = arg_sets if arg_sets is not None else [args]

    e2e_samples: list[float] = []
    npu_samples: list[float] = []
    for i in range(warmup + iters):
        call_args = rotation[i % len(rotation)]
        start = time.perf_counter()
        ret = fn(*call_args, **kwargs)
        e2e_us = (time.perf_counter() - start) * 1_000_000
        if i < warmup:
            continue
        e2e_samples.append(e2e_us)
        npu_ns = _extract_npu_time_ns(ret)
        if npu_ns is not None:
            npu_samples.append(npu_ns / 1_000.0)

    return BenchmarkResult(
        e2e=_stats(e2e_samples),
        npu=_stats(npu_samples) if npu_samples else None,
    )

preflight

preflight() -> Preflight

Describe the device through the host runtime (any backend).

Source code in python/utils/benchmark.py

def preflight() -> Preflight:
    """Describe the device through the host runtime (any backend)."""
    import aie.utils as aie_utils
    from aie.utils.compile.utils import resolve_target_arch

    runtime = aie_utils.DefaultNPURuntime
    if runtime is None:
        raise RuntimeError("no NPU runtime is available (XRT, HRX or HSA)")
    device = runtime.device()
    arch = resolve_target_arch(device)
    return Preflight(
        npu="npu2" if arch == "aie2p" else "npu1",
        arch=arch,
        device=runtime.device_name() or type(device).__name__,
        pmode=runtime.power_mode(),
    )

provenance

provenance(**extra: str | None) -> str

Return a one-line description of what produced a measurement.

The git commit (GITHUB_SHA or git rev-parse HEAD), the Peano that compiles the kernels (peano_version()), the installed mlir_aie version, the kernel tree (MLIR_AIE_KERNEL_SOURCES when set, and kernel_tree_digest()), the selected runtime's provenance, the runner the job ran on (RUNNER_NAME, under GitHub Actions: nightly runs move between hosts of a generation, and a host swap moves npu_us as a regression would), and any extra fields (device="NPU Strix", pmode="performance") as key value pairs. A benchmark row records it so a number can be traced to a toolchain, the kernel sources and the machine.

A package that is not installed is left out rather than recorded as unknown. CI builds mlir_aie from source and puts it on PYTHONPATH, so there is no distribution to read a version from, and every published row would otherwise carry a word that reads like a lookup failure. The commit already identifies that build.

Source code in python/utils/benchmark.py

def provenance(**extra: str | None) -> str:
    """Return a one-line description of what produced a measurement.

    The git commit (``GITHUB_SHA`` or ``git rev-parse HEAD``), the Peano that
    compiles the kernels (``peano_version()``), the installed ``mlir_aie``
    version, the kernel tree
    (``MLIR_AIE_KERNEL_SOURCES`` when set, and ``kernel_tree_digest()``),
    the selected runtime's provenance, the runner the job ran
    on (``RUNNER_NAME``, under GitHub Actions: nightly runs move between
    hosts of a generation, and a host swap moves ``npu_us`` as a regression
    would), and any ``extra`` fields (``device="NPU Strix"``,
    ``pmode="performance"``) as ``key value`` pairs. A benchmark row records
    it so a number can be traced to a toolchain, the kernel sources and
    the machine.

    A package that is not installed is left out rather than recorded as
    unknown. CI builds ``mlir_aie`` from source and puts it on ``PYTHONPATH``,
    so there is no distribution to read a version from, and every published
    row would otherwise carry a word that reads like a lookup failure. The
    commit already identifies that build.
    """
    import aie.utils as aie_utils

    def pkg(name: str) -> str | None:
        try:
            return _pkg_version(name)
        except PackageNotFoundError:
            return None

    commit = os.environ.get("GITHUB_SHA")
    if not commit:
        try:
            commit = subprocess.run(
                ["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True
            ).stdout.strip()
        except (OSError, subprocess.CalledProcessError):
            commit = "unknown"
    runtime_fields = {}
    try:
        runtime = aie_utils.DefaultNPURuntime
        if runtime is not None:
            runtime_fields = runtime.provenance()
    except Exception:  # noqa: BLE001 - optional metadata must not fail a measurement
        pass
    fields = {
        "commit": commit[:10],
        "peano": peano_version(),
        "mlir_aie": pkg("mlir_aie"),
        "kernel_sources": os.environ.get("MLIR_AIE_KERNEL_SOURCES"),
        "kernels": kernel_tree_digest(),
        **runtime_fields,
        "host": os.environ.get("RUNNER_NAME"),
        **extra,
    }
    return " | ".join(f"{k} {v}" for k, v in fields.items() if v)

kernel_tree_digest

kernel_tree_digest() -> str | None

Return a 12-hex digest of the kernel sources the library factories compile.

Every file under aie_kernels_dir() and aie_runtime_lib_dir(), by relative path and content. The commit alone cannot say which kernels ran: MLIR_AIE_KERNEL_SOURCES can name another tree, and a checkout can carry uncommitted edits. None when neither directory exists.

Source code in python/utils/benchmark.py

def kernel_tree_digest() -> str | None:
    """Return a 12-hex digest of the kernel sources the library factories compile.

    Every file under ``aie_kernels_dir()`` and ``aie_runtime_lib_dir()``, by
    relative path and content. The commit alone cannot say which kernels
    ran: ``MLIR_AIE_KERNEL_SOURCES`` can name another tree, and a checkout
    can carry uncommitted edits. ``None`` when neither directory exists.
    """
    h = hashlib.sha256()
    found = False
    for top in (config.aie_kernels_dir(), config.aie_runtime_lib_dir()):
        root = Path(top)
        if not root.is_dir():
            continue
        found = True
        for path in sorted(p for p in root.rglob("*") if p.is_file()):
            h.update(f"{root.name}/{path.relative_to(root)}\0".encode())
            h.update(path.read_bytes())
    return h.hexdigest()[:12] if found else None

Static checks

These compiler-remark checks are available on demand through python -m aie.utils.compile.remarks; there is no static-check CI workflow. The trace-marker check (test/python/test_kernel_trace_markers.py) runs in lit on every PR, through trace_markers below.

Static kernel checks: compile every library kernel with Peano and read its remarks.

python -m aie.utils.compile.remarks --target aie2p --out static.json         --out-pm static-pm.json --meta static-meta.json
python -m aie.utils.compile.remarks --target aie2p --only '^gelu'         --out static.json --sources . --baseline-sources ../mlir-aie-base
python -m aie.utils.compile.remarks --target aie2p --only '^tanh/'         --cases test/python/npu/kernel_cases.py --out static.json
python -m aie.utils.compile.remarks --target aie2p --out static.json         --build cascade_mm:dim_m=16,dim_k=24,dim_n=32 --build cascade_mm

CPU-only. Every factory in aie.iron.kernels (at its defaults and for each entry of its .dtypes table) is compiled exactly as the JIT compiles it (aie.utils.compile.utils.cxx_core_compile_command), plus the optimization-record flags below, and the records become per-kernel series for benchmark-action: a Peano bump that changes a loop's schedule shows up here before anyone looks at device numbers. With --sources DIR (or MLIR_AIE_KERNEL_SOURCES=DIR) naming a checkout, that checkout's aie_kernels/ and aie_runtime_lib/ are compiled instead of the installed copies. With --cases, the builds are the ones a cases file's tests run (shape and options baked in), each named by its case; with --build, the factory builds named on the command line, for a shape that picks a code path no default reaches.

Record shapes, as llvm-aie 22.0.0.2026090201 emits them (they are Peano's, not LLVM's documented ones):

Pass	Kind / Name	Args	Tracked as
pipeliner	Passed / schedule	II, NS, Loop, Pipeliner, prologue/epilogue bundles	loop/<fn>/<bb>/II (rest as hover text)
pipeliner	Missed / canPipelineLoop	"Failed to pipeline loop"; located by DebugLoc only	unpipelined_loops (keyed L<line>)
pipeliner	Analysis / schedule	MII, SwpMaxMii, "Unable to find schedule"	schedule_notes in the meta file
aie-hardware-loops	Analysis / analysis	BasicBlock, Zero-Overhead-Loop	non_zol_loops, and loop/<fn>/<bb>/not_zol per loop
aie-asm-printer	Analysis / analysis	BasicBlock, BundleCount, ByteCount	pm_bytes (summed over the shipped functions)
aie-multi-slot-pseudo	Missed / missing-memory-bank	Instruction	missing_bank_loads
stderr	-Wpass-failed	a #pragma clang loop / AIE_* macro the compiler dropped	pass_failed_warnings, text kept
the object	llvm-readobj sections, symbols, relocations	what the entry symbol reaches	the shipped functions; libcalls (e.g. __divsf3)
the object	llvm-readobj --stack-sizes (-fstack-size-section)	frame bytes per function	kernel_stack_bytes, the deepest path from the entry
the IR (-emit-llvm)	opt print<scalar-evolution>	a constant backedge-taken count	loop/<fn>/<bb>/II_x_trips

The loop counts and pm_bytes cover only the functions the entry symbol reaches in the object, which are the ones the core link keeps. kernel_stack_bytes counts the kernel's frames only: the core that calls it also holds main's, which aiecc's measured stack size includes (dwconv1d_channels_last on AIE2P: 64 here, 256 for the core). Over the contract's stack_bytes (else the device default) it prints a warning: the design reserves that much for the whole core, and an overflow corrupts the neighbouring memory without a fault. --baseline-sources DIR compiles everything a second time from DIR and prints each row that differs, for a before/after of a kernel change; a loop LLVM only renamed, its rows unchanged, is counted but not listed, and a build either tree fails to compile is named and left out. It names the aie_kernels/ this tree compiled and warns when that is the installed copy (MLIR_AIE_KERNEL_SOURCES unset) or the baseline itself. --keep DIR keeps the objects, which --meta names per build.

The loop-scheduling pass reports as pipeliner (a postpipeliner filter records nothing); it names loops by machine basic block (bb.1.for.body.i) while the other passes use the IR block (for.body.i), so the prefix is stripped to join them. unpipelined_loops counts every loop the pipeliner declined, outer loops included, so its change is the signal, not its value.

The integer series alert on any increase; pm_bytes goes to its own file (--out-pm) so it can carry a percentage threshold. Nothing gates: under GitHub Actions a dropped pragma is a ::warning at its file and line, a kernel that fails to compile an ::error, and the run then exits 3 with "RESULTS INVALID", after writing the rows of the builds that compiled.

LoopInfo dataclass

LoopInfo(
    function: str,
    block: str,
    ii: int | None = None,
    ns: int | None = None,
    prologue_bundles: int | None = None,
    epilogue_bundles: int | None = None,
    pipelined: bool | None = None,
    pipeliner: str | None = None,
    missed_reason: str | None = None,
    zol: bool | None = None,
    bundle_count: int | None = None,
    byte_count: int | None = None,
    trips: int | None = None,
    file: str | None = None,
    line: int | None = None,
)

StaticReport dataclass

StaticReport(
    loops: dict[tuple[str, str], LoopInfo] = dict(),
    pm_bytes_by_function: dict[str, int] = dict(),
    missing_bank_loads: int = 0,
    pass_failed_warnings: int = 0,
    pass_failed: list[str] = list(),
    schedule_notes: list[str] = list(),
    shipped: set[str] | None = None,
    libcalls: list[str] = list(),
    kernel_stack_bytes: int | None = None,
)

Linked dataclass

Linked(
    functions: set[str],
    undefined: list[str],
    stack: int | None = None,
)

parse_yaml

parse_yaml(
    path: str | Path, report: StaticReport | None = None
) -> StaticReport

Source code in python/utils/compile/remarks.py

def parse_yaml(path: str | Path, report: StaticReport | None = None) -> StaticReport:
    r = report or StaticReport()
    with open(path) as f:
        docs = list(yaml.load_all(f, Loader=_Loader))
    for d in docs:
        if not isinstance(d, dict):
            continue
        p, name, kind = d.get("Pass"), d.get("Name"), d.get("_kind")
        fn = d.get("Function", "?")
        a = _args(d)
        if p == "pipeliner":
            if kind == "Passed" and name == "schedule":
                # A schedule with no Loop arg has been seen (zero_scalar_*);
                # fall back to the source line rather than a shared "?" key.
                loop = r.loop(fn, _block(a.get("Loop")) or _loc(d))
                loop.pipelined = True
                loop.ii = _int(a.get("II"))
                loop.ns = _int(a.get("NS"))
                loop.prologue_bundles = _int(a.get("PrologueBundles"))
                loop.epilogue_bundles = _int(a.get("EpilogueBundles"))
                loop.pipeliner = str(a.get("Pipeliner", "")) or None
                loop.file, loop.line = _source(d)
            elif kind == "Missed" and name == "canPipelineLoop":
                loop = r.loop(fn, _loc(d))
                loop.pipelined = False
                loop.missed_reason = str(a.get("String", "")) or None
                loop.file, loop.line = _source(d)
            elif kind == "Analysis" and name == "schedule":
                # The sentence as clang prints it: every Arg value in order,
                # since the typed values (MII, II, ...) sit between the
                # String fragments and joining only the strings would print
                # "Minimal Initiation Interval too large:  > ." without them.
                message = "".join(
                    str(v)
                    for arg in d.get("Args", []) or []
                    if isinstance(arg, dict)
                    for k, v in arg.items()
                    if k != "DebugLoc"
                )
                r.schedule_notes.append(f"{fn}@{_loc(d)}: {message}")
        elif p == "aie-hardware-loops":
            bb = str(a.get("BasicBlock", "?"))
            r.loop(fn, bb).zol = str(a.get("Zero-Overhead-Loop", "")).lower() == "true"
        elif p == "aie-asm-printer":
            bb = str(a.get("BasicBlock", "?"))
            bc, byc = _int(a.get("BundleCount")), _int(a.get("ByteCount"))
            r.pm_bytes_by_function[fn] = r.pm_bytes_by_function.get(fn, 0) + (byc or 0)
            if (fn, bb) in r.loops:  # only annotate known loop blocks
                r.loops[(fn, bb)].bundle_count = bc
                r.loops[(fn, bb)].byte_count = byc
        elif p == "aie-multi-slot-pseudo" and name == "missing-memory-bank":
            r.missing_bank_loads += 1
    return r

parse_stderr

parse_stderr(
    text: str, report: StaticReport
) -> StaticReport

Source code in python/utils/compile/remarks.py

def parse_stderr(text: str, report: StaticReport) -> StaticReport:
    dropped = [line.strip() for line in _WARN_PASS_FAILED.findall(text)]
    report.pass_failed_warnings += len(dropped)
    report.pass_failed.extend(dropped)
    # The stderr warning duplicates the YAML remark; only count it when no
    # YAML record was produced for the build (keeps the two channels consistent).
    if report.missing_bank_loads == 0:
        report.missing_bank_loads = len(_WARN_NO_BANK.findall(text))
    return report

report_rows

report_rows(
    report: StaticReport, prefix: str, extra: str
) -> list[dict]

benchmark-action rows for one kernel build. Smaller is better throughout.

Source code in python/utils/compile/remarks.py

def report_rows(report: StaticReport, prefix: str, extra: str) -> list[dict]:
    """benchmark-action rows for one kernel build. Smaller is better throughout."""
    out = [
        _row(f"{prefix}/unpipelined_loops", "loops", report.unpipelined_loops, extra),
        _row(f"{prefix}/non_zol_loops", "loops", report.non_zol_loops, extra),
        _row(f"{prefix}/missing_bank_loads", "loads", report.missing_bank_loads, extra),
        _row(
            f"{prefix}/pass_failed_warnings",
            "warnings",
            report.pass_failed_warnings,
            extra,
        ),
        _row(f"{prefix}/pm_bytes", "bytes", report.pm_bytes, extra),
        _row(
            f"{prefix}/libcalls",
            "symbols",
            len(report.libcalls),
            extra,
            " ".join(report.libcalls) or None,
        ),
    ]
    if report.kernel_stack_bytes is not None:
        out.append(
            _row(
                f"{prefix}/kernel_stack_bytes",
                "bytes",
                report.kernel_stack_bytes,
                extra,
                "the kernel's frames; the core adds main's"
                + (" and the runtime routines'" if report.libcalls else ""),
            )
        )
    for (fn, bb), loop in sorted(report.loops.items()):
        if loop.ii is not None:
            # The hover text names the source line, so a reader of an II
            # alert can open the loop without decoding a basic-block name.
            where = (
                f" at {Path(loop.file).name}:{loop.line}"
                if loop.file and loop.line is not None
                else ""
            )
            out.append(
                _row(
                    f"{prefix}/loop/{fn}/{bb}/II",
                    "cycles",
                    loop.ii,
                    extra,
                    f"NS={loop.ns} pro={loop.prologue_bundles} "
                    f"epi={loop.epilogue_bundles} zol={loop.zol} "
                    f"trips={loop.trips} via={loop.pipeliner}{where}",
                )
            )
            if loop.trips is not None:
                # II alone reads a loop that does 4 blocks per iteration as
                # slower than one that does 1; II x trips compares them.
                out.append(
                    _row(
                        f"{prefix}/loop/{fn}/{bb}/II_x_trips",
                        "cycles",
                        loop.ii * loop.trips,
                        extra,
                        f"{loop.trips} trips per entry{where}",
                    )
                )
        if loop.zol is not None:
            # 1 when the loop is not a zero-overhead loop: an inner loop that
            # falls off the hardware loop unit costs its bundle count in
            # branches, whatever its II says.
            out.append(
                _row(
                    f"{prefix}/loop/{fn}/{bb}/not_zol",
                    "loops",
                    int(not loop.zol),
                    extra,
                )
            )
    return out

workflow_annotations

workflow_annotations(
    name: str,
    report: StaticReport | None,
    detail: str,
    root: str | None = None,
    seen: set[tuple] | None = None,
) -> list[str]

Annotations for one kernel build: its dropped pragmas, or its compile failure.

One source is compiled once per factory build and per target, so the same dropped pragma comes back several times a run; seen (shared across builds) keeps each file, line and message to its first annotation.

Source code in python/utils/compile/remarks.py

def workflow_annotations(
    name: str,
    report: StaticReport | None,
    detail: str,
    root: str | None = None,
    seen: set[tuple] | None = None,
) -> list[str]:
    """Annotations for one kernel build: its dropped pragmas, or its compile failure.

    One source is compiled once per factory build and per target, so the same
    dropped pragma comes back several times a run; ``seen`` (shared across
    builds) keeps each file, line and message to its first annotation.
    """
    if report is None:
        first = next((ln for ln in detail.splitlines() if "error:" in ln), None)
        return [
            annotation(
                "error",
                first or detail.strip()[-500:],
                title=f"{name}: kernel failed to compile",
            )
        ]
    out = []
    for warning in report.pass_failed:
        m = _DIAG_RE.match(warning)
        key = (
            (m.group("file"), m.group("line"), warning[m.end() :]) if m else (warning,)
        )
        if seen is not None:
            if key in seen:
                continue
            seen.add(key)
        out.append(
            annotation(
                "warning",
                warning[m.end() :] if m else warning,
                title=f"{name}: pragma dropped by the compiler",
                file=m.group("file") if m else None,
                line=int(m.group("line")) if m else None,
                root=root,
            )
        )
    return out

compile_command

compile_command(
    ext_fn, target: str, out_dir: Path
) -> tuple[list[str], Path]

Return the exact Peano command the JIT would run for ext_fn, plus remark flags.

Inline-source kernels (the aie2 LUT activations) are written out under the kernel's symbol name first, as the JIT does. Kernels built with use_chess are rejected: the remarks are Peano's.

Source code in python/utils/compile/remarks.py

def compile_command(ext_fn, target: str, out_dir: Path) -> tuple[list[str], Path]:
    """Return the exact Peano command the JIT would run for ``ext_fn``, plus remark flags.

    Inline-source kernels (the aie2 LUT activations) are written out under the
    kernel's symbol name first, as the JIT does. Kernels built with
    ``use_chess`` are rejected: the remarks are Peano's.
    """
    src, include_dirs = _kernel_file(ext_fn, out_dir)
    yaml_out = out_dir / f"{ext_fn.name}.opt.yaml"
    cmd = cxx_core_compile_command(
        str(src),
        target,
        str(out_dir / f"{ext_fn.name}.o"),
        include_dirs=include_dirs,
        compile_args=[
            *ext_fn.compile_flags,
            *REMARK_FLAGS,
            f"-foptimization-record-file={yaml_out}",
            *_EXTRA_WARNINGS,
            # A non-allocated section of frame sizes; the code is unchanged.
            "-fstack-size-section",
        ],
    )
    return cmd, yaml_out

analyze

analyze(
    ext_fn, target: str, workdir: Path
) -> tuple[StaticReport | None, str]

Compile one kernel and parse its records; (None, reason) when it fails to compile.

Source code in python/utils/compile/remarks.py

def analyze(ext_fn, target: str, workdir: Path) -> tuple[StaticReport | None, str]:
    """Compile one kernel and parse its records; ``(None, reason)`` when it fails to compile."""
    cmd, yaml_out = compile_command(ext_fn, target, workdir)
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        return None, f"compile failed: {compile_failure(p.stderr)}"
    rep = parse_yaml(yaml_out) if yaml_out.exists() else StaticReport()
    for key, trips in trip_counts(ext_fn, target, workdir).items():
        if key in rep.loops:
            rep.loops[key].trips = trips
    parse_stderr(p.stderr, rep)
    reached = linked(workdir / f"{ext_fn.name}.o", entry_symbol(ext_fn))
    rep.shipped, rep.libcalls = reached.functions, reached.undefined
    rep.kernel_stack_bytes = reached.stack
    return rep, "ok"

kernel_builds

kernel_builds()

Yield (name, ExternalFunction) for every factory build the library offers.

Each exported factory at its defaults, plus one build per entry of its .dtypes table; a factory that refuses the current device (NotImplementedError) is skipped. Remarks depend on source and flags, not on the shape a test runs, so this is the whole surface.

Source code in python/utils/compile/remarks.py

def kernel_builds():
    """Yield ``(name, ExternalFunction)`` for every factory build the library offers.

    Each exported factory at its defaults, plus one build per entry of its
    ``.dtypes`` table; a factory that refuses the current device
    (``NotImplementedError``) is skipped. Remarks depend on source and flags,
    not on the shape a test runs, so this is the whole surface.
    """
    from aie.iron import kernels

    for name in kernels.factories():
        f = getattr(kernels, name)
        combos = [{}] + [dict(c) for c in getattr(f, "dtypes", ()) if c]
        seen: set[str] = set()
        for combo in combos:
            try:
                ef = f(**combo)
            except NotImplementedError:
                continue  # exists only for the other architecture
            if ef.object_file_name in seen:
                continue  # the default build is one of the dtypes entries
            seen.add(ef.object_file_name)
            yield _build_name(name, combo), ef

linked

linked(obj: Path, entry: str) -> Linked

Functions and undefined symbols entry reaches in the object obj.

Source code in python/utils/compile/remarks.py

def linked(obj: Path, entry: str) -> Linked:
    """Functions and undefined symbols ``entry`` reaches in the object ``obj``."""
    out = subprocess.run(
        [
            config.readobj_path(),
            "--elf-output-style=JSON",
            "--sections",
            "--symbols",
            "--relocations",
            "--stack-sizes",
            str(obj),
        ],
        capture_output=True,
        text=True,
        check=True,
    ).stdout
    return parse_readobj(json.loads(out)[0], entry)

parse_readobj

parse_readobj(doc: dict, entry: str) -> Linked

linked on the parsed llvm-readobj --elf-output-style=JSON document.

Source code in python/utils/compile/remarks.py

def parse_readobj(doc: dict, entry: str) -> Linked:
    """``linked`` on the parsed ``llvm-readobj --elf-output-style=JSON`` document."""
    alloc = {}
    for s in doc["Sections"]:
        s = s["Section"]
        alloc[s["Index"]] = any(f["Name"] == "SHF_ALLOC" for f in s["Flags"]["Flags"])
    symbols = [s["Symbol"] for s in doc["Symbols"]]
    # A symbol names its section by index; an undefined one names index 0.
    home = [s["Section"]["Value"] or None for s in symbols]
    functions = {}
    for s, sec in zip(symbols, home):
        if s["Type"]["Name"] == "Function" and sec is not None:
            functions.setdefault(sec, set()).add(s["Name"]["Name"])
    edges: dict[int, set[int]] = {}
    calls: dict[int, set[str]] = {}
    for rel in doc.get("Relocations", []):
        # A relocation section applies to the section its sh_info names.
        src = next(
            s["Section"]["Info"]
            for s in doc["Sections"]
            if s["Section"]["Index"] == rel["SectionIndex"]
        )
        for r in rel["Relocs"]:
            i = r["Relocation"]["Symbol"]["Value"]
            if home[i] is not None:
                edges.setdefault(src, set()).add(home[i])
            elif symbols[i]["Section"]["Name"] == "Undefined":
                calls.setdefault(src, set()).add(symbols[i]["Name"]["Name"])
    frames = {
        name: e["Entry"]["Size"]
        for e in doc.get("StackSizes", [])
        for name in e["Entry"]["Functions"]
    }

    def deepest(sec: int, path: frozenset) -> int | None:
        if sec in path:
            return None
        below = 0
        # Only code sections: a jump table in .rodata points back at its
        # function and would read as recursion.
        for callee in edges.get(sec, set()) - {sec}:
            if callee in functions:
                d = deepest(callee, path | {sec})
                if d is None:
                    return None
                below = max(below, d)
        return (
            max((frames.get(n, 0) for n in functions.get(sec, ())), default=0) + below
        )

    roots = [sec for sec, names in functions.items() if entry in names]
    todo = list(roots)
    seen: set[int] = set()
    while todo:
        sec = todo.pop()
        if sec in seen or not alloc.get(sec):
            continue
        seen.add(sec)
        todo.extend(edges.get(sec, ()))
    return Linked(
        functions={n for sec in seen for n in functions.get(sec, ())},
        undefined=sorted({u for sec in seen for u in calls.get(sec, ())}),
        stack=deepest(roots[0], frozenset()) if frames and roots else None,
    )

trace_markers

trace_markers(ir: str, entry: str) -> str

Classify how event0()/event1() bracket a call of entry.

"whole_call" when one pair brackets every call and nothing else emits a marker, "none" when a call emits no marker at all, and otherwise a sentence saying what the markers do instead.

Source code in python/utils/compile/remarks.py

def trace_markers(ir: str, entry: str) -> str:
    """Classify how ``event0()``/``event1()`` bracket a call of ``entry``.

    ``"whole_call"`` when one pair brackets every call and nothing else
    emits a marker, ``"none"`` when a call emits no marker at all, and
    otherwise a sentence saying what the markers do instead.
    """
    functions = _ir_functions(ir)
    if entry not in functions:
        raise ValueError(f"{entry} is not defined in the compiled IR")
    paths = _marker_paths(functions, entry, {})
    if isinstance(paths, str):
        return paths
    if paths == {"01"}:
        return "whole_call"
    if paths == {""}:
        return "none"
    shown = ", ".join(repr(p) for p in sorted(paths))
    return (
        f"a call of {entry} emits marker sequences {shown}, not one event0 then event1"
    )

trace_shape

trace_shape(ext_fn, target: str, out_dir: Path) -> str

Compile ext_fn to -O2 IR and classify its entry's markers (trace_markers).

Source code in python/utils/compile/remarks.py

def trace_shape(ext_fn, target: str, out_dir: Path) -> str:
    """Compile ``ext_fn`` to -O2 IR and classify its entry's markers (``trace_markers``)."""
    src, include_dirs = _kernel_file(ext_fn, out_dir)
    cmd = cxx_core_compile_command(
        str(src),
        target,
        str(out_dir / f"{ext_fn.name}.ll"),
        include_dirs=include_dirs,
        compile_args=list(ext_fn.compile_flags),
        inline=True,
    )
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        raise RuntimeError(f"{ext_fn.name}: {compile_failure(p.stderr)}")
    ir = (out_dir / f"{ext_fn.name}.ll").read_text()
    return trace_markers(ir, entry_symbol(ext_fn))

entry_symbol

entry_symbol(ext_fn) -> str

Return the symbol the kernel source defines, before the JIT's per-build prefix.

Source code in python/utils/compile/remarks.py

def entry_symbol(ext_fn) -> str:
    """Return the symbol the kernel source defines, before the JIT's per-build prefix."""
    return getattr(ext_fn, "_original_name", ext_fn.name)

Host-side helpers

aie.utils.bfp is the host side of the block-floating-point kernels: the bfp16ebs8 codec and the tile shuffle the mm_bfp DMA layout needs. It is the Python counterpart of programming_examples/ml/block_datatypes/helper.h, which the examples' C++ hosts use; a Python host encodes and checks with this module.

bfp16ebs8 on the host: encode, decode and the mmul-block shuffle.

v8bfp16ebs8 is the AIE2P block-floating-point type: eight values share one 8-bit exponent and each carries an 8-bit two's-complement mantissa, so a block of 8 values is 9 bytes ([exponent, m0, ..., m7]). These are numpy ports of floatToBfp16, bfp16ebs8ToFloat and shuffleMatrixForBfp16ebs8 in programming_examples/ml/block_datatypes/helper.h, bit for bit (the host test compiles that header and compares), including the header's rounding: mantissas truncate toward negative infinity, and a value more than 31 binades below its block's maximum becomes 0 (positive) or -1 LSB (negative).

NumPy structured arrays describe the packed storage. Neither NumPy nor ml_dtypes provides shared-exponent arithmetic for this format; the codec below supplies that conversion, not a replacement scalar dtype.

The block-floating-point matmul kernels (aie.iron.kernels.mm_bfp) load 8x8 sub-tiles as one 72-byte block vector, which a DMA cannot gather at 9-byte granularity, so tiles are rearranged by shuffle on the host: within each (tile_height, tile_width) tile the 8-row by 8-block sub-tiles are made contiguous in raster order. quantize is what a kernel sees of a float input, and what a reference should multiply.

encode

encode(x, *, rounding: str = 'floor') -> ndarray

float32 (..., n) with n % 8 == 0 -> uint8 (..., n * 9 // 8).

Blocks are taken along the last axis. Inputs must be finite (the C++ silently drops inf and NaN, which shifts every later value).

floor is the header's truncation. conv_even rounds each mantissa to nearest, ties to even, as amd/IRON's f32_to_bfp16ebs8 packs weights, byte for byte. A mantissa that rounds to +128 saturates to 127 there, so decode(encode(x, rounding="conv_even")) differs from quantize, which models the core raising the exponent instead.

Source code in python/utils/bfp.py

def encode(x, *, rounding: str = "floor") -> np.ndarray:
    """float32 ``(..., n)`` with ``n % 8 == 0`` -> ``uint8`` ``(..., n * 9 // 8)``.

    Blocks are taken along the last axis. Inputs must be finite (the C++
    silently drops inf and NaN, which shifts every later value).

    ``floor`` is the header's truncation. ``conv_even`` rounds each mantissa
    to nearest, ties to even, as amd/IRON's ``f32_to_bfp16ebs8`` packs
    weights, byte for byte. A mantissa that rounds to +128 saturates to 127
    there, so ``decode(encode(x, rounding="conv_even"))`` differs from
    ``quantize``, which models the core raising the exponent instead.
    """
    if rounding not in ("floor", "conv_even"):
        raise ValueError(
            f"bfp.encode: rounding must be 'floor' or 'conv_even', got {rounding!r}"
        )
    x = np.ascontiguousarray(x, dtype=np.float32)
    n = x.shape[-1]
    if n % BLOCK:
        raise ValueError(f"bfp.encode: last axis {n} is not a multiple of {BLOCK}")
    if not np.isfinite(x).all():
        raise ValueError("bfp.encode: inputs must be finite")
    lead = x.shape[:-1]
    bits = x.view(np.uint32).reshape(*lead, n // BLOCK, BLOCK)
    sign = (bits >> 31).astype(bool)
    exp = (bits >> 23) & 0xFF
    mant = (bits & 0x7FFFFF) | np.where(exp != 0, np.uint32(0x800000), np.uint32(0))
    max_exp = exp.max(axis=-1, keepdims=True)
    shift = (max_exp - exp).astype(np.int64)
    far = shift >= 32
    if rounding == "conv_even":
        # The quotient of a 24-bit magnitude by a power of two is exact in
        # float64, so rint is the only rounding.
        signed = np.where(sign, -mant.astype(np.int64), mant.astype(np.int64))
        v = np.rint(signed / np.exp2(np.minimum(_MANTISSA_SHIFT + shift, 62)))
        v = np.clip(v, -128, 127)
    else:
        # Two's complement in 32 bits, logical shift, low byte: the header's
        # `(uint8_t)((sign ? ~m + 1 : m) >> 17)`.
        m32 = np.where(sign, (-mant.astype(np.int64)) & 0xFFFFFFFF, mant).astype(
            np.uint32
        )
        v = (
            ((m32 >> _MANTISSA_SHIFT) & 0xFF)
            .astype(np.uint8)
            .view(np.int8)
            .astype(np.int32)
        )
        v = np.right_shift(v, np.minimum(shift, 31).astype(np.int32))  # arithmetic
    out = np.empty(lead + (n // BLOCK,), dtype=_BLOCK_DTYPE)
    out["exponent"] = max_exp[..., 0]
    out["mantissas"] = np.where(far, np.where(sign, -1, 0), v)
    return out.view(np.uint8).reshape(*lead, n // BLOCK * BLOCK_BYTES)

decode

decode(b) -> ndarray

uint8 (..., n * 9 // 8) -> float32 (..., n); the exact inverse map of a block.

Source code in python/utils/bfp.py

def decode(b) -> np.ndarray:
    """``uint8`` ``(..., n * 9 // 8)`` -> float32 ``(..., n)``; the exact inverse map of a block."""
    b = np.ascontiguousarray(b, dtype=np.uint8)
    nb = b.shape[-1]
    if nb % BLOCK_BYTES:
        raise ValueError(
            f"bfp.decode: last axis {nb} is not a multiple of {BLOCK_BYTES}"
        )
    lead = b.shape[:-1]
    blk = b.view(_BLOCK_DTYPE)
    scale = np.ldexp(1.0, blk["exponent"].astype(np.int32) - 127 - 6)
    vals = blk["mantissas"].astype(np.float64) * scale[..., None]
    return vals.astype(np.float32).reshape(*lead, nb // BLOCK_BYTES * BLOCK)

quantize

quantize(x, *, rounding: str = 'floor') -> ndarray

Return what a kernel reads of x, for a given conversion rounding mode.

Which mode applies is a property of who converts, not of where. floor is decode(encode(x)): this module's encoder truncates toward negative infinity, and a core converting in floor mode agrees with it -- the q4nx_dequant kernel pins floor and its reference matches the device byte for byte. conv_even models a core converting with round-to-nearest-ties-to-even, which is what mm_bfp's mixed kernel pins so its K reduction does not accumulate a one-sided bias.

The mode is load-bearing, not a detail: on 64x64x64 mixed tiles of random and large inputs, the right one reproduces the kernel's bf16 output bit for bit, and floor mismatches about 3550 of each tile's 4096 outputs.

A mantissa that rounds up to +128 does not fit the 8-bit field, so the block's exponent goes up by one and the block is requantized. Clamping instead would cost a whole step to the one element the shared exponent was chosen for. The carry is one-sided, as on the core: -128 fits, so a block whose most negative value rounds to -128 keeps its exponent.

Source code in python/utils/bfp.py

def quantize(x, *, rounding: str = "floor") -> np.ndarray:
    """Return what a kernel reads of ``x``, for a given conversion rounding mode.

    Which mode applies is a property of who converts, not of where. ``floor``
    is ``decode(encode(x))``: this module's encoder truncates toward negative
    infinity, and a *core* converting in floor mode agrees with it -- the
    ``q4nx_dequant`` kernel pins floor and its reference matches the device
    byte for byte. ``conv_even`` models a core converting with
    round-to-nearest-ties-to-even, which is what ``mm_bfp``'s mixed kernel
    pins so its K reduction does not accumulate a one-sided bias.

    The mode is load-bearing, not a detail: on 64x64x64 mixed tiles of random
    and large inputs, the right one reproduces the kernel's bf16 output bit
    for bit, and floor mismatches about 3550 of each tile's 4096 outputs.

    A mantissa that rounds up to +128 does not fit the 8-bit field, so the
    block's exponent goes up by one and the block is requantized. Clamping
    instead would cost a whole step to the one element the shared exponent
    was chosen for. The carry is one-sided, as on the core: -128 fits, so a
    block whose most negative value rounds to -128 keeps its exponent.
    """
    if rounding == "floor":
        return decode(encode(x))
    if rounding != "conv_even":
        raise ValueError(
            f"bfp.quantize: rounding must be 'floor' or 'conv_even', got {rounding!r}"
        )
    x = np.ascontiguousarray(x, dtype=np.float32)
    n = x.shape[-1]
    if n % BLOCK:
        raise ValueError(f"bfp.quantize: last axis {n} is not a multiple of {BLOCK}")
    blocks = x.reshape(*x.shape[:-1], n // BLOCK, BLOCK)
    exp = (blocks.view(np.uint32) >> 23) & 0xFF
    scale = np.ldexp(1.0, exp.max(axis=-1, keepdims=True).astype(np.int32) - 127 - 6)
    mant = np.rint(blocks.astype(np.float64) / scale)
    carry = (mant.max(axis=-1, keepdims=True) > 127)[..., 0]
    if carry.any():
        scale = np.where(carry[..., None], scale * 2, scale)
        mant = np.rint(blocks.astype(np.float64) / scale)
    return (np.clip(mant, -128, 127) * scale).astype(np.float32).reshape(x.shape)

shuffle

shuffle(
    b,
    width: int,
    height: int,
    tile_width: int,
    tile_height: int,
    *,
    unshuffle=False
) -> ndarray

Reorder an encoded (height, width) matrix into (or out of) the mmul tile layout.

b is the encoded matrix (height rows of width * 9 // 8 bytes, flat or 2-D; width and tile_width count values, so both are multiples of 8 and the tile is tile_height rows, a multiple of 8). Within each tile, every 8-row by 8-value sub-tile (72 bytes) becomes contiguous, sub-tiles in raster order; the tiles themselves stay where they are, so a DMA that copies a tile row by row delivers the sub-tiles the kernel's block-vector loads expect. unshuffle is the inverse.

Source code in python/utils/bfp.py

def shuffle(
    b, width: int, height: int, tile_width: int, tile_height: int, *, unshuffle=False
) -> np.ndarray:
    """Reorder an encoded ``(height, width)`` matrix into (or out of) the mmul tile layout.

    ``b`` is the encoded matrix (``height`` rows of ``width * 9 // 8`` bytes,
    flat or 2-D; ``width`` and ``tile_width`` count values, so both are
    multiples of 8 and the tile is ``tile_height`` rows, a multiple of 8).
    Within each tile, every 8-row by 8-value sub-tile (72 bytes) becomes
    contiguous, sub-tiles in raster order; the tiles themselves stay where
    they are, so a DMA that copies a tile row by row delivers the sub-tiles
    the kernel's block-vector loads expect. ``unshuffle`` is the inverse.
    """
    for name, v, mult in (
        ("width", width, BLOCK),
        ("tile_width", tile_width, BLOCK),
        ("tile_height", tile_height, BLOCK),
    ):
        if v <= 0 or v % mult:
            raise ValueError(
                f"bfp.shuffle: {name} must be a positive multiple of {mult}"
            )
    if width % tile_width or height % tile_height:
        raise ValueError("bfp.shuffle: the tile must divide the matrix")
    W, tw = width * BLOCK_BYTES // BLOCK, tile_width * BLOCK_BYTES // BLOCK
    a = np.asarray(b, dtype=np.uint8).reshape(height, W)
    ty, tx = height // tile_height, W // tw
    sy, sx = tile_height // BLOCK, tw // BLOCK_BYTES
    if not unshuffle:
        # [tileY, subY, i, tileX, subX, j] -> [tileY, tileX, subY, subX, i, j]
        v = a.reshape(ty, sy, BLOCK, tx, sx, BLOCK_BYTES).transpose(0, 3, 1, 4, 2, 5)
        # ... flattened raster-wise into the tile's (tile_height, tw) bytes
        v = v.reshape(ty, tx, tile_height, tw).transpose(0, 2, 1, 3)
    else:
        v = a.reshape(ty, tile_height, tx, tw).transpose(0, 2, 1, 3)
        v = v.reshape(ty, tx, sy, sx, BLOCK, BLOCK_BYTES).transpose(0, 2, 4, 1, 3, 5)
    return np.ascontiguousarray(v.reshape(height, W))

Tolerance-based output verification helpers for examples and tests.

Mirrors the canonical test_utils::nearly_equal semantics used across the C++ testbenches so Python migrations of those examples behave identically:

|a - b|  <  max(atol, rtol * (|a| + |b|))

Defaults match the C++ default of rtol=0.128, which is the widely-used relative tolerance for bfloat16 / LUT-approximated kernels (exp, softmax, gelu, silu, swiglu, ...).

Tolerance dataclass

Tolerance(
    rtol: float | None = None,
    atol: float | None = None,
    ulps: int | None = None,
    range_frac: float | None = None,
    max_mismatch_frac: float = 0.0,
    note: str = "",
    bound: Callable | None = None,
)

How close a kernel's output must be to its reference.

Exactly one of four kinds, chosen by which fields are set:

• exact -- no field set: bit-equal after casting the reference to the output dtype. Integers, selections (relu, max), lossless copies.
• ulps -- ulps set: bf16 outputs within ulps units in the last place of the correctly rounded reference. atol may be set alongside as a floor, admitting an element that meets either -- what a kernel needs when the device flushes subnormals to zero, since a flushed value is a full 100% relative and dozens of ulps from the reference but absolutely negligible. rtol stays unset, or the kind is relative.
• relative -- rtol/atol set: the canonical |a - b| < max(atol, rtol * (|a| + |b|)) of nearly_equal. Integer outputs are compared with the same formula in exact integer arithmetic, so lsb (atol = n + 0.5) admits an n-LSB slack for fixed-point pixel kernels whose rounding shift is not modeled; under exact and ulps integers stay bit-equal.
• bound -- bound set: |a - b| <= bound(*args) element by element, where args are the reference's own arguments. For a floating-point kernel whose error is set by where its input falls rather than by what it produced -- an approximation exact in one range and a few ulps off in another -- so any one rtol/atol is either loose everywhere or unsound somewhere. compare takes the evaluated bound; ExternalFunction.judge evaluates it from the inputs it is given.

Non-finite values are never skipped: NaN must meet NaN, and an infinity must meet an infinity of the same sign, under every kind.

range_frac adds a floor scaled to the reference's own range: an element also passes at |a - b| <= range_frac * max|b|. It is for a kernel whose error is set by the magnitudes it worked from rather than by the magnitude it produced -- a dot product whose terms cancel to near zero is no less accurate than its neighbours, but an elementwise relative bound reads it as 100% wrong. The scale comes from expected, never from actual, so a kernel cannot widen its own tolerance by returning something large. Unlike atol it follows the data: the same fraction holds whether the outputs run to 34 or to 3.4e9, where a fixed floor would be either dead or permissive. Set it from a measured worst case, and say so in note.

max_mismatch_frac allows that fraction of elements to miss (LUT tails, saturation edges). note records where the number came from -- a docstring, a device run, a testbench default -- so a reviewer can tell an evidenced tolerance from a guessed one.

bf16_ulps classmethod

bf16_ulps(
    n: int = 1,
    *,
    atol: float | None = None,
    max_mismatch_frac: float = 0.0,
    note: str = ""
) -> "Tolerance"

bf16 outputs within n ulps of the correctly rounded reference.

atol is an optional strict bound (error < atol) an element may meet instead of the ulp bound. Use it for the device's subnormal flush to zero, set to the smallest normal bf16 so it admits the flushed values and nothing above them.

Source code in python/utils/verify.py

@classmethod
def bf16_ulps(
    cls,
    n: int = 1,
    *,
    atol: float | None = None,
    max_mismatch_frac: float = 0.0,
    note: str = "",
) -> "Tolerance":
    """bf16 outputs within ``n`` ulps of the correctly rounded reference.

    ``atol`` is an optional strict bound (``error < atol``) an element may
    meet instead of the ulp bound. Use it for the device's subnormal flush
    to zero, set to the smallest normal bf16 so it admits the flushed
    values and nothing above them.
    """
    return cls(ulps=n, atol=atol, max_mismatch_frac=max_mismatch_frac, note=note)

bounded classmethod

bounded(
    bound: Callable,
    *,
    max_mismatch_frac: float = 0.0,
    note: str = ""
) -> "Tolerance"

Each element within bound(*args) of the reference, absolutely.

bound takes the reference's arguments and returns one non-negative bound per output element, shaped like the reference's result (a tuple of them for several outputs). Derive it from the kernel's arithmetic and say in note which parts were measured.

Source code in python/utils/verify.py

@classmethod
def bounded(
    cls, bound: Callable, *, max_mismatch_frac: float = 0.0, note: str = ""
) -> "Tolerance":
    """Each element within ``bound(*args)`` of the reference, absolutely.

    ``bound`` takes the reference's arguments and returns one
    non-negative bound per output element, shaped like the reference's
    result (a tuple of them for several outputs). Derive it from the
    kernel's arithmetic and say in ``note`` which parts were measured.
    """
    return cls(bound=bound, max_mismatch_frac=max_mismatch_frac, note=note)

lsb classmethod

lsb(
    n: int = 1,
    *,
    max_mismatch_frac: float = 0.0,
    note: str = ""
) -> "Tolerance"

Integer outputs within n least-significant bits of the reference.

For fixed-point kernels whose final saturating shift may round or truncate (the AIE srs rounding mode is a core setting the kernel does not fix). rtol is zero: the slack is absolute.

Source code in python/utils/verify.py

@classmethod
def lsb(
    cls, n: int = 1, *, max_mismatch_frac: float = 0.0, note: str = ""
) -> "Tolerance":
    """Integer outputs within ``n`` least-significant bits of the reference.

    For fixed-point kernels whose final saturating shift may round or
    truncate (the AIE ``srs`` rounding mode is a core setting the kernel
    does not fix). ``rtol`` is zero: the slack is absolute.
    """
    return cls(
        rtol=0.0, atol=n + 0.5, max_mismatch_frac=max_mismatch_frac, note=note
    )

default_for classmethod

default_for(dtype) -> 'Tolerance'

Return the contract a kernel gets when it declares none.

Integer and boolean outputs are bit-exact. bfloat16 outputs get the repository's canonical rtol=0.128, the C++ testbench default that the LUT-approximated kernels document; float32 outputs are held to rtol=1e-4 (a few float32 ULPs of accumulation-order slack, far inside what a bf16 tolerance would hide) and float16 to 1e-2.

Source code in python/utils/verify.py

@classmethod
def default_for(cls, dtype) -> "Tolerance":
    """Return the contract a kernel gets when it declares none.

    Integer and boolean outputs are bit-exact. bfloat16 outputs get the
    repository's canonical ``rtol=0.128``, the C++ testbench default that
    the LUT-approximated kernels document; float32 outputs are held to
    ``rtol=1e-4`` (a few float32 ULPs of accumulation-order slack, far
    inside what a bf16 tolerance would hide) and float16 to ``1e-2``.
    """
    dt = np.dtype(dtype)
    if np.issubdtype(dt, np.integer) or dt == np.bool_:
        return cls.exact(note="default: integer output")
    if dt == np.dtype(np.float32) or dt == np.dtype(np.float64):
        return cls.relative(1e-4, note="default: float32 output")
    if dt == np.dtype(np.float16):
        return cls.relative(1e-2, note="default: float16 output")
    return cls.relative(_DEFAULT_RTOL, note="default: canonical bf16/LUT rtol")

Verdict dataclass

Verdict(
    ok: bool,
    n_checked: int,
    n_mismatch: int,
    max_abs_err: float,
    max_ulp_err: int | None,
    first_bad_index: int | None,
    detail: str,
)

Outcome of compare. Truthy when the comparison passed.

compare

compare(
    actual,
    expected,
    tol: Tolerance | None = None,
    *,
    range_axis: int | None = None,
    bound=None
) -> Verdict

Compare a kernel's actual output with a reference under tol.

expected may be higher precision than actual (a float64 sum, an int64 product); it is cast to actual.dtype, so the kernel is held to what a correctly rounded implementation would produce. With tol=None the output dtype's Tolerance.default_for applies.

range_axis selects the axis reduced to compute range_frac's reference scale. For (calls, tile) arrays, use 1 to scale each call independently. The default uses the whole reference array.

bound is a bound tolerance's per-element limit, already evaluated on the inputs and broadcastable to expected.

This measures; it does not model. What a kernel does on overflow, on a narrowing store, or with subnormal inputs belongs in the reference that produced expected -- a saturating kernel's reference clips, a denormal-flushing kernel's reference flushes. A reference that leaves the output range is reported as such when the comparison fails, since it means the reference is under-specified rather than the kernel wrong.

Source code in python/utils/verify.py

def compare(
    actual,
    expected,
    tol: Tolerance | None = None,
    *,
    range_axis: int | None = None,
    bound=None,
) -> Verdict:
    """Compare a kernel's ``actual`` output with a reference under ``tol``.

    ``expected`` may be higher precision than ``actual`` (a float64 sum, an
    int64 product); it is cast to ``actual.dtype``, so the kernel is held to
    what a correctly rounded implementation would produce. With ``tol=None``
    the output dtype's ``Tolerance.default_for`` applies.

    ``range_axis`` selects the axis reduced to compute ``range_frac``'s
    reference scale. For ``(calls, tile)`` arrays, use 1 to scale each call
    independently. The default uses the whole reference array.

    ``bound`` is a **bound** tolerance's per-element limit, already
    evaluated on the inputs and broadcastable to ``expected``.

    This measures; it does not model. What a kernel does on overflow, on a
    narrowing store, or with subnormal inputs belongs in the reference that
    produced ``expected`` -- a saturating kernel's reference clips, a
    denormal-flushing kernel's reference flushes. A reference that leaves the
    output range is reported as such when the comparison fails, since it
    means the reference is under-specified rather than the kernel wrong.
    """
    actual = np.asarray(actual)
    expected = np.asarray(expected)
    if tol is None:
        tol = Tolerance.default_for(actual.dtype)
    if actual.shape != expected.shape:
        return Verdict(
            False,
            0,
            0,
            float("inf"),
            None,
            None,
            f"shape mismatch {actual.shape} vs {expected.shape}",
        )
    if tol.kind == "bound":
        if bound is None:
            raise ValueError(
                "a bound tolerance needs its bound evaluated on the inputs; "
                "pass bound=, or judge through the kernel with inputs="
            )
        if not np.issubdtype(actual.dtype, np.floating) and actual.dtype != bfloat16:
            raise ValueError(
                f"a bound tolerance is for floating-point outputs, got {actual.dtype}"
            )
    a, e, n = actual.ravel(), expected.ravel(), actual.size

    def ref_scale(ref) -> float | np.ndarray:
        """Return max|ref| over finite entries.

        This is what ``range_frac`` is a fraction of, or zero when no range
        floor is in play.
        """
        if tol.range_frac is None or not n:
            return 0.0
        mag = np.abs(np.asarray(ref, dtype=np.float64).reshape(expected.shape))
        mag = np.where(np.isfinite(mag), mag, 0.0)
        scale = mag.max(axis=range_axis, keepdims=True, initial=0.0)
        if range_axis is None:
            return float(scale.item())
        return np.broadcast_to(scale, expected.shape).ravel()

    # Integers: bit-exact under exact / ulps; under a relative tolerance the
    # nearly_equal formula in int64 (Tolerance.lsb sets atol = n + 0.5).
    if np.issubdtype(actual.dtype, np.integer) or actual.dtype == np.bool_:
        n_over = 0
        if actual.dtype != np.bool_ and np.issubdtype(e.dtype, np.integer):
            info = np.iinfo(actual.dtype)
            e_wide = e.astype(np.int64)
            n_over = int(np.count_nonzero((e_wide < info.min) | (e_wide > info.max)))
        e_cast = e.astype(actual.dtype)
        a64, e64 = a.astype(np.int64), e_cast.astype(np.int64)
        err = np.abs(a64 - e64)
        scale = ref_scale(e64)
        if tol.kind == "relative":
            bound = np.maximum(
                tol.atol or 0.0,
                (tol.rtol or 0.0) * (np.abs(a64) + np.abs(e64)),
            )
            close = err < bound
            if tol.range_frac is not None:
                close |= err <= tol.range_frac * scale
            bad = ~close
        else:
            bad = a != e_cast
        v = _verdict(bad, err, None, tol, n, ref_range=scale)
        if n_over and not v.ok:
            # Not a policy, a diagnostic: the reference left the output range,
            # so what the device did there says nothing about the kernel.
            v = replace(
                v,
                detail=f"{v.detail}; the reference overflows "
                f"{np.dtype(actual.dtype).name} in {n_over} of {n} elements, so "
                "it does not model what the kernel does there (clip for a "
                "saturating kernel, cast for a wrapping one)",
            )
        return v

    a32, e32 = a.astype(np.float32), e.astype(np.float32)
    a_nan, e_nan = np.isnan(a32), np.isnan(e32)
    a_inf, e_inf = np.isinf(a32), np.isinf(e32)
    nonfinite_bad = (a_nan != e_nan) | (a_inf != e_inf) | (a_inf & e_inf & (a32 != e32))
    finite = ~(a_nan | e_nan | a_inf | e_inf)
    err = np.zeros(n, np.float64)

    if tol.kind == "ulps":
        if actual.dtype != bfloat16:
            raise ValueError(
                f"Tolerance in ULPs is defined for bfloat16 outputs, got {actual.dtype}"
            )
        # A cast to bfloat16 goes through float32 and can round twice;
        # round_to rounds the reference once, as a correctly rounded bf16
        # implementation would.
        e_bf = round_to(e, bfloat16)
        ulp = np.zeros(n, np.int64)
        ulp[finite] = bf16_ulp_distance(a[finite], e_bf[finite])
        err[finite] = np.abs(a32[finite] - e_bf[finite].astype(np.float32))
        max_ulps = tol.ulps if tol.ulps is not None else 0
        within = ulp <= max_ulps
        scale = ref_scale(e_bf.astype(np.float32))
        if tol.atol is not None:
            within |= err < tol.atol
        if tol.range_frac is not None:
            within |= err <= tol.range_frac * scale
        bad = nonfinite_bad | (finite & ~within)
        return _verdict(bad, err, ulp, tol, n, nonfinite_bad, ref_range=scale)

    if tol.kind == "bound":
        limit = np.broadcast_to(np.asarray(bound, np.float64), expected.shape).ravel()
        err[finite] = np.abs(a32[finite].astype(np.float64) - e32[finite])
        bad = nonfinite_bad | (finite & ~(err <= limit))
        return _verdict(bad, err, None, tol, n, nonfinite_bad, limit=limit)

    if tol.kind == "exact":
        e_cast = e32.astype(actual.dtype).astype(np.float32)
        err[finite] = np.abs(a32[finite] - e_cast[finite])
        bad = nonfinite_bad | (finite & (a32 != e_cast))
        return _verdict(bad, err, None, tol, n, nonfinite_bad)

    err[finite] = np.abs(a32[finite].astype(np.float64) - e32[finite])
    scale = ref_scale(e32)
    close = nearly_equal(a32, e32, rtol=tol.rtol or 0.0, atol=tol.atol)
    if tol.range_frac is not None:
        close |= err <= tol.range_frac * scale
    bad = nonfinite_bad | (finite & ~close)
    return _verdict(bad, err, None, tol, n, nonfinite_bad, ref_range=scale)

bf16_ulp_distance

bf16_ulp_distance(a, b) -> ndarray

Element-wise distance between two bf16 arrays in units in the last place.

Bit patterns are mapped to a monotonic integer scale (sign-magnitude to two's-complement style) so the distance is a plain subtraction. -0 and +0 map to the same point, so a kernel that produces the other zero is not penalized. aie.utils.accuracy.ulp_distance takes other dtypes.

Source code in python/utils/verify.py

def bf16_ulp_distance(a, b) -> np.ndarray:
    """Element-wise distance between two bf16 arrays in units in the last place.

    Bit patterns are mapped to a monotonic integer scale (sign-magnitude to
    two's-complement style) so the distance is a plain subtraction. -0 and +0
    map to the same point, so a kernel that produces the other zero is not
    penalized. ``aie.utils.accuracy.ulp_distance`` takes other dtypes.
    """
    return ulp_distance(a, b, bfloat16)

Sampling

sample_select and sample_combine draw the next token from a row of bf16 logits on the device: each sample_select core reduces one slice to its top k, and one sample_combine core draws from those summaries. sample_ref is the bit-exact host reference.

Top-k next-token sampling on the device, and its bit-exact host reference.

A row of bf16 logits is split into columns of slice_size. Each column's sample_select core reduces its slice to a summary (its top k, as entries above its k-th largest value and a bitmap of the ties at it); one sample_combine core draws the token from all the summaries. The draw is sample_ref:

1. Temperature 0 (or -0): the first index of the largest logit.
2. tau is the k-th largest logit, with multiplicity, compared on order_keys (-0 and +0 are one key). The candidates are the logits >= tau, in index order.
3. Each candidate value v weighs exp64_ref(float64(fl32(fl32(v / T) - fl32(max / T)))).
4. u = n53 * 2**-53, integer 0 <= n53 < 2**53. The token is the first candidate, in index order, whose exact prefix sum of the weights P exceeds u * S, S the exact total: P * 2**53 > n53 * S in integers. That is the inverse-CDF draw over the weights, with no rounding after them.

Every weight step is an IEEE float32 division or subtraction, or a float64 addition, subtraction, multiplication or comparison, never fused; the core computes the float32 steps with the soft-float builtins and float64 has only those, so numpy reproduces the device's bits. Each weight is an integer multiple of 2**-1074, so the core sums them exactly in fixed point. A temperature must satisfy check_order_preserving, which makes the threshold on bf16 keys select what a threshold on v / T would.

A draw is passed to the device as a four-word int32 row, draw_row.

summary_words

summary_words(slice_size: int, k_max: int) -> int

int32 words of one column's summary: header, entries, tie bitmap.

Source code in python/iron/kernels/sample.py

def summary_words(slice_size: int, k_max: int) -> int:
    """int32 words of one column's summary: header, entries, tie bitmap."""
    return SUMMARY_HEADER + 2 * k_max + (slice_size + 31) // 32

select_streams

select_streams(slice_size: int, chunk: int) -> int

How often sample_select takes its slice per position.

Once when a chunk is the whole slice, which one call passes over SELECT_PASSES times; SELECT_PASSES times otherwise.

Source code in python/iron/kernels/sample.py

def select_streams(slice_size: int, chunk: int) -> int:
    """How often ``sample_select`` takes its slice per position.

    Once when a chunk is the whole slice, which one call passes over
    ``SELECT_PASSES`` times; ``SELECT_PASSES`` times otherwise.
    """
    return 1 if chunk == slice_size else SELECT_PASSES

sample_select

sample_select(
    *, slice_size=32064, chunk=5344, k_max=64
) -> ExternalFunction

One column's half of sampling: its slice's summary, for sample_combine.

sample_select(x, row, state, summary) takes chunk bf16 logits per call; the slice arrives select_streams(slice_size, chunk) times, so a position is that many times slice_size // chunk calls with the same row and summary. state is SELECT_STATE_WORDS int32 of worker-local memory, zero before the first call; the last call of a position leaves it ready for the next. summary is summary_words(slice_size, k_max) int32.

Source code in python/iron/kernels/sample.py

def sample_select(*, slice_size=32064, chunk=5344, k_max=64) -> ExternalFunction:
    """One column's half of sampling: its slice's summary, for ``sample_combine``.

    ``sample_select(x, row, state, summary)`` takes ``chunk`` bf16 logits per
    call; the slice arrives ``select_streams(slice_size, chunk)`` times, so a
    position is that many times ``slice_size // chunk`` calls with the same
    ``row`` and ``summary``.
    ``state`` is ``SELECT_STATE_WORDS`` int32 of worker-local memory, zero
    before the first call; the last call of a position leaves it ready for the
    next. ``summary`` is ``summary_words(slice_size, k_max)`` int32.
    """
    _check_slice("sample_select", slice_size, k_max)
    _positive("sample_select", chunk=chunk)
    if slice_size % chunk:
        raise ValueError("sample_select: chunk must divide slice_size")
    if chunk % 2:
        raise ValueError("sample_select: chunk must be even (4-byte DMA)")
    return _make_extern(
        "sample_select",
        _kernel_source("sample/sample_select.cc"),
        [
            np.ndarray[(chunk,), np.dtype[bfloat16]],
            np.ndarray[(ROW_WORDS,), np.dtype[np.int32]],
            np.ndarray[(SELECT_STATE_WORDS,), np.dtype[np.int32]],
            np.ndarray[(summary_words(slice_size, k_max),), np.dtype[np.int32]],
        ],
        compile_flags=_STRICT_FP
        + [
            f"-DSAMPLE_SLICE={slice_size}",
            f"-DSAMPLE_CHUNK={chunk}",
            f"-DSAMPLE_K_MAX={k_max}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, InOut, InOut),
        ),
    )

sample_combine

sample_combine(
    *, columns=4, slice_size=32064, k_max=64
) -> ExternalFunction

Draw from columns summaries of sample_select, as sample_ref does.

sample_combine(summaries, row, token, record) reads the summaries column 0 first, and writes the drawn token (an index into the whole row of columns * slice_size) to both one-word outputs, so a design can send it two places.

Source code in python/iron/kernels/sample.py

def sample_combine(*, columns=4, slice_size=32064, k_max=64) -> ExternalFunction:
    """Draw from ``columns`` summaries of ``sample_select``, as ``sample_ref`` does.

    ``sample_combine(summaries, row, token, record)`` reads the summaries
    column 0 first, and writes the drawn token (an index into the whole row
    of ``columns * slice_size``) to both one-word outputs, so a design can
    send it two places.
    """
    _check_slice("sample_combine", slice_size, k_max)
    _positive("sample_combine", columns=columns)
    if columns * slice_size >= 1 << 31:
        raise ValueError("sample_combine: a token must fit int32")
    return _make_extern(
        "sample_combine",
        _kernel_source("sample/sample_combine.cc"),
        [
            np.ndarray[
                (columns * summary_words(slice_size, k_max),), np.dtype[np.int32]
            ],
            np.ndarray[(ROW_WORDS,), np.dtype[np.int32]],
            np.ndarray[(1,), np.dtype[np.int32]],
            np.ndarray[(1,), np.dtype[np.int32]],
        ],
        compile_flags=_STRICT_FP
        + [
            f"-DSAMPLE_SLICE={slice_size}",
            f"-DSAMPLE_COLUMNS={columns}",
            f"-DSAMPLE_K_MAX={k_max}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out, Out),
            # The weight table (12 bytes a k_max entry) and the 256-bin
            # histogram below it: 2240 bytes at k_max 64, 3008 at
            # K_MAX_LIMIT. The soft-float builtins carry no stack sizes for
            # aiecc to count; the deepest call that reaches them, 320 bytes
            # with their 64, stays above the histogram's 1024.
            stack_bytes=4096,
        ),
    )

exp64_ref

exp64_ref(x)

Compute exp over float64, bit for bit as aie_kernels/sample/exp64.h.

numpy's elementwise float64 +, -, * and comparisons are correctly rounded, so evaluating exp64.h's sequence of them gives its bits; np.exp is never used. Every branch is evaluated for every element and the results selected with np.where, which is exact.

Source code in python/iron/kernels/sample.py

def exp64_ref(x):
    """Compute exp over float64, bit for bit as ``aie_kernels/sample/exp64.h``.

    numpy's elementwise float64 +, -, * and comparisons are correctly
    rounded, so evaluating exp64.h's sequence of them gives its bits; np.exp
    is never used. Every branch is evaluated for every element and the
    results selected with np.where, which is exact.
    """
    c, table = _exp64_constants()
    n = np.uint64(table.size // 2)
    table_bits = int(n).bit_length() - 1
    u64, f64 = np.uint64, np.float64
    x = np.ascontiguousarray(x, dtype=np.float64)
    shape = x.shape
    x = x.reshape(-1)
    bits = x.view(np.uint64)
    abstop = (bits >> u64(52)).astype(np.int64) & 0x7FF
    tiny = abstop < 0x3C9  # |x| < 2^-54
    huge = abstop >= 0x409  # |x| >= 1024, inf, nan
    special = abstop == 0x408  # 512 <= |x| < 1024
    negative = (bits >> u64(63)) == u64(1)

    with np.errstate(all="ignore"):
        z = c["INV_LN2_N"] * x
        kd = z + c["SHIFT"]
        ki = kd.view(np.uint64).copy()
        kd = kd - c["SHIFT"]
        khi = kd * c["NEG_LN2_HI_N"]
        klo = kd * c["NEG_LN2_LO_N"]
        r = x + khi
        r = r + klo
        idx = (ki % n) * u64(2)
        top = ki << u64(52 - table_bits)
        tail = table[idx].view(np.float64)
        sbits = table[idx + u64(1)] + top
        r2 = r * r
        p23 = r * c["C3"]
        p23 = c["C2"] + p23
        p23 = r2 * p23
        p45 = r * c["C5"]
        p45 = c["C4"] + p45
        r4 = r2 * r2
        p45 = r4 * p45
        tmp = tail + r
        tmp = tmp + p23
        tmp = tmp + p45

        scale = sbits.view(np.float64)
        st = scale * tmp
        main = scale + st

        # exp64_special, k > 0.
        scale_pos = (sbits - (u64(1009) << u64(52))).view(np.float64)
        st_pos = scale_pos * tmp
        y_pos = scale_pos + st_pos
        y_pos = f64(2.0**1009) * y_pos
        # exp64_special, k < 0.
        scale_neg = (sbits + (u64(1022) << u64(52))).view(np.float64)
        st_neg = scale_neg * tmp
        y = scale_neg + st_neg
        lo = scale_neg - y
        lo = lo + st_neg
        hi = f64(1.0) + y
        lo2 = f64(1.0) - hi
        lo2 = lo2 + y
        lo2 = lo2 + lo
        y_sub = hi + lo2
        y_sub = y_sub - f64(1.0)
        y_sub = np.where(y_sub == f64(0.0), f64(0.0), y_sub)
        y_neg = np.where(y < f64(1.0), y_sub, y)
        y_neg = f64(2.0**-1022) * y_neg
        k_positive = (ki & u64(0x80000000)) == u64(0)
        special_value = np.where(k_positive, y_pos, y_neg)

        one_plus_x = f64(1.0) + x
        huge_value = np.where(
            bits == np.float64(-np.inf).view(np.uint64),
            f64(0.0),
            np.where(
                abstop >= 0x7FF,
                one_plus_x,
                np.where(negative, f64(0.0), f64(np.inf)),
            ),
        )

    out = np.where(special, special_value, main)
    out = np.where(huge, huge_value, out)
    out = np.where(tiny, one_plus_x, out)
    return out.reshape(shape)

order_keys

order_keys(logits) -> ndarray

uint16 keys whose order is the bf16 logits' numeric order; -0 and +0 are one key.

Source code in python/iron/kernels/sample.py

def order_keys(logits) -> np.ndarray:
    """uint16 keys whose order is the bf16 logits' numeric order; -0 and +0 are one key."""
    bits = np.ascontiguousarray(logits, dtype=bfloat16).view(np.uint16)
    bits = np.where(bits == 0x8000, np.uint16(0), bits)
    negative = (bits & np.uint16(0x8000)) != 0
    return np.where(negative, ~bits, bits | np.uint16(0x8000)).astype(np.uint16)

check_order_preserving

check_order_preserving(temperature) -> None

Raise unless l -> fl32(l / T) strictly increases over every finite bf16 l with a finite image.

Then a threshold on bf16 keys selects what a threshold on l / T would, up to logits whose l / T is -inf and whose weight is 0 anyway.

Source code in python/iron/kernels/sample.py

def check_order_preserving(temperature) -> None:
    """Raise unless ``l -> fl32(l / T)`` strictly increases over every finite bf16 l with a finite image.

    Then a threshold on bf16 keys selects what a threshold on ``l / T`` would,
    up to logits whose ``l / T`` is -inf and whose weight is 0 anyway.
    """
    temperature = np.float32(temperature)
    if not 0 < temperature < np.inf:
        raise ValueError(f"temperature {temperature} is not finite and positive")
    bits = np.arange(1 << 16, dtype=np.uint32).astype(np.uint16)
    values = bits.view(bfloat16).astype(np.float32)
    values = np.sort(values[np.isfinite(values) & (bits != 0x8000)])
    with np.errstate(over="ignore"):
        x = values / temperature
    x = x[np.isfinite(x)]
    if not np.all(x[1:] > x[:-1]):
        where = int(np.flatnonzero(x[1:] <= x[:-1])[0])
        raise ValueError(
            f"temperature {temperature} maps two bf16 logits to one x "
            f"({x[where]!r}); the bf16-key threshold would not be exact"
        )

draw_row

draw_row(temperature, top_k: int, n53: int) -> ndarray

Pack the four int32 words sample_select and sample_combine read for one draw.

The kernels clamp top_k to their k_max.

Source code in python/iron/kernels/sample.py

def draw_row(temperature, top_k: int, n53: int) -> np.ndarray:
    """Pack the four int32 words ``sample_select`` and ``sample_combine`` read for one draw.

    The kernels clamp ``top_k`` to their ``k_max``.
    """
    _positive("draw_row", top_k=top_k)
    if top_k >= 1 << 31:
        raise ValueError(f"draw_row: top_k {top_k} does not fit int32")
    n53 = _check_n53(n53)
    row = np.empty(ROW_WORDS, dtype=np.uint32)
    row[0] = np.float32(temperature).view(np.uint32)
    row[1] = top_k
    row[2] = n53 & 0xFFFFFFFF
    row[3] = n53 >> 32
    return row.view(np.int32)

sample_weights

sample_weights(
    logits, temperature, top_k: int
) -> tuple[ndarray, ndarray]

(candidates, weights) of a draw at a positive temperature.

The candidates are the indices of every logit at or above the top_k-th largest (ties with it included), in index order; each weight is exp64(fl32(v / T) - fl32(max / T)), unnormalised.

Source code in python/iron/kernels/sample.py

def sample_weights(logits, temperature, top_k: int) -> tuple[np.ndarray, np.ndarray]:
    """``(candidates, weights)`` of a draw at a positive temperature.

    The candidates are the indices of every logit at or above the
    ``top_k``-th largest (ties with it included), in index order; each weight
    is ``exp64(fl32(v / T) - fl32(max / T))``, unnormalised.
    """
    _positive("sample_weights", top_k=top_k)
    row = _row(logits)
    keys = order_keys(row)
    values = row.astype(np.float32)
    temperature = np.float32(temperature)
    n = keys.size
    k = min(top_k, n)
    tau = np.partition(keys, n - k)[n - k]
    candidates = np.flatnonzero(keys >= tau)
    with np.errstate(over="ignore"):
        xm = values[int(np.argmax(keys))] / temperature
        xv = values[candidates] / temperature
    if not np.isfinite(xm):
        raise ValueError("the largest logit / T is not finite in float32")
    return candidates, exp64_ref((xv - xm).astype(np.float64))

sample_ref

sample_ref(
    logits,
    temperature,
    top_k: int,
    n53: int,
    *,
    k_max: int | None = None
) -> int

Return the token sample_combine draws from one row of bf16 logits.

With k_max, top_k is clamped to it, as kernels built with that k_max clamp it; without, top_k is taken as it is.

Source code in python/iron/kernels/sample.py

def sample_ref(
    logits, temperature, top_k: int, n53: int, *, k_max: int | None = None
) -> int:
    """Return the token ``sample_combine`` draws from one row of bf16 logits.

    With ``k_max``, ``top_k`` is clamped to it, as kernels built with that
    ``k_max`` clamp it; without, ``top_k`` is taken as it is.
    """
    _positive("sample_ref", top_k=top_k)
    n53 = _check_n53(n53)
    if k_max is not None:
        _positive("sample_ref", k_max=k_max)
        top_k = min(top_k, k_max)
    if np.float32(temperature).view(np.uint32) & 0x7FFFFFFF == 0:
        return int(np.argmax(order_keys(_row(logits))))
    candidates, weights = sample_weights(logits, temperature, top_k)
    prefix = list(accumulate(_units(weights)))
    # P * 2**53 > n53 * S is P > floor(n53 * S / 2**53) for an integer P. The
    # prefixes never step down, so the first above is a bisection; S >= 1
    # (the maximum weighs 1) and u < 1, so the last always is.
    target = n53 * prefix[-1] >> 53
    return int(candidates[bisect_right(prefix, target)])

FastFlowLM Gemma 4

Kernels extracted from FastFlowLM's Gemma 4 implementation.

These are not general building blocks. Each source is the kernel code of one core in one of FastFlowLM's Gemma 4 designs (prefill attention, the LM head and the decode layer), called from that core's Worker body. Several kernels acquire and release core locks themselves, by the lock ids the factories take. The decode kernels build for the whole model's geometry (FlmGemma4DecodeGeometry) and for FastFlowLM's array layout. All are AIE2P only; their sources are in aie_kernels/flm_gemma4/.

A factory returns one ExternalFunction. As cascade_mm does with its put_only and put_get, every entry point of its object, its own included, is an attribute of it named by its C symbol: fn.attn_qk_round.

FlmGemma4DecodeGeometry dataclass

FlmGemma4DecodeGeometry(
    model_dim: int,
    num_attn_heads: int,
    num_kv_heads: int,
    intermediate_size: int,
    glu_slice: int,
    pli_d: int,
    dh: int,
    swa_dh: int,
    attn_scale: float,
    swa_attn_scale: float,
    pli_projection_scale: float,
    pli_input_scale: float,
    gelu: bool = True,
    qk_norm: bool = True,
    double_wide_mlp: bool = False,
    name: str = "custom",
)

The model geometry the flm_gemma4/decode_*.cc kernels build for.

The kernels are Gemma 4's fused decode layer, one core per stage, as FastFlowLM builds it. Every field reaches the kernels as a -DFLM_GEMMA4_DECODE_* flag. FLM_GEMMA4_E2B_DECODE and FLM_GEMMA4_E4B_DECODE are the two variants' values, and the factories build only those: the kernels keep assumptions other values break, such as eight query heads, the GELU table, the q/k norm and projections that divide into whole rounds.

flags

flags() -> list[str]

Return the -DFLM_GEMMA4_DECODE_* flags, floats as shortest round-trip literals.

Source code in python/iron/kernels/flm_gemma4.py

def flags(self) -> list[str]:
    """Return the ``-DFLM_GEMMA4_DECODE_*`` flags, floats as shortest round-trip literals."""
    values = {
        "MODEL_DIM": self.model_dim,
        "NUM_ATTN_HEADS": self.num_attn_heads,
        "NUM_KV_HEADS": self.num_kv_heads,
        "INTERMEDIATE_SIZE": self.intermediate_size,
        "GLU_SLICE": self.glu_slice,
        "PLI_D": self.pli_d,
        "DH": self.dh,
        "SWA_DH": self.swa_dh,
        "ATTN_SCALE": f"{float(self.attn_scale)!r}f",
        "SWA_ATTN_SCALE": f"{float(self.swa_attn_scale)!r}f",
        "PLI_PROJECTION_SCALE": f"{float(self.pli_projection_scale)!r}f",
        "PLI_INPUT_SCALE": f"{float(self.pli_input_scale)!r}f",
        "GELU": int(self.gelu),
        "QK_NORM": int(self.qk_norm),
        "DOUBLE_WIDE_MLP": int(self.double_wide_mlp),
    }
    return [f"-DFLM_GEMMA4_DECODE_{k}={v}" for k, v in values.items()]

flm_gemma4_attn_prefill

flm_gemma4_attn_prefill(
    *, in_prod_lock: int = 2, in_cons_lock: int = 3
) -> ExternalFunction

Causal flash-attention prefill at a head dim of 512, from flm_gemma4/prefill.cc.

linalg.prefill_fv's algorithm with FastFlowLM's numerics and synchronization: the softmax scales by log2(e) alone, the epilogue multiplies by a bf16 1 / l, and k and v arrive in a ping-pong pair that the core's own locks guard.

One core holds 8 query rows of a 128-row round and folds the keys in one 8-row step at a time, with one entry point per step of the caller's round, block and step loops. The returned kernel is attn_epilogue, which writes one 64-element chunk of o = y / l; its chunk index is bound to 0. The entry points, all attributes of the kernel:

• attn_rounds(L_begin, L_end, n_out): rounds in the dispatch
• attn_blocks(L_begin, i, n_out): key blocks in round i
• attn_round_begin(prev_m, new_m, c, l, y)
• attn_block_begin(m, prev_m)
• attn_qk_step(s, q, k_ping, k_pong, m, L_begin, window_size, row, col, i, block, j)
• attn_block_mid(s, m, new_m, prev_m, c, l, y)
• attn_fv_step(y, s, v_ping, v_pong, j)
• attn_block_end(prev_m, new_m)
• attn_finalize(l, l_bf16)
• attn_epilogue(o, l_bf16, y, chunk)

The design must declare the k/v locks and fill the pair with a tile DMA.

Parameters:

Name	Type	Description	Default
in_prod_lock	int	Core lock the k/v DMA acquires to fill a buffer.	2
in_cons_lock	int	Core lock the steps acquire to read one.	3

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_prefill(
    *, in_prod_lock: int = 2, in_cons_lock: int = 3
) -> ExternalFunction:
    """Causal flash-attention prefill at a head dim of 512, from ``flm_gemma4/prefill.cc``.

    [`linalg.prefill_fv`][iron.kernels.linalg.prefill_fv]'s algorithm with
    FastFlowLM's numerics and synchronization: the softmax scales by log2(e)
    alone, the epilogue multiplies by a bf16 ``1 / l``, and k and v arrive in
    a ping-pong pair that the core's own locks guard.

    One core holds 8 query rows of a 128-row round and folds the keys in one
    8-row step at a time, with one entry point per step of the caller's
    round, block and step loops. The returned kernel is ``attn_epilogue``,
    which writes one 64-element chunk of ``o = y / l``; its chunk index is
    bound to 0. The entry points, all attributes of the kernel:

    - ``attn_rounds(L_begin, L_end, n_out)``: rounds in the dispatch
    - ``attn_blocks(L_begin, i, n_out)``: key blocks in round ``i``
    - ``attn_round_begin(prev_m, new_m, c, l, y)``
    - ``attn_block_begin(m, prev_m)``
    - ``attn_qk_step(s, q, k_ping, k_pong, m, L_begin, window_size, row, col, i, block, j)``
    - ``attn_block_mid(s, m, new_m, prev_m, c, l, y)``
    - ``attn_fv_step(y, s, v_ping, v_pong, j)``
    - ``attn_block_end(prev_m, new_m)``
    - ``attn_finalize(l, l_bf16)``
    - ``attn_epilogue(o, l_bf16, y, chunk)``

    The design must declare the k/v locks and fill the pair with a tile DMA.

    Args:
        in_prod_lock: Core lock the k/v DMA acquires to fill a buffer.
        in_cons_lock: Core lock the steps acquire to read one.
    """
    return _prefill(
        "flm_gemma4_attn_prefill",
        _ATTN_PREFILL_DH,
        _ATTN_PREFILL_LQ,
        0,
        flm_gemma4_attn_prefill_ref,
        in_prod_lock,
        in_cons_lock,
    )

flm_gemma4_swa_prefill

flm_gemma4_swa_prefill(
    *, in_prod_lock: int = 2, in_cons_lock: int = 3
) -> ExternalFunction

Sliding-window causal flash-attention prefill at a head dim of 256, from flm_gemma4/prefill.cc.

The sliding-window build of flm_gemma4_attn_prefill: one core holds 16 query rows of a 128-row round and folds in, 16 rows at a time, the keys inside the window its window_size RTP gives. The entry points are those of flm_gemma4_attn_prefill, with attn_blocks(L_begin, window_size, i, n_out) also reading the window. The returned kernel is attn_epilogue, bound to chunk 33: rows 8 to 15, columns 64 to 127 of the (16, 256) y.

Parameters:

Name	Type	Description	Default
in_prod_lock	int	Core lock the k/v DMA acquires to fill a buffer.	2
in_cons_lock	int	Core lock the steps acquire to read one.	3

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_swa_prefill(
    *, in_prod_lock: int = 2, in_cons_lock: int = 3
) -> ExternalFunction:
    """Sliding-window causal flash-attention prefill at a head dim of 256, from ``flm_gemma4/prefill.cc``.

    The sliding-window build of [`flm_gemma4_attn_prefill`][iron.kernels.flm_gemma4.flm_gemma4_attn_prefill]:
    one core holds 16 query rows of a 128-row round and folds in, 16 rows at
    a time, the keys inside the window its ``window_size`` RTP gives. The
    entry points are those of ``flm_gemma4_attn_prefill``, with
    ``attn_blocks(L_begin, window_size, i, n_out)`` also reading the window.
    The returned kernel is ``attn_epilogue``, bound to chunk 33: rows 8 to 15,
    columns 64 to 127 of the (16, 256) ``y``.

    Args:
        in_prod_lock: Core lock the k/v DMA acquires to fill a buffer.
        in_cons_lock: Core lock the steps acquire to read one.
    """
    return _prefill(
        "flm_gemma4_swa_prefill",
        _SWA_PREFILL_DH,
        _SWA_PREFILL_LQ,
        _SWA_PREFILL_CHUNK,
        flm_gemma4_swa_prefill_ref,
        in_prod_lock,
        in_cons_lock,
    )

flm_gemma4_prefill_block_begin

flm_gemma4_prefill_block_begin(
    *, head_dim: int = 512
) -> ExternalFunction

attn_block_begin of the prefill kernels: row j of m is prev_m[j].

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_block_begin(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_block_begin`` of the prefill kernels: row ``j`` of ``m`` is ``prev_m[j]``."""
    return _prefill_sibling(
        "attn_block_begin",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, In),
            reference=flm_gemma4_prefill_block_begin_ref,
            tolerance=Tolerance.exact(note="bf16 broadcast"),
            ops_per_call=0,
        ),
    )

flm_gemma4_prefill_block_begin_ref

flm_gemma4_prefill_block_begin_ref(prev_m)

Numpy reference for flm_gemma4_prefill_block_begin.

m is (LQ, LK) with LK == LQ in both builds.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_block_begin_ref(prev_m):
    """Numpy reference for [`flm_gemma4_prefill_block_begin`][iron.kernels.flm_gemma4.flm_gemma4_prefill_block_begin].

    ``m`` is ``(LQ, LK)`` with ``LK == LQ`` in both builds.
    """
    prev_m = np.asarray(prev_m)
    return np.repeat(prev_m, prev_m.shape[-1], axis=-1)

flm_gemma4_prefill_block_end

flm_gemma4_prefill_block_end(
    *, head_dim: int = 512
) -> ExternalFunction

attn_block_end of the prefill kernels: prev_m = new_m.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_block_end(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_block_end`` of the prefill kernels: ``prev_m = new_m``."""
    return _prefill_sibling(
        "attn_block_end",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, In),
            reference=flm_gemma4_prefill_block_end_ref,
            tolerance=Tolerance.exact(note="bf16 copy"),
            ops_per_call=0,
        ),
    )

flm_gemma4_prefill_block_end_ref

flm_gemma4_prefill_block_end_ref(new_m)

Numpy reference for flm_gemma4_prefill_block_end.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_block_end_ref(new_m):
    """Numpy reference for [`flm_gemma4_prefill_block_end`][iron.kernels.flm_gemma4.flm_gemma4_prefill_block_end]."""
    return np.array(new_m, copy=True)

flm_gemma4_prefill_rounds_ref

flm_gemma4_prefill_rounds_ref(l_begin, l_end)

(L_end >> 7) - (L_begin >> 7): the 128-row rounds a dispatch spans.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_rounds_ref(l_begin, l_end):
    """``(L_end >> 7) - (L_begin >> 7)``: the 128-row rounds a dispatch spans."""
    n = (_word0(l_end) >> 7) - (_word0(l_begin) >> 7)
    return n.reshape(-1, 1)

flm_gemma4_prefill_rounds

flm_gemma4_prefill_rounds(
    *, head_dim: int = 512
) -> ExternalFunction

attn_rounds of the prefill kernels: the rounds between L_begin and L_end.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_rounds(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_rounds`` of the prefill kernels: the rounds between ``L_begin`` and ``L_end``."""
    return _prefill_sibling(
        "attn_rounds",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            sample=_prefill_rounds_sample,
            reference=flm_gemma4_prefill_rounds_ref,
            tolerance=Tolerance.exact(note="integer shifts and a subtraction"),
            out_valid=1,
            ops_per_call=1,
        ),
    )

flm_gemma4_prefill_blocks_ref

flm_gemma4_prefill_blocks_ref(
    l_begin, i, *, window_size=None
)

Key blocks round i folds in: all up to its own, or those in the window.

window_size is the sliding-window build's RTP buffer, None for the head-dim-512 build.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_blocks_ref(l_begin, i, *, window_size=None):
    """Key blocks round ``i`` folds in: all up to its own, or those in the window.

    ``window_size`` is the sliding-window build's RTP buffer, ``None`` for the
    head-dim-512 build.
    """
    begin = _word0(l_begin)
    if window_size is None:
        n = (begin >> 7) + np.int32(i) + 1
    else:
        q = begin + np.int32(i) * 128
        k = np.maximum(q - _word0(window_size), 0)
        n = ((q - k) >> 7) + 1
    return n.astype(np.int32).reshape(-1, 1)

flm_gemma4_prefill_blocks

flm_gemma4_prefill_blocks(
    *, head_dim: int = 512
) -> ExternalFunction

attn_blocks of the prefill kernels: the key blocks round i folds in.

At a head dim of 256, the sliding-window build, it also reads the window size: attn_blocks(L_begin, window_size, i, n_out).

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_blocks(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_blocks`` of the prefill kernels: the key blocks round ``i`` folds in.

    At a head dim of 256, the sliding-window build, it also reads the window
    size: ``attn_blocks(L_begin, window_size, i, n_out)``.
    """
    window = head_dim == 256
    return _prefill_sibling(
        "attn_blocks",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Param, Out) if window else (In, Param, Out),
            sample=partial(_prefill_blocks_sample, window=window),
            reference=(
                _prefill_swa_blocks_ref if window else flm_gemma4_prefill_blocks_ref
            ),
            tolerance=Tolerance.exact(note="integer shifts, adds and a clamp"),
            out_valid=1,
            ops_per_call=1,
        ),
    )

flm_gemma4_prefill_finalize

flm_gemma4_prefill_finalize(
    *, head_dim: int = 512
) -> ExternalFunction

attn_finalize of the prefill kernels: inv_l = bf16(1 / l) per query row.

Parameters:

Name	Type	Description	Default
head_dim	int	512 builds the global kernel's 8 rows; 256 builds the sliding-window kernel's 16 rows.	512

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_finalize(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_finalize`` of the prefill kernels: ``inv_l = bf16(1 / l)`` per query row.

    Args:
        head_dim: 512 builds the global kernel's 8 rows; 256 builds the
            sliding-window kernel's 16 rows.
    """
    return _prefill_sibling(
        "attn_finalize",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_prefill_finalize_ref,
            sample=partial(
                _prefill_finalize_sample,
                lq=_ATTN_PREFILL_LQ if head_dim == 512 else _SWA_PREFILL_LQ,
            ),
            # Floor narrowing lands at most one bf16 step from the nearest-even
            # cast of 1 / l if aie::inv errs by less than half a step (2**-9
            # relative).
            tolerance=Tolerance.bf16_ulps(
                1,
                note="floor bf16 narrowing of AIE2P's scalar inv instruction, "
                "whose error the AIE-ML v2 intrinsics guide does not document",
            ),
        ),
    )

flm_gemma4_prefill_finalize_ref

flm_gemma4_prefill_finalize_ref(row_sums)

Numpy reference for flm_gemma4_prefill_finalize: 1 / l.

The tolerance covers the kernel's hardware reciprocal and its floor narrowing to bf16.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_finalize_ref(row_sums):
    """Numpy reference for [`flm_gemma4_prefill_finalize`][iron.kernels.flm_gemma4.flm_gemma4_prefill_finalize]: ``1 / l``.

    The tolerance covers the kernel's hardware reciprocal and its floor
    narrowing to bf16.
    """
    return 1.0 / np.asarray(row_sums, dtype=np.float32)

flm_gemma4_prefill_round_begin_ref

flm_gemma4_prefill_round_begin_ref(*, head_dim=512)

Numpy reference for flm_gemma4_prefill_round_begin.

The bytes of y | c | l | prev_m | new_m: zeros, ones, zeros and two rows of NEG_INF.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_round_begin_ref(*, head_dim=512):
    """Numpy reference for [`flm_gemma4_prefill_round_begin`][iron.kernels.flm_gemma4.flm_gemma4_prefill_round_begin].

    The bytes of ``y | c | l | prev_m | new_m``: zeros, ones, zeros and two
    rows of NEG_INF.
    """
    lq = _prefill_lq(head_dim)
    return np.concatenate(
        [
            np.zeros(lq * head_dim, np.float32).view(np.uint8),
            np.ones(lq, np.float32).view(np.uint8),
            np.zeros(lq, np.float32).view(np.uint8),
            np.full(2 * lq, _PREFILL_NEG_INF, bfloat16).view(np.uint8),
        ]
    )

flm_gemma4_prefill_round_begin

flm_gemma4_prefill_round_begin(
    *, head_dim: int = 512
) -> ExternalFunction

attn_round_begin of the prefill kernels, its five outputs packed as bytes.

out is y | c | l | prev_m | new_m.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_round_begin(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_round_begin`` of the prefill kernels, its five outputs packed as bytes.

    ``out`` is ``y | c | l | prev_m | new_m``.
    """
    lq = _prefill_lq(head_dim)
    n = 4 * lq * head_dim + 12 * lq
    return _prefill_core(
        "attn_round_begin_core",
        head_dim,
        [np.ndarray[(n,), np.dtype[np.uint8]]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out,),
            reference=partial(flm_gemma4_prefill_round_begin_ref, head_dim=head_dim),
            tolerance=Tolerance.exact(note="constant stores"),
            ops_per_call=0,
        ),
    )

flm_gemma4_prefill_qk_core_ref

flm_gemma4_prefill_qk_core_ref(
    q,
    k,
    m,
    inner_k,
    inner_q,
    inner_k_current,
    *,
    head_dim=512
)

Numpy reference for flm_gemma4_prefill_qk_core: s | m.

s = q @ k.T, NEG_INF where masked, and m folds in s lanewise. The tolerance covers the bfp16 operands and the bf16 narrowing.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_qk_core_ref(
    q, k, m, inner_k, inner_q, inner_k_current, *, head_dim=512
):
    """Numpy reference for [`flm_gemma4_prefill_qk_core`][iron.kernels.flm_gemma4.flm_gemma4_prefill_qk_core]: ``s | m``.

    ``s = q @ k.T``, NEG_INF where masked, and ``m`` folds in ``s``
    lanewise. The tolerance covers the bfp16 operands and the bf16 narrowing.
    """
    lq = lk = _prefill_lq(head_dim)
    q, k = _tiles(q, lq, head_dim), _tiles(k, lk, head_dim, col_major=True)
    s = q @ k.transpose(0, 2, 1)
    keep = _prefill_keep(lq, lk, inner_k, inner_q, inner_k_current)
    s = np.where(keep, s, _PREFILL_NEG_INF)
    m = np.maximum(np.asarray(m, np.float64).reshape(-1, lq, lk), s)
    return np.concatenate([s.reshape(len(s), -1), m.reshape(len(m), -1)], axis=1)

flm_gemma4_prefill_qk_core

flm_gemma4_prefill_qk_core(
    *, head_dim: int = 512
) -> ExternalFunction

attn_qk_step's arithmetic for one key chunk, without its k lock.

attn_qk_core(q, k, m, inner_k, inner_q, inner_k_current, out): s = q @ k.T through the bfp16-emulated mmul, masked, and m folded with it, into out = s | m. The three positions are apply_mask_and_get_max's, which attn_qk_step derives from its RTPs. q and k arrive in the mmul's 8x8 tiles, k's tiles in column-major order.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_qk_core(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_qk_step``'s arithmetic for one key chunk, without its k lock.

    ``attn_qk_core(q, k, m, inner_k, inner_q, inner_k_current, out)``:
    ``s = q @ k.T`` through the bfp16-emulated mmul, masked, and ``m``
    folded with it, into ``out = s | m``. The three positions are
    ``apply_mask_and_get_max``'s, which attn_qk_step derives from its RTPs.
    q and k arrive in the mmul's 8x8 tiles, k's tiles in column-major order.
    """
    lq = lk = _prefill_lq(head_dim)
    kv = np.ndarray[(lq * head_dim,), _BF16]
    m = np.ndarray[(lq * lk,), _BF16]
    return _prefill_core(
        "attn_qk_core",
        head_dim,
        [kv, kv, m, np.int32, np.int32, np.int32, np.ndarray[(2 * lq * lk,), _BF16]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Param, Param, Param, Out),
            sample=partial(_prefill_qk_sample, head_dim=head_dim),
            reference=partial(flm_gemma4_prefill_qk_core_ref, head_dim=head_dim),
            tolerance=Tolerance.bounded(
                partial(_prefill_qk_bound, head_dim=head_dim),
                note="bfp16ebs8 operand rounding (zero on the sample's grid), "
                "float32 accumulation, and the bf16 narrowing of s",
            ),
            ops_per_call=2 * lq * lk * head_dim,
        ),
    )

flm_gemma4_prefill_fv_core_ref

flm_gemma4_prefill_fv_core_ref(y, sv, *, head_dim=512)

Numpy reference for flm_gemma4_prefill_fv_core: y + s @ v.

y and s are in the mmul's 8x8 tiles, v in column-major tile order; the tolerance covers the bfp16 operands.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_fv_core_ref(y, sv, *, head_dim=512):
    """Numpy reference for [`flm_gemma4_prefill_fv_core`][iron.kernels.flm_gemma4.flm_gemma4_prefill_fv_core]: ``y + s @ v``.

    ``y`` and ``s`` are in the mmul's 8x8 tiles, ``v`` in column-major tile
    order; the tolerance covers the bfp16 operands.
    """
    lq = lk = _prefill_lq(head_dim)
    s, v = _prefill_sv(sv, lq, lk, head_dim)
    return _to_tiles(_tiles(y, lq, head_dim) + s @ v)

flm_gemma4_prefill_fv_core

flm_gemma4_prefill_fv_core(
    *, head_dim: int = 512
) -> ExternalFunction

attn_fv_step's arithmetic for one key chunk, without its v lock.

attn_fv_core(y, sv, y_out): y_out = y + s @ v through flm_attn_fv's bfp16-emulated mmul, with sv = s | v. y and s are in the mmul's 8x8 tiles, v's tiles in column-major order.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_fv_core(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_fv_step``'s arithmetic for one key chunk, without its v lock.

    ``attn_fv_core(y, sv, y_out)``: ``y_out = y + s @ v`` through
    ``flm_attn_fv``'s bfp16-emulated mmul, with ``sv = s | v``. y and s are in
    the mmul's 8x8 tiles, v's tiles in column-major order.
    """
    lq = lk = _prefill_lq(head_dim)
    y = np.ndarray[(lq * head_dim,), _F32]
    return _prefill_core(
        "attn_fv_core",
        head_dim,
        [y, np.ndarray[(lq * lk + lk * head_dim,), _BF16], y],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            sample=partial(_prefill_fv_sample, head_dim=head_dim),
            reference=partial(flm_gemma4_prefill_fv_core_ref, head_dim=head_dim),
            tolerance=Tolerance.bounded(
                partial(_prefill_fv_core_bound, head_dim=head_dim),
                note="bfp16ebs8 operand rounding (zero on the sample's grid) "
                "and float32 accumulation",
            ),
            ops_per_call=2 * lq * lk * head_dim,
        ),
    )

flm_gemma4_prefill_block_mid_core_ref

flm_gemma4_prefill_block_mid_core_ref(
    in_bf16, in_f32, *, head_dim=512
)

Numpy reference for flm_gemma4_prefill_block_mid_core.

(s | new_m, y | l | c): new_m the row max of m, p = 2**((s - new_m) * log2e) and c = 2**((prev_m - new_m) * log2e) with FastFlowLM's bf16 log2(e), l = sum(p) + c * l and y = c * y.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_block_mid_core_ref(in_bf16, in_f32, *, head_dim=512):
    """Numpy reference for [`flm_gemma4_prefill_block_mid_core`][iron.kernels.flm_gemma4.flm_gemma4_prefill_block_mid_core].

    ``(s | new_m, y | l | c)``: ``new_m`` the row max of ``m``,
    ``p = 2**((s - new_m) * log2e)`` and ``c = 2**((prev_m - new_m) * log2e)``
    with FastFlowLM's bf16 log2(e), ``l = sum(p) + c * l`` and ``y = c * y``.
    """
    lq = _prefill_lq(head_dim)
    p, new_m, y, row_sums, c, _, _ = _prefill_block_mid_math(in_bf16, in_f32, head_dim)
    return (
        np.concatenate([_prefill_block_mid_s(p, lq), new_m], axis=1),
        np.concatenate([y, row_sums, c], axis=1),
    )

flm_gemma4_prefill_block_mid_core

flm_gemma4_prefill_block_mid_core(
    *, head_dim: int = 512
) -> ExternalFunction

attn_block_mid of the prefill kernels, its state passed in and out.

attn_block_mid_core(in_bf16, in_f32, out_bf16, out_f32) with in_bf16 = s | m | prev_m, in_f32 = y | l, out_bf16 = s | new_m and out_f32 = y | l | c: the row max, FastFlowLM's softmax, the correction c and the rescaled l and y.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_block_mid_core(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_block_mid`` of the prefill kernels, its state passed in and out.

    ``attn_block_mid_core(in_bf16, in_f32, out_bf16, out_f32)`` with
    ``in_bf16 = s | m | prev_m``, ``in_f32 = y | l``, ``out_bf16 = s | new_m``
    and ``out_f32 = y | l | c``: the row max, FastFlowLM's softmax, the
    correction ``c`` and the rescaled ``l`` and ``y``.
    """
    lq = lk = _prefill_lq(head_dim)
    return _prefill_core(
        "attn_block_mid_core",
        head_dim,
        [
            np.ndarray[(lq * 128 + lq * lk + lq,), _BF16],
            np.ndarray[(lq * head_dim + lq,), _F32],
            np.ndarray[(lq * 128 + lq,), _BF16],
            np.ndarray[(lq * head_dim + 2 * lq,), _F32],
        ],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out, Out),
            sample=partial(_prefill_block_mid_sample, head_dim=head_dim),
            reference=partial(flm_gemma4_prefill_block_mid_core_ref, head_dim=head_dim),
            tolerance=Tolerance.bounded(
                partial(_prefill_block_mid_bound, head_dim=head_dim),
                note="aie::exp2's 6.98% (linalg.mha_softmax), bf16 rounding of "
                "the exp2 arguments, in the kernel (zero on the sample's grid) "
                "and inside aie::exp2 (measured on npu2), bfp16ebs8 "
                "rounding of p in the row sums, the emulated float32 multiply",
            ),
            # Per score: subtract, scale, exp2, sum; per y element: one multiply.
            ops_per_call=4 * lq * 128 + lq * head_dim,
        ),
    )

flm_gemma4_decode_glu

flm_gemma4_decode_glu(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the decode layer's gated linear unit, from flm_gemma4/decode_glu.cc.

glu(y, x_ping, x_pong, y_ping, y_pong, skip): waits on the RTP lock, then, unless skip[0] is set, turns 2 * intermediate_size gate and up values, arriving glu_slice at a time in the x ping-pong pair, into intermediate_size activations in the core-local y (twice that with double_wide_mlp), and sends them out glu_slice // 2 at a time through the y pair once per down-projection repeat. skip is an int32 RTP buffer of 16 words.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults x_prod_lock=0, x_cons_lock=1, y_prod_lock=2, y_cons_lock=3 and rtp_lock=6.	{}

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_glu(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the decode layer's gated linear unit, from ``flm_gemma4/decode_glu.cc``.

    ``glu(y, x_ping, x_pong, y_ping, y_pong, skip)``: waits on the RTP lock,
    then, unless ``skip[0]`` is set, turns ``2 * intermediate_size`` gate and
    up values, arriving ``glu_slice`` at a time in the ``x`` ping-pong pair,
    into ``intermediate_size`` activations in the core-local ``y`` (twice that
    with ``double_wide_mlp``), and sends them out ``glu_slice // 2`` at a time
    through the ``y`` pair once per down-projection repeat. ``skip`` is an
    int32 RTP buffer of 16 words.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``x_prod_lock=0``,
            ``x_cons_lock=1``, ``y_prod_lock=2``, ``y_cons_lock=3`` and
            ``rtp_lock=6``.
    """
    bf16 = np.dtype[bfloat16]
    hid = geometry.intermediate_size * (2 if geometry.double_wide_mlp else 1)
    x_ty = np.ndarray[(geometry.glu_slice,), bf16]
    y_ty = np.ndarray[(geometry.glu_slice // 2,), bf16]
    return _decode_kernel(
        "glu",
        "glu",
        [
            np.ndarray[(hid,), bf16],
            x_ty,
            x_ty,
            y_ty,
            y_ty,
            np.ndarray[(16,), np.dtype[np.int32]],
        ],
        (Out, In, In, Out, Out, Param),
        dict(x_prod_lock=0, x_cons_lock=1, y_prod_lock=2, y_cons_lock=3, rtp_lock=6),
        locks,
        geometry,
        lut=True,
    )

flm_gemma4_glu_core_ref

flm_gemma4_glu_core_ref(x)

Numpy reference for flm_gemma4_glu_core: gelu(gate) * up.

x holds, per call, the up half then the gate half. GELU is the tanh form Gemma 4 uses; the tolerance covers the kernel's table.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_glu_core_ref(x):
    """Numpy reference for [`flm_gemma4_glu_core`][iron.kernels.flm_gemma4.flm_gemma4_glu_core]: ``gelu(gate) * up``.

    ``x`` holds, per call, the up half then the gate half. GELU is the tanh
    form Gemma 4 uses; the tolerance covers the kernel's table.
    """
    up, gate = _glu_split(x)
    return (_gelu_tanh(gate) * up).astype(np.float32)

flm_gemma4_glu_core

flm_gemma4_glu_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

pseduo_glu of flm_gemma4_decode_glu: gelu(gate) * up for one slice.

glu_core(x, y): x is glu_slice bf16, the up half then the gate half; y is glu_slice // 2 bf16. GELU goes through aie_runtime_lib's table (getGeluBf16).

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_glu_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``pseduo_glu`` of [`flm_gemma4_decode_glu`][iron.kernels.flm_gemma4.flm_gemma4_decode_glu]: ``gelu(gate) * up`` for one slice.

    ``glu_core(x, y)``: ``x`` is ``glu_slice`` bf16, the up half then the
    gate half; ``y`` is ``glu_slice // 2`` bf16. GELU goes through
    aie_runtime_lib's table (``getGeluBf16``).

    Args:
        geometry: The model the kernel builds for.
    """
    n = geometry.glu_slice
    return _wrapper(
        flm_gemma4_decode_glu(geometry=geometry),
        "decode_glu_core.cc",
        "glu_core",
        [np.ndarray[(n,), _BF16], np.ndarray[(n // 2,), _BF16]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_glu_core_ref,
            sample=partial(_glu_core_sample, n=n),
            tolerance=Tolerance.bounded(
                _glu_core_bound,
                note="the GELU table's error against tanh-GELU, measured on the "
                "host from aie_runtime_lib's table, plus a bf16 ulp of its slope "
                "and the floor bf16 stores of the GELU and the product",
            ),
            uses_lut=True,
            # Per output: the table line's multiply and add, then the product.
            ops_per_call=3 * (n // 2),
        ),
    )

flm_gemma4_decode_attn_kv

flm_gemma4_decode_attn_kv(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the global-attention score-times-value core for one KV head, from flm_gemma4/decode_attn_kv.cc.

The decode layer's global attention runs on two cores. This one folds each round's scores and running-max corrections from the qk core into the softmax denominator l and the float32 accumulator y, and at the end writes o = y / l. The returned kernel is attn_kv_begin, which zeroes y and l and waits for the RTPs. Per round the Worker calls attn_kv_round, then once attn_kv_finish. y and o need 64-byte-aligned base pointers.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; num_kv_heads must be 1.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults v_prod_lock=2, v_cons_lock=3, o_prod_lock=0, o_cons_lock=1 and l_cons_lock=8 (a lock of the qk core, on the tile below).	{}

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_attn_kv(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the global-attention score-times-value core for one KV head, from ``flm_gemma4/decode_attn_kv.cc``.

    The decode layer's global attention runs on two cores. This one folds
    each round's scores and running-max corrections from the qk core into
    the softmax denominator ``l`` and the float32 accumulator ``y``, and at
    the end writes ``o = y / l``. The returned kernel is ``attn_kv_begin``,
    which zeroes ``y`` and ``l`` and waits for the RTPs. Per round the
    Worker calls ``attn_kv_round``, then once ``attn_kv_finish``. ``y`` and
    ``o`` need 64-byte-aligned base pointers.

    Args:
        geometry: The model the kernel builds for; ``num_kv_heads`` must be 1.
        **locks: Core lock ids overriding the defaults ``v_prod_lock=2``,
            ``v_cons_lock=3``, ``o_prod_lock=0``, ``o_cons_lock=1`` and
            ``l_cons_lock=8`` (a lock of the qk core, on the tile below).
    """
    _decode_kv_heads("flm_gemma4_decode_attn_kv", geometry, 1)
    t = _decode_attn_types(geometry, geometry.dh, geometry.dh)
    fn = _decode_kernel(
        "attn_kv",
        "attn_kv_begin",
        [t["y"], t["l"]],
        (Out, Out),
        _DECODE_KV_LOCKS,
        locks,
        geometry,
        lut=True,
        cls=_AttnKvKernel,
    )
    bind = fn.object_file.bind
    fn.attn_kv_round = bind("attn_kv_round", [t["s"], t["kv"], t["kv"], t["y"], t["l"]])
    fn.attn_kv_finish = bind("attn_kv_finish", [t["y"], t["o"], t["l"]])
    return fn

flm_gemma4_attn_kv_core

flm_gemma4_attn_kv_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One attn_kv_round of flm_gemma4_decode_attn_kv without its v lock.

attn_kv_round_core(sv, ly_in, ly_out): ly_out is ly_in with l = c * l + rowsum(s) and y = c * y + s @ v per query head. sv holds s, then v. ly holds y, then l in a 16-float slot.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_kv_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """One ``attn_kv_round`` of [`flm_gemma4_decode_attn_kv`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv] without its v lock.

    ``attn_kv_round_core(sv, ly_in, ly_out)``: ``ly_out`` is ``ly_in`` with
    ``l = c * l + rowsum(s)`` and ``y = c * y + s @ v`` per query head. ``sv``
    holds ``s``, then ``v``. ``ly`` holds ``y``, then ``l`` in a 16-float slot.
    """
    base = flm_gemma4_decode_attn_kv(geometry=geometry)
    dh = geometry.dh
    return _attn_kv_round_core(
        base,
        "decode_attn_kv_core.cc",
        "attn_kv_round_core",
        dh=dh,
        kv_width=dh,
        impl="1x8x1",
        rows=8,
        ref=partial(flm_gemma4_attn_kv_core_ref, dh=dh),
        stack_bytes=2304,  # aiecc measured_stack_size
    )

flm_gemma4_attn_kv_core_ref

flm_gemma4_attn_kv_core_ref(sv, ly, *, dh=512)

Numpy reference for flm_gemma4_attn_kv_core, in float64.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_kv_core_ref(sv, ly, *, dh=512):
    """Numpy reference for [`flm_gemma4_attn_kv_core`][iron.kernels.flm_gemma4.flm_gemma4_attn_kv_core], in float64."""
    return _attn_kv_round_ref(sv, ly, dh=dh, impl="1x8x1")

flm_gemma4_decode_attn_kv_kvh2

flm_gemma4_decode_attn_kv_kvh2(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE,
    **locks: int
) -> ExternalFunction

Build the global-attention score-times-value core for two KV heads, from flm_gemma4/decode_attn_kv_kvh2.cc.

The two-KV-head (E4B) sibling of flm_gemma4_decode_attn_kv: the Worker runs a round as one attn_kv_s_begin, which folds the scores into l and rescales y, then one attn_kv_v_half per KV head.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; num_kv_heads must be 2.	FLM_GEMMA4_E4B_DECODE
**locks	int	Core lock ids overriding the defaults v_prod_lock=2, v_cons_lock=3, o_prod_lock=0, o_cons_lock=1 and l_cons_lock=8 (a lock of the qk core, on the tile below).	{}

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_attn_kv_kvh2(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the global-attention score-times-value core for two KV heads, from ``flm_gemma4/decode_attn_kv_kvh2.cc``.

    The two-KV-head (E4B) sibling of
    [`flm_gemma4_decode_attn_kv`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv]:
    the Worker runs a round as one ``attn_kv_s_begin``, which folds the
    scores into ``l`` and rescales ``y``, then one ``attn_kv_v_half`` per
    KV head.

    Args:
        geometry: The model the kernel builds for; ``num_kv_heads`` must be 2.
        **locks: Core lock ids overriding the defaults ``v_prod_lock=2``,
            ``v_cons_lock=3``, ``o_prod_lock=0``, ``o_cons_lock=1`` and
            ``l_cons_lock=8`` (a lock of the qk core, on the tile below).
    """
    _decode_kv_heads("flm_gemma4_decode_attn_kv_kvh2", geometry, 2)
    t = _decode_attn_types(geometry, geometry.dh, geometry.dh)
    fn = _decode_kernel(
        "attn_kv_kvh2",
        "attn_kv_begin",
        [t["y"], t["l"]],
        (Out, Out),
        _DECODE_KV_LOCKS,
        locks,
        geometry,
        lut=True,
        cls=_AttnKvKvh2Kernel,
    )
    bind = fn.object_file.bind
    fn.attn_kv_s_begin = bind("attn_kv_s_begin", [t["s"], t["y"], t["l"]])
    fn.attn_kv_v_half = bind(
        "attn_kv_v_half", [t["s"], t["kv"], t["kv"], t["y"], np.int32]
    )
    fn.attn_kv_finish = bind("attn_kv_finish", [t["y"], t["o"], t["l"]])
    return fn

flm_gemma4_attn_kv_kvh2_core

flm_gemma4_attn_kv_kvh2_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction

attn_kv_s_begin, then one KV head's attn_kv_v_half, without the v lock.

The entry points are those of flm_gemma4_decode_attn_kv_kvh2. attn_kv_kvh2_round_core(sv, ly_in, ly_out, kv_head): ly_out is ly_in with l = c * l + rowsum(s) and y = c * y, plus s @ v for query heads 4 * kv_head to 4 * kv_head + 3. sv holds s, then v. ly holds y, then l in a 16-float slot. Both KV heads' v buffers do not fit in core memory beside ly, so a call reads one.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_kv_kvh2_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction:
    """``attn_kv_s_begin``, then one KV head's ``attn_kv_v_half``, without the v lock.

    The entry points are those of
    [`flm_gemma4_decode_attn_kv_kvh2`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv_kvh2].
    ``attn_kv_kvh2_round_core(sv, ly_in, ly_out, kv_head)``: ``ly_out`` is
    ``ly_in`` with ``l = c * l + rowsum(s)`` and ``y = c * y``, plus ``s @ v``
    for query heads ``4 * kv_head`` to ``4 * kv_head + 3``. ``sv`` holds
    ``s``, then ``v``. ``ly`` holds ``y``, then ``l`` in a 16-float slot.
    Both KV heads' v buffers do not fit in core memory beside ``ly``, so a
    call reads one.
    """
    base = flm_gemma4_decode_attn_kv_kvh2(geometry=geometry)
    dh = geometry.dh
    return _attn_kv_round_core(
        base,
        "decode_attn_kv_kvh2_core.cc",
        "attn_kv_kvh2_round_core",
        dh=dh,
        kv_width=dh,
        impl="kvh2",
        rows=4,
        ref=partial(flm_gemma4_attn_kv_kvh2_core_ref, dh=dh),
    )

flm_gemma4_attn_kv_kvh2_core_ref

flm_gemma4_attn_kv_kvh2_core_ref(
    sv, ly, kv_head, *, dh=512
)

Numpy reference for flm_gemma4_attn_kv_kvh2_core, in float64.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_kv_kvh2_core_ref(sv, ly, kv_head, *, dh=512):
    """Numpy reference for [`flm_gemma4_attn_kv_kvh2_core`][iron.kernels.flm_gemma4.flm_gemma4_attn_kv_kvh2_core], in float64."""
    return _attn_kv_round_ref(sv, ly, int(kv_head), dh=dh, impl="kvh2")

flm_gemma4_decode_attn_qk

flm_gemma4_decode_attn_qk(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build an attention query-times-key core, from flm_gemma4/decode_attn_qk.cc.

The other half of flm_gemma4_decode_attn_kv: it multiplies the query heads into each round's 16 key rows, masks the keys past the sequence length, keeps the running row maximum m and hands the kv core the exponentiated scores with their corrections. The returned kernel is attn_qk_begin, which sets m to -inf and releases the kv core; per round the Worker calls attn_qk_round. The global build serves one KV head (E2B), the sliding-window build both.

Parameters:

Name	Type	Description	Default
sliding_window	bool	Build for the sliding-window layers' head dim.	False
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; without sliding_window num_kv_heads must be 1.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults k_prod_lock=2, k_cons_lock=3 and l_cons_lock=8.	{}

Source code in python/iron/kernels/flm_gemma4.py

@dtypes(
    [
        {"sliding_window": True},
        {"sliding_window": True, "geometry": FLM_GEMMA4_E4B_DECODE},
    ]
)
def flm_gemma4_decode_attn_qk(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int,
) -> ExternalFunction:
    """Build an attention query-times-key core, from ``flm_gemma4/decode_attn_qk.cc``.

    The other half of [`flm_gemma4_decode_attn_kv`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv]:
    it multiplies the query heads into each round's 16 key rows, masks the
    keys past the sequence length, keeps the running row maximum ``m`` and
    hands the kv core the exponentiated scores with their corrections. The
    returned kernel is ``attn_qk_begin``, which sets ``m`` to -inf and
    releases the kv core; per round the Worker calls ``attn_qk_round``.
    The global build serves one KV head (E2B), the sliding-window build both.

    Args:
        sliding_window: Build for the sliding-window layers' head dim.
        geometry: The model the kernel builds for; without ``sliding_window``
            ``num_kv_heads`` must be 1.
        **locks: Core lock ids overriding the defaults ``k_prod_lock=2``,
            ``k_cons_lock=3`` and ``l_cons_lock=8``.
    """
    if not sliding_window:
        _decode_kv_heads("flm_gemma4_decode_attn_qk", geometry, 1)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    t = _decode_attn_types(geometry, dh, geometry.num_kv_heads * dh)
    fn = _decode_kernel(
        "attn_qk",
        "attn_qk_begin",
        [t["m"]],
        (Out,),
        dict(k_prod_lock=2, k_cons_lock=3, l_cons_lock=8),
        locks,
        geometry,
        flags=(f"-DFLM_GEMMA4_DECODE_ATTN_QK_SWA={int(sliding_window)}",),
        cls=_AttnQkKernel,
    )
    fn.attn_qk_round = fn.object_file.bind(
        "attn_qk_round",
        [t["q"], t["kv"], t["kv"], t["s"], t["m"], t["c"], np.int32, np.int32],
    )
    return fn

flm_gemma4_decode_attn_qk_kvh2

flm_gemma4_decode_attn_qk_kvh2(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE,
    **locks: int
) -> ExternalFunction

Build the global-attention query-times-key core for two KV heads, from flm_gemma4/decode_attn_qk_kvh2.cc.

The two-KV-head (E4B) sibling of flm_gemma4_decode_attn_qk: the Worker runs a round as one attn_qk_half per KV head, then attn_qk_store_c, which copies the corrections into the scores.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; num_kv_heads must be 2.	FLM_GEMMA4_E4B_DECODE
**locks	int	Core lock ids overriding the defaults k_prod_lock=2, k_cons_lock=3 and l_prod_lock=8.	{}

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_attn_qk_kvh2(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the global-attention query-times-key core for two KV heads, from ``flm_gemma4/decode_attn_qk_kvh2.cc``.

    The two-KV-head (E4B) sibling of
    [`flm_gemma4_decode_attn_qk`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_qk]:
    the Worker runs a round as one ``attn_qk_half`` per KV head, then
    ``attn_qk_store_c``, which copies the corrections into the scores.

    Args:
        geometry: The model the kernel builds for; ``num_kv_heads`` must be 2.
        **locks: Core lock ids overriding the defaults ``k_prod_lock=2``,
            ``k_cons_lock=3`` and ``l_prod_lock=8``.
    """
    _decode_kv_heads("flm_gemma4_decode_attn_qk_kvh2", geometry, 2)
    t = _decode_attn_types(geometry, geometry.dh, geometry.dh)
    fn = _decode_kernel(
        "attn_qk_kvh2",
        "attn_qk_begin",
        [t["m"]],
        (Out,),
        dict(k_prod_lock=2, k_cons_lock=3, l_prod_lock=8),
        locks,
        geometry,
        cls=_AttnQkKvh2Kernel,
    )
    bind = fn.object_file.bind
    fn.attn_qk_half = bind(
        "attn_qk_half",
        [
            t["q"],
            t["kv"],
            t["kv"],
            t["s"],
            t["m"],
            t["c"],
            np.int32,
            np.int32,
            np.int32,
        ],
    )
    fn.attn_qk_store_c = bind("attn_qk_store_c", [t["s"], t["c"]])
    return fn

flm_gemma4_attn_qk_core_ref

flm_gemma4_attn_qk_core_ref(qm, k, iter, L0)

Numpy reference for flm_gemma4_attn_qk_core.

Keys j with j < L0 - 16 * iter are valid. Per query row the new running max is max(m_in, valid scores); s holds exp(score - max) for valid keys and 0 for the others, then the 16 running maxima (rows 8 to 15 copied from m_in). The second output is c = exp(m_in - max). Both exponents are clamped at -87, as the kernel does.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_qk_core_ref(qm, k, iter, L0):
    """Numpy reference for [`flm_gemma4_attn_qk_core`][iron.kernels.flm_gemma4.flm_gemma4_attn_qk_core].

    Keys ``j`` with ``j < L0 - 16 * iter`` are valid. Per query row the new
    running max is ``max(m_in, valid scores)``; ``s`` holds ``exp(score -
    max)`` for valid keys and 0 for the others, then the 16 running maxima
    (rows 8 to 15 copied from ``m_in``). The second output is ``c = exp(m_in
    - max)``. Both exponents are clamped at -87, as the kernel does.
    """
    _, _, m_in, m_new, _, e, c = _attn_qk_parts(qm, k, iter, L0)
    m_out = m_in.copy()
    m_out[:, :8] = m_new
    return np.concatenate([e.reshape(len(e), -1), m_out], axis=1), c

flm_gemma4_attn_qk_core

flm_gemma4_attn_qk_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One attn_qk_round of flm_gemma4_decode_attn_qk, without its k lock.

attn_qk_round_core(qm, k, s, c, iter, L0) takes the running max after q in qm and writes it after the scores in s. c holds the corrections the round stores in the tail of its s object.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes(
    [
        {"sliding_window": True},
        {"sliding_window": True, "geometry": FLM_GEMMA4_E4B_DECODE},
    ]
)
def flm_gemma4_attn_qk_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> ExternalFunction:
    """One ``attn_qk_round`` of [`flm_gemma4_decode_attn_qk`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_qk], without its k lock.

    ``attn_qk_round_core(qm, k, s, c, iter, L0)`` takes the running max
    after q in ``qm`` and writes it after the scores in ``s``. ``c`` holds the
    corrections the round stores in the tail of its s object.
    """
    base = flm_gemma4_decode_attn_qk(sliding_window=sliding_window, geometry=geometry)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    kv_heads = geometry.num_kv_heads
    return _wrapper(
        base,
        "decode_attn_qk_core.cc",
        "attn_qk_round_core",
        _attn_qk_core_types(geometry, dh, kv_heads),
        _attn_qk_core_contract(dh, kv_heads, flm_gemma4_attn_qk_core_ref),
    )

flm_gemma4_attn_qk_kvh2_core

flm_gemma4_attn_qk_kvh2_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction

One round of flm_gemma4_decode_attn_qk_kvh2, without its k lock.

attn_qk_kvh2_round_core(qm, k, s, c, iter, L0) runs attn_qk_half for both KV heads, whose k objects k holds back to back. Otherwise as flm_gemma4_attn_qk_core.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_qk_kvh2_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction:
    """One round of [`flm_gemma4_decode_attn_qk_kvh2`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_qk_kvh2], without its k lock.

    ``attn_qk_kvh2_round_core(qm, k, s, c, iter, L0)`` runs
    ``attn_qk_half`` for both KV heads, whose k objects ``k`` holds back to
    back. Otherwise as
    [`flm_gemma4_attn_qk_core`][iron.kernels.flm_gemma4.flm_gemma4_attn_qk_core].
    """
    base = flm_gemma4_decode_attn_qk_kvh2(geometry=geometry)
    return _wrapper(
        base,
        "decode_attn_qk_kvh2_core.cc",
        "attn_qk_kvh2_round_core",
        _attn_qk_core_types(geometry, geometry.dh, 2),
        _attn_qk_core_contract(geometry.dh, 2, flm_gemma4_attn_qk_kvh2_core_ref),
    )

flm_gemma4_attn_qk_kvh2_core_ref

flm_gemma4_attn_qk_kvh2_core_ref(qm, k, iter, L0)

Numpy reference for flm_gemma4_attn_qk_kvh2_core; see flm_gemma4_attn_qk_core_ref.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_qk_kvh2_core_ref(qm, k, iter, L0):
    """Numpy reference for [`flm_gemma4_attn_qk_kvh2_core`][iron.kernels.flm_gemma4.flm_gemma4_attn_qk_kvh2_core]; see ``flm_gemma4_attn_qk_core_ref``."""
    return flm_gemma4_attn_qk_core_ref(qm, k, iter, L0)

flm_gemma4_decode_swa_attn_kv

flm_gemma4_decode_swa_attn_kv(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the sliding-window-attention score-times-value core, from flm_gemma4/decode_swa_attn_kv.cc.

flm_gemma4_decode_attn_kv at head dim swa_dh, with entry points swa_attn_kv_begin, swa_attn_kv_round and swa_attn_kv_finish. One round covers every KV head, so it builds for one and for two.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults v_prod_lock=2, v_cons_lock=3, o_prod_lock=0, o_cons_lock=1 and l_cons_lock=8 (a lock of the qk core, on the tile below).	{}

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_swa_attn_kv(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the sliding-window-attention score-times-value core, from ``flm_gemma4/decode_swa_attn_kv.cc``.

    [`flm_gemma4_decode_attn_kv`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv]
    at head dim ``swa_dh``, with entry points ``swa_attn_kv_begin``,
    ``swa_attn_kv_round`` and ``swa_attn_kv_finish``. One round covers every
    KV head, so it builds for one and for two.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``v_prod_lock=2``,
            ``v_cons_lock=3``, ``o_prod_lock=0``, ``o_cons_lock=1`` and
            ``l_cons_lock=8`` (a lock of the qk core, on the tile below).
    """
    t = _decode_attn_types(
        geometry, geometry.swa_dh, geometry.num_kv_heads * geometry.swa_dh
    )
    fn = _decode_kernel(
        "swa_attn_kv",
        "swa_attn_kv_begin",
        [t["y"], t["l"]],
        (Out, Out),
        _DECODE_KV_LOCKS,
        locks,
        geometry,
        lut=True,
        cls=_SwaAttnKvKernel,
    )
    bind = fn.object_file.bind
    fn.swa_attn_kv_round = bind(
        "swa_attn_kv_round", [t["s"], t["kv"], t["kv"], t["y"], t["l"]]
    )
    fn.swa_attn_kv_finish = bind("swa_attn_kv_finish", [t["y"], t["o"], t["l"]])
    return fn

flm_gemma4_swa_attn_kv_core

flm_gemma4_swa_attn_kv_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One swa_attn_kv_round of flm_gemma4_decode_swa_attn_kv without its v lock.

swa_attn_kv_round_core(sv, ly_in, ly_out): ly_out is ly_in with l = c * l + rowsum(s) and y = c * y + s @ v per query head, each head reading its KV head's half of v when there are two. sv holds s, then v. ly holds y, then l in a 16-float slot.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_swa_attn_kv_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """One ``swa_attn_kv_round`` of [`flm_gemma4_decode_swa_attn_kv`][iron.kernels.flm_gemma4.flm_gemma4_decode_swa_attn_kv] without its v lock.

    ``swa_attn_kv_round_core(sv, ly_in, ly_out)``: ``ly_out`` is ``ly_in``
    with ``l = c * l + rowsum(s)`` and ``y = c * y + s @ v`` per query head,
    each head reading its KV head's half of ``v`` when there are two. ``sv``
    holds ``s``, then ``v``. ``ly`` holds ``y``, then ``l`` in a 16-float
    slot.
    """
    base = flm_gemma4_decode_swa_attn_kv(geometry=geometry)
    dh, kv_heads = geometry.swa_dh, geometry.num_kv_heads
    return _attn_kv_round_core(
        base,
        "decode_swa_attn_kv_core.cc",
        "swa_attn_kv_round_core",
        dh=dh,
        kv_width=kv_heads * dh,
        impl="1x8x1" if kv_heads == 1 else "2x4x1",
        rows=8,
        ref=partial(flm_gemma4_swa_attn_kv_core_ref, dh=dh, kv_heads=kv_heads),
    )

flm_gemma4_swa_attn_kv_core_ref

flm_gemma4_swa_attn_kv_core_ref(
    sv, ly, *, dh=256, kv_heads=1
)

Numpy reference for flm_gemma4_swa_attn_kv_core, in float64.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_swa_attn_kv_core_ref(sv, ly, *, dh=256, kv_heads=1):
    """Numpy reference for [`flm_gemma4_swa_attn_kv_core`][iron.kernels.flm_gemma4.flm_gemma4_swa_attn_kv_core], in float64."""
    impl = "1x8x1" if kv_heads == 1 else "2x4x1"
    return _attn_kv_round_ref(sv, ly, dh=dh, impl=impl)

flm_gemma4_decode_rope

flm_gemma4_decode_rope(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build a layer's q/k norm and rotary embedding, from flm_gemma4/decode_rope.cc.

rope(q, k, v, qkv_ping, qkv_pong, rope_w, skip_kv): takes the q, k and v projections one head (dh bf16) at a time through the qkv ping-pong pair, RMS-normalizes each q and k head in place against its weight and rotates it into q or k, and RMS-normalizes each v head into v. Unless skip_kv[0] is set, when the layer reuses another layer's KV cache, it hands k and v on. q is q_heads_padded * dh bf16, k and v num_kv_heads * dh, rope_w the cos and sin halves then, with qk_norm, the q and k norm weights (3 * dh), and skip_kv an int32 RTP buffer of 16 words. dh is the geometry's swa_dh with sliding_window, else its dh.

Parameters:

Name	Type	Description	Default
sliding_window	bool	Build for the sliding-window layers' head dim.	False
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults qkv_prod_lock=0, qkv_cons_lock=1, k_prod_lock=4, k_cons_lock=5, v_prod_lock=6, v_cons_lock=7, rope_prod_lock=8 and rope_cons_lock=9.	{}

Source code in python/iron/kernels/flm_gemma4.py

@dtypes(
    [
        {"geometry": FLM_GEMMA4_E4B_DECODE},
        {"sliding_window": True},
        {"sliding_window": True, "geometry": FLM_GEMMA4_E4B_DECODE},
    ]
)
def flm_gemma4_decode_rope(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int,
) -> ExternalFunction:
    """Build a layer's q/k norm and rotary embedding, from ``flm_gemma4/decode_rope.cc``.

    ``rope(q, k, v, qkv_ping, qkv_pong, rope_w, skip_kv)``: takes the q, k
    and v projections one head (``dh`` bf16) at a time through the
    ``qkv`` ping-pong pair, RMS-normalizes each q and k head in place against
    its weight and rotates it into ``q`` or ``k``, and RMS-normalizes each v
    head into ``v``. Unless ``skip_kv[0]`` is set, when the layer reuses
    another layer's KV cache, it hands ``k`` and ``v`` on. ``q`` is
    ``q_heads_padded * dh`` bf16, ``k`` and ``v`` ``num_kv_heads * dh``,
    ``rope_w`` the cos and sin halves then, with ``qk_norm``, the q and k
    norm weights (``3 * dh``), and ``skip_kv`` an int32 RTP buffer of 16
    words. ``dh`` is the geometry's ``swa_dh`` with ``sliding_window``, else
    its ``dh``.

    Args:
        sliding_window: Build for the sliding-window layers' head dim.
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``qkv_prod_lock=0``,
            ``qkv_cons_lock=1``, ``k_prod_lock=4``, ``k_cons_lock=5``,
            ``v_prod_lock=6``, ``v_cons_lock=7``, ``rope_prod_lock=8`` and
            ``rope_cons_lock=9``.
    """
    dh = geometry.swa_dh if sliding_window else geometry.dh
    bf16 = np.dtype[bfloat16]
    kv = geometry.num_kv_heads * dh
    qkv_ty = np.ndarray[(dh,), bf16]
    return _decode_kernel(
        "rope",
        "rope",
        [
            np.ndarray[(_decode_q_heads_padded(geometry) * dh,), bf16],
            np.ndarray[(kv,), bf16],
            np.ndarray[(kv,), bf16],
            qkv_ty,
            qkv_ty,
            np.ndarray[(dh * (3 if geometry.qk_norm else 1),), bf16],
            _DECODE_RTP,
        ],
        (Out, Out, Out, InOut, InOut, In, Param),
        _DECODE_ROPE_LOCKS,
        locks,
        geometry,
        flags=(f"-DFLM_GEMMA4_DECODE_ROPE_SWA={int(sliding_window)}",),
    )

flm_gemma4_rope_core_ref

flm_gemma4_rope_core_ref(x, rope_w)

Numpy reference for flm_gemma4_rope_core.

Per call, x is a q head then a k head and rope_w is [cos | sin | q weight | k weight]. Each head n = x * w / sqrt(mean(x^2) + 1e-6) is rotated by halves: y1 = n1 cos - n2 sin, y2 = n1 sin + n2 cos, with n1, n2 its first and second half.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_rope_core_ref(x, rope_w):
    """Numpy reference for [`flm_gemma4_rope_core`][iron.kernels.flm_gemma4.flm_gemma4_rope_core].

    Per call, ``x`` is a q head then a k head and ``rope_w`` is ``[cos |
    sin | q weight | k weight]``. Each head ``n = x * w / sqrt(mean(x^2) +
    1e-6)`` is rotated by halves: ``y1 = n1 cos - n2 sin``, ``y2 = n1 sin +
    n2 cos``, with ``n1``, ``n2`` its first and second half.
    """
    a, b = _rope_core_terms(x, rope_w)
    y = (a + b).reshape(len(a), -1)
    return y.astype(np.float32).astype(bfloat16)

flm_gemma4_rope_core

flm_gemma4_rope_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

_rotate_t of flm_gemma4_decode_rope on one q and one k head.

rope_head_core(x, rope_w, y): x is a q head then a k head, rope_w the base kernel's [cos | sin | q weight | k weight]. The kernel normalizes x in place; the write lands in the core's input element, which the next DMA fill overwrites.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes(_ROPE_CORE_VARIANTS)
def flm_gemma4_rope_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> ExternalFunction:
    """``_rotate_t`` of [`flm_gemma4_decode_rope`][iron.kernels.flm_gemma4.flm_gemma4_decode_rope] on one q and one k head.

    ``rope_head_core(x, rope_w, y)``: ``x`` is a q head then a k head,
    ``rope_w`` the base kernel's ``[cos | sin | q weight | k weight]``.
    The kernel normalizes ``x`` in place; the write lands in the core's input
    element, which the next DMA fill overwrites.
    """
    base = flm_gemma4_decode_rope(sliding_window=sliding_window, geometry=geometry)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    heads = np.ndarray[(2 * dh,), _BF16]
    return _wrapper(
        base,
        "decode_rope_core.cc",
        "rope_head_core",
        [heads, np.ndarray[(3 * dh,), _BF16], heads],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            reference=flm_gemma4_rope_core_ref,
            sample=partial(_rope_core_sample, dh=dh),
            tolerance=Tolerance.bounded(
                _rope_core_bound,
                note="bf16 ulp <= 2**-7 |v|. The kernel floors each normalized "
                "value to bf16 (1 ulp, plus the rms scale's error, under "
                "2**-12), so each rotation product is off by that much; it "
                "floors the rotated sum to bf16 (1 ulp) and the reference "
                "rounds it to nearest (1/2 ulp)",
            ),
            # Per element: square-add and two multiplies for the norm, then
            # a multiply and a multiply-add for the rotation.
            ops_per_call=2 * 7 * dh,
        ),
    )

flm_gemma4_v_norm_core_ref

flm_gemma4_v_norm_core_ref(x)

Numpy reference for flm_gemma4_v_norm_core: x / sqrt(mean(x^2) + 1e-6).

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_v_norm_core_ref(x):
    """Numpy reference for [`flm_gemma4_v_norm_core`][iron.kernels.flm_gemma4.flm_gemma4_v_norm_core]: ``x / sqrt(mean(x^2) + 1e-6)``."""
    x = np.asarray(x, np.float64)
    return (x * _rms(x)).astype(np.float32).astype(bfloat16)

flm_gemma4_v_norm_core

flm_gemma4_v_norm_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

rms_norm_unweighted of flm_gemma4_decode_rope on one v head.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes(_ROPE_CORE_VARIANTS)
def flm_gemma4_v_norm_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> ExternalFunction:
    """``rms_norm_unweighted`` of [`flm_gemma4_decode_rope`][iron.kernels.flm_gemma4.flm_gemma4_decode_rope] on one v head."""
    base = flm_gemma4_decode_rope(sliding_window=sliding_window, geometry=geometry)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    head = np.ndarray[(dh,), _BF16]
    return _wrapper(
        base,
        "decode_rope_core.cc",
        "v_norm_core",
        [head, head],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_v_norm_core_ref,
            sample=partial(_v_norm_core_sample, dh=dh),
            tolerance=Tolerance.bounded(
                _v_norm_core_bound,
                note="bf16 ulp <= 2**-7 |v|: the kernel floors to bf16 (1 ulp), "
                "the reference rounds to nearest (1/2 ulp), and the rms scale "
                "is within 2**-12",
            ),
            # Per element: square-add, then one multiply.
            ops_per_call=3 * dh,
        ),
    )

flm_gemma4_decode_rms_residual

flm_gemma4_decode_rms_residual(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the decode layer's four RMS norms and two residual adds, from flm_gemma4/decode_rms_residual.cc.

rms_residual(y, x_ping, x_pong, y_final, w, x_temp_buf, is_swa, skip_kv): one layer's worth of norms. It normalizes the layer input into y for the QKV projection, then normalizes the attention output, adds the residual and normalizes the sum into y for the up/gate projection, and finally normalizes the MLP output and adds the second residual into y_final. The attention and MLP outputs arrive in the x ping-pong pair (model_dim bf16 each). y is model_dim + 16 bf16, a 16-element packet header then the row; w holds the four norm weights (4 * model_dim); x_temp_buf is 2 * model_dim of scratch. is_swa and skip_kv are int32 RTP buffers of 16 words that select how many projection repeats y is released for.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults w_prod_lock=0, w_cons_lock=1, y_prod_lock=2, y_cons_lock=3, x_prod_lock=4, x_cons_lock=5, rtp_available_lock=6, lm_head_out_prod_lock=7 and lm_head_out_cons_lock=8.	{}

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_rms_residual(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the decode layer's four RMS norms and two residual adds, from ``flm_gemma4/decode_rms_residual.cc``.

    ``rms_residual(y, x_ping, x_pong, y_final, w, x_temp_buf, is_swa, skip_kv)``:
    one layer's worth of norms. It normalizes the layer input into ``y`` for
    the QKV projection, then normalizes the attention output, adds the
    residual and normalizes the sum into ``y`` for the up/gate projection,
    and finally normalizes the MLP output and adds the second residual into
    ``y_final``. The attention and MLP outputs arrive in the ``x`` ping-pong
    pair (``model_dim`` bf16 each). ``y`` is ``model_dim + 16`` bf16, a
    16-element packet header then the row; ``w`` holds the four norm weights
    (``4 * model_dim``); ``x_temp_buf`` is ``2 * model_dim`` of scratch.
    ``is_swa`` and ``skip_kv`` are int32 RTP buffers of 16 words that select
    how many projection repeats ``y`` is released for.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``w_prod_lock=0``,
            ``w_cons_lock=1``, ``y_prod_lock=2``, ``y_cons_lock=3``,
            ``x_prod_lock=4``, ``x_cons_lock=5``, ``rtp_available_lock=6``,
            ``lm_head_out_prod_lock=7`` and ``lm_head_out_cons_lock=8``.
    """
    bf16 = np.dtype[bfloat16]
    d = geometry.model_dim
    x_ty = np.ndarray[(d,), bf16]
    return _decode_kernel(
        "rms_residual",
        "rms_residual",
        [
            np.ndarray[(d + 16,), bf16],
            x_ty,
            x_ty,
            x_ty,
            np.ndarray[(4, d), bf16],
            np.ndarray[(2, d), bf16],
            _DECODE_RTP,
            _DECODE_RTP,
        ],
        (Out, In, In, Out, In, Out, Param, Param),
        dict(
            w_prod_lock=0,
            w_cons_lock=1,
            y_prod_lock=2,
            y_cons_lock=3,
            x_prod_lock=4,
            x_cons_lock=5,
            rtp_available_lock=6,
            lm_head_out_prod_lock=7,
            lm_head_out_cons_lock=8,
        ),
        locks,
        geometry,
    )

flm_gemma4_rms_residual_core_ref

flm_gemma4_rms_residual_core_ref(x, w, residual)

Numpy reference for flm_gemma4_rms_residual_core.

residual + x * w / sqrt(mean(x^2) + 1e-6), rounded to bf16.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_rms_residual_core_ref(x, w, residual):
    """Numpy reference for [`flm_gemma4_rms_residual_core`][iron.kernels.flm_gemma4.flm_gemma4_rms_residual_core].

    ``residual + x * w / sqrt(mean(x^2) + 1e-6)``, rounded to bf16.
    """
    return (
        _rms_residual_core_terms(x, w, residual)[1].astype(np.float32).astype(bfloat16)
    )

flm_gemma4_rms_residual_core

flm_gemma4_rms_residual_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

rms_norm then residual_add of flm_gemma4_decode_rms_residual.

The post-attention norm and residual add. rms_residual_core(x, w, residual, y), all model_dim bf16: y = residual + rms_norm(x) * w.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_rms_residual_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``rms_norm`` then ``residual_add`` of [`flm_gemma4_decode_rms_residual`][iron.kernels.flm_gemma4.flm_gemma4_decode_rms_residual].

    The post-attention norm and residual add. ``rms_residual_core(x, w,
    residual, y)``, all ``model_dim`` bf16: ``y = residual + rms_norm(x) * w``.
    """
    base = flm_gemma4_decode_rms_residual(geometry=geometry)
    row = np.ndarray[(geometry.model_dim,), _BF16]
    return _wrapper(
        base,
        "decode_rms_residual_core.cc",
        "rms_residual_core",
        [row, row, row, row],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Out),
            reference=flm_gemma4_rms_residual_core_ref,
            tolerance=Tolerance.bounded(
                _rms_residual_core_bound,
                note="bf16 ulp <= 2**-7 |v|. The kernel floors the normalized "
                "value to bf16 (1 ulp, plus the fast inverse sqrt's and the "
                "float32 sum of squares' error, under 2**-12), floors the sum "
                "to bf16 (1 ulp, plus its float32 add's error, under 2**-12), "
                "and the reference rounds the sum to nearest (1/2 ulp). Each "
                "flush of a subnormal adds the smallest normal",
            ),
            # Per element: square-add and two multiplies, then the add.
            ops_per_call=5 * geometry.model_dim,
            stack_bytes=1792,  # aiecc measured_stack_size
        ),
    )

flm_gemma4_decode_proj_main

flm_gemma4_decode_proj_main(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

One of the decode layer's 16 q4nx projection cores, from flm_gemma4/decode_proj_main.cc.

proj_main(y_ping, w_ping, x_ping, y_pong, w_pong, x_pong, is_swa, skip_kv, send_x_output): runs this core's share of every projection in a layer (QKV, output, up/gate and down), multiplying q4nx weight blocks from the w ping-pong pair into 256-element input slices from the x pair and writing each block's 32 outputs into the y pair. y_ping and y_pong are 2 * 32 + 16 bf16: a 16-element packet header, then two 32-output slots. With send_x_output set the core fills the first slot of its own pair and sends it; with it 0 the pair is the tile below's and the core fills the second slot. w_ping and w_pong are one q4nx block each (5120 bytes: 32 by 256 weights, their bf16 scales and mins, then the 4-bit codes). is_swa and skip_kv are int32 RTP buffers of 16 words that select the projection shapes.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults x_prod_lock=0, x_cons_lock=1, w_prod_lock=2, w_cons_lock=3, y_prod_ping_lock=4, y_prod_pong_lock=5, rtp_available_lock=6, y_cons_ping_lock=7 and y_cons_pong_lock=8. When send_x_output is 0 the four y locks are the tile below's.	{}

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_proj_main(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """One of the decode layer's 16 q4nx projection cores, from ``flm_gemma4/decode_proj_main.cc``.

    ``proj_main(y_ping, w_ping, x_ping, y_pong, w_pong, x_pong, is_swa, skip_kv, send_x_output)``:
    runs this core's share of every projection in a layer (QKV, output,
    up/gate and down), multiplying q4nx weight blocks from the ``w``
    ping-pong pair into 256-element input slices from the ``x`` pair and
    writing each block's 32 outputs into the ``y`` pair. ``y_ping`` and
    ``y_pong`` are ``2 * 32 + 16`` bf16: a 16-element packet header, then
    two 32-output slots. With ``send_x_output`` set the core fills the first
    slot of its own pair and sends it; with it 0 the pair is the tile
    below's and the core fills the second slot. ``w_ping`` and ``w_pong`` are one q4nx
    block each (5120 bytes: 32 by 256 weights, their bf16 scales and mins,
    then the 4-bit codes). ``is_swa`` and ``skip_kv`` are int32 RTP buffers
    of 16 words that select the projection shapes.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``x_prod_lock=0``,
            ``x_cons_lock=1``, ``w_prod_lock=2``, ``w_cons_lock=3``,
            ``y_prod_ping_lock=4``, ``y_prod_pong_lock=5``,
            ``rtp_available_lock=6``, ``y_cons_ping_lock=7`` and
            ``y_cons_pong_lock=8``. When ``send_x_output`` is 0 the four
            ``y`` locks are the tile below's.
    """
    bf16 = np.dtype[bfloat16]
    y_ty = np.ndarray[(2 * _DECODE_Q4NX_ROWS + 16,), bf16]
    # One q4nx block, as the bf16 words the weight streams move.
    w_ty = np.ndarray[(_DECODE_Q4NX_BLOCK_BYTES // 2,), bf16]
    x_ty = np.ndarray[(_DECODE_Q4NX_COLS,), bf16]
    return _decode_kernel(
        "proj_main",
        "proj_main",
        [
            y_ty,
            w_ty,
            x_ty,
            y_ty,
            w_ty,
            x_ty,
            _DECODE_RTP,
            _DECODE_RTP,
            np.int32,
        ],
        (Out, In, In, Out, In, In, Param, Param, Param),
        dict(
            x_prod_lock=0,
            x_cons_lock=1,
            w_prod_lock=2,
            w_cons_lock=3,
            y_prod_ping_lock=4,
            y_prod_pong_lock=5,
            rtp_available_lock=6,
            y_cons_ping_lock=7,
            y_cons_pong_lock=8,
        ),
        locks,
        geometry,
    )

flm_gemma4_decode_per_layer_up

flm_gemma4_decode_per_layer_up(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the per-layer-input up projection, from flm_gemma4/decode_per_layer_up.cc.

per_layer_up(x, proj_w_ping, proj_w_pong, y): gates the layer's per-layer input, projects it from pli_d up to model_dim with bf16 weight blocks (32 by 256 each) streamed through the proj_w ping-pong pair, RMS-normalizes the result, adds the residual and scales it by the layer scale into y (model_dim bf16). x is 2 * (pli_d + model_dim) + 32 bf16: the norm weight, the layer scale padded to 32, the per-layer input, the residual and the gate. The kernel gates the per-layer input in place.

y needs a 64-byte-aligned base pointer.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults x_prod_lock=0, x_cons_lock=1, proj_w_prod_lock=2, proj_w_cons_lock=3, y_prod_lock=4 and y_cons_lock=5.	{}

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_per_layer_up(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the per-layer-input up projection, from ``flm_gemma4/decode_per_layer_up.cc``.

    ``per_layer_up(x, proj_w_ping, proj_w_pong, y)``: gates the layer's
    per-layer input, projects it from ``pli_d`` up to ``model_dim`` with bf16
    weight blocks (32 by 256 each) streamed through the ``proj_w`` ping-pong
    pair, RMS-normalizes the result, adds the residual and scales it by the
    layer scale into ``y`` (``model_dim`` bf16). ``x`` is
    ``2 * (pli_d + model_dim) + 32`` bf16: the norm weight, the layer scale
    padded to 32, the per-layer input, the residual and the gate. The kernel
    gates the per-layer input in place.

    ``y`` needs a 64-byte-aligned base pointer.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``x_prod_lock=0``,
            ``x_cons_lock=1``, ``proj_w_prod_lock=2``, ``proj_w_cons_lock=3``,
            ``y_prod_lock=4`` and ``y_cons_lock=5``.
    """
    bf16 = np.dtype[bfloat16]
    d, pli = geometry.model_dim, geometry.pli_d
    w_ty = np.ndarray[(_DECODE_BF16_BLOCK,), bf16]
    return _decode_kernel(
        "per_layer_up",
        "per_layer_up",
        [
            np.ndarray[(2 * (pli + d) + 32,), bf16],
            w_ty,
            w_ty,
            np.ndarray[(d,), bf16],
        ],
        (InOut, In, In, Out),
        dict(
            x_prod_lock=0,
            x_cons_lock=1,
            proj_w_prod_lock=2,
            proj_w_cons_lock=3,
            y_prod_lock=4,
            y_cons_lock=5,
        ),
        locks,
        geometry,
    )

flm_gemma4_bf16_proj_core_ref

flm_gemma4_bf16_proj_core_ref(w, x)

Numpy reference for flm_gemma4_bf16_proj_core: one weight block times x.

The block is stored column by column: w[k * 32 + m] multiplies x[k] into out-feature m. Sums in float64.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_bf16_proj_core_ref(w, x):
    """Numpy reference for [`flm_gemma4_bf16_proj_core`][iron.kernels.flm_gemma4.flm_gemma4_bf16_proj_core]: one weight block times ``x``.

    The block is stored column by column: ``w[k * 32 + m]`` multiplies
    ``x[k]`` into out-feature ``m``. Sums in float64.
    """
    w = np.asarray(w, np.float64).reshape(-1, _BF16_PROJ_K, _BF16_PROJ_M)
    x = np.asarray(x, np.float64).reshape(-1, _BF16_PROJ_K)
    return np.einsum("ckm,ck->cm", w, x).astype(np.float32)

flm_gemma4_bf16_proj_core

flm_gemma4_bf16_proj_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

_mvm_bf16_bf16 of the per-layer-input projections, one block.

bf16_proj_block_core(w, x, y): zeroes a float32 accumulator in its own frame, accumulates one 32 by 256 bf16 weight block times 256 inputs into it, as linear_proj does, and copies it to y. The block does not depend on geometry.

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_bf16_proj_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``_mvm_bf16_bf16`` of the per-layer-input projections, one block.

    ``bf16_proj_block_core(w, x, y)``: zeroes a float32 accumulator in its own
    frame, accumulates one 32 by 256 bf16 weight block times 256 inputs into
    it, as ``linear_proj`` does, and copies it to ``y``. The block does not
    depend on ``geometry``.
    """
    base = flm_gemma4_decode_per_layer_up(geometry=geometry)
    return _wrapper(
        base,
        "decode_per_layer_up_core.cc",
        "bf16_proj_block_core",
        [
            np.ndarray[(_DECODE_BF16_BLOCK,), _BF16],
            np.ndarray[(_BF16_PROJ_K,), _BF16],
            np.ndarray[(_BF16_PROJ_M,), _F32],
        ],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            reference=flm_gemma4_bf16_proj_core_ref,
            sample=_bf16_proj_core_sample,
            tolerance=Tolerance.exact(
                note="the sample keeps every product and sum exact in float32"
            ),
            ops_per_call=2 * _DECODE_BF16_BLOCK,
            acc_dtype=np.float32,
            reduction=_BF16_PROJ_K,
        ),
    )

flm_gemma4_decode_proj_layer_embedding

flm_gemma4_decode_proj_layer_embedding(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the per-layer-input embedding projection, from flm_gemma4/decode_proj_layer_embedding.cc.

proj_layer_embedding(norm_w, x0_per_layer, x0, x_proj, y, proj_w_ping, proj_w_pong): projects the layer input x0 (model_dim bf16) down to pli_d with bf16 weight blocks (32 by 256 each) streamed through the proj_w ping-pong pair, scales, RMS-normalizes and adds the token's per-layer embedding x0_per_layer (pli_d), and scales again, in the scratch x_proj (pli_d). norm_w (pli_d + model_dim + 32 bf16) holds the norm weight, then model_dim + 32 values the kernel copies to the start of y (same size) for the up projection; the result follows them.

x_proj needs a 64-byte-aligned base pointer.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults norm_w_prod_lock=0, x0_per_layer_prod_lock=1, x0_prod_lock=2, xw_cons_lock=3, proj_w_prod_lock=4, proj_w_cons_lock=5, y_prod_lock=6 and y_cons_lock=7.	{}

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_proj_layer_embedding(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the per-layer-input embedding projection, from ``flm_gemma4/decode_proj_layer_embedding.cc``.

    ``proj_layer_embedding(norm_w, x0_per_layer, x0, x_proj, y, proj_w_ping, proj_w_pong)``:
    projects the layer input ``x0`` (``model_dim`` bf16) down to ``pli_d``
    with bf16 weight blocks (32 by 256 each) streamed through the ``proj_w``
    ping-pong pair, scales, RMS-normalizes and adds the token's per-layer
    embedding ``x0_per_layer`` (``pli_d``), and scales again, in the scratch
    ``x_proj`` (``pli_d``). ``norm_w`` (``pli_d + model_dim + 32`` bf16)
    holds the norm weight, then ``model_dim + 32`` values the kernel copies
    to the start of ``y`` (same size) for the up projection; the result
    follows them.

    ``x_proj`` needs a 64-byte-aligned base pointer.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``norm_w_prod_lock=0``,
            ``x0_per_layer_prod_lock=1``, ``x0_prod_lock=2``,
            ``xw_cons_lock=3``, ``proj_w_prod_lock=4``,
            ``proj_w_cons_lock=5``, ``y_prod_lock=6`` and ``y_cons_lock=7``.
    """
    bf16 = np.dtype[bfloat16]
    d, pli = geometry.model_dim, geometry.pli_d
    pli_ty = np.ndarray[(pli,), bf16]
    y_ty = np.ndarray[(pli + d + 32,), bf16]
    w_ty = np.ndarray[(_DECODE_BF16_BLOCK,), bf16]
    return _decode_kernel(
        "proj_layer_embedding",
        "proj_layer_embedding",
        [
            y_ty,
            pli_ty,
            np.ndarray[(d,), bf16],
            pli_ty,
            y_ty,
            w_ty,
            w_ty,
        ],
        (In, In, In, Out, Out, In, In),
        dict(
            norm_w_prod_lock=0,
            x0_per_layer_prod_lock=1,
            x0_prod_lock=2,
            xw_cons_lock=3,
            proj_w_prod_lock=4,
            proj_w_cons_lock=5,
            y_prod_lock=6,
            y_cons_lock=7,
        ),
        locks,
        geometry,
    )

flm_gemma4_decode_gate_layer_embedding

flm_gemma4_decode_gate_layer_embedding(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the per-layer-input gate, from flm_gemma4/decode_gate_layer_embedding.cc.

gate_layer_embedding(x, proj_w_ping, proj_w_pong, y): copies the layer output x (model_dim bf16, a buffer of the tile to the left) into y, projects it to pli_d with bf16 weight blocks (32 by 256 each) streamed through the proj_w ping-pong pair, and applies the activation into y + model_dim. y is model_dim + pli_d bf16.

y needs a 64-byte-aligned base pointer; the kernel also stores 512 bits at a time at y + model_dim.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults x_prod_lock=0, x_cons_lock=1, proj_w_prod_lock=2, proj_w_cons_lock=3, y_prod_lock=4, y_cons_lock=5, final_x_prod_lock=7 and final_x_cons_lock=8 (the last two locks of the tile to the left).	{}

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_gate_layer_embedding(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the per-layer-input gate, from ``flm_gemma4/decode_gate_layer_embedding.cc``.

    ``gate_layer_embedding(x, proj_w_ping, proj_w_pong, y)``: copies the
    layer output ``x`` (``model_dim`` bf16, a buffer of the tile to the left)
    into ``y``, projects it to ``pli_d`` with bf16 weight blocks (32 by 256
    each) streamed through the ``proj_w`` ping-pong pair, and applies the
    activation into ``y + model_dim``. ``y`` is ``model_dim + pli_d`` bf16.

    ``y`` needs a 64-byte-aligned base pointer; the kernel also stores
    512 bits at a time at ``y + model_dim``.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``x_prod_lock=0``,
            ``x_cons_lock=1``, ``proj_w_prod_lock=2``, ``proj_w_cons_lock=3``,
            ``y_prod_lock=4``, ``y_cons_lock=5``, ``final_x_prod_lock=7`` and
            ``final_x_cons_lock=8`` (the last two locks of the tile to the
            left).
    """
    bf16 = np.dtype[bfloat16]
    d = geometry.model_dim
    w_ty = np.ndarray[(_DECODE_BF16_BLOCK,), bf16]
    return _decode_kernel(
        "gate_layer_embedding",
        "gate_layer_embedding",
        [
            np.ndarray[(d,), bf16],
            w_ty,
            w_ty,
            np.ndarray[(d + geometry.pli_d,), bf16],
        ],
        (In, In, In, Out),
        dict(
            x_prod_lock=0,
            x_cons_lock=1,
            proj_w_prod_lock=2,
            proj_w_cons_lock=3,
            y_prod_lock=4,
            y_cons_lock=5,
            final_x_prod_lock=7,
            final_x_cons_lock=8,
        ),
        locks,
        geometry,
        lut=True,
    )

flm_gemma4_pli_gelu_core_ref

flm_gemma4_pli_gelu_core_ref(x)

Numpy reference for flm_gemma4_pli_gelu_core: tanh-GELU.

The tolerance covers the kernel's table.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_pli_gelu_core_ref(x):
    """Numpy reference for [`flm_gemma4_pli_gelu_core`][iron.kernels.flm_gemma4.flm_gemma4_pli_gelu_core]: tanh-GELU.

    The tolerance covers the kernel's table.
    """
    return _gelu_tanh(x).astype(np.float32)

flm_gemma4_pli_gelu_core

flm_gemma4_pli_gelu_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

_activate of flm_gemma4_decode_gate_layer_embedding: GELU over pli_d values.

pli_gelu_core(x, y): copies x into y and applies GELU there, through aie_runtime_lib's table (getGeluBf16). Both are pli_d bf16.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

Source code in python/iron/kernels/flm_gemma4.py

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_pli_gelu_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``_activate`` of [`flm_gemma4_decode_gate_layer_embedding`][iron.kernels.flm_gemma4.flm_gemma4_decode_gate_layer_embedding]: GELU over ``pli_d`` values.

    ``pli_gelu_core(x, y)``: copies ``x`` into ``y`` and applies GELU there,
    through aie_runtime_lib's table (``getGeluBf16``). Both are ``pli_d``
    bf16.

    Args:
        geometry: The model the kernel builds for.
    """
    n = geometry.pli_d
    return _wrapper(
        flm_gemma4_decode_gate_layer_embedding(geometry=geometry),
        "decode_gate_layer_embedding_core.cc",
        "pli_gelu_core",
        [np.ndarray[(n,), _BF16], np.ndarray[(n,), _BF16]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_pli_gelu_core_ref,
            sample=partial(_pli_gelu_core_sample, n=n),
            tolerance=Tolerance.bounded(
                _pli_gelu_core_bound,
                note="the GELU table's error against tanh-GELU, measured on the "
                "host from aie_runtime_lib's table, plus a bf16 ulp of its slope "
                "and the floor bf16 store",
            ),
            uses_lut=True,
            # Per value: the table line's multiply and add.
            ops_per_call=2 * n,
        ),
    )

flm_gemma4_q4nx_lm_head_ref

flm_gemma4_q4nx_lm_head_ref(
    w, x, sums, *, m_tile=32, k_tile=256, group=32
)

Numpy reference for flm_gemma4_q4nx_lm_head: one q4nx_lm_head_block call.

w is one q4nx block per call, as bytes or bf16 words: bf16 scales, then bf16 minima, both indexed [k // group, m], then 4-bit codes, low nibble first, indexed [m // 16, k // 32, (k % 32) // 8, k % 8, m % 16]. x is the token then its RMS weight, and sums the token's per-32 column sums, both shared by every call; slice 0 of the token is the one the call reads. Returns the float32 accumulator, zero before the call, holding sum_k (min + scale * code) * x with the minima folded in through sums.

Per 16 rows and 32 columns the kernel sums the codes times x in float32, narrows that dot product to bf16 with the floor rounding a fresh core uses, then accumulates it times the scale, and the minimum times the column sum, into float32. Every product is exact, so each multiply-accumulate rounds once.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_ref(w, x, sums, *, m_tile=32, k_tile=256, group=32):
    """Numpy reference for [`flm_gemma4_q4nx_lm_head`][iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head]: one ``q4nx_lm_head_block`` call.

    ``w`` is one q4nx block per call, as bytes or bf16 words: bf16 scales,
    then bf16 minima, both indexed ``[k // group, m]``, then 4-bit codes, low
    nibble first, indexed ``[m // 16, k // 32, (k % 32) // 8, k % 8, m % 16]``.
    ``x`` is the token then its RMS weight, and ``sums`` the token's per-32
    column sums, both shared by every call; slice 0 of the token is the one
    the call reads. Returns the float32
    accumulator, zero before the call, holding ``sum_k (min + scale * code) *
    x`` with the minima folded in through ``sums``.

    Per 16 rows and 32 columns the kernel sums the codes times ``x`` in
    float32, narrows that dot product to bf16 with the floor rounding a fresh
    core uses, then accumulates it times the scale, and the minimum times the
    column sum, into float32. Every product is exact, so each
    multiply-accumulate rounds once.
    """
    w = np.ascontiguousarray(w).view(np.uint8).reshape(-1, m_tile * k_tile * 5 // 8)
    calls, n_groups = len(w), k_tile // group
    params = np.ascontiguousarray(w[:, : 4 * m_tile * n_groups]).view("<u2")
    params = (params.astype(np.uint32) << 16).view(np.float32)
    scales, mins = params.reshape(calls, 2, n_groups, m_tile).transpose(1, 0, 2, 3)
    packed = w[:, 4 * m_tile * n_groups :]
    codes = np.empty((calls, m_tile * k_tile), np.float64)
    codes[:, 0::2], codes[:, 1::2] = packed & 15, packed >> 4
    codes = codes.reshape(calls, m_tile // 16, n_groups, 32, 16)
    x = np.asarray(x, np.float32).reshape(-1)[:k_tile]
    x = x.astype(np.float64).reshape(1, n_groups, 32)
    sums = np.asarray(sums, np.float32).reshape(-1)[:n_groups]
    sums = sums.astype(np.float64).reshape(1, n_groups)
    acc = np.zeros((calls, m_tile // 16, 16), np.float32)
    for g in range(n_groups):
        dot = np.zeros((calls, m_tile // 16, 16), np.float32)
        for c in range(32):
            dot = _f32(dot + codes[:, :, g, c, :] * x[:, g, c, None, None])
        dot = _bf16_floor(dot).astype(np.float64)
        scale = scales[:, g].reshape(calls, m_tile // 16, 16).astype(np.float64)
        low = mins[:, g].reshape(calls, m_tile // 16, 16).astype(np.float64)
        acc = _f32(acc + dot * scale)
        acc = _f32(acc + low * sums[:, g, None, None])
    return acc.reshape(calls, m_tile)

flm_gemma4_q4nx_lm_head

flm_gemma4_q4nx_lm_head(
    *,
    dim: int = 1536,
    m_tile: int = 32,
    k_tile: int = 256,
    group: int = 32
) -> ExternalFunction

AIE2P logits from a q4nx vocabulary, from flm_gemma4/q4nx_lm_head.cc.

One core owns m_tile out-features at a time and streams their q4nx blocks past a token of dim values. The returned kernel is q4nx_lm_head_block, which accumulates one m_tile by k_tile block times one slice of the token into float32; its slice index is bound to 0. See flm_gemma4_q4nx_lm_head_ref for the block layout.

The other entry points, attributes of the kernel:

• q4nx_lm_head_rms(x, sums): once per token; normalizes the token in x (the token, then its RMS weight) in place and writes its per-32 column sums.
• q4nx_lm_head_zero(y_acc): before a tile's k loop.
• q4nx_lm_head_epilogue(y, y_acc, softcap): after the k loop, y = c * tanh(y_acc / c), with the float32 c in RTP word 0.

Parameters:

Name	Type	Description	Default
dim	int	The token's length, a multiple of k_tile; sets the RMS norm's span.	1536
m_tile	int	Out-features per block, a multiple of 16.	32
k_tile	int	In-features per block, a multiple of group.	256
group	int	In-features per scale and minimum; must be 32.	32

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head(
    *, dim: int = 1536, m_tile: int = 32, k_tile: int = 256, group: int = 32
) -> ExternalFunction:
    """AIE2P logits from a q4nx vocabulary, from ``flm_gemma4/q4nx_lm_head.cc``.

    One core owns ``m_tile`` out-features at a time and streams their q4nx
    blocks past a token of ``dim`` values. The returned kernel is
    ``q4nx_lm_head_block``, which accumulates one ``m_tile`` by ``k_tile``
    block times one slice of the token into float32; its slice index is bound
    to 0. See ``flm_gemma4_q4nx_lm_head_ref`` for the block layout.

    The other entry points, attributes of the kernel:

    - ``q4nx_lm_head_rms(x, sums)``: once per token; normalizes the token in
      ``x`` (the token, then its RMS weight) in place and writes its per-32
      column sums.
    - ``q4nx_lm_head_zero(y_acc)``: before a tile's k loop.
    - ``q4nx_lm_head_epilogue(y, y_acc, softcap)``: after the k loop,
      ``y = c * tanh(y_acc / c)``, with the float32 ``c`` in RTP word 0.

    Args:
        dim: The token's length, a multiple of ``k_tile``; sets the RMS
            norm's span.
        m_tile: Out-features per block, a multiple of 16.
        k_tile: In-features per block, a multiple of ``group``.
        group: In-features per scale and minimum; must be 32.
    """
    _lm_head_geometry(dim, m_tile, k_tile, group)
    if not _arch_traits().bfp16:
        raise NotImplementedError(
            "flm_gemma4_q4nx_lm_head() is only available on aie2p."
        )
    geometry = dict(m_tile=m_tile, k_tile=k_tile, group=group)
    x = np.ndarray[(2, dim), _BF16]
    sums = np.ndarray[(dim // group,), _BF16]
    y_acc = np.ndarray[(m_tile,), _F32]
    fn = _make_extern(
        "q4nx_lm_head_block",
        _kernel_source("flm_gemma4/q4nx_lm_head.cc"),
        [np.ndarray[(m_tile * k_tile * 5 // 8 // 2,), _BF16], x, y_acc, sums, np.int32],
        compile_flags=[
            f"-DQ4NX_M_TILE={m_tile}",
            f"-DQ4NX_K_TILE={k_tile}",
            f"-DQ4NX_GROUP={group}",
            f"-DFLM_GEMMA4_LM_HEAD_DIM={dim}",
            # Without it the bf16 mmul emulation runs about 8x slower.
            "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16",
        ],
        cls=_LmHeadKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            # The token and its sums are fixed while the blocks stream past.
            roles=(In, Param, InOut, Param, Param),
            parameter_bindings=((4, 0),),
            initializers=((2, _lm_head_zero),),
            reference=partial(flm_gemma4_q4nx_lm_head_ref, **geometry),
            sample=partial(_q4nx_lm_head_sample, dim=dim, **geometry),
            tolerance=Tolerance.exact(
                note="the sample keeps every sum exact in float32; the floor "
                "bf16 narrowing of each dot product is modeled"
            ),
            ops_per_call=2 * m_tile * k_tile,
            acc_dtype=np.float32,
            reduction=k_tile,
        ),
    )
    bind = fn.object_file.bind
    fn.q4nx_lm_head_block = fn
    fn.q4nx_lm_head_rms = bind("q4nx_lm_head_rms", [x, sums])
    fn.q4nx_lm_head_zero = bind("q4nx_lm_head_zero", [y_acc])
    rtp = np.ndarray[(_LM_HEAD_RTP_WORDS,), _RTP_WORD]
    fn.q4nx_lm_head_epilogue = bind(
        "q4nx_lm_head_epilogue", [np.ndarray[(m_tile,), _BF16], y_acc, rtp]
    )
    return fn

flm_gemma4_q4nx_lm_head_epilogue_ref

flm_gemma4_q4nx_lm_head_epilogue_ref(y_acc, rtp)

Numpy reference for flm_gemma4_q4nx_lm_head_epilogue: c * tanh(y_acc / c).

rtp holds the float32 bits of c in word 0. The reference computes in float64 and rounds to bf16 once; the tolerance covers the kernel's narrowings and its tanh.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_epilogue_ref(y_acc, rtp):
    """Numpy reference for [`flm_gemma4_q4nx_lm_head_epilogue`][iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head_epilogue]: ``c * tanh(y_acc / c)``.

    ``rtp`` holds the float32 bits of ``c`` in word 0. The reference computes
    in float64 and rounds to bf16 once; the tolerance covers the kernel's
    narrowings and its tanh.
    """
    c = _lm_head_softcap(rtp)
    y = np.asarray(y_acc, np.float64)
    return (c * np.tanh(y / c)).astype(bfloat16)

flm_gemma4_q4nx_lm_head_epilogue

flm_gemma4_q4nx_lm_head_epilogue(
    *,
    dim: int = 1536,
    m_tile: int = 32,
    k_tile: int = 256,
    group: int = 32
) -> ExternalFunction

q4nx_lm_head_epilogue of the LM head: y = c * tanh(y_acc / c).

The softcap c is float32 bits in word 0 of a 32-word RTP buffer. The tolerance assumes c is a bf16 value, so the kernel's narrowing of it is exact; Gemma 4's 30 is.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_epilogue(
    *, dim: int = 1536, m_tile: int = 32, k_tile: int = 256, group: int = 32
) -> ExternalFunction:
    """``q4nx_lm_head_epilogue`` of the LM head: ``y = c * tanh(y_acc / c)``.

    The softcap ``c`` is float32 bits in word 0 of a 32-word RTP buffer. The
    tolerance assumes ``c`` is a bf16 value, so the kernel's narrowing of it
    is exact; Gemma 4's 30 is.
    """
    return _lm_head_sibling(
        "q4nx_lm_head_epilogue",
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, In, Param),
            reference=flm_gemma4_q4nx_lm_head_epilogue_ref,
            sample=partial(_lm_head_epilogue_sample, m_tile=m_tile),
            tolerance=Tolerance.bounded(
                _lm_head_epilogue_bound,
                note="c times vtanh's error, measured on npu2 (see "
                "activation._vtanh_error), plus the floor bf16 roundings",
            ),
        ),
        dim=dim,
        m_tile=m_tile,
        k_tile=k_tile,
        group=group,
    )

flm_gemma4_q4nx_lm_head_rms_ref

flm_gemma4_q4nx_lm_head_rms_ref(x, *, dim=1536)

Numpy reference for flm_gemma4_q4nx_lm_head_rms: column sums of the normalized token.

x holds, per call, the token then its RMS weight. The reference narrows each normalized value to bf16 and sums each 32 of them in float64.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_rms_ref(x, *, dim=1536):
    """Numpy reference for [`flm_gemma4_q4nx_lm_head_rms`][iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head_rms]: column sums of the normalized token.

    ``x`` holds, per call, the token then its RMS weight. The reference
    narrows each normalized value to bf16 and sums each 32 of them in float64.
    """
    y = _lm_head_rms_y(x, dim).astype(np.float32).astype(bfloat16)
    sums = y.astype(np.float64).reshape(len(y), dim // 32, 32).sum(axis=2)
    return sums.astype(np.float32)

flm_gemma4_q4nx_lm_head_rms

flm_gemma4_q4nx_lm_head_rms(
    *,
    dim: int = 1536,
    m_tile: int = 32,
    k_tile: int = 256,
    group: int = 32
) -> ExternalFunction

q4nx_lm_head_rms of flm_gemma4_q4nx_lm_head: the token's RMS norm and per-32 column sums.

The kernel normalizes the token in place. The contract declares the token In: the write lands in the core's input element, and the next DMA fill overwrites it. The column sums depend on every normalized value.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_rms(
    *, dim: int = 1536, m_tile: int = 32, k_tile: int = 256, group: int = 32
) -> ExternalFunction:
    """``q4nx_lm_head_rms`` of [`flm_gemma4_q4nx_lm_head`][iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head]: the token's RMS norm and per-32 column sums.

    The kernel normalizes the token in place. The contract declares the token
    ``In``: the write lands in the core's input element, and the next DMA fill
    overwrites it. The column sums depend on every normalized value.
    """
    return _lm_head_sibling(
        "q4nx_lm_head_rms",
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=partial(flm_gemma4_q4nx_lm_head_rms_ref, dim=dim),
            tolerance=Tolerance.bounded(
                partial(_lm_head_rms_bound, dim=dim),
                note="bf16 ulp <= 2**-7 |v|. Per term: the kernel floors to bf16 "
                "(1 ulp), the reference rounds to nearest (1/2 ulp), and the fast "
                "inverse sqrt after two Newton steps adds 5e-6 relative. The "
                "kernel floors the sum to bf16 (1 ulp of the sum). 2.5 ulps of "
                "the sum of |y| plus that 5e-6 stays under 3 * 2**-7 of it",
            ),
            # Per element: square and add, two multiplies, one column-sum add.
            ops_per_call=5 * dim,
        ),
        dim=dim,
        m_tile=m_tile,
        k_tile=k_tile,
        group=group,
    )

flm_gemma4_attn_prefill_ref

flm_gemma4_attn_prefill_ref(l_bf16, y)

Numpy reference for flm_gemma4_attn_prefill: attn_epilogue at chunk 0.

y is the round's (8, 512) float32 accumulator, flat; chunk 0 is its first 64 values, eight per row of l_bf16. The kernel narrows y to bf16, multiplies by the row's 1 / l exactly and narrows again, both times with the floor rounding a fresh core uses.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_prefill_ref(l_bf16, y):
    """Numpy reference for [`flm_gemma4_attn_prefill`][iron.kernels.flm_gemma4.flm_gemma4_attn_prefill]: ``attn_epilogue`` at chunk 0.

    ``y`` is the round's (8, 512) float32 accumulator, flat; chunk 0 is its
    first 64 values, eight per row of ``l_bf16``. The kernel narrows ``y`` to
    bf16, multiplies by the row's ``1 / l`` exactly and narrows again, both
    times with the floor rounding a fresh core uses.
    """
    return _prefill_epilogue(
        l_bf16, y, lq=_ATTN_PREFILL_LQ, dh=_ATTN_PREFILL_DH, chunk=0
    )

flm_gemma4_swa_prefill_ref

flm_gemma4_swa_prefill_ref(l_bf16, y)

Numpy reference for flm_gemma4_swa_prefill: attn_epilogue at chunk 33.

As flm_gemma4_attn_prefill_ref, over a (16, 256) y: rows 8 to 15, columns 64 to 127, each row scaled by its own 1 / l.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_swa_prefill_ref(l_bf16, y):
    """Numpy reference for [`flm_gemma4_swa_prefill`][iron.kernels.flm_gemma4.flm_gemma4_swa_prefill]: ``attn_epilogue`` at chunk 33.

    As [`flm_gemma4_attn_prefill_ref`][iron.kernels.flm_gemma4.flm_gemma4_attn_prefill_ref], over a
    (16, 256) ``y``: rows 8 to 15, columns 64 to 127, each row scaled by its
    own ``1 / l``.
    """
    return _prefill_epilogue(
        l_bf16, y, lq=_SWA_PREFILL_LQ, dh=_SWA_PREFILL_DH, chunk=_SWA_PREFILL_CHUNK
    )