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(),
            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: 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.

axpy_ref

axpy_ref(x, y, a)

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

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."""
    return np.float32(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)

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)"
            ),
        ),
    )

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(),
            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(),
            # aiecc measured_stack_size (Peano 22). A group size that is not a
            # power of two makes the `/ GROUP` in the inner loop call __muldi3,
            # which needs 64 bytes more than the power-of-two geometries.
            stack_bytes=1280,
            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.

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."""
    return x.astype(np.int64).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.

Parameters:

Name	Type	Description	Default
tile_size	int	Number of elements in the input tile.	1024
dtype	type	Element data type (only np.int32 supported).	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.

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

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.

    Args:
        tile_size: Number of elements in the input tile.
        dtype: Element data type (only ``np.int32`` supported).
        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``.
    """
    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.

Parameters:

Name	Type	Description	Default
tile_size	int	Number of elements in the input tile.	1024
dtype	type	Element data type (only np.int32 supported).	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.

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

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

    Args:
        tile_size: Number of elements in the input tile.
        dtype: Element data type (only ``np.int32`` supported).
        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``.
    """
    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 dims_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 dims_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 dims_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 ``dims_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
    # Walking an operand as (r x s) blocks is what TensorTiler2D generates, so
    # A and B ask for it rather than restating it.
    a = _blocked(m, k, r, s)
    b = _blocked(n, k, t, s) if b_col_maj else _blocked(k, n, s, t)
    # C is not expressible that way. The DMA reads a core-blocked buffer and
    # writes a differently ordered stream, so the intra-tile row term comes
    # *outside* the tile index -- (r, t) before (n//t, r*t). Every
    # TensorTiler2D classmethod iterates tiles outermost and elements within
    # them, and no combination of tile_col_major / iter_col_major /
    # prune_step produces this order. Closing the gap needs an un-blocking
    # tiler in taplib, which is its own change.
    if c_col_maj:
        c = [(n // t, t * m), (t, r), (m // r, r * t), (r, 1)]
    else:
        c = [(m // r, r * n), (r, t), (n // t, r * t), (t, 1)]
    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 dims_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 dims_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 ``dims_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 ``dims_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,
            stack_bytes=0xD00,  # programming_examples/basic/matrix_multiplication
            # 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,
) -> 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.

Parameters:

Name	Type	Description	Default
dim_m	int	Number of rows of A (output vector length).	32
dim_k	int	Number of columns of A (input vector length).	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.	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

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,
) -> 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.

    Args:
        dim_m: Number of rows of A (output vector length).
        dim_k: Number of columns of A (input vector length).
        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.
        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.

    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)
    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 dims_from_stream on the hop into the core, reading it
    # from the layout (programming_examples/basic/matrix_multiplication/
    # matrix_vector does).
    a_dims_from_stream = (
        [(dim_m, 2), (dim_k // 2, 2 * dim_m), (2, 1)] 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_dims_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,
            stack_bytes=0xF00,  # programming_examples/ml/block_datatypes
            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.cc and mm.cc 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.cc'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.

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.cc`` and ``mm.cc`` 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.cc``'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.

    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.cc without C_COL_MAJ, and without B_COL_MAJ unless
    # b_col_maj. matmul_bf16_bf16_rowmaj is always row-major and expands
    # aie::mmul<8, 8, 8, bf16, bf16> directly, not the micro-tile
    # _MM_MAC_DIMS records for mm.cc.
    b_col_maj = b_col_maj and not pv
    if pv:
        r, s, t = (8, 8, 8)
    elif emulate_bf16_mmul_with_bfp16:
        r, s, t = _MM_EMULATED_BF16_MAC_DIMS_AIE2P[(bfloat16, bfloat16)]
    else:
        r, s, t = _MM_MAC_DIMS["aie2p"][(bfloat16, bfloat16)]
    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,)),
            ),
            stack_bytes=0xD00,  # mm.cc's product: programming_examples/basic/matrix_multiplication
            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),
            # aiecc measured_stack_size of 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() is None 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 comes off _blocked rather than 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.
    k_blocks, n_blocks, *within = _blocked(lk, head_dim, s, t)
    v_dims = [n_blocks, k_blocks, *within]
    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,  # aiecc measured_stack_size
            # 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). The partial sum crosses the cascade as a 32-bit integer lane: with a floating-point output type the PUT half's product is truncated toward zero.

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``). The partial sum
    crosses the cascade as a 32-bit integer lane: with a floating-point
    output type the PUT half's product is truncated toward zero.

    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),
            # Scalar on both targets: row-major operands, 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,
        ),
    )

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=(
                # aiecc measured_stack_size at 17 taps
                1888
                if _detect_arch() == "aie2" and _tuned_arch() is None
                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(
            trace=Trace.whole_call(),
            stack_bytes=1280,  # aiecc measured_stack_size
            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(
            trace=Trace.whole_call(),
            # aiecc measured_stack_size of the untuned code on aie2p
            # (1504 B on aie2); 288 B tuned for aie2
            stack_bytes=None if _tuned_arch() == "aie2" else 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.

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``.

    """
    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(
            trace=Trace.whole_call(),
            # 0 B tuned for aie2; see conv2dk1 for the other figure's source
            stack_bytes=None if _tuned_arch() == "aie2" else 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(
            trace=Trace.whole_call(),
            # With an int8 skip; 32 B tuned for aie2. Measured as conv2dk1's
            stack_bytes=None if _tuned_arch() == "aie2" else 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(),
            stack_bytes=1504,  # aiecc measured_stack_size
            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.	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.	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.
        input_channels: Number of input channels.
        output_channels: Number of output channels.
        kernel_width: Width (and height) of the convolution kernel.

    Returns:
        ExternalFunction configured for the conv2dk14 kernel.
    """
    tiles = input_width // 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(
            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,
        contract=KernelContract(
            trace=Trace.whole_call(),
            # aie2p: >=2144 measured; __modsi3 has no .stack_sizes. 288 B tuned
            # for aie2
            stack_bytes=None if _tuned_arch() == "aie2" else 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))

Both shifts must be at least 1.

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))
    ```

    Both shifts must be at least 1.
    """
    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; output_channels_padd must not exceed output_channels.

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; ``output_channels_padd`` must not
    exceed ``output_channels``.
    """
    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, OC), dtype=np.uint16)
    start = tile * int(weight_index)
    out[..., start : start + tile] = total.astype(np.uint16)
    del output_channels_padd
    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"],
        contract=KernelContract(
            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)],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Param, Out, *((Param,) * 8)),
            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"],
        contract=KernelContract(
            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,
        contract=KernelContract(
            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}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Param, Out, *((Param,) * 8)),
            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(
            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
) -> 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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the PUT tile.

Raises:

Type	Description
ValueError	When block_index is not 13 or 14.

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,
) -> 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.

    Returns:
        ExternalFunction configured for the PUT tile.

    Raises:
        ValueError: When ``block_index`` is not 13 or 14.
    """
    _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"],
    )

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
) -> 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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the GET tile.

Raises:

Type	Description
ValueError	When block_index is not 13 or 14.

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,
) -> 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.

    Returns:
        ExternalFunction configured for the GET tile.

    Raises:
        ValueError: When ``block_index`` is not 13 or 14.
    """
    _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"],
    )

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"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Param, Out, Out, *((Param,) * 8)),
            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
) -> 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

Returns:

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

Raises:

Type	Description
ValueError	When block_index is not 13 or 14.

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,
) -> 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.

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

    Raises:
        ValueError: When ``block_index`` is not 13 or 14.
    """
    _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"
        ],
    )

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
) -> 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

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.

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,
) -> 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.

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

    Raises:
        ValueError: When ``block_index`` is not 13 or 14.
    """
    _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"
        ],
    )

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(
            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_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_lut_ref if use_lut_model else swiglu_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 and integer exponent reconstruction replace the hardware exp2 approximation, preserving subnormal outputs. 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 and integer exponent reconstruction replace
    the hardware exp2 approximation, preserving subnormal outputs. 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 (exp2_poly.h), 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,
            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 degree-5 minimax poly, not a LUT.

A float32-output alternative to bf16_exp, sharing its AIE2P range-reduced polynomial but with a separately configurable input domain. See aie_kernels/activation/exp2f_vec.cc for the accuracy rationale and the noinline codegen hazard this kernel carries.

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. The default -111 is the lowest exponent that still holds the kernel's 8.9e-5 relative error; -126 is the hard floor (one f32 exponent field), reachable at up to 6.5e-3. See aie_kernels/activation/exp2f_vec.cc for the measured table.	-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 degree-5 minimax poly, not a LUT.

    A float32-output alternative to [`bf16_exp`]
    [iron.kernels.activation.bf16_exp], sharing its AIE2P range-reduced
    polynomial but with a separately configurable input domain. See
    ``aie_kernels/activation/exp2f_vec.cc`` for the accuracy rationale and the
    ``noinline`` codegen hazard this kernel carries.

    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. The default
            -111 is the lowest exponent that still holds the kernel's
            8.9e-5 relative error; -126 is the hard floor (one f32
            exponent field), reachable at up to 6.5e-3. See
            ``aie_kernels/activation/exp2f_vec.cc`` for the measured table.

    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(),
            setup=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="minimax poly targets 8.9e-5 relative error; see exp2f_vec_ref",
            ),
            stack_bytes=2048 if _detect_arch() == "aie2" else None,
        ),
    )

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_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 = x.astype(np.float32)
    return (xf / (1.0 + np.exp(-xf))).astype(x.dtype)

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 = x.astype(np.float64)
    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 built with use_lut=True.

Follows activation/sigmoid.cc 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] built with ``use_lut=True``.

    Follows activation/sigmoid.cc 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)

silu_lut_ref

silu_lut_ref(x)

Model of silu built with use_lut=True.

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.

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

def silu_lut_ref(x):
    """Model of [`silu`][iron.kernels.activation.silu] built with ``use_lut=True``.

    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.
    """
    xf = np.asarray(x).astype(np.float32)
    sig = np.asarray(sigmoid_lut_ref(xf), np.float32)
    return _bf16(np.maximum(xf, -8.0) * sig).astype(np.asarray(x).dtype)

swiglu_lut_ref

swiglu_lut_ref(x, w1, w2)

Model of swiglu built with use_lut=True.

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.

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

def swiglu_lut_ref(x, w1, w2):
    """Model of [`swiglu`][iron.kernels.activation.swiglu] built with ``use_lut=True``.

    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``.
    """
    with np.errstate(over="ignore", invalid="ignore"):
        xw1 = _bf16(np.asarray(x, np.float32) * np.asarray(w1, np.float32))
        xw2 = _bf16(np.asarray(x, np.float32) * np.asarray(w2, np.float32))
        sig = np.asarray(sigmoid_lut_ref(xw2), np.float32)
        silu_out = _bf16(np.maximum(xw2, -8.0) * sig)
        out = np.where(silu_out == 0, np.float32(0.0), _bf16(xw1 * silu_out))
    return out.astype(np.asarray(x).dtype)

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 np.tanh(x.astype(np.float32)).astype(x.dtype)

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.
    """
    xf = x.astype(np.float32)
    return (1.0 / (1.0 + np.exp(-xf))).astype(x.dtype)

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. 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.  Exact up to bf16 rounding; pair with a
    small ``rtol`` when verifying.
    """
    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 float32; LUT-approximation territory, pair with rtol=0.128.

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
    float32; LUT-approximation territory, pair with ``rtol=0.128``.
    """
    xf = x.astype(np.float32)
    xw1 = (xf * w1.astype(np.float32)).astype(bfloat16).astype(np.float32)
    xw2 = (xf * w2.astype(np.float32)).astype(bfloat16).astype(np.float32)
    return (xw1 * (xw2 / (1.0 + np.exp(-xw2)))).astype(x.dtype)

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 reference itself in range: exp(88) = 1.65e+38 fits 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 reference itself in range:
    ``exp(88) = 1.65e+38`` fits float32 where ``exp(89)`` would not.
    """
    xf = np.clip(x.astype(np.float32), -_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 np.exp(xf).astype(x.dtype)

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 targets ~8.9e-5 relative error by design, 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 targets ~8.9e-5
    relative error by design, 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 float32 softmax. 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 float32 softmax.  ``x.size`` must be a
    multiple of ``tile_size``.
    """
    xf = x.astype(np.float32)
    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 out.reshape(x.shape).astype(x.dtype)

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=_RMS_NORM_BF16_AIE2 if _tuned_arch() == "aie2" else _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=_RMS_NORM_BF16_AIE2 if _tuned_arch() == "aie2" else _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 on aie2 or 32 on aie2p (the source processes whole vectors, 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 on aie2 or 32 on
    aie2p (the source processes whole vectors, without a scalar tail).
    """
    tile_size = _row_size("layer_norm", tile_size, cols, _arch_traits().bf16_lanes)
    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=(
                _LAYER_NORM_BF16_AIE2 if _tuned_arch() == "aie2" else _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 1e-3 (2e-6 on aie2) 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 1e-3
    (2e-6 on aie2) 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_AIE2 if _tuned_arch() == "aie2" else _NORM_F32,
        6 * cols,
        # Headroom: the frame measures 256 bytes on AIE2P and 64 on AIE2, and
        # neither build calls a soft-float helper.
        stack_bytes=2048,
    )

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=(
                _LAYER_NORM_BF16_AIE2 if _tuned_arch() == "aie2" else _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:
        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()}"]
    if not use_chess and _tuned_arch() == "aie2p":
        # LICM hoists the three accumulator constants out of the loop, where
        # they spill. Capping its MemorySSA walk at zero keeps them in the
        # loop.
        flags += ["-mllvm", "--licm-mssa-optimization-cap=0"]
    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."""
    return _color_convert_kernel(
        "rgba2grayLine",
        "rgba2gray.cc",
        line_width * 4,
        line_width,
        use_chess=use_chess,
        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,
)

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.
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.
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).

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: list | 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 (dims_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,
    **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.

With guard=True the core writes GUARD_BYTES of 0x55 after each output tile in its memory before every call and drains them with the tile, so a kernel that writes past its output shows up on the host: size the outputs with output_size(..., guard=True) and split them with strip_guard. bfp outputs carry no guard.

Source code in python/iron/algorithms/kernel_design.py

def design(
    factory,
    *,
    calls=1,
    scalars=(),
    shape=None,
    params=None,
    aiecc_flags=None,
    guard=False,
    **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.

    With ``guard=True`` the core writes ``GUARD_BYTES`` of ``0x55`` after each
    output tile in its memory before every call and drains them with the
    tile, so a kernel that writes past its output shows up on the host:
    size the outputs with ``output_size(..., guard=True)`` and split them
    with ``strip_guard``. bfp outputs carry no guard.
    """
    calls = _calls(calls, shape)
    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
    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(fn),
        scalars=tuple(scalars),
        params=_encode_params(fn, params or ()),
        guard=guard,
        **({"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)
    )
    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)

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). A stream longer than that is some kernel emitting markers it does not declare.

Source code in python/iron/algorithms/kernel_design.py

def split_intervals(durations, *, calls, per_call, setup=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)``. A stream longer
    than that is some kernel emitting markers it does not declare.
    """
    durations = [int(d) for d in durations]
    expected = setup + calls * per_call
    if len(durations) > expected:
        raise RuntimeError(
            f"expected {expected} trace intervals, got {len(durations)}; a kernel "
            "on the core emits markers its contract's trace does not declare"
        )
    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()), 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 and to the kernel sources.

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()``),
    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 and to the kernel sources.

    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.
    """

    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"
    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(),
        **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 --baseline-sources ../mlir-aie-base

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 MLIR_AIE_KERNEL_SOURCES set to a checkout, that checkout's aie_kernels/ and aie_runtime_lib/ are compiled instead of the installed copies.

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	stack_bytes on the deepest path from the entry

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. stack_bytes over the contract's stack_bytes (else the device default) prints a warning: the design reserves that much 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; --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 and writes nothing.

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,
    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(),
    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.stack_bytes is not None:
        out.append(
            _row(
                f"{prefix}/stack_bytes",
                "bytes",
                report.stack_bytes,
                extra,
                "plus the runtime routines' own" if report.libcalls else None,
            )
        )
    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"via={loop.pipeliner}{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: {p.stderr.strip()[-2000:]}"
    rep = parse_yaml(yaml_out) if yaml_out.exists() else StaticReport()
    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.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
    from aie.utils.bfp import dtype_name

    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)
            suffix = "".join(
                f"/{k}={dtype_name(v) if isinstance(v, type) else v}"
                for k, v in sorted(combo.items())
            )
            yield f"{name}{suffix}", 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``."""
    from aie.utils import config

    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}: {p.stderr.strip()[-2000:]}")
    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}"
            )
        e_bf = e32.astype(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.

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.
    """

    def ordinal(x):
        bits = np.asarray(x).astype(bfloat16).view(np.uint16).astype(np.int32)
        return np.where(bits & 0x8000, 0x8000 - bits, bits)

    return np.abs(ordinal(a) - ordinal(b))