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, clamp, magnitude_f32.

add_ref

add_ref(a, b)

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

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

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

add_i8_ref

add_i8_ref(a, b, a_shift, b_shift, out_shift)

Numpy reference for add on int8.

sat_i8(round_half_even(((a << a_shift) + (b << b_shift)) / 2**out_shift)).

Raises:

Type	Description
ValueError	A shift is outside what the kernel takes: a_shift and b_shift in [0, 23], so a shifted int8 fits the int32 accumulator, and out_shift in [0, 31].

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

def add_i8_ref(a, b, a_shift, b_shift, out_shift):
    """Numpy reference for [`add`][aie.iron.kernels.eltwise.add] on int8.

    ``sat_i8(round_half_even(((a << a_shift) + (b << b_shift)) / 2**out_shift))``.

    Raises:
        ValueError: A shift is outside what the kernel takes: ``a_shift`` and
            ``b_shift`` in [0, 23], so a shifted int8 fits the int32
            accumulator, and ``out_shift`` in [0, 31].
    """
    if not (0 <= int(a_shift) <= 23 and 0 <= int(b_shift) <= 23):
        raise ValueError(
            f"add: a_shift ({a_shift}) and b_shift ({b_shift}) must be in [0, 23]"
        )
    if not 0 <= int(out_shift) <= 31:
        raise ValueError(f"add: out_shift ({out_shift}) must be in [0, 31]")
    total = (np.asarray(a, np.int64) << int(a_shift)) + (
        np.asarray(b, np.int64) << int(b_shift)
    )
    return np.clip(np.round(total / 2.0 ** int(out_shift)), -128, 127).astype(np.int8)

mul_ref

mul_ref(a, b)

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

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

def mul_ref(a, b):
    """Numpy reference for [`mul`][aie.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 utils/pyright_stubs/aie/iron/kernels/eltwise.py

def scale_ref(x, factor):
    """Numpy reference for [`scale`][aie.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.int8, 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.int8, 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 utils/pyright_stubs/aie/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.int8``, ``np.uint8``, ``np.int16``, or
            ``np.int32``).

    Returns:
        ExternalFunction configured for ``passThroughLine``.

    Raises:
        ValueError: When ``dtype`` is not ``np.int8``, ``np.uint8``,
            ``np.int16``, or ``np.int32``, or when ``tile_size`` is not a whole number of
            64-byte vectors, or is nonpositive.
    """
    # The copy is bitwise, so int8 shares the 8-bit build.
    bit_width = (
        8
        if np.dtype(dtype) == np.dtype(np.int8)
        else _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 utils/pyright_stubs/aie/iron/kernels/eltwise.py

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 addition: bf16, or int8 on power-of-two scales.

The int8 kernel is ONNX QLinearAdd with power-of-two scales: it takes a_shift, b_shift and out_shift after the output tile, shifts each operand up by its shift onto a common scale, and the sum down by out_shift, rounding half to even and saturating. a_shift and b_shift are in [0, 23], so a shifted int8 fits the int32 accumulator, and out_shift in [0, 31]: an output scale finer than both inputs' (a negative out_shift) is not supported.

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per tile: 1024 for bf16 (hard-coded in C++); a multiple of 32 for int8.	1024
dtype	type	bfloat16 or np.int8.	bfloat16
vectorized	bool	If True use vectorized path; False selects scalar (bf16 only).	True

Returns:

Type	Description
ExternalFunction	ExternalFunction for eltwise_add_bf16 or eltwise_add_i8.

Raises:

Type	Description
ValueError	When dtype is neither, or an int8 tile_size is not a multiple of 32.

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

def add(
    tile_size: int = 1024, dtype: type = bfloat16, vectorized: bool = True
) -> ExternalFunction:
    """Element-wise addition: bf16, or int8 on power-of-two scales.

    The int8 kernel is ONNX QLinearAdd with power-of-two scales: it takes
    ``a_shift``, ``b_shift`` and ``out_shift`` after the output tile, shifts
    each operand up by its shift onto a common scale, and the sum down by
    ``out_shift``, rounding half to even and saturating. ``a_shift`` and
    ``b_shift`` are in [0, 23], so a shifted int8 fits the int32
    accumulator, and ``out_shift`` in [0, 31]: an output scale finer than
    both inputs' (a negative ``out_shift``) is not supported.

    Args:
        tile_size: Elements per tile: 1024 for bf16 (hard-coded in C++); a
            multiple of 32 for int8.
        dtype: ``bfloat16`` or ``np.int8``.
        vectorized: If ``True`` use vectorized path; ``False`` selects scalar
            (bf16 only).

    Returns:
        ExternalFunction for eltwise_add_bf16 or eltwise_add_i8.

    Raises:
        ValueError: When ``dtype`` is neither, or an int8 ``tile_size`` is not
            a multiple of 32.
    """
    if np.dtype(dtype) != np.dtype(np.int8):
        return _eltwise_bf16_kernel("add", tile_size, dtype, vectorized)
    if tile_size % 32 or tile_size <= 0 or not vectorized:
        raise ValueError(
            f"add(): an int8 tile_size must be a positive multiple of 32 and "
            f"vectorized, got {tile_size}"
        )
    tile_ty = np.ndarray[(tile_size,), np.dtype[np.int8]]
    return _make_extern(
        "eltwise_add_i8_vector",
        _kernel_source("eltwise/add.cc"),
        [tile_ty, tile_ty, tile_ty, np.int32, np.int32, np.int32],
        compile_flags=["-DADD_INT8", f"-DADD_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out, Param, Param, Param),
            reference=add_i8_ref,
            acc_dtype=np.int32,
            tolerance=Tolerance.exact(note="integer arithmetic, rounding modeled"),
            alignments=((0, 32), (1, 32), (2, 32)),
        ),
    )

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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/iron/kernels/eltwise.py

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

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

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

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

mul_add_ref

mul_add_ref(a, b, is_mul)

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

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

def mul_add_ref(a, b, is_mul):
    """Numpy reference for [`mul_add`][aie.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)

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 utils/pyright_stubs/aie/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`][aie.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 utils/pyright_stubs/aie/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`][aie.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

relu(tile_size: int = 1024) -> ExternalFunction

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

The ABI keeps (in, out, size). Not LUT-based. Positive multiples of 32 elements are supported.

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

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

    The ABI keeps ``(in, out, size)``. Not LUT-based. Positive multiples of
    32 elements are supported.
    """
    _require_vector_alignment("relu", 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,
        ),
    )

clamp

clamp(tile_size: int = 1024) -> ExternalFunction

Element-wise bf16 clamp to [low, high], with a compiled-in element count.

The design passes (in, out, size, low, high), each bound as its bf16 bits in an int32 (int(np.array(b, bfloat16).view(np.uint16))), the type a runtime parameter word holds, so the bounds stay runtime values.

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call, a positive multiple of 32.	1024

Returns:

Type	Description
ExternalFunction	ExternalFunction for clamp_bf16.

Raises:

Type	Description
ValueError	When tile_size is not a positive multiple of 32.

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

def clamp(tile_size: int = 1024) -> ExternalFunction:
    """Element-wise bf16 clamp to ``[low, high]``, with a compiled-in element count.

    The design passes ``(in, out, size, low, high)``, each bound as its bf16
    bits in an int32 (``int(np.array(b, bfloat16).view(np.uint16))``), the
    type a runtime parameter word holds, so the bounds stay runtime values.

    Args:
        tile_size: Elements per call, a positive multiple of 32.

    Returns:
        ExternalFunction for ``clamp_bf16``.

    Raises:
        ValueError: When ``tile_size`` is not a positive multiple of 32.
    """
    _require_vector_alignment("clamp", tile_size, 32)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _make_extern(
        "clamp_bf16",
        _kernel_source("eltwise/clamp.cc"),
        [tile_ty, tile_ty, np.int32, np.int32, np.int32],
        compile_flags=[f"-DCLAMP_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param, Param, Param),
            parameter_bindings=((2, tile_size),),
            reference=clamp_ref,
            tolerance=Tolerance.exact(note="selection: min and max are exact in bf16"),
            ops_per_call=tile_size,
        ),
    )

clamp_ref

clamp_ref(x, low, high)

Numpy reference for clamp: min(max(x, low), high).

low and high are the bounds' bf16 bits, as the kernel takes them.

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

def clamp_ref(x, low, high):
    """Numpy reference for [`clamp`][aie.iron.kernels.eltwise.clamp]: ``min(max(x, low), high)``.

    ``low`` and ``high`` are the bounds' bf16 bits, as the kernel takes them.
    """
    low, high = (np.uint16(b).view(bfloat16).astype(np.float32) for b in (low, high))
    return np.minimum(np.maximum(x.astype(np.float32), low), high)

magnitude_f32

magnitude_f32(
    tile_size: int = 320, power: int = 1
) -> ExternalFunction

sqrt(re^2 + im^2) of complex float32, or its square, float32-accurate.

With power=2 it is re^2 + im^2, the power spectrum. Takes (x, y, size): x holds tile_size real parts then as many imaginary parts, and y gets tile_size float32 results. Every product is of bf16 limbs, exact in the float32 accumulator, and the square root is three Newton steps from a bit-trick estimate. Defined for |z| zero or from 2**-51 to below 2**63, where |z|^2 and the steps' residuals stay normal and finite float32.

The source sets conv_even itself and restores the core's mode on exit, and builds for aie2 and aie2p.

Parameters:

Name	Type	Description	Default
tile_size	int	Results per call, a positive multiple of 32.	320
power	int	1 for the magnitude, 2 for its square.	1

Returns:

Type	Description
ExternalFunction	ExternalFunction for magnitude_f32 or power_f32.

Raises:

Type	Description
ValueError	When tile_size is not a positive multiple of 32, or power is not 1 or 2.

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

def magnitude_f32(tile_size: int = 320, power: int = 1) -> ExternalFunction:
    """``sqrt(re^2 + im^2)`` of complex float32, or its square, float32-accurate.

    With ``power=2`` it is ``re^2 + im^2``, the power spectrum. Takes ``(x, y, size)``: ``x`` holds ``tile_size`` real parts then as many
    imaginary parts, and ``y`` gets ``tile_size`` float32 results. Every
    product is of bf16 limbs, exact in the float32 accumulator, and the
    square root is three Newton steps from a bit-trick estimate. Defined for
    ``|z|`` zero or from ``2**-51`` to below ``2**63``, where ``|z|^2`` and
    the steps' residuals stay normal and finite float32.

    The source sets ``conv_even`` itself and restores the core's mode on
    exit, and builds for aie2 and aie2p.

    Args:
        tile_size: Results per call, a positive multiple of 32.
        power: 1 for the magnitude, 2 for its square.

    Returns:
        ExternalFunction for ``magnitude_f32`` or ``power_f32``.

    Raises:
        ValueError: When ``tile_size`` is not a positive multiple of 32, or
            ``power`` is not 1 or 2.
    """
    _require_vector_alignment("magnitude_f32", tile_size, 32)
    if power not in (1, 2):
        raise ValueError(f"magnitude_f32() power must be 1 or 2, got {power}")
    return _make_extern(
        "magnitude_f32" if power == 1 else "power_f32",
        _kernel_source("eltwise/magnitude_f32.cc"),
        [
            np.ndarray[(2 * tile_size,), np.dtype[np.float32]],
            np.ndarray[(tile_size,), np.dtype[np.float32]],
            np.int32,
        ],
        compile_flags=[f"-DMAGNITUDE_ELEMS={tile_size}"]
        + (["-DMAGNITUDE_POWER"] if power == 2 else []),
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            reference=partial(magnitude_f32_ref, power=power),
            sample=lambda rng, calls: magnitude_f32_sample(
                rng, calls, tile_size=tile_size
            ),
            tolerance=(
                _MAGNITUDE_F32_TOLERANCE if power == 1 else _POWER_F32_TOLERANCE
            ),
            ops_per_call=tile_size,
        ),
    )

magnitude_f32_ref

magnitude_f32_ref(x, power=1)

Numpy reference for magnitude_f32.

hypot of each call's two halves, or its square, in float64.

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

def magnitude_f32_ref(x, power=1):
    """Numpy reference for [`magnitude_f32`][aie.iron.kernels.eltwise.magnitude_f32].

    ``hypot`` of each call's two halves, or its square, in float64.
    """
    re, im = np.split(np.asarray(x, np.float64), 2, axis=-1)
    return np.hypot(re, im) if power == 1 else re * re + im * im

magnitude_f32_sample

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

Complex float32 of every phase over the whole domain, the first call led by edges.

Magnitudes are log-uniform from 2**-51 to 2**62, and one in eight is a zero bin. The edges are zero, the axes, 3 + 4i in each quadrant, the domain's ends, and parts far apart in scale.

Source code in utils/pyright_stubs/aie/iron/kernels/eltwise.py

def magnitude_f32_sample(rng, calls: int, *, tile_size: int) -> list:
    """Complex float32 of every phase over the whole domain, the first call led by edges.

    Magnitudes are log-uniform from ``2**-51`` to ``2**62``, and one in
    eight is a zero bin. The edges are zero, the axes, ``3 + 4i`` in
    each quadrant, the domain's ends, and parts far apart in scale.
    """
    n = calls * tile_size
    r = np.exp2(rng.uniform(-51, 62, n)) * (rng.uniform(size=n) >= 0.125)
    phase = rng.uniform(0, 2 * np.pi, n)
    re, im = r * np.cos(phase), r * np.sin(phase)
    edges = [
        (0.0, 0.0),
        (1.0, 0.0),
        (0.0, -1.0),
        (3.0, 4.0),
        (-3.0, 4.0),
        (-3.0, -4.0),
        (3.0, -4.0),
        (2.0**-51, 0.0),
        (0.0, 2.0**-51),
        (2.0**63 * (1 - 2.0**-24), 0.0),
        (0.0, -(2.0**63) * (1 - 2.0**-24)),
        (1e15, 1e-15),
        (1e-15, 1e15),
        (1e15, 1e15),
        (1e-15, 1e-15),
        (1.0, 2.0**-12),
    ]
    k = min(len(edges), n)
    re[:k], im[:k] = np.array(edges[:k]).T
    re, im = re.reshape(calls, tile_size), im.reshape(calls, tile_size)
    return [np.concatenate((re, im), axis=-1).astype(np.float32)]

Data movement

Data-movement / conversion kernel factories.

affine_cast, axpy, convert_copy, decode_coded, expand, limbs_f32, merge_rows, patch_positions, row_addresses, transpose.

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

axpy_ref

axpy_ref(x, y, a)

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

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

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

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

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

convert_copy_ref

convert_copy_ref(x)

Numpy reference for convert_copy.

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

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

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

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

affine_cast_ref

affine_cast_ref(x, gamma_beta)

Numpy reference for affine_cast.

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

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

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

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

expand_ref

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

Numpy reference for expand.

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

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def expand_ref(payload, *, tile_size: int, group_size: int):
    """Numpy reference for [`expand`][aie.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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/iron/kernels/datamovement.py

def transpose_ref(x, *, dim_m: int, dim_n: int, subtile: int):
    """Numpy reference for [`transpose`][aie.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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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)"
            ),
        ),
    )

limbs_f32_split

limbs_f32_split(x)

Split float32 x into bf16 limbs (hi, mid, lo) that sum to it exactly.

Each limb is its residual rounded half-to-even, as the kernel's conv_even rounds: hi = bf16(x), mid = bf16(x - hi) and lo = bf16(x - hi - mid), every subtraction exact in float32.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def limbs_f32_split(x):
    """Split float32 ``x`` into bf16 limbs ``(hi, mid, lo)`` that sum to it exactly.

    Each limb is its residual rounded half-to-even, as the kernel's
    ``conv_even`` rounds: ``hi = bf16(x)``, ``mid = bf16(x - hi)`` and
    ``lo = bf16(x - hi - mid)``, every subtraction exact in float32.
    """
    x = np.asarray(x, dtype=np.float32)
    hi = x.astype(bfloat16)
    r = x - hi.astype(np.float32)
    mid = r.astype(bfloat16)
    lo = (r - mid.astype(np.float32)).astype(bfloat16)
    return hi, mid, lo

limbs_f32_ref

limbs_f32_ref(x)

Numpy reference for limbs_f32.

The six planes (hi, mid, hi, lo, mid, hi) of each call's limbs, concatenated along the last axis.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def limbs_f32_ref(x):
    """Numpy reference for [`limbs_f32`][aie.iron.kernels.datamovement.limbs_f32].

    The six planes ``(hi, mid, hi, lo, mid, hi)`` of each call's limbs,
    concatenated along the last axis.
    """
    hi, mid, lo = limbs_f32_split(x)
    return np.concatenate((hi, mid, hi, lo, mid, hi), axis=-1)

limbs_f32_sample

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

Float32 of either sign over 1e-30 to 1e38, the first call led by edges.

The edges are signed zeros and ones, half-to-even ties in hi and in mid, the smallest magnitudes whose lo is still a normal bf16, and the largest whose hi is still finite.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def limbs_f32_sample(rng, calls: int, *, tile_size: int) -> list:
    """Float32 of either sign over ``1e-30`` to ``1e38``, the first call led by edges.

    The edges are signed zeros and ones, half-to-even ties in ``hi`` and in
    ``mid``, the smallest magnitudes whose ``lo`` is still a normal bf16,
    and the largest whose ``hi`` is still finite.
    """
    x = rng.choice([-1.0, 1.0], calls * tile_size) * 10.0 ** rng.uniform(
        -30, 38, calls * tile_size
    )
    x = x.astype(np.float32)
    edges = np.array(
        [
            0x00000000,
            0x80000000,
            0x3F800000,
            0xBF800000,
            0x3F800001,
            0x3F808000,
            0x3F818000,
            0xBF80FFFF,
            0x3F800181,
            0x3F800183,
            0x0C000101,
            0x8C7F80FF,
            0x7F7F7FFF,
            0xFF7F7FFF,
            0x3EAAAAAB,
            0xBDCCCCCD,
        ],
        dtype=np.uint32,
    ).view(np.float32)
    x[: edges.size] = edges[: x.size]
    return [x.reshape(calls, tile_size)]

limbs_f32

limbs_f32(tile_size: int = 320) -> ExternalFunction

Float32 as six planes of its bf16 limbs, for a float32-accurate bf16 matmul.

Takes (x, y, size): tile_size float32 in, 6 * tile_size bf16 out. Each element splits exactly into hi + mid + lo (limbs_f32_split), and y holds the planes (hi, mid, hi, lo, mid, hi), tile_size each. A bf16 matmul of those planes against a second operand's limbs stacked (hi, hi, mid, hi, mid, lo) along K sums the six largest of a float32 product's nine limb products, missing only mid * lo, lo * mid and lo * lo. The split is exact for |x| from 2**-103, where lo may be the smallest normal bf16 (the core flushes a subnormal one), to below 3.396e38, where hi rounds to inf.

The source sets conv_even itself and restores the core's mode on exit, and builds for aie2 and aie2p.

Parameters:

Name	Type	Description	Default
tile_size	int	Float32 elements per call, a positive multiple of 32.	320

Returns:

Type	Description
ExternalFunction	ExternalFunction for limbs_f32.

Raises:

Type	Description
ValueError	When tile_size is not a positive multiple of 32.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def limbs_f32(tile_size: int = 320) -> ExternalFunction:
    """Float32 as six planes of its bf16 limbs, for a float32-accurate bf16 matmul.

    Takes ``(x, y, size)``: ``tile_size`` float32 in, ``6 * tile_size`` bf16
    out. Each element splits exactly into ``hi + mid + lo``
    ([`limbs_f32_split`][aie.iron.kernels.datamovement.limbs_f32_split]), and
    ``y`` holds the planes ``(hi, mid, hi, lo, mid, hi)``, ``tile_size``
    each. A bf16 matmul of those planes against a second operand's limbs
    stacked ``(hi, hi, mid, hi, mid, lo)`` along K sums the six largest of a
    float32 product's nine limb products, missing only ``mid * lo``,
    ``lo * mid`` and ``lo * lo``. The split is exact for ``|x|`` from
    ``2**-103``, where ``lo`` may be the smallest normal bf16 (the core
    flushes a subnormal one), to below ``3.396e38``, where ``hi`` rounds to
    inf.

    The source sets ``conv_even`` itself and restores the core's mode on
    exit, and builds for aie2 and aie2p.

    Args:
        tile_size: Float32 elements per call, a positive multiple of 32.

    Returns:
        ExternalFunction for ``limbs_f32``.

    Raises:
        ValueError: When ``tile_size`` is not a positive multiple of 32.
    """
    _require_vector_alignment("limbs_f32", tile_size, 32)
    return _make_extern(
        "limbs_f32",
        _kernel_source("datamovement/limbs_f32.cc"),
        [
            np.ndarray[(tile_size,), np.dtype[np.float32]],
            np.ndarray[(6 * tile_size,), np.dtype[bfloat16]],
            np.int32,
        ],
        compile_flags=[f"-DLIMBS_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            reference=limbs_f32_ref,
            sample=lambda rng, calls: limbs_f32_sample(rng, calls, tile_size=tile_size),
            tolerance=Tolerance.exact(
                note="each residual is exact in the accumulator, and so is each "
                "limb of it"
            ),
            ops_per_call=tile_size,
        ),
    )

affine_cast

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction for affine_cast_f32_bf16.

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

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

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

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

    Returns:
        ExternalFunction for ``affine_cast_f32_bf16``.

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

expand

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction for expand_uint4_to_bfloat16.

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

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

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

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

    Returns:
        ExternalFunction for ``expand_uint4_to_bfloat16``.

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

merge_rows_ref

merge_rows_ref(
    ids,
    *,
    audio_token: int,
    image_token: int,
    audio_at: int,
    vision_at: int
)

Numpy reference for merge_rows.

ids is (calls, block), block b of one sequence per call: the j-th audio_token of the sequence takes row audio_at + j, the j-th image_token row vision_at + j, and every other id its position.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def merge_rows_ref(
    ids, *, audio_token: int, image_token: int, audio_at: int, vision_at: int
):
    """Numpy reference for [`merge_rows`][aie.iron.kernels.datamovement.merge_rows].

    ``ids`` is ``(calls, block)``, block ``b`` of one sequence per call: the
    j-th ``audio_token`` of the sequence takes row ``audio_at + j``, the j-th
    ``image_token`` row ``vision_at + j``, and every other id its position.
    """
    ids = np.asarray(ids)
    flat = ids.reshape(-1)
    rows = np.arange(flat.size, dtype=np.int64)
    for token, at in ((audio_token, audio_at), (image_token, vision_at)):
        places = np.flatnonzero(flat == token)
        rows[places] = at + np.arange(places.size)
    return rows.astype(np.int32).reshape(ids.shape)

merge_rows_sample

merge_rows_sample(
    rng,
    calls: int,
    *,
    block: int,
    audio_token: int,
    image_token: int
) -> list

Return ids, a third of them placeholders and the rest any int32, led by the ids beside the placeholders.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def merge_rows_sample(
    rng, calls: int, *, block: int, audio_token: int, image_token: int
) -> list:
    """Return ids, a third of them placeholders and the rest any int32, led by the ids beside the placeholders."""
    tokens = np.array([audio_token, image_token], np.int64)
    ids = rng.integers(-(2**31), 2**31, size=calls * block, dtype=np.int64)
    ids = np.where(rng.random(ids.size) < 1 / 3, rng.choice(tokens, ids.size), ids)
    edges = [audio_token - 1, audio_token + 1, image_token - 1, image_token + 1]
    edges = [e for e in edges if e not in tokens and -(2**31) <= e < 2**31]
    ids[: len(edges)] = edges[: ids.size]
    return [ids.reshape(calls, block).astype(np.int32)]

merge_rows

merge_rows(
    block: int = 64,
    *,
    audio_token: int = 258881,
    image_token: int = 258880,
    audio_at: int = 2048,
    vision_at: int = 2816
) -> ExternalFunction

Merge-rows kernel: the row each position of a token sequence takes from a table.

The table holds the text's embeddings followed by two towers' soft tokens. Takes (ids, rows, b): block b of a sequence's int32 ids, block of them, and the row each takes, int32. A position takes its own row, except that the j-th audio_token of the sequence takes row audio_at + j and the j-th image_token row vision_at + j. The counts carry from one call to the next and restart at b = 0, so one core calls it for blocks 0, 1, ... in order. That the placeholders are as many as the soft tokens is the caller's check. amd/IRON's Merge gathers a multimodal prompt's embeddings by these rows. The defaults are EmbeddingGemma 2's placeholders at 2048 text rows and 768 audio rows.

Parameters:

Name	Type	Description	Default
block	int	Ids per call.	64
audio_token	int	The audio placeholder id.	258881
image_token	int	The image placeholder id.	258880
audio_at	int	The audio tower's first row.	2048
vision_at	int	The vision tower's first row.	2816

Returns:

Type	Description
ExternalFunction	ExternalFunction for merge_rows.

Raises:

Type	Description
ValueError	When block is below 1, the placeholders are equal, or a row is negative.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def merge_rows(
    block: int = 64,
    *,
    audio_token: int = 258881,
    image_token: int = 258880,
    audio_at: int = 2048,
    vision_at: int = 2816,
) -> ExternalFunction:
    """Merge-rows kernel: the row each position of a token sequence takes from a table.

    The table holds the text's embeddings followed by two towers' soft
    tokens. Takes ``(ids, rows, b)``: block ``b`` of a sequence's int32 ids,
    ``block`` of them, and the row each takes, int32. A position takes its
    own row, except that the j-th ``audio_token`` of the sequence takes row
    ``audio_at + j`` and the j-th ``image_token`` row ``vision_at + j``. The
    counts carry from one call to the next and restart at ``b = 0``, so one
    core calls it for blocks 0, 1, ... in order. That the placeholders are
    as many as the soft tokens is the caller's check. amd/IRON's ``Merge``
    gathers a multimodal prompt's embeddings by these rows. The defaults
    are EmbeddingGemma 2's placeholders at 2048 text rows and 768 audio
    rows.

    Args:
        block: Ids per call.
        audio_token: The audio placeholder id.
        image_token: The image placeholder id.
        audio_at: The audio tower's first row.
        vision_at: The vision tower's first row.

    Returns:
        ExternalFunction for ``merge_rows``.

    Raises:
        ValueError: When ``block`` is below 1, the placeholders are equal, or
            a row is negative.
    """
    if block < 1 or audio_token == image_token or min(audio_at, vision_at) < 0:
        raise ValueError(
            "merge_rows() needs block >= 1, two distinct placeholders and rows "
            f">= 0, got block={block}, audio_token={audio_token}, "
            f"image_token={image_token}, audio_at={audio_at}, vision_at={vision_at}."
        )
    return _make_extern(
        "merge_rows",
        _kernel_source("datamovement/merge_rows.cc"),
        [
            np.ndarray[(block,), np.dtype[np.int32]],
            np.ndarray[(block,), np.dtype[np.int32]],
            np.int32,
        ],
        compile_flags=[
            f"-DBLOCK={block}",
            f"-DAUDIO_TOKEN={audio_token}",
            f"-DIMAGE_TOKEN={image_token}",
            f"-DAUDIO_AT={audio_at}",
            f"-DVISION_AT={vision_at}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, 0),),
            call_index=2,
            reference=lambda ids: merge_rows_ref(
                ids,
                audio_token=audio_token,
                image_token=image_token,
                audio_at=audio_at,
                vision_at=vision_at,
            ),
            tolerance=Tolerance.exact(note="integer row arithmetic"),
            ops_per_call=0,
            sample=lambda rng, calls: merge_rows_sample(
                rng,
                calls,
                block=block,
                audio_token=audio_token,
                image_token=image_token,
            ),
        ),
    )

patch_positions_ref

patch_positions_ref(
    count,
    ho,
    wo,
    *,
    block: int,
    side: int,
    pool: int,
    positions: int,
    cores: int
)

Numpy reference for patch_positions.

Returns (count, 5 * block), one record per block: xy (2 * block), the x positions, the y positions plus positions, and the raster rows.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def patch_positions_ref(
    count, ho, wo, *, block: int, side: int, pool: int, positions: int, cores: int
):
    """Numpy reference for [`patch_positions`][aie.iron.kernels.datamovement.patch_positions].

    Returns ``(count, 5 * block)``, one record per block: ``xy`` (``2 * block``),
    the x positions, the y positions plus ``positions``, and the raster rows.
    """
    count, ho, wo = int(count), int(ho), int(wo)
    rows = count * block
    bands, strips = max(ho, 0) // side, max(wo, 0) // side
    pad = (strips // cores + 1) * cores
    across = strips // pool
    total = bands * strips
    # The kernel's unsigned 64-bit bound, and n's truncation to int.
    limit = (count % 2**64) * block % 2**64
    ok = (
        ho % side == 0
        and wo % side == 0
        and 0 < total <= limit
        and bands % pool == 0
        and strips % pool == 0
    )
    n = (total + 2**31) % 2**32 - 2**31 if ok else 0
    k = np.arange(min(max(n, 0), rows), dtype=np.int64)
    token, w = np.divmod(k, pool * pool)
    ty, tx = np.divmod(token, across) if across else (token, token)
    y, x = pool * ty + w // pool, pool * tx + w % pool
    xy = np.full((rows, 2), positions, np.int64)
    position_ids = np.full((2, rows), 2 * positions, np.int64)
    order = np.full(rows, pad - 1, np.int64)
    xy[: k.size, 0], xy[: k.size, 1] = x, y
    position_ids[0, : k.size], position_ids[1, : k.size] = x, positions + y
    order[: k.size] = y * pad + x
    record = np.concatenate(
        [
            xy.reshape(count, 2 * block),
            position_ids[0].reshape(count, block),
            position_ids[1].reshape(count, block),
            order.reshape(count, block),
        ],
        axis=1,
    )
    return record.astype(np.int32)

patch_positions

patch_positions(
    block: int = 256,
    *,
    side: int = 16,
    pool: int = 3,
    positions: int = 10240,
    cores: int = 16
) -> ExternalFunction

Patch-position kernel: where each patch of a vision tower's image comes from.

Takes (record, b, count, ho, wo): block b of count of the patches of an ho by wo pixel image, cut into side-pixel patches, in pooling-window order (each soft token's pool ** 2 patches consecutive and row-major within the window, the tokens row-major over the image). The record, 5 * block int32, holds each patch's x and y side by side, then every patch's x, then every patch's y plus positions, then its row of a raster of the image padded past a multiple of cores patch columns. Rows past the image, and every row of an image that is not whole windows or that count blocks do not hold, take positions, 2 * positions and the raster's last row, the padding rows of the position tables and the raster. The state carries from one call to the next and restarts at b = 0, so one core calls it for blocks 0 to count - 1 in order. amd/IRON's PatchPositions writes these for EmbeddingGemma 2's vision tower.

Parameters:

Name	Type	Description	Default
block	int	Patches per call.	256
side	int	Pixels on a patch's side.	16
pool	int	Patches on a pooling window's side.	3
positions	int	Rows of each axis's position table.	10240
cores	int	The raster's columns are padded past a multiple of this.	16

Returns:

Type	Description
ExternalFunction	ExternalFunction for patch_positions.

Raises:

Type	Description
ValueError	When block, side, pool or cores is below 1, or positions is negative or 2 * positions past int32.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def patch_positions(
    block: int = 256,
    *,
    side: int = 16,
    pool: int = 3,
    positions: int = 10240,
    cores: int = 16,
) -> ExternalFunction:
    """Patch-position kernel: where each patch of a vision tower's image comes from.

    Takes ``(record, b, count, ho, wo)``: block ``b`` of ``count`` of the
    patches of an ``ho`` by ``wo`` pixel image, cut into ``side``-pixel
    patches, in pooling-window order (each soft token's ``pool ** 2``
    patches consecutive and row-major within the window, the tokens
    row-major over the image). The record, ``5 * block`` int32, holds each
    patch's x and y side by side, then every patch's x, then every patch's
    y plus ``positions``, then its row of a raster of the image padded past
    a multiple of ``cores`` patch columns. Rows past the image, and every
    row of an image that is not whole windows or that ``count`` blocks do
    not hold, take ``positions``, ``2 * positions`` and the raster's last
    row, the padding rows of the position tables and the raster. The state
    carries from one call to the next and restarts at ``b = 0``, so one core
    calls it for blocks 0 to ``count - 1`` in order. amd/IRON's
    ``PatchPositions`` writes these for EmbeddingGemma 2's vision tower.

    Args:
        block: Patches per call.
        side: Pixels on a patch's side.
        pool: Patches on a pooling window's side.
        positions: Rows of each axis's position table.
        cores: The raster's columns are padded past a multiple of this.

    Returns:
        ExternalFunction for ``patch_positions``.

    Raises:
        ValueError: When ``block``, ``side``, ``pool`` or ``cores`` is below 1,
            or ``positions`` is negative or ``2 * positions`` past int32.
    """
    if min(block, side, pool, cores) < 1 or not 0 <= 2 * positions < 2**31:
        raise ValueError(
            "patch_positions() needs block, side, pool and cores >= 1 and "
            f"0 <= 2 * positions < 2**31, got block={block}, side={side}, "
            f"pool={pool}, positions={positions}, cores={cores}."
        )
    return _make_extern(
        "patch_positions",
        _kernel_source("datamovement/patch_positions.cc"),
        [np.ndarray[(5 * block,), np.dtype[np.int32]]] + [np.int32] * 4,
        compile_flags=[
            f"-DBLOCK={block}",
            f"-DSIDE={side}",
            f"-DPOOL={pool}",
            f"-DPOSITIONS={positions}",
            f"-DCORES={cores}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, Param, Param, Param, Param),
            parameter_bindings=((1, 0),),
            call_index=1,
            reference=lambda count, ho, wo: patch_positions_ref(
                count,
                ho,
                wo,
                block=block,
                side=side,
                pool=pool,
                positions=positions,
                cores=cores,
            ),
            tolerance=Tolerance.exact(note="integer index arithmetic"),
            ops_per_call=0,
        ),
    )

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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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)

row_addresses_ref

row_addresses_ref(
    ids,
    lo,
    hi,
    *,
    table_rows: int,
    row_bytes: int,
    low_bits: int,
    aperture: int
)

Numpy reference for row_addresses.

The table sits at (hi << low_bits) + lo, each word read as unsigned; an id past the table is clipped to its first or last row and a negative one counts from the end, np.clip(ids, -n, n - 1) % n.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def row_addresses_ref(
    ids, lo, hi, *, table_rows: int, row_bytes: int, low_bits: int, aperture: int
):
    """Numpy reference for [`row_addresses`][aie.iron.kernels.datamovement.row_addresses].

    The table sits at ``(hi << low_bits) + lo``, each word read as unsigned; an
    id past the table is clipped to its first or last row and a negative one
    counts from the end, ``np.clip(ids, -n, n - 1) % n``.
    """
    ids = np.asarray(ids, dtype=np.int64)
    rows = np.clip(ids, -table_rows, table_rows - 1) % table_rows
    base = ((int(hi) & 0xFFFFFFFF) << low_bits) + (int(lo) & 0xFFFFFFFF) + aperture
    address = (base + rows * row_bytes).astype(np.uint64)
    words = np.empty(ids.shape[:-1] + (2 * ids.shape[-1],), np.uint32)
    words[..., 0::2] = address & np.uint64(0xFFFFFFFC)
    words[..., 1::2] = (address >> np.uint64(32)) & np.uint64(0xFFFF)
    return words

row_addresses_sample

row_addresses_sample(
    rng, calls: int, *, rows: int, table_rows: int
) -> list

Ids over twice the table each way, the first calls' led by the edges of its range.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def row_addresses_sample(rng, calls: int, *, rows: int, table_rows: int) -> list:
    """Ids over twice the table each way, the first calls' led by the edges of its range."""
    n = table_rows
    ids = rng.integers(-2 * n, 2 * n, size=calls * rows, dtype=np.int64)
    edges = [-n - 1, -n, -1, 0, n - 1, n, -(2**31), 2**31 - 1]
    ids[: len(edges)] = edges[: ids.size]
    return [ids.reshape(calls, rows).astype(np.int32)]

row_addresses

row_addresses(
    rows: int = 8,
    table_rows: int = 1024,
    row_bytes: int = 4096,
    *,
    low_bits: int = 29,
    aperture: int = 2147483648
) -> ExternalFunction

Row-address kernel: each id's table row as a shim buffer descriptor's address words.

Takes (ids, out, lo, hi): rows int32 ids, and for each the low and high address words of its row, 2 * rows uint32. The table sits at (hi << low_bits) + lo, split so each part fits the 30 bits a core reads of a runtime value; aperture is added to every address. An id past the table is clipped to its first or last row and a negative one counts from the end, so every id names a row. The low word is rounded down to 4 bytes and the high word keeps 16 bits, as a shim buffer descriptor's address fields hold them. amd/IRON's GatherWords writes these into the control packets of a gather whose ids the device made.

Parameters:

Name	Type	Description	Default
rows	int	Ids per call.	8
table_rows	int	Rows of the table.	1024
row_bytes	int	Bytes per row (a positive multiple of 4).	4096
low_bits	int	Bits of the address lo holds, from 0 to 30.	29
aperture	int	Added to every address; the default is the offset a DDR address carries in a shim buffer descriptor (kDDRAIEAddrOffset).	2147483648

Returns:

Type	Description
ExternalFunction	ExternalFunction for row_addresses.

Raises:

Type	Description
ValueError	When rows or table_rows is below 1, row_bytes is not a positive multiple of 4, or low_bits is outside 0 to 30.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def row_addresses(
    rows: int = 8,
    table_rows: int = 1024,
    row_bytes: int = 4096,
    *,
    low_bits: int = 29,
    aperture: int = 0x80000000,
) -> ExternalFunction:
    """Row-address kernel: each id's table row as a shim buffer descriptor's address words.

    Takes ``(ids, out, lo, hi)``: ``rows`` int32 ids, and for each the low
    and high address words of its row, ``2 * rows`` uint32. The table sits at
    ``(hi << low_bits) + lo``, split so each part fits the 30 bits a core
    reads of a runtime value; ``aperture`` is added to every address. An id
    past the table is clipped to its first or last row and a negative one
    counts from the end, so every id names a row. The low word is rounded
    down to 4 bytes and the high word keeps 16 bits, as a shim buffer
    descriptor's address fields hold them. amd/IRON's ``GatherWords`` writes
    these into the control packets of a gather whose ids the device made.

    Args:
        rows: Ids per call.
        table_rows: Rows of the table.
        row_bytes: Bytes per row (a positive multiple of 4).
        low_bits: Bits of the address ``lo`` holds, from 0 to 30.
        aperture: Added to every address; the default is the offset a DDR
            address carries in a shim buffer descriptor (kDDRAIEAddrOffset).

    Returns:
        ExternalFunction for ``row_addresses``.

    Raises:
        ValueError: When ``rows`` or ``table_rows`` is below 1, ``row_bytes``
            is not a positive multiple of 4, or ``low_bits`` is outside 0 to 30.
    """
    if rows < 1 or table_rows < 1 or row_bytes < 4 or row_bytes % 4:
        raise ValueError(
            "row_addresses() needs rows and table_rows >= 1 and row_bytes a "
            f"positive multiple of 4, got rows={rows}, table_rows={table_rows}, "
            f"row_bytes={row_bytes}."
        )
    if not 0 <= low_bits <= 30:
        raise ValueError(
            "row_addresses() needs low_bits from 0 to 30, since lo must fit the "
            f"30 bits a core reads of a runtime value, got low_bits={low_bits}."
        )
    return _make_extern(
        "row_addresses",
        _kernel_source("datamovement/row_addresses.cc"),
        [
            np.ndarray[(rows,), np.dtype[np.int32]],
            np.ndarray[(2 * rows,), np.dtype[np.uint32]],
            np.int32,
            np.int32,
        ],
        compile_flags=[
            f"-DROWS={rows}",
            f"-DTABLE_ROWS={table_rows}",
            f"-DROW_BYTES={row_bytes}",
            f"-DLOW_BITS={low_bits}",
            f"-DAPERTURE={aperture:#x}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param, Param),
            reference=lambda ids, lo, hi: row_addresses_ref(
                ids,
                lo,
                hi,
                table_rows=table_rows,
                row_bytes=row_bytes,
                low_bits=low_bits,
                aperture=aperture,
            ),
            tolerance=Tolerance.exact(note="integer address arithmetic"),
            ops_per_call=0,
            sample=lambda rng, calls: row_addresses_sample(
                rng, calls, rows=rows, table_rows=table_rows
            ),
        ),
    )

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 utils/pyright_stubs/aie/iron/kernels/datamovement.py

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

coded_rows_sample

coded_rows_sample(
    rng, calls: int, *, dim_m: int, dim_k: int
) -> ndarray

Draw calls tiles of dim_m coded rows of dim_k weights.

Rows are normal draws held to at least 2^-8 in magnitude, so they span 11 binades and fit their windows, with up to CODED_CAPACITY entries pushed far outside them. The first row of each call holds the full count, its first and last weight among them, so every call patches exceptions and the patch reaches both ends of a row.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def coded_rows_sample(rng, calls: int, *, dim_m: int, dim_k: int) -> np.ndarray:
    """Draw ``calls`` tiles of ``dim_m`` coded rows of ``dim_k`` weights.

    Rows are normal draws held to at least 2^-8 in magnitude, so they span
    11 binades and fit their windows, with up to ``CODED_CAPACITY`` entries
    pushed far outside them. The first row of each call holds the full
    count, its first and last weight among them, so every call patches
    exceptions and the patch reaches both ends of a row.
    """
    a = rng.standard_normal((calls, dim_m, dim_k)).astype(np.float32)
    a = np.copysign(np.maximum(np.abs(a), np.float32(2**-8)), a)
    inner = np.arange(1, dim_k - 1)
    for c in range(calls):
        for m in range(dim_m):
            if m == 0:
                at = np.r_[0, dim_k - 1, rng.choice(inner, CODED_CAPACITY - 2, False)]
            else:
                n = int(rng.integers(0, CODED_CAPACITY + 1))
                at = rng.choice(dim_k, size=n, replace=False)
            a[c, m, at] *= np.exp2(rng.choice([-60.0, 60.0], size=at.size))
    return coded.encode_rows(a.astype(bfloat16), CODED_CAPACITY)

decode_coded_ref

decode_coded_ref(a, *, dim_k: int)

Numpy reference for decode_coded: aie.utils.coded.decode_rows.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def decode_coded_ref(a, *, dim_k: int):
    """Numpy reference for [`decode_coded`][aie.iron.kernels.datamovement.decode_coded]: ``aie.utils.coded.decode_rows``."""
    return coded.decode_rows(np.asarray(a), dim_k)

decode_coded

decode_coded(
    dim_m: int = 4, dim_k: int = 2048
) -> ExternalFunction

Coded rows in, bf16 rows out.

aie_kernels/datamovement/decode_coded.cc, signature (m, A, out): A is dim_m rows coded as aie.utils.coded writes them (capacity 15), coded.row_bytes(dim_k) bytes each, and out their dim_m * dim_k bf16 weights, bit for bit. It is the decode the coded matvec runs ahead of its products (common/coded_rows.h), as a kernel of its own.

Parameters:

Name	Type	Description	Default
dim_m	int	Rows a call decodes.	4
dim_k	int	Weights in a row, a multiple of 64.	2048

Returns:

Type	Description
ExternalFunction	ExternalFunction for decode_coded_bf16.

Raises:

Type	Description
ValueError	dim_k is not a positive multiple of 64.

Source code in utils/pyright_stubs/aie/iron/kernels/datamovement.py

def decode_coded(dim_m: int = 4, dim_k: int = 2048) -> ExternalFunction:
    """Coded rows in, bf16 rows out.

    ``aie_kernels/datamovement/decode_coded.cc``, signature ``(m, A, out)``:
    ``A`` is ``dim_m`` rows coded as ``aie.utils.coded`` writes them (capacity
    15), ``coded.row_bytes(dim_k)`` bytes each, and ``out`` their
    ``dim_m * dim_k`` bf16 weights, bit for bit. It is the decode the coded
    matvec runs ahead of its products (``common/coded_rows.h``), as a kernel
    of its own.

    Args:
        dim_m: Rows a call decodes.
        dim_k: Weights in a row, a multiple of 64.

    Returns:
        ExternalFunction for ``decode_coded_bf16``.

    Raises:
        ValueError: ``dim_k`` is not a positive multiple of 64.
    """
    if dim_k <= 0 or dim_k % 64:
        raise ValueError(f"decode_coded(): dim_k ({dim_k}) must be a multiple of 64")
    row_bytes = coded.row_bytes(dim_k, CODED_CAPACITY)
    return _make_extern(
        "decode_coded_bf16",
        _kernel_source("datamovement/decode_coded.cc"),
        [
            np.int32,
            np.ndarray[(dim_m * row_bytes,), np.dtype[np.uint8]],
            np.ndarray[(dim_m * dim_k,), np.dtype[bfloat16]],
        ],
        compile_flags=[f"-DDIM_K={dim_k}", f"-DCAPACITY={CODED_CAPACITY}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(Param, In, Out),
            parameter_bindings=((0, dim_m),),
            layouts=(
                None,
                TensorLayout((dim_m, row_bytes)),
                TensorLayout((dim_m, dim_k)),
            ),
            reference=lambda a: decode_coded_ref(a, dim_k=dim_k),
            tolerance=Tolerance.exact(note="a lossless decode"),
            ops_per_call=dim_m * dim_k,
            sample=lambda rng, calls: [
                coded_rows_sample(rng, calls, dim_m=dim_m, dim_k=dim_k)
            ],
        ),
    )

Independent, target-native zero-fill kernel.

zero

zero(
    tile_size: int | tuple[int, ...] = 1024,
    dtype: type | dtype = int32,
    *,
    vectorized: bool = True,
    window: int | None = None,
    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.

With window, each call zeroes window elements of the one tile at a runtime offset, call * window, so the calls fill it between them from starts of every alignment.

Source code in utils/pyright_stubs/aie/iron/kernels/zero.py

def zero(
    tile_size: int | tuple[int, ...] = 1024,
    dtype: type | np.dtype = np.int32,
    *,
    vectorized: bool = True,
    window: int | None = None,
    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.

    With ``window``, each call zeroes ``window`` elements of the one tile
    at a runtime offset, ``call * window``, so the calls fill it between
    them from starts of every alignment.
    """
    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")
    arg_types = [np.ndarray[shape, np.dtype[dtype]]]
    roles, layouts, bindings, out_offset = (Out,), (layout,), (), None
    if window is not None:
        if block or not vectorized:
            raise ValueError("zero: window needs a plain dtype and vectorized=True")
        if window <= 0 or size % window:
            raise ValueError(f"zero: window {window} must divide tile_size {size}")
        arg_types += [np.int32, np.int32]
        layout = TensorLayout((window,))
        roles, layouts = (Out, Param, Param), (layout, None, None)
        bindings, out_offset = ((1, 0), (2, window)), (1, window)
    return _make_extern(
        "zero" if window is None else "zero_window",
        _kernel_source("zero/zero.cc"),
        arg_types,
        compile_flags=flags,
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=roles,
            layouts=layouts,
            parameter_bindings=bindings,
            out_offset=out_offset,
            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_unpack

q4nx_unpack(
    payload,
    *,
    m_tile=Q4NX_M_TILE,
    k_tile=Q4NX_K_TILE,
    group=Q4NX_GROUP,
    ct_k=128,
    s=8,
    t=8
) -> tuple[ndarray, ndarray, ndarray]

Unpack packed q4nx to its codes, scales and minima.

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]. The weight at [..., k, n] is mins[..., k // group, n] + scales[..., k // group, n] * codes[..., k, n]. Finite scales and minima are required.

Returns:

Type	Description
tuple[ndarray, ndarray, ndarray]	tuple[np.ndarray, np.ndarray, np.ndarray]: (codes, scales, mins), with the same leading dimensions as payload. codes is uint8 with shape (..., k_tile, m_tile); scales and mins are float32 with shape (..., k_tile // group, m_tile).

Source code in utils/pyright_stubs/aie/iron/kernels/quant.py

def q4nx_unpack(
    payload,
    *,
    m_tile=Q4NX_M_TILE,
    k_tile=Q4NX_K_TILE,
    group=Q4NX_GROUP,
    ct_k=128,
    s=8,
    t=8,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Unpack packed q4nx to its codes, scales and minima.

    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]``. The weight at ``[..., k, n]`` is
    ``mins[..., k // group, n] + scales[..., k // group, n] * codes[..., k, n]``.
    Finite scales and minima are required.

    Returns:
        tuple[np.ndarray, np.ndarray, np.ndarray]: ``(codes, scales, mins)``,
            with the same leading dimensions as ``payload``. ``codes`` is
            ``uint8`` with shape ``(..., k_tile, m_tile)``; ``scales`` and
            ``mins`` are ``float32`` with shape ``(..., k_tile // group,
            m_tile)``.
    """
    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_unpack: expected {input_bytes} bytes per tile")
    lead = payload.shape[:-1]
    data = payload.reshape(-1, input_bytes)
    if not len(data):
        codes = np.empty((*lead, k_tile, m_tile), dtype=np.uint8)
        params = np.empty((*lead, k_tile // group, m_tile), dtype=np.float32)
        return codes, params, params.copy()
    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_unpack: 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(*lead, k_tile, m_tile)
    scales = scales.reshape(*lead, k_tile // group, m_tile)
    mins = mins.reshape(*lead, k_tile // group, m_tile)
    return q, scales, mins

q4nx_dequant_ref

q4nx_dequant_ref(
    payload,
    *,
    m_tile=Q4NX_M_TILE,
    k_tile=Q4NX_K_TILE,
    group=Q4NX_GROUP,
    ct_k=128,
    s=8,
    t=8
)

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

See q4nx_unpack for the packed input layout.

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 dequantized bf16 results are required.

Source code in utils/pyright_stubs/aie/iron/kernels/quant.py

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

    See ``q4nx_unpack`` for the packed input layout.

    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 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]
    if not math.prod(lead):
        return np.empty((*lead, m_tile * k_tile * 9 // 8), dtype=np.uint8)
    q, scales, mins = q4nx_unpack(payload, **geometry)
    q = q.reshape(-1, k_tile, m_tile)
    scales = scales.reshape(-1, k_tile // group, m_tile)
    mins = mins.reshape(-1, k_tile // group, 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(math.prod(lead), -1))
    return encoded.reshape(*lead, m_tile * k_tile * 9 // 8)

q4nx_dequant

q4nx_dequant(
    *,
    m_tile=Q4NX_M_TILE,
    k_tile=Q4NX_K_TILE,
    group=Q4NX_GROUP,
    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 utils/pyright_stubs/aie/iron/kernels/quant.py

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

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

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

Core state

set_rounding: the core's rounding-mode register.

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

RoundingMode

Bases: str, Enum

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

set_rounding

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

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

Parameters:

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

Returns:

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

Source code in utils/pyright_stubs/aie/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`][aie.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, argmax, maxpool2d.

reduce_add_ref

reduce_add_ref(x)

Numpy reference for reduce_add: per-tile sum.

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

Source code in utils/pyright_stubs/aie/iron/kernels/reduce.py

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

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

reduce_min_ref

reduce_min_ref(x)

Numpy reference for reduce_min: per-tile minimum.

Source code in utils/pyright_stubs/aie/iron/kernels/reduce.py

def reduce_min_ref(x):
    """Numpy reference for [`reduce_min`][aie.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 utils/pyright_stubs/aie/iron/kernels/reduce.py

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

reduce_add

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the reduce_add kernel.

Raises:

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

Source code in utils/pyright_stubs/aie/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 (int32 or bfloat16).

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

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

    Returns:
        ExternalFunction configured for the reduce_add kernel.

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

reduce_min

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the reduce_min kernel.

Raises:

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

Source code in utils/pyright_stubs/aie/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 (int32 or bfloat16).

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

    Returns:
        ExternalFunction configured for the reduce_min kernel.

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

reduce_max

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

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

Parameters:

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

Source code in utils/pyright_stubs/aie/iron/kernels/reduce.py

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, bfloat16 or float32).

    Args:
        tile_size: Number of elements in the input tile.
        dtype: Element data type (``np.int32``, ``bfloat16`` or ``np.float32``).
        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``, ``bfloat16`` or
            ``np.float32``.
    """
    suffixes = {
        np.dtype(np.int32): "",
        np.dtype(bfloat16): "_bfloat16",
        np.dtype(np.float32): "_float",
    }
    if np.dtype(dtype) not in suffixes:
        raise ValueError(
            f"reduce_max() dtype must be np.int32, bfloat16 or np.float32, got {dtype}"
        )
    actual_dtype = np.dtype(dtype).type
    if vectorized:
        _require_vector_alignment(
            "reduce_max",
            tile_size,
            (_REDUCE_VEC_ELEMS_BF16 if actual_dtype is bfloat16 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"
    return _make_extern(
        f"reduce_max_{func_variant}{suffixes[np.dtype(dtype)]}",
        _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 utils/pyright_stubs/aie/iron/kernels/reduce.py

def compute_max(dtype: type = np.int32) -> ExternalFunction:
    """Pairwise scalar max — companion to [`reduce_max`][aie.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`][aie.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 utils/pyright_stubs/aie/iron/kernels/reduce.py

def compute_max_ref(a, b):
    """Numpy reference for [`compute_max`][aie.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])

argmax_ref

argmax_ref(x, index_offset: int = 0)

Numpy reference for argmax: one record per tile.

Returns the (..., 2) int32 records the kernel writes, so a host can compare them verbatim. A NaN reads as -inf and a -0 as +0, as they do in the kernel's comparisons; that is where this differs from a plain numpy.argmax.

Source code in utils/pyright_stubs/aie/iron/kernels/reduce.py

def argmax_ref(x, index_offset: int = 0):
    """Numpy reference for [`argmax`][aie.iron.kernels.reduce.argmax]: one record per tile.

    Returns the ``(..., 2)`` int32 records the kernel writes, so a host can
    compare them verbatim. A NaN reads as -inf and a -0 as +0, as they do in
    the kernel's comparisons; that is where this differs from a plain
    ``numpy.argmax``.
    """
    x = np.asarray(x)
    if np.issubdtype(x.dtype, np.integer):
        keys = x.astype(np.int32)
    else:
        keys = x.astype(np.float32)
        keys = np.where(np.isnan(keys), np.float32(-np.inf), keys)
        keys = np.where(keys == 0, np.float32(0), keys)
    index = keys.argmax(axis=-1)[..., None]
    value = np.take_along_axis(keys, index, axis=-1).view(np.int32)
    return np.concatenate([value, (index + index_offset).astype(np.int32)], axis=-1)

argmax

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

Reduction kernel: the largest element of a tile and its index.

The partial half of a distributed argmax, in the same shape as reduce_max + compute_max: each core runs this over its own slice and a tree merges the records with argmax_combine.

The kernel writes a 2-element int32 record: out[0] the winning value (int32 as itself, bfloat16 widened to float and bit-cast) and out[1] its index plus the index_offset runtime argument. A caller that passes its slice's start gets global indices, and the combine order stops mattering.

Ties resolve to the lowest index, matching numpy.argmax. A NaN never compares greater, so it reads as -inf, and a -0 reads as +0, so -0 and +0 tie.

Parameters:

Name	Type	Description	Default
tile_size	int	Number of elements in the input slice, 1..32767. The vector path takes any remainder in a scalar tail.	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 argmax kernel; signature is
ExternalFunction	(in_ty, out_ty, int32 tile_size, int32 index_offset).

Raises:

Type	Description
ValueError	When dtype is not np.int32 or bfloat16, or tile_size is not in 1..32767.

Source code in utils/pyright_stubs/aie/iron/kernels/reduce.py

def argmax(
    tile_size: int = 1024, dtype: type = np.int32, vectorized: bool = True
) -> ExternalFunction:
    """Reduction kernel: the largest element of a tile and its index.

    The partial half of a distributed argmax, in the same shape as
    [`reduce_max`][aie.iron.kernels.reduce.reduce_max] +
    [`compute_max`][aie.iron.kernels.reduce.compute_max]: each core runs this over
    its own slice and a tree merges the records with
    [`argmax_combine`][aie.iron.kernels.reduce.argmax_combine].

    The kernel writes a 2-element int32 record: ``out[0]`` the winning value
    (int32 as itself, bfloat16 widened to float and bit-cast) and ``out[1]``
    its index plus the ``index_offset`` runtime argument. A caller that passes
    its slice's start gets global indices, and the combine order stops
    mattering.

    Ties resolve to the lowest index, matching ``numpy.argmax``. A NaN never
    compares greater, so it reads as -inf, and a -0 reads as +0, so -0 and +0
    tie.

    Args:
        tile_size: Number of elements in the input slice, 1..32767. The
            vector path takes any remainder in a scalar tail.
        dtype: Element data type (``np.int32`` or ``bfloat16``).
        vectorized: If ``True`` use vectorized path; ``False`` selects scalar.

    Returns:
        ExternalFunction configured for the argmax kernel; signature is
        ``(in_ty, out_ty, int32 tile_size, int32 index_offset)``.

    Raises:
        ValueError: When ``dtype`` is not ``np.int32`` or ``bfloat16``, or
            ``tile_size`` is not in 1..32767.
    """
    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"argmax() dtype must be np.int32 or bfloat16, got {dtype}")
    actual_dtype = bfloat16 if is_bf16 else np.int32
    if not 0 < tile_size <= _ARGMAX_MAX_ELEMS:
        raise ValueError(
            f"argmax() tile_size must be in 1..{_ARGMAX_MAX_ELEMS} (the kernel "
            f"indexes a tile with int16 lanes), got {tile_size}"
        )
    in_ty = np.ndarray[(tile_size,), np.dtype[actual_dtype]]
    out_ty = np.ndarray[(2,), np.dtype[np.int32]]
    func_variant = "vector" if vectorized else "scalar"
    suffix = "_bfloat16" if is_bf16 else ""
    return _make_extern(
        f"argmax_{func_variant}{suffix}",
        _kernel_source("reduce/argmax.cc"),
        [in_ty, out_ty, np.int32, np.int32],
        compile_flags=[f"-DARGMAX_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param, Param),
            parameter_bindings=((2, tile_size),),
            reference=argmax_ref,
            tolerance=Tolerance.exact(note="an exact selection"),
            ops_per_call=tile_size,
        ),
    )

argmax_combine_ref

argmax_combine_ref(a, b, dtype: type = int32)

Numpy reference for argmax_combine.

Keeps, per call, the record with the larger value, or the lower index when the values are equal. dtype is the original input's and says how out[0] reads.

Source code in utils/pyright_stubs/aie/iron/kernels/reduce.py

def argmax_combine_ref(a, b, dtype: type = np.int32):
    """Numpy reference for [`argmax_combine`][aie.iron.kernels.reduce.argmax_combine].

    Keeps, per call, the record with the larger value, or the lower index when
    the values are equal. ``dtype`` is the original input's and says how
    ``out[0]`` reads.
    """
    a, b = np.asarray(a, np.int32), np.asarray(b, np.int32)
    is_bf16 = np.dtype(dtype) == np.dtype(bfloat16)
    value_dtype = np.float32 if is_bf16 else np.int32
    va, vb = a[..., :1].view(value_dtype), b[..., :1].view(value_dtype)
    take_b = (vb > va) | ((vb == va) & (b[..., 1:] < a[..., 1:]))
    return np.where(take_b, b, a)

argmax_combine_sample

argmax_combine_sample(
    rng, calls: int, *, dtype: type
) -> list

Two tiles of records argmax can write, with frequent value ties.

The values come from a handful of small integers plus the extremes a record can hold (int32's limits, or ±inf and -0.0), so equal values (and the index tie-break) come up in most calls; a record never carries a NaN.

Source code in utils/pyright_stubs/aie/iron/kernels/reduce.py

def argmax_combine_sample(rng, calls: int, *, dtype: type) -> list:
    """Two tiles of records ``argmax`` can write, with frequent value ties.

    The values come from a handful of small integers plus the extremes a
    record can hold (int32's limits, or ±inf and -0.0), so equal values (and
    the index tie-break) come up in most calls; a record never carries a NaN.
    """
    if np.dtype(dtype) == np.dtype(bfloat16):
        pool = np.array([-np.inf, -3, -1, -0.0, 0, 2, np.inf], np.float32)
        pool = pool.view(np.int32)
    else:
        info = np.iinfo(np.int32)
        pool = np.array([info.min, -3, -1, 0, 2, info.max], np.int32)
    values = rng.choice(pool, size=(2, calls, 1))
    indices = rng.integers(0, 8, size=(2, calls, 1))
    records = np.concatenate([values, indices], axis=-1).astype(np.int32)
    return [records[0], records[1]]

argmax_combine

argmax_combine(dtype: type = int32) -> ExternalFunction

Pairwise record merge — companion to argmax.

Takes two of the records described in argmax and keeps the one with the larger value, or the lower index when the values are equal. When both operands carry global indices, the merge order does not affect the result.

Parameters:

Name	Type	Description	Default
dtype	type	Element data type of the original input (np.int32 or bfloat16); it selects how out[0] is compared.	int32

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the argmax_combine kernel;
ExternalFunction	signature is (rec_ty, rec_ty, rec_ty) where rec_ty is a
ExternalFunction	2-element int32 record.

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/reduce.py

def argmax_combine(dtype: type = np.int32) -> ExternalFunction:
    """Pairwise record merge — companion to [`argmax`][aie.iron.kernels.reduce.argmax].

    Takes two of the records described in [`argmax`][aie.iron.kernels.reduce.argmax]
    and keeps the one with the larger value, or the lower index when the values
    are equal. When both operands carry global indices, the merge order does
    not affect the result.

    Args:
        dtype: Element data type of the original input (``np.int32`` or
            ``bfloat16``); it selects how ``out[0]`` is compared.

    Returns:
        ExternalFunction configured for the ``argmax_combine`` kernel;
        signature is ``(rec_ty, rec_ty, rec_ty)`` where ``rec_ty`` is a
        2-element int32 record.

    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"argmax_combine() dtype must be np.int32 or bfloat16, got {dtype}"
        )
    actual_dtype = bfloat16 if is_bf16 else np.int32
    rec_ty = np.ndarray[(2,), np.dtype[np.int32]]
    suffix = "_bfloat16" if is_bf16 else ""
    return _make_extern(
        f"argmax_combine{suffix}",
        _kernel_source("reduce/argmax.cc"),
        [rec_ty, rec_ty, rec_ty],
        object_file_name=_ARGMAX_OBJ,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            reference=partial(argmax_combine_ref, dtype=actual_dtype),
            sample=partial(argmax_combine_sample, dtype=actual_dtype),
            tolerance=Tolerance.exact(note="selection"),
            ops_per_call=1,
        ),
    )

maxpool2d_ref

maxpool2d_ref(
    line0, line1, line2, line3, line4, input_width, channels
)

Numpy reference for maxpool2d: a 5x5 stride-1 max pool row.

Lines are [C/8][W][8] int8, the output row is the max over the five lines and the five pixels around each one, pixels past the row's ends counting as -inf.

Source code in utils/pyright_stubs/aie/iron/kernels/reduce.py

def maxpool2d_ref(line0, line1, line2, line3, line4, input_width, channels):
    """Numpy reference for [`maxpool2d`][aie.iron.kernels.reduce.maxpool2d]: a 5x5 stride-1 max pool row.

    Lines are ``[C/8][W][8]`` int8, the output row is the max over the five
    lines and the five pixels around each one, pixels past the row's ends
    counting as -inf.
    """
    W, C = int(input_width), int(channels)
    lines = [np.asarray(line) for line in (line0, line1, line2, line3, line4)]
    lead = lines[0].shape[:-1]
    column = np.max(np.stack(lines), axis=0).reshape(*lead, C // 8, W, 8)
    padded = np.pad(
        column, [(0, 0)] * len(lead) + [(0, 0), (2, 2), (0, 0)], constant_values=-128
    )
    out = np.max(np.stack([padded[..., k : k + W, :] for k in range(5)]), axis=0)
    return out.astype(np.int8).reshape(*lead, C * W)

maxpool2d

maxpool2d(
    input_width: int = 20, channels: int = 128
) -> ExternalFunction

5x5, stride-1 max pool of one int8 row, padded with -inf (YOLO's SPPF).

Takes the five [C/8][W][8] rows around the output row, then the output row. At the top or bottom of the image the caller passes an edge row again in place of a missing one, which a max leaves unchanged.

Parameters:

Name	Type	Description	Default
input_width	int	Pixels per row, a multiple of 4.	20
channels	int	Channels, a multiple of 8.	128

Returns:

Type	Description
ExternalFunction	ExternalFunction for maxpool2d_5x5_i8.

Raises:

Type	Description
ValueError	On a width or channel count the kernel cannot tile.

Source code in utils/pyright_stubs/aie/iron/kernels/reduce.py

def maxpool2d(input_width: int = 20, channels: int = 128) -> ExternalFunction:
    """5x5, stride-1 max pool of one int8 row, padded with -inf (YOLO's SPPF).

    Takes the five ``[C/8][W][8]`` rows around the output row, then the
    output row. At the top or bottom of the image the caller passes an edge
    row again in place of a missing one, which a max leaves unchanged.

    Args:
        input_width: Pixels per row, a multiple of 4.
        channels: Channels, a multiple of 8.

    Returns:
        ExternalFunction for maxpool2d_5x5_i8.

    Raises:
        ValueError: On a width or channel count the kernel cannot tile.
    """
    if input_width % 4 or channels % 8 or input_width < 4:
        raise ValueError(
            f"maxpool2d(): needs a width of 4k and channels of 8k, got "
            f"{input_width} and {channels}"
        )
    line_ty = np.ndarray[(input_width * channels,), np.dtype[np.int8]]
    align = 64 if input_width % 8 == 0 else 32
    return _make_extern(
        "maxpool2d_5x5_i8",
        _kernel_source("reduce/maxpool2d.cc"),
        [line_ty] * 6,
        compile_flags=[
            f"-DMAXPOOL_WIDTH={input_width}",
            f"-DMAXPOOL_CHANNELS={channels}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, In, In, Out),
            reference=partial(
                maxpool2d_ref, input_width=input_width, channels=channels
            ),
            tolerance=Tolerance.exact(note="a max is exact"),
            alignments=tuple((i, align) for i in range(6)),
            ops_per_call=24 * input_width * channels,
        ),
    )

Linear algebra

Linear algebra kernel factories: mm, mv, cascade_mm.

StreamDimsABC

Bases: NamedTuple

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

None for an operand a build streams untransformed.

MatrixKernel

MatrixKernel(
    name: str,
    object_file_name: str | None = None,
    source_file: str | None = None,
    source_string: str | None = None,
    arg_types: list[ArgType] | None = None,
    include_dirs: list[str] | None = None,
    compile_flags: list[str] | None = None,
    *,
    symbol_prefix: str | None = None,
    use_chess: bool = False,
    inline: bool = False,
    stack_size_override: int | None = None,
    contract: Any = None
)

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.

Source code in utils/pyright_stubs/aie/iron/kernel.py

def __init__(
    self,
    name: str,
    object_file_name: str | None = None,
    source_file: str | None = None,
    source_string: str | None = None,
    arg_types: list[ArgType] | None = None,
    include_dirs: list[str] | None = None,
    compile_flags: list[str] | None = None,
    *,
    symbol_prefix: str | None = None,
    use_chess: bool = False,
    inline: bool = False,
    stack_size_override: int | None = None,
    contract: Any = None,
) -> None:
    """Construct an ExternalFunction compiled from C/C++ source at JIT time.

    Args:
        name: Symbol name of the function as it will appear in the object
            file.
        object_file_name: Output artifact name. Defaults to
            ``<effective_name>.o``, or ``<effective_name>.ll`` with
            ``inline=True``. With ``inline=True`` an explicit name must end
            in ``.ll`` (textual LLVM IR) or ``.bc`` (bitcode) -- that suffix
            selects the emitted format -- and is otherwise rejected.
        source_file: Path to a C/C++ source file on disk.  Mutually
            exclusive with ``source_string``.
        source_string: Inline C/C++ source code.  Mutually exclusive with
            ``source_file``.
        arg_types: Type signature of the function arguments.  Defaults to
            None (empty list).
        include_dirs: Additional ``-I`` directories passed to the chosen
            compiler (Peano by default; xchesscc when ``use_chess=True``).
            Relative paths resolve against the current directory at construction.
            Defaults to None (empty list).
        compile_flags: Additional flags passed verbatim to the chosen
            compiler.  Defaults to None (empty list).
        symbol_prefix: Optional prefix for the exported symbol name.  When
            set, the effective symbol name becomes ``<symbol_prefix>_<name>``
            and the object file is named accordingly.  The original name is
            preserved in ``_original_name`` for source file naming.
        use_chess: When ``True``, this ExternalFunction's source is
            compiled with ``xchesscc_wrapper`` instead of Peano's
            ``clang++``.  The JIT compile orchestration auto-detects the
            design-level toolchain from the registered EFs and switches
            aiecc's front-end accordingly; mixing chess + peano EFs in
            one design is rejected loudly because aiecc only invokes one
            front-end per compile.
        inline: When True, compile the kernel to ``alwaysinline`` LLVM IR
            (``.ll``) and declare it with ``link_with_mode = "merge"`` so
            aiecc llvm-links it into the core and inlines it, instead of
            object-linking a separate ``.o``. Removes the ``func.call``
            boundary and the separate object. Peano path only (the
            Chess/xchesscc toolchain cannot llvm-link).
        stack_size_override: Declared upper bound, in bytes, on the stack
            that this kernel's call subtree uses. See
            [`Kernel.stack_size_override`][aie.iron.kernel.Kernel.stack_size_override].
            With ``inline=True``, the merged kernel has no separate object,
            so this bound is the one input aiecc's stack analysis reads for
            this kernel.
        contract: What the kernel computes, as an
            ``aie.iron.kernels.KernelContract``: argument roles, a host
            reference, a tolerance and the operand layouts. The library
            factories always give one; a hand-built kernel may leave it
            ``None`` and then cannot be built or judged generically.
    """
    if inline and use_chess:
        raise ValueError(
            f"ExternalFunction '{name}': inline=True requires the Peano "
            "toolchain and cannot be combined with use_chess=True."
        )
    if inline and symbol_prefix:
        raise NotImplementedError(
            f"ExternalFunction '{name}': inline=True combined with symbol_prefix is "
            "not supported (an inline kernel is emitted as LLVM IR and cannot be "
            "symbol-renamed). Use inline without a symbol_prefix, or drop inline for "
            "this kernel."
        )

    self._original_name = name
    self.contract = contract
    effective_name = f"{symbol_prefix}_{name}" if symbol_prefix else name
    object_file_name_explicit = object_file_name is not None
    if not object_file_name:
        object_file_name = (
            f"{effective_name}.ll" if inline else f"{effective_name}.o"
        )
    elif inline and Path(object_file_name).suffix.lower() not in (".ll", ".bc"):
        # An inline kernel is emitted as LLVM IR, and the suffix picks the
        # format (textual vs bitcode), so a wrong one has no valid reading.
        # Reject it instead of silently renaming the caller's artifact --
        # aiecc routes on the `link_with_mode` attribute, not the suffix, so
        # a rename would buy nothing.  Compared case-insensitively, matching
        # compile_cxx_core_function's own suffix check; the caller's exact
        # spelling is preserved either way.
        raise ValueError(
            f"ExternalFunction '{name}': inline=True emits LLVM IR, so "
            f"object_file_name must end in '.ll' (textual LLVM IR) or "
            f"'.bc' (bitcode); got '{object_file_name}'."
        )
    super().__init__(
        effective_name,
        object_file_name,
        arg_types,
        link_with_mode="merge" if inline else None,
        stack_size_override=stack_size_override,
    )

    if source_file is None and source_string is None:
        raise ValueError("source_file or source_string must be provided.")
    if source_file is not None:
        try:
            source_bytes = Path(source_file).read_bytes()
        except OSError:
            source_bytes = f"<unreadable:{source_file}>".encode()
    else:
        assert source_string is not None
        source_bytes = source_string.encode()
    self._object_file = KernelObject(
        object_file_name,
        "merge" if inline else None,
        _KernelSource(
            str(Path(source_file).resolve()) if source_file is not None else None,
            source_string if source_file is None else None,
            hashlib.sha256(source_bytes).hexdigest(),
            tuple(str(Path(d).absolute()) for d in include_dirs or ()),
            tuple(compile_flags or ()),
            use_chess,
            symbol_prefix,
            name if inline else None,
        ),
    )
    self._cached_digest: str | None = None

    # A translation unit may export several symbols, but an output path
    # must have only one recipe. Auto-suffix default names on conflict;
    # explicit names are never silently changed.
    registry = get_compile_arg("_iron_kernel_registry", ExternalFunction._instances)
    for existing in registry:
        if (
            existing.object_file_name == object_file_name
            and existing.object_file._source != self.object_file._source
        ):
            if object_file_name_explicit:
                raise ValueError(
                    f"ExternalFunction '{effective_name}' would collide with "
                    f"an already-registered instance: same "
                    f"explicit object_file_name='{object_file_name}' but "
                    f"different compile_flags / source.  Distinguish them "
                    f"by passing a distinct `object_file_name=...`."
                )
            suffix = self._content_digest()[:8]
            output_path = Path(object_file_name)
            object_file_name = str(
                output_path.with_name(
                    f"{output_path.stem}_{suffix}{output_path.suffix}"
                )
            )
            self._object_file = replace(self.object_file, name=object_file_name)
            self._cached_digest = None
            break
    for existing in registry:
        if existing.object_file_name == self.object_file_name:
            if existing.object_file._source != self.object_file._source:
                raise ValueError(
                    f"ExternalFunction '{effective_name}' would collide on '{self.object_file_name}'"
                )
            self._object_file = existing.object_file
            break
    registry.add(self)

mac_dims property

mac_dims: tuple[int, int, int]

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

stream_dims property

stream_dims: StreamDimsABC

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

mm_acc_dtype

mm_acc_dtype(input_dtype)

Return what mm.cc accumulates in for input_dtype.

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

Source code in utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/iron/kernels/linalg.py

def mm_ref(a, b):
    """Numpy reference for [`mm`][aie.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 utils/pyright_stubs/aie/iron/kernels/linalg.py

def mv_ref(a, b):
    """Numpy reference for [`mv`][aie.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 utils/pyright_stubs/aie/iron/kernels/linalg.py

def mm_bfp_ref(a, b):
    """Numpy reference for [`mm_bfp`][aie.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 utils/pyright_stubs/aie/iron/kernels/linalg.py

def mm_bfp_mixed_ref(a, b):
    """Numpy reference for [`mm_bfp`][aie.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`][aie.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 utils/pyright_stubs/aie/iron/kernels/linalg.py

def mm_tile_ref(a, b, *, dim_m: int, dim_k: int, dim_n: int):
    """One [`mm`][aie.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 utils/pyright_stubs/aie/iron/kernels/linalg.py

def mv_tile_ref(a, b, *, dim_m: int, dim_k: int):
    """One [`mv`][aie.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 utils/pyright_stubs/aie/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`][aie.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`][aie.iron.kernels.linalg.mm_bfp_mixed_ref] for
    why modelling it as a host-side quantize is worse, not better.
    """
    from aie.utils import bfp

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

mm_stream_dims

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

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

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

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

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

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

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

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

mm

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

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.

.k_tail(A, B, k_valid) zeroes A's columns and B's rows from k_valid (clamped to [0, dim_k]) on, in the blocking this kernel reads, so a call after it sums only the first k_valid of the tile's K, whatever the rest holds: the last tile of a reduction whose length is not a multiple of dim_k.

Parameters:

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

Returns:

Type	Description
_MatMulKernel	ExternalFunction configured for the matmul kernel, with its .zero
_MatMulKernel	and .k_tail.

Raises:

Type	Description
ValueError	When (input_dtype, output_dtype) is not a supported combination, or the vectorized kernel's loops cannot step the tile.

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

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

    ``.k_tail(A, B, k_valid)`` zeroes A's columns and B's rows from
    ``k_valid`` (clamped to ``[0, dim_k]``) on, in the blocking this kernel
    reads, so a call after it sums only the first ``k_valid`` of the tile's
    K, whatever the rest holds: the last tile of a reduction whose length is
    not a multiple of ``dim_k``.

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

    Returns:
        ExternalFunction configured for the matmul kernel, with its ``.zero``
        and ``.k_tail``.

    Raises:
        ValueError: When ``(input_dtype, output_dtype)`` is not a supported
            combination, or the vectorized kernel's loops cannot step the tile.
    """
    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,
    )
    if vectorized:
        step_m, step_n = _MM_TILE_STEPS[arch][key]
        if dim_m % (step_m * r) or dim_k % s or dim_n % (step_n * t):
            raise ValueError(
                f"mm(): on {arch} a {input_dtype.__name__} tile is a multiple of "
                f"{step_m * r} x {s} x {step_n * t}, got {dim_m} x {dim_k} x {dim_n}"
            )
    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),
        ),
    )
    extern = _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=_MatMulKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            layouts=layouts,
            # mm_aie2p.h's entry points set conv_even for a bf16 C and restore
            # the caller's mode; 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=(
                Tolerance.bounded(
                    partial(_mm_nearest_bound, dim_m=dim_m, dim_k=dim_k, dim_n=dim_n),
                    note="the f32 sum rounded once to nearest even",
                )
                if vectorized
                and not bf16_emulated
                and input_dtype is bfloat16
                and output_dtype is bfloat16
                else _linalg_tolerance(input_dtype)
            ),
            ops_per_call=2 * dim_m * dim_k * dim_n,
        ),
    )
    extern.k_tail = extern.object_file.bind(
        f"{prefix}_{suffix}_k_tail", [a_ty, b_ty, np.int32]
    )
    return extern

mv

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

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

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

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

a_coded=True builds it with -DA_CODED, signature (m, row_offset, A, scratch, b, c): A is dim_m coded rows (aie.utils.coded, capacity 15), uint8 of coded.row_bytes(dim_k) each, which the core decodes into scratch (dim_m * dim_k bf16) before the row-major kernel runs on it. So it returns plain mv's bits on the decoded matrix at the same dim_m and vec_size.

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the matvec kernel.

Raises:

Type	Description
ValueError	When the dtype combination is not supported, or both a_col_maj and a_coded are given.

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

def mv(
    dim_m: int = 32,
    dim_k: int = 32,
    input_dtype: type = np.int16,
    output_dtype: type = np.int32,
    vectorized: bool = True,
    use_chess: bool = False,
    vec_size: int = 64,
    output_rows: int | None = None,
    a_col_maj: bool = False,
    a_coded: bool = False,
) -> ExternalFunction:
    """Matrix-vector multiply kernel: c += A * b.

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

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

    ``a_coded=True`` builds it with ``-DA_CODED``, signature ``(m,
    row_offset, A, scratch, b, c)``: ``A`` is ``dim_m`` coded rows
    (``aie.utils.coded``, capacity 15), ``uint8`` of
    ``coded.row_bytes(dim_k)`` each, which the core decodes into ``scratch``
    (``dim_m * dim_k`` bf16) before the row-major kernel runs on it. So it
    returns plain ``mv``'s bits on the decoded matrix at the same ``dim_m``
    and ``vec_size``.

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

    Returns:
        ExternalFunction configured for the matvec kernel.

    Raises:
        ValueError: When the dtype combination is not supported, or both
            ``a_col_maj`` and ``a_coded`` are given.
    """
    if a_col_maj and a_coded:
        raise ValueError("mv(): a_col_maj and a_coded are two layouts of A; pick one")
    if (input_dtype, output_dtype) == (bfloat16, bfloat16):
        return _mv_bf16(
            dim_m,
            dim_k,
            vectorized,
            use_chess,
            vec_size,
            output_rows,
            a_col_maj,
            a_coded,
        )
    if a_col_maj or a_coded:
        raise ValueError("mv(): a_col_maj and a_coded are bf16 layouts")
    if output_rows is not None:
        raise ValueError("mv(): output_rows needs the bf16 kernel's row_offset")
    if input_dtype != np.int16 or output_dtype != np.int32:
        raise ValueError(
            f"mv(): only (np.int16, np.int32) and (bfloat16, bfloat16) are supported, got ({input_dtype}, {output_dtype})"
        )

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

mv_coded_sample

mv_coded_sample(
    rng, calls, *, dim_m: int, dim_k: int, tiled: bool
)

Coded A tiles (coded_rows_sample) and a b for mv with a_coded.

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

def mv_coded_sample(rng, calls, *, dim_m: int, dim_k: int, tiled: bool):
    """Coded A tiles (``coded_rows_sample``) and a b for [`mv`][aie.iron.kernels.linalg.mv] with ``a_coded``."""
    a = coded_rows_sample(rng, calls, dim_m=dim_m, dim_k=dim_k)
    b_shape = (dim_k,) if tiled else (calls, dim_k)
    return [a, rng.standard_normal(b_shape).astype(bfloat16)]

mv_bf16_ref

mv_bf16_ref(m, row_offset, a, b, *, a_coded: bool = False)

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. With a_coded it is the coded rows mv(a_coded=True) reads, decoded first.

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

def mv_bf16_ref(m, row_offset, a, b, *, a_coded: bool = False):
    """Numpy reference for the bf16 [`mv`][aie.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``. With ``a_coded``
    it is the coded rows ``mv(a_coded=True)`` reads, decoded first.
    """
    K = np.asarray(b).shape[-1]
    if a_coded:
        rows = np.asarray(a).reshape(-1, coded.row_bytes(K, CODED_CAPACITY))
        a = coded.decode_rows(rows, K)
    a = np.asarray(a).reshape(-1, K)
    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 utils/pyright_stubs/aie/iron/kernels/linalg.py

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

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

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

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

mm_bfp_shuffle

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

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

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

Parameters:

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

Source code in utils/pyright_stubs/aie/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

attention_i8_ref

attention_i8_ref(
    q_row,
    kt,
    v,
    exp_table,
    q_mul,
    q_shift,
    score_shift,
    prob_shift,
    out_shift,
    *,
    seq_len: int,
    key_dim: int,
    head_dim: int,
    width: int
)

Numpy reference for attention_i8: one row of queries.

q_row is [key_dim/8][width][8] int8 ((calls, n) of them), kt the keys as [seq_len/8][key_dim][8], v the values as [head_dim/8][seq_len][8]; the output row is [head_dim/8][width][8]. Every requantization rounds half to even and saturates; each probability is 2**prob_shift * e / sum(e) rounded exactly, e the exp table at the score's distance below the row's maximum.

Raises:

Type	Description
ValueError	q_mul is not an int8 or prob_shift is not in [0, 30], which the kernel does not take.

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

def attention_i8_ref(
    q_row,
    kt,
    v,
    exp_table,
    q_mul,
    q_shift,
    score_shift,
    prob_shift,
    out_shift,
    *,
    seq_len: int,
    key_dim: int,
    head_dim: int,
    width: int,
):
    """Numpy reference for [`attention_i8`][aie.iron.kernels.linalg.attention_i8]: one row of queries.

    ``q_row`` is ``[key_dim/8][width][8]`` int8 (``(calls, n)`` of them),
    ``kt`` the keys as ``[seq_len/8][key_dim][8]``, ``v`` the values as
    ``[head_dim/8][seq_len][8]``; the output row is ``[head_dim/8][width][8]``.
    Every requantization rounds half to even and saturates; each probability
    is ``2**prob_shift * e / sum(e)`` rounded exactly, ``e`` the exp table at
    the score's distance below the row's maximum.

    Raises:
        ValueError: ``q_mul`` is not an int8 or ``prob_shift`` is not in
            [0, 30], which the kernel does not take.
    """
    if not -128 <= int(q_mul) <= 127 or not 0 <= int(prob_shift) <= 30:
        raise ValueError(
            f"attention_i8: q_mul ({q_mul}) must be an int8 and prob_shift "
            f"({prob_shift}) in [0, 30]"
        )
    N, KD, HD, W = int(seq_len), int(key_dim), int(head_dim), int(width)
    q = np.asarray(q_row, np.int64)
    lead = q.shape[:-1]
    q = q.reshape(*lead, KD // 8, W, 8).swapaxes(-1, -2).reshape(*lead, KD, W)
    q = np.clip(np.round(q * int(q_mul) / 2.0 ** int(q_shift)), -128, 127)
    keys = (
        np.asarray(kt, np.int64)
        .reshape(N // 8, KD, 8)
        .transpose(1, 0, 2)
        .reshape(KD, N)
    )
    vals = (
        np.asarray(v, np.int64).reshape(HD // 8, N, 8).transpose(0, 2, 1).reshape(HD, N)
    )
    scores = np.einsum("...dw,dn->...wn", q, keys)
    scores = np.clip(np.round(scores / 2.0 ** int(score_shift)), -128, 127).astype(
        np.int64
    )
    e = np.asarray(exp_table, np.int64)[scores.max(axis=-1, keepdims=True) - scores]
    num = e << int(prob_shift)
    total = e.sum(axis=-1, keepdims=True)
    quotient, remainder = np.divmod(2 * num, 2 * total)
    up = (remainder > total) | ((remainder == total) & (quotient % 2 == 1))
    probs = np.minimum(quotient + up, 127)
    out = np.einsum("...wn,dn->...dw", probs, vals)
    out = np.clip(np.round(out / 2.0 ** int(out_shift)), -128, 127)
    out = out.reshape(*lead, HD // 8, 8, W).swapaxes(-1, -2)
    return out.astype(np.int8).reshape(*lead, HD * W)

attention_i8_sample

attention_i8_sample(
    rng,
    calls: int,
    *,
    seq_len: int,
    key_dim: int,
    head_dim: int,
    width: int
) -> list

Random int8 queries, keys and values, and the exp table of a 2**-3 score scale.

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

def attention_i8_sample(
    rng, calls: int, *, seq_len: int, key_dim: int, head_dim: int, width: int
) -> list:
    """Random int8 queries, keys and values, and the exp table of a 2**-3 score scale."""
    table = np.round(2.0**26 * np.exp(-np.arange(256) / 8.0)).astype(np.int32)
    return [
        rng.integers(-128, 128, (calls, key_dim * width)).astype(np.int8),
        rng.integers(-128, 128, (seq_len * key_dim,)).astype(np.int8),
        rng.integers(-128, 128, (seq_len * head_dim,)).astype(np.int8),
        table,
    ]

attention_i8

attention_i8(
    seq_len: int = 400,
    key_dim: int = 32,
    head_dim: int = 64,
    width: int = 20,
) -> _AttentionI8Kernel

One head of int8 attention over a 2-D map, one row of queries per call.

YOLO26n's PSA quantized as ONNX QDQ with power-of-two scales: queries scaled by q_mul and shifted, scores shifted to int8, an integer softmax (an exp table indexed by each score's distance below the row's maximum, each probability rounded exactly to 2**-prob_shift steps), and the value product shifted to int8, all rounding half to even. Not flash attention: every key's scores of a query fit the tile at once, so there is no running state, unlike mha.

Arguments: a [key_dim/8][width][8] row of queries, the keys [seq_len/8][key_dim][8], the values [head_dim/8][seq_len][8], the 256-entry int32 exp table (non-increasing, as an exp table is, and each entry from 0 to below 2**29), the output row, then q_mul, q_shift, score_shift, prob_shift, out_shift: q_mul an int8, which the kernel multiplies int8 queries by, and prob_shift from 0 to 30. fn.load_kv(k_row, v_row, kt, v, row) stores one row of the map's keys and values in those layouts.

Parameters:

Name	Type	Description	Default
seq_len	int	Keys: the map's pixels, a multiple of 8.	400
key_dim	int	Query and key channels, a multiple of 8.	32
head_dim	int	Value channels, a multiple of 8.	64
width	int	Pixels per row, a multiple of 4.	20

Returns:

Type	Description
_AttentionI8Kernel	ExternalFunction for attention_i8, with load_kv bound.

Raises:

Type	Description
ValueError	On a shape the kernel cannot tile.

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

def attention_i8(
    seq_len: int = 400, key_dim: int = 32, head_dim: int = 64, width: int = 20
) -> _AttentionI8Kernel:
    """One head of int8 attention over a 2-D map, one row of queries per call.

    YOLO26n's PSA quantized as ONNX QDQ with power-of-two scales: queries
    scaled by ``q_mul`` and shifted, scores shifted to int8, an integer
    softmax (an exp table indexed by each score's distance below the row's
    maximum, each probability rounded exactly to ``2**-prob_shift`` steps),
    and the value product shifted to int8, all rounding half to even. Not
    flash attention: every key's scores of a query fit the tile at once, so
    there is no running state, unlike [`mha`][aie.iron.kernels.linalg.mha].

    Arguments: a ``[key_dim/8][width][8]`` row of queries, the keys
    ``[seq_len/8][key_dim][8]``, the values ``[head_dim/8][seq_len][8]``, the
    256-entry int32 exp table (non-increasing, as an exp table is, and each
    entry from 0 to below ``2**29``), the output row, then
    ``q_mul, q_shift, score_shift, prob_shift, out_shift``: ``q_mul`` an
    int8, which the kernel multiplies int8 queries by, and ``prob_shift``
    from 0 to 30.
    ``fn.load_kv(k_row, v_row, kt, v, row)`` stores one row of the map's keys
    and values in those layouts.

    Args:
        seq_len: Keys: the map's pixels, a multiple of 8.
        key_dim: Query and key channels, a multiple of 8.
        head_dim: Value channels, a multiple of 8.
        width: Pixels per row, a multiple of 4.

    Returns:
        ExternalFunction for attention_i8, with ``load_kv`` bound.

    Raises:
        ValueError: On a shape the kernel cannot tile.
    """
    if seq_len % 8 or key_dim % 8 or head_dim % 8 or width % 4 or seq_len % width:
        raise ValueError(
            f"attention_i8(): needs seq_len, key_dim, head_dim of 8k and a width of 4k "
            f"dividing seq_len, got {seq_len}, {key_dim}, {head_dim}, {width}"
        )
    shapes = dict(seq_len=seq_len, key_dim=key_dim, head_dim=head_dim, width=width)
    q_ty = np.ndarray[(key_dim * width,), np.dtype[np.int8]]
    kt_ty = np.ndarray[(seq_len * key_dim,), np.dtype[np.int8]]
    v_ty = np.ndarray[(seq_len * head_dim,), np.dtype[np.int8]]
    table_ty = np.ndarray[(256,), np.dtype[np.int32]]
    out_ty = np.ndarray[(head_dim * width,), np.dtype[np.int8]]
    compile_flags = [
        f"-DATTN_SEQ={seq_len}",
        f"-DATTN_KEY_DIM={key_dim}",
        f"-DATTN_HEAD_DIM={head_dim}",
        f"-DATTN_WIDTH={width}",
    ]
    if _detect_arch() == "aie2p":
        # Peano hangs at -O2/-O3 compiling the aie2p LUT path below for some
        # shapes (seq_len=256, key_dim=64, head_dim=32, width=16 reproduces
        # it); -O1 is verified bit-exact and compiles in seconds.
        compile_flags.append("-O1")
    fn = _make_extern(
        "attention_i8",
        _kernel_source("linalg/attention_i8.cc"),
        [q_ty, kt_ty, v_ty, table_ty, out_ty] + [np.int32] * 5,
        compile_flags=compile_flags,
        include_dirs=[*_include_dirs(), _runtime_lib_include()],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Param, Param, Param, Out) + (Param,) * 5,
            reference=partial(attention_i8_ref, **shapes),
            sample=partial(attention_i8_sample, **shapes),
            acc_dtype=np.int32,
            tolerance=Tolerance.exact(
                note="integer arithmetic, every rounding modeled"
            ),
            alignments=((0, 32), (1, 64), (2, 64), (4, 32)),
            ops_per_call=2 * width * seq_len * (key_dim + head_dim),
        ),
        cls=_AttentionI8Kernel,
    )
    fn.load_kv = fn.object_file.bind(
        "attention_i8_load_kv",
        [
            np.ndarray[(key_dim * width,), np.dtype[np.int8]],
            np.ndarray[(head_dim * width,), np.dtype[np.int8]],
            kt_ty,
            v_ty,
            np.int32,
        ],
    )
    return fn

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,
    accurate_exp2: bool = False,
    causal: bool = True,
    window: int = 0,
    packed: bool = False,
    f32_scores: bool = False,
    float_o: bool = False,
    fused: bool = False,
    group: int = 1,
    kv_heads: int = 1,
    soft_cap: float | None = None,
    relative_band: bool = False,
    pair: bool = False,
) -> ExternalFunction

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

One translation unit that includes softmax_aie2p.h and the architecture's own GEMM header (mm_aie2.h or mm_aie2p.h) and exports the symbols an attention dataflow composes over one micro-tile. Unless fused, 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, GEMM's bf16 product on its micro-tile (AIE2 4x8x4, AIE2P 4x8x8) behind an idx_buffer gate (the call runs when key block idx[0] holds a key some row of query block idx[1], its first row idx[2] positions past the block's start, keeps) bound here to [0, 0, 0]. With pv it is the P*V product matmul_bf16_bf16_rowmaj, ungated: GEMM's routine on AIE2's 4x8x4 micro-tile, AIE2P's native 8x8x8. With float_o it is matmul_PV, that product added into a float32 O after O's rows are rescaled by the running state's factor, here 1 on a block every row keeps; test_mha_e2e.py exercises it across key blocks. rescale_O divides that O by the row sums into a bf16 tile. mha.mac_dims(pv=...) answers either product's micro-tile without building a kernel, as mm.mac_dims does.

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

Parameters:

Name	Type	Description	Default
dim_m	int	Rows of the micro-tile (multiple of 8, or of 16 with emulate_bf16_mmul_with_bfp16). An 8-row block, a decode step's KV group of up to 8 heads, runs the online state 8 lanes wide.	64
dim_k	int	Depth of the micro-tile, a multiple of 16: every entry of the object is built, and P*V takes it as its columns. With pair each core's half is that, so a multiple of 32.	64
dim_n	int	Columns of the micro-tile, the key block: 16, 32 or 64.	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
accurate_exp2	bool	As for mha_softmax: partial_softmax and the running-state update take 2^x from a polynomial rather than AIE2P's exp2 instruction or AIE2's cubic.	False
causal	bool	As for mha_softmax: a query keeps no key past its own position.	True
window	int	As for mha_softmax: a query keeps no key more than window positions from it; 0 is no window.	0
packed	bool	As for mha_softmax: every row is a query at the last key. matmul_bf16_bf16_wrapper then computes every block.	False
f32_scores	bool	Keep the QK^T scores in float32 (-DMHA_F32_SCORES): matmul_bf16_bf16_wrapper writes, and partial_softmax reads, a float32 tile, held to float32 summation order. Not with emulate_bf16_mmul_with_bfp16. At 64-key blocks a design places only with P V split over two cores.	False
float_o	bool	With pv, return matmul_PV: P, dim_m by dim_n, times V, dim_n by dim_k, into a float32 O, as IRON's MHA operator binds it.	False
fused	bool	Return attend_block: one key block on one core, S = Q K^T, the online softmax over S in the product's tiles and O += P V, bit for bit the three entries above. KV is the block's keys, each its K row then its V row as a paired cache holds them, in QK^T's B tiling (b_col_maj), so one stream carries both; S and P, dim_m by dim_n each, hold the scores (float32 with f32_scores) and the weights. On AIE2, accurate_exp2 only with f32_scores; on AIE2P, accurate_exp2 (the pipeline's), and no pair, soft_cap or relative_band.	False
group	int	Each position's query heads a block packs, position-major: row i is at the block's first position plus i // group. dim_m is whole positions of it.	1
kv_heads	int	With fused, the KV heads a key block's KV holds, each key's rows of every head in turn as a position-major cache holds them, tiled as one dim_n by 2 dim_k kv_heads block; attend_block takes the first, attend_head any (not built with relative_band or pair). 1, 2, 4 or 8: Peano's backend fails on others.	1
soft_cap	float | None	With fused and f32_scores: each score s becomes soft_cap * tanh(s / soft_cap) * log2(e) (-DMHA_SOFT_CAP, softmax's float32 cap), so the scale attend_block takes is 1: base-e logits, as Gemma's capped attention has them.	None
relative_band	bool	With fused and f32_scores, a query a row (group 1, not packed): attend_block takes QR after KV and adds Q_i . R_d to score (i, j) at distance d = pos_i - key_j, 0 <= d < 16, before the cap (-DMHA_REL_BAND, Gemma 4 audio's relative-position logits). QR is relative_rows(Q, R, QR)'s: dim_m rows of 32 floats, distance d at column 23 - d, from R a key block's KV whose first 16 keys' K rows are the head's relative keys (dim_n at least 16), so they ride the KV stream.	False
pair	bool	With fused, return attend_pair: the block at head dim dim_k on two neighbouring cores, each dim_k / 2 of it (-DMHA_PAIR). Each core's pair_scores(Q, KV, S_part, idx, S_q_eff, S_kv_eff) writes its half's Q K^T in float32, and attend_pair(S0, S1, KV, O, S, P, state, ...) adds the first half's partial and the second's, in that order on both cores, so both hold one S and one P, then runs attend_block's steps over its half of V into its half of O. Q, KV and O are a half's, tiled as attend_block's at dim_k / 2. With bf16 scores, at 16 rows or more: Peano's backend fails on attend_pair at 8.	False

Returns:

Type	Description
ExternalFunction	A MatrixKernel for
ExternalFunction	matmul_bf16_bf16_wrapper, with pv matmul_bf16_bf16_rowmaj
ExternalFunction	and with float_o matmul_PV; an ExternalFunction for
ExternalFunction	attend_block with fused, or attend_pair with pair.

Raises:

Type	Description
ValueError	A dim_m not a positive multiple of 8 (of 16 with emulate_bf16_mmul_with_bfp16), a dim_k not a positive multiple of 16, a dim_n not 16, 32 or 64, a negative window, float_o without pv, fused with pv, without b_col_maj, or with accurate_exp2 but bf16 scores, f32_scores with emulate_bf16_mmul_with_bfp16, a group that does not divide dim_m, kv_heads below 1 or past 1 without fused, soft_cap without fused and f32_scores or outside [2^-100, 2^10], relative_band without fused and f32_scores or with a group or packed, or pair without fused, with kv_heads or relative_band, or at a dim_k not a multiple of 32.
NotImplementedError	fused off AIE2 and AIE2P, in an AIE_KERNELS_PORTABLE build, or on AIE2P without accurate_exp2 or with pair, soft_cap, relative_band or emulate_bf16_mmul_with_bfp16; kv_heads other than 1, 2, 4 or 8, or pair with bf16 scores at 8 rows, which Peano's backend cannot build.

Source code in utils/pyright_stubs/aie/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,
    accurate_exp2: bool = False,
    causal: bool = True,
    window: int = 0,
    packed: bool = False,
    f32_scores: bool = False,
    float_o: bool = False,
    fused: bool = False,
    group: int = 1,
    kv_heads: int = 1,
    soft_cap: float | None = None,
    relative_band: bool = False,
    pair: bool = False,
) -> ExternalFunction:
    """Flash-attention toolkit from ``aie_kernels/linalg/mha.cc``.

    One translation unit that includes ``softmax_aie2p.h`` and the
    architecture's own GEMM header (``mm_aie2.h`` or ``mm_aie2p.h``) and
    exports the symbols an attention dataflow composes over one micro-tile.
    Unless ``fused``, the returned kernel is one of the toolkit's two
    matmuls, both accumulating into ``C``, so it is a
    [`MatrixKernel`][aie.iron.kernels.linalg.MatrixKernel] judged like
    [`mm`][aie.iron.kernels.linalg.mm]. By default that is the ``QK^T`` product
    ``matmul_bf16_bf16_wrapper``, GEMM's bf16 product on its micro-tile
    (AIE2 4x8x4, AIE2P 4x8x8) behind an ``idx_buffer`` gate (the call runs when key block
    ``idx[0]`` holds a key some row of query block ``idx[1]``, its first row
    ``idx[2]`` positions past the block's start, keeps) bound here to
    ``[0, 0, 0]``. With ``pv`` it is the
    ``P*V`` product ``matmul_bf16_bf16_rowmaj``, ungated: GEMM's routine on
    AIE2's 4x8x4 micro-tile, AIE2P's native 8x8x8. With ``float_o`` it is
    ``matmul_PV``, that product added into a float32 ``O`` after ``O``'s
    rows are rescaled by the running state's factor, here 1 on a block every
    row keeps; ``test_mha_e2e.py`` exercises it across key blocks.
    ``rescale_O`` divides that ``O`` by the row sums into a bf16 tile.
    ``mha.mac_dims(pv=...)`` answers either product's micro-tile without
    building a kernel, as ``mm.mac_dims`` does.

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

    Args:
        dim_m: Rows of the micro-tile (multiple of 8, or of 16 with
            ``emulate_bf16_mmul_with_bfp16``). An 8-row block, a decode
            step's KV group of up to 8 heads, runs the online state 8 lanes
            wide.
        dim_k: Depth of the micro-tile, a multiple of 16: every entry of the
            object is built, and ``P*V`` takes it as its columns. With
            ``pair`` each core's half is that, so a multiple of 32.
        dim_n: Columns of the micro-tile, the key block: 16, 32 or 64.
        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`][aie.iron.kernels.linalg.mm]:
            both products use BFP16-based emulation, and ``QK^T``'s
            micro-tile becomes (8, 8, 8).  Ignored on AIE2.
        accurate_exp2: As for [`mha_softmax`][aie.iron.kernels.linalg.mha_softmax]:
            ``partial_softmax`` and the running-state update take 2^x from a
            polynomial rather than AIE2P's exp2 instruction or AIE2's cubic.
        causal: As for [`mha_softmax`][aie.iron.kernels.linalg.mha_softmax]: a
            query keeps no key past its own position.
        window: As for [`mha_softmax`][aie.iron.kernels.linalg.mha_softmax]: a
            query keeps no key more than ``window`` positions from it; 0 is
            no window.
        packed: As for [`mha_softmax`][aie.iron.kernels.linalg.mha_softmax]:
            every row is a query at the last key. ``matmul_bf16_bf16_wrapper``
            then computes every block.
        f32_scores: Keep the ``QK^T`` scores in float32 (``-DMHA_F32_SCORES``):
            ``matmul_bf16_bf16_wrapper`` writes, and ``partial_softmax``
            reads, a float32 tile, held to float32 summation order. Not with
            ``emulate_bf16_mmul_with_bfp16``. At 64-key blocks a design places
            only with ``P V`` split over two cores.
        float_o: With ``pv``, return ``matmul_PV``: ``P``, ``dim_m`` by
            ``dim_n``, times ``V``, ``dim_n`` by ``dim_k``, into a float32
            ``O``, as IRON's MHA operator binds it.
        fused: Return ``attend_block``: one key block on one core, ``S = Q
            K^T``, the online softmax over ``S`` in the product's tiles and
            ``O += P V``, bit for bit the three entries above. ``KV`` is the
            block's keys, each its ``K`` row then its ``V`` row as a paired
            cache holds them, in ``QK^T``'s B tiling (``b_col_maj``), so one
            stream carries both; ``S`` and ``P``, ``dim_m`` by ``dim_n`` each,
            hold the scores (float32 with ``f32_scores``) and the weights.
            On AIE2, ``accurate_exp2`` only with ``f32_scores``; on AIE2P,
            ``accurate_exp2`` (the pipeline's), and no ``pair``, ``soft_cap``
            or ``relative_band``.
        group: Each position's query heads a block packs, position-major:
            row ``i`` is at the block's first position plus ``i // group``.
            ``dim_m`` is whole positions of it.
        kv_heads: With ``fused``, the KV heads a key block's ``KV`` holds,
            each key's rows of every head in turn as a position-major cache
            holds them, tiled as one ``dim_n`` by ``2 dim_k kv_heads`` block;
            ``attend_block`` takes the first, ``attend_head`` any (not built
            with ``relative_band`` or ``pair``). 1, 2, 4 or 8: Peano's backend
            fails on others.
        soft_cap: With ``fused`` and ``f32_scores``: each score ``s`` becomes
            ``soft_cap * tanh(s / soft_cap) * log2(e)`` (``-DMHA_SOFT_CAP``,
            ``softmax``'s float32 cap), so the scale ``attend_block`` takes
            is 1: base-e logits, as Gemma's capped attention has them.
        relative_band: With ``fused`` and ``f32_scores``, a query a row
            (``group`` 1, not ``packed``): ``attend_block`` takes ``QR``
            after ``KV`` and adds ``Q_i . R_d`` to score ``(i, j)`` at
            distance ``d = pos_i - key_j``, ``0 <= d < 16``, before the cap
            (``-DMHA_REL_BAND``, Gemma 4 audio's relative-position logits).
            ``QR`` is ``relative_rows(Q, R, QR)``'s: ``dim_m`` rows of 32
            floats, distance ``d`` at column ``23 - d``, from ``R`` a key
            block's ``KV`` whose first 16 keys' K rows are the head's
            relative keys (``dim_n`` at least 16), so they ride the KV
            stream.
        pair: With ``fused``, return ``attend_pair``: the block at head dim
            ``dim_k`` on two neighbouring cores, each ``dim_k / 2`` of it
            (``-DMHA_PAIR``). Each core's ``pair_scores(Q, KV, S_part, idx,
            S_q_eff, S_kv_eff)`` writes its half's ``Q K^T`` in float32, and
            ``attend_pair(S0, S1, KV, O, S, P, state, ...)`` adds the first
            half's partial and the second's, in that order on both cores, so
            both hold one ``S`` and one ``P``, then runs ``attend_block``'s
            steps over its half of ``V`` into its half of ``O``. ``Q``,
            ``KV`` and ``O`` are a half's, tiled as ``attend_block``'s at
            ``dim_k / 2``. With bf16 scores, at 16 rows or more: Peano's
            backend fails on ``attend_pair`` at 8.

    Returns:
        A [`MatrixKernel`][aie.iron.kernels.linalg.MatrixKernel] for
        ``matmul_bf16_bf16_wrapper``, with ``pv`` ``matmul_bf16_bf16_rowmaj``
        and with ``float_o`` ``matmul_PV``; an ``ExternalFunction`` for
        ``attend_block`` with ``fused``, or ``attend_pair`` with ``pair``.

    Raises:
        ValueError: A ``dim_m`` not a positive multiple of 8 (of 16 with
            ``emulate_bf16_mmul_with_bfp16``), a ``dim_k`` not a positive
            multiple of 16, a ``dim_n`` not 16, 32 or 64, a negative window, ``float_o`` without ``pv``, ``fused`` with ``pv``,
            without ``b_col_maj``, or with ``accurate_exp2`` but bf16 scores,
            ``f32_scores`` with ``emulate_bf16_mmul_with_bfp16``, a ``group``
            that does not divide ``dim_m``, ``kv_heads`` below 1 or past 1
            without ``fused``, ``soft_cap`` without ``fused`` and
            ``f32_scores`` or outside [2^-100, 2^10], ``relative_band``
            without ``fused`` and ``f32_scores`` or with a ``group`` or
            ``packed``, or ``pair`` without ``fused``, with ``kv_heads`` or
            ``relative_band``, or at a ``dim_k`` not a multiple of 32.
        NotImplementedError: ``fused`` off AIE2 and AIE2P, in an
            ``AIE_KERNELS_PORTABLE`` build, or on AIE2P without
            ``accurate_exp2`` or with ``pair``, ``soft_cap``,
            ``relative_band`` or ``emulate_bf16_mmul_with_bfp16``; ``kv_heads`` other
            than 1, 2, 4 or 8, or ``pair`` with bf16 scores at 8 rows, which
            Peano's backend cannot build.
    """
    if float_o and not pv:
        raise ValueError("mha: float_o is the P*V product's, so it needs pv")
    if fused and pv:
        raise ValueError("mha: fused returns its own entry, not pv's")
    if fused and not b_col_maj:
        raise ValueError("mha: fused is QK^T on K as stored, so it needs b_col_maj")
    arch = _tuned_arch()
    if fused and arch == "aie2" and accurate_exp2 and not f32_scores:
        raise ValueError("mha: fused takes the accurate exponential over f32_scores")
    if soft_cap is not None and not (fused and f32_scores):
        raise ValueError("mha: soft_cap is the fused block's, over f32_scores")
    if relative_band and not (fused and f32_scores and group == 1 and not packed):
        raise ValueError(
            "mha: relative_band is the fused block's, over f32_scores, a query "
            "a row (group 1, not packed)"
        )
    if soft_cap is not None and not 2.0**-100 <= soft_cap <= 2.0**10:
        raise ValueError(f"mha: soft_cap must be in [2^-100, 2^10], got {soft_cap}")
    if pair and (not fused or kv_heads != 1 or relative_band or dim_k % 32):
        raise ValueError(
            "mha: pair is the fused block's, one KV head, no relative band, "
            f"over two halves of dim_k that are multiples of 16, got {dim_k}"
        )
    if fused and arch not in ("aie2", "aie2p"):
        raise NotImplementedError("mha: fused is written for tuned AIE2 and AIE2P")
    if (
        fused
        and arch == "aie2p"
        and (
            not accurate_exp2
            or pair
            or soft_cap is not None
            or relative_band
            or emulate_bf16_mmul_with_bfp16
        )
    ):
        raise NotImplementedError(
            "mha: fused on AIE2P takes accurate_exp2, the pipeline's, and no "
            "pair, soft_cap, relative_band or emulate_bf16_mmul_with_bfp16"
        )
    if kv_heads < 1:
        raise ValueError(f"mha: kv_heads must be at least 1, got {kv_heads}")
    if kv_heads > 1 and not fused:
        raise ValueError(f"mha: kv_heads ({kv_heads}) past 1 is fused's")
    if kv_heads not in (1, 2, 4, 8):
        raise NotImplementedError(
            f"mha: kv_heads {kv_heads}: Peano's backend cannot legalize "
            "attend_block's KV pitch (G_UNMERGE_VALUES s20); use 1, 2, 4 or 8"
        )
    if pair and not f32_scores and dim_m == 8:
        raise NotImplementedError(
            "mha: pair with bf16 scores at 8 rows: Peano's backend cannot "
            "legalize attend_pair (G_BITCAST <16 x s16>); use f32_scores or 16 rows"
        )
    if group < 1 or dim_m % group:
        raise ValueError(f"mha: group ({group}) must divide dim_m ({dim_m})")
    for name, v, mult in (
        ("dim_m", dim_m, 16 if emulate_bf16_mmul_with_bfp16 else 8),
        ("dim_k", dim_k, 16),
    ):
        if v <= 0 or v % mult:
            raise ValueError(
                f"mha: {name} must be a positive multiple of {mult}, got {v}"
            )
    if dim_n not in _MHA_KEY_BLOCKS:
        raise ValueError(f"mha: dim_n must be one of {_MHA_KEY_BLOCKS}, got {dim_n}")
    if f32_scores and emulate_bf16_mmul_with_bfp16:
        raise ValueError(
            "mha: f32_scores keeps the scores to float32 summation order, which "
            "a bfp16-emulated product gives up"
        )
    score = np.float32 if f32_scores and not pv else bfloat16
    tile = np.ndarray[(dim_m * dim_n,), np.dtype[score]]
    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[(3,), 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 // 2 if pair else dim_k}",
        f"-DDIM_N={dim_n}",
        *(
            ["-DMHA_F32_SCORES", "-Dbf16_f32_ONLY"]
            if f32_scores
            else ["-Dbf16_bf16_ONLY"]
        ),
    ]
    if b_col_maj:
        flags.append("-DB_COL_MAJ")
    if accurate_exp2:
        flags.append("-DEXP2_BF16_ACCURATE=1")
    flags += _mha_band_flags("mha", causal, window, packed)
    if group != 1:
        flags.append(f"-DMHA_GROUP={group}")
    if kv_heads > 1:
        flags.append(f"-DMHA_KV_HEADS={kv_heads}")
    if soft_cap is not None:
        flags.append(f"-DMHA_SOFT_CAP={float(soft_cap)!r}f")
    if relative_band:
        flags.append("-DMHA_REL_BAND")
    if pair:
        flags.append("-DMHA_PAIR")
        return _mha_pair(
            dim_m,
            dim_k,
            dim_n,
            flags,
            causal,
            window,
            packed,
            group,
            f32_scores,
            accurate_exp2,
            soft_cap,
        )
    if fused:
        return _mha_fused(
            dim_m,
            dim_k,
            dim_n,
            flags,
            causal,
            window,
            packed,
            group,
            kv_heads,
            f32_scores,
            accurate_exp2,
            soft_cap,
            relative_band,
        )
    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 its arch's mm header without C_COL_MAJ, and without
    # B_COL_MAJ unless b_col_maj. matmul_bf16_bf16_rowmaj is always row-major.
    b_col_maj = b_col_maj and not pv
    r, s, t = _MhaFactory.mac_dims(
        pv=pv, emulate_bf16_mmul_with_bfp16=emulate_bf16_mmul_with_bfp16
    )
    if float_o:
        streams = mm_stream_dims(dim_m, dim_n, dim_k, (r, s, t))
        state = np.zeros(4 * dim_m, np.float32)
        state[3 * dim_m :] = 1
        return _make_extern(
            "matmul_PV",
            _kernel_source("linalg/mha.cc"),
            [
                np.ndarray[(dim_m * dim_n,), np.dtype[bfloat16]],
                np.ndarray[(dim_n * dim_k,), np.dtype[bfloat16]],
                np.ndarray[(dim_m * dim_k,), np.dtype[np.float32]],
                np.ndarray[(4 * dim_m,), np.dtype[np.float32]],
                np.int32,
                np.int32,
                idx,
                np.int32,
            ],
            compile_flags=flags,
            cls=MatrixKernel,
            contract=KernelContract(
                trace=Trace.whole_call(),
                layouts=(
                    _tile_layout((dim_m, dim_n), streams.A, block=(r, s)),
                    _tile_layout((dim_n, dim_k), streams.B, block=(s, t)),
                    _tile_layout((dim_m, dim_k), streams.C, inverse=True, block=(r, t)),
                    *([None] * 5),
                ),
                roles=(In, In, InOut, *([Param] * 5)),
                # A later key block (first_iter 1) of a block every row keeps:
                # the rescale runs, by 1, and so does the whole product.
                parameter_bindings=(
                    (3, state),
                    (4, dim_m),
                    (5, 1),
                    (6, np.array([0, 0, 0], np.int32)),
                    (7, dim_n),
                ),
                reference=partial(mm_tile_ref, dim_m=dim_m, dim_k=dim_n, dim_n=dim_k),
                initializers=((2, _zero_output),),
                acc_dtype=np.float32,
                reduction=dim_n,
                tolerance=Tolerance.bounded(
                    partial(
                        _f32_sum_order_bound, dim_m=dim_m, dim_k=dim_n, dim_n=dim_k
                    ),
                    note="float32 O: f32 summation order over dim_n exact bf16 "
                    "products",
                ),
                ops_per_call=2 * dim_m * dim_k * dim_n,
            ),
        )
    streams = mm_stream_dims(dim_m, dim_k, dim_n, (r, s, t), b_col_maj=b_col_maj)
    return _make_extern(
        "matmul_bf16_bf16_rowmaj" if pv else "matmul_bf16_bf16_wrapper",
        _kernel_source("linalg/mha.cc"),
        [a_ty, b_ty, tile] if pv else [a_ty, b_ty, tile, idx],
        compile_flags=flags,
        cls=MatrixKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            layouts=(
                _tile_layout((dim_m, dim_k), streams.A, block=(r, s)),
                _tile_layout(
                    (dim_k, dim_n),
                    streams.B,
                    axes=(1, 0) if b_col_maj else None,
                    block=(s, t),
                ),
                _tile_layout((dim_m, dim_n), streams.C, inverse=True, block=(r, t)),
                *(() if pv else (None,)),
            ),
            roles=(In, In, InOut) if pv else (In, In, InOut, Param),
            parameter_bindings=(() if pv else ((3, np.array([0, 0, 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=(
                Tolerance.bounded(
                    partial(
                        _f32_sum_order_bound, dim_m=dim_m, dim_k=dim_k, dim_n=dim_n
                    ),
                    note="float32 scores: f32 summation order over dim_k exact "
                    "bf16 products",
                )
                if f32_scores and not pv
                else _linalg_tolerance(bfloat16)
            ),
            ops_per_call=2 * dim_m * dim_k * dim_n,
        ),
    )

mha_attend_block_terms

mha_attend_block_terms(
    q,
    kv,
    qr=None,
    *,
    idx,
    s_q_eff,
    s_kv_eff,
    scale,
    dim_m,
    dim_k,
    dim_n,
    causal=True,
    window=0,
    packed=False,
    group=1,
    kv_heads=1,
    f32_scores=False,
    soft_cap=None
)

Numpy reference for mha(fused=True)'s attend_block.

From a zeroed state, O = P V with P = exp2((S - m) s), true exp2 of the scores S as the core keeps them (rounded to bf16, or float32 with f32_scores), m each row's largest kept score or 0, whichever is larger, as a zeroed running maximum leaves it.

Parameters:

Name	Type	Description	Default
q		Query blocks, (calls, dim_m * dim_k) or reshapeable to it.	required
kv		Key blocks as attend_block reads them, each key's K row then its V row for each of kv_heads heads; the first head's are attended.	required
qr		relative_rows' output, which adds each score's relative term, or None.	None
idx		(key block, query block, row offset).	required
s_q_eff		Query rows that exist, from the offset.	required
s_kv_eff		Keys that exist.	required
scale		The call's scale, applied after the maximum.	required
dim_m		Query rows of the block.	required
dim_k		Head dimension.	required
dim_n		Keys of the block.	required
causal		A query keeps no key past its position.	True
window		Keys more than this far from a query are dropped; 0 is none.	0
packed		Every row is a query at the last key.	False
group		Query heads a position packs.	1
kv_heads		Heads kv interleaves.	1
f32_scores		Keep the scores in float32 rather than bf16.	False
soft_cap		Take each score to soft_cap * tanh(s / soft_cap) * log2(e), or None.	None

Returns:

Type	Description
	(O, P): O as (calls, dim_m * dim_k), and the weights,
	(calls, dim_m, dim_n).

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

def mha_attend_block_terms(
    q,
    kv,
    qr=None,
    *,
    idx,
    s_q_eff,
    s_kv_eff,
    scale,
    dim_m,
    dim_k,
    dim_n,
    causal=True,
    window=0,
    packed=False,
    group=1,
    kv_heads=1,
    f32_scores=False,
    soft_cap=None,
):
    """Numpy reference for ``mha(fused=True)``'s ``attend_block``.

    From a zeroed state, ``O = P V`` with ``P = exp2((S - m) s)``, true
    ``exp2`` of the scores ``S`` as the core keeps them (rounded to bf16, or
    float32 with ``f32_scores``), ``m`` each row's largest kept score or 0,
    whichever is larger, as a zeroed running maximum leaves it.

    Args:
        q: Query blocks, ``(calls, dim_m * dim_k)`` or reshapeable to it.
        kv: Key blocks as ``attend_block`` reads them, each key's ``K`` row
            then its ``V`` row for each of ``kv_heads`` heads; the first
            head's are attended.
        qr: ``relative_rows``' output, which adds each score's relative term,
            or None.
        idx: ``(key block, query block, row offset)``.
        s_q_eff: Query rows that exist, from the offset.
        s_kv_eff: Keys that exist.
        scale: The call's scale, applied after the maximum.
        dim_m: Query rows of the block.
        dim_k: Head dimension.
        dim_n: Keys of the block.
        causal: A query keeps no key past its position.
        window: Keys more than this far from a query are dropped; 0 is none.
        packed: Every row is a query at the last key.
        group: Query heads a position packs.
        kv_heads: Heads ``kv`` interleaves.
        f32_scores: Keep the scores in float32 rather than bf16.
        soft_cap: Take each score to ``soft_cap * tanh(s / soft_cap) *
            log2(e)``, or None.

    Returns:
        ``(O, P)``: ``O`` as ``(calls, dim_m * dim_k)``, and the weights,
        ``(calls, dim_m, dim_n)``.
    """
    q = np.asarray(q, np.float64).reshape(-1, dim_m, dim_k)
    kv = np.asarray(kv, np.float64).reshape(-1, dim_n, kv_heads, 2, dim_k)[:, :, 0]
    return _mha_block_terms(
        q @ kv[:, :, 0].swapaxes(1, 2),
        kv[:, :, 1],
        qr,
        idx=idx,
        s_q_eff=s_q_eff,
        s_kv_eff=s_kv_eff,
        scale=scale,
        dim_m=dim_m,
        dim_n=dim_n,
        causal=causal,
        window=window,
        packed=packed,
        group=group,
        f32_scores=f32_scores,
        soft_cap=soft_cap,
    )

mha_attend_pair_terms

mha_attend_pair_terms(
    s0,
    s1,
    kv,
    *,
    idx,
    s_q_eff,
    s_kv_eff,
    scale,
    dim_m,
    dim_k,
    dim_n,
    causal=True,
    window=0,
    packed=False,
    group=1,
    f32_scores=False,
    soft_cap=None
)

Numpy reference for mha(fused=True, pair=True)'s attend_pair.

mha_attend_block_terms from a zeroed state, over the scores s0 + s1.

Parameters:

Name	Type	Description	Default
s0		The float32 scores of the head's first half.	required
s1		The float32 scores of the head's second half.	required
kv		The core's half of the head, each key's K row then its V row; V is attended.	required
idx		(key block, query block, row offset).	required
s_q_eff		Query rows that exist, from the offset.	required
s_kv_eff		Keys that exist.	required
scale		The call's scale, applied after the maximum.	required
dim_m		Query rows of the block.	required
dim_k		The half's head dimension.	required
dim_n		Keys of the block.	required
causal		A query keeps no key past its position.	True
window		Keys more than this far from a query are dropped; 0 is none.	0
packed		Every row is a query at the last key.	False
group		Query heads a position packs.	1
f32_scores		Keep the scores in float32 rather than bf16.	False
soft_cap		Take each score to soft_cap * tanh(s / soft_cap) * log2(e), or None.	None

Returns:

Type	Description
	(O, P): the core's half of O as (calls, dim_m * dim_k),
	and the weights, (calls, dim_m, dim_n).

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

def mha_attend_pair_terms(
    s0,
    s1,
    kv,
    *,
    idx,
    s_q_eff,
    s_kv_eff,
    scale,
    dim_m,
    dim_k,
    dim_n,
    causal=True,
    window=0,
    packed=False,
    group=1,
    f32_scores=False,
    soft_cap=None,
):
    """Numpy reference for ``mha(fused=True, pair=True)``'s ``attend_pair``.

    ``mha_attend_block_terms`` from a zeroed state, over the scores
    ``s0 + s1``.

    Args:
        s0: The float32 scores of the head's first half.
        s1: The float32 scores of the head's second half.
        kv: The core's half of the head, each key's ``K`` row then its ``V``
            row; ``V`` is attended.
        idx: ``(key block, query block, row offset)``.
        s_q_eff: Query rows that exist, from the offset.
        s_kv_eff: Keys that exist.
        scale: The call's scale, applied after the maximum.
        dim_m: Query rows of the block.
        dim_k: The half's head dimension.
        dim_n: Keys of the block.
        causal: A query keeps no key past its position.
        window: Keys more than this far from a query are dropped; 0 is none.
        packed: Every row is a query at the last key.
        group: Query heads a position packs.
        f32_scores: Keep the scores in float32 rather than bf16.
        soft_cap: Take each score to ``soft_cap * tanh(s / soft_cap) *
            log2(e)``, or None.

    Returns:
        ``(O, P)``: the core's half of ``O`` as ``(calls, dim_m * dim_k)``,
        and the weights, ``(calls, dim_m, dim_n)``.
    """
    s0, s1 = (np.asarray(x, np.float64).reshape(-1, dim_m, dim_n) for x in (s0, s1))
    v = np.asarray(kv, np.float64).reshape(-1, dim_n, 2, dim_k)[:, :, 1]
    return _mha_block_terms(
        s0 + s1,
        v,
        idx=idx,
        s_q_eff=s_q_eff,
        s_kv_eff=s_kv_eff,
        scale=scale,
        dim_m=dim_m,
        dim_n=dim_n,
        causal=causal,
        window=window,
        packed=packed,
        group=group,
        f32_scores=f32_scores,
        soft_cap=soft_cap,
    )

mha_softmax

mha_softmax(
    accurate_exp2: bool = False,
    dim_m: int = 64,
    dim_n: int = 64,
    causal: bool = True,
    window: int = 0,
    packed: bool = False,
    scale: float = log2(e) / 8,
    f32_scores: bool = False,
) -> ExternalFunction

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

Writes the block's unnormalized weights P = exp2(A * s - m * s), with s the bf16 scale and m each query row's running maximum of A, unscaled, and updates the running state scale_buffer: [m, m, l, exp2((m_prev - m) * s)], dim_m float32 rows each. Both products are of bf16s, so exact in float32, and the key that sets a row's running maximum takes an exponent of exactly 0. 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, row offset), the query block's first row at position query block * dim_m + row offset, so the queries may start anywhere among the keys. A block no query row keeps a key of is skipped, and S_q_eff/S_kv_eff are the query rows (in blocks' rows, from the offset) and keys whose tails pad the block. Every block takes one path: eight rows at a time, a row one key block of 16, 32 or 64 keys.

Parameters:

Name	Type	Description	Default
accurate_exp2	bool	Take 2^x from a polynomial within 2.8e-6 before its bf16 rounding (-DEXP2_BF16_ACCURATE), in place of AIE2P's exp2 instruction, which is off by up to 6.15%, or AIE2's cubic, 0.74% with the bf16 store. s stays the bf16 scalar the kernel takes, 1.8e-3 from log2(e) / 8.	False
dim_m	int	Query rows of the block (B_q), a multiple of 8.	64
dim_n	int	Keys of the block (B_kv): 16, 32 or 64.	64
causal	bool	A query keeps no key past its own position.	True
window	int	A query keeps no key more than window positions before it, nor, when not causal, after it; 0 is no window. For self-attention, where every query row keeps its own position.	0
packed	bool	Every row is one query at the last key, S_kv_eff - 1, as for one decode step with a KV group's heads as the rows; idx's query block and row offset are then unused.	False
scale	float	The log2-domain factor s, rounded to bf16: log2(e) over the square root of the head dimension. Positive and finite, so a masked score, A's lowest times s, stays below every kept one.	log2(e) / 8
f32_scores	bool	A is float32, as mha's f32_scores leaves it (-DMHA_F32_SCORES): A - m is taken in float32 and times s as three bf16 limbs.	False

Returns:

Type	Description
ExternalFunction	ExternalFunction for partial_softmax.

Raises:

Type	Description
ValueError	A dim_m that is not a positive multiple of 8, a dim_n not 16, 32 or 64, a scale that is not positive and finite, or a negative window.

Source code in utils/pyright_stubs/aie/iron/kernels/linalg.py

def mha_softmax(
    accurate_exp2: bool = False,
    dim_m: int = 64,
    dim_n: int = 64,
    causal: bool = True,
    window: int = 0,
    packed: bool = False,
    scale: float = np.log2(np.e) / 8,
    f32_scores: bool = False,
) -> ExternalFunction:
    """One block of ``mha.cc``'s online softmax, ``partial_softmax``.

    Writes the block's unnormalized weights ``P = exp2(A * s - m * s)``,
    with ``s`` the bf16 ``scale`` and ``m`` each query row's running maximum
    of ``A``, unscaled, and updates the running state ``scale_buffer``:
    ``[m, m, l, exp2((m_prev - m) * s)]``, ``dim_m`` float32 rows each. Both
    products are of bf16s, so exact in float32, and the key that sets a row's
    running maximum takes an exponent of exactly 0. 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, row offset)``, the query block's first row at
    position ``query block * dim_m + row offset``, so the queries may start
    anywhere among the keys. A block no query row keeps a key of is skipped,
    and ``S_q_eff``/``S_kv_eff`` are the query rows (in blocks' rows, from
    the offset) and keys whose tails pad the block. Every block takes one
    path: eight rows at a time, a row one key block of 16, 32 or 64 keys.

    Args:
        accurate_exp2: Take 2^x from a polynomial within 2.8e-6 before its
            bf16 rounding (``-DEXP2_BF16_ACCURATE``), in place of AIE2P's
            exp2 instruction, which is off by up to 6.15%, or AIE2's cubic,
            0.74% with the bf16 store. ``s`` stays the bf16 scalar the
            kernel takes, 1.8e-3 from log2(e) / 8.
        dim_m: Query rows of the block (``B_q``), a multiple of 8.
        dim_n: Keys of the block (``B_kv``): 16, 32 or 64.
        causal: A query keeps no key past its own position.
        window: A query keeps no key more than ``window`` positions before
            it, nor, when not causal, after it; 0 is no window. For
            self-attention, where every query row keeps its own position.
        packed: Every row is one query at the last key, ``S_kv_eff - 1``,
            as for one decode step with a KV group's heads as the rows;
            ``idx``'s query block and row offset are then unused.
        scale: The log2-domain factor ``s``, rounded to bf16: ``log2(e)``
            over the square root of the head dimension. Positive and finite,
            so a masked score, ``A``'s lowest times ``s``, stays below every
            kept one.
        f32_scores: ``A`` is float32, as [`mha`][aie.iron.kernels.linalg.mha]'s
            ``f32_scores`` leaves it (``-DMHA_F32_SCORES``): ``A - m`` is
            taken in float32 and times ``s`` as three bf16 limbs.

    Returns:
        ExternalFunction for ``partial_softmax``.

    Raises:
        ValueError: A ``dim_m`` that is not a positive multiple of 8, a
            ``dim_n`` not 16, 32 or 64, a ``scale`` that is not positive and
            finite, or a negative window.
    """
    if dim_m <= 0 or dim_m % 8:
        raise ValueError(
            f"mha_softmax: dim_m must be a positive multiple of 8, got {dim_m}"
        )
    if dim_n not in _MHA_KEY_BLOCKS:
        raise ValueError(
            f"mha_softmax: dim_n must be one of {_MHA_KEY_BLOCKS}, got {dim_n}"
        )
    scale = float(bfloat16(scale))
    if not 0 < scale < np.inf:
        raise ValueError(f"mha_softmax: scale must be positive and finite, got {scale}")
    band_flags = _mha_band_flags("mha_softmax", causal, window, packed)
    band = dict(dim_m=dim_m, dim_n=dim_n, causal=causal, window=window, packed=packed)
    scores = np.ndarray[
        (dim_m * dim_n,), np.dtype[np.float32 if f32_scores else bfloat16]
    ]
    tile = np.ndarray[(dim_m * dim_n,), np.dtype[bfloat16]]
    state = np.ndarray[(4 * dim_m,), np.dtype[np.float32]]
    idx = np.ndarray[(3,), np.dtype[np.int32]]
    return _make_extern(
        "partial_softmax",
        _kernel_source("linalg/mha.cc"),
        [scores, tile, state, idx, bfloat16, *([np.int32] * 4)],
        # mha.cc calls only mm.cc's bf16 products; the others assert on an
        # 8-row block.
        compile_flags=[
            f"-DDIM_M={dim_m}",
            "-DDIM_K=64",
            f"-DDIM_N={dim_n}",
            *(
                ["-DMHA_F32_SCORES", "-Dbf16_f32_ONLY"]
                if f32_scores
                else ["-Dbf16_bf16_ONLY"]
            ),
            *(["-DEXP2_BF16_ACCURATE=1"] if accurate_exp2 else []),
            *band_flags,
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, InOut, *([Param] * 6)),
            parameter_bindings=((4, scale), (5, dim_m), (6, dim_n)),
            initializers=((2, _zero_output),),
            # 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.
            per_build=tuple(
                (
                    build,
                    partial(mha_softmax_ref, scale=scale, **band),
                    (
                        Tolerance.bounded(
                            partial(
                                _mha_softmax_accurate_bound,
                                scale=scale,
                                # npu1 measured P within 0.5012 ulp, and 0.63
                                # with exp2's f_lo term dropped; npu2 measured
                                # 0.73.
                                half_ulp_p=build.arch == "aie2",
                                **band,
                            ),
                            note="exp2 1e-5 past half an ulp, then each store's "
                            "rounding; P measured 0.5012 ulp on npu1, 0.73 on npu2",
                        )
                        if accurate_exp2
                        else (
                            Tolerance.relative(
                                0.004,
                                2.0**-120,
                                note="exp2_bf16.h cubic, 0.74% worst element on npu1",
                            )
                            if not ARCH_TRAITS[build.arch].native_exp2
                            else Tolerance.relative(
                                0.035,
                                4 * 2.0**-7,
                                note="aie::exp2 interpolant envelope, 6.98% "
                                "(test_mha_e2e.py)",
                            )
                        )
                    ),
                )
                for build in BUILDS
            ),
        ),
    )

mha_softmax_ref

mha_softmax_ref(
    a,
    idx,
    s_q_eff,
    s_kv_eff,
    *,
    scale,
    dim_m=64,
    dim_n=64,
    causal=True,
    window=0,
    packed=False
)

Numpy reference for mha_softmax: (P, scale_buffer).

True exp2 rather than the device's interpolant, from a zeroed running state. m is the unscaled maximum, 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 utils/pyright_stubs/aie/iron/kernels/linalg.py

def mha_softmax_ref(
    a,
    idx,
    s_q_eff,
    s_kv_eff,
    *,
    scale,
    dim_m=64,
    dim_n=64,
    causal=True,
    window=0,
    packed=False,
):
    """Numpy reference for [`mha_softmax`][aie.iron.kernels.linalg.mha_softmax]: ``(P, scale_buffer)``.

    True ``exp2`` rather than the device's interpolant, from a zeroed
    running state. ``m`` is the unscaled maximum, as the kernel stores it. A
    padded row keeps ``m = l = 0``; a skipped or wholly padded block leaves
    the state untouched.
    """
    kv, q, offset = (int(i) for i in np.asarray(idx).ravel())
    rows, cols = np.indices((dim_m, dim_n))
    key = kv * dim_n + cols - (s_kv_eff - 1 if packed else q * dim_m + offset + rows)
    keep = (rows < s_q_eff - q * dim_m) & (cols < s_kv_eff - kv * dim_n)
    if causal or window:
        keep &= key <= (0 if causal else window)
    if window:
        keep &= key >= -window
    if not keep.any():
        return np.zeros((len(a), dim_m * dim_n)), np.zeros(
            (len(a), 4 * dim_m), np.float32
        )
    scores = a.astype(np.float32).reshape(-1, dim_m, dim_n)
    m = np.maximum(np.where(keep, scores, -np.inf).max(axis=2), 0)
    s = np.float32(scale)
    # Centred in float32 before scaling, as the kernel takes A - m: scaled
    # first, float32 scores far from 0 would round apart from their maximum.
    centred = (scores - m[..., None]).astype(np.float64)
    p = np.exp2(np.where(keep, centred * np.float64(s), -np.inf))
    state = np.concatenate([m, m, p.sum(axis=2), np.exp2(-m * s)], 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 utils/pyright_stubs/aie/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`][aie.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 = _PREFILL_GEOM[head_dim]
    # flash_attn_prefill.h's MMUL is aie::mmul<8, 8, 8, bf16, bf16>, the native
    # bf16 micro-tile, not the (4, 8, 8) _MM_MAC_DIMS records for mm.cc.
    r = s = t = 8
    y_ty = np.ndarray[(lq * head_dim,), np.dtype[np.float32]]
    s_ty = np.ndarray[(lq * lk,), np.dtype[bfloat16]]
    v_ty = np.ndarray[(lk * head_dim,), np.dtype[bfloat16]]
    streams = mm_stream_dims(lq, lk, head_dim, (r, s, t))
    # attn_fv walks V n-block-outer, k-block-inner -- the transpose of mm.cc's
    # B block order, so this is built here rather than taken from streams.B. At
    # head_dim 512 LK is 8, making the k term degenerate, so only the 256
    # geometry can tell the two orders apart.
    v_dims = TensorAccessPattern.full((lk, head_dim)).tile((s, t)).permute((1, 0, 2, 3))
    return _make_extern(
        "prefill_fv_step",
        _kernel_source("linalg/flash_attn_prefill.cc"),
        [y_ty, s_ty, v_ty, np.int32],
        compile_flags=[f"-DPREFILL_HEAD_DIM={head_dim}"],
        cls=_ZeroInitializedKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            layouts=(
                _tile_layout((lq, head_dim), streams.C, inverse=True, block=(r, t)),
                _tile_layout((lq, lk), streams.A, block=(r, s)),
                _tile_layout((lk, head_dim), v_dims, block=(s, t)),
                None,
            ),
            roles=(InOut, In, In, Param),
            parameter_bindings=((3, 0),),
            initializers=((0, partial(_zero_output, index=0)),),
            acc_dtype=np.float32,
            reduction=lk,
            # 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.
            per_build=tuple(
                (
                    build,
                    partial(prefill_fv_ref, dim_m=lq, dim_k=lk, dim_n=head_dim),
                    (
                        Tolerance.bounded(
                            partial(
                                _f32_sum_order_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; npu2: unmeasured",
                        )
                        if build.tuned != "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",
                        )
                    ),
                )
                for build in BUILDS
            ),
            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 utils/pyright_stubs/aie/iron/kernels/linalg.py

def prefill_fv_ref(s, v, *, dim_m: int, dim_k: int, dim_n: int):
    """Numpy reference for [`prefill_fv`][aie.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`][aie.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 utils/pyright_stubs/aie/iron/kernels/linalg.py

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

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

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

cascade_mm

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

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

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

Parameters:

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

Raises:

Type	Description
ValueError	When the dtype combination is not supported.

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

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

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

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

cascade_mm_put

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

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

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

Source code in utils/pyright_stubs/aie/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`][aie.iron.kernels.linalg.cascade_mm]: ``A * B`` onto the cascade stream.

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

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

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

fused_mm

fused_mm(
    *,
    dim_m: int = 32,
    dim_k: int = 32,
    dim_n: int = 16,
    band_m: int = 16,
    chunk_k: int = 16,
    out_chunk: int = 64,
    epilogue: str = "none",
    clamp: tuple[float, float] | None = None,
    bfp16_b: bool = False,
    b_col_maj: bool = False,
    emulate_bf16_mmul_with_bfp16: bool = False,
    epilogue_modes: tuple[str, ...] | None = None,
    rounding: str = "conv_even",
    gelu: str = "fp32",
    mmul_shape: tuple[int, int, int] | None = None,
    c_depth: int = 2,
    step_markers: bool = False
) -> _FusedMMKernel

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

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

Parameters:

Name	Type	Description	Default
dim_m	int	Rows of A and C.	32
dim_k	int	The reduction length.	32
dim_n	int	Columns of B and C.	16
band_m	int	Rows of A per reduction band.	16
chunk_k	int	Columns of A per reduction step.	16
out_chunk	int	Elements of C per drain chunk, drained at a depth of two.	64
epilogue	str	none, gelu (the kernel's sigmoid approximation, not the tanh GELU curve), silu or sigmoid.	'none'
clamp	tuple[float, float] | None	An optional finite (min, max) after the activation, before the bf16 conversion.	None
bfp16_b	bool	(aie2p) The prepacked-B form amd/IRON's flm GEMM builds: B arrives as bfp16ebs8 blocks the host packs once, in the rounding mode as IRON's pack_b does, the mmul is 8x8x8, and the core converts A to bfp16 itself. Implies emulate_bf16_mmul_with_bfp16.	False
b_col_maj	bool	Take B as stored transposed, the (n, k) layout of a checkpoint's weights, with an (r, 8, 8) mmul (r 8 on bfp16 macs): each k chunk holds k panels of s rows, every n row of a panel s contiguous k, which a DMA builds from (n, k) storage in 16-byte runs.	False
emulate_bf16_mmul_with_bfp16	bool	(aie2p) Multiply on bfp16 macs, as aie_api's emulated bf16 mmul does: the core converts A and B to bfp16ebs8, each shared exponent spanning 8 consecutive k. The mmul defaults to 8x8x8, which b_col_maj needs.	False
epilogue_modes	tuple[str, ...] | None	The activations compiled in, for a caller that picks one at run time through mm_fused_epilogue_chunk's mode argument; (epilogue,) by default, and must hold epilogue. On aie2 an activation links the tanh tables, 5 KB of L1.	None
rounding	str	The core's rounding mode for every f32 -> bf16 and bf16 -> bfp16 conversion: conv_even, or floor, the mode a core powers up in.	'conv_even'
gelu	str	fp32 applies gelu to the f32 accumulator and rounds once; bf16_steps rounds to bf16 and again after each step of x * sigmoid(1.702x), as FastFlowLM's shipped mm overlay does (bit for bit on AIE2P with rounding="floor").	'fp32'
mmul_shape	tuple[int, int, int] | None	Overrides the mmul's (r, s, t).	None
c_depth	int	The depth of the C fifo, which a caller's drain loop unrolls by.	2
step_markers	bool	Bracket each init, k step and drain chunk with event markers, for a caller that calls those entry points itself, rather than one pair around the whole fused_mm_tile call.	False

Returns:

Type	Description
_FusedMMKernel	fused_mm_tile, carrying every entry point of its object as an
_FusedMMKernel	attribute named by its C symbol: mm_fused_acc_init,
_FusedMMKernel	mm_fused_k_step, mm_fused_epilogue_chunk and
_FusedMMKernel	fused_mm_tile (itself). The reference multiplies the operands as
_FusedMMKernel	the core sees them.

Source code in utils/pyright_stubs/aie/iron/kernels/fused.py

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

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

    Args:
        dim_m: Rows of A and C.
        dim_k: The reduction length.
        dim_n: Columns of B and C.
        band_m: Rows of A per reduction band.
        chunk_k: Columns of A per reduction step.
        out_chunk: Elements of C per drain chunk, drained at a depth of two.
        epilogue: ``none``, ``gelu`` (the kernel's sigmoid approximation,
            not the tanh GELU curve), ``silu`` or ``sigmoid``.
        clamp: An optional finite ``(min, max)`` after the activation,
            before the bf16 conversion.
        bfp16_b: (aie2p) The prepacked-B form amd/IRON's flm GEMM builds: B
            arrives as bfp16ebs8 blocks the host packs once, in the
            ``rounding`` mode as IRON's ``pack_b`` does, the mmul is 8x8x8,
            and the core converts A to bfp16 itself. Implies
            ``emulate_bf16_mmul_with_bfp16``.
        b_col_maj: Take B as stored transposed, the ``(n, k)`` layout of a
            checkpoint's weights, with an ``(r, 8, 8)`` mmul (r 8 on bfp16
            macs): each k chunk holds k panels of ``s`` rows, every n row of
            a panel ``s`` contiguous k, which a DMA builds from ``(n, k)``
            storage in 16-byte runs.
        emulate_bf16_mmul_with_bfp16: (aie2p) Multiply on bfp16 macs, as
            aie_api's emulated bf16 mmul does: the core converts A and B to
            bfp16ebs8, each shared exponent spanning 8 consecutive k. The
            mmul defaults to 8x8x8, which ``b_col_maj`` needs.
        epilogue_modes: The activations compiled in, for a caller that picks
            one at run time through ``mm_fused_epilogue_chunk``'s mode
            argument; ``(epilogue,)`` by default, and must hold
            ``epilogue``. On aie2 an activation links the tanh tables, 5 KB
            of L1.
        rounding: The core's rounding mode for every f32 -> bf16 and bf16 ->
            bfp16 conversion: ``conv_even``, or ``floor``, the mode a core
            powers up in.
        gelu: ``fp32`` applies gelu to the f32 accumulator and rounds once;
            ``bf16_steps`` rounds to bf16 and again after each step of ``x *
            sigmoid(1.702x)``, as FastFlowLM's shipped mm overlay does (bit
            for bit on AIE2P with ``rounding="floor"``).
        mmul_shape: Overrides the mmul's ``(r, s, t)``.
        c_depth: The depth of the C fifo, which a caller's drain loop
            unrolls by.
        step_markers: Bracket each init, k step and drain chunk with event
            markers, for a caller that calls those entry points itself,
            rather than one pair around the whole ``fused_mm_tile`` call.

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

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

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

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

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

    # Block (i, j) of a k chunk at (i * colB + j), each B^T: n-major.
    def pack_b_t(b):
        return (
            b.reshape(-1, dim_k // chunk_k, chunk_k // s, s, dim_n // t, t)
            .transpose(0, 1, 2, 4, 5, 3)
            .reshape(len(b), -1)
        )

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

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

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

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

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

    def operands(a, b):
        a = a.reshape(-1, dim_m, dim_k).astype(np.float32)
        b = b.reshape(-1, dim_k, dim_n).astype(np.float32)
        if emulated:
            a = bfp.quantize(a, rounding=rounding)
        if bfp16_b:
            b = unpack_b_bfp(pack_b_bfp(b))
        elif emulated:
            b = bfp.quantize(b.swapaxes(1, 2), rounding=rounding).swapaxes(1, 2)
        return a.astype(np.float64), b.astype(np.float64)

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

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

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

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

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

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

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

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

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

Convolution

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

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

conv2dk1_ref

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

Numpy reference for conv2dk1: 1x1 conv, requantized.

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

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk1_ref(x, weights, input_width, input_channels, output_channels, scale):
    """Numpy reference for [`conv2dk1`][aie.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(round(sum_ic x * w / 2**scale)), rounding half away from zero (symmetric_inf).

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

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

    The layouts of [`conv2dk1_ref`][aie.iron.kernels.conv.conv2dk1_ref] with an
    ``int8`` output and no ReLU: ``out = sat_i8(round(sum_ic x * w /
    2**scale))``, rounding half away from zero (``symmetric_inf``).
    """
    W, IC, OC = int(input_width), int(input_channels), int(output_channels)
    acc, lead = _conv1x1_acc(x, weights, W, IC, OC)
    return _requant_bias(acc, scale, "symmetric_inf").reshape(*lead, W * OC)

conv2dk1_i8_bias_ref

conv2dk1_i8_bias_ref(
    x,
    weights,
    bias,
    input_width,
    input_channels,
    output_channels,
    scale,
    rounding,
)

Numpy reference for conv2dk1_i8 with bias.

The layouts of conv2dk1_ref; bias is one int32 per output channel in accumulator units, added before the shift: out = sat_i8(round((sum_ic x * w + bias) / 2**scale)), rounding half to even (conv_even) or away from zero (symmetric_inf).

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk1_i8_bias_ref(
    x, weights, bias, input_width, input_channels, output_channels, scale, rounding
):
    """Numpy reference for [`conv2dk1_i8`][aie.iron.kernels.conv.conv2dk1_i8] with ``bias``.

    The layouts of [`conv2dk1_ref`][aie.iron.kernels.conv.conv2dk1_ref]; ``bias``
    is one int32 per output channel in accumulator units, added before the
    shift: ``out = sat_i8(round((sum_ic x * w + bias) / 2**scale))``,
    rounding half to even (``conv_even``) or away from zero
    (``symmetric_inf``).
    """
    W, IC, OC = int(input_width), int(input_channels), int(output_channels)
    acc, lead = _conv1x1_acc(x, weights, W, IC, OC)
    acc = acc + np.asarray(bias, dtype=np.int64).reshape(OC // 8, 1, 8)
    return _requant_bias(acc, scale, rounding).reshape(*lead, W * OC)

conv2dk1_i8_even_ref

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

conv2dk1_i8_ref rounding half to even (conv_even).

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk1_i8_even_ref(
    x, weights, input_width, input_channels, output_channels, scale
):
    """[`conv2dk1_i8_ref`][aie.iron.kernels.conv.conv2dk1_i8_ref] rounding half to even (``conv_even``)."""
    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)

conv2dk1_i8_bias_table_ref

conv2dk1_i8_bias_table_ref(
    x,
    weights,
    bias,
    table,
    input_width,
    input_channels,
    output_channels,
    scale,
    rounding,
)

conv2dk1_i8_bias_ref, then table[out + 128].

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk1_i8_bias_table_ref(
    x,
    weights,
    bias,
    table,
    input_width,
    input_channels,
    output_channels,
    scale,
    rounding,
):
    """[`conv2dk1_i8_bias_ref`][aie.iron.kernels.conv.conv2dk1_i8_bias_ref], then ``table[out + 128]``."""
    out = conv2dk1_i8_bias_ref(
        x, weights, bias, input_width, input_channels, output_channels, scale, rounding
    )
    return np.asarray(table, dtype=np.int8)[out.astype(np.int64) + 128]

conv2dk1_i8_bias_slices_ref

conv2dk1_i8_bias_slices_ref(
    x,
    weights,
    bias,
    input_width,
    input_channels,
    output_channels,
    scale,
    *,
    rounding,
    input_slice
)

conv2dk1_i8_bias_ref summed over input-channel slices.

Call s reads its own row of the input_slice channels [s * input_slice, (s + 1) * input_slice) through weights ([OC/8][input_slice/8][ic8][oc8]); the sums of the calls, seeded with bias, give the one output tile. The kernel writes it on slice input_channels // input_slice - 1, so a run makes that many calls.

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk1_i8_bias_slices_ref(
    x,
    weights,
    bias,
    input_width,
    input_channels,
    output_channels,
    scale,
    *,
    rounding,
    input_slice,
):
    """[`conv2dk1_i8_bias_ref`][aie.iron.kernels.conv.conv2dk1_i8_bias_ref] summed over input-channel slices.

    Call ``s`` reads its own row of the ``input_slice`` channels
    ``[s * input_slice, (s + 1) * input_slice)`` through ``weights``
    (``[OC/8][input_slice/8][ic8][oc8]``); the sums of the calls, seeded
    with ``bias``, give the one output tile.
    The kernel writes it on slice ``input_channels // input_slice - 1``, so
    a run makes that many calls.
    """
    del input_channels
    W, OC = int(input_width), int(output_channels)
    x = np.asarray(x)
    acc = 0
    for row in x:
        acc = acc + _conv1x1_acc(row.ravel(), weights, W, input_slice, OC)[0]
    acc = acc + np.asarray(bias, dtype=np.int64).reshape(OC // 8, 1, 8)
    if rounding == "conv_even":
        out = _requant_even(acc, scale, -128, 127, np.int8)
    else:
        half = (1 << (int(scale) - 1)) if int(scale) > 0 else 0
        out = np.clip(np.sign(acc) * ((np.abs(acc) + half) >> int(scale)), -128, 127)
    return out.astype(np.int8).reshape(1, W * OC)

conv2dk1_i8_bias_table_slices_ref

conv2dk1_i8_bias_table_slices_ref(
    x,
    weights,
    bias,
    table,
    input_width,
    input_channels,
    output_channels,
    scale,
    *,
    rounding,
    input_slice
)

conv2dk1_i8_bias_slices_ref, then table[out + 128].

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk1_i8_bias_table_slices_ref(
    x,
    weights,
    bias,
    table,
    input_width,
    input_channels,
    output_channels,
    scale,
    *,
    rounding,
    input_slice,
):
    """[`conv2dk1_i8_bias_slices_ref`][aie.iron.kernels.conv.conv2dk1_i8_bias_slices_ref], then ``table[out + 128]``."""
    out = conv2dk1_i8_bias_slices_ref(
        x,
        weights,
        bias,
        input_width,
        input_channels,
        output_channels,
        scale,
        rounding=rounding,
        input_slice=input_slice,
    )
    return np.asarray(table, dtype=np.int8)[out.astype(np.int64) + 128]

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 utils/pyright_stubs/aie/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`][aie.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 utils/pyright_stubs/aie/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`][aie.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 utils/pyright_stubs/aie/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`][aie.iron.kernels.conv.conv2dk1_skip_init]: 1x1 conv plus a projected residual.

    Like [`conv2dk1_skip_ref`][aie.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 utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk14_ref(
    x, weights, input_width, input_channels, output_channels, kernel_width, scale
):
    """Numpy reference for [`conv2dk14`][aie.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 utils/pyright_stubs/aie/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`][aie.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,
        ),
    )

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 utils/pyright_stubs/aie/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`][aie.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 utils/pyright_stubs/aie/iron/kernels/conv.py

def dwconv1d(
    seq_len: int = 1024, kernel_size: int = 9, bias: bool = True
) -> ExternalFunction:
    """Compatibility alias for [`dwconv1d_channels_first`][aie.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 utils/pyright_stubs/aie/iron/kernels/conv.py

def dwconv1d_ref(x_pad, w, seq_len, *, kernel_size: int, bias: bool):
    """Compatibility alias for [`dwconv1d_channels_first_ref`][aie.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 utils/pyright_stubs/aie/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`][aie.iron.kernels.conv.dwconv1d_channels_first];
    layout picks the vectorization axis, so neither subsumes the other. See
    the "Choosing a depthwise conv1d" section of ``aie_kernels/README.md``.

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

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

    Args:
        channels: Channels per call (multiple of 32).
        clamp: Clamp the result to the runtime ``lo``/``hi`` buffers.
    """
    if channels <= 0 or channels % 32:
        raise ValueError(
            "dwconv1d_channels_last: channels must be a positive multiple of the "
            f"32-lane store, got {channels}"
        )
    _TAPS = 5
    plane_ty = np.ndarray[(channels,), np.dtype[bfloat16]]
    lim_ty = np.ndarray[(1,), np.dtype[np.float32]]
    return _make_extern(
        "dwconv1d_channels_last_k5_bf16",
        _kernel_source("conv/dwconv1d_channels_last.cc"),
        [*([plane_ty] * 2 * _TAPS), plane_ty, lim_ty, lim_ty],
        compile_flags=[
            f"-DDWCONV1D_CL_C={channels}",
            f"-DDWCONV1D_CL_CLAMP={int(clamp)}",
        ],
        contract=KernelContract(
            alignments=_vector_args(*range(2 * _TAPS + 1)),
            trace=Trace.whole_call(),
            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; npu1: "
                "unmeasured",
            ),
            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 utils/pyright_stubs/aie/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`][aie.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 utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

    Returns:
        ExternalFunction configured for the conv2dk1 kernel.

    Raises:
        ValueError: When ``act_dtype`` is not ``np.int8`` or ``np.uint8``.
    """
    func_name, flags = _conv_act_dtype_info(
        "conv2dk1", act_dtype, factory_name="conv2dk1"
    )
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[act_dtype]]
    wt_ty = np.ndarray[(input_channels * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    return _make_extern(
        func_name,
        _kernel_source("conv/conv2dk1.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(4)],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            alignments=_vector_args(0, 1, 2),
            trace=Trace.whole_call(),
            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,
    bias: bool = False,
    table: bool = False,
    stride: int = 1,
    rounding: str | None = None,
    table_copies: int = 2,
    input_slice: int | None = None,
) -> ExternalFunction

3x3 convolution kernel.

With bias it is ONNX QLinearConv's arithmetic on int8 activations for a power-of-two per-tensor output scale and zero zero-points: an int32 per output channel, in accumulator units, seeds the sum, the output is int8, the stride may be 2 and the width any multiple of 4 (8 at stride 2). table then maps each result through a 256-entry int8 table indexed by value + 128.

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
bias	bool	Take an int32 bias per output channel after the weights, and write int8.	False
table	bool	Take a 256-entry int8 table after the bias. Needs bias.	False
stride	int	1, or 2 with bias.	1
rounding	str | None	How the shift rounds: "positive_inf" (the uint8 kernel), or with bias "symmetric_inf" or "conv_even" (half to even, as ONNX QuantizeLinear). By default "conv_even" with bias, "positive_inf" without.	None
table_copies	int	With table, how many copies of the table the kernel keeps in data memory: 2 (two banks, the LUT read fetches both halves at once) or 1 (1 KB less, the reads serialize).	2
input_slice	int | None	With bias, the input channels one call takes, so a streamed weight chunk holds 9 * input_slice * output_channels bytes. A call's lines are then that slice's own rows. The kernel takes an int32 sums buffer, one per output, before the output, and its slice in the last scalar: slice 0 starts from the bias, the last ends in the output, and the others carry the sums. None takes every input channel in one call.	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, when input_width is not 32 without bias, or on an option only the bias variant has.

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

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,
    bias: bool = False,
    table: bool = False,
    stride: int = 1,
    rounding: str | None = None,
    table_copies: int = 2,
    input_slice: int | None = None,
) -> ExternalFunction:
    """3x3 convolution kernel.

    With ``bias`` it is ONNX QLinearConv's arithmetic on ``int8``
    activations for a power-of-two per-tensor output scale and zero
    zero-points: an ``int32`` per output channel, in accumulator units, seeds
    the sum, the output is ``int8``, the stride may be 2 and the width any
    multiple of 4 (8 at stride 2). ``table`` then maps each result through a
    256-entry ``int8`` table indexed by ``value + 128``.

    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).
        bias: Take an ``int32`` bias per output channel after the weights,
            and write ``int8``.
        table: Take a 256-entry ``int8`` table after the bias. Needs ``bias``.
        stride: 1, or 2 with ``bias``.
        rounding: How the shift rounds: ``"positive_inf"`` (the uint8 kernel),
            or with ``bias`` ``"symmetric_inf"`` or ``"conv_even"`` (half to
            even, as ONNX QuantizeLinear). By default ``"conv_even"`` with
            ``bias``, ``"positive_inf"`` without.
        table_copies: With ``table``, how many copies of the table the
            kernel keeps in data memory: 2 (two banks, the LUT read fetches
            both halves at once) or 1 (1 KB less, the reads serialize).
        input_slice: With ``bias``, the input channels one call takes, so a
            streamed weight chunk holds ``9 * input_slice * output_channels``
            bytes. A call's lines are then that slice's own rows. The kernel
            takes an ``int32`` sums buffer, one per output, before the output,
            and its slice in the last scalar:
            slice 0 starts from the bias, the last ends in the output, and
            the others carry the sums. ``None`` takes every input channel in
            one call.

    Returns:
        ExternalFunction configured for the conv2dk3 kernel.

    Raises:
        ValueError: When ``act_dtype`` is not ``np.int8`` or ``np.uint8``,
            when ``input_width`` is not 32 without ``bias``, or on an option
            only the ``bias`` variant has.

    """
    if bias:
        return _conv2dk3_bias(
            input_width,
            input_channels,
            output_channels,
            act_dtype,
            table,
            stride,
            rounding or "conv_even",
            table_copies,
            input_slice,
        )
    if table or stride != 1 or rounding not in (None, "positive_inf") or input_slice:
        raise ValueError("conv2dk3: table, stride, rounding and input_slice need bias")
    if input_width != 32:
        # Every compiled variant of the vector kernel (aie2, aie2p, and the
        # portable branch) hard-codes a 32-pixel row inside conv2dk3.cc and
        # ignores the runtime_input_width arg; a different width would
        # silently compute over the wrong pixels instead of raising.
        raise ValueError(f"conv2dk3: input_width must be 32, got {input_width}")
    func_name, flags = _conv_act_dtype_info(
        "conv2dk3", act_dtype, factory_name="conv2dk3"
    )
    if weight_output_channels is None:
        weight_output_channels = output_channels
    line_size = input_width * input_channels
    line_ty = np.ndarray[(line_size,), np.dtype[act_dtype]]
    wt_ty = np.ndarray[
        (3 * 3 * input_channels * weight_output_channels,), np.dtype[np.int8]
    ]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    return _make_extern(
        func_name,
        _kernel_source("conv/conv2dk3.cc"),
        [line_ty, line_ty, line_ty, wt_ty, out_ty, *_i32s(8)],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels)
        + ["-DCONV_KERNEL_WIDTH=3", "-DCONV_KERNEL_HEIGHT=3"],
        include_dirs=[*_include_dirs(), _runtime_lib_include()],
        contract=KernelContract(
            alignments=_vector_args(0, 1, 2, 3, 4),
            trace=Trace.whole_call(),
            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,
        ),
    )

conv2dk3_bias_ref

conv2dk3_bias_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
    *,
    stride,
    rounding
)

Numpy reference for conv2dk3 with bias.

The layouts of conv2dk3_ref at stride 1 or 2; bias is one int32 per output channel in accumulator units: out = sat_i8(round((sum + bias) / 2**scale)), rounding half to even (conv_even) or away from zero (symmetric_inf).

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk3_bias_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
    *,
    stride,
    rounding,
):
    """Numpy reference for [`conv2dk3`][aie.iron.kernels.conv.conv2dk3] with ``bias``.

    The layouts of [`conv2dk3_ref`][aie.iron.kernels.conv.conv2dk3_ref] at
    ``stride`` 1 or 2; ``bias`` is one int32 per output channel in
    accumulator units: ``out = sat_i8(round((sum + bias) / 2**scale))``,
    rounding half to even (``conv_even``) or away from zero
    (``symmetric_inf``).
    """
    acc, lead, W, OC = _conv3x3_acc(
        (line0, line1, line2),
        weights,
        input_width,
        input_channels,
        output_channels,
        kernel_width,
        check,
        channel_offset,
        stride,
    )
    del kernel_height
    acc = acc + np.asarray(bias, dtype=np.int64).reshape(OC // 8, 1, 8)
    return _requant_bias(acc, scale, rounding).reshape(*lead, W * OC)

conv2dk3_bias_table_ref

conv2dk3_bias_table_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    table,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
    *,
    stride,
    rounding
)

conv2dk3_bias_ref, then table[out + 128].

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk3_bias_table_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    table,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
    *,
    stride,
    rounding,
):
    """[`conv2dk3_bias_ref`][aie.iron.kernels.conv.conv2dk3_bias_ref], then ``table[out + 128]``."""
    out = conv2dk3_bias_ref(
        line0,
        line1,
        line2,
        weights,
        bias,
        input_width,
        input_channels,
        output_channels,
        kernel_width,
        kernel_height,
        check,
        scale,
        channel_offset,
        stride=stride,
        rounding=rounding,
    )
    return np.asarray(table, dtype=np.int8)[out.astype(np.int64) + 128]

conv2dk3_bias_slices_ref

conv2dk3_bias_slices_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    *,
    stride,
    rounding,
    input_slice
)

conv2dk3_bias_ref summed over input-channel slices.

Call s reads its own lines of the input_slice channels [s * input_slice, (s + 1) * input_slice) through weights ([OC/8][input_slice/8][3][3][ic8][oc8]); the sums of the calls, seeded with bias, give the one output tile. The kernel writes it on slice input_channels // input_slice - 1, so a run makes that many calls.

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk3_bias_slices_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    *,
    stride,
    rounding,
    input_slice,
):
    """[`conv2dk3_bias_ref`][aie.iron.kernels.conv.conv2dk3_bias_ref] summed over input-channel slices.

    Call ``s`` reads its own lines of the ``input_slice`` channels
    ``[s * input_slice, (s + 1) * input_slice)`` through ``weights``
    (``[OC/8][input_slice/8][3][3][ic8][oc8]``); the sums of the calls,
    seeded with ``bias``, give the one output tile. The kernel writes it on
    slice ``input_channels // input_slice - 1``, so a run makes that many
    calls.
    """
    del kernel_height, input_channels
    OC = int(output_channels)
    Wo = int(input_width) // stride
    lines = [np.asarray(line) for line in (line0, line1, line2)]
    acc = 0
    for s in range(len(lines[0])):
        sums = _conv3x3_acc(
            [line[s].ravel() for line in lines],
            weights,
            input_width,
            input_slice,
            output_channels,
            kernel_width,
            check,
            0,
            stride,
        )[0]
        acc = acc + sums
    acc = acc + np.asarray(bias, dtype=np.int64).reshape(OC // 8, 1, 8)
    if rounding == "conv_even":
        out = _requant_even(acc, scale, -128, 127, np.int8)
    else:
        half = (1 << (int(scale) - 1)) if int(scale) > 0 else 0
        out = np.clip(np.sign(acc) * ((np.abs(acc) + half) >> int(scale)), -128, 127)
    return out.astype(np.int8).reshape(1, Wo * OC)

conv2dk3_bias_table_slices_ref

conv2dk3_bias_table_slices_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    table,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    *,
    stride,
    rounding,
    input_slice
)

conv2dk3_bias_slices_ref, then table[out + 128].

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk3_bias_table_slices_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    table,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    *,
    stride,
    rounding,
    input_slice,
):
    """[`conv2dk3_bias_slices_ref`][aie.iron.kernels.conv.conv2dk3_bias_slices_ref], then ``table[out + 128]``."""
    out = conv2dk3_bias_slices_ref(
        line0,
        line1,
        line2,
        weights,
        bias,
        input_width,
        input_channels,
        output_channels,
        kernel_width,
        kernel_height,
        check,
        scale,
        stride=stride,
        rounding=rounding,
        input_slice=input_slice,
    )
    return np.asarray(table, dtype=np.int8)[out.astype(np.int64) + 128]

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 utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

    Returns:
        ExternalFunction configured for the conv2dk1_skip kernel.

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

    Note:
        The activations are ``uint8`` in two half-channel tensors whatever
        ``act_dtype``, which types the residual (``skip``) only. The generic
        harness packs tensors of one type into one fifo, so an ``int8``
        residual streams beside the ``uint8`` activations in a second fifo.
    """
    func_name, flags = _conv_act_dtype_info(
        "conv2dk1_skip", act_dtype, factory_name="conv2dk1_skip"
    )
    half_ch = input_channels // 2
    in0_ty = np.ndarray[(input_width * half_ch,), np.dtype[np.uint8]]
    in1_ty = np.ndarray[(input_width * half_ch,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(input_channels * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    skip_ty = np.ndarray[(input_width * output_channels,), np.dtype[act_dtype]]
    return _make_extern(
        func_name,
        _kernel_source("conv/conv2dk1_skip.cc"),
        [in0_ty, in1_ty, wt_ty, out_ty, skip_ty, *_i32s(5)],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            alignments=_vector_args(0, 1, 2, 3, 4),
            trace=Trace.whole_call(),
            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,
    bias: bool = False,
    table: bool = False,
    rounding: str | None = None,
    table_copies: int = 2,
    input_slice: int | None = None,
) -> ExternalFunction

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

With bias it is ONNX QLinearConv's arithmetic for a power-of-two per-tensor output scale and zero zero-points: an int32 per output channel, in accumulator units, seeds the sum before the shift. table then maps each int8 result through a 256-entry int8 table indexed by value + 128 (any element-wise activation, such as a quantized SiLU).

Parameters:

Name	Type	Description	Default
input_width	int	Spatial width of the input; a multiple of 32, or of 4 with bias.	32
input_channels	int	Number of input channels.	64
output_channels	int	Number of output channels.	64
bias	bool	Take an int32 bias per output channel after the weights.	False
table	bool	Take a 256-entry int8 table after the bias. Needs bias.	False
rounding	str | None	How the shift rounds: "symmetric_inf" (half away from zero) or "conv_even" (half to even, as ONNX QuantizeLinear). By default "conv_even" with bias, "symmetric_inf" without.	None
table_copies	int	With table, how many copies of the table the kernel keeps in data memory: 2 (two banks, the LUT read fetches both halves at once) or 1 (1 KB less, the reads serialize).	2
input_slice	int | None	With bias, the input channels one call takes, so a streamed weight chunk holds input_slice * output_channels bytes. A call's input is then that slice's own row, so a concat's parts can be passed apart. The kernel takes an int32 sums buffer, one per output, before the output, and its slice in the last scalar: slice 0 starts from the bias, the last ends in the output, and the others carry the sums. None takes every input channel in one call.	None

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk1_i8 kernel.

Raises:

Type	Description
ValueError	On an unknown rounding, table or input_slice without bias, a slice that does not split the input channels, or a width the variant cannot tile.

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def conv2dk1_i8(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    bias: bool = False,
    table: bool = False,
    rounding: str | None = None,
    table_copies: int = 2,
    input_slice: int | None = None,
) -> ExternalFunction:
    """1x1 convolution kernel with int8 activations/weights/output.

    With ``bias`` it is ONNX QLinearConv's arithmetic for a power-of-two
    per-tensor output scale and zero zero-points: an int32 per output
    channel, in accumulator units, seeds the sum before the shift. ``table``
    then maps each int8 result through a 256-entry int8 table indexed by
    ``value + 128`` (any element-wise activation, such as a quantized SiLU).

    Args:
        input_width: Spatial width of the input; a multiple of 32, or of 4
            with ``bias``.
        input_channels: Number of input channels.
        output_channels: Number of output channels.
        bias: Take an ``int32`` bias per output channel after the weights.
        table: Take a 256-entry ``int8`` table after the bias. Needs ``bias``.
        rounding: How the shift rounds: ``"symmetric_inf"`` (half away from
            zero) or ``"conv_even"`` (half to even, as ONNX QuantizeLinear).
            By default ``"conv_even"`` with ``bias``, ``"symmetric_inf"``
            without.
        table_copies: With ``table``, how many copies of the table the
            kernel keeps in data memory: 2 (two banks, the LUT read fetches
            both halves at once) or 1 (1 KB less, the reads serialize).
        input_slice: With ``bias``, the input channels one call takes, so a
            streamed weight chunk holds ``input_slice * output_channels``
            bytes. A call's input is then that slice's own row, so a concat's
            parts can be passed apart. The kernel takes an ``int32`` sums
            buffer, one per output, before the output, and its slice in the
            last scalar:
            slice 0 starts from the bias, the last ends in the output, and
            the others carry the sums. ``None`` takes every input channel in
            one call.

    Returns:
        ExternalFunction configured for the conv2dk1_i8 kernel.

    Raises:
        ValueError: On an unknown ``rounding``, ``table`` or ``input_slice``
            without ``bias``, a slice that does not split the input channels,
            or a width the variant cannot tile.
    """
    rounding = rounding or ("conv_even" if bias else "symmetric_inf")
    if rounding not in ("symmetric_inf", "conv_even"):
        raise ValueError(f"conv2dk1_i8: unknown rounding {rounding!r}")
    if (table or input_slice) and not bias:
        raise ValueError("conv2dk1_i8: table and input_slice need bias")
    if input_width % (4 if bias else 32):
        raise ValueError(
            f"conv2dk1_i8: input_width {input_width} must be a multiple of "
            f"{4 if bias else 32}"
        )
    if input_slice is not None and (
        input_slice % 8 or input_channels % input_slice or input_slice == input_channels
    ):
        raise ValueError(
            f"conv2dk1_i8: input_slice {input_slice} must be a multiple of 8 "
            f"that splits the {input_channels} input channels at least twice"
        )
    name = "conv2dk1_i8" + "_bias" * bias + "_table" * table
    name += "_even" * (rounding == "conv_even")
    name += f"_slice{input_slice}" * (input_slice is not None)
    flags = ["-DINT8_ACT", f"-DCONV2DK1_I8_SYMBOL={name}"]
    flags += ["-DCONV_BIAS"] * bias + _table_flags(table, table_copies)
    flags += ["-DCONV_ROUND_EVEN"] * (rounding == "conv_even")
    flags += [f"-DCONV_IC_SLICE={input_slice}"] * (input_slice is not None)
    call_channels = input_slice or input_channels
    in_ty = np.ndarray[(input_width * call_channels,), np.dtype[np.int8]]
    wt_ty = np.ndarray[(call_channels * output_channels,), np.dtype[np.int8]]
    bias_ty = np.ndarray[(output_channels,), np.dtype[np.int32]]
    table_ty = np.ndarray[(256,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.int8]]
    sums_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.int32]]
    extra = [bias_ty] * bias + [table_ty] * table + [sums_ty] * bool(input_slice)
    out_index = 2 + len(extra)
    bindings, call_index = (), None
    if input_slice:
        reference = partial(
            conv2dk1_i8_bias_table_slices_ref if table else conv2dk1_i8_bias_slices_ref,
            rounding=rounding,
            input_slice=input_slice,
        )
        tolerance = Tolerance.exact(note="the reference models the rounding mode")
        # Each call takes one slice into sums the core keeps, the last ending
        # in the output tile.
        call_index = out_index + 5
        bindings = (
            (out_index - 1, np.zeros(input_width * output_channels, np.int32)),
            (call_index, 0),
        )
        out_align = ((out_index, 32),) if input_width % 8 else _vector_args(out_index)
        alignments = (
            _vector_args(0, 1)
            + ((2, 32),)
            + _vector_32b(*[3] * table)
            + _vector_args(out_index - 1)
            + out_align
        )
    elif bias:
        reference = partial(
            conv2dk1_i8_bias_table_ref if table else conv2dk1_i8_bias_ref,
            rounding=rounding,
        )
        tolerance = Tolerance.exact(note="the reference models the rounding mode")
        # A row of 8k + 4 pixels goes 4 pixels (32 bytes) at a time.
        out_align = (
            ((2 + len(extra), 32),) if input_width % 8 else _vector_args(2 + len(extra))
        )
        alignments = (
            _vector_args(0, 1) + ((2, 32),) + _vector_32b(*[3] * table) + out_align
        )
    else:
        reference = conv2dk1_i8_even_ref if rounding == "conv_even" else conv2dk1_i8_ref
        tolerance = Tolerance.exact(note="the reference models the rounding mode")
        alignments = ()
    return _make_extern(
        name,
        _kernel_source("conv/conv2dk1_i8.cc"),
        [in_ty, wt_ty, *extra, out_ty, *_i32s(4 + bool(input_slice))],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels),
        include_dirs=[*_include_dirs(), _runtime_lib_include()],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(
                In,
                Param,
                *((Param,) * len(extra)),
                Out,
                *((Param,) * (4 + bool(input_slice))),
            ),
            reference=reference,
            acc_dtype=np.int32,
            reduction=input_channels,
            tolerance=tolerance,
            alignments=alignments,
            ops_per_call=2 * input_width * call_channels * output_channels,
            parameter_bindings=bindings,
            call_index=call_index,
            out_last=bool(input_slice),
            required_calls=input_channels // input_slice if input_slice else None,
        ),
    )

conv2dk14

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

14x14 convolution kernel.

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk14 kernel.

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

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

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

conv2dk1_skip_init

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

1x1 convolution kernel with skip-init connection.

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk1_skip_init kernel.

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

    Returns:
        ExternalFunction configured for the conv2dk1_skip_init kernel.

    Raises:
        ValueError: When ``act_dtype`` is not ``np.int8`` or ``np.uint8``,
            when ``input_width`` is not a positive multiple of 32, or when a
            channel count is not a positive multiple of what the source
            steps by (16 input channels, 8 output and skip channels).
    """
    func_name, flags = _conv_act_dtype_info(
        "conv2dk1_skip_init", act_dtype, factory_name="conv2dk1_skip_init"
    )
    if input_width <= 0 or input_width % 32:
        # The kernel computes whole 32-wide blocks and its tail path was never
        # implemented, so a width that is not a multiple of 32 silently leaves
        # its last columns unwritten. Refuse it rather than return part of an
        # answer.
        raise ValueError(
            f"conv2dk1_skip_init: input_width must be a positive multiple of 32, "
            f"got {input_width}"
        )
    if skip_input_channels is None:
        skip_input_channels = input_channels
    # Both paths in the source count channels in whole steps: the two input
    # halves in 16s (an 8x8 weight block per half), the output and the skip
    # projection in 8s. A count that is not a whole number of steps is
    # silently truncated by the integer division in the loop bounds.
    for label, count, step in (
        ("input_channels", input_channels, 16),
        ("output_channels", output_channels, 8),
        ("skip_input_channels", skip_input_channels, 8),
    ):
        if count <= 0 or count % step:
            raise ValueError(
                f"conv2dk1_skip_init: {label} must be a positive multiple of "
                f"{step}, got {count}"
            )
    half_ch = input_channels // 2
    total_in_ch = input_channels + skip_input_channels
    in0_ty = np.ndarray[(input_width * half_ch,), np.dtype[np.uint8]]
    in1_ty = np.ndarray[(input_width * half_ch,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(total_in_ch * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    skip_ty = np.ndarray[(input_width * skip_input_channels,), np.dtype[act_dtype]]
    return _make_extern(
        func_name,
        _kernel_source("conv/conv2dk1_skip_init.cc"),
        [in0_ty, in1_ty, wt_ty, out_ty, skip_ty, *_i32s(7)],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels)
        + [f"-DCONV_SKIP_INPUT_CHANNELS={skip_input_channels}"],
        contract=KernelContract(
            alignments=_vector_args(0, 1, 2, 3, 4),
            trace=Trace.whole_call(),
            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 utils/pyright_stubs/aie/iron/kernels/conv.py

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

    The layouts of [`conv2dk1_ref`][aie.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 utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

bn_conv2dk1_skip_ref

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

Numpy reference for bn_conv2dk1_skip.

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

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

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

Source code in utils/pyright_stubs/aie/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`][aie.iron.kernels.conv.bn_conv2dk1_skip].

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

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

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

bn_conv2dk3_ref

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

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

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

Source code in utils/pyright_stubs/aie/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`][aie.iron.kernels.conv.bn_conv2dk3]: 3x3 stride-2 conv.

    The layouts and ``check`` of [`conv2dk3_ref`][aie.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 utils/pyright_stubs/aie/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`][aie.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`][aie.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 utils/pyright_stubs/aie/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`][aie.iron.kernels.conv.bn_conv2dk3_dw_out_split].

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

bn_conv2dk1_relu_xy_pool_padded_ref

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

Numpy reference for bn_conv2dk1_relu_xy_pool_padded into a zeroed output.

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

Source code in utils/pyright_stubs/aie/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`][aie.iron.kernels.conv.bn_conv2dk1_relu_xy_pool_padded] into a zeroed output.

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

bn_fc_relu_ui16_ic_chunk_ref

bn_fc_relu_ui16_ic_chunk_ref(
    x,
    weights,
    images,
    scale,
    *,
    input_channels: int,
    chunk_input_channels: int,
    chunk_output_channels: int,
    max_images: int
)

Numpy reference for bn_fc_relu_ui16_ic_chunk.

Call k of weights ((calls, n), one [OC/8][ICc][8] chunk each) adds input channels [k * ICc, (k + 1) * ICc) of each of the first images images in x (at most max_images) to their running sums, and stores sat_u8(round_even(sum, scale)) as uint16.

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def bn_fc_relu_ui16_ic_chunk_ref(
    x,
    weights,
    images,
    scale,
    *,
    input_channels: int,
    chunk_input_channels: int,
    chunk_output_channels: int,
    max_images: int,
):
    """Numpy reference for [`bn_fc_relu_ui16_ic_chunk`][aie.iron.kernels.conv.bn_fc_relu_ui16_ic_chunk].

    Call ``k`` of ``weights`` (``(calls, n)``, one ``[OC/8][ICc][8]`` chunk
    each) adds input channels ``[k * ICc, (k + 1) * ICc)`` of each of the
    first ``images`` images in ``x`` (at most ``max_images``) to their
    running sums, and stores ``sat_u8(round_even(sum, scale))`` as
    ``uint16``.
    """
    IC, ICc, OCc = input_channels, chunk_input_channels, chunk_output_channels
    n = min(int(images), max_images)
    w = np.asarray(weights, dtype=np.int8).astype(np.int64)
    calls = w.shape[0]
    w = w.reshape(calls, OCc // 8, ICc, 8).transpose(0, 2, 1, 3)
    w = w.reshape(calls, ICc, OCc)
    xi = np.asarray(x).astype(np.int64).reshape(max_images, IC)[:n]
    out = np.zeros((calls, max_images * OCc), dtype=np.uint16)
    acc = np.zeros((n, OCc), dtype=np.int64)
    for k in range(calls):
        acc += xi[:, k * ICc : (k + 1) * ICc] @ w[k]
        out[k, : n * OCc] = _requant_even(acc, scale, dtype=np.uint16).reshape(-1)
    return out

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 utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

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

bn_conv2dk3

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk3 kernel.

Source code in utils/pyright_stubs/aie/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`][aie.iron.kernels.conv.conv2dk3].
            Defaults to ``output_channels``.

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

bn_conv2dk1_i8

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk1_i8 kernel.

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

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

bn_conv2dk1_skip

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk1_skip kernel.

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

    Returns:
        ExternalFunction configured for the bn_conv2dk1_skip kernel.

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

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

bn_conv2dk3_dw

bn_conv2dk3_dw(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    stride: int = 1,
    bias: bool = False,
    table: bool = False,
    rounding: str = "conv_even",
    table_copies: int = 2,
) -> ExternalFunction

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

With bias it is instead ONNX QLinearConv's depthwise arithmetic on int8 for a power-of-two per-tensor output scale and zero zero-points: an int32 per channel, in accumulator units, seeds the sum, the output is int8 with no ReLU, the stride is 1 and the width any multiple of 4. table then maps each result through a 256-entry int8 table indexed by value + 128.

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 with bias).	1
bias	bool	Take an int32 bias per channel after the weights.	False
table	bool	Take a 256-entry int8 table after the bias. Needs bias.	False
rounding	str	With bias, how the shift rounds: "conv_even" (half to even, as ONNX QuantizeLinear) or "symmetric_inf". The uint8 kernel always rounds half to even.	'conv_even'
table_copies	int	With table, how many copies of the table the kernel keeps in data memory: 2 (two banks, the LUT read fetches both halves at once) or 1 (1 KB less, the reads serialize).	2

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk3_dw kernel.

Raises:

Type	Description
ValueError	When stride is not 1 or 2, or on an option only the bias variant has.

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def bn_conv2dk3_dw(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    stride: int = 1,
    bias: bool = False,
    table: bool = False,
    rounding: str = "conv_even",
    table_copies: int = 2,
) -> ExternalFunction:
    """Bottleneck depthwise 3x3 conv + ReLU kernel (uint8 in/out).

    With ``bias`` it is instead ONNX QLinearConv's depthwise arithmetic on
    ``int8`` for a power-of-two per-tensor output scale and zero
    zero-points: an ``int32`` per channel, in accumulator units, seeds the sum,
    the output is ``int8`` with no ReLU, the stride is 1 and the width any
    multiple of 4. ``table`` then maps each result through a 256-entry
    ``int8`` table indexed by ``value + 128``.

    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; 1 with ``bias``).
        bias: Take an ``int32`` bias per channel after the weights.
        table: Take a 256-entry ``int8`` table after the bias. Needs ``bias``.
        rounding: With ``bias``, how the shift rounds: ``"conv_even"`` (half
            to even, as ONNX QuantizeLinear) or ``"symmetric_inf"``. The
            uint8 kernel always rounds half to even.
        table_copies: With ``table``, how many copies of the table the
            kernel keeps in data memory: 2 (two banks, the LUT read fetches
            both halves at once) or 1 (1 KB less, the reads serialize).

    Returns:
        ExternalFunction configured for the bn_conv2dk3_dw kernel.

    Raises:
        ValueError: When ``stride`` is not 1 or 2, or on an option only the
            ``bias`` variant has.
    """
    if bias:
        return _bn_conv2dk3_dw_bias(
            input_width, output_channels, table, stride, rounding, table_copies
        )
    if table:
        raise ValueError("bn_conv2dk3_dw(): table needs bias")
    if stride not in (1, 2):
        raise ValueError(f"bn_conv2dk3_dw(): stride must be 1 or 2, got {stride}")

    func_name = f"conv2dk3_dw_stride{stride}_relu_ui8_ui8"

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

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

bn_conv2dk3_dw_bias_ref

bn_conv2dk3_dw_bias_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
    *,
    rounding
)

Numpy reference for bn_conv2dk3_dw with bias.

The layouts of bn_conv2dk3_dw_ref on int8 lines: out = sat_i8(round((sum + bias) / 2**scale)), rounding half to even (conv_even) or away from zero (symmetric_inf).

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def bn_conv2dk3_dw_bias_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
    *,
    rounding,
):
    """Numpy reference for [`bn_conv2dk3_dw`][aie.iron.kernels.conv.bn_conv2dk3_dw] with ``bias``.

    The layouts of [`bn_conv2dk3_dw_ref`][aie.iron.kernels.conv.bn_conv2dk3_dw_ref]
    on ``int8`` lines: ``out = sat_i8(round((sum + bias) / 2**scale))``,
    rounding half to even (``conv_even``) or away from zero
    (``symmetric_inf``).
    """
    del input_channels, kernel_width, kernel_height, channel_offset
    C = int(output_channels)
    acc, lead, Wo = _dw3x3_acc((line0, line1, line2), weights, input_width, C, check)
    acc = acc + np.asarray(bias, dtype=np.int64).reshape(C // 8, 1, 8)
    return _requant_bias(acc, scale, rounding).reshape(*lead, Wo * C)

bn_conv2dk3_dw_bias_table_ref

bn_conv2dk3_dw_bias_table_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    table,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
    *,
    rounding
)

bn_conv2dk3_dw_bias_ref, then table[out + 128].

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def bn_conv2dk3_dw_bias_table_ref(
    line0,
    line1,
    line2,
    weights,
    bias,
    table,
    input_width,
    input_channels,
    output_channels,
    kernel_width,
    kernel_height,
    check,
    scale,
    channel_offset,
    *,
    rounding,
):
    """[`bn_conv2dk3_dw_bias_ref`][aie.iron.kernels.conv.bn_conv2dk3_dw_bias_ref], then ``table[out + 128]``."""
    out = bn_conv2dk3_dw_bias_ref(
        line0,
        line1,
        line2,
        weights,
        bias,
        input_width,
        input_channels,
        output_channels,
        kernel_width,
        kernel_height,
        check,
        scale,
        channel_offset,
        rounding=rounding,
    )
    return np.asarray(table, dtype=np.int8)[out.astype(np.int64) + 128]

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 utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

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

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

bn_conv2dk1_partial_put_i8

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

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the PUT tile.

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

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

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

    Returns:
        ExternalFunction configured for the PUT tile.

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

bn_conv2dk1_partial_get_relu_i8

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

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the GET tile.

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

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

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

    Returns:
        ExternalFunction configured for the GET tile.

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

bn_conv2dk3_dw_out_split

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

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

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

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

Parameters:

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

Returns:

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

Raises:

Type	Description
ValueError	When block_index is not 13 or 14.

Source code in utils/pyright_stubs/aie/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`][aie.iron.kernels.conv.bn_conv2dk3_dw] (stride=1) that writes its output
    to TWO separate buffers — the channel dimension is split in half so
    downstream cascade-PUT tiles can each consume one slice.  Used by
    MobileNet V3's bn13 / bn14 depthwise stage to feed the L3 cascade.

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

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

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

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

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

bn_conv2dk1_input_split_partial_put_ui8

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

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

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

Parameters:

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

Returns:

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

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

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

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

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

bn_conv2dk1_input_split_partial_skip_get

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

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

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

Parameters:

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

Returns:

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

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

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

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

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

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

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

bn_fc_relu_ui16_ic_chunk

bn_fc_relu_ui16_ic_chunk(
    input_channels: int = 1280,
    chunk_input_channels: int = 320,
    chunk_output_channels: int = 32,
    max_images: int = 1,
) -> ExternalFunction

Fully-connected layer + ReLU over one input-channel chunk, uint16 in and out.

MobileNet V3's classifier head streams each FC layer's weights in chunks of chunk_output_channels outputs by chunk_input_channels inputs, so a layer whose input width is any multiple of the chunk (FC1's 960, FC2's 1280) streams no padding. Call chunk adds that slice of each image's input_channels uint16 inputs to the outputs' int32 sums, which the kernel keeps from one call to the next and restarts at chunk 0, and stores them requantized, so the last chunk's store is the result. Takes (input, weights, output, chunk, images, scale); the weights are [chunk_output_channels / 8][chunk_input_channels][8].

Parameters:

Name	Type	Description	Default
input_channels	int	Input channels per image in input.	1280
chunk_input_channels	int	Input channels per call.	320
chunk_output_channels	int	Output channels per call.	32
max_images	int	Images the buffers hold; images may be fewer, and more run as max_images.	1

Returns:

Type	Description
ExternalFunction	ExternalFunction for post_L2_fc_ui16_ic_chunk.

Raises:

Type	Description
ValueError	When a channel count is not a multiple of the vector width (32 inputs, 8 outputs), chunk_input_channels does not divide input_channels or max_images is below 1.

Source code in utils/pyright_stubs/aie/iron/kernels/conv.py

def bn_fc_relu_ui16_ic_chunk(
    input_channels: int = 1280,
    chunk_input_channels: int = 320,
    chunk_output_channels: int = 32,
    max_images: int = 1,
) -> ExternalFunction:
    """Fully-connected layer + ReLU over one input-channel chunk, uint16 in and out.

    MobileNet V3's classifier head streams each FC layer's weights in chunks
    of ``chunk_output_channels`` outputs by ``chunk_input_channels`` inputs,
    so a layer whose input width is any multiple of the chunk (FC1's 960,
    FC2's 1280) streams no padding. Call ``chunk`` adds that slice of each
    image's ``input_channels`` uint16 inputs to the outputs' int32 sums,
    which the kernel keeps from one call to the next and restarts at chunk
    0, and stores them requantized, so the last chunk's store is the
    result. Takes ``(input, weights, output, chunk, images, scale)``; the
    weights are ``[chunk_output_channels / 8][chunk_input_channels][8]``.

    Args:
        input_channels: Input channels per image in ``input``.
        chunk_input_channels: Input channels per call.
        chunk_output_channels: Output channels per call.
        max_images: Images the buffers hold; ``images`` may be fewer, and
            more run as ``max_images``.

    Returns:
        ExternalFunction for ``post_L2_fc_ui16_ic_chunk``.

    Raises:
        ValueError: When a channel count is not a multiple of the vector
            width (32 inputs, 8 outputs), ``chunk_input_channels`` does not
            divide ``input_channels`` or ``max_images`` is below 1.
    """
    if input_channels % chunk_input_channels:
        raise ValueError(
            f"bn_fc_relu_ui16_ic_chunk: chunk_input_channels "
            f"({chunk_input_channels}) must divide input_channels "
            f"({input_channels}), or the last chunk reads past each image"
        )
    if (
        input_channels % 32
        or chunk_input_channels % 32
        or chunk_output_channels % 8
        or max_images < 1
    ):
        raise ValueError(
            "bn_fc_relu_ui16_ic_chunk() needs input channel counts that are "
            "multiples of 32, an output chunk that is a multiple of 8 and "
            f"max_images >= 1, got input_channels={input_channels}, "
            f"chunk_input_channels={chunk_input_channels}, "
            f"chunk_output_channels={chunk_output_channels}, "
            f"max_images={max_images}."
        )
    in_ty = np.ndarray[(max_images * input_channels,), np.dtype[np.uint16]]
    wt_ty = np.ndarray[
        (chunk_input_channels * chunk_output_channels,), np.dtype[np.int8]
    ]
    out_ty = np.ndarray[(max_images * chunk_output_channels,), np.dtype[np.uint16]]
    return _make_extern(
        "post_L2_fc_ui16_ic_chunk",
        _kernel_source("conv/bn_conv2dk1_relu.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(3)],
        compile_flags=[
            "-DPOSTL2_IC_CHUNK",
            f"-DFC_INPUT_CHANNELS={input_channels}",
            f"-DFC_CHUNK_IC={chunk_input_channels}",
            f"-DFC_CHUNK_OC={chunk_output_channels}",
            f"-DFC_MAX_IMAGES={max_images}",
        ],
        contract=KernelContract(
            alignments=_vector_loads(1),
            trace=Trace.whole_call(),
            roles=(Param, In, Out, Param, Param, Param),
            parameter_bindings=((3, 0),),
            call_index=3,
            reference=partial(
                bn_fc_relu_ui16_ic_chunk_ref,
                input_channels=input_channels,
                chunk_input_channels=chunk_input_channels,
                chunk_output_channels=chunk_output_channels,
                max_images=max_images,
            ),
            acc_dtype=np.int32,
            reduction=input_channels,
            tolerance=_BN_EXACT,
            ops_per_call=2 * max_images * chunk_input_channels * chunk_output_channels,
        ),
    )

Activation functions

gelu_erf is GELU as defined, x * Phi(x) (Hendrycks and Gimpel), its erfc in the form of Numerical Recipes' erfcc (Press et al., 2nd ed., §6.2) with a refit polynomial. log10_clamp_f32 takes log-mel features as Hugging Face's WhisperFeatureExtractor does.

Activation kernel factories and NumPy reference implementations.

softmax

softmax(
    tile_size: int = 1024,
    accurate_exp2: bool = False,
    input_dtype: type = bfloat16,
    soft_cap: float | None = None,
    masked: bool = False,
    addend: bool = False,
    window: int | None = None,
    period: int | None = None,
    scale: float | None = None,
    span: bool = False,
) -> ExternalFunction

Softmax over a row of tile_size elements, into bf16.

With bf16 input this is softmax_bf16(x, y, n), or with span softmax_span_bf16(x, y, n, lo, hi). With float32 input it is softmax_f32_bf16(x, [addend], [mask], y, n), which computes softmax(cap * tanh((x + addend) / cap) + mask) with the logits in float32 and overwrites x. With window and period it is instead softmax_rows_banded_f32_bf16(x, addend, y, 1, n, window, period, row_offset), one row a call: the mask is a band the kernel computes from window and the row's own query position (row_offset, wrapped every period rows), not read from memory.

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
accurate_exp2	bool	Scale by log2(e) to f32 and take 2^x with a polynomial within 2.8e-6 (-DEXP2_BF16_ACCURATE), in place of bf16(log2(e)) and AIE2P's exp2 instruction, or AIE2's exp lookup table. The block entry points (softmax_sum_bf16, softmax_scale_bf16) bound from this object follow it too. Always on for float32 input.	False
input_dtype	type	bfloat16, or np.float32 for the float path.	bfloat16
soft_cap	float | None	Float input only: each logit x + addend becomes cap * tanh((x + addend) / cap), within 6e-6 * cap. In [2^-100, 2^10].	None
masked	bool	Float input only: a bf16 row added to the capped logits; lowest or -inf drops an element. A row keeps at least one. Exclusive with window/period.	False
addend	bool	Float input only: a float32 row added to x before the cap, such as relative-position logits. Always on with window/period, whose caller (EG2's audio attention) always has one.	False
window	int | None	Float input only, with period: the band's width: a row's query position q = row_offset % period keeps column j where j >= max(q % 2, window - 1 - q + q % 2) and j < window + q % 2.	None
period	int | None	Float input only, with window: the band wraps every period rows (EG2 tiles one band over every attention head).	None
scale	float | None	bf16 input only: fold a scale (1/sqrt(d) for attention) into the exp2 argument as scale * log2(e), host-split into the same bf16 limbs (or the single one, without accurate_exp2) log2(e) alone takes, exact to the same 2^-26 relative bound for any value -- a pre-scaled q is exact only at a power-of-two scale. None keeps log2(e) itself, bit for bit the unscaled kernel; so does scale=1.0, computed rather than skipped. Finite and non-negative; 0 weighs every element alike.	None
span	bool	bf16 input only: softmax_span_bf16(x, y, n, lo, hi) in place of softmax_bf16. lo/hi are per-call scalars bounding the row's live span: [0, lo) and [hi, n) come out exactly 0 at any scale, the span normalised alone, so a causal row gives only hi short of n and a sliding window gives both.	False

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the softmax kernel.

Raises:

Type	Description
ValueError	A tile_size that is not a positive multiple of 32, an input_dtype other than the two, a float-only option with bf16 input, scale or span with float32 input_dtype or with an aie2 non-accurate build, a negative or non-finite scale, a soft_cap outside [2^-100, 2^10], window or period without the other, or either with masked.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def softmax(
    tile_size: int = 1024,
    accurate_exp2: bool = False,
    input_dtype: type = bfloat16,
    soft_cap: float | None = None,
    masked: bool = False,
    addend: bool = False,
    window: int | None = None,
    period: int | None = None,
    scale: float | None = None,
    span: bool = False,
) -> ExternalFunction:
    """Softmax over a row of ``tile_size`` elements, into bf16.

    With bf16 input this is ``softmax_bf16(x, y, n)``, or with ``span``
    ``softmax_span_bf16(x, y, n, lo, hi)``. With float32 input it is
    ``softmax_f32_bf16(x, [addend], [mask], y, n)``, which computes
    ``softmax(cap * tanh((x + addend) / cap) + mask)`` with the logits in
    float32 and overwrites ``x``. With ``window`` and ``period`` it is
    instead ``softmax_rows_banded_f32_bf16(x, addend, y, 1, n, window,
    period, row_offset)``, one row a call: the mask is a band the kernel
    computes from ``window`` and the row's own query position
    (``row_offset``, wrapped every ``period`` rows), not read from memory.

    Args:
        tile_size: Number of elements per tile, passed at run time; a
            positive multiple of 32, the kernel's vector step.
        accurate_exp2: Scale by log2(e) to f32 and take 2^x with a polynomial
            within 2.8e-6 (`-DEXP2_BF16_ACCURATE`), in place of bf16(log2(e))
            and AIE2P's exp2 instruction, or AIE2's exp lookup table. The
            block entry points (`softmax_sum_bf16`, `softmax_scale_bf16`)
            bound from this object follow it too. Always on for float32
            input.
        input_dtype: ``bfloat16``, or ``np.float32`` for the float path.
        soft_cap: Float input only: each logit ``x + addend`` becomes
            ``cap * tanh((x + addend) / cap)``, within ``6e-6 * cap``. In
            [2^-100, 2^10].
        masked: Float input only: a bf16 row added to the capped logits;
            lowest or -inf drops an element. A row keeps at least one.
            Exclusive with ``window``/``period``.
        addend: Float input only: a float32 row added to ``x`` before the
            cap, such as relative-position logits. Always on with
            ``window``/``period``, whose caller (EG2's audio attention)
            always has one.
        window: Float input only, with ``period``: the band's width: a row's
            query position ``q = row_offset % period`` keeps column ``j``
            where ``j >= max(q % 2, window - 1 - q + q % 2)`` and
            ``j < window + q % 2``.
        period: Float input only, with ``window``: the band wraps every
            ``period`` rows (EG2 tiles one band over every attention head).
        scale: bf16 input only: fold a scale (1/sqrt(d) for attention) into
            the exp2 argument as ``scale * log2(e)``, host-split into the
            same bf16 limbs (or the single one, without ``accurate_exp2``)
            log2(e) alone takes, exact to the same 2^-26 relative bound for
            any value -- a pre-scaled q is exact only at a power-of-two
            scale. None keeps log2(e) itself, bit for bit the unscaled
            kernel; so does ``scale=1.0``, computed rather than skipped.
            Finite and non-negative; 0 weighs every element alike.
        span: bf16 input only: ``softmax_span_bf16(x, y, n, lo, hi)`` in
            place of ``softmax_bf16``. ``lo``/``hi`` are per-call scalars
            bounding the row's live span: ``[0, lo)`` and ``[hi, n)`` come
            out exactly 0 at any ``scale``, the span normalised alone, so a
            causal row gives only ``hi`` short of ``n`` and a sliding window
            gives both.

    Returns:
        ExternalFunction configured for the softmax kernel.

    Raises:
        ValueError: A ``tile_size`` that is not a positive multiple of 32, an
            ``input_dtype`` other than the two, a float-only option with bf16
            input, ``scale`` or ``span`` with float32 ``input_dtype`` or with
            an aie2 non-accurate build, a negative or non-finite ``scale``,
            a ``soft_cap`` outside
            [2^-100, 2^10], ``window`` or ``period`` without the other, or
            either with ``masked``.
    """
    _require_vector_alignment("softmax", tile_size, _RUNTIME_VECTOR_WIDTH)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    if np.dtype(input_dtype) == np.dtype(np.float32):
        if scale is not None or span:
            raise ValueError(
                "softmax: scale and span are the bf16 input path's options; "
                "the float32 path masks through masked/window instead"
            )
        if window is not None or period is not None:
            if masked:
                raise ValueError(
                    "softmax: window/period and masked are different bands; " "give one"
                )
            return _softmax_f32_banded(tile_size, soft_cap, window, period)
        return _softmax_f32(tile_size, soft_cap, masked, addend)
    if np.dtype(input_dtype) != np.dtype(bfloat16):
        raise ValueError(
            f"softmax: input_dtype must be bfloat16 or float32, got {input_dtype}"
        )
    if (
        soft_cap is not None
        or masked
        or addend
        or window is not None
        or period is not None
    ):
        raise ValueError(
            "softmax: soft_cap, masked, addend, window and period are "
            "options of the float32 input path; pass input_dtype=np.float32"
        )
    if scale is not None and not (math.isfinite(scale) and scale >= 0):
        raise ValueError(
            f"softmax: scale must be finite and non-negative, got {scale}: the "
            "kernel centres on the row's largest logit, which is the largest "
            "scaled one only at such a scale"
        )
    if (scale is not None or span) and _tuned_arch() == "aie2" and not accurate_exp2:
        raise ValueError(
            "softmax: scale and span fold into softmax_aie2p.h's scale_log2e "
            "and mask_head_bf16, which aie2's non-accurate build does not "
            "compile (softmax_aie2.h instead); pass accurate_exp2=True, or "
            "build for aie2p"
        )
    scale_flags = _softmax_scale_flags(scale, accurate_exp2)
    if span:
        return _softmax_span_bf16(tile_size, accurate_exp2, scale, scale_flags)
    return _create_lut_kernel(
        "softmax_bf16",
        "softmax.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=(["-DEXP2_BF16_ACCURATE=1"] if accurate_exp2 else [])
        + scale_flags,
        contract=_unary_lut_contract(
            lambda x: softmax_ref(x, tile_size=tile_size, scale=scale),
            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=(
                _SOFTMAX_ACCURATE_TOLERANCE
                if accurate_exp2
                else Tolerance.relative(
                    0.04,
                    0.1 / tile_size,
                    note="AIE2P exp instruction through the softmax "
                    "normalisation, measured on npu2 at 2.94e-2 relative "
                    "(npu1: unmeasured); "
                    "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" and not accurate_exp2 else None
            ),
        ),
    )

gelu

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

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

The ABI keeps (in, out, size). Positive whole vectors only: multiples of 16 on aie2 or 32 on aie2p.

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call.	1024
use_lut	bool	Compute tanh from the interpolated LUT rather than AIE2P's vtanh instruction, which returns its argument up to 0.5, and be judged against an exact model of the kernel, gelu_lut_ref. Moot on aie2, which always takes the LUT and is judged against gelu_aie2_lut_ref.	False

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def gelu(tile_size: int = 1024, use_lut: bool = False) -> ExternalFunction:
    """GELU (tanh approximation) for bf16 tiles, with a compiled-in element count.

    The ABI keeps ``(in, out, size)``. Positive whole vectors only: multiples
    of 16 on aie2 or 32 on aie2p.

    Args:
        tile_size: Elements per call.
        use_lut: Compute tanh from the interpolated LUT rather than AIE2P's
            vtanh instruction, which returns its argument up to 0.5, and be
            judged against an exact model of the kernel,
            [`gelu_lut_ref`][aie.iron.kernels.activation.gelu_lut_ref]. Moot on
            aie2, which always takes the LUT and is judged against
            [`gelu_aie2_lut_ref`][aie.iron.kernels.activation.gelu_aie2_lut_ref].
    """
    _require_vector_alignment("gelu", 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_VTANH_TOLERANCE,
            use_lut=use_lut,
            elementwise=gelu_lut_ref,
            branch_models={"aie2": gelu_aie2_lut_ref},
        ),
        use_lut_tanh=use_lut,
    )

gelu_erf

gelu_erf(tile_size: int = 1024) -> ExternalFunction

Exact GELU, x * Phi(x), for bf16 tiles; (in, out, size).

gelu takes the tanh approximation, which rounds differently on 405 finite bf16 inputs, 280 below -3. This one evaluates erfc in f32 and rounds once.

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call, a positive multiple of 32.	1024

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def gelu_erf(tile_size: int = 1024) -> ExternalFunction:
    """Exact GELU, ``x * Phi(x)``, for bf16 tiles; ``(in, out, size)``.

    [`gelu`][aie.iron.kernels.activation.gelu] takes the tanh approximation, which
    rounds differently on 405 finite bf16 inputs, 280 below -3. This one
    evaluates erfc in f32 and rounds once.

    Args:
        tile_size: Elements per call, a positive multiple of 32.
    """
    _require_vector_alignment("gelu_erf", tile_size, 32)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _make_extern(
        "gelu_erf_bf16_size",
        _kernel_source("activation/gelu_erf.cc"),
        [tile_ty, tile_ty, np.int32],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            reference=gelu_erf_ref,
            tolerance=_GELU_ERF_ROUNDED,
            acc_dtype=np.float32,
        ),
    )

silu

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

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

The ABI keeps (in, out, size). Positive whole vectors only: multiples of 16 on aie2 or 32 on aie2p.

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call.	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 utils/pyright_stubs/aie/iron/kernels/activation.py

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

    The ABI keeps ``(in, out, size)``. Positive whole vectors only: multiples
    of 16 on aie2 or 32 on aie2p.

    Args:
        tile_size: Elements per call.
        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.
    """
    _require_vector_alignment("silu", tile_size, _arch_traits().bf16_lanes)
    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,
            use_lut=use_lut,
            elementwise=silu_table_ref,
            tolerance=_vtanh_family_tolerance("silu"),
        ),
        use_lut_tanh=use_lut,
    )

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 utils/pyright_stubs/aie/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`][aie.iron.kernels.activation.swiglu_ref].
    """
    return _bf16_lut_factory(
        "swiglu",
        "swiglu_bf16",
        "swiglu.cc",
        tile_size,
        arg_arity=4,
        compile_flags=[f"-DSWIGLU_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, In, In, Out),
            true_function=swiglu_ref,
            per_build=tuple(
                (
                    (build, swiglu_table_ref, _LUT_MODEL_TOLERANCE)
                    if use_lut or not ARCH_TRAITS[build.arch].native_tanh
                    else (build, swiglu_ref, _vtanh_family_tolerance("swiglu"))
                )
                for build in BUILDS
            ),
            acc_dtype=bfloat16,
            ops_per_call=6 * tile_size,
            uses_lut=True,
        ),
        use_lut_tanh=use_lut,
    )

bf16_exp

bf16_exp(tile_size: int = 1024) -> ExternalFunction

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

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

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

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

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

exp2f_vec

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

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

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

The same source builds for aie2.

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the exp2f_vec kernel.

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

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

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

    The same source builds for aie2.

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

    Returns:
        ExternalFunction configured for the exp2f_vec kernel.

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

log_f32

log_f32(tile_size: int = 1024) -> ExternalFunction

log(x + offset) of float32, rounded once to bf16.

Takes (x, y, size, offset): tile_size float32 in, as many bf16 out, and a float32 offset added to every element first, in float32 (a log-mel spectrogram's floor). Defined where x + offset is a positive normal float32. Every product is of bf16 limbs, exact in the float32 accumulator, so the log is float32-accurate before its one bf16 rounding.

The source sets conv_even itself and restores the core's mode on exit, and builds for aie2 and aie2p.

Parameters:

Name	Type	Description	Default
tile_size	int	Float32 elements per call, a positive multiple of 32.	1024

Returns:

Type	Description
ExternalFunction	ExternalFunction for log_f32_bf16.

Raises:

Type	Description
ValueError	When tile_size is not a positive multiple of 32.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def log_f32(tile_size: int = 1024) -> ExternalFunction:
    """``log(x + offset)`` of float32, rounded once to bf16.

    Takes ``(x, y, size, offset)``: ``tile_size`` float32 in, as many bf16
    out, and a float32 ``offset`` added to every element first, in float32
    (a log-mel spectrogram's floor). Defined where ``x + offset`` is a
    positive normal float32. Every product is of bf16 limbs, exact in the
    float32 accumulator, so the log is float32-accurate before its one bf16
    rounding.

    The source sets ``conv_even`` itself and restores the core's mode on
    exit, and builds for aie2 and aie2p.

    Args:
        tile_size: Float32 elements per call, a positive multiple of 32.

    Returns:
        ExternalFunction for ``log_f32_bf16``.

    Raises:
        ValueError: When ``tile_size`` is not a positive multiple of 32.
    """
    _require_vector_alignment("log_f32", tile_size, 32)
    return _make_extern(
        "log_f32_bf16",
        _kernel_source("activation/log_f32.cc"),
        [
            np.ndarray[(tile_size,), np.dtype[np.float32]],
            np.ndarray[(tile_size,), np.dtype[bfloat16]],
            np.int32,
            np.float32,
        ],
        compile_flags=[f"-DLOG_ELEMS={tile_size}"]
        # log_f32.cc stages AIE2P tiles of whole 256-element blocks, whose
        # passes need eight stages to pipeline; under that cap Peano
        # miscompiled one-trip passes, so other tiles keep the nested loop and
        # the default cap.
        + (
            ["-mllvm", "--aie-pipeliner-max-stagecount=8"]
            if tile_size % 256 == 0 and _detect_arch() == "aie2p"
            else []
        ),
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param, Param),
            parameter_bindings=((2, tile_size),),
            reference=log_f32_ref,
            sample=lambda rng, calls: log_f32_sample(rng, calls, tile_size=tile_size),
            tolerance=_LOG_F32_TOLERANCE,
            ops_per_call=tile_size,
        ),
    )

log10_clamp_f32

log10_clamp_f32(tile_size: int = 1024) -> ExternalFunction

Whisper's log-mel features from mel energies, rounded once to bf16.

Takes (x, peak, y, size, mel_floor, decades, shift, scale): tile_size float32 energies and the clip's largest, one float32, in, and tile_size bf16 out, (max(log10(max(x, mel_floor)), log10(max(peak, mel_floor)) - decades) + shift) * scale, each step in float32 as WhisperFeatureExtractor takes it (its mel_floor 1e-10, decades 8, shift 4 and scale 1/4). The log is log_f32's, kept in float32, times 1 / ln 10 on bf16 limbs.

The source sets conv_even itself and restores the core's mode on exit, and builds for aie2 and aie2p.

Parameters:

Name	Type	Description	Default
tile_size	int	Float32 elements per call, a positive multiple of 32.	1024

Returns:

Type	Description
ExternalFunction	ExternalFunction for log10_clamp_f32_bf16.

Raises:

Type	Description
ValueError	When tile_size is not a positive multiple of 32.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def log10_clamp_f32(tile_size: int = 1024) -> ExternalFunction:
    """Whisper's log-mel features from mel energies, rounded once to bf16.

    Takes ``(x, peak, y, size, mel_floor, decades, shift, scale)``:
    ``tile_size`` float32 energies and the clip's largest, one float32, in,
    and ``tile_size`` bf16 out, ``(max(log10(max(x, mel_floor)),
    log10(max(peak, mel_floor)) - decades) + shift) * scale``, each step in
    float32 as WhisperFeatureExtractor takes it (its ``mel_floor`` 1e-10,
    ``decades`` 8, ``shift`` 4 and ``scale`` 1/4). The log is
    ``log_f32``'s, kept in float32, times ``1 / ln 10`` on bf16 limbs.

    The source sets ``conv_even`` itself and restores the core's mode on
    exit, and builds for aie2 and aie2p.

    Args:
        tile_size: Float32 elements per call, a positive multiple of 32.

    Returns:
        ExternalFunction for ``log10_clamp_f32_bf16``.

    Raises:
        ValueError: When ``tile_size`` is not a positive multiple of 32.
    """
    _require_vector_alignment("log10_clamp_f32", tile_size, 32)
    return _make_extern(
        "log10_clamp_f32_bf16",
        _kernel_source("activation/log_f32.cc"),
        [
            np.ndarray[(tile_size,), np.dtype[np.float32]],
            np.ndarray[(1,), np.dtype[np.float32]],
            np.ndarray[(tile_size,), np.dtype[bfloat16]],
            np.int32,
            np.float32,
            np.float32,
            np.float32,
            np.float32,
        ],
        compile_flags=[f"-DLOG_ELEMS={tile_size}", "-DLOG10_CLAMP"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out, Param, Param, Param, Param, Param),
            parameter_bindings=((3, tile_size),),
            reference=log10_clamp_f32_ref,
            sample=lambda rng, calls: log10_clamp_f32_sample(
                rng, calls, tile_size=tile_size
            ),
            tolerance=_LOG10_CLAMP_F32_TOLERANCE,
            ops_per_call=tile_size,
        ),
    )

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 utils/pyright_stubs/aie/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`][aie.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 utils/pyright_stubs/aie/iron/kernels/activation.py

def sigmoid(tile_size: int = 1024, use_lut: bool = False) -> ExternalFunction:
    """Sigmoid for bf16 tiles of a positive multiple of 32 elements.

    The count is compiled in; retain ``tile_size`` as a trailing ABI argument.
    """
    _require_vector_alignment("sigmoid", tile_size, _RUNTIME_VECTOR_WIDTH)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "sigmoid_bf16",
        "sigmoid.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DSIGMOID_ELEMS={tile_size}"],
        contract=_unary_lut_contract(
            sigmoid_ref,
            count=tile_size,
            use_lut=use_lut,
            elementwise=sigmoid_table_ref,
            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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/iron/kernels/activation.py

def silu_ref(x):
    """Numpy reference for [`silu`][aie.iron.kernels.activation.silu] (Swish) — ``x * sigmoid(x)``.

    LUT-approximation territory; pair with ``rtol=0.128`` (the default
    in `count_mismatches`) when verifying.
    """
    xf = _no_neg_inf(_f64(x))
    with np.errstate(over="ignore"):
        return _rounded(xf / (1.0 + np.exp(-xf)), x)

gelu_ref

gelu_ref(x)

Numpy reference for gelu.

Tanh approximation 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))). Matches the C++ kernel's tanh-GELU formula. It is evaluated in float64: in float32, 1 + tanh cancels for x below about -4.5 and leaves values over ten times too large.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def gelu_ref(x):
    """Numpy reference for [`gelu`][aie.iron.kernels.activation.gelu].

    Tanh approximation ``0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))``.
    Matches the C++ kernel's tanh-GELU formula. It is evaluated in float64:
    in float32, ``1 + tanh`` cancels for x below about -4.5 and leaves
    values over ten times too large.
    """
    xf = _no_neg_inf(x.astype(np.float64))
    with np.errstate(over="ignore"):
        inner = math.sqrt(2.0 / math.pi) * (xf + 0.044715 * xf**3)
    return (0.5 * xf * (1.0 + np.tanh(inner))).astype(x.dtype)

gelu_erf_ref

gelu_erf_ref(x)

Numpy reference for gelu_erf: x * Phi(x).

0.5 x erfc(-x / sqrt 2) in float64, rounded once.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def gelu_erf_ref(x):
    """Numpy reference for [`gelu_erf`][aie.iron.kernels.activation.gelu_erf]: ``x * Phi(x)``.

    ``0.5 x erfc(-x / sqrt 2)`` in float64, rounded once.
    """
    xf = _no_neg_inf(_f64(x))
    erfc = np.asarray(_erfc(-xf / math.sqrt(2.0)), dtype=np.float64)
    return _rounded(0.5 * xf * erfc, x)

gelu_lut_ref

gelu_lut_ref(x)

Model of AIE2P's gelu built with use_lut=True.

gelu_aie2p.h step for step: x*x and x*s_beta are stored to bf16, the tanh argument x*s + s_beta*x*x^2 is summed in f32 and narrowed to bf16 for tanh_lut_ref, and x/2 + t * x/2 is one f32 mac, rounded once to bf16 with a subnormal flushed to zero. x/2 takes x clamped at -8, where 1 + t is already 0, so -inf gives 0.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def gelu_lut_ref(x):
    """Model of AIE2P's [`gelu`][aie.iron.kernels.activation.gelu] built with ``use_lut=True``.

    gelu_aie2p.h step for step: ``x*x`` and ``x*s_beta`` are stored to bf16,
    the tanh argument ``x*s + s_beta*x*x^2`` is summed in f32 and narrowed to
    bf16 for [`tanh_lut_ref`][aie.iron.kernels.activation.tanh_lut_ref], and
    ``x/2 + t * x/2`` is one f32 mac, rounded once to bf16 with a subnormal
    flushed to zero. ``x/2`` takes x clamped at -8, where ``1 + t`` is
    already 0, so -inf gives 0.
    """
    xf = np.asarray(x).astype(np.float32)
    s = np.float32(bfloat16(math.sqrt(2.0 / math.pi)))
    s_beta = np.float32(bfloat16(s * np.float32(0.044715)))
    with np.errstate(over="ignore", invalid="ignore"):
        x2 = _bf16_ftz(xf * xf)
        sbeta_x = _bf16_ftz(xf * s_beta)
        t = np.asarray(tanh_lut_ref(_bf16(xf * s + sbeta_x * x2)), np.float32)
    half_x = np.maximum(xf, -8.0) * np.float32(0.5)
    return _bf16_ftz(half_x + t * half_x).astype(np.asarray(x).dtype)

gelu_aie2_lut_ref

gelu_aie2_lut_ref(x)

Model of AIE2's gelu.

gelu_aie2.h step for step: x is clamped at -8 first, where 1 + t is already 0, then x*x, c + d*x^2 (one f32 mac with c in f32) and x * (c + d*x^2) are each stored to bf16, and the last is the argument of tanh_lut_ref. x/2 + t * bf16(x/2) is one f32 mac, rounded once to bf16. Each store flushes a subnormal to zero, and a zero result is +0, as on npu1, even for x = -0.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def gelu_aie2_lut_ref(x):
    """Model of AIE2's [`gelu`][aie.iron.kernels.activation.gelu].

    gelu_aie2.h step for step: x is clamped at -8 first, where ``1 + t`` is
    already 0, then ``x*x``, ``c + d*x^2`` (one f32 mac with c in f32) and
    ``x * (c + d*x^2)`` are each stored to bf16, and the last is the
    argument of [`tanh_lut_ref`][aie.iron.kernels.activation.tanh_lut_ref].
    ``x/2 + t * bf16(x/2)`` is one f32 mac, rounded once to bf16. Each store
    flushes a subnormal to zero, and a zero result is +0, as on npu1, even
    for x = -0.
    """
    xl = np.maximum(np.asarray(x).astype(np.float32), np.float32(-8.0))
    c = np.float32(0.79788456)
    d = np.float32(bfloat16(c * np.float32(0.044715)))
    with np.errstate(over="ignore", invalid="ignore"):
        x2 = _bf16_ftz(xl * xl)
        p = _bf16_ftz(c + x2 * d)
        t = np.asarray(tanh_lut_ref(_bf16_ftz(xl * p)), np.float32)
        half_x = xl * np.float32(0.5)
        out = _bf16_ftz(half_x + t * _bf16_ftz(half_x)) + np.float32(0)
    return out.astype(np.asarray(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 utils/pyright_stubs/aie/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_table_ref

sigmoid_table_ref(x)

Model of sigmoid's LUT build.

activation/sigmoid.cc reads its own table: getTanhBf16's segments rewritten for 0.5 + 0.5 * tanh(x/2), so segment e has slope slope[e] / 4 and offset 0.5 + 0.5 * offset[e] and covers 0.5 of x. x is clamped to [-8, 8 - 1/32]; the one rounding is the store to bf16, as in tanh_lut_ref.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def sigmoid_table_ref(x):
    """Model of [`sigmoid`][aie.iron.kernels.activation.sigmoid]'s LUT build.

    activation/sigmoid.cc reads its own table: getTanhBf16's segments
    rewritten for ``0.5 + 0.5 * tanh(x/2)``, so segment ``e`` has slope
    ``slope[e] / 4`` and offset ``0.5 + 0.5 * offset[e]`` and covers 0.5 of x.
    x is clamped to ``[-8, 8 - 1/32]``; the one rounding is the store to bf16,
    as in [`tanh_lut_ref`][aie.iron.kernels.activation.tanh_lut_ref].
    """
    xf = np.clip(np.asarray(x).astype(np.float32), -8.0, 8.0 - 1.0 / 32)
    e = np.clip(np.floor(xf * 2.0).astype(np.int64), -16, 15) + 16
    slope = np.asarray(_TANH_LUT_SLOPE, np.float32)[e] / 4
    offset = 0.5 + 0.5 * np.asarray(_TANH_LUT_OFFSET, np.float32)[e]
    return (slope * xf + offset).astype(bfloat16).astype(np.asarray(x).dtype)

silu_table_ref

silu_table_ref(x)

Model of silu's LUT build.

activation/silu.cc multiplies x by the bf16 sigmoid of sigmoid_table_ref, so that rounding is modelled too, not folded away. The sigmoid is exactly 0 from x = -8 down, and x is clamped there before the multiply, so -inf gives 0 rather than NaN. A subnormal product is flushed to zero.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def silu_table_ref(x):
    """Model of [`silu`][aie.iron.kernels.activation.silu]'s LUT build.

    activation/silu.cc multiplies x by the bf16 sigmoid of
    [`sigmoid_table_ref`][aie.iron.kernels.activation.sigmoid_table_ref], so that
    rounding is modelled too, not folded away. The sigmoid is exactly 0 from
    x = -8 down, and x is clamped there before the multiply, so -inf gives 0
    rather than NaN. A subnormal product is flushed to zero.
    """
    xf = np.asarray(x).astype(np.float32)
    sig = np.asarray(sigmoid_table_ref(xf), np.float32)
    return _bf16_ftz(np.maximum(xf, -8.0) * sig).astype(np.asarray(x).dtype)

swiglu_table_ref

swiglu_table_ref(x, w1, w2)

Model of swiglu's LUT build.

activation/swiglu.cc narrows after every multiply -- x*w1, x*w2, the sigmoid factor of sigmoid_table_ref 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_table_ref, and where the silu product is 0 the output is 0, so an overflowed x*w1 does not make inf * 0. Each subnormal product is flushed to zero, so a subnormal x*w2 zeroes the output through the gate.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def swiglu_table_ref(x, w1, w2):
    """Model of [`swiglu`][aie.iron.kernels.activation.swiglu]'s LUT build.

    activation/swiglu.cc narrows after every multiply -- ``x*w1``, ``x*w2``, the
    sigmoid factor of
    [`sigmoid_table_ref`][aie.iron.kernels.activation.sigmoid_table_ref] 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_table_ref, and where the silu product is 0 the output is 0, so an
    overflowed ``x*w1`` does not make ``inf * 0``. Each subnormal product is
    flushed to zero, so a subnormal ``x*w2`` zeroes the output through the
    gate.
    """
    return _swiglu_model(x, w1, w2, sigmoid_table_ref)

tanh_lut_ref

tanh_lut_ref(x)

Numpy model of getTanhBf16, the interpolated-LUT tanh.

The kernel clamps x to the table's range [-4, 4 - 1/64], then evaluates slope[e] * x + offset[e] for e = floor(4x) + 16: 32 segments of width 0.25 over [-4, 4). The end segments are the constants -1 and +1, so the clamp changes no finite result, and +-inf gives +-1 rather than 0 * inf. The product is exact in f32 (bf16 carries 8 mantissa bits and 8 + 8 < 24), so the only rounding is the accumulator's store back to bf16, which is why the build using this is judged at one ulp rather than a percentage.

This is what tanh computes with use_lut=True, and what it always computes on aie2. The default aie2p build uses the vtanh instruction instead, which is a coarser approximation with no published spec, so it is judged against tanh_ref and a measured bound.

Source code in utils/pyright_stubs/aie/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`][aie.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`][aie.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 utils/pyright_stubs/aie/iron/kernels/activation.py

def tanh_ref(x):
    """Numpy reference for [`tanh`][aie.iron.kernels.activation.tanh] — element-wise ``tanh(x)``.

    LUT/native-approximation territory; pair with ``rtol=0.128`` when verifying.
    """
    return _rounded(np.tanh(_f64(x)), x)

sigmoid_ref

sigmoid_ref(x)

Numpy reference for sigmoid — 1 / (1 + exp(-x)).

LUT-approximation territory; pair with rtol=0.128 when verifying.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def sigmoid_ref(x):
    """Numpy reference for [`sigmoid`][aie.iron.kernels.activation.sigmoid] — ``1 / (1 + exp(-x))``.

    LUT-approximation territory; pair with ``rtol=0.128`` when verifying.
    """
    with np.errstate(over="ignore"):
        return _rounded(1.0 / (1.0 + np.exp(-_f64(x))), x)

leaky_relu_ref

leaky_relu_ref(x, alpha=0.01)

Numpy reference for leaky_relu.

x if x > 0 else alpha * x. alpha must match the slope the design passes to the kernel at runtime; the kernel takes it as bf16, so it is rounded to bf16 here. Exact up to bf16 rounding; pair with a small rtol when verifying.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def leaky_relu_ref(x, alpha=0.01):
    """Numpy reference for [`leaky_relu`][aie.iron.kernels.activation.leaky_relu].

    ``x if x > 0 else alpha * x``.  ``alpha`` must match the slope the design
    passes to the kernel at runtime; the kernel takes it as bf16, so it is
    rounded to bf16 here.  Exact up to bf16 rounding; pair with a small
    ``rtol`` when verifying.
    """
    alpha = np.float32(bfloat16(alpha))
    xf = x.astype(np.float32)
    return np.where(xf > 0.0, xf, alpha * xf).astype(x.dtype)

swiglu_ref

swiglu_ref(x, w1, w2)

Numpy reference for swiglu: (x * w1) * silu(x * w2).

swiglu.cc forms the two products in bf16, then silu of the second through the tanh LUT (0.5 * (1 + tanh(z / 2))). The reference rounds the two products to bf16 as the kernel does and computes the rest in float64; LUT-approximation territory, pair with rtol=0.128.

Where silu is 0 (x * w2 is -inf, or low enough to underflow) the output is 0, as swiglu_table_ref's is, even where x * w1 overflowed: 0 is the limit as x goes to -inf, and inf * 0 would be NaN.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def swiglu_ref(x, w1, w2):
    """Numpy reference for [`swiglu`][aie.iron.kernels.activation.swiglu]: ``(x * w1) * silu(x * w2)``.

    ``swiglu.cc`` forms the two products in bf16, then ``silu`` of the second
    through the tanh LUT (``0.5 * (1 + tanh(z / 2))``). The reference rounds
    the two products to bf16 as the kernel does and computes the rest in
    float64; LUT-approximation territory, pair with ``rtol=0.128``.

    Where silu is 0 (``x * w2`` is -inf, or low enough to underflow) the
    output is 0, as swiglu_table_ref's is, even where ``x * w1`` overflowed:
    0 is the limit as x goes to -inf, and ``inf * 0`` would be NaN.
    """
    xf = _f64(x)
    xw1 = _f64(round_to(xf * _f64(w1), bfloat16))
    xw2 = _no_neg_inf(_f64(round_to(xf * _f64(w2), bfloat16)))
    with np.errstate(over="ignore", invalid="ignore"):
        silu = xw2 / (1.0 + np.exp(-xw2))
        out = np.where(silu == 0, np.copysign(0.0, xw1) * silu, xw1 * silu)
    return _rounded(out, x)

bf16_exp_ref

bf16_exp_ref(x)

Numpy reference for bf16_exp — element-wise exp(x).

LUT approximation territory; pair with the canonical 12.8% relative tolerance and stop_at_nonfinite=True (the default in count_mismatches) when verifying.

exp(clip(x, -88, 88)), not plain exp(x): the kernel clamps to EXP_BF16_CLAMP before its Q8 fixed-point table lookup (see aie_runtime_lib/AIE2/lut_based_ops.h), so it saturates rather than overflowing. +88 is the largest value the tables carry. exp(-88) is a nonzero bf16 subnormal (about 6.06e-39), not zero: AIE2P preserves it through integer exponent reconstruction, whereas the AIE2 LUT may flush the tail under the absolute tolerance. The clamp also keeps the result in range: exp(88) = 1.65e+38 fits bf16 and float32 where exp(89) would not.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def bf16_exp_ref(x):
    """Numpy reference for [`bf16_exp`][aie.iron.kernels.activation.bf16_exp] — element-wise ``exp(x)``.

    LUT approximation territory; pair with the canonical 12.8% relative
    tolerance and ``stop_at_nonfinite=True`` (the default in
    `count_mismatches`) when verifying.

    ``exp(clip(x, -88, 88))``, not plain ``exp(x)``: the kernel clamps to
    ``EXP_BF16_CLAMP`` before its Q8 fixed-point table lookup (see
    ``aie_runtime_lib/AIE2/lut_based_ops.h``), so it saturates rather than
    overflowing. ``+88`` is the largest value the tables carry.
    ``exp(-88)`` is a nonzero bf16 subnormal (about ``6.06e-39``), not
    zero: AIE2P preserves it through integer exponent reconstruction,
    whereas the AIE2 LUT may flush the tail under the absolute tolerance.
    The clamp also keeps the result in range: ``exp(88) = 1.65e+38`` fits
    bf16 and float32 where ``exp(89)`` would not.
    """
    xf = np.clip(_f64(x), -_EXP_BF16_CLAMP, _EXP_BF16_CLAMP)
    # The clamp rules out overflow, so that warning stays un-suppressed and
    # would now be a real signal. A NaN input still reaches exp -- the kernel
    # declares nonfinite="unspecified" and callers do feed raw bit patterns
    # (programming_examples/basic/vector_exp sweeps all 65536 of them).
    with np.errstate(invalid="ignore"):
        return _rounded(np.exp(xf), x)

exp2f_vec_ref

exp2f_vec_ref(x, min_x: float = -111.0)

Numpy reference for exp2f_vec: exact 2**x.

Unlike the LUT-based refs above, this is float64 2**x (not a reimplementation of the on-device poly): the kernel holds 8.9e-5 relative error or better, several orders tighter than the LUT-based kernels' 12.8% default, so pair with a correspondingly tight tolerance (e.g. rtol=1e-3) rather than the LUT default.

The kernel clamps its input to min_x before evaluating (see the factory's min_x), so the reference does the same: 2**-5000 is 2**min_x on the device, not zero. Pass the factory's min_x.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def exp2f_vec_ref(x, min_x: float = -111.0):
    """Numpy reference for [`exp2f_vec`][aie.iron.kernels.activation.exp2f_vec]: exact ``2**x``.

    Unlike the LUT-based refs above, this is float64 ``2**x`` (not a
    reimplementation of the on-device poly): the kernel holds 8.9e-5
    relative error or better, several orders tighter than the LUT-based
    kernels' 12.8% default, so pair with a correspondingly tight
    tolerance (e.g. ``rtol=1e-3``) rather than the LUT default.

    The kernel clamps its input to ``min_x`` before evaluating (see the
    factory's ``min_x``), so the reference does the same: ``2**-5000`` is
    ``2**min_x`` on the device, not zero. Pass the factory's ``min_x``.
    """
    xf = np.maximum(x.astype(np.float64), min_x)
    return np.exp2(xf).astype(x.dtype)

log_f32_ref

log_f32_ref(x, offset=0.0)

Numpy reference for log_f32.

The kernel adds offset in float32, so the sum rounds as float32 here too; its log is taken in float64 and rounded once to bf16.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def log_f32_ref(x, offset=0.0):
    """Numpy reference for [`log_f32`][aie.iron.kernels.activation.log_f32].

    The kernel adds ``offset`` in float32, so the sum rounds as float32 here
    too; its log is taken in float64 and rounded once to bf16.
    """
    z = np.asarray(x, dtype=np.float32) + np.float32(offset)
    return np.log(z.astype(np.float64)).astype(bfloat16)

log_f32_sample

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

Positive float32 over every binade, the first call led by edges.

The edges are one and its neighbours, the reduction's split point sqrt(1/2) and sqrt(2) with theirs, e, a log-mel floor (1e-3), and the smallest normal and largest finite float32.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def log_f32_sample(rng, calls: int, *, tile_size: int) -> list:
    """Positive float32 over every binade, the first call led by edges.

    The edges are one and its neighbours, the reduction's split point
    ``sqrt(1/2)`` and ``sqrt(2)`` with theirs, ``e``, a log-mel floor
    (``1e-3``), and the smallest normal and largest finite float32.
    """
    x = (10.0 ** rng.uniform(-37, 38, calls * tile_size)).astype(np.float32)
    edges = np.array(
        [
            0x3F800000,
            0x3F800001,
            0x3F7FFFFF,
            0x3F000000,
            0x40000000,
            0x3F3504F2,
            0x3F3504F3,
            0x3F3504F4,
            0x3FB504F2,
            0x3FB504F3,
            0x3FB504F4,
            0x402DF854,
            0x3A83126F,
            0x00800000,
            0x7F7FFFFF,
            0x7149F2CA,
        ],
        dtype=np.uint32,
    ).view(np.float32)
    x[: edges.size] = edges[: x.size]
    return [x.reshape(calls, tile_size)]

log10_clamp_f32_ref

log10_clamp_f32_ref(
    x,
    peak,
    mel_floor=1e-10,
    decades=8.0,
    shift=4.0,
    scale=0.25,
)

Numpy reference for log10_clamp_f32.

WhisperFeatureExtractor's float32 steps after its log10, each log correctly rounded, the result rounded once to bf16.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def log10_clamp_f32_ref(x, peak, mel_floor=1e-10, decades=8.0, shift=4.0, scale=0.25):
    """Numpy reference for [`log10_clamp_f32`][aie.iron.kernels.activation.log10_clamp_f32].

    WhisperFeatureExtractor's float32 steps after its log10, each log
    correctly rounded, the result rounded once to bf16.
    """
    v, _ = _log10_clamped(x, peak, mel_floor, decades)
    return ((v + np.float32(shift)) * np.float32(scale)).astype(bfloat16)

log10_clamp_f32_sample

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

Mel energies over seventeen decades and each call's largest as its peak.

An eighth of the energies are zero, as a padded clip's silent frames are.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def log10_clamp_f32_sample(rng, calls: int, *, tile_size: int) -> list:
    """Mel energies over seventeen decades and each call's largest as its peak.

    An eighth of the energies are zero, as a padded clip's silent frames are.
    """
    n = calls * tile_size
    x = 10.0 ** rng.uniform(-14, 3, n) * (rng.uniform(size=n) >= 0.125)
    x = x.astype(np.float32).reshape(calls, tile_size)
    return [x, x.max(axis=-1, keepdims=True)]

softmax_ref

softmax_ref(
    x,
    addend=None,
    mask=None,
    *,
    tile_size: int = 1024,
    soft_cap: float | None = None,
    scale: float | None = None
)

Numpy reference for softmax.

The AIE kernel computes softmax independently per tile_size-element tile (no cross-tile reduction), so the reference splits x the same way before applying the softmax, in float64. x.size must be a multiple of tile_size. The logits are x + addend, as cap * tanh(. / cap) under a soft_cap, plus mask -- float path options, never given alongside scale -- then scale * logits where scale is given (bf16 input only); the result is rounded to x's dtype, float32 for the float path, whose bf16 output its contract judges by a bound.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def softmax_ref(
    x,
    addend=None,
    mask=None,
    *,
    tile_size: int = 1024,
    soft_cap: float | None = None,
    scale: float | None = None,
):
    """Numpy reference for [`softmax`][aie.iron.kernels.activation.softmax].

    The AIE kernel computes softmax independently per ``tile_size``-element
    tile (no cross-tile reduction), so the reference splits ``x`` the same
    way before applying the softmax, in float64.  ``x.size`` must be a
    multiple of ``tile_size``. The logits are ``x + addend``, as
    ``cap * tanh(. / cap)`` under a ``soft_cap``, plus ``mask`` -- float path
    options, never given alongside ``scale`` -- then ``scale * logits`` where
    ``scale`` is given (bf16 input only); the result is rounded to ``x``'s
    dtype, float32 for the float path, whose bf16 output its contract judges
    by a bound.
    """
    flat = _softmax_logits(x, addend, mask, tile_size=tile_size, soft_cap=soft_cap)
    if scale is not None:
        flat = flat * scale
    flat = flat - flat.max(axis=1, keepdims=True)
    exp = np.exp(flat)
    out = exp / exp.sum(axis=1, keepdims=True)
    return _rounded(out.reshape(np.shape(x)), x)

softmax_span_ref

softmax_span_ref(
    x, lo, hi, *, tile_size: int, scale: float | None = None
)

Numpy reference for softmax's span path.

softmax_ref of each row's live [lo, hi) alone: scale * logits there, and exactly 0 outside it at any scale, as the kernel zeroes those exponentials. lo and hi are one bound for every row, as the kernel's scalars, or one per row.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def softmax_span_ref(x, lo, hi, *, tile_size: int, scale: float | None = None):
    """Numpy reference for [`softmax`][aie.iron.kernels.activation.softmax]'s ``span`` path.

    ``softmax_ref`` of each row's live ``[lo, hi)`` alone: ``scale *
    logits`` there, and exactly 0 outside it at any scale, as the kernel
    zeroes those exponentials. ``lo`` and ``hi`` are one bound for every
    row, as the kernel's scalars, or one per row.
    """
    logits = _softmax_logits(x, None, None, tile_size=tile_size, soft_cap=None)
    if scale is not None:
        logits = logits * scale
    cols = np.arange(tile_size)
    live = (cols >= np.reshape(lo, (-1, 1))) & (cols < np.reshape(hi, (-1, 1)))
    peak = np.where(live, logits, -np.inf).max(axis=1, keepdims=True)
    exp = np.exp(np.where(live, logits - peak, -np.inf))
    out = exp / exp.sum(axis=1, keepdims=True)
    return _rounded(out.reshape(np.shape(x)), x)

softmax_span_sample

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

span path inputs: a tile a call, with a random live [lo, hi) each.

The first call (when there are at least two) is unmasked, [0, tile_size); the second keeps a single element, [0, 1), the narrowest live span a row has.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def softmax_span_sample(rng, calls: int, *, tile_size: int) -> list:
    """``span`` path inputs: a tile a call, with a random live ``[lo, hi)`` each.

    The first call (when there are at least two) is unmasked, ``[0,
    tile_size)``; the second keeps a single element, ``[0, 1)``, the
    narrowest live span a row has.
    """
    x = (rng.uniform(-1.0, 1.0, (calls, tile_size)) * 20.0).astype(bfloat16)
    lo = rng.integers(0, tile_size, calls)
    span = rng.integers(1, tile_size + 1 - lo)
    hi = lo + span
    if calls >= 2:
        lo[0], hi[0] = 0, tile_size
        lo[1], hi[1] = 0, 1
    return [x, lo.astype(np.int32), hi.astype(np.int32)]

softmax_f32_bound

softmax_f32_bound(
    x,
    addend=None,
    mask=None,
    *,
    tile_size: int,
    soft_cap: float | None = None
)

Return the float path's per-element bound on |y - softmax_ref|.

Four bf16 roundings, 2^x's and aie::inv's errors and the sum's float32 rounding, plus twice the logits' error, 2^-21 of a kept logit's size. Capped, that is the row's largest logit plus the cap's tanh. Uncapped, the kernel centres the logits first, so the size is 128 (the logits near the max) or a logit whose x + addend or + mask rounds in float32. Below 2^-125 an output may be +0.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def softmax_f32_bound(
    x, addend=None, mask=None, *, tile_size: int, soft_cap: float | None = None
):
    """Return the float path's per-element bound on ``|y - softmax_ref|``.

    Four bf16 roundings, 2^x's and ``aie::inv``'s errors and the sum's
    float32 rounding, plus twice the logits' error, ``2^-21`` of a kept
    logit's size. Capped, that is the row's largest logit plus the cap's
    tanh. Uncapped, the kernel centres the logits first, so the size is 128
    (the logits near the max) or a logit whose ``x + addend`` or ``+ mask``
    rounds in float32. Below ``2^-125`` an output may be +0.
    """
    z = _softmax_logits(x, addend, mask, tile_size=tile_size, soft_cap=soft_cap)
    if soft_cap is not None:
        size = np.abs(z)
    else:
        rows = z.shape
        s = _f64(x).reshape(rows)
        if addend is not None:
            s = s + _f64(addend).reshape(rows)
        size = np.maximum(np.where(s.astype(np.float32) != s, np.abs(s), 0), 128.0)
        if mask is not None:
            sm = s.astype(np.float32).astype(np.float64) + _f64(mask).reshape(rows)
            size = np.maximum(
                size, np.where(sm.astype(np.float32) != sm, np.abs(sm), 0)
            )
    kept = np.where(z > -1e30, size, 0.0)
    logit = 2.0**-21 * kept.max(axis=1, keepdims=True)
    if soft_cap is not None:
        logit = logit + soft_cap * _SOFT_CAP_ERR
    rel = _SOFTMAX_F32_REL + (tile_size / 32 + 5) * 2.0**-24 + 2 * logit
    p = softmax_ref(
        _f64(x), addend, mask, tile_size=tile_size, soft_cap=soft_cap
    ).reshape(z.shape)
    return np.maximum(rel * p, 2.0**-125).reshape(np.shape(x))

softmax_f32_sample

softmax_f32_sample(
    rng,
    calls: int,
    *,
    tile_size: int,
    soft_cap: float | None,
    masked: bool,
    addend: bool
) -> list

Float path inputs: logits to |x| >= 32, near ties, and masked-out tails.

Each call's logits are uniform on a span that grows over the calls, from 4 to 48 (and to 400 under a cap, which saturates it). The first row leads with +/-32 and two near ties, the second is all 3e38 and the third leads with +/-2e38. The addend is uniform on [-8, 8]. A mask keeps a random leading run of at least one element and masks the tail with lowest and -inf, and one kept element in eight with lowest; uncapped, the fourth row is 2^127 and 2^127 - 2^104 under a mask of -2^127.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def softmax_f32_sample(
    rng,
    calls: int,
    *,
    tile_size: int,
    soft_cap: float | None,
    masked: bool,
    addend: bool,
) -> list:
    """Float path inputs: logits to |x| >= 32, near ties, and masked-out tails.

    Each call's logits are uniform on a span that grows over the calls, from
    4 to 48 (and to 400 under a cap, which saturates it). The first row leads
    with +/-32 and two near ties, the second is all 3e38 and the third leads
    with +/-2e38. The addend is uniform on [-8, 8]. A mask keeps a random
    leading run of at least one element and masks the tail with lowest and
    -inf, and one kept element in eight with lowest; uncapped, the fourth row
    is 2^127 and 2^127 - 2^104 under a mask of -2^127.
    """
    top = 400.0 if soft_cap is not None else 48.0
    spans = np.geomspace(4.0, top, calls)[:, None]
    x = (rng.uniform(-1.0, 1.0, (calls, tile_size)) * spans).astype(np.float32)
    lead = np.array([32, -32, 5, np.nextafter(np.float32(5), 6)], np.float32)
    x[0, : lead.size] = lead[:tile_size]
    if calls >= 3:
        x[1] = 3e38
        x[2, :2] = [2e38, -2e38]
    inputs = [x]
    if addend:
        inputs.append(rng.uniform(-8.0, 8.0, (calls, tile_size)).astype(np.float32))
    if masked:
        lowest = float(finfo(bfloat16).min)
        keep = rng.integers(1, tile_size + 1, calls)
        cols = np.arange(tile_size)
        m = np.where(cols[None, :] < keep[:, None], 0.0, lowest)
        m[:, 1::2] = np.where(m[:, 1::2] < 0, -np.inf, 0.0)
        m[:, 1:][rng.random((calls, tile_size - 1)) < 1 / 8] = lowest
        if soft_cap is None and calls >= 4:
            x[3] = 2.0**127 - 2.0**104
            x[3, 0] = 2.0**127
            m[3] = -(2.0**127)
        inputs.append(m.astype(bfloat16))
    return inputs

softmax_banded_mask

softmax_banded_mask(
    row_offset, *, window: int, period: int, tile_size: int
)

softmax_rows_banded_f32_bf16's band, one call a row.

Restated as softmax_ref's explicit bf16 mask: query position q = row_offset % period keeps column j where j >= max(q % 2, window - 1 - q + q % 2) and j < window + q % 2 (an odd query's band starts a column later).

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def softmax_banded_mask(row_offset, *, window: int, period: int, tile_size: int):
    """``softmax_rows_banded_f32_bf16``'s band, one call a row.

    Restated as [`softmax_ref`][aie.iron.kernels.activation.softmax_ref]'s
    explicit bf16 mask: query position ``q = row_offset % period`` keeps
    column ``j`` where ``j >= max(q % 2, window - 1 - q + q % 2)`` and
    ``j < window + q % 2`` (an odd query's band starts a column later).
    """
    q = np.asarray(row_offset) % period
    parity = q % 2
    lo = np.maximum(parity, window - 1 - q + parity)
    hi = window + parity
    j = np.arange(tile_size)[None, :]
    band = (j >= lo[:, None]) & (j < hi[:, None])
    lowest = float(finfo(bfloat16).min)
    return np.where(band, 0.0, lowest).astype(bfloat16)

softmax_banded_ref

softmax_banded_ref(
    x,
    addend,
    row_offset,
    *,
    tile_size: int,
    window: int,
    period: int,
    soft_cap
)

Numpy reference for softmax's banded path.

softmax_ref with the band computed by softmax_banded_mask.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def softmax_banded_ref(
    x, addend, row_offset, *, tile_size: int, window: int, period: int, soft_cap
):
    """Numpy reference for ``softmax``'s banded path.

    ``softmax_ref`` with the band computed by
    [`softmax_banded_mask`][aie.iron.kernels.activation.softmax_banded_mask].
    """
    mask = softmax_banded_mask(
        row_offset, window=window, period=period, tile_size=tile_size
    )
    return softmax_ref(x, addend, mask, tile_size=tile_size, soft_cap=soft_cap)

softmax_banded_bound

softmax_banded_bound(
    x,
    addend,
    row_offset,
    *,
    tile_size: int,
    window: int,
    period: int,
    soft_cap
)

softmax_f32_bound.

Takes the same band softmax_banded_ref does.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def softmax_banded_bound(
    x, addend, row_offset, *, tile_size: int, window: int, period: int, soft_cap
):
    """[`softmax_f32_bound`][aie.iron.kernels.activation.softmax_f32_bound].

    Takes the same band ``softmax_banded_ref`` does.
    """
    mask = softmax_banded_mask(
        row_offset, window=window, period=period, tile_size=tile_size
    )
    return softmax_f32_bound(x, addend, mask, tile_size=tile_size, soft_cap=soft_cap)

softmax_banded_sample

softmax_banded_sample(
    rng,
    calls: int,
    *,
    tile_size: int,
    window: int,
    period: int,
    soft_cap
) -> list

Float path inputs, with a row_offset a call.

softmax_f32_sample's x/addend (masked=False, since the band replaces a data mask), plus row_offset: the sequence's own start (the band's boundary effect), then positions spanning past one period, so the wrap is covered too.

Source code in utils/pyright_stubs/aie/iron/kernels/activation.py

def softmax_banded_sample(
    rng, calls: int, *, tile_size: int, window: int, period: int, soft_cap
) -> list:
    """Float path inputs, with a ``row_offset`` a call.

    ``softmax_f32_sample``'s ``x``/``addend`` (``masked=False``, since the
    band replaces a data mask), plus ``row_offset``: the sequence's own
    start (the band's boundary effect), then positions spanning past one
    ``period``, so the wrap is covered too.
    """
    x, addend = softmax_f32_sample(
        rng, calls, tile_size=tile_size, soft_cap=soft_cap, masked=False, addend=True
    )
    start = np.arange(min(calls, period))
    rest = rng.integers(0, 4 * period, max(0, calls - period))
    row_offset = np.concatenate([start, rest])[:calls].astype(np.int32)
    return [x, addend, row_offset]

Normalization

rms_norm, rms_norm_eps, layer_norm, and layer_norm_eps 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.

Each norm has two C entries, one fixing eps at 1e-5 with (in, out, cols) and an _eps one taking (in, out, cols, epsilon); a factory names the entry it builds, so the two keep their names rather than one ABI changing.

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 utils/pyright_stubs/aie/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.
    """
    return _norm("rms_norm", tile_size, cols, eps=False, layer=False)

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 utils/pyright_stubs/aie/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)``."""
    return _norm("rms_norm_eps", tile_size, cols, eps=True, layer=False)

layer_norm

layer_norm(
    tile_size: int = 1024, *, cols: int | None = None
) -> ExternalFunction

Layer-norm a bf16 row; (in, out, cols), gamma=1, beta=0, eps=1e-5.

cols aliases tile_size, a positive multiple of 16 (the source processes whole 16-lane halves, without a scalar tail).

Source code in utils/pyright_stubs/aie/iron/kernels/norm.py

def layer_norm(tile_size: int = 1024, *, cols: int | None = None) -> ExternalFunction:
    """Layer-norm a bf16 row; ``(in, out, cols)``, gamma=1, beta=0, eps=1e-5.

    ``cols`` aliases ``tile_size``, a positive multiple of 16 (the source
    processes whole 16-lane halves, without a scalar tail).
    """
    return _norm("layer_norm", tile_size, cols, eps=False, layer=True, multiple=16)

layer_norm_eps

layer_norm_eps(
    tile_size: int = 1024, *, cols: int | None = None
) -> ExternalFunction

Layer-norm a bf16 row (gamma=1, beta=0); design passes (in, out, cols, epsilon).

Parameters:

Name	Type	Description	Default
tile_size	int	The row length, a positive multiple of 16.	1024
cols	int | None	An alias for tile_size.	None

Returns:

Type	Description
ExternalFunction	The kernel.

Source code in utils/pyright_stubs/aie/iron/kernels/norm.py

def layer_norm_eps(
    tile_size: int = 1024, *, cols: int | None = None
) -> ExternalFunction:
    """Layer-norm a bf16 row (gamma=1, beta=0); design passes ``(in, out, cols, epsilon)``.

    Args:
        tile_size: The row length, a positive multiple of 16.
        cols: An alias for ``tile_size``.

    Returns:
        The kernel.
    """
    return _norm("layer_norm_eps", tile_size, cols, eps=True, layer=True, multiple=16)

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 utils/pyright_stubs/aie/iron/kernels/norm.py

def rms_norm_ref(x, *, eps: float = 1e-5):
    """Numpy reference for [`rms_norm`][aie.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 utils/pyright_stubs/aie/iron/kernels/norm.py

def layer_norm_ref(x, *, eps: float = 1e-5):
    """Numpy reference for [`layer_norm`][aie.iron.kernels.norm.layer_norm]: ``(x - mean) / sqrt(var + eps)``."""
    xf = x.astype(np.float32)
    mean = np.mean(xf, axis=-1, keepdims=True)
    var = np.mean((xf - mean) ** 2, axis=-1, keepdims=True)
    return ((xf - mean) / np.sqrt(var + eps)).astype(x.dtype)

Transformer blocks

Transformer building blocks: rms_norm, layer_norm (bf16, f32, affine+cast), rope, mm_activation_epilogue.

The bf16 norms and RoPE are re-exported from norm and datamovement; they support aie2 and aie2p, with cols as an alias for tile_size. The activation epilogue and the f32/affine norms take their aie2p source on both generations. Each processes one row (cols elements) per call; the row length is a scalar Param the factory binds to cols. These are the kernels programming_examples/ml/{norm,rope,mm_activation_epilogue} build.

layer_norm_f32

layer_norm_f32(cols: int = 4096) -> ExternalFunction

Row-wise LayerNorm on float32 in and out (gamma = 1, beta = 0, eps 1e-5).

A separate factory rather than a dtype of layer_norm: this one is held to atol 2e-6 (1e-3 in the portable build) instead of the bf16 tolerance, which its reference meets only by computing the variance two-pass in float64. Merging them would put that numerical difference behind a dtype switch.

Parameters:

Name	Type	Description	Default
cols	int	Elements per row (multiple of 16).	4096

Source code in utils/pyright_stubs/aie/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`][aie.iron.kernels.norm.layer_norm]: this one is held to atol 2e-6
    (1e-3 in the portable build) instead of the bf16 tolerance, which its
    reference meets only by computing the variance two-pass in float64.
    Merging them would put that numerical difference behind a dtype switch.

    Args:
        cols: Elements per row (multiple of 16).
    """
    return _row_kernel(
        "layer_norm_f32",
        "layer_norm_f32",
        "transformer/layer_norm_f32.cc",
        cols,
        np.float32,
        np.float32,
        tuple(
            (build, layer_norm_f32_ref, _NORM_F32_TOLERANCE.get(build.tuned, _NORM_F32))
            for build in BUILDS
        ),
        6 * cols,
        # 896 B tuned for aie2p, 160 B tuned for aie2, 832 B untuned on aie2p
        # (672 B on aie2); the 1024 B default covers every build.
    )

layer_norm_affine_cast

layer_norm_affine_cast(
    cols: int = 4096,
) -> ExternalFunction

Row-wise LayerNorm, f32 in, per-column gamma/beta, bf16 out.

The second argument holds gamma (cols values) followed by beta (cols values) as float32: a tensor Param, which the generic builder bakes into a core buffer.

Parameters:

Name	Type	Description	Default
cols	int	Elements per row (multiple of 16).	4096

Source code in utils/pyright_stubs/aie/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),),
            per_build=tuple(
                (
                    build,
                    layer_norm_affine_cast_ref,
                    _AFFINE_TOLERANCE.get(build.tuned, _NORM_BF16),
                )
                for build in BUILDS
            ),
            acc_dtype=np.float32,
            reduction=cols,
            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 utils/pyright_stubs/aie/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
    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),),
            true_function=mm_activation_epilogue_ref,
            per_build=tuple(
                (
                    (build, mm_activation_epilogue_lut_ref, _EPILOGUE_LUT_TOLERANCE)
                    if build.tuned == "aie2"
                    else (build, mm_activation_epilogue_ref, _EPILOGUE_TOLERANCE)
                )
                for build in BUILDS
            ),
            acc_dtype=np.float32,
            reduction=1,
            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 utils/pyright_stubs/aie/iron/kernels/transformer.py

def layer_norm_f32_ref(x):
    """Numpy reference for [`layer_norm_f32`][aie.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 utils/pyright_stubs/aie/iron/kernels/transformer.py

def layer_norm_affine_cast_ref(x, gamma_beta):
    """Numpy reference for [`layer_norm_affine_cast`][aie.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 utils/pyright_stubs/aie/iron/kernels/transformer.py

def mm_activation_epilogue_ref(x, mode):
    """Numpy reference for [`mm_activation_epilogue`][aie.iron.kernels.transformer.mm_activation_epilogue].

    ``mode`` 0 identity, 1 ``x * sigmoid(x)``, 2 tanh-approximation GELU,
    3 ``max(x, 0)``.
    """
    x32 = x.astype(np.float32)
    mode = int(mode)
    if mode == 0:
        return x32.astype(x.dtype)
    if mode == 1:
        with np.errstate(over="ignore"):
            return (x32 / (1.0 + np.exp(-x32))).astype(x.dtype)
    if mode == 2:
        # gelu's limit at -inf is 0, which the most negative float gives
        # rather than -inf * 0.
        x32 = np.maximum(x32, -np.finfo(np.float32).max)
        inner = 0.7978845608 * (x32 + 0.044715 * x32**3)
        return (0.5 * x32 * (1.0 + np.tanh(inner))).astype(x.dtype)
    if mode == 3:
        return np.maximum(x32, 0.0).astype(x.dtype)
    raise ValueError(f"mm_activation_epilogue mode must be 0, 1, 2 or 3, got {mode}")

mm_activation_epilogue_lut_ref

mm_activation_epilogue_lut_ref(x, mode)

Model of mm_activation_epilogue on aie2.

Follows mm_activation_epilogue.cc's roundings around getTanhBf16 (tanh_lut_ref). SiLU splits x into hi, its top 16 bits, and lo, bf16(x - hi), and multiplies each by the bf16 sigmoid (bf16(t + 1)) / 2, where t is the table's tanh of bf16(x) / 2 narrowed to bf16. hi is finite for any finite x, so huge inputs give about x or 0 rather than NaN; +-inf still gives NaN. GELU runs in bf16: x, x * x and the inner polynomial are each rounded before the next step, and the output is bf16(x / 2) * bf16(t + 1), all reading x clamped at -8 so -inf gives 0. Both return +0 where IEEE arithmetic gives -0, as the accumulator does. Identity and ReLU are exact.

Source code in utils/pyright_stubs/aie/iron/kernels/transformer.py

def mm_activation_epilogue_lut_ref(x, mode):
    """Model of [`mm_activation_epilogue`][aie.iron.kernels.transformer.mm_activation_epilogue] on aie2.

    Follows mm_activation_epilogue.cc's roundings around getTanhBf16
    ([`tanh_lut_ref`][aie.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, and the filter tables of torch's antialiased bicubic resize.

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 utils/pyright_stubs/aie/iron/kernels/vision.py

def rgba2hue(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Convert a line of RGBA pixels to hue values (full-range, 0..255)."""
    _require_vector_alignment("rgba2hue", line_width, 32, param="line_width")
    # lut_inv.h pins its gather pair with AIE_BANK_A/AIE_BANK_B.
    flags = [f"-I{_runtime_lib_include()}"]
    return _color_convert_kernel(
        "rgba2hueLine",
        "rgba2hue.cc",
        line_width * 4,
        line_width,
        use_chess=use_chess,
        compile_flags=flags,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, line_width),),
            reference=rgba2hue_ref,
            acc_dtype=np.int32,
            reduction=1,
            tolerance=Tolerance.exact(note="integer reciprocal, no rounding slack"),
            uses_lut=True,
        ),
    )

threshold

threshold(
    line_width: int = 1920,
    dtype: type = uint8,
    use_chess: bool = False,
) -> ExternalFunction

Apply a threshold operation to a line of pixels.

Parameters:

Name	Type	Description	Default
line_width	int	Number of elements per line.	1920
dtype	type	Element data type, np.uint8 or np.int16. The source's 32-bit branch multiplies int32 data by int16 coefficients, a MAC AIE2 does not have, so it does not compile and is not offered.	uint8
use_chess	bool	When True, build the .o with xchesscc_wrapper instead of Peano.	False

Raises:

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

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

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 utils/pyright_stubs/aie/iron/kernels/vision.py

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 utils/pyright_stubs/aie/iron/kernels/vision.py

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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/iron/kernels/vision.py

def rgba2gray(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Convert an RGBA line to grayscale."""
    flags = []
    if not use_chess and _tuned_arch() == "aie2p":
        # The pre-RA pipeliner's schedule of the 64-pixel loop ends up at II13
        # after register allocation; the postpipeliner finds II11 (llvm-aie#1066).
        flags += ["-mllvm", "--aie-force-postpipeliner"]
    return _color_convert_kernel(
        "rgba2grayLine",
        "rgba2gray.cc",
        line_width * 4,
        line_width,
        use_chess=use_chess,
        compile_flags=flags,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, line_width),),
            reference=rgba2gray_ref,
            tolerance=Tolerance.exact(
                note="measured bit-exact against the reference over every data case"
            ),
            acc_dtype=np.int32,
            reduction=3,
        ),
    )

filter2d

filter2d(
    line_width: int = 1920, use_chess: bool = False
) -> ExternalFunction

Apply a 3x3 2D convolution filter across three input lines.

Source code in utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/iron/kernels/vision.py

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

levels

levels(
    in_bytes: int = 3 * 1920,
    out_width: int = 1920,
    *,
    bgr: bool = False
) -> ExternalFunction

Map a row of uint8 RGB pixels through per-channel tables to bf16 pixels of four channels.

levelsLine(in, table, out, first, count, at, row, lo, hi) writes count pixels of the row in, from pixel first, to the row out from pixel at, each channel's level looked up in its own 256 entries of table; every other pixel of out is the four fill values after them. Output row row places its run only for lo <= row < hi, so a design's core can take every row of an image whatever rows a call places. A run past either row is cut at its end, and a negative first or at writes the fill alone. The tables hold whatever a host maps a level to, such as an image processor's rescale and normalization, computed once and exactly; with bgr the first output channel reads the input's blue.

Parameters:

Name	Type	Description	Default
in_bytes	int	The input row's bytes, three a pixel; a multiple of 4.	3 * 1920
out_width	int	The output row's pixels.	1920
bgr	bool	Whether the output's channels are blue, green, red.	False

Returns:

Type	Description
ExternalFunction	ExternalFunction for levelsLine.

Raises:

Type	Description
ValueError	When in_bytes is not a positive multiple of 4 or out_width is below 1.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def levels(
    in_bytes: int = 3 * 1920, out_width: int = 1920, *, bgr: bool = False
) -> ExternalFunction:
    """Map a row of uint8 RGB pixels through per-channel tables to bf16 pixels of four channels.

    ``levelsLine(in, table, out, first, count, at, row, lo, hi)`` writes
    ``count`` pixels of the row ``in``, from pixel ``first``, to the row
    ``out`` from pixel ``at``, each channel's level looked up in its own 256
    entries of ``table``; every other pixel of ``out`` is the four fill
    values after them. Output row ``row`` places its run only for ``lo <=
    row < hi``, so a design's core can take every row of an image whatever
    rows a call places. A run past either row is cut at its end, and a
    negative ``first`` or ``at`` writes the fill alone. The tables hold whatever a host maps a
    level to, such as an image processor's rescale and normalization,
    computed once and exactly; with ``bgr`` the first output channel reads
    the input's blue.

    Args:
        in_bytes: The input row's bytes, three a pixel; a multiple of 4.
        out_width: The output row's pixels.
        bgr: Whether the output's channels are blue, green, red.

    Returns:
        ExternalFunction for ``levelsLine``.

    Raises:
        ValueError: When ``in_bytes`` is not a positive multiple of 4 or
            ``out_width`` is below 1.
    """
    if in_bytes < 4 or in_bytes % 4:
        raise ValueError(
            f"levels() needs in_bytes a positive multiple of 4, got {in_bytes}."
        )
    if out_width < 1:
        raise ValueError(f"levels() needs out_width >= 1, got {out_width}.")
    return _make_extern(
        "levelsLine",
        _kernel_source("vision/levels.cc"),
        [
            np.ndarray[(in_bytes,), np.dtype[np.uint8]],
            np.ndarray[(LEVELS_TABLE,), np.dtype[bfloat16]],
            np.ndarray[(4 * out_width,), np.dtype[bfloat16]],
        ]
        + [np.int32] * 6,
        compile_flags=[
            f"-DIN_BYTES={in_bytes}",
            f"-DOUT_WIDTH={out_width}",
            f"-DBGR={int(bgr)}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Param, Out) + (Param,) * 6,
            reference=lambda rows, table, first, count, at, row, lo, hi: levels_ref(
                rows,
                table,
                first,
                count,
                at,
                row,
                lo,
                hi,
                out_width=out_width,
                bgr=bgr,
            ),
            tolerance=Tolerance.exact(note="a table lookup"),
            ops_per_call=0,
        ),
    )

grid_decode

grid_decode(
    width: int = 52,
    channels: int = 144,
    *,
    obj_at: int = 16,
    cls_at: int = 48,
    classes: int = 80
) -> ExternalFunction

Decode a row of an anchor-free detector's prediction map into float32 boxes.

grid_decode(in, out, y, stride) reads width cells of channels bf16, a cell's box (dx, dy, dw, dh) first, its objectness at obj_at and classes scores from cls_at, and writes a row of 5 + classes float32 for each: (dx + x) * stride, (dy + y) * stride, exp(dw) * stride, exp(dh) * stride, then the objectness and scores, as YOLOX's decode_outputs computes them in float32 for grid cell (x, y) of a map of stride, a power of two (so the multiply is exact); the exponential is exp_f32. Each sum is one native float add, not correctly rounded on AIE2P (nms's own kernel says why), which the contract's tolerance covers.

Parameters:

Name	Type	Description	Default
width	int	The cells of a row.	52
channels	int	The bf16 channels of a cell.	144
obj_at	int	The objectness's channel.	16
cls_at	int	The first class score's channel.	48
classes	int	The class scores of a cell.	80

Returns:

Type	Description
ExternalFunction	ExternalFunction for grid_decode.

Raises:

Type	Description
ValueError	When a field lies past a cell's channels or inside its box, or width or classes is below 1.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def grid_decode(
    width: int = 52,
    channels: int = 144,
    *,
    obj_at: int = 16,
    cls_at: int = 48,
    classes: int = 80,
) -> ExternalFunction:
    """Decode a row of an anchor-free detector's prediction map into float32 boxes.

    ``grid_decode(in, out, y, stride)`` reads ``width`` cells of
    ``channels`` bf16, a cell's box ``(dx, dy, dw, dh)`` first, its
    objectness at ``obj_at`` and ``classes`` scores from ``cls_at``, and
    writes a row of ``5 + classes`` float32 for each: ``(dx + x) * stride``,
    ``(dy + y) * stride``, ``exp(dw) * stride``, ``exp(dh) * stride``, then
    the objectness and scores, as YOLOX's ``decode_outputs`` computes them
    in float32 for grid cell ``(x, y)`` of a map of ``stride``, a power of
    two (so the multiply is exact); the exponential is ``exp_f32``. Each
    sum is one native float add, not correctly rounded on AIE2P (``nms``'s
    own kernel says why), which the contract's tolerance covers.

    Args:
        width: The cells of a row.
        channels: The bf16 channels of a cell.
        obj_at: The objectness's channel.
        cls_at: The first class score's channel.
        classes: The class scores of a cell.

    Returns:
        ExternalFunction for ``grid_decode``.

    Raises:
        ValueError: When a field lies past a cell's channels or inside its
            box, or ``width`` or ``classes`` is below 1.
    """
    if width < 1 or classes < 1:
        raise ValueError(
            f"grid_decode() needs width >= 1 and classes >= 1, got {width} and {classes}."
        )
    if max(obj_at + 1, cls_at + classes) > channels or min(obj_at, cls_at) < 4:
        raise ValueError(
            f"grid_decode() needs objectness at {obj_at} and {classes} scores from "
            f"{cls_at} inside {channels} channels, after the box."
        )
    return _make_extern(
        "grid_decode",
        _kernel_source("vision/grid_decode.cc"),
        [
            np.ndarray[(width * channels,), np.dtype[bfloat16]],
            np.ndarray[(width * (5 + classes),), np.dtype[np.float32]],
            np.int32,
            np.int32,
        ],
        compile_flags=[
            f"-DWIDTH={width}",
            f"-DCHANNELS={channels}",
            f"-DOBJ={obj_at}",
            f"-DCLS={cls_at}",
            f"-DCLASSES={classes}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param, Param),
            reference=lambda cells, y, stride: grid_decode_ref(
                cells,
                y,
                stride,
                width=width,
                channels=channels,
                obj_at=obj_at,
                cls_at=cls_at,
                classes=classes,
            ),
            acc_dtype=np.float32,
            tolerance=Tolerance.relative(
                1e-6,
                note="exp_f32 is within 9.41 float32 ulps (5.6e-7) over every "
                "input, numpy's exp within one; each sum is one native float "
                "add, not correctly rounded on AIE2P (nms.cc), and the stride "
                "and the copies are exact",
            ),
        ),
    )

nms

nms(
    width: int = 52, classes: int = 80, cap: int = 320
) -> _NmsKernel

Score-threshold and class-aware NMS over an anchor-free detector's boxes.

Bit for bit torchvision's CPU nms under batched_nms's own choice of path (see nms_ref, the full arithmetic model). Three entry points share one core's state, and a core runs, once per image:

nms_setup(conf_thre)
for each row of each pyramid level: nms_score(row, width)
nms_finish(out, count, nms_thre, class_agnostic)

row is grid_decode's output row: width cells of 5 + classes float32, a cell's box (cx, cy, w, h), its objectness, then classes scores; width is a runtime count up to this build's width, so one build serves every pyramid level. A cell survives at objectness * its best class's score, ties to the lowest class id. Past cap survivors, a new one only replaces the held one of the lowest score (ties to the latest arrival), so a core always holds the top cap survivors by score and nms_finish's output equals torchvision's unbounded one for every detection ranked within that top cap, at any survivor count. out is cap rows of x1, y1, x2, y2, score, class (class as float32), zeros past count[0]'s kept detections, in descending score order; count[1] is every survivor seen, cap or not, and also what picks class_agnostic 0's path between batched_nms's coordinate-offset trick (1000 survivors or fewer) and its vanilla per-class one (nonzero class_agnostic always runs plain nms instead, regardless of count[1], as YOLOX does).

Parameters:

Name	Type	Description	Default
width	int	The widest pyramid level's cells; a row's compile size.	52
classes	int	The class scores of a cell.	80
cap	int	The candidates one core holds. 320 leaves headroom under the ~373 the default width and classes allow in 64 KiB of core memory (measured; see the commit for the arithmetic).	320

Returns:

Type	Description
_NmsKernel	ExternalFunction for nms_score, with nms_setup and
_NmsKernel	nms_finish bound on its object.

Raises:

Type	Description
ValueError	When width, classes or cap is below 1, or cap is above 1000, torchvision's own cutover to its vanilla path.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def nms(width: int = 52, classes: int = 80, cap: int = 320) -> _NmsKernel:
    """Score-threshold and class-aware NMS over an anchor-free detector's boxes.

    Bit for bit torchvision's CPU ``nms`` under ``batched_nms``'s own
    choice of path (see ``nms_ref``, the full arithmetic model). Three
    entry points share one core's state, and a core runs, once per image:

    ```text
    nms_setup(conf_thre)
    for each row of each pyramid level: nms_score(row, width)
    nms_finish(out, count, nms_thre, class_agnostic)
    ```

    ``row`` is [`grid_decode`][aie.iron.kernels.vision.grid_decode]'s output
    row: ``width`` cells of ``5 + classes`` float32, a cell's box ``(cx,
    cy, w, h)``, its objectness, then ``classes`` scores; ``width`` is a
    runtime count up to this build's ``width``, so one build serves every
    pyramid level. A cell survives at ``objectness * its best class's
    score``, ties to the lowest class id. Past ``cap`` survivors, a new
    one only replaces the held one of the lowest score (ties to the
    latest arrival), so a core always holds the top ``cap`` survivors by
    score and ``nms_finish``'s output equals torchvision's unbounded one
    for every detection ranked within that top ``cap``, at any survivor
    count. ``out`` is ``cap`` rows of ``x1, y1, x2, y2, score, class``
    (class as float32), zeros past ``count[0]``'s kept detections, in
    descending score order; ``count[1]`` is every survivor seen, cap or
    not, and also what picks ``class_agnostic`` 0's path between
    ``batched_nms``'s coordinate-offset trick (1000 survivors or fewer)
    and its vanilla per-class one (nonzero ``class_agnostic`` always runs
    plain ``nms`` instead, regardless of ``count[1]``, as YOLOX does).

    Args:
        width: The widest pyramid level's cells; a row's compile size.
        classes: The class scores of a cell.
        cap: The candidates one core holds. 320 leaves headroom under the
            ~373 the default `width` and `classes` allow in 64 KiB of
            core memory (measured; see the commit for the arithmetic).

    Returns:
        ExternalFunction for ``nms_score``, with ``nms_setup`` and
        ``nms_finish`` bound on its object.

    Raises:
        ValueError: When ``width``, ``classes`` or ``cap`` is below 1, or
            ``cap`` is above 1000, torchvision's own cutover to its
            vanilla path.
    """
    if min(width, classes, cap) < 1:
        raise ValueError(
            "nms() needs width, classes and cap >= 1, got "
            f"width={width}, classes={classes}, cap={cap}."
        )
    if cap > 1000:
        raise ValueError(f"nms() needs cap <= 1000, got {cap}.")
    row_ty = np.ndarray[(width * (5 + classes),), np.dtype[np.float32]]
    out_ty = np.ndarray[(cap * 6,), np.dtype[np.float32]]
    count_ty = np.ndarray[(2,), np.dtype[np.int32]]
    fn = _make_extern(
        "nms_score",
        _kernel_source("vision/nms.cc"),
        [row_ty, np.int32],
        compile_flags=_STRICT_FP
        + [
            f"-DNMS_WIDTH={width}",
            f"-DNMS_CLASSES={classes}",
            f"-DNMS_CAP={cap}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Param),
            unsupported=(
                "three entry points share one core's state, driven by a "
                "running candidate count nms_finish reads; test_nms_e2e.py "
                "drives and judges them"
            ),
            ops_per_call=0,
            even_calls=False,
        ),
        cls=_NmsKernel,
    )
    lib = fn.object_file
    fn.nms_setup = lib.bind("nms_setup", [np.float32])
    fn.nms_finish = lib.bind("nms_finish", [out_ty, count_ty, np.float32, np.int32])
    return fn

resize_linear

resize_linear(
    in_bytes: int = 3072, out_bytes: int = 1248
) -> _ResizeLinearKernel

Resize a uint8 RGB image as OpenCV's INTER_LINEAR does, a row at a time.

Bit for bit OpenCV's fixed-point path on the CPU (resizeGeneric_ with 11-bit weights, as cv2.setUseOptimized(False) runs it): each source row's horizontal pass exact in int32, each output row's vertical pass OpenCV's 8-bit one, ((b0 * (s0 >> 4)) >> 16) + ((b1 * (s1 >> 4)) >> 16) + 2 >> 2, the same on every element. Four entry points share one core's state, and a core runs:

resize_linear_setup(counts, rows, columns, out_rows, out_columns)
for each of counts[0] source rows:
    resize_linear_take(row, counts)
    for counts[1] output rows: resize_linear_emit(out)
resize_linear_finish(counts)
for counts[1] output rows: resize_linear_emit(out)

An output row is out_bytes, zeros past 3 * out_columns. Sizes the kernel refuses (below 1, or past in_bytes or out_bytes) take every source row and emit out_rows rows of zeros, so a design's streams stay in step; every call emits out_rows rows. resize_linear_ref is the arithmetic.

Parameters:

Name	Type	Description	Default
in_bytes	int	A source row's bytes; a multiple of 4 up to 32764.	3072
out_bytes	int	An output row's bytes; a multiple of 4 holding whole pixels' worth up to 32764.	1248

Returns:

Type	Description
_ResizeLinearKernel	ExternalFunction for resize_linear_take, with
_ResizeLinearKernel	resize_linear_setup, resize_linear_emit and
_ResizeLinearKernel	resize_linear_finish bound on its object.

Raises:

Type	Description
ValueError	When either row's bytes are not a multiple of 4 from 4 to 32764.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def resize_linear(in_bytes: int = 3072, out_bytes: int = 1248) -> _ResizeLinearKernel:
    """Resize a uint8 RGB image as OpenCV's INTER_LINEAR does, a row at a time.

    Bit for bit OpenCV's fixed-point path on the CPU (resizeGeneric_ with
    11-bit weights, as ``cv2.setUseOptimized(False)`` runs it): each source
    row's horizontal pass exact in int32, each output row's vertical pass
    OpenCV's 8-bit one, ``((b0 * (s0 >> 4)) >> 16) + ((b1 * (s1 >> 4)) >>
    16) + 2 >> 2``, the same on every element. Four entry points
    share one core's state, and a core runs:

    ```text
    resize_linear_setup(counts, rows, columns, out_rows, out_columns)
    for each of counts[0] source rows:
        resize_linear_take(row, counts)
        for counts[1] output rows: resize_linear_emit(out)
    resize_linear_finish(counts)
    for counts[1] output rows: resize_linear_emit(out)
    ```

    An output row is ``out_bytes``, zeros past ``3 * out_columns``. Sizes
    the kernel refuses (below 1, or past ``in_bytes`` or ``out_bytes``)
    take every source row and emit ``out_rows`` rows of zeros, so a
    design's streams stay in step; every call emits ``out_rows`` rows.
    [`resize_linear_ref`][aie.iron.kernels.vision.resize_linear_ref] is the
    arithmetic.

    Args:
        in_bytes: A source row's bytes; a multiple of 4 up to 32764.
        out_bytes: An output row's bytes; a multiple of 4 holding whole
            pixels' worth up to 32764.

    Returns:
        ExternalFunction for ``resize_linear_take``, with
        ``resize_linear_setup``, ``resize_linear_emit`` and
        ``resize_linear_finish`` bound on its object.

    Raises:
        ValueError: When either row's bytes are not a multiple of 4 from 4
            to 32764.
    """
    for name, n in (("in_bytes", in_bytes), ("out_bytes", out_bytes)):
        if n % 4 or not 4 <= n <= 32764:
            raise ValueError(
                f"resize_linear() needs {name} a multiple of 4 from 4 to 32764, got {n}."
            )
    counts_ty = np.ndarray[(2,), np.dtype[np.int32]]
    fn = _make_extern(
        "resize_linear_take",
        _kernel_source("vision/resize_linear.cc"),
        [np.ndarray[(in_bytes,), np.dtype[np.uint8]], counts_ty],
        compile_flags=[f"-DIN_BYTES={in_bytes}", f"-DOUT_BYTES={out_bytes}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            unsupported=(
                "four entry points share one core's state, driven by counts "
                "they write; test_resize_linear_e2e.py drives and judges them"
            ),
            ops_per_call=0,
        ),
        cls=_ResizeLinearKernel,
    )
    lib = fn.object_file
    fn.resize_linear_setup = lib.bind(
        "resize_linear_setup", [counts_ty, np.int32, np.int32, np.int32, np.int32]
    )
    fn.resize_linear_emit = lib.bind(
        "resize_linear_emit", [np.ndarray[(out_bytes,), np.dtype[np.uint8]]]
    )
    fn.resize_linear_finish = lib.bind("resize_linear_finish", [counts_ty])
    return fn

resample_peak

resample_peak(
    words: int = 256,
    *,
    cores: int = 16,
    slots: int | None = None
) -> _ResamplePeakKernel

Return the largest normalized weight of one core's share of a resample table.

One axis of torch's antialiased bicubic resize (in samples to out) is a table of words-int32 chunks, cores cores taking chunks core, core + cores, ...; see resample_quantize. resample_peak(peak, in, out, core, chunks) writes the largest normalized weight over the outputs of the core's chunks chunks as float64 bits in peak (two int32), 0 when the table cannot hold the sizes. fn.resample_join(a, b, out) writes the larger of two peaks, so a chain of cores reduces them to the table's precision. The weights are float64 without contraction, as torch computes them on the CPU; amd/IRON's ResampleTaps runs both for EmbeddingGemma 2's image processor.

Parameters:

Name	Type	Description	Default
words	int	The int32 words of a chunk.	256
cores	int	The cores the chunks are dealt to.	16
slots	int | None	The outputs a chunk is fixed to, or None for as many as the window lets fit.	None

Returns:

Type	Description
_ResamplePeakKernel	ExternalFunction for resample_peak, with resample_join bound
_ResamplePeakKernel	on its object.

Raises:

Type	Description
ValueError	When words holds no slot, cores is below 1, or slots is below 1.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def resample_peak(
    words: int = 256, *, cores: int = 16, slots: int | None = None
) -> _ResamplePeakKernel:
    """Return the largest normalized weight of one core's share of a resample table.

    One axis of torch's antialiased bicubic resize (``in`` samples to
    ``out``) is a table of ``words``-int32 chunks, ``cores`` cores taking
    chunks ``core, core + cores, ...``; see
    [`resample_quantize`][aie.iron.kernels.vision.resample_quantize].
    ``resample_peak(peak, in, out, core, chunks)`` writes the largest
    normalized weight over the outputs of the core's ``chunks`` chunks as
    float64 bits in ``peak`` (two int32), 0 when the table cannot hold the
    sizes. ``fn.resample_join(a, b, out)`` writes the larger of two peaks,
    so a chain of cores reduces them to the table's precision. The weights
    are float64 without contraction, as torch computes them on the CPU;
    amd/IRON's ``ResampleTaps`` runs both for EmbeddingGemma 2's image
    processor.

    Args:
        words: The int32 words of a chunk.
        cores: The cores the chunks are dealt to.
        slots: The outputs a chunk is fixed to, or ``None`` for as many as
            the window lets fit.

    Returns:
        ExternalFunction for ``resample_peak``, with ``resample_join`` bound
        on its object.

    Raises:
        ValueError: When ``words`` holds no slot, ``cores`` is below 1, or
            ``slots`` is below 1.
    """
    _check_resample("resample_peak", words, cores, slots)
    peak_ty = np.ndarray[(2,), np.dtype[np.int32]]
    fn = _make_extern(
        "resample_peak",
        _kernel_source("vision/resample_peak.cc"),
        [peak_ty] + [np.int32] * 4,
        compile_flags=[f"-DWORDS={words}", f"-DCORES={cores}", f"-DPER={slots or 0}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, Param, Param, Param, Param),
            reference=lambda in_size, out_size, core, chunks: resample_peak_ref(
                in_size, out_size, core, chunks, words=words, cores=cores, slots=slots
            ),
            tolerance=Tolerance.exact(note="IEEE float64, as torch on the CPU"),
            ops_per_call=0,
        ),
        cls=_ResamplePeakKernel,
    )
    fn.resample_join = fn.object_file.bind("resample_join", [peak_ty, peak_ty, peak_ty])
    return fn

resample_quantize

resample_quantize(
    words: int = 256,
    *,
    cores: int = 16,
    slots: int | None = None
) -> ExternalFunction

One chunk of the int16 filter table torch resamples an axis with.

resample_quantize(peak, chunk, in, out, core, k, chunks) writes chunk core + k * cores of the table of in samples to out whose cores cores each write chunks chunks. A chunk is words int32: a header [precision, window, first, outputs], then outputs slots from output first on, each [start, count] and the window's int16 weights two to a word, the rest zero. The precision is the one the largest normalized weight peak (float64 bits, from resample_peak) sets, as torch picks it. A size the table cannot hold (a size below 1, a window no slot of words holds, too few chunks) writes [-1, window, 0, 0].

Parameters:

Name	Type	Description	Default
words	int	The int32 words of a chunk.	256
cores	int	The cores the chunks are dealt to.	16
slots	int | None	The outputs a chunk is fixed to, or None for as many as the window lets fit.	None

Returns:

Type	Description
ExternalFunction	ExternalFunction for resample_quantize.

Raises:

Type	Description
ValueError	When words holds no slot, cores is below 1, or slots is below 1.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def resample_quantize(
    words: int = 256, *, cores: int = 16, slots: int | None = None
) -> ExternalFunction:
    """One chunk of the int16 filter table torch resamples an axis with.

    ``resample_quantize(peak, chunk, in, out, core, k, chunks)`` writes
    chunk ``core + k * cores`` of the table of ``in`` samples to ``out``
    whose ``cores`` cores each write ``chunks`` chunks. A chunk is ``words``
    int32: a header ``[precision, window, first, outputs]``, then
    ``outputs`` slots from output ``first`` on, each ``[start, count]`` and
    the window's int16 weights two to a word, the rest zero. The precision
    is the one the largest normalized weight ``peak`` (float64 bits, from
    [`resample_peak`][aie.iron.kernels.vision.resample_peak]) sets, as torch
    picks it. A size the table cannot hold (a size below 1, a window no
    slot of ``words`` holds, too few chunks) writes ``[-1, window, 0, 0]``.

    Args:
        words: The int32 words of a chunk.
        cores: The cores the chunks are dealt to.
        slots: The outputs a chunk is fixed to, or ``None`` for as many as
            the window lets fit.

    Returns:
        ExternalFunction for ``resample_quantize``.

    Raises:
        ValueError: When ``words`` holds no slot, ``cores`` is below 1, or
            ``slots`` is below 1.
    """
    _check_resample("resample_quantize", words, cores, slots)
    return _make_extern(
        "resample_quantize",
        _kernel_source("vision/resample_quantize.cc"),
        [
            np.ndarray[(2,), np.dtype[np.int32]],
            np.ndarray[(words,), np.dtype[np.int32]],
        ]
        + [np.int32] * 5,
        compile_flags=[f"-DWORDS={words}", f"-DCORES={cores}", f"-DPER={slots or 0}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param, Param, Param, Param, Param),
            parameter_bindings=((5, 0),),
            call_index=5,
            even_calls=False,
            sample=_resample_peaks,
            reference=lambda peak, in_size, out_size, core, chunks: (
                resample_quantize_ref(
                    peak,
                    in_size,
                    out_size,
                    core,
                    chunks,
                    words=words,
                    cores=cores,
                    slots=slots,
                )
            ),
            tolerance=Tolerance.exact(note="IEEE float64, as torch on the CPU"),
            ops_per_call=0,
        ),
    )

resize

resize(
    words: int = 356,
    *,
    chunk: int = 4096,
    patch_columns: int = 8,
    cores: int = 16
) -> _ResizeKernel

Resize a uint8 RGB image as torch's antialiased bicubic resize, as patches.

The image is resized across, then down, each pass rounded to uint8 as torch does it on the CPU, and written as 16x16 patches of u8 / 255 in bf16, channels last. The filters are resample_quantize tables at 16 slots, one chunk a patch column (width) or a band of 16 rows (height). cores cores receive every chunk-byte image chunk, each image row padded to whole chunks; core c owns patch columns c, c + cores, ..., at most patch_columns of them. Six entry points share one core's state, and a core runs:

resize_setup(counts, rows, columns, out_rows, out_columns, c)
for k in range(counts[0]): resize_take(width chunk k, k)
for each of counts[1] bands:
    resize_band(height chunk, counts)
    for counts[2] image chunks: resize_consume(image chunk, height chunk)
    for i in range(counts[3]): resize_emit(patch, i)
resize_finish(counts)
then counts[4] image chunks, unread

counts is five int32. Band b's patch i is patch column c + i * cores, zeros past the image: a core writes out_columns // 16 // cores + 1 patches a band, so the last patch column of every band is zeros. The counts are the same on every core whatever the input, so the broadcast streams stay in step; sizes or a table the kernel refuses write zeros from there on, where resize_ref says. amd/IRON's Resample runs it for EmbeddingGemma 2's image processor.

Parameters:

Name	Type	Description	Default
words	int	The int32 words of a table chunk; a 16-slot chunk holds a window of 2 * ((words - 4) // 16 - 2) - 1 taps.	356
chunk	int	The bytes of an image chunk.	4096
patch_columns	int	The most patch columns one core holds.	8
cores	int	The cores the patch columns are dealt to.	16

Returns:

Name	Type	Description
	_ResizeKernel	ExternalFunction for resize_consume(image_chunk, height_chunk),
	_ResizeKernel	with resize_setup, resize_take, resize_band,
	_ResizeKernel	resize_emit, resize_emit_row and resize_finish bound on
	_ResizeKernel	its object. resize_emit_row(out, i) writes the same band and
	_ResizeKernel	patch column as resize_emit(out, i) -- a 16-row, 16-pixel-wide
	_ResizeKernel	block, not a slice of a full-width raster row -- as 16 rows of
	_ResizeKernel	SIDE interleaved uint8 RGB pixels (row 0's 3 * SIDE bytes,
	_ResizeKernel	then row 1's, ...) rather than a bf16 u8 / 255 patch, for
	_ResizeKernel	vision.levels to read. Because core
	_ResizeKernel	c owns patch columns c, c + cores, ... (strided, not
	_ResizeKernel	contiguous), a true output row needs every core's blocks for that
band	_ResizeKernel	row y's pixels at columns (core + i * cores) * SIDE
	_ResizeKernel	on are call i's byte range 3 * SIDE * y on. Every block is
	_ResizeKernel	whole, 16 by 16, even past a target that is not a multiple of 16,
	_ResizeKernel	and each core emits strips // cores + 1 of them a band, zeros
	_ResizeKernel	past the real strips. So the drain's destination is a padded
	_ResizeKernel	staging raster of bands * SIDE rows and ``pad = (strips //
	_ResizeKernel	cores + 1) * cores`` blocks across: a per-core strided drain --
	_ResizeKernel	call i's 16 chunks of 3 * SIDE bytes to staging row ``band
	_ResizeKernel	SIDE + jat column byte offset(core + i * cores) * 3 *
	_ResizeKernel	SIDE, stride3 * SIDE * pad`` between rows -- fills it the way
	_ResizeKernel	the patches path gathers patches from every core's stream. The
	_ResizeKernel	resized image is that raster's first out_rows rows and
	_ResizeKernel	out_columns pixels, and a crop of it (levels's first /
	_ResizeKernel	count / at on columns, row / lo / hi on rows)
	_ResizeKernel	is free-form, not limited to this kernel's SIDE- or ``cores *
	_ResizeKernel	SIDE``-pixel block boundaries.

Raises:

Type	Description
ValueError	When words holds no window of 1 to 64 taps, chunk is not a multiple of 4 from 4 to 98300, or patch_columns or cores is below 1.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def resize(
    words: int = 356, *, chunk: int = 4096, patch_columns: int = 8, cores: int = 16
) -> _ResizeKernel:
    """Resize a uint8 RGB image as torch's antialiased bicubic resize, as patches.

    The image is resized across, then down, each pass rounded to uint8 as
    torch does it on the CPU, and written as 16x16 patches of ``u8 / 255``
    in bf16, channels last. The filters are
    [`resample_quantize`][aie.iron.kernels.vision.resample_quantize] tables at
    16 slots, one chunk a patch column (width) or a band of 16 rows
    (height). ``cores`` cores receive every ``chunk``-byte image chunk, each
    image row padded to whole chunks; core ``c`` owns patch columns ``c, c +
    cores, ...``, at most ``patch_columns`` of them. Six entry points share
    one core's state, and a core runs:

    ```text
    resize_setup(counts, rows, columns, out_rows, out_columns, c)
    for k in range(counts[0]): resize_take(width chunk k, k)
    for each of counts[1] bands:
        resize_band(height chunk, counts)
        for counts[2] image chunks: resize_consume(image chunk, height chunk)
        for i in range(counts[3]): resize_emit(patch, i)
    resize_finish(counts)
    then counts[4] image chunks, unread
    ```

    ``counts`` is five int32. Band ``b``'s patch ``i`` is patch column ``c
    + i * cores``, zeros past the image: a core writes ``out_columns // 16
    // cores + 1`` patches a band, so the last patch column of every band is
    zeros. The counts are the same on every core whatever the input, so the
    broadcast streams stay in step; sizes or a table the kernel refuses
    write zeros from there on, where
    [`resize_ref`][aie.iron.kernels.vision.resize_ref] says. amd/IRON's
    ``Resample`` runs it for EmbeddingGemma 2's image processor.

    Args:
        words: The int32 words of a table chunk; a 16-slot chunk holds a
            window of ``2 * ((words - 4) // 16 - 2) - 1`` taps.
        chunk: The bytes of an image chunk.
        patch_columns: The most patch columns one core holds.
        cores: The cores the patch columns are dealt to.

    Returns:
        ExternalFunction for ``resize_consume(image_chunk, height_chunk)``,
        with ``resize_setup``, ``resize_take``, ``resize_band``,
        ``resize_emit``, ``resize_emit_row`` and ``resize_finish`` bound on
        its object. ``resize_emit_row(out, i)`` writes the same band and
        patch column as ``resize_emit(out, i)`` -- a 16-row, 16-pixel-wide
        block, not a slice of a full-width raster row -- as 16 rows of
        ``SIDE`` interleaved uint8 RGB pixels (row 0's ``3 * SIDE`` bytes,
        then row 1's, ...) rather than a bf16 u8 / 255 patch, for
        [`vision.levels`][aie.iron.kernels.vision.levels] to read. Because core
        ``c`` owns patch columns ``c, c + cores, ...`` (strided, not
        contiguous), a true output row needs every core's blocks for that
        band: row ``y``'s pixels at columns ``(core + i * cores) * SIDE``
        on are call ``i``'s byte range ``3 * SIDE * y`` on. Every block is
        whole, 16 by 16, even past a target that is not a multiple of 16,
        and each core emits ``strips // cores + 1`` of them a band, zeros
        past the real strips. So the drain's destination is a padded
        staging raster of ``bands * SIDE`` rows and ``pad = (strips //
        cores + 1) * cores`` blocks across: a per-core strided drain --
        call ``i``'s 16 chunks of ``3 * SIDE`` bytes to staging row ``band
        * SIDE + j`` at column byte offset ``(core + i * cores) * 3 *
        SIDE``, stride ``3 * SIDE * pad`` between rows -- fills it the way
        the patches path gathers ``patches`` from every core's stream. The
        resized image is that raster's first ``out_rows`` rows and
        ``out_columns`` pixels, and a crop of it (``levels``'s ``first`` /
        ``count`` / ``at`` on columns, ``row`` / ``lo`` / ``hi`` on rows)
        is free-form, not limited to this kernel's ``SIDE``- or ``cores *
        SIDE``-pixel block boundaries.

    Raises:
        ValueError: When ``words`` holds no window of 1 to 64 taps,
            ``chunk`` is not a multiple of 4 from 4 to 98300, or
            ``patch_columns`` or ``cores`` is below 1.
    """
    window = 2 * ((words - _RESAMPLE_HEADER) // _RESIZE_SIDE - 2) - 1
    if not 1 <= window <= 64:
        raise ValueError(
            f"resize() needs words from 52 to 563 (a window of 1 to 64 taps), "
            f"got words={words}."
        )
    if chunk % 4 or not 4 <= chunk < 98304:
        raise ValueError(
            f"resize() needs chunk a multiple of 4 from 4 to 98300, got chunk={chunk}."
        )
    if patch_columns < 1 or cores < 1:
        raise ValueError(
            "resize() needs patch_columns >= 1 and cores >= 1, got "
            f"patch_columns={patch_columns}, cores={cores}."
        )
    counts_ty = np.ndarray[(5,), np.dtype[np.int32]]
    taps_ty = np.ndarray[(words,), np.dtype[np.int32]]
    fn = _make_extern(
        "resize_consume",
        _kernel_source("vision/resize.cc"),
        [np.ndarray[(chunk,), np.dtype[np.uint8]], taps_ty],
        compile_flags=[
            f"-DWORDS={words}",
            f"-DCHUNK={chunk}",
            f"-DCOLS={patch_columns}",
            f"-DCORES={cores}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In),
            unsupported=(
                "six entry points share one core's state, driven by counts "
                "they write; test_resize_e2e.py drives and judges them"
            ),
            ops_per_call=0,
        ),
        cls=_ResizeKernel,
    )
    lib = fn.object_file
    fn.resize_setup = lib.bind(
        "resize_setup",
        [counts_ty, np.int32, np.int32, np.int32, np.int32, np.int32],
    )
    fn.resize_take = lib.bind("resize_take", [taps_ty, np.int32])
    fn.resize_band = lib.bind("resize_band", [taps_ty, counts_ty])
    fn.resize_emit = lib.bind(
        "resize_emit",
        [np.ndarray[(_RESIZE_SIDE * _RESIZE_SIDE * 3,), np.dtype[bfloat16]], np.int32],
    )
    fn.resize_emit_row = lib.bind(
        "resize_emit_row",
        [np.ndarray[(_RESIZE_SIDE * _RESIZE_SIDE * 3,), np.dtype[np.uint8]], np.int32],
    )
    fn.resize_finish = lib.bind("resize_finish", [counts_ty])
    return fn

gray2rgba_ref

gray2rgba_ref(y)

Numpy reference for gray2rgba: (y, y, y, 255) per pixel.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def gray2rgba_ref(y):
    """Numpy reference for [`gray2rgba`][aie.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 utils/pyright_stubs/aie/iron/kernels/vision.py

def rgba2gray_ref(rgba):
    """Numpy reference for [`rgba2gray`][aie.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 utils/pyright_stubs/aie/iron/kernels/vision.py

def rgba2hue_ref(rgba):
    """Numpy reference for [`rgba2hue`][aie.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)

grid_decode_ref

grid_decode_ref(
    cells,
    y,
    stride,
    *,
    width: int,
    channels: int,
    obj_at: int,
    cls_at: int,
    classes: int
)

Numpy reference for grid_decode.

YOLOX's decode_outputs in float32.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def grid_decode_ref(
    cells,
    y,
    stride,
    *,
    width: int,
    channels: int,
    obj_at: int,
    cls_at: int,
    classes: int,
):
    """Numpy reference for [`grid_decode`][aie.iron.kernels.vision.grid_decode].

    YOLOX's ``decode_outputs`` in float32.
    """
    p = np.asarray(cells).astype(np.float32).reshape(-1, width, channels)
    s = np.float32(stride)
    out = np.empty((p.shape[0], width, 5 + classes), np.float32)
    out[..., 0] = (p[..., 0] + np.arange(width, dtype=np.float32)) * s
    out[..., 1] = (p[..., 1] + np.float32(y)) * s
    with np.errstate(over="ignore"):
        out[..., 2:4] = np.exp(p[..., 2:4]) * s
    out[..., 4] = p[..., obj_at]
    out[..., 5:] = p[..., cls_at : cls_at + classes]
    return out.reshape(p.shape[0], width * (5 + classes))

nms_ref

nms_ref(
    rows,
    conf_thre,
    nms_thre,
    class_agnostic,
    *,
    classes: int,
    cap: int
) -> tuple[ndarray, int, int]

Numpy reference for nms.

torchvision's CPU nms under batched_nms's own choice of path: the coordinate-offset trick at 1000 survivors or fewer, its vanilla per-class one past it, as YOLOX's postprocess calls batched_nms (class_agnostic calls plain nms instead, regardless of survivor count).

rows is (row, width) pairs, each row width cells of 5 + classes float32 as grid_decode writes them, in call order. A cell survives at conf_thre or above (objectness * its best class's confidence, ties to the lowest class id); past cap survivors, only the top cap by score survive (ties to the earliest of them). Suppression then visits them descending by score (ties to the lowest index), each unsuppressed one kept and checked against every later one it can suppress, strict ovr > nms_thre. The coordinate-offset trick shifts every box by its class times one past every survivor's largest coordinate (boxes.max() over all of them, not only the top cap) before this, so cross-class pairs never overlap; the vanilla path runs no shift and only ever compares same-class pairs instead; class_agnostic runs neither. Kept detections (the original, unshifted boxes) are written in score order.

Returns:

Name	Type	Description
detections	ndarray	cap * 6 float32 (x1, y1, x2, y2, score, class), zeros past count.
count	int	How many detections were kept.
survivors	int	How many candidates reached conf_thre, cap or not.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def nms_ref(
    rows, conf_thre, nms_thre, class_agnostic, *, classes: int, cap: int
) -> tuple[np.ndarray, int, int]:
    """Numpy reference for [`nms`][aie.iron.kernels.vision.nms].

    torchvision's CPU ``nms`` under ``batched_nms``'s own choice of path:
    the coordinate-offset trick at 1000 survivors or fewer, its vanilla
    per-class one past it, as YOLOX's ``postprocess`` calls ``batched_nms``
    (``class_agnostic`` calls plain ``nms`` instead, regardless of
    survivor count).

    ``rows`` is ``(row, width)`` pairs, each row ``width`` cells of ``5 +
    classes`` float32 as [`grid_decode`][aie.iron.kernels.vision.grid_decode]
    writes them, in call order. A cell survives at ``conf_thre`` or above
    (``objectness * its best class's confidence``, ties to the lowest
    class id); past ``cap`` survivors, only the top ``cap`` by score
    survive (ties to the earliest of them). Suppression then visits them
    descending by score (ties to the lowest index), each unsuppressed one
    kept and checked against every later one it can suppress, strict
    ``ovr > nms_thre``. The coordinate-offset trick shifts every box by
    its class times one past every survivor's largest coordinate
    (``boxes.max()`` over all of them, not only the top ``cap``) before
    this, so cross-class pairs never overlap; the vanilla path runs no
    shift and only ever compares same-class pairs instead;
    ``class_agnostic`` runs neither. Kept detections (the original,
    unshifted boxes) are written in score order.

    Returns:
        detections: ``cap * 6`` float32 (``x1, y1, x2, y2, score, class``),
            zeros past ``count``.
        count: How many detections were kept.
        survivors: How many candidates reached ``conf_thre``, cap or not.
    """
    x1s, y1s, x2s, y2s, scores, classes_ = [], [], [], [], [], []
    thre = np.float32(conf_thre)
    for row, row_width in rows:
        cells = np.asarray(row, np.float32).reshape(-1, 5 + classes)
        for i in range(int(row_width)):
            cx, cy, w, h, obj = (
                cells[i, 0],
                cells[i, 1],
                cells[i, 2],
                cells[i, 3],
                cells[i, 4],
            )
            class_scores = cells[i, 5:]
            best_c = int(np.argmax(class_scores))
            score = obj * class_scores[best_c]
            if score >= thre:
                x1s.append(cx - w / np.float32(2))
                y1s.append(cy - h / np.float32(2))
                x2s.append(cx + w / np.float32(2))
                y2s.append(cy + h / np.float32(2))
                scores.append(score)
                classes_.append(np.float32(best_c))

    out = np.zeros(cap * 6, np.float32)
    survivors = len(scores)
    if survivors == 0:
        return out, 0, survivors
    x1 = np.array(x1s, np.float32)
    y1 = np.array(y1s, np.float32)
    x2 = np.array(x2s, np.float32)
    y2 = np.array(y2s, np.float32)
    score_arr = np.array(scores, np.float32)
    class_arr = np.array(classes_, np.float32)

    # torchvision's boxes.max(): over every survivor, before the top-cap
    # selection below picks which of them this kernel ever holds.
    max_coord = max(x1.max(), y1.max(), x2.max(), y2.max())
    vanilla = not class_agnostic and survivors > 1000
    offset_trick = not class_agnostic and not vanilla

    # Every survivor, stably by descending score; the first `cap` of them
    # are what nms_score's eviction keeps, in the same order.
    rank = np.argsort(-score_arr, kind="stable")[:cap]
    x1, y1, x2, y2 = x1[rank], y1[rank], x2[rank], y2[rank]
    score_arr, class_arr = score_arr[rank], class_arr[rank]
    n = len(rank)

    if offset_trick:
        offset = class_arr * (max_coord + np.float32(1))
        bx1, by1, bx2, by2 = x1 + offset, y1 + offset, x2 + offset, y2 + offset
    else:
        bx1, by1, bx2, by2 = x1, y1, x2, y2
    area = (bx2 - bx1) * (by2 - by1)

    # rank is already stable descending, so the suppression order is i, j itself.
    suppressed = np.zeros(n, dtype=bool)
    kept = []
    with np.errstate(invalid="ignore", divide="ignore"):
        for i in range(n):
            if suppressed[i]:
                continue
            kept.append(i)
            for j in range(i + 1, n):
                if suppressed[j]:
                    continue
                if vanilla and class_arr[j] != class_arr[i]:
                    continue
                xx1, yy1 = max(bx1[i], bx1[j]), max(by1[i], by1[j])
                xx2, yy2 = min(bx2[i], bx2[j]), min(by2[i], by2[j])
                w = max(np.float32(0), xx2 - xx1)
                h = max(np.float32(0), yy2 - yy1)
                inter = w * h
                ovr = inter / (area[i] + area[j] - inter)
                if ovr > nms_thre:
                    suppressed[j] = True

    for k, i in enumerate(kept):
        out[k * 6 : k * 6 + 6] = (
            x1[i],
            y1[i],
            x2[i],
            y2[i],
            score_arr[i],
            class_arr[i],
        )
    return out, len(kept), survivors

resize_linear_ref

resize_linear_ref(
    rows,
    columns,
    out_rows,
    out_columns,
    image,
    *,
    out_bytes: int
)

Numpy reference for resize_linear.

image holds the source rows, 3 * columns bytes of each used; returns out_rows rows of out_bytes. Exact.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def resize_linear_ref(rows, columns, out_rows, out_columns, image, *, out_bytes: int):
    """Numpy reference for [`resize_linear`][aie.iron.kernels.vision.resize_linear].

    ``image`` holds the source rows, ``3 * columns`` bytes of each used;
    returns ``out_rows`` rows of ``out_bytes``. Exact.
    """
    image = np.asarray(image, np.uint8)
    h, w, oh, ow = int(rows), int(columns), int(out_rows), int(out_columns)
    out = np.zeros((max(oh, 0), out_bytes), np.uint8)
    if min(h, w, oh, ow) < 1 or 3 * w > image.shape[1] or 3 * ow > out_bytes:
        return out
    coef = np.float32(2048)

    def axis(n_in, n_out):
        scale = 1.0 / (n_out / n_in)
        f = ((np.arange(n_out) + 0.5) * scale - 0.5).astype(np.float32)
        s = np.floor(f).astype(np.int64)
        return s, (f - s.astype(np.float32)).astype(np.float32)

    sx, fx = axis(w, ow)
    under = sx < 0
    fx[under], sx[under] = 0, 0
    past = sx + 1 >= w
    xmax = int(np.argmax(past)) if past.any() else ow
    over = sx >= w - 1
    fx[over], sx[over] = 0, w - 1
    a0 = np.rint((np.float32(1) - fx) * coef).astype(np.int64)
    a1 = np.rint(fx * coef).astype(np.int64)
    src = image[:h, : 3 * w].reshape(h, w, 3).astype(np.int64)
    nxt = np.minimum(sx + 1, w - 1)
    inner = (np.arange(ow) < xmax)[None, :, None]
    hor = np.where(
        inner,
        src[:, sx] * a0[None, :, None] + src[:, nxt] * a1[None, :, None],
        src[:, sx] * 2048,
    ).reshape(h, 3 * ow)
    sy, fy = axis(h, oh)
    b0 = np.rint((np.float32(1) - fy) * coef).astype(np.int64)[:, None]
    b1 = np.rint(fy * coef).astype(np.int64)[:, None]
    s0, s1 = hor[np.clip(sy, 0, h - 1)], hor[np.clip(sy + 1, 0, h - 1)]
    resized = (((b0 * (s0 >> 4)) >> 16) + ((b1 * (s1 >> 4)) >> 16) + 2) >> 2
    out[:, : 3 * ow] = resized
    return out

levels_ref

levels_ref(
    rows,
    table,
    first,
    count,
    at,
    row,
    lo,
    hi,
    *,
    out_width: int,
    bgr: bool = False
)

Numpy reference for levels. Exact.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def levels_ref(
    rows, table, first, count, at, row, lo, hi, *, out_width: int, bgr: bool = False
):
    """Numpy reference for [`levels`][aie.iron.kernels.vision.levels]. Exact."""
    rows = np.asarray(rows, np.uint8)
    table = np.asarray(table).astype(bfloat16)
    first, count, at = int(first), int(count), int(at)
    out = np.empty((rows.shape[0], out_width, 4), bfloat16)
    out[:] = table[3 * 256 : 3 * 256 + 4]
    if first < 0 or at < 0 or not int(lo) <= int(row) < int(hi):
        count = 0
    count = max(0, min(count, out_width - at, rows.shape[1] // 3 - first))
    pixels = rows[:, 3 * first : 3 * (first + count)].reshape(rows.shape[0], count, 3)
    if bgr:
        pixels = pixels[..., ::-1]
    for c in range(3):
        out[:, at : at + count, c] = table[c * 256 + pixels[..., c].astype(np.int64)]
    return out.reshape(rows.shape[0], 4 * out_width)

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 utils/pyright_stubs/aie/iron/kernels/vision.py

def threshold_ref(x, thresh, maxval, ttype):
    """Numpy reference for [`threshold`][aie.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 utils/pyright_stubs/aie/iron/kernels/vision.py

def bitwise_or_ref(a, b):
    """Numpy reference for [`bitwise_or`][aie.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 utils/pyright_stubs/aie/iron/kernels/vision.py

def bitwise_and_ref(a, b):
    """Numpy reference for [`bitwise_and`][aie.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 utils/pyright_stubs/aie/iron/kernels/vision.py

def add_weighted_ref(a, b, alpha, beta, gamma):
    """Numpy reference for [`add_weighted`][aie.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 utils/pyright_stubs/aie/iron/kernels/vision.py

def filter2d_ref(line0, line1, line2, kernel):
    """Numpy reference for [`filter2d`][aie.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)

resample_peak_ref

resample_peak_ref(
    in_size: int,
    out_size: int,
    core: int,
    chunks: int,
    *,
    words: int = 256,
    cores: int = 16,
    slots: int | None = None
) -> ndarray

Numpy reference for resample_peak.

Parameters:

Name	Type	Description	Default
in_size	int	The input samples on the axis.	required
out_size	int	The output samples on the axis.	required
core	int	The core whose chunks core, core + cores, ... are read.	required
chunks	int	The chunks each core writes.	required
words	int	The int32 words of a chunk.	256
cores	int	The cores the chunks are dealt to.	16
slots	int | None	The outputs a chunk is fixed to, or None.	None

Returns:

Type	Description
ndarray	The peak's float64 bits as two int32.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def resample_peak_ref(
    in_size: int,
    out_size: int,
    core: int,
    chunks: int,
    *,
    words: int = 256,
    cores: int = 16,
    slots: int | None = None,
) -> np.ndarray:
    """Numpy reference for [`resample_peak`][aie.iron.kernels.vision.resample_peak].

    Args:
        in_size: The input samples on the axis.
        out_size: The output samples on the axis.
        core: The core whose chunks ``core, core + cores, ...`` are read.
        chunks: The chunks each core writes.
        words: The int32 words of a chunk.
        cores: The cores the chunks are dealt to.
        slots: The outputs a chunk is fixed to, or ``None``.

    Returns:
        The peak's float64 bits as two int32.
    """
    axis = _ResampleAxis.of(in_size, out_size, chunks * cores, words=words, slots=slots)
    m = 0.0
    for k in range(chunks if axis.ok else 0):
        first = (core + k * cores) * axis.per
        for i in range(first, min(first + axis.per, out_size)):
            _, weights, total = axis.taps(i)
            if total == 0.0:
                continue
            e = weights[0]
            for v in weights[1:]:
                if (e < v) if total > 0.0 else (v < e):
                    e = v
            if m < e / total:
                m = e / total
    return np.array([m], np.float64).view(np.int32)

resample_quantize_ref

resample_quantize_ref(
    peak: ndarray,
    in_size: int,
    out_size: int,
    core: int,
    chunks: int,
    *,
    words: int = 256,
    cores: int = 16,
    slots: int | None = None
) -> ndarray

Numpy reference for resample_quantize.

Parameters:

Name	Type	Description	Default
peak	ndarray	(calls, 2) int32, each call's peak as float64 bits; call k writes chunk core + k * cores.	required
in_size	int	The input samples on the axis.	required
out_size	int	The output samples on the axis.	required
core	int	The core whose chunks are written.	required
chunks	int	The chunks each core writes.	required
words	int	The int32 words of a chunk.	256
cores	int	The cores the chunks are dealt to.	16
slots	int | None	The outputs a chunk is fixed to, or None.	None

Returns:

Type	Description
ndarray	(calls, words) int32, a chunk per call. A weight past int16
ndarray	wraps, as the kernel's store does.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def resample_quantize_ref(
    peak: np.ndarray,
    in_size: int,
    out_size: int,
    core: int,
    chunks: int,
    *,
    words: int = 256,
    cores: int = 16,
    slots: int | None = None,
) -> np.ndarray:
    """Numpy reference for [`resample_quantize`][aie.iron.kernels.vision.resample_quantize].

    Args:
        peak: ``(calls, 2)`` int32, each call's peak as float64 bits; call
            ``k`` writes chunk ``core + k * cores``.
        in_size: The input samples on the axis.
        out_size: The output samples on the axis.
        core: The core whose chunks are written.
        chunks: The chunks each core writes.
        words: The int32 words of a chunk.
        cores: The cores the chunks are dealt to.
        slots: The outputs a chunk is fixed to, or ``None``.

    Returns:
        ``(calls, words)`` int32, a chunk per call. A weight past int16
        wraps, as the kernel's store does.
    """
    peaks = np.ascontiguousarray(peak, np.int32).reshape(-1, 2).view(np.float64)
    axis = _ResampleAxis.of(in_size, out_size, chunks * cores, words=words, slots=slots)
    table = np.zeros((len(peaks), words), np.int32)
    for k, chunk in enumerate(table):
        if not axis.ok:
            chunk[:2] = -1, axis.window
            continue
        wt_max = float(peaks[k, 0])
        p = next(
            (p for p in range(22) if int(0.5 + wt_max * (1 << (p + 1))) >= 1 << 15),
            22,
        )
        first = (core + k * cores) * axis.per
        n = min(max(out_size - first, 0), axis.per)
        chunk[:_RESAMPLE_HEADER] = p, axis.window, first, n
        for s in range(n):
            xmin, weights, total = axis.taps(first + s)
            at = _RESAMPLE_HEADER + s * axis.slot
            chunk[at : at + 2] = xmin, len(weights)
            q = np.zeros(2 * (axis.slot - 2), np.int64)
            for j, v in enumerate(weights):
                v = (v / total if total != 0.0 else v) * (1 << p)
                q[j] = int(-0.5 + v) if v < 0 else int(0.5 + v)
            chunk[at + 2 : at + axis.slot] = q.astype(np.int16).view(np.int32)
    return table

resize_ref

resize_ref(
    image: ndarray,
    taps_w: ndarray,
    taps_h: ndarray,
    rows: int,
    columns: int,
    out_rows: int,
    out_columns: int,
    *,
    words: int = 356,
    chunk: int = 4096,
    patch_columns: int = 8,
    cores: int = 16,
    as_rows: bool = False
) -> ndarray

Numpy reference for resize.

Parameters:

Name	Type	Description	Default
image	ndarray	The image chunks, (n, chunk) uint8: row r's 3 * columns bytes start chunk r * ceil(3 * columns / chunk).	required
taps_w	ndarray	(n, words) int32, the width table, chunk k patch column k.	required
taps_h	ndarray	(n, words) int32, the height table, chunk b band b.	required
rows	int	The image's rows.	required
columns	int	The image's columns.	required
out_rows	int	The resized rows.	required
out_columns	int	The resized columns.	required
words	int	The int32 words of a table chunk.	356
chunk	int	The bytes of an image chunk.	4096
patch_columns	int	The most patch columns one core holds.	8
cores	int	The cores the patch columns are dealt to.	16
as_rows	bool	When True, resize_emit_row's raster instead of resize_emit's patches.	False

Returns:

Name	Type	Description
	ndarray	With as_rows, (bands * 16, strips * 16, 3) uint8 -- out_rows
	ndarray	and out_columns rounded up to whole 16-pixel blocks, as the kernel's
	ndarray	own granularity is -- zero past out_rows or out_columns (the
	ndarray	last band's or strip's unreal rows or columns) and past a band or
	ndarray	patch column the kernel refuses. Otherwise (bands * pad, 768)
bf16	ndarray	band b's patch column j is row b * pad + j, ``pad
	ndarray	= (strips // cores + 1) * cores``, and patch columns past the image
	ndarray	are zeros. Sizes the kernel refuses zero every patch, a width chunk
	ndarray	a core refuses zeros that core's patch columns, and a height chunk
	ndarray	refused, or a band whose rows read past the image's, zeros every
	ndarray	patch from that band on. Neither out_rows nor out_columns
	ndarray	need be a multiple of 16: the last strip or band is however many of
	ndarray	its 16 outputs are real, the rest zero, as the kernel's last chunk
	ndarray	of either axis is.

Source code in utils/pyright_stubs/aie/iron/kernels/vision.py

def resize_ref(
    image: np.ndarray,
    taps_w: np.ndarray,
    taps_h: np.ndarray,
    rows: int,
    columns: int,
    out_rows: int,
    out_columns: int,
    *,
    words: int = 356,
    chunk: int = 4096,
    patch_columns: int = 8,
    cores: int = 16,
    as_rows: bool = False,
) -> np.ndarray:
    """Numpy reference for [`resize`][aie.iron.kernels.vision.resize].

    Args:
        image: The image chunks, ``(n, chunk)`` uint8: row ``r``'s ``3 *
            columns`` bytes start chunk ``r * ceil(3 * columns / chunk)``.
        taps_w: ``(n, words)`` int32, the width table, chunk ``k`` patch
            column ``k``.
        taps_h: ``(n, words)`` int32, the height table, chunk ``b`` band
            ``b``.
        rows: The image's rows.
        columns: The image's columns.
        out_rows: The resized rows.
        out_columns: The resized columns.
        words: The int32 words of a table chunk.
        chunk: The bytes of an image chunk.
        patch_columns: The most patch columns one core holds.
        cores: The cores the patch columns are dealt to.
        as_rows: When ``True``, [`resize_emit_row`][aie.iron.kernels.vision.resize]'s
            raster instead of [`resize_emit`][aie.iron.kernels.vision.resize]'s patches.

    Returns:
        With ``as_rows``, ``(bands * 16, strips * 16, 3)`` uint8 -- out_rows
        and out_columns rounded up to whole 16-pixel blocks, as the kernel's
        own granularity is -- zero past ``out_rows`` or ``out_columns`` (the
        last band's or strip's unreal rows or columns) and past a band or
        patch column the kernel refuses. Otherwise ``(bands * pad, 768)``
        bf16: band ``b``'s patch column ``j`` is row ``b * pad + j``, ``pad
        = (strips // cores + 1) * cores``, and patch columns past the image
        are zeros. Sizes the kernel refuses zero every patch, a width chunk
        a core refuses zeros that core's patch columns, and a height chunk
        refused, or a band whose rows read past the image's, zeros every
        patch from that band on. Neither ``out_rows`` nor ``out_columns``
        need be a multiple of 16: the last strip or band is however many of
        its 16 outputs are real, the rest zero, as the kernel's last chunk
        of either axis is.
    """
    side = _RESIZE_SIDE
    win_cap = 2 * ((words - _RESAMPLE_HEADER) // side - 2) - 1
    strips = -(-out_columns // side) if out_columns > 0 else 0
    bands = -(-out_rows // side) if out_rows > 0 else 0
    pad = (strips // cores + 1) * cores
    out = np.zeros((bands, pad, side, side, 3), np.uint8)
    sized = (
        rows >= 1
        and columns >= 1
        and out_rows >= 1
        and out_columns >= 1
        and pad // cores <= patch_columns
    )
    if not sized:
        if as_rows:
            return np.zeros((max(bands, 0) * side, max(strips, 0) * side, 3), np.uint8)
        return _RESIZE_PIXELS[out.reshape(bands * pad, -1)]

    cpr = -(-3 * columns // chunk)
    lines = np.asarray(image, np.uint8).reshape(-1)[: rows * cpr * chunk]
    pixels = lines.reshape(rows, cpr * chunk)[:, : 3 * columns]
    pixels = pixels.reshape(rows, columns, 3).astype(np.int64)

    # Across: each core's patch columns at the precision and window of the
    # last width chunk it takes.
    ok = [True] * cores
    mid = np.zeros((rows, strips, side, 3), np.int64)
    for c in range(cores):
        filters = []
        for k in range(c, strips, cores):
            t = np.asarray(taps_w[k], np.int32)
            p, win = int(t[0]), int(t[1])
            # The last strip's real outputs, past which a chunk's slots are
            # unwritten, zero-headered padding, not taps to read.
            n = side if (k + 1) * side <= out_columns else out_columns - k * side
            if not (1 <= p <= 22 and 1 <= win <= win_cap) or (
                t[2] != k * side or t[3] != n
            ):
                ok[c] = False
                break
            start, count, weights = _resize_slots(t, win)
            first, end = int(start[0]), int(start[n - 1] + count[n - 1])
            if (
                (start[:n] < first).any()
                or (count[:n] < 0).any()
                or (count[:n] > win).any()
                or (start[:n] + count[:n] > end).any()
                or first < 0
                or end > columns
                or end - first > 5 * win_cap + 8
            ):
                ok[c] = False
                break
            hweight = np.zeros((side, 64), np.int64)
            for o in range(n):
                hweight[o, : count[o]] = weights[o][: count[o]]
            filters.append((k, start, hweight, p, win))
        if not ok[c] or not filters:
            continue
        _, _, _, hp, win = filters[-1]
        taps = 64 if win > 32 else 32
        for k, start, hweight, _, _ in filters:
            index = np.minimum(start[:, None] + np.arange(taps), columns - 1)
            acc = np.einsum("rotc,ot->roc", pixels[:, index], hweight[:, :taps])
            mid[:, k] = np.clip((acc + (1 << (hp - 1))) >> hp, 0, 255)

    # Down: a band's output rows as their input rows arrive, every core alike.
    arrived = 0
    for b in range(bands):
        vt = np.asarray(taps_h[b], np.int32)
        p, win = int(vt[0]), int(vt[1])
        # The last band's real rows, past which a chunk's slots are
        # unwritten, zero-headered padding, not taps to read.
        real_rows = side if (b + 1) * side <= out_rows else out_rows - b * side
        if not (1 <= p <= 22 and 1 <= win <= win_cap) or (
            vt[2] != b * side or vt[3] != real_rows
        ):
            break
        start, count, weights = _resize_slots(vt, win)
        e = min(max(int(start[real_rows - 1] + count[real_rows - 1]), 0), rows)
        hold = np.zeros((side, strips, side, 3), np.int64)
        done, failed = 0, False
        for now in range(arrived, max(arrived, e) + 1):
            for o in range(done, side):
                if o >= real_rows:
                    done = o + 1  # past the real rows: already zero
                    continue
                s, n = int(start[o]), int(count[o])
                if s + n > now:
                    break
                if s < 0 or s < now - win_cap or n < 0 or n > win_cap:
                    failed = True
                    break
                w = weights[o][:n].astype(np.int64)
                acc = np.einsum("k,kjxc->jxc", w, mid[s : s + n])
                hold[o] = np.clip((acc + (1 << (p - 1))) >> p, 0, 255)
                done = o + 1
            if failed:
                break
        arrived = max(arrived, e)
        if failed or done < side:
            break
        for c in range(cores):
            for k in range(c, strips, cores) if ok[c] else ():
                out[b, k] = hold[:, k]
    if as_rows:
        raster = out[:, :strips].transpose(0, 2, 1, 3, 4)
        return raster.reshape(bands * side, strips * side, 3)
    return _RESIZE_PIXELS[out.reshape(bands * pad, -1)]

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 in kernels.DTYPES["<factory>"], and each factory's submodule defines its references as *_ref functions (eltwise.add_ref, eltwise.add_i8_ref, reduce.reduce_max_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,
    true_function: (
        Callable[..., ndarray | tuple[ndarray, ...]] | 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,
    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,
    call_index: int | None = None,
    out_last: bool = False,
    required_calls: int | None = None,
    even_calls: bool = True,
    trace: Trace | None = None,
    uses_lut: bool = False,
    alignments: tuple[tuple[int, int], ...] = (),
    per_build: tuple[
        tuple[Build, Callable | None, Tolerance | None], ...
    ] = (),
)

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, without per_build, 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.
per_build	tuple[tuple[Build, Callable | None, Tolerance | None], ...]	(build, reference, tolerance) for each of BUILDS, for a kernel whose arithmetic or measured bound differs by build: a LUT one build reads and another does not, a bound measured on each arch. reference and tolerance are then None, and for_build picks the row of the build that produced an output, once that is known.
true_function	Callable[..., ndarray | tuple[ndarray, ...]] | None	The function the kernel approximates, for error reports, where the reference (or a build's row) models the kernel instead, as a LUT build's does. None: the reference is it.
ops_per_call	int | None	Arithmetic operations per call; None means one per output element.
out_valid	int | None	Meaningful leading elements of a DMA-padded output tile; None means the whole tile.
sample	Callable[..., list] | None	sample(rng, calls) -> list[np.ndarray] for inputs with structure a dtype cannot express; None draws random data.
acc_dtype	type | None	The accumulator type, or None when nothing accumulates. With reduction it bounds the inputs so the accumulator cannot overflow.
reduction	int | None	Terms summed into one output element per call; None means one.
setup	Callable[[], object] | None	A kernel to run once on the core first (conv_even sets the rounding mode a bf16 store needs); None when the source sets its own mode or narrows nothing. A Worker handed the kernel calls it before its loop.
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 and tensor Param ones the kernel writes, such as a running state; factory(fn) returns the kernel that initializes the buffer before each call.
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.
call_index	int | None	The bound scalar Param the builder sets to each call's index, 0 first, for a kernel whose state carries from one call to the next and restarts at 0. The reference sees every call at once, in order. None: the bound value holds.
out_last	bool	Whether only the last of the calls writes the one Out, the earlier ones building up state the core keeps (with call_index), as a reduction split across calls does. The builder hands every call the same output tile and drains it once, and the reference returns that one tile.
required_calls	int | None	The exact positive number of calls needed to produce the output of an out_last kernel; required when out_last is set, otherwise None.
even_calls	bool	Whether every call does the same work. Timing then reports the cheapest call, the rest being the core waiting; calls of unequal work (each its own chunk of a resample) report the median call.
trace	Trace | None	The Trace shape of the kernel's markers. Every library factory declares one; None (undeclared) is only for ad-hoc kernels, and cycles_per_call refuses it.
uses_lut	bool	Whether the kernel gathers through an aie::lut<4> table pair, so a build should verify the two tables land in different banks. Set it on the contract, not per source file: the LUT often comes in through a header (lut_based_ops.h, lut_inv.h).
alignments	tuple[tuple[int, int], ...]	(index, bytes) pairs for arguments the kernel loads as whole vectors from their start, so they must begin at a multiple of bytes. A call handed a memref.view at a constant offset that breaks one raises.

Overflow, rounding and NaN handling are not declared twice: the reference is the arithmetic model and the tolerance the slack against it.

out_indices property

out_indices: tuple[int, ...]

Output argument positions, in declaration order.

out_index property

out_index: int

Position of the one output, written (Out) or accumulated into (InOut).

Raises for a kernel with several outputs: code that must handle any kernel reads out_indices.

accumulates property

accumulates: bool

Whether the kernel reads its output back (InOut), as C += A * B does.

for_build

for_build(build: Build) -> KernelContract

Return this contract with the reference and tolerance of build's row.

A contract without per_build is the same for every build.

Raises:

Type	Description
ValueError	When per_build has no row for build.

Source code in utils/pyright_stubs/aie/iron/kernels/_common.py

def for_build(self, build: Build) -> "KernelContract":
    """Return this contract with the reference and tolerance of ``build``'s row.

    A contract without ``per_build`` is the same for every build.

    Raises:
        ValueError: When ``per_build`` has no row for ``build``.
    """
    if not self.per_build:
        return self
    for row, reference, tolerance in self.per_build:
        if row == build:
            return replace(
                self, reference=reference, tolerance=tolerance, per_build=()
            )
    raise ValueError(f"no per_build row for {build}")

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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/iron/kernels/_common.py

def validate_types(self, arg_types):
    """Validate contracts against NumPy tensor aliases and scalar dtypes.

    Raw MLIR types remain usable by ExternalFunction, but the host contract
    requires NumPy declarations for sampling, layouts and references.
    """
    if len(arg_types) != len(self.roles):
        raise ValueError("roles must have one entry per argument")
    bound = dict(self.parameter_bindings)
    for i, (role, arg_type) in enumerate(zip(self.roles, arg_types)):
        tensor = _is_tensor_type(arg_type)
        if tensor:
            try:
                np_ndarray_type_get_shape(arg_type)
                dtype = np_ndarray_type_get_dtype(arg_type)
            except (AssertionError, IndexError, TypeError) as exc:
                raise ValueError(
                    f"argument {i}: expected np.ndarray[shape, np.dtype[dtype]]"
                ) from exc
        else:
            dtype = arg_type
        if dtype not in get_args(NpuDType):
            raise ValueError(
                f"argument {i}: kernel contracts require NumPy tensor aliases "
                "or supported NumPy scalar dtypes"
            )
        if not tensor and role is not Param:
            raise ValueError(f"argument {i}: scalar arguments require Param")
        if self.layouts and self.layouts[i] is not None and not tensor:
            raise ValueError(f"argument {i}: layouts require tensor arguments")
        if i not in bound:
            continue
        value = bound[i]
        if tensor:
            value = np.asarray(value)
            layout = self.layouts[i] if self.layouts else None
            shape = (
                layout.shape
                if layout is not None
                else np_ndarray_type_get_shape(arg_type)
            )
            if value.ndim == 0 or value.size != int(np.prod(shape)):
                raise ValueError(
                    f"argument {i}: tensor parameter must contain {shape} elements"
                )
        elif not isinstance(value, (int, float, np.integer, np.floating)):
            raise ValueError(f"argument {i}: expected scalar parameter")

TensorLayout dataclass

TensorLayout(
    shape: tuple[int, ...],
    pack: Callable | None = None,
    unpack: Callable | None = None,
    stream: TensorAccessPattern | None = None,
    block: tuple[int, ...] | None = None,
)

How a kernel wants one tensor operand laid out.

shape is the logical tile. pack and unpack are the reversible host codec between (calls, *shape) and (calls, storage_elements); identity is the default. stream is the DMA transform (to_stream) a design applies on the hop that feeds this operand to the kernel or drains it, None when the operand streams as stored; block is the micro-tile the kernel consumes or produces, (r, s) for an MMUL operand. The codec is built from the same two facts, so the host and the design agree by construction. None of this is an algorithm or a whole-problem iteration schedule.

Param

Read-only test-fixture parameter; its ABI determines scalar or tensor.

The generic harness holds its value fixed across calls. This is not a C++ operand lifetime: direct designs may pass a new value on every kernel call.

Trace dataclass

Trace(shape: str, reason: str | None = None)

How a kernel's event0()/event1() markers bracket one call.

Trace.whole_call(): one pair brackets every call of the entry symbol and nothing it calls emits another, so each trace interval is one call. Trace.none(reason): a call emits no marker. Trace.partial(reason): markers exist but do not bracket each call exactly once (around an inner loop, or skipped on an early return), so intervals cannot be attributed to calls. test_kernel_trace_markers.py checks the declaration against the compiled IR of every library build.

aie.iron.algorithms.KernelDesign turns any contract-bearing kernel 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.

Bind the target device before calling the kernel factory, and give the call count and guard once, to the KernelDesign; every method sizes, builds and checks from them:

from aie.iron import kernels
from aie.iron.algorithms import KernelDesign

kd = KernelDesign(kernels.relu(tile_size=1024), calls=4, guard=True)
inputs = kd.sample_inputs()
outputs = kd.run(inputs)
verdict = kd.check(outputs, inputs)
assert verdict, verdict.detail

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.

KernelDesign

KernelDesign(
    fn,
    *,
    calls=1,
    shape=None,
    guard=False,
    scalars=(),
    stack_bytes=None,
    arg_byte_offsets=(),
    aiecc_flags=None
)

calls independent calls of one contract-bearing kernel.

The configuration is fixed at construction, and every method sizes, lays out, builds and checks from it; construct another to change it.

scalars supply the unbound scalar Params. The unbound tensor Params are embedded in the design: run takes them from its logical inputs, and design(params) from params. Changing them recompiles the design. Direct designs can supply different operands on each call.

stack_bytes fixes the core stack. Leave it unset and aiecc sizes the stack from the linked core; set it where aiecc cannot measure, as for a Chess build. Unset, choosing between one and two sets of tiles reads the kernels' stack frames, so generating a design whose tiles could double buffer compiles its kernels.

With guard=True the core fills each output tile and GUARD_BYTES after it with 0x55 before every call and drains the guard with the tile, so a kernel that skips part of its output or writes past it shows up on the host: output_size and host_args count the guards, strip_guard splits them off and run raises when one changed. bfp outputs carry no guard.

arg_byte_offsets binds tensor Params at a byte offset: ((1, 16),) hands argument 1 a view 16 bytes into a buffer VIEW_PAD larger, poisoned around the data, as a design that packs several weights into one buffer does. The shift is loaded at run time, so the kernel sees an address no design-time check could have folded.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def __init__(
    self,
    fn,
    *,
    calls=1,
    shape=None,
    guard=False,
    scalars=(),
    stack_bytes=None,
    arg_byte_offsets=(),
    aiecc_flags=None,
):
    if fn.contract is None:
        raise ValueError(
            f"kernel '{fn.name}' declares no contract; add a KernelContract"
        )
    fn.contract.validate_types(fn.arg_types())
    self.fn = fn
    self.calls = _calls(calls, shape)
    self.guard = guard
    self.scalars = tuple(scalars)
    self.stack_bytes = stack_bytes
    self.aiecc_flags = tuple(aiecc_flags or ())
    self.arg_byte_offsets = self._byte_offsets(arg_byte_offsets)

design

design(params=())

Return the JIT design with params as its unbound tensor Params.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def design(self, params=()):
    """Return the JIT design with ``params`` as its unbound tensor Params."""
    _device()
    c = self.fn.contract
    if c.unsupported:
        raise ValueError(
            f"{self.fn.name}: the generic harness cannot build this kernel: "
            f"{c.unsupported}"
        )
    self._out_tiles
    self._fifo_plan  # the DMA channel budget is checked before anything builds
    flags: list[str] = list(self.aiecc_flags)
    if any(
        bfp.is_bfp(shape_dtype(t)[1]) for t in self._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 self.fn.use_chess:
        if not any(f.startswith("--check-lut-banks") for f in flags):
            flags.append("--check-lut-banks")
    return _stream.specialize(
        kernel=self.fn,
        plan=self._plan,
        calls=self.calls,
        scalars=self.scalars,
        params=self._encode_params(params),
        guard=self.guard,
        # Only when given, so every other design keeps its cache key.
        **({"stack_bytes": self.stack_bytes} if self.stack_bytes else {}),
        **(
            {"arg_byte_offsets": self.arg_byte_offsets}
            if self.arg_byte_offsets
            else {}
        ),
        **({"aiecc_flags": flags} if flags else {}),
    )

as_mlir

as_mlir()

Lower design() to MLIR.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def as_mlir(self):
    """Lower ``design()`` to MLIR."""
    return self.design().as_mlir()

sample_inputs

sample_inputs(rng=None)

One array per unbound In/tensor Param; only In has a call dimension.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def sample_inputs(self, rng=None):
    """One array per unbound In/tensor Param; only In has a call dimension."""
    rng = np.random.default_rng(0) if rng is None else rng
    c = self.fn.contract
    if c.sample is not None:
        return c.sample(rng, self.calls)
    result = []
    for i in self._tensor_positions:
        s, dt = shape_dtype(self._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 = ((self.calls,) if c.roles[i] is In else ()) + s
        if np.issubdtype(np.dtype(dt), np.integer):
            r = self.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(inputs)

Pack declared layouts and interleave same-type In tiles; omit Param buffers.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def host_layout(self, inputs):
    """Pack declared layouts and interleave same-type In tiles; omit Param buffers."""
    c = self.fn.contract
    positions = self._tensor_positions
    if len(inputs) != len(positions):
        raise ValueError(f"{self.fn.name}: expected {len(positions)} input arrays")
    values = dict(zip(positions, inputs))
    result = []
    for group in self._fifo_plan[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(self._types[i]))
            )
        result.append(
            np.ascontiguousarray(
                arrays[0] if len(arrays) == 1 else np.stack(arrays, axis=1)
            )
        )
    return result

output_size

output_size()

Storage elements per output; a tuple for multiple outputs.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def output_size(self):
    """Storage elements per output; a tuple for multiple outputs."""
    types = self._types
    sizes = tuple(
        (elems(types[i]) + (_guard_elems(types[i]) if i in self._guarded else 0))
        * self._out_tiles
        * (bfp.BLOCK_BYTES if bfp.is_bfp(shape_dtype(types[i])[1]) else 1)
        for i in self.fn.contract.out_indices
    )
    return sizes[0] if len(sizes) == 1 else sizes

strip_guard

strip_guard(outputs)

Split the outputs of a run into data and overrun.

Returns the outputs without their guards, shaped as output_size without guard would size them, and per output the number of guard bytes the kernel changed (0 for an unguarded output).

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def strip_guard(self, outputs):
    """Split the outputs of a run into data and overrun.

    Returns the outputs without their guards, shaped as ``output_size``
    without ``guard`` would size them, and per output the number of guard
    bytes the kernel changed (0 for an unguarded output).
    """
    multiple = isinstance(outputs, tuple)
    data, overrun = [], []
    for i, out in zip(
        self.fn.contract.out_indices, outputs if multiple else (outputs,)
    ):
        if i not in self._guarded:
            data.append(out)
            overrun.append(0)
            continue
        rows = np.ascontiguousarray(out).reshape(self._out_tiles, -1)
        raw = rows.view(np.uint8)
        n = raw.shape[1] - GUARD_BYTES
        overrun.append(int(np.count_nonzero(raw[:, n:] != 0x55)))
        data.append(np.ascontiguousarray(raw[:, :n]).view(out.dtype).reshape(-1))
    if multiple:
        return tuple(data), tuple(overrun)
    return data[0], overrun[0]

host_args

host_args()

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 utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def host_args(self):
    """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``.
    """
    types = self._types
    result = []
    entries = [(In, g) for g in self._fifo_plan[0]]
    entries += [(Out, [i]) for i in self.fn.contract.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 self._guarded:
            n += _guard_elems(types[i])
        rows = self.calls if direction is In else self._out_tiles
        s = (rows, n) if len(indices) == 1 else (rows, len(indices), n)
        result.append(_HostBuffer(direction, s, dt))
    return result

upload

upload(inputs, *, poison=False)

Return design inputs and output tensor(s); splat multiple outputs when calling.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def upload(self, inputs, *, poison=False):
    """Return design inputs and output tensor(s); splat multiple outputs when calling."""
    packed = self.host_layout(inputs)
    args = self.host_args()
    for j, (a, arg) in enumerate(zip(packed, args)):
        if a.size != arg.n_elements or a.dtype != np.dtype(arg.dtype):
            raise ValueError(
                f"{self.fn.name}: input buffer {j} holds {a.size} {a.dtype}; "
                f"{self.calls} call(s) take {arg.n_elements} "
                f"{np.dtype(arg.dtype)}"
            )
    ins = [tensor(a.reshape(-1), dtype=a.dtype, device="npu") for a in packed]
    outs = [
        tensor(
            (
                poisoned(arg.n_elements, arg.dtype)
                if poison
                else np.zeros(arg.n_elements, dtype=arg.dtype)
            ),
            dtype=arg.dtype,
            device="npu",
        )
        for arg in args
        if arg.direction is Out
    ]
    return ins, tuple(outs) if len(outs) > 1 else outs[0]

run

run(inputs, *, trace_config=None)

Run on inputs into poisoned outputs and return them without guards.

The tensor Params come from inputs. The outputs are copies: Tensor.numpy() views the XRT buffer's mapped host memory, which dies with the run's tensors.

Raises:

Type	Description
RuntimeError	When the kernel changed a guard.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def run(self, inputs, *, trace_config=None):
    """Run on ``inputs`` into poisoned outputs and return them without guards.

    The tensor Params come from ``inputs``. The outputs are copies:
    ``Tensor.numpy()`` views the XRT buffer's mapped host memory, which
    dies with the run's tensors.

    Raises:
        RuntimeError: When the kernel changed a guard.
    """
    design = self.design(self.fn.param_values(inputs))
    ins, out = self.upload(inputs, poison=True)
    outs = out if isinstance(out, tuple) else (out,)
    design(
        *ins,
        *outs,
        **({"trace_config": trace_config} if trace_config else {}),
    )
    got = [o.numpy().copy() for o in outs]
    data, overrun = self.strip_guard(tuple(got) if len(got) > 1 else got[0])
    if any(np.atleast_1d(overrun)):
        raise RuntimeError(
            f"{self.fn.name}: changed {overrun} guard byte(s) past its output "
            f"({GUARD_BYTES} after each tile)"
        )
    return data

check

check(outputs, inputs)

Judge run's outputs against the contract's reference for inputs.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def check(self, outputs, inputs):
    """Judge ``run``'s outputs against the contract's reference for ``inputs``."""
    return self.fn.judge(
        outputs,
        self.fn.expected(inputs, scalars=self.scalars),
        inputs=inputs,
        scalars=self.scalars,
        calls=self.calls,
    )

traced_intervals

traced_intervals() -> int

How many event0/event1 intervals a traced run emits.

A trace_size that holds fewer truncates the run; 0 when the kernel itself is not timed.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def traced_intervals(self) -> int:
    """How many ``event0``/``event1`` intervals a traced run emits.

    A ``trace_size`` that holds fewer truncates the run; ``0`` when the
    kernel itself is not timed.
    """
    trace = self.fn.contract.trace
    if trace is None or trace.shape != "whole_call":
        return 0
    setup, inits = self._traced()
    return int(setup) + self.calls * (len(inits) + 1)

cycles_per_call

cycles_per_call(
    inputs, *, trace_size, workdir
) -> 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 utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def cycles_per_call(self, inputs, *, trace_size, workdir) -> 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 = self.fn.contract
    if c.trace is None:
        raise ValueError(f"{self.fn.name}: the contract declares no trace")
    if c.trace.shape != "whole_call":
        return CallCycles(untimed=c.trace.reason)
    setup, inits = self._traced()
    workdir = Path(workdir)
    workdir.mkdir(parents=True, exist_ok=True)
    cfg = TraceConfig(trace_size=trace_size, trace_file=str(workdir / "trace.txt"))
    self.run(inputs, 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=self.calls,
        per_call=len(inits) + 1,
        setup=int(setup),
        flush=TRACE_FLUSH,
    )
    if not per_kernel[-1]:
        raise RuntimeError(
            f"{self.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,
    )

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.

split_intervals

split_intervals(
    durations, *, calls, per_call, setup=0, flush=0
)

Label an interval stream: setup intervals, then per_call per call.

The harness runs the setup kernel once, then each call runs its traced initializers in contract order and the kernel last, so interval j of call n sits at setup + n * per_call + j. Returns (setup intervals, one tuple per per-call kernel, truncated). Up to flush intervals of at most FLUSH_CYCLES may follow, the pairs the core emits to push the trace out. Anything else is some kernel emitting markers it does not declare; one whose extra intervals are that short and that few passes as a flush.

Source code in utils/pyright_stubs/aie/iron/algorithms/kernel_design.py

def split_intervals(durations, *, calls, per_call, setup=0, flush=0):
    """Label an interval stream: ``setup`` intervals, then ``per_call`` per call.

    The harness runs the setup kernel once, then each call runs its traced
    initializers in contract order and the kernel last, so interval ``j`` of
    call ``n`` sits at ``setup + n * per_call + j``. Returns ``(setup
    intervals, one tuple per per-call kernel, truncated)``. Up to ``flush``
    intervals of at most ``FLUSH_CYCLES`` may follow, the pairs the core
    emits to push the trace out. Anything else is some kernel emitting
    markers it does not declare; one whose extra intervals are that short
    and that few passes as a flush.
    """
    durations = [int(d) for d in durations]
    expected = setup + calls * per_call
    tail = durations[expected:]
    if len(tail) > flush or any(d > FLUSH_CYCLES for d in tail):
        raise RuntimeError(
            f"expected {expected} trace intervals and up to {flush} flush pairs, "
            f"got {len(durations)}, the extra {tail[:flush + 1]} cycles; a "
            "kernel on the core emits markers its contract's trace does not "
            "declare"
        )
    durations = durations[:expected]
    head, stream = durations[:setup], durations[setup:]
    return (
        tuple(head),
        [tuple(stream[j::per_call]) for j in range(per_call)],
        len(durations) < expected,
    )

test/python/npu/test_kernels_perf.py times the library: correctness first, then cycles, wall time and build size, gated by a device preflight and a measurement-sanity test. It is an ordinary pytest module, so -k selects cases and the session's exit status decides whether any numbers are written. --baseline-sources DIR measures every selected case a second time with its kernels from DIR and compares the two runs' raw output words and cycles in --perf-meta. The timing helpers live here:

Benchmarking helpers for NPU kernel callables.

Stats keeps the raw sample list and robust statistics (median, MAD, p95, coefficient of variation, the median's standard error) beside avg/min/max, so callers can report jitter, not just central tendency. The numbers are numpy's.

run_iters times one callable; run_interleaved times several against each other, round by round, so a drift in the host or the NPU lands on all of them alike. wait_idle holds a measurement until no other process uses the NPU.

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,
    median_se_us: float = 0.0,
    samples_us: list[float] = list(),
)

as_dict

as_dict(prefix: str = '') -> dict

Flat dict for JSON emission (samples omitted).

Source code in utils/pyright_stubs/aie/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,
    dispatch: Stats | None = None,
    dispatch_overhead: Stats | None = None,
    npu: Stats | None = None,
)

One callable's timings, every one on the host's clock.

Attributes:

Name	Type	Description
e2e	Stats	The whole Python call.
dispatch	Stats | None	Submit to completion of the call's dispatch, or None for a callable that returns no KernelResult.
dispatch_overhead	Stats | None	Submit to completion of a dispatch that does nothing, on the same hardware context, one after each timed call; or None where the runtime has no empty dispatch for the kernel (KernelHandle.empty_dispatch_ns), as for an ELF module, or a full ELF without aiecc's empty sequence.
npu	Stats | None	Each dispatch less the overhead's median: the time the NPU spent on the call's work. None without both.

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.

CompletionCores

CompletionCores(before: Mapping[str, int])

Bases: AbstractContextManager

While entered, the calling thread held on the CPUs the NPU's completions interrupted.

A core in deep idle adds its exit latency to long runs but not short ones. A nice-19 spinner on each of those CPUs' SMT siblings keeps the core shallow. Linux only; elsewhere nothing is held.

Parameters:

Name	Type	Description	Default
before	Mapping[str, int]	interrupts() before the calls whose completions choose the CPUs.	required

Source code in utils/pyright_stubs/aie/utils/benchmark.py

def __init__(self, before: Mapping[str, int]):
    self.cpus: set[int] = set()
    for irq, count in self.interrupts().items():
        if count > before.get(irq, 0):
            self.cpus |= self._cpus(f"/proc/irq/{irq}/effective_affinity_list")
    # A container's cpuset may leave out the cores the completions
    # interrupt: only the ones the process may run on can be held.
    allowed = os.sched_getaffinity(0) if self.cpus else set()
    self.cpus &= allowed
    self.siblings: set[int] = set()
    for cpu in self.cpus:
        self.siblings |= self._cpus(
            f"/sys/devices/system/cpu/cpu{cpu}/topology/thread_siblings_list"
        )
    self.siblings = (self.siblings & allowed) - self.cpus
    self._affinity: set[int] | None = None
    self._spinners: list[subprocess.Popen] = []

interrupts classmethod

interrupts() -> dict[str, int]

Each NPU completion interrupt's count over every CPU, by IRQ number.

Returns:

Type	Description
dict[str, int]	dict[str, int]: Empty where /proc/interrupts does not exist.

Source code in utils/pyright_stubs/aie/utils/benchmark.py

@classmethod
def interrupts(cls) -> dict[str, int]:
    """Each NPU completion interrupt's count over every CPU, by IRQ number.

    Returns:
        dict[str, int]: Empty where ``/proc/interrupts`` does not exist.
    """
    path = Path("/proc/interrupts")
    if not path.is_file():
        return {}
    counts = {}
    for line in path.read_text().splitlines():
        fields = line.split()
        if fields and fields[-1] == cls.IRQ:
            counts[fields[0].rstrip(":")] = sum(
                int(f) for f in fields[1:] if f.isdigit()
            )
    return counts

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 a KernelResult, its npu_time (the runtime's host clock from submit to completion) is the call's dispatch time; with the kernel's handle, an empty dispatch after each timed call measures the overhead every dispatch carries, and the NPU's time is the difference.

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 utils/pyright_stubs/aie/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 a ``KernelResult``, its ``npu_time`` (the runtime's host
    clock from submit to completion) is the call's dispatch time; with the
    kernel's handle, an empty dispatch after each timed call measures the
    overhead every dispatch carries, and the NPU's time is the difference.

    ``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] = []
    dispatch_samples: list[float] = []
    overhead_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)
        handle, result = _kernel_call(ret)
        if result is not None:
            dispatch_samples.append(result.npu_time / 1_000.0)
        if handle is not None and (empty_ns := handle.empty_dispatch_ns()) is not None:
            overhead_samples.append(empty_ns / 1_000.0)

    dispatch = _stats(dispatch_samples) if dispatch_samples else None
    overhead = _stats(overhead_samples) if overhead_samples else None
    return BenchmarkResult(
        e2e=_stats(e2e_samples),
        dispatch=dispatch,
        dispatch_overhead=overhead,
        npu=(
            _stats([d - overhead.median_us for d in dispatch_samples])
            if dispatch is not None and overhead is not None
            else None
        ),
    )

run_interleaved

run_interleaved(
    fns: Sequence[Callable[[], object]],
    *,
    rounds: int = 8,
    calls: int = 50,
    settle: int = 2,
    cutoff: float = 1.5,
    groups: Sequence[Hashable | None] | None = None,
    fastest_us: float = inf,
    completion_cores: bool = False,
    max_cpu_wait_us: float = CPU_WAIT_US,
    idle: bool = False,
    log: Callable[[str], None] = print
) -> list[Stats]

Time callables against each other in interleaved rounds.

Each round runs every callable still timed calls times in turn and keeps each one's median call. A call's time is the npu_time its return value carries (run_iters's rule), else the time around it. Each callable is called twice first, untimed, so first-run setup lands on nobody's figure.

Parameters:

Name	Type	Description	Default
fns	Sequence[Callable[[], object]]	The callables, each taking no arguments.	required
rounds	int	Rounds to time.	8
calls	int	Calls per callable per round.	50
settle	int	Rounds after which cutoff applies.	2
cutoff	float	With groups, a group whose slowest callable takes more than cutoff times the fastest group's is timed no further; its figure is from the rounds it had.	1.5
groups	Sequence[Hashable | None] | None	Per callable, the setting it measures (callables of one group are stopped together), or None for one never stopped. Every callable is timed every round if not given.	None
fastest_us	float	The fastest group's slowest median from callables timed before these, which cutoff also measures against.	inf
completion_cores	bool	Hold the timing thread on the CPUs the callables' NPU completions interrupt (CompletionCores), for callables that dispatch to the NPU.	False
max_cpu_wait_us	float	The most a callable's median call may wait for a CPU in its median round before the figures are refused; math.inf accepts figures from a busy host.	CPU_WAIT_US
idle	bool	Time only while no other process holds the NPU: wait for it to be idle first (wait_idle), and time a round again if another holds it when the round ends. A use that starts and ends within one round goes unseen, and one held but idle context stops the measurement with wait_idle's TimeoutError. Only this user's processes are seen, as other_npu_users says.	False
log	Callable[[str], None]	Where idle reports waiting or timing a round again.	print

Returns:

Type	Description
list[Stats]	list[Stats]: Per callable, the statistics of its per-round medians.

Raises:

Type	Description
TimeoutError	With idle, another process held the NPU for IDLE_WAIT seconds, or used it during more than 3 * rounds rounds; the message names its pids.
RuntimeError	A callable's median call waited for a CPU past max_cpu_wait_us (cpu_wait_us) in its median round: the host was busy enough to hold the timing thread off its core, so the figures are refused. Measure with the host otherwise idle.
ValueError	rounds or calls is below one, or groups is not one per callable.

Source code in utils/pyright_stubs/aie/utils/benchmark.py

def run_interleaved(
    fns: Sequence[Callable[[], object]],
    *,
    rounds: int = 8,
    calls: int = 50,
    settle: int = 2,
    cutoff: float = 1.5,
    groups: Sequence[Hashable | None] | None = None,
    fastest_us: float = math.inf,
    completion_cores: bool = False,
    max_cpu_wait_us: float = CPU_WAIT_US,
    idle: bool = False,
    log: Callable[[str], None] = print,
) -> list[Stats]:
    """Time callables against each other in interleaved rounds.

    Each round runs every callable still timed ``calls`` times in turn and
    keeps each one's median call. A call's time is the ``npu_time`` its
    return value carries (``run_iters``'s rule), else the time around it.
    Each callable is called twice first, untimed, so first-run setup lands on
    nobody's figure.

    Args:
        fns: The callables, each taking no arguments.
        rounds: Rounds to time.
        calls: Calls per callable per round.
        settle: Rounds after which `cutoff` applies.
        cutoff: With `groups`, a group whose slowest callable takes more than
            `cutoff` times the fastest group's is timed no further; its
            figure is from the rounds it had.
        groups: Per callable, the setting it measures (callables of one group
            are stopped together), or None for one never stopped. Every
            callable is timed every round if not given.
        fastest_us: The fastest group's slowest median from callables timed
            before these, which `cutoff` also measures against.
        completion_cores: Hold the timing thread on the CPUs the callables'
            NPU completions interrupt (`CompletionCores`), for callables that
            dispatch to the NPU.
        max_cpu_wait_us: The most a callable's median call may wait for a CPU
            in its median round before the figures are refused; ``math.inf``
            accepts figures from a busy host.
        idle: Time only while no other process holds the NPU: wait for it to
            be idle first (`wait_idle`), and time a round again if another
            holds it when the round ends. A use that starts and ends within
            one round goes unseen, and one held but idle context stops the
            measurement with `wait_idle`'s ``TimeoutError``. Only this
            user's processes are seen, as `other_npu_users` says.
        log: Where `idle` reports waiting or timing a round again.

    Returns:
        list[Stats]: Per callable, the statistics of its per-round medians.

    Raises:
        TimeoutError: With `idle`, another process held the NPU for
            ``IDLE_WAIT`` seconds, or used it during more than ``3 * rounds``
            rounds; the message names its pids.
        RuntimeError: A callable's median call waited for a CPU past
            `max_cpu_wait_us` (`cpu_wait_us`) in its median round: the host was
            busy enough to hold the timing thread off its core, so the
            figures are refused. Measure with the host otherwise idle.
        ValueError: `rounds` or `calls` is below one, or `groups` is not one
            per callable.
    """
    if rounds < 1 or calls < 1:
        raise ValueError(f"{rounds} rounds of {calls} calls: both must be positive")
    if groups is not None and len(groups) != len(fns):
        raise ValueError(f"{len(groups)} groups for {len(fns)} callables")
    for fn in fns:
        fn()
    before = CompletionCores.interrupts()
    for fn in fns:
        fn()
    medians: list[list[float]] = [[] for _ in fns]
    waited: list[list[float]] = [[] for _ in fns]
    live = set(range(len(fns)))
    if idle:
        wait_idle(log, holding=True)
    r = retries = 0
    with CompletionCores(before) if completion_cores else nullcontext():
        while r < rounds:
            timed = {}
            for i in sorted(live):
                times, waits = [], []
                for _ in range(calls):
                    wait = cpu_wait_us()
                    start = time.perf_counter()
                    ret = fns[i]()
                    elapsed_us = (time.perf_counter() - start) * 1e6
                    waited_after = cpu_wait_us()
                    if wait is not None and waited_after is not None:
                        waits.append(waited_after - wait)
                    _, result = _kernel_call(ret)
                    times.append(
                        elapsed_us if result is None else result.npu_time / 1e3
                    )
                timed[i] = (
                    float(np.median(times)),
                    float(np.median(waits)) if waits else 0.0,
                )
            if idle and (held := other_npu_users(holding=True)):
                retries += 1
                if retries > rounds * 3:
                    raise TimeoutError(
                        f"pid {sorted(held)} used the NPU during {retries} rounds; "
                        "time with the NPU otherwise idle"
                    )
                log(
                    f"    pid {sorted(held)} used the NPU during a round: timing it again"
                )
                wait_idle(log, holding=True)
                continue
            for i, (us, wait) in timed.items():
                medians[i].append(us)
                waited[i].append(wait)
            r += 1
            if groups is None or r < settle:
                continue
            slowest: dict[Hashable | None, float] = {}
            for i in live:
                us = float(np.median(medians[i]))
                slowest[groups[i]] = max(slowest.get(groups[i], 0.0), us)
            bar = cutoff * min([fastest_us, *slowest.values()])
            live = {i for i in live if groups[i] is None or slowest[groups[i]] <= bar}
    late = max(float(np.median(w)) for w in waited)
    if late > max_cpu_wait_us:
        raise RuntimeError(
            f"a callable's median call waited {late:.0f} us for a CPU in its "
            "median round; time with the host otherwise idle"
        )
    return [Stats.from_samples(m) for m in medians]

cpu_wait_us

cpu_wait_us() -> float | None

How long the calling thread has waited, runnable, for a CPU, microseconds in all.

Read from /proc/thread-self/schedstat.

Returns:

Type	Description
float | None	float | None: The run delay so far, or None where the kernel does not
float | None	report it (not Linux, or schedstats off).

Source code in utils/pyright_stubs/aie/utils/benchmark.py

def cpu_wait_us() -> float | None:
    """How long the calling thread has waited, runnable, for a CPU, microseconds in all.

    Read from ``/proc/thread-self/schedstat``.

    Returns:
        float | None: The run delay so far, or None where the kernel does not
        report it (not Linux, or schedstats off).
    """
    try:
        return int(Path("/proc/thread-self/schedstat").read_text().split()[1]) / 1e3
    except OSError:
        return None

preflight

preflight() -> Preflight

Describe the device through the host runtime (any backend).

Source code in utils/pyright_stubs/aie/utils/benchmark.py

def preflight() -> Preflight:
    """Describe the device through the host runtime (any backend)."""
    import aie.utils as aie_utils
    from aie.utils.compile.utils import resolve_target_arch

    runtime = aie_utils.DefaultNPURuntime
    if runtime is None:
        raise RuntimeError("no NPU runtime is available (XRT, HRX or HSA)")
    device = runtime.device()
    arch = resolve_target_arch(device)
    return Preflight(
        npu="npu2" if arch == "aie2p" else "npu1",
        arch=arch,
        device=runtime.device_name() or type(device).__name__,
        pmode=runtime.power_mode(),
    )

provenance

provenance(**extra: str | None) -> str

Return a one-line description of what produced a measurement.

The git commit (GITHUB_SHA or git rev-parse HEAD), the Peano that compiles the kernels (peano_version()), the installed mlir_aie version, the kernel tree (MLIR_AIE_KERNEL_SOURCES when set, and kernel_tree_digest()), the selected runtime's provenance, the runner the job ran on (RUNNER_NAME, under GitHub Actions: nightly runs move between hosts of a generation, and a host swap moves dispatch_us as a regression would), and any extra fields (device="NPU Strix", pmode="performance") as key value pairs. A benchmark row records it so a number can be traced to a toolchain, the kernel sources and the machine.

A package that is not installed is left out rather than recorded as unknown. CI builds mlir_aie from source and puts it on PYTHONPATH, so there is no distribution to read a version from, and every published row would otherwise carry a word that reads like a lookup failure. The commit already identifies that build.

Source code in utils/pyright_stubs/aie/utils/benchmark.py

def provenance(**extra: str | None) -> str:
    """Return a one-line description of what produced a measurement.

    The git commit (``GITHUB_SHA`` or ``git rev-parse HEAD``), the Peano that
    compiles the kernels (``peano_version()``), the installed ``mlir_aie``
    version, the kernel tree
    (``MLIR_AIE_KERNEL_SOURCES`` when set, and ``kernel_tree_digest()``),
    the selected runtime's provenance, the runner the job ran
    on (``RUNNER_NAME``, under GitHub Actions: nightly runs move between
    hosts of a generation, and a host swap moves ``dispatch_us`` as a regression
    would), and any ``extra`` fields (``device="NPU Strix"``,
    ``pmode="performance"``) as ``key value`` pairs. A benchmark row records
    it so a number can be traced to a toolchain, the kernel sources and
    the machine.

    A package that is not installed is left out rather than recorded as
    unknown. CI builds ``mlir_aie`` from source and puts it on ``PYTHONPATH``,
    so there is no distribution to read a version from, and every published
    row would otherwise carry a word that reads like a lookup failure. The
    commit already identifies that build.
    """
    import aie.utils as aie_utils

    def pkg(name: str) -> str | None:
        try:
            return _pkg_version(name)
        except PackageNotFoundError:
            return None

    commit = os.environ.get("GITHUB_SHA")
    if not commit:
        try:
            commit = subprocess.run(
                ["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True
            ).stdout.strip()
        except (OSError, subprocess.CalledProcessError):
            commit = "unknown"
    runtime_fields = {}
    try:
        runtime = aie_utils.DefaultNPURuntime
        if runtime is not None:
            runtime_fields = runtime.provenance()
    except Exception:  # noqa: BLE001 - optional metadata must not fail a measurement
        pass
    fields = {
        "commit": commit[:10],
        "peano": peano_version(),
        "mlir_aie": pkg("mlir_aie"),
        "kernel_sources": os.environ.get("MLIR_AIE_KERNEL_SOURCES"),
        "kernels": kernel_tree_digest(),
        **runtime_fields,
        "host": os.environ.get("RUNNER_NAME"),
        **extra,
    }
    return " | ".join(f"{k} {v}" for k, v in fields.items() if v)

kernel_tree_digest

kernel_tree_digest() -> str | None

Return a 12-hex digest of the kernel sources the library factories compile.

Every file under aie_kernels_dir() and aie_runtime_lib_dir(), by relative path and content. The commit alone cannot say which kernels ran: MLIR_AIE_KERNEL_SOURCES can name another tree, and a checkout can carry uncommitted edits. None when neither directory exists.

Source code in utils/pyright_stubs/aie/utils/benchmark.py

def kernel_tree_digest() -> str | None:
    """Return a 12-hex digest of the kernel sources the library factories compile.

    Every file under ``aie_kernels_dir()`` and ``aie_runtime_lib_dir()``, by
    relative path and content. The commit alone cannot say which kernels
    ran: ``MLIR_AIE_KERNEL_SOURCES`` can name another tree, and a checkout
    can carry uncommitted edits. ``None`` when neither directory exists.
    """
    h = hashlib.sha256()
    found = False
    for top in (config.aie_kernels_dir(), config.aie_runtime_lib_dir()):
        root = Path(top)
        if not root.is_dir():
            continue
        found = True
        for path in sorted(p for p in root.rglob("*") if p.is_file()):
            h.update(f"{root.name}/{path.relative_to(root)}\0".encode())
            h.update(path.read_bytes())
    return h.hexdigest()[:12] if found else None

Static checks

These compiler-remark checks are available on demand through python -m aie.utils.compile.remarks; there is no static-check CI workflow. The trace-marker check (test/python/test_kernel_trace_markers.py) runs in lit on every PR, through trace_markers below.

Static kernel checks: compile every library kernel with Peano and read its remarks.

python -m aie.utils.compile.remarks --target aie2p --out static.json         --out-pm static-pm.json --meta static-meta.json
python -m aie.utils.compile.remarks --target aie2p --only '^gelu'         --out static.json --sources . --baseline-sources ../mlir-aie-base
python -m aie.utils.compile.remarks --target aie2p --only '^tanh/'         --cases test/python/npu/kernel_cases.py --out static.json
python -m aie.utils.compile.remarks --target aie2p --out static.json         --build cascade_mm:dim_m=16,dim_k=24,dim_n=32 --build cascade_mm

CPU-only. Every factory in aie.iron.kernels (at its defaults and for each entry of its kernels.DTYPES table) is compiled exactly as the JIT compiles it (aie.utils.compile.utils.cxx_core_compile_command), plus the optimization-record flags below, and the records become per-kernel series for benchmark-action: a Peano bump that changes a loop's schedule shows up here before anyone looks at device numbers. With --sources DIR (or MLIR_AIE_KERNEL_SOURCES=DIR) naming a checkout, that checkout's aie_kernels/ and aie_runtime_lib/ are compiled instead of the installed copies. With --cases, the builds are the ones a cases file's tests run (shape and options baked in), each named by its case; with --build, the factory builds named on the command line, for a shape that picks a code path no default reaches.

Record shapes, as llvm-aie 22.0.0.2026090201 emits them (they are Peano's, not LLVM's documented ones):

Pass	Kind / Name	Args	Tracked as
pipeliner	Passed / schedule	II, NS, Loop, Pipeliner, prologue/epilogue bundles	loop/<fn>/<bb>/II (rest as hover text)
pipeliner	Missed / canPipelineLoop	"Failed to pipeline loop"; located by DebugLoc only	unpipelined_loops (keyed L<line>)
pipeliner	Analysis / schedule	MII, SwpMaxMii, "Unable to find schedule"	schedule_notes in the meta file
aie-hardware-loops	Analysis / analysis	BasicBlock, Zero-Overhead-Loop	non_zol_loops, and loop/<fn>/<bb>/not_zol per loop
aie-asm-printer	Analysis / analysis	BasicBlock, BundleCount, ByteCount	pm_bytes (summed over the shipped functions)
aie-multi-slot-pseudo	Missed / missing-memory-bank	Instruction	missing_bank_loads
stderr	-Wpass-failed	a #pragma clang loop / AIE_* macro the compiler dropped	pass_failed_warnings, text kept
the object	llvm-readobj sections, symbols, relocations	what the entry symbol reaches	the shipped functions; libcalls (e.g. __divsf3)
the object	llvm-readobj --stack-sizes (-fstack-size-section)	frame bytes per function	kernel_stack_bytes, the deepest path from the entry
the IR (-emit-llvm)	opt print<scalar-evolution>	a constant backedge-taken count	loop/<fn>/<bb>/II_x_trips

The loop counts and pm_bytes cover only the functions the entry symbol reaches in the object, which are the ones the core link keeps. kernel_stack_bytes counts the kernel's frames only: the core that calls it also holds main's, which aiecc's measured stack size includes (dwconv1d_channels_last on AIE2P: 64 here, 256 for the core). --baseline-sources DIR compiles everything a second time from DIR and prints each row that differs, for a before/after of a kernel change; a loop LLVM only renamed, its rows unchanged, is counted but not listed, and a build either tree fails to compile is named and left out. It names the aie_kernels/ this tree compiled and warns when that is the installed copy (MLIR_AIE_KERNEL_SOURCES unset) or the baseline itself. --keep DIR keeps the objects, which --meta names per build.

The loop-scheduling pass reports as pipeliner (a postpipeliner filter records nothing); it names loops by machine basic block (bb.1.for.body.i) while the other passes use the IR block (for.body.i), so the prefix is stripped to join them. unpipelined_loops counts every loop the pipeliner declined, outer loops included, so its change is the signal, not its value.

The integer series alert on any increase; pm_bytes goes to its own file (--out-pm) so it can carry a percentage threshold. Nothing gates: under GitHub Actions a dropped pragma is a ::warning at its file and line, a kernel that fails to compile an ::error, and the run then exits 3 with "RESULTS INVALID", after writing the rows of the builds that compiled.

LoopInfo dataclass

LoopInfo(
    function: str,
    block: str,
    ii: int | None = None,
    ns: int | None = None,
    prologue_bundles: int | None = None,
    epilogue_bundles: int | None = None,
    pipelined: bool | None = None,
    pipeliner: str | None = None,
    missed_reason: str | None = None,
    zol: bool | None = None,
    bundle_count: int | None = None,
    byte_count: int | None = None,
    trips: int | None = None,
    file: str | None = None,
    line: int | None = None,
)

StaticReport dataclass

StaticReport(
    loops: dict[tuple[str, str], LoopInfo] = dict(),
    pm_bytes_by_function: dict[str, int] = dict(),
    missing_bank_loads: int = 0,
    pass_failed_warnings: int = 0,
    pass_failed: list[str] = list(),
    schedule_notes: list[str] = list(),
    shipped: set[str] | None = None,
    libcalls: list[str] = list(),
    kernel_stack_bytes: int | None = None,
)

parse_yaml

parse_yaml(
    path: str | Path, report: StaticReport | None = None
) -> StaticReport

Source code in utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/utils/compile/remarks.py

def report_rows(report: StaticReport, prefix: str, extra: str) -> list[dict]:
    """benchmark-action rows for one kernel build. Smaller is better throughout."""
    out = [
        _row(f"{prefix}/unpipelined_loops", "loops", report.unpipelined_loops, extra),
        _row(f"{prefix}/non_zol_loops", "loops", report.non_zol_loops, extra),
        _row(f"{prefix}/missing_bank_loads", "loads", report.missing_bank_loads, extra),
        _row(
            f"{prefix}/pass_failed_warnings",
            "warnings",
            report.pass_failed_warnings,
            extra,
        ),
        _row(f"{prefix}/pm_bytes", "bytes", report.pm_bytes, extra),
        _row(
            f"{prefix}/libcalls",
            "symbols",
            len(report.libcalls),
            extra,
            " ".join(report.libcalls) or None,
        ),
    ]
    if report.kernel_stack_bytes is not None:
        out.append(
            _row(
                f"{prefix}/kernel_stack_bytes",
                "bytes",
                report.kernel_stack_bytes,
                extra,
                "the kernel's frames; the core adds main's"
                + (" and the runtime routines'" if report.libcalls else ""),
            )
        )
    for (fn, bb), loop in sorted(report.loops.items()):
        if loop.ii is not None:
            # The hover text names the source line, so a reader of an II
            # alert can open the loop without decoding a basic-block name.
            where = (
                f" at {Path(loop.file).name}:{loop.line}"
                if loop.file and loop.line is not None
                else ""
            )
            out.append(
                _row(
                    f"{prefix}/loop/{fn}/{bb}/II",
                    "cycles",
                    loop.ii,
                    extra,
                    f"NS={loop.ns} pro={loop.prologue_bundles} "
                    f"epi={loop.epilogue_bundles} zol={loop.zol} "
                    f"trips={loop.trips} via={loop.pipeliner}{where}",
                )
            )
            if loop.trips is not None:
                # II alone reads a loop that does 4 blocks per iteration as
                # slower than one that does 1; II x trips compares them.
                out.append(
                    _row(
                        f"{prefix}/loop/{fn}/{bb}/II_x_trips",
                        "cycles",
                        loop.ii * loop.trips,
                        extra,
                        f"{loop.trips} trips per entry{where}",
                    )
                )
        if loop.zol is not None:
            # 1 when the loop is not a zero-overhead loop: an inner loop that
            # falls off the hardware loop unit costs its bundle count in
            # branches, whatever its II says.
            out.append(
                _row(
                    f"{prefix}/loop/{fn}/{bb}/not_zol",
                    "loops",
                    int(not loop.zol),
                    extra,
                )
            )
    return out

workflow_annotations

workflow_annotations(
    name: str,
    report: StaticReport | None,
    detail: str,
    root: str | None = None,
    seen: set[tuple] | None = None,
) -> list[str]

Annotations for one kernel build: its dropped pragmas, or its compile failure.

One source is compiled once per factory build and per target, so the same dropped pragma comes back several times a run; seen (shared across builds) keeps each file, line and message to its first annotation.

Source code in utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/utils/compile/remarks.py

def analyze(ext_fn, target: str, workdir: Path) -> tuple[StaticReport | None, str]:
    """Compile one kernel and parse its records; ``(None, reason)`` when it fails to compile."""
    cmd, yaml_out = compile_command(ext_fn, target, workdir)
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        return None, f"compile failed: {compile_failure(p.stderr)}"
    rep = parse_yaml(yaml_out) if yaml_out.exists() else StaticReport()
    for key, trips in trip_counts(ext_fn, target, workdir).items():
        if key in rep.loops:
            rep.loops[key].trips = trips
    parse_stderr(p.stderr, rep)
    reached = linked(workdir / f"{ext_fn.name}.o", entry_symbol(ext_fn))
    rep.shipped, rep.libcalls = reached.functions, reached.undefined
    rep.kernel_stack_bytes = reached.stack
    return rep, "ok"

kernel_builds

kernel_builds()

Yield (name, ExternalFunction) for every factory build the library offers.

Each exported factory at its defaults, plus one build per entry of its kernels.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 utils/pyright_stubs/aie/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
    ``kernels.DTYPES`` table; a factory that refuses the current device
    (``NotImplementedError``) is skipped. Remarks depend on source and flags,
    not on the shape a test runs, so this is the whole surface.
    """
    from aie.iron import kernels

    for name in kernels.factories():
        f = getattr(kernels, name)
        combos = [{}] + [dict(c) for c in kernels.DTYPES.get(name, ()) if c]
        seen: set[str] = set()
        for combo in combos:
            try:
                ef = f(**combo)
            except NotImplementedError:
                continue  # exists only for the other architecture
            if ef.object_file_name in seen:
                continue  # the default build is one of the dtypes entries
            seen.add(ef.object_file_name)
            yield _build_name(name, combo), ef

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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/utils/compile/remarks.py

def trace_shape(ext_fn, target: str, out_dir: Path) -> str:
    """Compile ``ext_fn`` to -O2 IR and classify its entry's markers (``trace_markers``)."""
    src, include_dirs = _kernel_file(ext_fn, out_dir)
    cmd = cxx_core_compile_command(
        str(src),
        target,
        str(out_dir / f"{ext_fn.name}.ll"),
        include_dirs=include_dirs,
        compile_args=list(ext_fn.compile_flags),
        inline=True,
    )
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        raise RuntimeError(f"{ext_fn.name}: {compile_failure(p.stderr)}")
    ir = (out_dir / f"{ext_fn.name}.ll").read_text()
    return trace_markers(ir, entry_symbol(ext_fn))

entry_symbol

entry_symbol(ext_fn) -> str

Return the symbol the kernel source defines, before the JIT's per-build prefix.

Source code in utils/pyright_stubs/aie/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)

What a linked kernel keeps, read from its object with llvm-readobj.

Peano gives every function its own section (.text.<fn>) and the core link drops unreferenced sections, so what ships is what the entry symbol reaches through relocations. A helper that is inlined at its call sites and also emitted standalone has asm-printer records for both; counting only the reached functions keeps it from counting twice. A relocation from a reached section to an undefined symbol is a call into the runtime library -- on AIE2P that is scalar float divide, int-to-float and the like, each a software routine of hundreds of cycles.

Linked dataclass

Linked(
    functions: set[str],
    undefined: list[str],
    stack: int | None = None,
    data: int = 0,
    reached_data: int = 0,
    linked_bytes: int = 0,
)

linked

linked(obj: Path, entry: str) -> Linked

Functions and undefined symbols entry reaches in the object obj.

Source code in utils/pyright_stubs/aie/utils/compile/readobj.py

def linked(obj: Path, entry: str) -> Linked:
    """Functions and undefined symbols ``entry`` reaches in the object ``obj``."""
    out = subprocess.run(
        [
            config.readobj_path(),
            "--elf-output-style=JSON",
            "--sections",
            "--symbols",
            "--relocations",
            "--stack-sizes",
            str(obj),
        ],
        capture_output=True,
        text=True,
        check=True,
    ).stdout
    return parse_readobj(json.loads(out)[0], entry)

parse_readobj

parse_readobj(doc: dict, entry: str) -> Linked

linked on the parsed llvm-readobj --elf-output-style=JSON document.

Source code in utils/pyright_stubs/aie/utils/compile/readobj.py

def parse_readobj(doc: dict, entry: str) -> Linked:
    """``linked`` on the parsed ``llvm-readobj --elf-output-style=JSON`` document."""
    alloc, stored, data_sections = {}, {}, []
    for s in doc["Sections"]:
        s = s["Section"]
        flags = {f["Name"] for f in s["Flags"]["Flags"]}
        alloc[s["Index"]] = "SHF_ALLOC" in flags
        stored[s["Index"]] = 0 if s["Type"]["Name"] == "SHT_NOBITS" else s["Size"]
        if "SHF_ALLOC" in flags and "SHF_EXECINSTR" not in flags:
            data_sections.append(
                (s["Index"], s["Size"], s["AddressAlignment"], s["Name"]["Name"])
            )
    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"]
    }

    depth: dict[int, int | None] = {}

    def deepest(sec: int) -> int | None:
        if sec in depth:
            return depth[sec]
        # None until the section finishes, so reaching it again is recursion.
        depth[sec] = 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)
                if d is None:
                    return None
                below = max(below, d)
        depth[sec] = (
            max((frames.get(n, 0) for n in functions.get(sec, ())), default=0) + below
        )
        return depth[sec]

    def laid_out(sections) -> int:
        # As aiecc's generated link script (AIETargetLdScript.cpp): every
        # .data* section, then every .rodata*, then every .bss*, in that
        # output-section order regardless of the object's own section-table
        # order; each from the next multiple of its alignment, as aiecc's
        # measureDataSectionDemand does over the linked ELF.
        def group(name: str) -> int:
            for i, prefix in enumerate(_DATA_SECTION_GROUPS):
                if name.startswith(prefix):
                    return i
            return len(_DATA_SECTION_GROUPS)

        end = 0
        for _, size, align, _name in sorted(sections, key=lambda s: group(s[3])):
            if size:
                end = -(-end // max(align, 1)) * max(align, 1) + size
        return end

    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]) if frames and roots else None,
        data=laid_out(data_sections),
        reached_data=laid_out(s for s in data_sections if s[0] in seen),
        linked_bytes=sum(stored[sec] for sec in seen),
    )

Host-side helpers

aie.utils.aie2p_emulation provides the arithmetic and runtime-library tables used by FastFlowLM host references. The q4nx block defaults (Q4NX_M_TILE, Q4NX_K_TILE, Q4NX_GROUP, Q4NX_BLOCK_BYTES) and Gemma 4 constants (RMS_EPS, BF16_PROJ_M, BF16_PROJ_K, DECODE_KEYS_PER_ROUND) are exported by aie.iron.kernels. quant.q4nx_unpack returns codes in (k, m) order.

Host arithmetic for FastFlowLM's AIE2P references.

Bf16 results are float32 values unless a function returns uint16 bits. Rounding modes are floor and conv_even, as in aie.utils.bfp. The limb products and reductions use IEEE float32 accumulation; these reference approximations can differ from device arithmetic near rounding boundaries. aie.utils.aie_fp32.aie_add models the scalar accumulator adder.

bf16_floor_bits

bf16_floor_bits(x)

float32 x rounded to bf16 toward minus infinity, preserving NaNs, as uint16 bits.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

def bf16_floor_bits(x):
    """float32 ``x`` rounded to bf16 toward minus infinity, preserving NaNs, as uint16 bits."""
    u = np.ascontiguousarray(x, dtype=np.float32).view(np.uint32)
    inexact = (u & 0xFFFF) != 0
    negative = (u >> 31) != 0
    bits = (u >> 16) + (inexact & negative)
    # Retain the NaN sign and high payload bits, and set the bf16 quiet bit.
    return np.where((u & 0x7FFFFFFF) > 0x7F800000, (u >> 16) | 0x40, bits).astype(
        np.uint16
    )

bf16_floor

bf16_floor(x)

float32 x rounded to bf16 toward minus infinity, as float32 values.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

def bf16_floor(x):
    """float32 ``x`` rounded to bf16 toward minus infinity, as float32 values."""
    return _from_bits(bf16_floor_bits(x))

f32

f32(x)

Round x to IEEE float32, returned as float32 values.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

def f32(x):
    """Round ``x`` to IEEE float32, returned as float32 values."""
    return np.asarray(x, np.float64).astype(np.float32)

round_bf16_bits

round_bf16_bits(x, rounding='floor')

x rounded to fp32, then to bf16 bit patterns (uint16).

rounding is "floor" (toward minus infinity, a core's default) or "conv_even" (round to nearest even). The first rounding models the fp32 register that holds a value before a core narrows it.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

def round_bf16_bits(x, rounding="floor"):
    """``x`` rounded to fp32, then to bf16 bit patterns (uint16).

    ``rounding`` is ``"floor"`` (toward minus infinity, a core's default) or
    ``"conv_even"`` (round to nearest even). The first rounding models the fp32
    register that holds a value before a core narrows it.
    """
    if rounding not in _ROUNDINGS:
        raise ValueError(f"rounding must be 'floor' or 'conv_even', got {rounding!r}")
    if rounding == "conv_even":
        return f32(x).astype(bfloat16).view(np.uint16)
    return bf16_floor_bits(f32(x))

round_bf16

round_bf16(x, rounding='floor')

x rounded to fp32, then to bf16, as float32 values.

See round_bf16_bits for rounding.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

def round_bf16(x, rounding="floor"):
    """``x`` rounded to fp32, then to bf16, as float32 values.

    See [`round_bf16_bits`][aie.utils.aie2p_emulation.round_bf16_bits] for
    ``rounding``.
    """
    return _from_bits(round_bf16_bits(x, rounding))

fmul

fmul(a, b)

Approximate an fp32 product by accumulating its nine bf16 limb products.

Products accumulate largest first with IEEE float32 rounding.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

def fmul(a, b):
    """Approximate an fp32 product by accumulating its nine bf16 limb products.

    Products accumulate largest first with IEEE float32 rounding.
    """
    from aie.iron.kernels.datamovement import limbs_f32_split

    a, b = np.broadcast_arrays(np.asarray(a, np.float64), np.asarray(b, np.float64))
    A, B = ([limb.astype(np.float32) for limb in limbs_f32_split(v)] for v in (a, b))
    acc = None
    for i, j in (
        (0, 0),
        (0, 1),
        (1, 0),
        (0, 2),
        (1, 1),
        (2, 0),
        (1, 2),
        (2, 1),
        (2, 2),
    ):
        p = A[i].astype(np.float64) * B[j]
        acc = f32(p) if acc is None else f32(acc + p)
    assert acc is not None
    return acc

tree_sum

tree_sum(x, lanes)

Sum the last axis in lanes lanes with IEEE float32 rounding.

Each lane keeps a running sum of every lanes-th value. A pairwise halving of the lanes then adds the lane sums. The last axis must be a multiple of lanes, and lanes a power of two.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

def tree_sum(x, lanes):
    """Sum the last axis in ``lanes`` lanes with IEEE float32 rounding.

    Each lane keeps a running sum of every ``lanes``-th value. A pairwise
    halving of the lanes then adds the lane sums. The last axis must be a
    multiple of ``lanes``, and ``lanes`` a power of two.
    """
    x = np.asarray(x, np.float64)
    xs = x.reshape(x.shape[:-1] + (-1, lanes))
    acc = f32(xs[..., 0, :])
    for i in range(1, xs.shape[-2]):
        acc = f32(acc + xs[..., i, :])
    while acc.shape[-1] > 1:
        h = acc.shape[-1] // 2
        acc = f32(acc[..., :h] + acc[..., h:])
    return acc[..., 0]

fast_rsqrt

fast_rsqrt(s)

1 / sqrt(s) in fp32: the 0x5f3759df seed, then two Newton steps in fmul.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

def fast_rsqrt(s):
    """``1 / sqrt(s)`` in fp32: the ``0x5f3759df`` seed, then two Newton steps in ``fmul``."""
    s = f32(s)
    half = fmul(s, np.float32(0.5))
    y = (
        (np.uint32(0x5F3759DF) - (s.view(np.uint32) >> 1))
        .astype(np.uint32)
        .view(np.float32)
    )
    for _ in range(2):
        y = fmul(y, f32(np.float32(1.5) - fmul(fmul(half, y), y)))
    return y

inv_bf16

inv_bf16(x)

getInvBf16: 1 / x as bf16, from the exponent and INV_MANTISSA.

x must be positive. The function rounds x to the nearest bf16 first. The relative error is up to 0.4% for a bf16 x and up to 0.6% for an fp32 x.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

def inv_bf16(x):
    """``getInvBf16``: ``1 / x`` as bf16, from the exponent and ``INV_MANTISSA``.

    ``x`` must be positive. The function rounds ``x`` to the nearest bf16
    first. The relative error is up to 0.4% for a bf16 ``x`` and up to 0.6%
    for an fp32 ``x``.
    """
    bits = f32(x).view(np.uint32).astype(np.uint64) + 0x8000
    exponent = (bits & 0x7F800000) >> 23
    mantissa = (bits & 0x007FFFFF) >> 16
    inv_exp = (mantissa == 0).astype(np.uint64) + (253 - exponent)
    return _from_bits(
        (((inv_exp << 7) + INV_MANTISSA[mantissa]) & 0xFFFF).astype(np.uint16)
    )

gelu_lut_segments cached

gelu_lut_segments()

getGeluBf16's table: 64 (slope, offset) rows, one per 1/8 of [-4, 4).

The function reads gelu_lut_ab from the aie_runtime_lib sources that the kernels compile against. No simple fit of GELU reproduces the table. The source holds each run of four segments twice, for the gather read.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

@cache
def gelu_lut_segments():
    """``getGeluBf16``'s table: 64 (slope, offset) rows, one per 1/8 of [-4, 4).

    The function reads ``gelu_lut_ab`` from the aie_runtime_lib sources that the
    kernels compile against. No simple fit of GELU reproduces the table. The
    source holds each run of four segments twice, for the gather read.
    """
    path = Path(config.aie_runtime_lib_dir()) / "AIE2P" / "lut_based_ops.cpp"
    table = re.search(r"gelu_lut_ab\[\d+\]\s*=\s*\{([^}]*)\}", path.read_text())
    if table is None:
        raise ValueError(f"{path} has no gelu_lut_ab table")
    body = table[1]
    values = [float(v.strip().rstrip("f")) for v in body.split(",") if v.strip()]
    return np.array(values, np.float32).reshape(-1, 2, 8)[:, 0].reshape(-1, 2)

gelu_bf16

gelu_bf16(x)

getGeluBf16: GELU of x from gelu_lut_segments, as bf16.

The device reads the slope as bf16 and the offset as fp32. Inputs outside [-4, 4) take the end segments.

Source code in utils/pyright_stubs/aie/utils/aie2p_emulation.py

def gelu_bf16(x):
    """``getGeluBf16``: GELU of ``x`` from ``gelu_lut_segments``, as bf16.

    The device reads the slope as bf16 and the offset as fp32. Inputs outside
    [-4, 4) take the end segments.
    """
    x = np.asarray(x, np.float64)
    k = np.clip(np.floor(x * 128), -512, 511)
    # Any segment works for NaN: the arithmetic below propagates it.
    k = np.nan_to_num(k, nan=0).astype(np.int64) >> 4
    pair = gelu_lut_segments()[k + 32]
    slope = (pair[..., 0].view(np.uint32) & 0xFFFF0000).view(np.float32)
    return round_bf16(f32(slope.astype(np.float64) * x + pair[..., 1]))

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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/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 utils/pyright_stubs/aie/utils/verify.py

def compare(
    actual,
    expected,
    tol: Tolerance | None = None,
    *,
    range_axis: int | None = None,
    bound=None,
) -> Verdict:
    """Compare a kernel's ``actual`` output with a reference under ``tol``.

    ``expected`` may be higher precision than ``actual`` (a float64 sum, an
    int64 product); it is cast to ``actual.dtype``, so the kernel is held to
    what a correctly rounded implementation would produce. With ``tol=None``
    the output dtype's ``Tolerance.default_for`` applies.

    ``range_axis`` selects the axis reduced to compute ``range_frac``'s
    reference scale. For ``(calls, tile)`` arrays, use 1 to scale each call
    independently. The default uses the whole reference array.

    ``bound`` is a **bound** tolerance's per-element limit, already
    evaluated on the inputs and broadcastable to ``expected``.

    This measures; it does not model. What a kernel does on overflow, on a
    narrowing store, or with subnormal inputs belongs in the reference that
    produced ``expected`` -- a saturating kernel's reference clips, a
    denormal-flushing kernel's reference flushes. A reference that leaves the
    output range is reported as such when the comparison fails, since it
    means the reference is under-specified rather than the kernel wrong.
    """
    actual = np.asarray(actual)
    expected = np.asarray(expected)
    if tol is None:
        tol = Tolerance.default_for(actual.dtype)
    if actual.shape != expected.shape:
        return Verdict(
            False,
            0,
            0,
            float("inf"),
            None,
            None,
            f"shape mismatch {actual.shape} vs {expected.shape}",
        )
    if tol.kind == "bound":
        if bound is None:
            raise ValueError(
                "a bound tolerance needs its bound evaluated on the inputs; "
                "pass bound=, or judge through the kernel with inputs="
            )
        if not np.issubdtype(actual.dtype, np.floating) and actual.dtype != bfloat16:
            raise ValueError(
                f"a bound tolerance is for floating-point outputs, got {actual.dtype}"
            )
    a, e, n = actual.ravel(), expected.ravel(), actual.size

    def ref_scale(ref) -> float | np.ndarray:
        """Return max|ref| over finite entries.

        This is what ``range_frac`` is a fraction of, or zero when no range
        floor is in play.
        """
        if tol.range_frac is None or not n:
            return 0.0
        mag = np.abs(np.asarray(ref, dtype=np.float64).reshape(expected.shape))
        mag = np.where(np.isfinite(mag), mag, 0.0)
        scale = mag.max(axis=range_axis, keepdims=True, initial=0.0)
        if range_axis is None:
            return float(scale.item())
        return np.broadcast_to(scale, expected.shape).ravel()

    # Integers: bit-exact under exact / ulps; under a relative tolerance the
    # nearly_equal formula in int64 (Tolerance.lsb sets atol = n + 0.5).
    if np.issubdtype(actual.dtype, np.integer) or actual.dtype == np.bool_:
        n_over = 0
        if actual.dtype != np.bool_ and np.issubdtype(e.dtype, np.integer):
            info = np.iinfo(actual.dtype)
            e_wide = e.astype(np.int64)
            n_over = int(np.count_nonzero((e_wide < info.min) | (e_wide > info.max)))
        e_cast = e.astype(actual.dtype)
        a64, e64 = a.astype(np.int64), e_cast.astype(np.int64)
        err = np.abs(a64 - e64)
        scale = ref_scale(e64)
        if tol.kind == "relative":
            bound = np.maximum(
                tol.atol or 0.0,
                (tol.rtol or 0.0) * (np.abs(a64) + np.abs(e64)),
            )
            close = err < bound
            if tol.range_frac is not None:
                close |= err <= tol.range_frac * scale
            bad = ~close
        else:
            bad = a != e_cast
        v = _verdict(bad, err, None, tol, n, ref_range=scale)
        if n_over and not v.ok:
            # Not a policy, a diagnostic: the reference left the output range,
            # so what the device did there says nothing about the kernel.
            v = replace(
                v,
                detail=f"{v.detail}; the reference overflows "
                f"{np.dtype(actual.dtype).name} in {n_over} of {n} elements, so "
                "it does not model what the kernel does there (clip for a "
                "saturating kernel, cast for a wrapping one)",
            )
        return v

    a32, e32 = a.astype(np.float32), e.astype(np.float32)
    a_nan, e_nan = np.isnan(a32), np.isnan(e32)
    a_inf, e_inf = np.isinf(a32), np.isinf(e32)
    nonfinite_bad = (a_nan != e_nan) | (a_inf != e_inf) | (a_inf & e_inf & (a32 != e32))
    finite = ~(a_nan | e_nan | a_inf | e_inf)
    err = np.zeros(n, np.float64)

    if tol.kind == "ulps":
        if actual.dtype != bfloat16:
            raise ValueError(
                f"Tolerance in ULPs is defined for bfloat16 outputs, got {actual.dtype}"
            )
        # A cast to bfloat16 goes through float32 and can round twice;
        # round_to rounds the reference once, as a correctly rounded bf16
        # implementation would.
        e_bf = round_to(e, bfloat16)
        ulp = np.zeros(n, np.int64)
        ulp[finite] = bf16_ulp_distance(a[finite], e_bf[finite])
        err[finite] = np.abs(a32[finite] - e_bf[finite].astype(np.float32))
        max_ulps = tol.ulps if tol.ulps is not None else 0
        within = ulp <= max_ulps
        scale = ref_scale(e_bf.astype(np.float32))
        if tol.atol is not None:
            within |= err < tol.atol
        if tol.range_frac is not None:
            within |= err <= tol.range_frac * scale
        bad = nonfinite_bad | (finite & ~within)
        return _verdict(bad, err, ulp, tol, n, nonfinite_bad, ref_range=scale)

    if tol.kind == "bound":
        limit = np.broadcast_to(np.asarray(bound, np.float64), expected.shape).ravel()
        err[finite] = np.abs(a32[finite].astype(np.float64) - e32[finite])
        bad = nonfinite_bad | (finite & ~(err <= limit))
        return _verdict(bad, err, None, tol, n, nonfinite_bad, limit=limit)

    if tol.kind == "exact":
        e_cast = e32.astype(actual.dtype).astype(np.float32)
        err[finite] = np.abs(a32[finite] - e_cast[finite])
        bad = nonfinite_bad | (finite & (a32 != e_cast))
        return _verdict(bad, err, None, tol, n, nonfinite_bad)

    err[finite] = np.abs(a32[finite].astype(np.float64) - e32[finite])
    scale = ref_scale(e32)
    close = nearly_equal(a32, e32, rtol=tol.rtol or 0.0, atol=tol.atol)
    if tol.range_frac is not None:
        close |= err <= tol.range_frac * scale
    bad = nonfinite_bad | (finite & ~close)
    return _verdict(bad, err, None, tol, n, nonfinite_bad, ref_range=scale)

bf16_ulp_distance

bf16_ulp_distance(a, b) -> ndarray

Element-wise distance between two bf16 arrays in units in the last place.

Bit patterns are mapped to a monotonic integer scale (sign-magnitude to two's-complement style) so the distance is a plain subtraction. -0 and +0 map to the same point, so a kernel that produces the other zero is not penalized. aie.utils.accuracy.ulp_distance takes other dtypes.

Source code in utils/pyright_stubs/aie/utils/verify.py

def bf16_ulp_distance(a, b) -> np.ndarray:
    """Element-wise distance between two bf16 arrays in units in the last place.

    Bit patterns are mapped to a monotonic integer scale (sign-magnitude to
    two's-complement style) so the distance is a plain subtraction. -0 and +0
    map to the same point, so a kernel that produces the other zero is not
    penalized. ``aie.utils.accuracy.ulp_distance`` takes other dtypes.
    """
    return ulp_distance(a, b, bfloat16)

Sampling

sample_select and sample_combine draw the next token from a row of bf16 logits on the device: each sample_select core reduces one slice to its top k, and one sample_combine core draws from those summaries. sample_ref is the bit-exact host reference.

Top-k next-token sampling on the device, and its bit-exact host reference.

A row of bf16 logits is split into columns of slice_size. Each column's sample_select core reduces its slice to a summary (its top k, as entries above its k-th largest value and a bitmap of the ties at it); one sample_combine core draws the token from all the summaries. The draw is sample_ref:

1. Temperature 0 (or -0): the first index of the largest logit.
2. tau is the k-th largest logit, with multiplicity, compared on order_keys (-0 and +0 are one key). The candidates are the logits >= tau, in index order.
3. Each candidate value v weighs exp64_ref(float64(fl32(fl32(v / T) - fl32(max / T)))).
4. u = n53 * 2**-53, integer 0 <= n53 < 2**53. The token is the first candidate, in index order, whose exact prefix sum of the weights P exceeds u * S, S the exact total: P * 2**53 > n53 * S in integers. That is the inverse-CDF draw over the weights, with no rounding after them.

Every weight step is an IEEE float32 division or subtraction, or a float64 addition, subtraction, multiplication or comparison, never fused; the core computes the float32 steps with the soft-float builtins and float64 has only those, so numpy reproduces the device's bits. Each weight is an integer multiple of 2**-1074, so the core sums them exactly in fixed point. A temperature must satisfy check_order_preserving, which makes the threshold on bf16 keys select what a threshold on v / T would.

A draw is passed to the device as a four-word int32 row, draw_row.

summary_words

summary_words(slice_size: int, k_max: int) -> int

int32 words of one column's summary: header, entries, tie bitmap.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def summary_words(slice_size: int, k_max: int) -> int:
    """int32 words of one column's summary: header, entries, tie bitmap."""
    return SUMMARY_HEADER + 2 * k_max + (slice_size + 31) // 32

rule_column_words

rule_column_words(slice_size: int, k_max: int) -> int

int32 words of one column's input to sample_combine(probability_rule=True).

That column's sample_select summary, restricted summary and reduction, back to back.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def rule_column_words(slice_size: int, k_max: int) -> int:
    """int32 words of one column's input to ``sample_combine(probability_rule=True)``.

    That column's ``sample_select`` summary, restricted summary and
    reduction, back to back.
    """
    return 2 * summary_words(slice_size, k_max) + REDUCTION_WORDS

select_row_words

select_row_words(
    max_ranges: int = 0, probability_rule: bool = False
) -> int

int32 words of sample_select's row.

draw_row's ROW_WORDS, then 2 * max_ranges range words, then one split word if probability_rule. 0, False (the default) is ROW_WORDS.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def select_row_words(max_ranges: int = 0, probability_rule: bool = False) -> int:
    """int32 words of ``sample_select``'s row.

    ``draw_row``'s ``ROW_WORDS``, then ``2 * max_ranges`` range words, then
    one split word if ``probability_rule``. ``0, False`` (the default) is
    ``ROW_WORDS``.
    """
    _non_negative("select_row_words", max_ranges=max_ranges)
    return ROW_WORDS + 2 * max_ranges + (1 if probability_rule else 0)

combine_row_words

combine_row_words(probability_rule: bool = False) -> int

int32 words of sample_combine's row.

draw_row's ROW_WORDS, then one split word if probability_rule.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def combine_row_words(probability_rule: bool = False) -> int:
    """int32 words of ``sample_combine``'s row.

    ``draw_row``'s ``ROW_WORDS``, then one split word if ``probability_rule``.
    """
    return ROW_WORDS + (1 if probability_rule else 0)

range_words

range_words(ranges, max_ranges: int) -> ndarray

Pack up to max_ranges (lo, hi) pairs as sample_select's row words.

Appended after draw_row's ROW_WORDS: global row coordinates, one pair per range, in order. An unused slot is (0, 0), which excludes nothing (lo >= hi, once clipped to a column's own slice, is empty).

Raises:

Type	Description
ValueError	When more than max_ranges ranges are given.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def range_words(ranges, max_ranges: int) -> np.ndarray:
    """Pack up to ``max_ranges`` ``(lo, hi)`` pairs as ``sample_select``'s row words.

    Appended after ``draw_row``'s ``ROW_WORDS``: global row coordinates, one
    pair per range, in order. An unused slot is ``(0, 0)``, which excludes
    nothing (``lo >= hi``, once clipped to a column's own slice, is empty).

    Raises:
        ValueError: When more than ``max_ranges`` ranges are given.
    """
    _non_negative("range_words", max_ranges=max_ranges)
    ranges = list(ranges)
    if len(ranges) > max_ranges:
        raise ValueError(
            f"range_words: {len(ranges)} ranges exceeds max_ranges ({max_ranges})"
        )
    words = np.zeros(2 * max_ranges, dtype=np.int32)
    for i, (lo, hi) in enumerate(ranges):
        words[2 * i] = lo
        words[2 * i + 1] = hi
    return words

split_word

split_word(split: int) -> ndarray

sample_select's and sample_combine's one extra row word.

Present when probability_rule: timestamp_begin, global row coordinates.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def split_word(split: int) -> np.ndarray:
    """``sample_select``'s and ``sample_combine``'s one extra row word.

    Present when ``probability_rule``: ``timestamp_begin``, global row
    coordinates.
    """
    return np.array([split], dtype=np.int32)

select_streams

select_streams(slice_size: int, chunk: int) -> int

How often sample_select takes its slice per position.

Once when a chunk is the whole slice, which one call passes over SELECT_PASSES times; SELECT_PASSES times otherwise.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def select_streams(slice_size: int, chunk: int) -> int:
    """How often ``sample_select`` takes its slice per position.

    Once when a chunk is the whole slice, which one call passes over
    ``SELECT_PASSES`` times; ``SELECT_PASSES`` times otherwise.
    """
    return 1 if chunk == slice_size else SELECT_PASSES

sample_select

sample_select(
    *,
    slice_size=32064,
    chunk=5344,
    k_max=64,
    valid=None,
    base=0,
    max_ranges=0,
    probability_rule=False
) -> ExternalFunction

One column's half of sampling: its slice's summary, for sample_combine.

sample_select(x, row, state, summary) takes chunk bf16 logits per call; the slice arrives select_streams(slice_size, chunk) times, so a position is that many times slice_size // chunk calls with the same row and summary. state is SELECT_STATE_WORDS int32 of worker-local memory, zero before the first call; the last call of a position leaves it ready for the next. summary is summary_words(slice_size, k_max) int32. valid (default slice_size) is how many of the slice's logits are tokens'; the rest pad it to whole words, and the kernel writes them -inf in x so none is drawn.

base is this column's global offset into the row, which row's ranges (max_ranges) and split (probability_rule) are given in; it does not change row's own shape. max_ranges (default 0: none) accepts that many (lo, hi) pairs per call (range_words, appended to row after draw_row's ROW_WORDS), each masking [lo, hi) (clamped to the row) to -inf before either pass reads it, as padding is: while any unmasked logit is finite, no masked one is drawn.

probability_rule (default off, bit-identical to a build without it) additionally takes row's one split word (split_word, after the ranges) and five more arguments: restricted_state/ restricted_summary (mirroring state/summary, a second top-k built over [split, vocab) only), restricted_x (this column's own chunk-wide bf16 scratch) and reduction (REDUCTION_WORDS int32, this column's running evidence for sample_combine's forced-timestamp compare). sample_combine (probability_rule=True) reads summary, restricted_summary and reduction back to back (rule_column_words), reduces every column's reduction into that compare and draws from summary or restricted_summary as it decides.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def sample_select(
    *,
    slice_size=32064,
    chunk=5344,
    k_max=64,
    valid=None,
    base=0,
    max_ranges=0,
    probability_rule=False,
) -> ExternalFunction:
    """One column's half of sampling: its slice's summary, for ``sample_combine``.

    ``sample_select(x, row, state, summary)`` takes ``chunk`` bf16 logits per
    call; the slice arrives ``select_streams(slice_size, chunk)`` times, so a
    position is that many times ``slice_size // chunk`` calls with the same
    ``row`` and ``summary``.
    ``state`` is ``SELECT_STATE_WORDS`` int32 of worker-local memory, zero
    before the first call; the last call of a position leaves it ready for the
    next. ``summary`` is ``summary_words(slice_size, k_max)`` int32.
    ``valid`` (default ``slice_size``) is how many of the slice's logits are
    tokens'; the rest pad it to whole words, and the kernel writes them -inf
    in ``x`` so none is drawn.

    ``base`` is this column's global offset into the row, which ``row``'s
    ranges (``max_ranges``) and split (``probability_rule``) are given in; it
    does not change ``row``'s own shape. ``max_ranges`` (default 0: none)
    accepts that many ``(lo, hi)`` pairs per call (``range_words``, appended
    to ``row`` after ``draw_row``'s ``ROW_WORDS``), each masking ``[lo, hi)``
    (clamped to the row) to -inf before either pass reads it, as padding is:
    while any unmasked logit is finite, no masked one is drawn.

    ``probability_rule`` (default off, bit-identical to a build without it)
    additionally takes ``row``'s one split word (``split_word``, after the
    ranges) and five more arguments: ``restricted_state``/
    ``restricted_summary`` (mirroring ``state``/``summary``, a second top-k
    built over ``[split, vocab)`` only), ``restricted_x`` (this column's own
    ``chunk``-wide bf16 scratch) and ``reduction`` (``REDUCTION_WORDS``
    int32, this column's running evidence for ``sample_combine``'s
    forced-timestamp compare). ``sample_combine`` (``probability_rule=True``)
    reads ``summary``, ``restricted_summary`` and ``reduction`` back to back
    (``rule_column_words``), reduces every column's ``reduction`` into that
    compare and draws from ``summary`` or ``restricted_summary`` as it
    decides.
    """
    _check_slice("sample_select", slice_size, k_max)
    _positive("sample_select", chunk=chunk)
    _non_negative("sample_select", base=base, max_ranges=max_ranges)
    valid = slice_size if valid is None else valid
    if not k_max <= valid <= slice_size:
        raise ValueError(
            f"sample_select: valid ({valid}) must be from k_max ({k_max}) to "
            f"slice_size ({slice_size})"
        )
    if slice_size % chunk:
        raise ValueError("sample_select: chunk must divide slice_size")
    if chunk % 2:
        raise ValueError("sample_select: chunk must be even (4-byte DMA)")
    row_words = select_row_words(max_ranges, probability_rule)
    summary_shape = (summary_words(slice_size, k_max),)
    arg_types = [
        np.ndarray[(chunk,), np.dtype[bfloat16]],
        np.ndarray[(row_words,), np.dtype[np.int32]],
        np.ndarray[(SELECT_STATE_WORDS,), np.dtype[np.int32]],
        np.ndarray[summary_shape, np.dtype[np.int32]],
    ]
    # x is written in place: by the padding fill below valid, and by
    # apply_ranges (sample_select.cc) masking a range's span to -inf bits,
    # whether or not padding also applies.
    roles = [In if valid == slice_size and not max_ranges else InOut, In, InOut, InOut]
    if probability_rule:
        arg_types += [
            np.ndarray[(SELECT_STATE_WORDS,), np.dtype[np.int32]],
            np.ndarray[summary_shape, np.dtype[np.int32]],
            np.ndarray[(chunk,), np.dtype[bfloat16]],
            np.ndarray[(REDUCTION_WORDS,), np.dtype[np.int32]],
        ]
        roles += [InOut, InOut, InOut, InOut]
    compile_flags = list(
        _STRICT_FP
        + [
            f"-DSAMPLE_SLICE={slice_size}",
            f"-DSAMPLE_CHUNK={chunk}",
            f"-DSAMPLE_K_MAX={k_max}",
            f"-DSAMPLE_VALID={valid}",
        ]
    )
    if base:
        compile_flags.append(f"-DSAMPLE_BASE={base}")
    if max_ranges:
        compile_flags.append(f"-DSAMPLE_MAX_RANGES={max_ranges}")
    if probability_rule:
        compile_flags.append("-DSAMPLE_PROB_RULE=1")
    return _make_extern(
        "sample_select",
        _kernel_source("sample/sample_select.cc"),
        arg_types,
        compile_flags=compile_flags,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=tuple(roles),
            ops_per_call=0,
            # A position's calls are not equal work: the first of each pass
            # (begin(), zeroing the bitmap, and a cut=0 scan keeping every
            # key) and the last (end_pass()) do more than the steady-state
            # calls between them, and pass 0 (scan) and pass 1 (collect)
            # are different code.
            even_calls=False,
        ),
    )

sample_combine

sample_combine(
    *,
    columns: int = 4,
    slice_size: int = 32064,
    k_max: int = 64,
    last_slice: int | None = None,
    probability_rule: bool = False
) -> ExternalFunction

Draw from columns summaries of sample_select, as sample_ref does.

sample_combine(summaries, row, token, record) reads the summaries column 0 first, and writes the drawn token (an index into the whole row of columns * slice_size) to both one-word outputs, so a design can send it two places.

Parameters:

Name	Type	Description	Default
columns	int	Summaries to combine, one per sample_select column.	4
slice_size	int	Each column's slice of the row.	32064
k_max	int	The largest top-k the summaries carry.	64
last_slice	int | None	The last column's slice, if shorter than slice_size: the row and the last summary end early.	None
probability_rule	bool	Matches the same argument of every column's sample_select (default off, bit-identical to a build without it). row then carries one more word, its split (split_word), and each column of summaries is rule_column_words long: its summary, then its second, [split, vocab)-only top-k (restricted_summary), then its reduction, so one stream per column carries all three. The forced-timestamp compare the reductions decide chooses which of the two summaries the draw comes from. Every column is slice_size: no last_slice.	False

Returns:

Type	Description
ExternalFunction	The kernel.

Raises:

Type	Description
ValueError	When a size is out of range, or probability_rule is given with a last_slice.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def sample_combine(
    *,
    columns: int = 4,
    slice_size: int = 32064,
    k_max: int = 64,
    last_slice: int | None = None,
    probability_rule: bool = False,
) -> ExternalFunction:
    """Draw from ``columns`` summaries of ``sample_select``, as ``sample_ref`` does.

    ``sample_combine(summaries, row, token, record)`` reads the summaries
    column 0 first, and writes the drawn token (an index into the whole row
    of ``columns * slice_size``) to both one-word outputs, so a design can
    send it two places.

    Args:
        columns: Summaries to combine, one per ``sample_select`` column.
        slice_size: Each column's slice of the row.
        k_max: The largest top-k the summaries carry.
        last_slice: The last column's slice, if shorter than
            ``slice_size``: the row and the last summary end early.
        probability_rule: Matches the same argument of every column's
            ``sample_select`` (default off, bit-identical to a build without
            it). ``row`` then carries one more word, its split
            (``split_word``), and each column of ``summaries`` is
            ``rule_column_words`` long: its summary, then its second,
            ``[split, vocab)``-only top-k (``restricted_summary``), then its
            ``reduction``, so one stream per column carries all three. The
            forced-timestamp compare the reductions decide chooses which of
            the two summaries the draw comes from. Every column is
            ``slice_size``: no ``last_slice``.

    Returns:
        The kernel.

    Raises:
        ValueError: When a size is out of range, or ``probability_rule`` is
            given with a ``last_slice``.
    """
    last = slice_size if last_slice is None else last_slice
    _check_slice("sample_combine", slice_size, k_max)
    _positive("sample_combine", columns=columns, last_slice=last)
    # A tie's index in its column must stay below SAMPLE_SLICE, its sentinel.
    if not k_max <= last <= slice_size:
        raise ValueError(
            "sample_combine: last_slice must hold k_max to slice_size logits"
        )
    if columns * slice_size >= 1 << 31:
        raise ValueError("sample_combine: a token must fit int32")
    if probability_rule:
        # The core strides by a whole column, so the last's restricted
        # summary must sit where a full column's would.
        if last != slice_size:
            raise ValueError("sample_combine: probability_rule needs equal slices")
        summaries_shape = (columns * rule_column_words(slice_size, k_max),)
    else:
        summaries_shape = (
            (columns - 1) * summary_words(slice_size, k_max)
            + summary_words(last, k_max),
        )
    row_words = combine_row_words(probability_rule)
    arg_types = [
        np.ndarray[summaries_shape, np.dtype[np.int32]],
        np.ndarray[(row_words,), np.dtype[np.int32]],
        np.ndarray[(1,), np.dtype[np.int32]],
        np.ndarray[(1,), np.dtype[np.int32]],
    ]
    roles = [In, In, Out, Out]
    compile_flags = list(
        _STRICT_FP
        + [
            f"-DSAMPLE_SLICE={slice_size}",
            f"-DSAMPLE_COLUMNS={columns}",
            f"-DSAMPLE_K_MAX={k_max}",
        ]
    )
    if probability_rule:
        compile_flags.append("-DSAMPLE_PROB_RULE=1")
    return _make_extern(
        "sample_combine",
        _kernel_source("sample/sample_combine.cc"),
        arg_types,
        compile_flags=compile_flags,
        contract=KernelContract(
            trace=Trace.whole_call(), roles=tuple(roles), ops_per_call=0
        ),
    )

exp64_ref

exp64_ref(x)

Compute exp over float64, bit for bit as aie_kernels/sample/exp64.h.

numpy's elementwise float64 +, -, * and comparisons are correctly rounded, so evaluating exp64.h's sequence of them gives its bits; np.exp is never used. Every branch is evaluated for every element and the results selected with np.where, which is exact.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def exp64_ref(x):
    """Compute exp over float64, bit for bit as ``aie_kernels/sample/exp64.h``.

    numpy's elementwise float64 +, -, * and comparisons are correctly
    rounded, so evaluating exp64.h's sequence of them gives its bits; np.exp
    is never used. Every branch is evaluated for every element and the
    results selected with np.where, which is exact.
    """
    c, table = _exp64_constants()
    n = np.uint64(table.size // 2)
    table_bits = int(n).bit_length() - 1
    u64, f64 = np.uint64, np.float64
    x = np.ascontiguousarray(x, dtype=np.float64)
    shape = x.shape
    x = x.reshape(-1)
    bits = x.view(np.uint64)
    abstop = (bits >> u64(52)).astype(np.int64) & 0x7FF
    tiny = abstop < 0x3C9  # |x| < 2^-54
    huge = abstop >= 0x409  # |x| >= 1024, inf, nan
    special = abstop == 0x408  # 512 <= |x| < 1024
    negative = (bits >> u64(63)) == u64(1)

    with np.errstate(all="ignore"):
        z = c["INV_LN2_N"] * x
        kd = z + c["SHIFT"]
        ki = kd.view(np.uint64).copy()
        kd = kd - c["SHIFT"]
        khi = kd * c["NEG_LN2_HI_N"]
        klo = kd * c["NEG_LN2_LO_N"]
        r = x + khi
        r = r + klo
        idx = (ki % n) * u64(2)
        top = ki << u64(52 - table_bits)
        tail = table[idx].view(np.float64)
        sbits = table[idx + u64(1)] + top
        r2 = r * r
        p23 = r * c["C3"]
        p23 = c["C2"] + p23
        p23 = r2 * p23
        p45 = r * c["C5"]
        p45 = c["C4"] + p45
        r4 = r2 * r2
        p45 = r4 * p45
        tmp = tail + r
        tmp = tmp + p23
        tmp = tmp + p45

        scale = sbits.view(np.float64)
        st = scale * tmp
        main = scale + st

        # exp64_special, k > 0.
        scale_pos = (sbits - (u64(1009) << u64(52))).view(np.float64)
        st_pos = scale_pos * tmp
        y_pos = scale_pos + st_pos
        y_pos = f64(2.0**1009) * y_pos
        # exp64_special, k < 0.
        scale_neg = (sbits + (u64(1022) << u64(52))).view(np.float64)
        st_neg = scale_neg * tmp
        y = scale_neg + st_neg
        lo = scale_neg - y
        lo = lo + st_neg
        hi = f64(1.0) + y
        lo2 = f64(1.0) - hi
        lo2 = lo2 + y
        lo2 = lo2 + lo
        y_sub = hi + lo2
        y_sub = y_sub - f64(1.0)
        y_sub = np.where(y_sub == f64(0.0), f64(0.0), y_sub)
        y_neg = np.where(y < f64(1.0), y_sub, y)
        y_neg = f64(2.0**-1022) * y_neg
        k_positive = (ki & u64(0x80000000)) == u64(0)
        special_value = np.where(k_positive, y_pos, y_neg)

        one_plus_x = f64(1.0) + x
        huge_value = np.where(
            bits == np.float64(-np.inf).view(np.uint64),
            f64(0.0),
            np.where(
                abstop >= 0x7FF,
                one_plus_x,
                np.where(negative, f64(0.0), f64(np.inf)),
            ),
        )

    out = np.where(special, special_value, main)
    out = np.where(huge, huge_value, out)
    out = np.where(tiny, one_plus_x, out)
    return out.reshape(shape)

order_keys

order_keys(logits) -> ndarray

uint16 keys whose order is the bf16 logits' numeric order; -0 and +0 are one key.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def order_keys(logits) -> np.ndarray:
    """uint16 keys whose order is the bf16 logits' numeric order; -0 and +0 are one key."""
    bits = np.ascontiguousarray(logits, dtype=bfloat16).view(np.uint16)
    bits = np.where(bits == 0x8000, np.uint16(0), bits)
    negative = (bits & np.uint16(0x8000)) != 0
    return np.where(negative, ~bits, bits | np.uint16(0x8000)).astype(np.uint16)

select_summary_ref

select_summary_ref(logits, k_max: int) -> ndarray

Return the summary sample_select writes for one column of bf16 logits at top_k k_max.

Words as sample.h lays them out: tau (the k-th largest key), the count of keys above it, the largest key, its first index and the count of keys equal to tau, three zero words, an (index, key) pair per key above tau in index order, zeros for the unused pairs (the device leaves them as they were), then the bitmap of the indices holding tau.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def select_summary_ref(logits, k_max: int) -> np.ndarray:
    """Return the summary ``sample_select`` writes for one column of bf16 ``logits`` at top_k ``k_max``.

    Words as ``sample.h`` lays them out: tau (the k-th largest key), the
    count of keys above it, the largest key, its first index and the count
    of keys equal to tau, three zero words, an (index, key) pair per key
    above tau in index order, zeros for the unused pairs (the device leaves
    them as they were), then the bitmap of the indices holding tau.
    """
    keys = order_keys(logits).astype(np.int64)
    n = keys.size
    _check_slice("select_summary_ref", n, k_max)
    tau = np.partition(keys, n - k_max)[n - k_max]
    above = np.flatnonzero(keys > tau)
    tied = keys == tau
    words = np.zeros(summary_words(n, k_max), np.int32)
    words[:5] = (tau, above.size, keys.max(), np.argmax(keys), tied.sum())
    words[SUMMARY_HEADER : SUMMARY_HEADER + 2 * above.size : 2] = above
    words[SUMMARY_HEADER + 1 : SUMMARY_HEADER + 2 * above.size : 2] = keys[above]
    bitmap = np.packbits(tied, bitorder="little")
    bitmap = np.pad(bitmap, (0, 4 * ((n + 31) // 32) - bitmap.size))
    words[SUMMARY_HEADER + 2 * k_max :] = bitmap.view("<i4")
    return words

apply_ranges_ref

apply_ranges_ref(logits, ranges) -> ndarray

Return bf16 logits, (lo, hi) ranges masked to -inf.

As sample_select's own masking does before either of its passes reads them, each endpoint clamped to the row as each column clamps it to its slice: select_summary_ref (or sample_ref) of the result draws as the device would.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def apply_ranges_ref(logits, ranges) -> np.ndarray:
    """Return bf16 ``logits``, ``(lo, hi)`` ranges masked to -inf.

    As ``sample_select``'s own masking does before either of its passes
    reads them, each endpoint clamped to the row as each column clamps it
    to its slice: ``select_summary_ref`` (or ``sample_ref``) of the result
    draws as the device would.
    """
    row = np.ascontiguousarray(logits, dtype=bfloat16).reshape(-1).copy()
    for lo, hi in ranges:
        row[min(max(lo, 0), row.size) : min(max(hi, 0), row.size)] = -np.inf
    return row

logsumexp_reduction_ref

logsumexp_reduction_ref(
    logits, split: int
) -> tuple[float32, float32, float32]

Return one column's sample_select reduction output (probability_rule).

(max_text, max_ts, sum_ts) over its own slice of bf16 logits, local split. A running max over logits[:split], and a running (max, sum of exp(v - max)) over logits[split:], each folded index by index as the device does (so a near-tie rounds the same way): float32 throughout, exp64_ref (float64) evaluating each exp before its single rounding into the float32 sum. A -inf logit weighs 0 and is not folded. -inf/-inf/0.0 where a side is empty (or all -inf), matching sample_select's own begin_reduction.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def logsumexp_reduction_ref(
    logits, split: int
) -> tuple[np.float32, np.float32, np.float32]:
    """Return one column's ``sample_select`` reduction output (``probability_rule``).

    ``(max_text, max_ts, sum_ts)`` over its own slice of bf16 ``logits``,
    local ``split``. A running max over ``logits[:split]``, and a running ``(max, sum of
    exp(v - max))`` over ``logits[split:]``, each folded index by index as
    the device does (so a near-tie rounds the same way): float32
    throughout, ``exp64_ref`` (float64) evaluating each exp before its single
    rounding into the float32 sum. A ``-inf`` logit weighs 0 and is not
    folded. ``-inf``/``-inf``/``0.0`` where a side is empty (or all
    ``-inf``), matching ``sample_select``'s own ``begin_reduction``.
    """
    values = np.ascontiguousarray(logits, dtype=bfloat16).astype(np.float32)
    max_text = np.float32(-np.inf)
    for v in values[:split]:
        if v > max_text:
            max_text = v
    max_ts = np.float32(-np.inf)
    sum_ts = np.float32(0.0)
    for v in values[split:]:
        if v == -np.inf:
            continue
        if v <= max_ts:
            sum_ts = np.float32(
                np.float64(sum_ts) + exp64_ref(np.float64(v - max_ts)).item()
            )
        else:
            sum_ts = np.float32(
                np.float64(sum_ts) * exp64_ref(np.float64(max_ts - v)).item() + 1.0
            )
            max_ts = v
    return max_text, max_ts, sum_ts

forced_timestamp_ref

forced_timestamp_ref(reductions) -> bool

Return sample_combine's forced-timestamp compare (probability_rule).

From every column's logsumexp_reduction_ref triple, column 0 first: logsumexp(logits[split:]) > max(logits[:split]), merged across columns without a log: with max_text the largest of every column's own text-range max, the compare holds exactly where sum_c sum_ts_c * exp(max_ts_c - max_text) > 1. A column with no fold on a side (sum_ts == 0) drops out of the sum. A column whose max_ts is above max_text holds at least that one logit (sum_ts >= 1), so the compare holds outright, before an exp that rounds a small excess to 1; exp64_ref (float64) of a float64 difference matches sample_combine's own merge for the rest.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def forced_timestamp_ref(reductions) -> bool:
    """Return ``sample_combine``'s forced-timestamp compare (``probability_rule``).

    From every column's ``logsumexp_reduction_ref`` triple, column 0 first:
    ``logsumexp(logits[split:]) > max(logits[:split])``, merged across
    columns without a log: with ``max_text`` the largest of every column's
    own text-range max, the compare holds exactly where ``sum_c sum_ts_c *
    exp(max_ts_c - max_text) > 1``. A column with no fold on a side
    (``sum_ts == 0``) drops out of the sum. A column whose ``max_ts`` is
    above ``max_text`` holds at least that one logit (``sum_ts >= 1``), so
    the compare holds outright, before an exp that rounds a small excess to
    1; ``exp64_ref`` (float64) of a float64 difference matches
    ``sample_combine``'s own merge for the rest.
    """
    max_text = np.float32(-np.inf)
    for text, _, _ in reductions:
        if text > max_text:
            max_text = text
    total = 0.0
    for _, max_ts, sum_ts in reductions:
        if sum_ts == np.float32(0.0):
            continue
        if max_ts > max_text:
            return True
        total += (
            float(sum_ts) * exp64_ref(np.float64(max_ts) - np.float64(max_text)).item()
        )
    return total > 1.0

check_order_preserving

check_order_preserving(temperature) -> None

Raise unless l -> fl32(l / T) strictly increases over every finite bf16 l with a finite image.

Then a threshold on bf16 keys selects what a threshold on l / T would, up to logits whose l / T is -inf and whose weight is 0 anyway.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def check_order_preserving(temperature) -> None:
    """Raise unless ``l -> fl32(l / T)`` strictly increases over every finite bf16 l with a finite image.

    Then a threshold on bf16 keys selects what a threshold on ``l / T`` would,
    up to logits whose ``l / T`` is -inf and whose weight is 0 anyway.
    """
    temperature = np.float32(temperature)
    if not 0 < temperature < np.inf:
        raise ValueError(f"temperature {temperature} is not finite and positive")
    bits = np.arange(1 << 16, dtype=np.uint32).astype(np.uint16)
    values = bits.view(bfloat16).astype(np.float32)
    values = np.sort(values[np.isfinite(values) & (bits != 0x8000)])
    with np.errstate(over="ignore"):
        x = values / temperature
    x = x[np.isfinite(x)]
    if not np.all(x[1:] > x[:-1]):
        where = int(np.flatnonzero(x[1:] <= x[:-1])[0])
        raise ValueError(
            f"temperature {temperature} maps two bf16 logits to one x "
            f"({x[where]!r}); the bf16-key threshold would not be exact"
        )

draw_row

draw_row(temperature, top_k: int, n53: int) -> ndarray

Pack the four int32 words sample_select and sample_combine read for one draw.

The kernels clamp top_k to their k_max.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def draw_row(temperature, top_k: int, n53: int) -> np.ndarray:
    """Pack the four int32 words ``sample_select`` and ``sample_combine`` read for one draw.

    The kernels clamp ``top_k`` to their ``k_max``.
    """
    _positive("draw_row", top_k=top_k)
    if top_k >= 1 << 31:
        raise ValueError(f"draw_row: top_k {top_k} does not fit int32")
    n53 = _check_n53(n53)
    row = np.empty(ROW_WORDS, dtype=np.uint32)
    row[0] = np.float32(temperature).view(np.uint32)
    row[1] = top_k
    row[2] = n53 & 0xFFFFFFFF
    row[3] = n53 >> 32
    return row.view(np.int32)

sample_weights

sample_weights(
    logits, temperature, top_k: int
) -> tuple[ndarray, ndarray]

(candidates, weights) of a draw at a positive temperature.

The candidates are the indices of every logit at or above the top_k-th largest (ties with it included), in index order; each weight is exp64(fl32(v / T) - fl32(max / T)), unnormalised.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def sample_weights(logits, temperature, top_k: int) -> tuple[np.ndarray, np.ndarray]:
    """``(candidates, weights)`` of a draw at a positive temperature.

    The candidates are the indices of every logit at or above the
    ``top_k``-th largest (ties with it included), in index order; each weight
    is ``exp64(fl32(v / T) - fl32(max / T))``, unnormalised.
    """
    _positive("sample_weights", top_k=top_k)
    row = _row(logits)
    keys = order_keys(row)
    values = row.astype(np.float32)
    temperature = np.float32(temperature)
    n = keys.size
    k = min(top_k, n)
    tau = np.partition(keys, n - k)[n - k]
    candidates = np.flatnonzero(keys >= tau)
    with np.errstate(over="ignore"):
        xm = values[int(np.argmax(keys))] / temperature
        xv = values[candidates] / temperature
    if not np.isfinite(xm):
        raise ValueError("the largest logit / T is not finite in float32")
    return candidates, exp64_ref((xv - xm).astype(np.float64))

sample_ref

sample_ref(
    logits,
    temperature,
    top_k: int,
    n53: int,
    *,
    k_max: int | None = None,
    ranges=(),
    probability_rule_split: int | None = None,
    columns: int = 1
) -> int

Return the token sample_select/sample_combine draw from one row of bf16 logits.

With k_max, top_k is clamped to it, as kernels built with that k_max clamp it; without, top_k is taken as it is.

ranges ([(lo, hi), ...], default none) and probability_rule_split (default None: off) are sample_select/sample_combine's own features: ranges masks each [lo, hi) first (apply_ranges_ref); then, if probability_rule_split is given, logits is split into columns equal slices (matching how a multi-core device build would, so the rule's fold order matches bit for bit), each slice's own logsumexp_reduction_ref, merged as forced_timestamp_ref does, additionally masks [0, probability_rule_split) (the split clamped to the row) where it fires.

Source code in utils/pyright_stubs/aie/iron/kernels/sample.py

def sample_ref(
    logits,
    temperature,
    top_k: int,
    n53: int,
    *,
    k_max: int | None = None,
    ranges=(),
    probability_rule_split: int | None = None,
    columns: int = 1,
) -> int:
    """Return the token ``sample_select``/``sample_combine`` draw from one row of bf16 logits.

    With ``k_max``, ``top_k`` is clamped to it, as kernels built with that
    ``k_max`` clamp it; without, ``top_k`` is taken as it is.

    ``ranges`` (``[(lo, hi), ...]``, default none) and ``probability_rule_split``
    (default ``None``: off) are ``sample_select``/``sample_combine``'s own
    features: ``ranges`` masks each ``[lo, hi)`` first
    (``apply_ranges_ref``); then, if ``probability_rule_split`` is given,
    ``logits`` is split into ``columns`` equal slices (matching how a
    multi-core device build would, so the rule's fold order matches bit for
    bit), each slice's own ``logsumexp_reduction_ref``, merged as
    ``forced_timestamp_ref`` does, additionally masks ``[0,
    probability_rule_split)`` (the split clamped to the row) where it fires.
    """
    _positive("sample_ref", top_k=top_k)
    n53 = _check_n53(n53)
    if k_max is not None:
        _positive("sample_ref", k_max=k_max)
        top_k = min(top_k, k_max)
    if ranges or probability_rule_split is not None:
        logits = apply_ranges_ref(logits, ranges)
        if probability_rule_split is not None:
            n = len(logits)
            split = min(max(probability_rule_split, 0), n)
            _positive("sample_ref", columns=columns)
            if n % columns:
                raise ValueError("sample_ref: columns must divide the row")
            slice_size = n // columns
            reductions = [
                logsumexp_reduction_ref(
                    logits[c * slice_size : (c + 1) * slice_size],
                    min(max(split - c * slice_size, 0), slice_size),
                )
                for c in range(columns)
            ]
            if forced_timestamp_ref(reductions):
                logits = logits.copy()
                logits[:split] = -np.inf
    if np.float32(temperature).view(np.uint32) & 0x7FFFFFFF == 0:
        return int(np.argmax(order_keys(_row(logits))))
    candidates, weights = sample_weights(logits, temperature, top_k)
    prefix = list(accumulate(_units(weights)))
    # P * 2**53 > n53 * S is P > floor(n53 * S / 2**53) for an integer P. The
    # prefixes never step down, so the first above is a bisection; S >= 1
    # (the maximum weighs 1) and u < 1, so the last always is.
    target = n53 * prefix[-1] >> 53
    return int(candidates[bisect_right(prefix, target)])

Mixture of experts

router_top_k turns one sample_select summary into a router's top_k expert ids and their softmax weights, deterministically. router_top_k_ref is the bit-exact host reference. experts runs one row through its chosen experts' gated feed-forward on one core, from the row and one stream of weight tiles.

Mixture of experts kernels: the router's top-k selection and the experts' FFN.

router_top_k reads one column's sample_select summary (sample.py, k_max == top_k, one call reduces the whole row) and writes the top_k expert ids and their softmax weights, deterministically: where sample_combine draws one token stochastically from every candidate at or above its threshold, routing instead takes exactly the top_k candidates the summary already holds (the entries strictly above tau in index order, then the lowest-index ties at tau, not ranked by value) and a softmax over their values alone. That softmax is the weighting of routers that normalize over the selected experts (Granite-MoE, Mixtral, norm_topk_prob=True).

router_top_k

router_top_k(*, slice_size=32, top_k=8) -> ExternalFunction

One column's router: sample_select's summary to top_k ids and weights.

router_top_k(summary, ids, weights) takes the summary_words(slice_size, top_k) int32 words a sample_select(slice_size=slice_size, chunk=slice_size, k_max=top_k) call produces for the whole row in one pass, and writes top_k int32 ids and 2 * top_k bf16 weights: ids[i] pairs with weights[2 * i], weights[2 * i + 1] a zeroed, never meaningful pad word. The padding keeps every real weight at a 4-byte-aligned bf16 index, which a graph's Copy of one selected weight into another buffer needs (found tracing a real iron/common/image/fusion.py alignment failure, not assumed: the allocator always places a buffer's own base 64-byte aligned, so an unpadded weights[i] at odd i is the one place the misalignment could come from).

Source code in utils/pyright_stubs/aie/iron/kernels/moe.py

def router_top_k(*, slice_size=32, top_k=8) -> ExternalFunction:
    """One column's router: ``sample_select``'s summary to ``top_k`` ids and weights.

    ``router_top_k(summary, ids, weights)`` takes the ``summary_words(slice_size,
    top_k)`` int32 words a ``sample_select(slice_size=slice_size, chunk=slice_size,
    k_max=top_k)`` call produces for the whole row in one pass, and writes
    ``top_k`` int32 ``ids`` and ``2 * top_k`` bf16 ``weights``: ``ids[i]``
    pairs with ``weights[2 * i]``, ``weights[2 * i + 1]`` a zeroed, never
    meaningful pad word. The padding keeps every real weight at a
    4-byte-aligned bf16 index, which a graph's ``Copy`` of one selected
    weight into another buffer needs (found tracing a real
    ``iron/common/image/fusion.py`` alignment failure, not assumed: the
    allocator always places a buffer's own base 64-byte aligned, so an
    unpadded ``weights[i]`` at odd ``i`` is the one place the misalignment
    could come from).
    """
    _positive("router_top_k", slice_size=slice_size, top_k=top_k)
    _check_slice("router_top_k", slice_size, top_k)
    return _make_extern(
        "router_top_k",
        _kernel_source("moe/router_top_k.cc"),
        [
            np.ndarray[(summary_words(slice_size, top_k),), np.dtype[np.int32]],
            np.ndarray[(top_k,), np.dtype[np.int32]],
            np.ndarray[(2 * top_k,), np.dtype[bfloat16]],
        ],
        # Unrolled, the float32 and float64 paths overflow aie2's 16 KB.
        compile_flags=(_STRICT_FP_ROLLED if _detect_arch() == "aie2" else _STRICT_FP)
        + [f"-DSAMPLE_SLICE={slice_size}", f"-DSAMPLE_K_MAX={top_k}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Out),
            reference=partial(router_top_k_summary_ref, top_k=top_k),
            sample=partial(router_top_k_sample, slice_size=slice_size, top_k=top_k),
            tolerance=Tolerance.exact(
                note="ids and weights bit for bit: the weights are the float64 "
                "softmax rounded once to bf16"
            ),
        ),
    )

router_top_k_ref

router_top_k_ref(
    logits, top_k: int
) -> tuple[ndarray, ndarray]

(ids, weights) as the device's router_top_k writes one row.

Bit for bit: see router_top_k.cc. weights is 2 * top_k wide, a real weight at every even index (weights[2 * i] pairs with ids[i]) and a zeroed pad word at every odd one.

Exact top-top_k of logits, the lowest indices winning a tie at the top_k-th value. ids holds the entries above that value in index order, then the tied ones in index order, not ranked by value as torch.topk(sorted=True) is. Then a softmax over the selected values: a soft-float IEEE float32 subtraction against the row's max (what __subsf3 gives the device, where Peano's native - would not), exp64 (sample.py's bit-exact float64 exponential, the same primitive sample_combine uses) and a plain float64 sum and divide, rounded once to bf16.

Source code in utils/pyright_stubs/aie/iron/kernels/moe.py

def router_top_k_ref(logits, top_k: int) -> tuple[np.ndarray, np.ndarray]:
    """``(ids, weights)`` as the device's ``router_top_k`` writes one row.

    Bit for bit: see ``router_top_k.cc``. ``weights`` is ``2 * top_k`` wide,
    a real weight at every even index (``weights[2 * i]`` pairs with
    ``ids[i]``) and a zeroed pad word at every odd one.

    Exact top-``top_k`` of ``logits``, the lowest indices winning a tie at
    the ``top_k``-th value. ``ids`` holds the entries above that value in
    index order, then the tied ones in index order, not ranked by value as
    ``torch.topk(sorted=True)`` is. Then a softmax over the selected values:
    a soft-float IEEE float32 subtraction against the row's max (what
    ``__subsf3`` gives the device, where Peano's native ``-`` would not),
    ``exp64`` (``sample.py``'s bit-exact float64 exponential, the same
    primitive ``sample_combine`` uses) and a plain float64 sum and divide,
    rounded once to bf16.
    """
    row = np.ascontiguousarray(logits, dtype=bfloat16).reshape(-1)
    keys = order_keys(row).astype(np.int64)
    n = keys.size
    if not 1 <= top_k <= n:
        raise ValueError(f"router_top_k_ref: top_k ({top_k}) must be in [1, {n}]")
    tau = np.partition(keys, n - top_k)[n - top_k]
    above = np.flatnonzero(keys > tau)  # ascending index, fewer than top_k
    ties = np.flatnonzero(keys == tau)  # ascending index
    needed = top_k - above.size
    ids = np.concatenate([above, ties[:needed]]).astype(np.int32)
    return ids, _router_weights(row[ids], row[int(np.argmax(keys))])

router_top_k_summary_ref

router_top_k_summary_ref(
    summary, *, top_k: int
) -> tuple[ndarray, ndarray]

(ids, weights) per call, as router_top_k writes them from each call's summary.

The summary's entries in order, then its lowest-index ties at tau, up to top_k; the softmax is router_top_k_ref's.

Source code in utils/pyright_stubs/aie/iron/kernels/moe.py

def router_top_k_summary_ref(summary, *, top_k: int) -> tuple[np.ndarray, np.ndarray]:
    """``(ids, weights)`` per call, as ``router_top_k`` writes them from each call's summary.

    The summary's entries in order, then its lowest-index ties at tau, up to
    ``top_k``; the softmax is ``router_top_k_ref``'s.
    """
    summary = np.asarray(summary, np.int32).reshape(-1, np.shape(summary)[-1])
    ids = np.zeros((len(summary), top_k), np.int32)
    weights = np.zeros((len(summary), 2 * top_k), bfloat16)
    for c, words in enumerate(summary):
        entries = int(words[1])
        pairs = words[SUMMARY_HEADER : SUMMARY_HEADER + 2 * entries]
        bitmap = words[SUMMARY_HEADER + 2 * top_k :].view(np.uint32)
        tied = np.flatnonzero(np.unpackbits(bitmap.view(np.uint8), bitorder="little"))
        ids[c] = np.concatenate([pairs[0::2], tied[: top_k - entries]])
        keys = np.concatenate([pairs[1::2], np.full(top_k - entries, words[0])])
        weights[c] = _router_weights(_key_values(keys), _key_values(words[2]))
    return ids, weights

router_top_k_sample

router_top_k_sample(
    rng: Generator,
    calls: int,
    *,
    slice_size: int,
    top_k: int
) -> list[ndarray]

One sample_select summary per call, of rows chosen where a ranking can go wrong.

Random rows, then by turns: ties at the k-th value, one value everywhere, signed zeros, exactly top_k finite logits among -inf, and the largest value last among the winners.

Source code in utils/pyright_stubs/aie/iron/kernels/moe.py

def router_top_k_sample(
    rng: np.random.Generator, calls: int, *, slice_size: int, top_k: int
) -> list[np.ndarray]:
    """One ``sample_select`` summary per call, of rows chosen where a ranking can go wrong.

    Random rows, then by turns: ties at the k-th value, one value everywhere,
    signed zeros, exactly ``top_k`` finite logits among -inf, and the largest
    value last among the winners.
    """
    n = slice_size
    rows = []
    for c in range(calls):
        kind = c % 8
        logits = rng.normal(0, 2, n)
        if kind == 1:
            logits = rng.normal(0, 0.5, n)
            logits[rng.choice(n, top_k, replace=False)] = np.sort(logits)[
                -1 - top_k // 2
            ]
        elif kind == 2:
            logits = np.full(n, 1.5)
        elif kind == 3:
            logits = np.where(rng.random(n) < 0.5, -0.0, 0.0)
            logits[rng.choice(n, min(3, n), replace=False)] = 1.0
        elif kind == 4:
            logits = np.full(n, -np.inf)
            logits[rng.choice(n, top_k, replace=False)] = rng.normal(0, 1, top_k)
        elif kind == 5:
            logits = np.arange(n) / 16.0
        rows.append(select_summary_ref(np.asarray(logits, bfloat16), top_k))
    return [np.stack(rows)]

experts_sample

experts_sample(
    rng: Generator,
    calls: int,
    *,
    hidden: int,
    width: int,
    num_experts: int,
    top_k: int,
    rows: int,
    cols: int
) -> list[ndarray]

Return random inputs for experts, as a design passes them.

Every third row repeats its first id when top_k > 1, so a device run checks that a duplicated id sums its weights.

Parameters:

Name	Type	Description	Default
rng	Generator	The generator to draw from.	required
calls	int	How many kernel calls, one step each, to draw for.	required
hidden	int	As for experts.	required
width	int	As for experts.	required
num_experts	int	As for experts.	required
top_k	int	As for experts.	required
rows	int	As for experts.	required
cols	int	As for experts.	required

Returns:

Type	Description
list[ndarray]	[row, tile], one entry per call: each row repeated for its
list[ndarray]	steps, and a random tile.

Source code in utils/pyright_stubs/aie/iron/kernels/moe.py

def experts_sample(
    rng: np.random.Generator,
    calls: int,
    *,
    hidden: int,
    width: int,
    num_experts: int,
    top_k: int,
    rows: int,
    cols: int,
) -> list[np.ndarray]:
    """Return random inputs for [`experts`][aie.iron.kernels.moe.experts], as a design passes them.

    Every third row repeats its first id when ``top_k`` > 1, so a device run
    checks that a duplicated id sums its weights.

    Args:
        rng: The generator to draw from.
        calls: How many kernel calls, one step each, to draw for.
        hidden: As for `experts`.
        width: As for `experts`.
        num_experts: As for `experts`.
        top_k: As for `experts`.
        rows: As for `experts`.
        cols: As for `experts`.

    Returns:
        ``[row, tile]``, one entry per call: each row repeated for its
        steps, and a random tile.
    """
    row_steps = num_experts * _experts_steps(hidden, width, rows, cols)[1]
    packed = []
    for row in range(-(-calls // row_steps)):
        weights = np.zeros(2 * top_k, np.float32)
        weights[::2] = rng.random(top_k)
        ids = rng.choice(num_experts, top_k, replace=False)
        if row % 3 == 2 and top_k > 1:
            ids[1] = ids[0]
        packed.append(_experts_row(rng.standard_normal(hidden) * 0.5, weights, ids))
    rows_in = np.repeat(np.stack(packed), row_steps, axis=0)[:calls]
    tiles = rng.standard_normal((calls, rows * cols)) * 0.1
    return [rows_in, tiles.astype(bfloat16)]

experts_ref

experts_ref(
    row: ndarray,
    tile: ndarray,
    *,
    hidden: int,
    width: int,
    num_experts: int,
    top_k: int,
    rows: int,
    cols: int
) -> ndarray

Numpy reference for experts.

Parameters:

Name	Type	Description	Default
row	ndarray	The rows, as experts_sample returns them.	required
tile	ndarray	The weight tiles, as experts_sample returns them.	required
hidden	int	As for experts.	required
width	int	As for experts.	required
num_experts	int	As for experts.	required
top_k	int	As for experts.	required
rows	int	As for experts.	required
cols	int	As for experts.	required

Returns:

Type	Description
ndarray	y after each call, call c being step c.

Source code in utils/pyright_stubs/aie/iron/kernels/moe.py

def experts_ref(
    row: np.ndarray,
    tile: np.ndarray,
    *,
    hidden: int,
    width: int,
    num_experts: int,
    top_k: int,
    rows: int,
    cols: int,
) -> np.ndarray:
    """Numpy reference for [`experts`][aie.iron.kernels.moe.experts].

    Args:
        row: The rows, as `experts_sample` returns them.
        tile: The weight tiles, as `experts_sample` returns them.
        hidden: As for `experts`.
        width: As for `experts`.
        num_experts: As for `experts`.
        top_k: As for `experts`.
        rows: As for `experts`.
        cols: As for `experts`.

    Returns:
        ``y`` after each call, call ``c`` being step ``c``.
    """
    gate_up_steps, expert_steps = _experts_steps(hidden, width, rows, cols)
    row = np.ascontiguousarray(row, bfloat16).reshape(-1, hidden + 4 * top_k)
    ids = row[:, hidden + 2 * top_k :].copy().view(np.int32)
    row = row.astype(np.float32)
    tile = np.asarray(tile, np.float32).reshape(-1, rows, cols)
    y = np.zeros((len(row), hidden), np.float32)
    total = np.zeros(hidden, np.float32)
    acc = np.zeros(2 * width, np.float32)
    for call, (r, t) in enumerate(zip(row, tile)):
        e, s = divmod(call % (num_experts * expert_steps), expert_steps)
        if e == 0 and s == 0:
            total = np.zeros(hidden, np.float32)
        chosen = ids[call] == e
        if s == 0:
            acc = np.zeros(2 * width, np.float32)
        if chosen.any() and s < gate_up_steps:
            r0, c0 = divmod(s, hidden // cols)
            r0, c0 = r0 * rows, c0 * cols
            acc[r0 : r0 + rows] += t @ r[c0 : c0 + cols]
            if s == gate_up_steps - 1:
                g = acc[:width]
                with np.errstate(over="ignore"):
                    acc[:width] = g / (1 + np.exp(-g)) * acc[width:]
        elif chosen.any():
            weight = r[hidden : hidden + 2 * top_k : 2][chosen].sum()
            r0, c0 = divmod(s - gate_up_steps, width // cols)
            r0, c0 = r0 * rows, c0 * cols
            total[r0 : r0 + rows] += weight * (t @ acc[c0 : c0 + cols])
        y[call] = total
    return y.astype(bfloat16)

experts

experts(
    *,
    hidden: int = 64,
    width: int = 32,
    num_experts: int = 4,
    top_k: int = 2,
    rows: int = 8,
    cols: int = 32
) -> ExternalFunction

Mixture of experts' feed-forward, a row at a time, from two input streams.

experts(row, tile, y, step): row is bf16 x (hidden), then the row's 2 * top_k router weights as router_top_k writes them (chosen expert j's at 2 * j), then its top_k int32 ids as pairs of bf16 words. Step s brings tile s modulo a row's steps, rows x cols bf16, of one fixed stream: for each of the num_experts experts in order, the blocks of its gate_up (2 * width x hidden) row block by row block, then those of its down (hidden x width). The kernel skips the tiles of experts the row did not choose. y is the row's bf16 output so far; after the row's last step it is the sum over j of weight j times down[ids[j]] @ (silu(g) * u), g and u the halves of gate_up[ids[j]] @ x, in float32 and rounded once.

Parameters:

Name	Type	Description	Default
hidden	int	The model's hidden size, the length of x and y.	64
width	int	Each expert's intermediate size.	32
num_experts	int	Experts in the weight stream.	4
top_k	int	Experts each row chooses.	2
rows	int	Rows of a weight tile. A shim BD repeats a dimension at most 64 times, so at Qwen3 shapes it is 32 or more.	8
cols	int	Columns of a weight tile, a multiple of 16.	32

Returns:

Type	Description
ExternalFunction	The kernel.

Raises:

Type	Description
ValueError	A dimension is not a positive integer, a tile does not divide the weights, cols is not a multiple of 16 (the kernel's vector width; hidden is then too, so the ids after x are word aligned), or top_k exceeds num_experts.

Source code in utils/pyright_stubs/aie/iron/kernels/moe.py

def experts(
    *,
    hidden: int = 64,
    width: int = 32,
    num_experts: int = 4,
    top_k: int = 2,
    rows: int = 8,
    cols: int = 32,
) -> ExternalFunction:
    """Mixture of experts' feed-forward, a row at a time, from two input streams.

    ``experts(row, tile, y, step)``: ``row`` is bf16 x (``hidden``), then
    the row's ``2 * top_k`` router weights as ``router_top_k`` writes them
    (chosen expert ``j``'s at ``2 * j``), then its ``top_k`` int32 ids as
    pairs of bf16 words. Step ``s`` brings tile ``s`` modulo a row's steps,
    ``rows`` x ``cols`` bf16, of one fixed stream: for each of the
    ``num_experts`` experts in order, the blocks of its ``gate_up`` (``2 *
    width`` x ``hidden``) row block by row block, then those of its ``down``
    (``hidden`` x ``width``). The kernel skips the tiles of experts the row
    did not choose. ``y`` is the row's bf16 output so far; after the row's
    last step it is the sum over ``j`` of weight ``j`` times ``down[ids[j]]
    @ (silu(g) * u)``, ``g`` and ``u`` the halves of ``gate_up[ids[j]] @
    x``, in float32 and rounded once.

    Args:
        hidden: The model's hidden size, the length of x and y.
        width: Each expert's intermediate size.
        num_experts: Experts in the weight stream.
        top_k: Experts each row chooses.
        rows: Rows of a weight tile. A shim BD repeats a dimension at most 64
            times, so at Qwen3 shapes it is 32 or more.
        cols: Columns of a weight tile, a multiple of 16.

    Returns:
        The kernel.

    Raises:
        ValueError: A dimension is not a positive integer, a tile does not
            divide the weights, ``cols`` is not a multiple of 16 (the
            kernel's vector width; ``hidden`` is then too, so the ids after
            x are word aligned), or ``top_k`` exceeds ``num_experts``.
    """
    geometry = dict(
        hidden=hidden,
        width=width,
        num_experts=num_experts,
        top_k=top_k,
        rows=rows,
        cols=cols,
    )
    _positive("experts", **geometry)
    if (2 * width) % rows or hidden % rows or hidden % cols or width % cols:
        raise ValueError(
            "experts: rows must divide 2 * width and hidden, and cols must "
            "divide hidden and width"
        )
    if cols % 16:
        raise ValueError(f"experts: cols ({cols}) must be a multiple of 16")
    if top_k > num_experts:
        raise ValueError("experts: top_k must be at most num_experts")
    return _make_extern(
        "experts",
        _kernel_source("moe/experts.cc"),
        [
            np.ndarray[(hidden + 4 * top_k,), np.dtype[bfloat16]],
            np.ndarray[(rows * cols,), np.dtype[bfloat16]],
            np.ndarray[(hidden,), np.dtype[bfloat16]],
            np.int32,
        ],
        compile_flags=_STRICT_FP_ROLLED
        + [
            f"-DEXPERTS_HIDDEN={hidden}",
            f"-DEXPERTS_WIDTH={width}",
            f"-DEXPERTS_EXPERTS={num_experts}",
            f"-DEXPERTS_TOP_K={top_k}",
            f"-DEXPERTS_TILE_ROWS={rows}",
            f"-DEXPERTS_TILE_COLS={cols}",
        ],
        contract=KernelContract(
            # A tile row, and x beside it, load 32 bf16s a vector where cols
            # is whole 32s, else 16; y stores 32 where hidden is whole 32s.
            alignments=(_vector_args(0, 1) if cols % 32 == 0 else ((0, 32), (1, 32)))
            + (_vector_args(2) if hidden % 32 == 0 else ((2, 32),)),
            trace=Trace.whole_call(),
            roles=(In, In, Out, Param),
            parameter_bindings=((3, 0),),
            call_index=3,
            reference=partial(experts_ref, **geometry),
            sample=partial(experts_sample, **geometry),
            tolerance=Tolerance(
                ulps=1,
                range_frac=1e-5,
                note="npu2: bit for bit on random rows, duplicated ids included; "
                "sums that cancel (edge-all-chosen) land 1 ulp off in 24 of 18432; "
                "npu1: unmeasured",
            ),
        ),
    )

State space models

selective_scan runs Mamba2's recurrence for one head over a block of rows, and gated_delta_rule Gated DeltaNet's for one value head, each with its float32 state carried from call to call.

• Mamba2: Dao and Gu; the reference is selective_state_update in state-spaces/mamba, less its optional silu(z) gate.
• Gated DeltaNet: Yang, Kautz and Hatamizadeh, on DeltaNet (Yang et al.); the references are fla-org's fused_recurrent_gated_delta_rule and Hugging Face's Qwen3-Next.
• Their exponentials, and router_top_k's, reduce by ln 2 in two parts (Cody and Waite, Software Manual for the Elementary Functions, 1980).

State space model kernels: Mamba2's selective scan and Gated DeltaNet's delta rule.

selective_scan_sample

selective_scan_sample(
    rng,
    calls: int,
    *,
    rows: int,
    head_dim: int,
    state: int,
    blocks: int
)

Return heads of blocks blocks, each from a random state.

A < 0, so the state decays.

Source code in utils/pyright_stubs/aie/iron/kernels/ssm.py

def selective_scan_sample(
    rng, calls: int, *, rows: int, head_dim: int, state: int, blocks: int
):
    """Return heads of ``blocks`` blocks, each from a random state.

    A < 0, so the state decays.
    """
    words = -(-rows * (head_dim + 2 * state + 1) // 2)
    block = rng.standard_normal((calls, 2 * words)) * 0.5
    heads = [
        np.concatenate(
            [
                [rng.uniform(-4, 0)],  # dt_bias
                [-rng.uniform(0.5, 16)],  # A = -exp(A_log)
                [rng.uniform(0, 2)],  # D
                rng.standard_normal(head_dim * state),
            ]
        )
        for _ in range(-(-calls // blocks))
    ]
    head = np.repeat(np.stack(heads), blocks, axis=0)[:calls]
    return [block.astype(bfloat16), head.astype(np.float32)]

selective_scan_ref

selective_scan_ref(
    block,
    head,
    valid,
    *,
    rows: int,
    head_dim: int,
    state: int,
    blocks: int
)

Numpy reference for selective_scan.

Call b is block b. Returns (final, y) per call.

Source code in utils/pyright_stubs/aie/iron/kernels/ssm.py

def selective_scan_ref(
    block, head, valid, *, rows: int, head_dim: int, state: int, blocks: int
):
    """Numpy reference for [`selective_scan`][aie.iron.kernels.ssm.selective_scan].

    Call ``b`` is block ``b``. Returns ``(final, y)`` per call.
    """
    head = np.asarray(head, np.float32).reshape(-1, 3 + head_dim * state)
    calls = len(head)
    block = np.asarray(block, np.float32).reshape(calls, -1)
    x = block[:, : rows * head_dim].reshape(calls, rows, head_dim)
    end = rows * (head_dim + 2 * state)
    bc = block[:, rows * head_dim : end].reshape(calls, rows, 2, state)
    dt = block[:, end : end + rows]
    final = np.zeros((calls, head_dim * state), np.float32)
    y = np.zeros((calls, rows, head_dim), np.float32)
    s = np.zeros((head_dim, state), np.float32)
    for c in range(calls):
        if c % blocks == 0:
            s = head[c, 3:].reshape(head_dim, state).copy()
        dt_bias, A, D = head[c, :3]
        for t in range(min(int(valid), rows)):
            B, C = bc[c, t]
            step = np.logaddexp(np.float32(0), dt[c, t] + dt_bias)
            s = np.exp(step * A) * s + (step * x[c, t])[:, None] * B
            y[c, t] = s @ C + D * x[c, t]
        final[c] = s.reshape(-1)
    return final, y.astype(bfloat16).reshape(calls, -1)

selective_scan

selective_scan(
    *,
    rows: int = 16,
    head_dim: int = 16,
    state: int = 32,
    blocks: int = 1
) -> ExternalFunction

Mamba2's state space recurrence for one head, a block of rows a call.

selective_scan(block, head, final, y, valid, b): block is bf16 x (rows x head_dim), then bc (rows rows of the group's state B then state C), then dt (rows), then a bf16 of padding where those are odd, so a DMA moves whole words. head is float32 (dt_bias, A, D) then the state before the head's first block, (head_dim, state). b counts blocks from 0, blocks to a head, and block b loads the state from head when it is a head's first. Each block runs its first valid rows: step = softplus(dt + dt_bias), state = exp(step * A) * state + step * outer(x, B) and y = state @ C + D * x, rounded once to bf16; rows from valid on leave the state alone and write zero y. The state stays in the core between blocks, and final is it after this one. The caller gathers head h's x and dt and group h // (heads // groups)'s B and C into block, so each group's B and C stream heads // groups times.

Parameters:

Name	Type	Description	Default
rows	int	Rows of a block.	16
head_dim	int	The head's x width, or a slice of it: the rows of state are independent. A multiple of 16.	16
state	int	The state size, a multiple of 16.	32
blocks	int	Blocks to a head.	1

Returns:

Type	Description
ExternalFunction	The kernel.

Raises:

Type	Description
ValueError	A dimension is not a positive integer, or head_dim or state is not a multiple of 16, the kernel's vector width.

Source code in utils/pyright_stubs/aie/iron/kernels/ssm.py

def selective_scan(
    *, rows: int = 16, head_dim: int = 16, state: int = 32, blocks: int = 1
) -> ExternalFunction:
    """Mamba2's state space recurrence for one head, a block of rows a call.

    ``selective_scan(block, head, final, y, valid, b)``: ``block`` is bf16
    x (``rows`` x ``head_dim``), then bc (``rows`` rows of the group's
    ``state`` B then ``state`` C), then dt (``rows``), then a bf16 of
    padding where those are odd, so a DMA moves whole words. ``head`` is float32
    ``(dt_bias, A, D)`` then the state before the head's first block,
    ``(head_dim, state)``. ``b`` counts blocks from 0, ``blocks`` to a head,
    and block ``b`` loads the state from ``head`` when it is a head's first.
    Each block runs its first ``valid`` rows:
    ``step = softplus(dt + dt_bias)``,
    ``state = exp(step * A) * state + step * outer(x, B)`` and
    ``y = state @ C + D * x``, rounded
    once to bf16; rows from ``valid`` on leave the state alone and write
    zero ``y``. The state stays in the core between blocks, and ``final``
    is it after this one. The caller gathers head ``h``'s x and dt and
    group ``h // (heads // groups)``'s B and C into ``block``, so each
    group's B and C stream ``heads // groups`` times.

    Args:
        rows: Rows of a block.
        head_dim: The head's x width, or a slice of it: the rows of state
            are independent. A multiple of 16.
        state: The state size, a multiple of 16.
        blocks: Blocks to a head.

    Returns:
        The kernel.

    Raises:
        ValueError: A dimension is not a positive integer, or ``head_dim`` or
            ``state`` is not a multiple of 16, the kernel's vector width.
    """
    geometry = dict(rows=rows, head_dim=head_dim, state=state, blocks=blocks)
    _positive("selective_scan", **geometry)
    if head_dim % 16 or state % 16:
        raise ValueError(
            f"selective_scan: head_dim ({head_dim}) and state ({state}) must be "
            "multiples of 16"
        )
    # One bf16 of padding after an odd block, so a DMA moves whole words.
    words = -(-rows * (head_dim + 2 * state + 1) // 2)
    return _make_extern(
        "selective_scan",
        _kernel_source("ssm/selective_scan.cc"),
        [
            np.ndarray[(2 * words,), np.dtype[bfloat16]],
            np.ndarray[(3 + head_dim * state,), np.dtype[np.float32]],
            np.ndarray[(head_dim * state,), np.dtype[np.float32]],
            np.ndarray[(rows * head_dim,), np.dtype[bfloat16]],
            np.int32,
            np.int32,
        ],
        compile_flags=_STRICT_FP_ROLLED
        + [
            # Scheduling the row update past the default pressure limits
            # measured 8% fewer cycles at IRON's prefill slice.
            "-mllvm",
            "--aie-premisched-reg-pressure-instr-threshold=10000",
            "-mllvm",
            "--aie-premisched-finer-rp-tracking=false",
            f"-DSCAN_ROWS={rows}",
            f"-DSCAN_HEAD_DIM={head_dim}",
            f"-DSCAN_STATE={state}",
            f"-DSCAN_BLOCKS={blocks}",
        ],
        contract=KernelContract(
            # B and C load 32 bf16s a vector where state is whole 32s and
            # start on a 64-byte line (rows * head_dim whole 32s), x 16.
            alignments=(
                _vector_args(0)
                if state % 32 == 0 and rows * head_dim % 32 == 0
                else ((0, 32),)
            )
            + _vector_args(2)
            + ((3, 32),),
            trace=Trace.whole_call(),
            roles=(In, In, Out, Out, Param, Param),
            parameter_bindings=((5, 0),),
            call_index=5,
            reference=partial(selective_scan_ref, **geometry),
            sample=partial(selective_scan_sample, **geometry),
            tolerance=Tolerance.bounded(
                partial(_selective_scan_bound, **geometry),
                note="derived: y exact but at a near tie, final within the "
                "float32 recurrence's error bound; npu2: within it on every "
                "case over three seeds; npu1: unmeasured",
            ),
        ),
    )

gated_delta_rule_sample

gated_delta_rule_sample(
    rng,
    calls: int,
    *,
    rows: int,
    key_dim: int,
    value_dim: int,
    blocks: int
)

Return heads of blocks blocks, A < 0 so the state decays, each from a random state.

Source code in utils/pyright_stubs/aie/iron/kernels/ssm.py

def gated_delta_rule_sample(
    rng, calls: int, *, rows: int, key_dim: int, value_dim: int, blocks: int
):
    """Return heads of ``blocks`` blocks, A < 0 so the state decays, each from a random state."""
    block = rng.standard_normal((calls, rows * (2 * key_dim + value_dim + 2)))
    # Row 1's q and k are tiny, so their norms depend on eps.
    for m in (0, rows):
        block[:, (m + 1) * key_dim : (m + 2) * key_dim] *= 1e-4
    heads = [
        np.concatenate(
            [
                [rng.uniform(-4, 0)],  # dt_bias
                [-rng.uniform(0.5, 16)],  # A = -exp(A_log)
                rng.standard_normal(key_dim * value_dim) * 0.5,
            ]
        )
        for _ in range(-(-calls // blocks))
    ]
    head = np.repeat(np.stack(heads), blocks, axis=0)[:calls]
    return [block.astype(bfloat16), head.astype(np.float32)]

gated_delta_rule_ref

gated_delta_rule_ref(
    block,
    head,
    valid,
    *,
    rows: int,
    key_dim: int,
    value_dim: int,
    blocks: int
)

Numpy reference for gated_delta_rule: call b is block b.

Returns (final, y) per call.

Source code in utils/pyright_stubs/aie/iron/kernels/ssm.py

def gated_delta_rule_ref(
    block, head, valid, *, rows: int, key_dim: int, value_dim: int, blocks: int
):
    """Numpy reference for [`gated_delta_rule`][aie.iron.kernels.ssm.gated_delta_rule]: call ``b`` is block ``b``.

    Returns ``(final, y)`` per call.
    """
    head = np.asarray(head, np.float32).reshape(-1, 2 + key_dim * value_dim)
    calls = len(head)
    block = np.asarray(block, np.float32).reshape(calls, -1)
    q, k, v, ab = np.split(
        block, np.cumsum([rows * key_dim, rows * key_dim, rows * value_dim]), axis=1
    )
    q, k = q.reshape(calls, rows, key_dim), k.reshape(calls, rows, key_dim)
    v, a, b = v.reshape(calls, rows, value_dim), ab[:, :rows], ab[:, rows:]
    final = np.zeros((calls, key_dim * value_dim), np.float32)
    y = np.zeros((calls, rows, value_dim), np.float32)
    s = np.zeros((key_dim, value_dim), np.float32)
    for c in range(calls):
        if c % blocks == 0:
            s = head[c, 2:].reshape(key_dim, value_dim).copy()
        dt_bias, A = head[c, :2]
        for t in range(min(int(valid), rows)):
            qt = q[c, t] / np.sqrt((q[c, t] ** 2).sum() + np.float32(1e-6))
            qt = qt / np.sqrt(np.float32(key_dim))
            kt = k[c, t] / np.sqrt((k[c, t] ** 2).sum() + np.float32(1e-6))
            g = A * np.logaddexp(np.float32(0), a[c, t] + dt_bias)
            beta = 1 / (1 + np.exp(-b[c, t]))
            s = np.exp(g) * s
            s = s + np.outer(kt, beta * (v[c, t] - kt @ s))
            y[c, t] = qt @ s
        final[c] = s.reshape(-1)
    return final, y.astype(bfloat16).reshape(calls, -1)

gated_delta_rule

gated_delta_rule(
    *,
    rows: int = 16,
    key_dim: int = 32,
    value_dim: int = 16,
    blocks: int = 1
) -> ExternalFunction

Gated DeltaNet's recurrence for one value head, a block of rows a call.

gated_delta_rule(block, head, final, y, valid, b): block is bf16 q and k (rows x key_dim each), v (rows x value_dim), then a and b (rows each). head is float32 (dt_bias, A) then the state before the head's first block, (key_dim, value_dim). b counts blocks from 0, blocks to a head, and a head's first block loads the state from head. Each block runs its first valid rows: q and k are l2-normalized (eps 1e-6) and q scaled by 1 / sqrt(key_dim), g = A * softplus(a + dt_bias), beta = sigmoid(b), state = exp(g) * state, then state += outer(k, beta * (v - k @ state)) and y = q @ state, rounded once to bf16. Rows from valid on leave the state alone and write zero y. The state stays in the core between blocks, and final is it after this one. The caller gathers key head h // (v_heads // k_heads)'s q and k and value head h's v, a and b into block.

Parameters:

Name	Type	Description	Default
rows	int	Rows of a block.	16
key_dim	int	The q and k width.	32
value_dim	int	The v width, or a slice of it: the state's columns are independent. rows * value_dim is even, so the block and y are whole 4-byte words, as a DMA moves them.	16
blocks	int	Blocks to a head.	1

Returns:

Type	Description
ExternalFunction	The kernel.

Raises:

Type	Description
ValueError	A dimension is not a positive integer, or rows * value_dim is odd.

Source code in utils/pyright_stubs/aie/iron/kernels/ssm.py

def gated_delta_rule(
    *, rows: int = 16, key_dim: int = 32, value_dim: int = 16, blocks: int = 1
) -> ExternalFunction:
    """Gated DeltaNet's recurrence for one value head, a block of rows a call.

    ``gated_delta_rule(block, head, final, y, valid, b)``: ``block`` is bf16
    q and k (``rows`` x ``key_dim`` each), v (``rows`` x ``value_dim``),
    then a and b (``rows`` each). ``head`` is float32 ``(dt_bias, A)`` then
    the state before the head's first block, ``(key_dim, value_dim)``.
    ``b`` counts blocks from 0, ``blocks`` to a head, and a head's first
    block loads the state from ``head``. Each block runs its first ``valid``
    rows: q and k are l2-normalized (eps 1e-6) and q scaled by ``1 /
    sqrt(key_dim)``, ``g = A * softplus(a + dt_bias)``, ``beta =
    sigmoid(b)``, ``state = exp(g) * state``, then ``state += outer(k, beta
    * (v - k @ state))`` and ``y = q @ state``, rounded once to bf16. Rows
    from ``valid`` on leave the state alone and write zero ``y``. The state
    stays in the core between blocks, and ``final`` is it after this one.
    The caller gathers key head ``h // (v_heads // k_heads)``'s q and k and
    value head ``h``'s v, a and b into ``block``.

    Args:
        rows: Rows of a block.
        key_dim: The q and k width.
        value_dim: The v width, or a slice of it: the state's columns are
            independent. ``rows * value_dim`` is even, so the block and y
            are whole 4-byte words, as a DMA moves them.
        blocks: Blocks to a head.

    Returns:
        The kernel.

    Raises:
        ValueError: A dimension is not a positive integer, or
            ``rows * value_dim`` is odd.
    """
    geometry = dict(rows=rows, key_dim=key_dim, value_dim=value_dim, blocks=blocks)
    _positive("gated_delta_rule", **geometry)
    if rows * value_dim % 2:
        raise ValueError(
            f"gated_delta_rule: rows * value_dim ({rows} * {value_dim}) must be "
            "even, so the block and y are whole 4-byte words"
        )
    return _make_extern(
        "gated_delta_rule",
        _kernel_source("ssm/gated_delta_rule.cc"),
        [
            np.ndarray[(rows * (2 * key_dim + value_dim + 2),), np.dtype[bfloat16]],
            np.ndarray[(2 + key_dim * value_dim,), np.dtype[np.float32]],
            np.ndarray[(key_dim * value_dim,), np.dtype[np.float32]],
            np.ndarray[(rows * value_dim,), np.dtype[bfloat16]],
            np.int32,
            np.int32,
        ],
        compile_flags=_STRICT_FP_ROLLED
        + [
            f"-DGDR_ROWS={rows}",
            f"-DGDR_KEY_DIM={key_dim}",
            f"-DGDR_VALUE_DIM={value_dim}",
            f"-DGDR_BLOCKS={blocks}",
        ],
        contract=KernelContract(
            # q, k and v rows load 16 bf16s a vector in place, y stores them.
            alignments=((0, 32),) + _vector_args(2) + ((3, 32),),
            trace=Trace.whole_call(),
            roles=(In, In, Out, Out, Param, Param),
            parameter_bindings=((5, 0),),
            call_index=5,
            reference=partial(gated_delta_rule_ref, **geometry),
            sample=partial(gated_delta_rule_sample, **geometry),
            tolerance=Tolerance.bounded(
                partial(
                    _state_bound,
                    gated_delta_rule_ref,
                    final_frac=1e-6,
                    y_ulps=1,
                    **geometry,
                ),
                note="npu2, against the float32 reference over five seeds: final "
                "within 9.4e-7 of its range at key_dim 2, 5.1e-7 at the default; y "
                "one bf16 ulp off in 1 of 10240, as the scalar kernel's; npu1: "
                "unmeasured",
            ),
        ),
    )

FastFlowLM Gemma 4

Kernels extracted from FastFlowLM's Gemma 4 implementation.

AIE2P kernels for prefill attention, decode and the LM head, with sources in aie_kernels/flm_gemma4/. Decode kernels use FlmGemma4DecodeGeometry for model dimensions and buffer layouts.

Factories expose related entry points as attributes named by their C symbols, such as fn.attn_qk_round.

FlmGemma4DecodeGeometry dataclass

FlmGemma4DecodeGeometry(
    model_dim: int,
    num_attn_heads: int,
    num_kv_heads: int,
    intermediate_size: int,
    glu_slice: int,
    pli_d: int,
    dh: int,
    swa_dh: int,
    attn_scale: float,
    swa_attn_scale: float,
    pli_projection_scale: float,
    pli_input_scale: float,
    gelu: bool = True,
    qk_norm: bool = True,
    double_wide_mlp: bool = False,
    name: str = "custom",
)

The model geometry the flm_gemma4/decode_*.cc kernels build for.

The kernels are Gemma 4's fused decode layer, one core per stage, as FastFlowLM builds it. Every field reaches the kernels as a -DFLM_GEMMA4_DECODE_* flag. FLM_GEMMA4_E2B_DECODE and FLM_GEMMA4_E4B_DECODE are the two variants' values, and the factories build only those: the kernels keep assumptions other values break, such as eight query heads, the GELU table, the q/k norm and projections that divide into whole rounds.

flags

flags() -> list[str]

Return the -DFLM_GEMMA4_DECODE_* flags, floats as shortest round-trip literals.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flags(self) -> list[str]:
    """Return the ``-DFLM_GEMMA4_DECODE_*`` flags, floats as shortest round-trip literals."""
    values = {
        "MODEL_DIM": self.model_dim,
        "NUM_ATTN_HEADS": self.num_attn_heads,
        "NUM_KV_HEADS": self.num_kv_heads,
        "INTERMEDIATE_SIZE": self.intermediate_size,
        "GLU_SLICE": self.glu_slice,
        "PLI_D": self.pli_d,
        "DH": self.dh,
        "SWA_DH": self.swa_dh,
        "ATTN_SCALE": f"{float(self.attn_scale)!r}f",
        "SWA_ATTN_SCALE": f"{float(self.swa_attn_scale)!r}f",
        "PLI_PROJECTION_SCALE": f"{float(self.pli_projection_scale)!r}f",
        "PLI_INPUT_SCALE": f"{float(self.pli_input_scale)!r}f",
        "GELU": int(self.gelu),
        "QK_NORM": int(self.qk_norm),
        "DOUBLE_WIDE_MLP": int(self.double_wide_mlp),
    }
    return [f"-DFLM_GEMMA4_DECODE_{k}={v}" for k, v in values.items()]

flm_gemma4_attn_prefill

flm_gemma4_attn_prefill() -> ExternalFunction

Causal flash-attention prefill at a head dim of 512, from flm_gemma4/prefill.cc.

linalg.prefill_fv's algorithm with FastFlowLM's numerics and synchronization: the softmax scales by log2(e) alone, and the epilogue multiplies by a bf16 1 / l.

One core holds 8 query rows of a 128-row round and folds the keys in one 8-row chunk at a time, with one entry point per step of the caller's round, block and chunk loops. The design acquires each k and v chunk and passes it in; no entry point takes a lock. The returned kernel is attn_epilogue, which writes one 64-element chunk of o = y / l; its chunk index is bound to 0. The entry points, all attributes of the kernel:

• attn_rounds(L_begin, L_end, n_out): rounds in the dispatch
• attn_blocks(L_begin, i, n_out): key blocks in round i
• attn_round_begin(prev_m, new_m, c, l, y)
• attn_block_begin(m, prev_m)
• attn_qk_chunk(s, q, k, m, L_begin, window_size, row, col, i, block, j)
• attn_block_mid(s, m, new_m, prev_m, c, l, y)
• attn_fv_chunk(y, s, v, j)
• attn_block_end(prev_m, new_m)
• attn_finalize(l, l_bf16)
• attn_epilogue(o, l_bf16, y, chunk)

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_prefill() -> ExternalFunction:
    """Causal flash-attention prefill at a head dim of 512, from ``flm_gemma4/prefill.cc``.

    [`linalg.prefill_fv`][aie.iron.kernels.linalg.prefill_fv]'s algorithm with
    FastFlowLM's numerics and synchronization: the softmax scales by log2(e)
    alone, and the epilogue multiplies by a bf16 ``1 / l``.

    One core holds 8 query rows of a 128-row round and folds the keys in one
    8-row chunk at a time, with one entry point per step of the caller's
    round, block and chunk loops. The design acquires each k and v chunk and
    passes it in; no entry point takes a lock. The returned kernel is ``attn_epilogue``,
    which writes one 64-element chunk of ``o = y / l``; its chunk index is
    bound to 0. The entry points, all attributes of the kernel:

    - ``attn_rounds(L_begin, L_end, n_out)``: rounds in the dispatch
    - ``attn_blocks(L_begin, i, n_out)``: key blocks in round ``i``
    - ``attn_round_begin(prev_m, new_m, c, l, y)``
    - ``attn_block_begin(m, prev_m)``
    - ``attn_qk_chunk(s, q, k, m, L_begin, window_size, row, col, i, block, j)``
    - ``attn_block_mid(s, m, new_m, prev_m, c, l, y)``
    - ``attn_fv_chunk(y, s, v, j)``
    - ``attn_block_end(prev_m, new_m)``
    - ``attn_finalize(l, l_bf16)``
    - ``attn_epilogue(o, l_bf16, y, chunk)``
    """
    return _prefill(
        "flm_gemma4_attn_prefill",
        _ATTN_PREFILL_DH,
        _ATTN_PREFILL_LQ,
        0,
        flm_gemma4_attn_prefill_ref,
    )

flm_gemma4_swa_prefill

flm_gemma4_swa_prefill() -> ExternalFunction

Sliding-window causal flash-attention prefill at a head dim of 256, from flm_gemma4/prefill.cc.

The sliding-window build of flm_gemma4_attn_prefill: one core holds 16 query rows of a 128-row round and folds in, 16 rows at a time, the keys inside the window its window_size RTP gives. The entry points are those of flm_gemma4_attn_prefill, with attn_blocks(L_begin, window_size, i, n_out) also reading the window. The returned kernel is attn_epilogue, bound to chunk 33: rows 8 to 15, columns 64 to 127 of the (16, 256) y.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_swa_prefill() -> ExternalFunction:
    """Sliding-window causal flash-attention prefill at a head dim of 256, from ``flm_gemma4/prefill.cc``.

    The sliding-window build of [`flm_gemma4_attn_prefill`][aie.iron.kernels.flm_gemma4.flm_gemma4_attn_prefill]:
    one core holds 16 query rows of a 128-row round and folds in, 16 rows at
    a time, the keys inside the window its ``window_size`` RTP gives. The
    entry points are those of ``flm_gemma4_attn_prefill``, with
    ``attn_blocks(L_begin, window_size, i, n_out)`` also reading the window.
    The returned kernel is ``attn_epilogue``, bound to chunk 33: rows 8 to 15,
    columns 64 to 127 of the (16, 256) ``y``.
    """
    return _prefill(
        "flm_gemma4_swa_prefill",
        _SWA_PREFILL_DH,
        _SWA_PREFILL_LQ,
        _SWA_PREFILL_CHUNK,
        flm_gemma4_swa_prefill_ref,
    )

flm_gemma4_prefill_block_begin

flm_gemma4_prefill_block_begin(
    *, head_dim: int = 512
) -> ExternalFunction

attn_block_begin of the prefill kernels: row j of m is prev_m[j].

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_block_begin(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_block_begin`` of the prefill kernels: row ``j`` of ``m`` is ``prev_m[j]``."""
    return _prefill_sibling(
        "attn_block_begin",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, In),
            reference=flm_gemma4_prefill_block_begin_ref,
            tolerance=Tolerance.exact(note="bf16 broadcast"),
            ops_per_call=0,
        ),
    )

flm_gemma4_prefill_block_begin_ref

flm_gemma4_prefill_block_begin_ref(prev_m)

Numpy reference for flm_gemma4_prefill_block_begin.

m is (LQ, LK) with LK == LQ in both builds.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_block_begin_ref(prev_m):
    """Numpy reference for [`flm_gemma4_prefill_block_begin`][aie.iron.kernels.flm_gemma4.flm_gemma4_prefill_block_begin].

    ``m`` is ``(LQ, LK)`` with ``LK == LQ`` in both builds.
    """
    prev_m = np.asarray(prev_m)
    return np.repeat(prev_m, prev_m.shape[-1], axis=-1)

flm_gemma4_prefill_block_end

flm_gemma4_prefill_block_end(
    *, head_dim: int = 512
) -> ExternalFunction

attn_block_end of the prefill kernels: prev_m = new_m.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_block_end(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_block_end`` of the prefill kernels: ``prev_m = new_m``."""
    return _prefill_sibling(
        "attn_block_end",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, In),
            reference=flm_gemma4_prefill_block_end_ref,
            tolerance=Tolerance.exact(note="bf16 copy"),
            ops_per_call=0,
        ),
    )

flm_gemma4_prefill_block_end_ref

flm_gemma4_prefill_block_end_ref(new_m)

Numpy reference for flm_gemma4_prefill_block_end.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_block_end_ref(new_m):
    """Numpy reference for [`flm_gemma4_prefill_block_end`][aie.iron.kernels.flm_gemma4.flm_gemma4_prefill_block_end]."""
    return np.array(new_m, copy=True)

flm_gemma4_prefill_rounds_ref

flm_gemma4_prefill_rounds_ref(l_begin, l_end)

(L_end >> 7) - (L_begin >> 7): the 128-row rounds a dispatch spans.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_rounds_ref(l_begin, l_end):
    """``(L_end >> 7) - (L_begin >> 7)``: the 128-row rounds a dispatch spans."""
    n = (_word0(l_end) >> 7) - (_word0(l_begin) >> 7)
    return n.reshape(-1, 1)

flm_gemma4_prefill_rounds

flm_gemma4_prefill_rounds(
    *, head_dim: int = 512
) -> ExternalFunction

attn_rounds of the prefill kernels: the rounds between L_begin and L_end.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_rounds(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_rounds`` of the prefill kernels: the rounds between ``L_begin`` and ``L_end``."""
    return _prefill_sibling(
        "attn_rounds",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            sample=_prefill_rounds_sample,
            reference=flm_gemma4_prefill_rounds_ref,
            tolerance=Tolerance.exact(note="integer shifts and a subtraction"),
            out_valid=1,
            ops_per_call=1,
        ),
    )

flm_gemma4_prefill_blocks_ref

flm_gemma4_prefill_blocks_ref(
    l_begin, i, *, window_size=None
)

Key blocks round i folds in: all up to its own, or those in the window.

window_size is the sliding-window build's RTP buffer, None for the head-dim-512 build.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_blocks_ref(l_begin, i, *, window_size=None):
    """Key blocks round ``i`` folds in: all up to its own, or those in the window.

    ``window_size`` is the sliding-window build's RTP buffer, ``None`` for the
    head-dim-512 build.
    """
    begin = _word0(l_begin)
    if window_size is None:
        n = (begin >> 7) + np.int32(i) + 1
    else:
        q = begin + np.int32(i) * 128
        k = np.maximum(q - _word0(window_size), 0)
        n = ((q - k) >> 7) + 1
    return n.astype(np.int32).reshape(-1, 1)

flm_gemma4_prefill_blocks

flm_gemma4_prefill_blocks(
    *, head_dim: int = 512
) -> ExternalFunction

attn_blocks of the prefill kernels: the key blocks round i folds in.

At a head dim of 256, the sliding-window build, it also reads the window size: attn_blocks(L_begin, window_size, i, n_out).

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_blocks(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_blocks`` of the prefill kernels: the key blocks round ``i`` folds in.

    At a head dim of 256, the sliding-window build, it also reads the window
    size: ``attn_blocks(L_begin, window_size, i, n_out)``.
    """
    window = head_dim == 256
    return _prefill_sibling(
        "attn_blocks",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Param, Out) if window else (In, Param, Out),
            sample=partial(_prefill_blocks_sample, window=window),
            reference=(
                _prefill_swa_blocks_ref if window else flm_gemma4_prefill_blocks_ref
            ),
            tolerance=Tolerance.exact(note="integer shifts, adds and a clamp"),
            out_valid=1,
            ops_per_call=1,
        ),
    )

flm_gemma4_prefill_finalize

flm_gemma4_prefill_finalize(
    *, head_dim: int = 512
) -> ExternalFunction

attn_finalize of the prefill kernels: inv_l = bf16(1 / l) per query row.

Parameters:

Name	Type	Description	Default
head_dim	int	512 builds the global kernel's 8 rows; 256 builds the sliding-window kernel's 16 rows.	512

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_finalize(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_finalize`` of the prefill kernels: ``inv_l = bf16(1 / l)`` per query row.

    Args:
        head_dim: 512 builds the global kernel's 8 rows; 256 builds the
            sliding-window kernel's 16 rows.
    """
    return _prefill_sibling(
        "attn_finalize",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_prefill_finalize_ref,
            sample=partial(
                _prefill_finalize_sample,
                lq=_ATTN_PREFILL_LQ if head_dim == 512 else _SWA_PREFILL_LQ,
            ),
            # Floor narrowing lands at most one bf16 step from the nearest-even
            # cast of 1 / l if aie::inv errs by less than half a step (2**-9
            # relative).
            tolerance=Tolerance.bf16_ulps(
                1,
                note="floor bf16 narrowing of AIE2P's scalar inv instruction, "
                "whose error the AIE-ML v2 intrinsics guide does not document",
            ),
        ),
    )

flm_gemma4_prefill_finalize_ref

flm_gemma4_prefill_finalize_ref(row_sums)

Numpy reference for flm_gemma4_prefill_finalize: 1 / l.

The tolerance covers the kernel's hardware reciprocal and its floor narrowing to bf16.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_finalize_ref(row_sums):
    """Numpy reference for [`flm_gemma4_prefill_finalize`][aie.iron.kernels.flm_gemma4.flm_gemma4_prefill_finalize]: ``1 / l``.

    The tolerance covers the kernel's hardware reciprocal and its floor
    narrowing to bf16.
    """
    return 1.0 / np.asarray(row_sums, dtype=np.float32)

flm_gemma4_prefill_round_begin_ref

flm_gemma4_prefill_round_begin_ref(*, head_dim=512)

Numpy reference for flm_gemma4_prefill_round_begin.

The bytes of y | c | l | prev_m | new_m: zeros, ones, zeros and two rows of NEG_INF.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_round_begin_ref(*, head_dim=512):
    """Numpy reference for [`flm_gemma4_prefill_round_begin`][aie.iron.kernels.flm_gemma4.flm_gemma4_prefill_round_begin].

    The bytes of ``y | c | l | prev_m | new_m``: zeros, ones, zeros and two
    rows of NEG_INF.
    """
    lq = _prefill_lq(head_dim)
    return np.concatenate(
        [
            np.zeros(lq * head_dim, np.float32).view(np.uint8),
            np.ones(lq, np.float32).view(np.uint8),
            np.zeros(lq, np.float32).view(np.uint8),
            np.full(2 * lq, _PREFILL_NEG_INF, bfloat16).view(np.uint8),
        ]
    )

flm_gemma4_prefill_round_begin

flm_gemma4_prefill_round_begin(
    *, head_dim: int = 512
) -> ExternalFunction

attn_round_begin of the prefill kernels, its five outputs packed as bytes.

out is y | c | l | prev_m | new_m.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_round_begin(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_round_begin`` of the prefill kernels, its five outputs packed as bytes.

    ``out`` is ``y | c | l | prev_m | new_m``.
    """
    lq = _prefill_lq(head_dim)
    n = 4 * lq * head_dim + 12 * lq
    return _prefill_core(
        "attn_round_begin_core",
        head_dim,
        [np.ndarray[(n,), np.dtype[np.uint8]]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out,),
            reference=partial(flm_gemma4_prefill_round_begin_ref, head_dim=head_dim),
            tolerance=Tolerance.exact(note="constant stores"),
            ops_per_call=0,
        ),
    )

flm_gemma4_prefill_qk_core_ref

flm_gemma4_prefill_qk_core_ref(
    q,
    k,
    m,
    inner_k,
    inner_q,
    inner_k_current,
    *,
    head_dim=512
)

Numpy reference for flm_gemma4_prefill_qk_core: s | m.

s = q @ k.T, NEG_INF where masked, and m folds in s lanewise. The tolerance covers the bfp16 operands and the bf16 narrowing.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_qk_core_ref(
    q, k, m, inner_k, inner_q, inner_k_current, *, head_dim=512
):
    """Numpy reference for [`flm_gemma4_prefill_qk_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_prefill_qk_core]: ``s | m``.

    ``s = q @ k.T``, NEG_INF where masked, and ``m`` folds in ``s``
    lanewise. The tolerance covers the bfp16 operands and the bf16 narrowing.
    """
    lq = lk = _prefill_lq(head_dim)
    q, k = _tiles(q, lq, head_dim), _tiles(k, lk, head_dim, col_major=True)
    s = q @ k.transpose(0, 2, 1)
    keep = _prefill_keep(lq, lk, inner_k, inner_q, inner_k_current)
    s = np.where(keep, s, _PREFILL_NEG_INF)
    m = np.maximum(np.asarray(m, np.float64).reshape(-1, lq, lk), s)
    return np.concatenate([s.reshape(len(s), -1), m.reshape(len(m), -1)], axis=1)

flm_gemma4_prefill_qk_core

flm_gemma4_prefill_qk_core(
    *, head_dim: int = 512
) -> ExternalFunction

attn_qk_chunk's arithmetic.

attn_qk_core(q, k, m, inner_k, inner_q, inner_k_current, out): s = q @ k.T through the bfp16-emulated mmul, masked, and m folded with it, into out = s | m. The three positions are apply_mask_and_get_max's, which attn_qk_chunk derives from its RTPs. q and k arrive in the mmul's 8x8 tiles, k's tiles in column-major order.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_qk_core(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_qk_chunk``'s arithmetic.

    ``attn_qk_core(q, k, m, inner_k, inner_q, inner_k_current, out)``:
    ``s = q @ k.T`` through the bfp16-emulated mmul, masked, and ``m``
    folded with it, into ``out = s | m``. The three positions are
    ``apply_mask_and_get_max``'s, which attn_qk_chunk derives from its RTPs.
    q and k arrive in the mmul's 8x8 tiles, k's tiles in column-major order.
    """
    lq = lk = _prefill_lq(head_dim)
    kv = np.ndarray[(lq * head_dim,), _BF16]
    m = np.ndarray[(lq * lk,), _BF16]
    return _prefill_core(
        "attn_qk_core",
        head_dim,
        [kv, kv, m, np.int32, np.int32, np.int32, np.ndarray[(2 * lq * lk,), _BF16]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Param, Param, Param, Out),
            sample=partial(_prefill_qk_sample, head_dim=head_dim),
            reference=partial(flm_gemma4_prefill_qk_core_ref, head_dim=head_dim),
            tolerance=Tolerance.bounded(
                partial(_prefill_qk_bound, head_dim=head_dim),
                note="bfp16ebs8 operand rounding (zero on the sample's grid), "
                "float32 accumulation, and the bf16 narrowing of s",
            ),
            ops_per_call=2 * lq * lk * head_dim,
        ),
    )

flm_gemma4_prefill_fv_core_ref

flm_gemma4_prefill_fv_core_ref(y, sv, *, head_dim=512)

Numpy reference for flm_gemma4_prefill_fv_core: y + s @ v.

y and s are in the mmul's 8x8 tiles, v in column-major tile order; the tolerance covers the bfp16 operands.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_fv_core_ref(y, sv, *, head_dim=512):
    """Numpy reference for [`flm_gemma4_prefill_fv_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_prefill_fv_core]: ``y + s @ v``.

    ``y`` and ``s`` are in the mmul's 8x8 tiles, ``v`` in column-major tile
    order; the tolerance covers the bfp16 operands.
    """
    lq = lk = _prefill_lq(head_dim)
    s, v = _prefill_sv(sv, lq, lk, head_dim)
    return _to_tiles(_tiles(y, lq, head_dim) + s @ v)

flm_gemma4_prefill_fv_core

flm_gemma4_prefill_fv_core(
    *, head_dim: int = 512
) -> ExternalFunction

attn_fv_chunk's arithmetic.

attn_fv_core(y, sv, y_out): y_out = y + s @ v through flm_attn_fv's bfp16-emulated mmul, with sv = s | v. y and s are in the mmul's 8x8 tiles, v's tiles in column-major order.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_fv_core(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_fv_chunk``'s arithmetic.

    ``attn_fv_core(y, sv, y_out)``: ``y_out = y + s @ v`` through
    ``flm_attn_fv``'s bfp16-emulated mmul, with ``sv = s | v``. y and s are in
    the mmul's 8x8 tiles, v's tiles in column-major order.
    """
    lq = lk = _prefill_lq(head_dim)
    y = np.ndarray[(lq * head_dim,), _F32]
    return _prefill_core(
        "attn_fv_core",
        head_dim,
        [y, np.ndarray[(lq * lk + lk * head_dim,), _BF16], y],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            sample=partial(_prefill_fv_sample, head_dim=head_dim),
            reference=partial(flm_gemma4_prefill_fv_core_ref, head_dim=head_dim),
            tolerance=Tolerance.bounded(
                partial(_prefill_fv_core_bound, head_dim=head_dim),
                note="bfp16ebs8 operand rounding (zero on the sample's grid) "
                "and float32 accumulation",
            ),
            ops_per_call=2 * lq * lk * head_dim,
        ),
    )

flm_gemma4_prefill_block_mid_core_ref

flm_gemma4_prefill_block_mid_core_ref(
    in_bf16, in_f32, *, head_dim=512
)

Numpy reference for flm_gemma4_prefill_block_mid_core.

(s | new_m, y | l | c): new_m the row max of m, p = 2**((s - new_m) * log2e) and c = 2**((prev_m - new_m) * log2e) with FastFlowLM's bf16 log2(e), l = sum(p) + c * l and y = c * y.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_block_mid_core_ref(in_bf16, in_f32, *, head_dim=512):
    """Numpy reference for [`flm_gemma4_prefill_block_mid_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_prefill_block_mid_core].

    ``(s | new_m, y | l | c)``: ``new_m`` the row max of ``m``,
    ``p = 2**((s - new_m) * log2e)`` and ``c = 2**((prev_m - new_m) * log2e)``
    with FastFlowLM's bf16 log2(e), ``l = sum(p) + c * l`` and ``y = c * y``.
    """
    lq = _prefill_lq(head_dim)
    p, new_m, y, row_sums, c, _, _ = _prefill_block_mid_math(in_bf16, in_f32, head_dim)
    return (
        np.concatenate([_prefill_block_mid_s(p, lq), new_m], axis=1),
        np.concatenate([y, row_sums, c], axis=1),
    )

flm_gemma4_prefill_block_mid_core

flm_gemma4_prefill_block_mid_core(
    *, head_dim: int = 512
) -> ExternalFunction

attn_block_mid of the prefill kernels, its state passed in and out.

attn_block_mid_core(in_bf16, in_f32, out_bf16, out_f32) with in_bf16 = s | m | prev_m, in_f32 = y | l, out_bf16 = s | new_m and out_f32 = y | l | c: the row max, FastFlowLM's softmax, the correction c and the rescaled l and y.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_prefill_block_mid_core(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_block_mid`` of the prefill kernels, its state passed in and out.

    ``attn_block_mid_core(in_bf16, in_f32, out_bf16, out_f32)`` with
    ``in_bf16 = s | m | prev_m``, ``in_f32 = y | l``, ``out_bf16 = s | new_m``
    and ``out_f32 = y | l | c``: the row max, FastFlowLM's softmax, the
    correction ``c`` and the rescaled ``l`` and ``y``.
    """
    lq = lk = _prefill_lq(head_dim)
    return _prefill_core(
        "attn_block_mid_core",
        head_dim,
        [
            np.ndarray[(lq * 128 + lq * lk + lq,), _BF16],
            np.ndarray[(lq * head_dim + lq,), _F32],
            np.ndarray[(lq * 128 + lq,), _BF16],
            np.ndarray[(lq * head_dim + 2 * lq,), _F32],
        ],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out, Out),
            sample=partial(_prefill_block_mid_sample, head_dim=head_dim),
            reference=partial(flm_gemma4_prefill_block_mid_core_ref, head_dim=head_dim),
            tolerance=Tolerance.bounded(
                partial(_prefill_block_mid_bound, head_dim=head_dim),
                note="aie::exp2's 6.98% (linalg.mha_softmax), bf16 rounding of "
                "the exp2 arguments, in the kernel (zero on the sample's grid) "
                "and inside aie::exp2 (measured on npu2), bfp16ebs8 "
                "rounding of p in the row sums, the emulated float32 multiply",
            ),
            # Per score: subtract, scale, exp2, sum; per y element: one multiply.
            ops_per_call=4 * lq * 128 + lq * head_dim,
        ),
    )

flm_gemma4_decode_glu

flm_gemma4_decode_glu(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _GluKernel

Build the decode layer's gated linear unit, from flm_gemma4/decode_glu.cc.

glu_chunk(hid, x, i) turns one glu_slice of gate and up values into glu_slice // 2 activations at slice i of hid, which holds intermediate_size of them (twice that with double_wide_mlp). glu_emit(y_out, hid, j) copies slice j back out for the down projection.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_glu(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _GluKernel:
    """Build the decode layer's gated linear unit, from ``flm_gemma4/decode_glu.cc``.

    ``glu_chunk(hid, x, i)`` turns one ``glu_slice`` of gate and up values
    into ``glu_slice // 2`` activations at slice ``i`` of
    ``hid``, which holds ``intermediate_size`` of them (twice that with
    ``double_wide_mlp``). ``glu_emit(y_out, hid, j)`` copies slice ``j`` back
    out for the down projection.

    Args:
        geometry: The model the kernel builds for.
    """
    bf16 = np.dtype[bfloat16]
    hid = geometry.intermediate_size * (2 if geometry.double_wide_mlp else 1)
    hid_ty = np.ndarray[(hid,), bf16]
    x_ty = np.ndarray[(geometry.glu_slice,), bf16]
    y_ty = np.ndarray[(geometry.glu_slice // 2,), bf16]
    fn = _decode_kernel(
        "glu",
        "glu_chunk",
        [hid_ty, x_ty, np.int32],
        (InOut, In, Param),
        geometry,
        lut=True,
        cls=_GluKernel,
    )
    fn.glu_emit = fn.object_file.bind("glu_emit", [y_ty, hid_ty, np.int32])
    return fn

flm_gemma4_glu_core_ref

flm_gemma4_glu_core_ref(x)

Numpy reference for flm_gemma4_glu_core: gelu(gate) * up.

x holds, per call, the up half then the gate half. GELU is the tanh form Gemma 4 uses; the tolerance covers the kernel's table.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_glu_core_ref(x):
    """Numpy reference for [`flm_gemma4_glu_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_glu_core]: ``gelu(gate) * up``.

    ``x`` holds, per call, the up half then the gate half. GELU is the tanh
    form Gemma 4 uses; the tolerance covers the kernel's table.
    """
    up, gate = _glu_split(x)
    return (_gelu_tanh(gate) * up).astype(np.float32)

flm_gemma4_glu_core

flm_gemma4_glu_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

pseduo_glu of flm_gemma4_decode_glu: gelu(gate) * up for one slice.

glu_core(x, y): x is glu_slice bf16, the up half then the gate half; y is glu_slice // 2 bf16. GELU goes through aie_runtime_lib's table (getGeluBf16).

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_glu_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``pseduo_glu`` of [`flm_gemma4_decode_glu`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_glu]: ``gelu(gate) * up`` for one slice.

    ``glu_core(x, y)``: ``x`` is ``glu_slice`` bf16, the up half then the
    gate half; ``y`` is ``glu_slice // 2`` bf16. GELU goes through
    aie_runtime_lib's table (``getGeluBf16``).

    Args:
        geometry: The model the kernel builds for.
    """
    n = geometry.glu_slice
    return _wrapper(
        flm_gemma4_decode_glu(geometry=geometry),
        "decode_glu_core.cc",
        "glu_core",
        [np.ndarray[(n,), _BF16], np.ndarray[(n // 2,), _BF16]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_glu_core_ref,
            sample=partial(_glu_core_sample, n=n),
            tolerance=Tolerance.bounded(
                _glu_core_bound,
                note="the GELU table's error against tanh-GELU, measured on the "
                "host from aie_runtime_lib's table, plus a bf16 ulp of its slope "
                "and the floor bf16 stores of the GELU and the product",
            ),
            uses_lut=True,
            # Per output: the table line's multiply and add, then the product.
            ops_per_call=3 * (n // 2),
        ),
    )

flm_gemma4_decode_attn_kv

flm_gemma4_decode_attn_kv(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _AttnKvKernel

Build the global-attention score-times-value core for one KV head, from flm_gemma4/decode_attn_kv.cc.

Accumulate scores and running-max corrections into the softmax denominator l and float32 accumulator y, then narrow and normalize y in place. The returned kernel is attn_kv_begin, which zeroes y and l. Call attn_kv_round per round, then attn_kv_finish(y, l). Finish leaves bf16 values in the first half of the float32 allocation, in [dh/8, head, 8] order. y needs a 64-byte-aligned base pointer.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; num_kv_heads must be 1.	FLM_GEMMA4_E2B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_attn_kv(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _AttnKvKernel:
    """Build the global-attention score-times-value core for one KV head, from ``flm_gemma4/decode_attn_kv.cc``.

    Accumulate scores and running-max corrections into the softmax denominator
    ``l`` and float32 accumulator ``y``, then narrow and normalize ``y`` in place.
    The returned kernel is ``attn_kv_begin``, which zeroes ``y`` and ``l``.
    Call ``attn_kv_round`` per round, then ``attn_kv_finish(y, l)``.
    Finish leaves bf16 values in the first half of the float32 allocation,
    in [dh/8, head, 8] order. ``y`` needs a 64-byte-aligned base pointer.

    Args:
        geometry: The model the kernel builds for; ``num_kv_heads`` must be 1.
    """
    _decode_kv_heads("flm_gemma4_decode_attn_kv", geometry, 1)
    t = _decode_attn_types(geometry, geometry.dh, geometry.dh)
    fn = _decode_kernel(
        "attn_kv",
        "attn_kv_begin",
        [t["y"], t["l"]],
        (Out, Out),
        geometry,
        lut=True,
        cls=_AttnKvKernel,
    )
    bind = fn.object_file.bind
    fn.attn_kv_round = bind("attn_kv_round", [t["s"], t["kv"], t["y"], t["l"]])
    fn.attn_kv_finish = bind("attn_kv_finish", [t["y"], t["l"]])
    return fn

flm_gemma4_attn_kv_core

flm_gemma4_attn_kv_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One attn_kv_round of flm_gemma4_decode_attn_kv with caller-owned state.

attn_kv_round_core(sv, ly_in, ly_out): ly_out is ly_in with l = c * l + rowsum(s) and y = c * y + s @ v per query head. sv holds s, then v. ly holds y, then l in a 16-float slot.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_kv_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """One ``attn_kv_round`` of [`flm_gemma4_decode_attn_kv`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv] with caller-owned state.

    ``attn_kv_round_core(sv, ly_in, ly_out)``: ``ly_out`` is ``ly_in`` with
    ``l = c * l + rowsum(s)`` and ``y = c * y + s @ v`` per query head. ``sv``
    holds ``s``, then ``v``. ``ly`` holds ``y``, then ``l`` in a 16-float slot.
    """
    base = flm_gemma4_decode_attn_kv(geometry=geometry)
    dh = geometry.dh
    return _attn_kv_round_core(
        base,
        "decode_attn_kv_core.cc",
        "attn_kv_round_core",
        dh=dh,
        kv_width=dh,
        impl="1x8x1",
        rows=8,
        ref=partial(flm_gemma4_attn_kv_core_ref, dh=dh),
    )

flm_gemma4_attn_kv_core_ref

flm_gemma4_attn_kv_core_ref(sv, ly, *, dh=512)

Numpy reference for flm_gemma4_attn_kv_core, in float64.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_kv_core_ref(sv, ly, *, dh=512):
    """Numpy reference for [`flm_gemma4_attn_kv_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_attn_kv_core], in float64."""
    return _attn_kv_round_ref(sv, ly, dh=dh, impl="1x8x1")

flm_gemma4_decode_attn_kv_kvh2

flm_gemma4_decode_attn_kv_kvh2(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> _AttnKvKvh2Kernel

Build the global-attention score-times-value core for two KV heads, from flm_gemma4/decode_attn_kv_kvh2.cc.

The two-KV-head (E4B) sibling of flm_gemma4_decode_attn_kv: each round starts with attn_kv_s_begin, which folds the scores into l and rescales y, then one attn_kv_v_half per KV head.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; num_kv_heads must be 2.	FLM_GEMMA4_E4B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_attn_kv_kvh2(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> _AttnKvKvh2Kernel:
    """Build the global-attention score-times-value core for two KV heads, from ``flm_gemma4/decode_attn_kv_kvh2.cc``.

    The two-KV-head (E4B) sibling of
    [`flm_gemma4_decode_attn_kv`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv]:
    each round starts with ``attn_kv_s_begin``, which folds the
    scores into ``l`` and rescales ``y``, then one ``attn_kv_v_half`` per
    KV head.

    Args:
        geometry: The model the kernel builds for; ``num_kv_heads`` must be 2.
    """
    _decode_kv_heads("flm_gemma4_decode_attn_kv_kvh2", geometry, 2)
    t = _decode_attn_types(geometry, geometry.dh, geometry.dh)
    fn = _decode_kernel(
        "attn_kv_kvh2",
        "attn_kv_begin",
        [t["y"], t["l"]],
        (Out, Out),
        geometry,
        lut=True,
        cls=_AttnKvKvh2Kernel,
    )
    bind = fn.object_file.bind
    fn.attn_kv_s_begin = bind("attn_kv_s_begin", [t["s"], t["y"], t["l"]])
    fn.attn_kv_v_half = bind("attn_kv_v_half", [t["s"], t["kv"], t["y"], np.int32])
    fn.attn_kv_finish = bind("attn_kv_finish", [t["y"], t["l"]])
    return fn

flm_gemma4_attn_kv_kvh2_core

flm_gemma4_attn_kv_kvh2_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction

attn_kv_s_begin, then one KV head's attn_kv_v_half, with caller-owned state.

The entry points are those of flm_gemma4_decode_attn_kv_kvh2. attn_kv_kvh2_round_core(sv, ly_in, ly_out, kv_head): ly_out is ly_in with l = c * l + rowsum(s) and y = c * y, plus s @ v for query heads 4 * kv_head to 4 * kv_head + 3. sv holds s, then v. ly holds y, then l in a 16-float slot.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_kv_kvh2_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction:
    """``attn_kv_s_begin``, then one KV head's ``attn_kv_v_half``, with caller-owned state.

    The entry points are those of
    [`flm_gemma4_decode_attn_kv_kvh2`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv_kvh2].
    ``attn_kv_kvh2_round_core(sv, ly_in, ly_out, kv_head)``: ``ly_out`` is
    ``ly_in`` with ``l = c * l + rowsum(s)`` and ``y = c * y``, plus ``s @ v``
    for query heads ``4 * kv_head`` to ``4 * kv_head + 3``. ``sv`` holds
    ``s``, then ``v``. ``ly`` holds ``y``, then ``l`` in a 16-float slot.
    """
    base = flm_gemma4_decode_attn_kv_kvh2(geometry=geometry)
    dh = geometry.dh
    return _attn_kv_round_core(
        base,
        "decode_attn_kv_kvh2_core.cc",
        "attn_kv_kvh2_round_core",
        dh=dh,
        kv_width=dh,
        impl="kvh2",
        rows=4,
        ref=partial(flm_gemma4_attn_kv_kvh2_core_ref, dh=dh),
    )

flm_gemma4_attn_kv_kvh2_core_ref

flm_gemma4_attn_kv_kvh2_core_ref(
    sv, ly, kv_head, *, dh=512
)

Numpy reference for flm_gemma4_attn_kv_kvh2_core, in float64.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_kv_kvh2_core_ref(sv, ly, kv_head, *, dh=512):
    """Numpy reference for [`flm_gemma4_attn_kv_kvh2_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_attn_kv_kvh2_core], in float64."""
    return _attn_kv_round_ref(sv, ly, int(kv_head), dh=dh, impl="kvh2")

flm_gemma4_decode_attn_qk

flm_gemma4_decode_attn_qk(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _AttnQkKernel

Build an attention query-times-key core, from flm_gemma4/decode_attn_qk.cc.

The other half of flm_gemma4_decode_attn_kv: it multiplies the query heads into each round's 16 key rows, masks the keys past the sequence length, keeps the running row maximum m and writes the exponentiated scores followed by float32 corrections in s. The returned kernel is attn_qk_begin, which initializes m to the lowest finite bf16 value. Call attn_qk_round once per round. The global build serves one KV head; the sliding-window build serves geometry.num_kv_heads.

Parameters:

Name	Type	Description	Default
sliding_window	bool	Build for the sliding-window layers' head dim.	False
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; without sliding_window num_kv_heads must be 1.	FLM_GEMMA4_E2B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_attn_qk(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> _AttnQkKernel:
    """Build an attention query-times-key core, from ``flm_gemma4/decode_attn_qk.cc``.

    The other half of [`flm_gemma4_decode_attn_kv`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv]:
    it multiplies the query heads into each round's 16 key rows, masks the
    keys past the sequence length, keeps the running row maximum ``m`` and
    writes the exponentiated scores followed by float32 corrections in ``s``.
    The returned kernel is ``attn_qk_begin``, which initializes ``m`` to the
    lowest finite bf16 value. Call ``attn_qk_round`` once per round.
    The global build serves one KV head; the sliding-window build serves
    ``geometry.num_kv_heads``.

    Args:
        sliding_window: Build for the sliding-window layers' head dim.
        geometry: The model the kernel builds for; without ``sliding_window``
            ``num_kv_heads`` must be 1.
    """
    if not sliding_window:
        _decode_kv_heads("flm_gemma4_decode_attn_qk", geometry, 1)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    t = _decode_attn_types(geometry, dh, geometry.num_kv_heads * dh)
    fn = _decode_kernel(
        "attn_qk",
        "attn_qk_begin",
        [t["m"]],
        (Out,),
        geometry,
        flags=(f"-DFLM_GEMMA4_DECODE_ATTN_QK_SWA={int(sliding_window)}",),
        cls=_AttnQkKernel,
    )
    fn.attn_qk_round = fn.object_file.bind(
        "attn_qk_round",
        [t["q"], t["kv"], t["s"], t["m"], t["c"], np.int32, np.int32],
    )
    return fn

flm_gemma4_decode_attn_qk_kvh2

flm_gemma4_decode_attn_qk_kvh2(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> _AttnQkKvh2Kernel

Build the global-attention query-times-key core for two KV heads, from flm_gemma4/decode_attn_qk_kvh2.cc.

The two-KV-head (E4B) sibling of flm_gemma4_decode_attn_qk: call attn_qk_half once per KV head per round, then attn_qk_store_c, which copies the corrections into the scores.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; num_kv_heads must be 2.	FLM_GEMMA4_E4B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_attn_qk_kvh2(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> _AttnQkKvh2Kernel:
    """Build the global-attention query-times-key core for two KV heads, from ``flm_gemma4/decode_attn_qk_kvh2.cc``.

    The two-KV-head (E4B) sibling of
    [`flm_gemma4_decode_attn_qk`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_attn_qk]:
    call ``attn_qk_half`` once per KV head per round, then
    ``attn_qk_store_c``, which copies the corrections into the scores.

    Args:
        geometry: The model the kernel builds for; ``num_kv_heads`` must be 2.
    """
    _decode_kv_heads("flm_gemma4_decode_attn_qk_kvh2", geometry, 2)
    t = _decode_attn_types(geometry, geometry.dh, geometry.dh)
    fn = _decode_kernel(
        "attn_qk_kvh2",
        "attn_qk_begin",
        [t["m"]],
        (Out,),
        geometry,
        cls=_AttnQkKvh2Kernel,
    )
    bind = fn.object_file.bind
    fn.attn_qk_half = bind(
        "attn_qk_half",
        [t["q"], t["kv"], t["s"], t["m"], t["c"], np.int32, np.int32, np.int32],
    )
    fn.attn_qk_store_c = bind("attn_qk_store_c", [t["s"], t["c"]])
    return fn

flm_gemma4_attn_qk_core_ref

flm_gemma4_attn_qk_core_ref(qm, k, iter, L0)

Numpy reference for flm_gemma4_attn_qk_core.

Keys j with j < L0 - 16 * iter are valid. Per query row the new running max is max(m_in, valid scores); s holds exp(score - max) for valid keys and 0 for the others, then the 16 running maxima (rows 8 to 15 copied from m_in). The second output is c = exp(m_in - max). Both exponents are clamped at -87, as the kernel does.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_qk_core_ref(qm, k, iter, L0):
    """Numpy reference for [`flm_gemma4_attn_qk_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_attn_qk_core].

    Keys ``j`` with ``j < L0 - 16 * iter`` are valid. Per query row the new
    running max is ``max(m_in, valid scores)``; ``s`` holds ``exp(score -
    max)`` for valid keys and 0 for the others, then the 16 running maxima
    (rows 8 to 15 copied from ``m_in``). The second output is ``c = exp(m_in
    - max)``. Both exponents are clamped at -87, as the kernel does.
    """
    _, _, m_in, m_new, _, e, c = _attn_qk_parts(qm, k, iter, L0)
    m_out = m_in.copy()
    m_out[:, :8] = m_new
    return np.concatenate([e.reshape(len(e), -1), m_out], axis=1), c

flm_gemma4_attn_qk_core

flm_gemma4_attn_qk_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One attn_qk_round of flm_gemma4_decode_attn_qk, with caller-owned state.

attn_qk_round_core(qm, k, s, c, iter, L0) takes the running max after q in qm and writes it after the scores in s. c holds the eight float32 correction factors separately from s.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_qk_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> ExternalFunction:
    """One ``attn_qk_round`` of [`flm_gemma4_decode_attn_qk`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_attn_qk], with caller-owned state.

    ``attn_qk_round_core(qm, k, s, c, iter, L0)`` takes the running max
    after q in ``qm`` and writes it after the scores in ``s``. ``c`` holds the
    eight float32 correction factors separately from ``s``.
    """
    base = flm_gemma4_decode_attn_qk(sliding_window=sliding_window, geometry=geometry)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    kv_heads = geometry.num_kv_heads
    return _wrapper(
        base,
        "decode_attn_qk_core.cc",
        "attn_qk_round_core",
        _attn_qk_core_types(geometry, dh, kv_heads),
        _attn_qk_core_contract(dh, kv_heads, flm_gemma4_attn_qk_core_ref),
    )

flm_gemma4_attn_qk_kvh2_core

flm_gemma4_attn_qk_kvh2_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction

One round of flm_gemma4_decode_attn_qk_kvh2, with caller-owned state.

attn_qk_kvh2_round_core(qm, k, s, c, iter, L0) runs attn_qk_half for both KV heads, whose key buffers k holds back to back. Otherwise as flm_gemma4_attn_qk_core.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_qk_kvh2_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction:
    """One round of [`flm_gemma4_decode_attn_qk_kvh2`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_attn_qk_kvh2], with caller-owned state.

    ``attn_qk_kvh2_round_core(qm, k, s, c, iter, L0)`` runs
    ``attn_qk_half`` for both KV heads, whose key buffers ``k`` holds back to
    back. Otherwise as
    [`flm_gemma4_attn_qk_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_attn_qk_core].
    """
    base = flm_gemma4_decode_attn_qk_kvh2(geometry=geometry)
    return _wrapper(
        base,
        "decode_attn_qk_kvh2_core.cc",
        "attn_qk_kvh2_round_core",
        _attn_qk_core_types(geometry, geometry.dh, 2),
        _attn_qk_core_contract(geometry.dh, 2, flm_gemma4_attn_qk_kvh2_core_ref),
    )

flm_gemma4_attn_qk_kvh2_core_ref

flm_gemma4_attn_qk_kvh2_core_ref(qm, k, iter, L0)

Numpy reference for flm_gemma4_attn_qk_kvh2_core; see flm_gemma4_attn_qk_core_ref.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_qk_kvh2_core_ref(qm, k, iter, L0):
    """Numpy reference for [`flm_gemma4_attn_qk_kvh2_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_attn_qk_kvh2_core]; see ``flm_gemma4_attn_qk_core_ref``."""
    return flm_gemma4_attn_qk_core_ref(qm, k, iter, L0)

flm_gemma4_decode_swa_attn_kv

flm_gemma4_decode_swa_attn_kv(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _SwaAttnKvKernel

Build the sliding-window-attention score-times-value core, from flm_gemma4/decode_swa_attn_kv.cc.

flm_gemma4_decode_attn_kv at head dim swa_dh, with entry points swa_attn_kv_begin, swa_attn_kv_round and swa_attn_kv_finish. One round covers every KV head, so it builds for one and for two.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_swa_attn_kv(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _SwaAttnKvKernel:
    """Build the sliding-window-attention score-times-value core, from ``flm_gemma4/decode_swa_attn_kv.cc``.

    [`flm_gemma4_decode_attn_kv`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv]
    at head dim ``swa_dh``, with entry points ``swa_attn_kv_begin``,
    ``swa_attn_kv_round`` and ``swa_attn_kv_finish``. One round covers every
    KV head, so it builds for one and for two.

    Args:
        geometry: The model the kernel builds for.
    """
    t = _decode_attn_types(
        geometry, geometry.swa_dh, geometry.num_kv_heads * geometry.swa_dh
    )
    fn = _decode_kernel(
        "swa_attn_kv",
        "swa_attn_kv_begin",
        [t["y"], t["l"]],
        (Out, Out),
        geometry,
        lut=True,
        cls=_SwaAttnKvKernel,
    )
    bind = fn.object_file.bind
    fn.swa_attn_kv_round = bind("swa_attn_kv_round", [t["s"], t["kv"], t["y"], t["l"]])
    fn.swa_attn_kv_finish = bind("swa_attn_kv_finish", [t["y"], t["l"]])
    return fn

flm_gemma4_swa_attn_kv_core

flm_gemma4_swa_attn_kv_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One swa_attn_kv_round of flm_gemma4_decode_swa_attn_kv with caller-owned state.

swa_attn_kv_round_core(sv, ly_in, ly_out): ly_out is ly_in with l = c * l + rowsum(s) and y = c * y + s @ v per query head, each head reading its KV head's half of v when there are two. sv holds s, then v. ly holds y, then l in a 16-float slot.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_swa_attn_kv_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """One ``swa_attn_kv_round`` of [`flm_gemma4_decode_swa_attn_kv`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_swa_attn_kv] with caller-owned state.

    ``swa_attn_kv_round_core(sv, ly_in, ly_out)``: ``ly_out`` is ``ly_in``
    with ``l = c * l + rowsum(s)`` and ``y = c * y + s @ v`` per query head,
    each head reading its KV head's half of ``v`` when there are two. ``sv``
    holds ``s``, then ``v``. ``ly`` holds ``y``, then ``l`` in a 16-float
    slot.
    """
    base = flm_gemma4_decode_swa_attn_kv(geometry=geometry)
    dh, kv_heads = geometry.swa_dh, geometry.num_kv_heads
    return _attn_kv_round_core(
        base,
        "decode_swa_attn_kv_core.cc",
        "swa_attn_kv_round_core",
        dh=dh,
        kv_width=kv_heads * dh,
        impl="1x8x1" if kv_heads == 1 else "2x4x1",
        rows=8,
        ref=partial(flm_gemma4_swa_attn_kv_core_ref, dh=dh, kv_heads=kv_heads),
    )

flm_gemma4_swa_attn_kv_core_ref

flm_gemma4_swa_attn_kv_core_ref(
    sv, ly, *, dh=256, kv_heads=1
)

Numpy reference for flm_gemma4_swa_attn_kv_core, in float64.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_swa_attn_kv_core_ref(sv, ly, *, dh=256, kv_heads=1):
    """Numpy reference for [`flm_gemma4_swa_attn_kv_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_swa_attn_kv_core], in float64."""
    impl = "1x8x1" if kv_heads == 1 else "2x4x1"
    return _attn_kv_round_ref(sv, ly, dh=dh, impl=impl)

flm_gemma4_decode_rope

flm_gemma4_decode_rope(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _RopeKernel

Build a layer's q/k norm and rotary embedding, from flm_gemma4/decode_rope.cc.

Three entry points, one read-only head (dh bf16) each. rope_q(q, qkv, rope_w, i) and rope_k(k, qkv, rope_w, i) RMS-normalize qkv against the head's weight and rotate it into output slot i; rope_v(v, qkv, i) RMS-normalizes it into slot i. q is q_heads_padded * dh bf16, k and v num_kv_heads * dh, and rope_w the cos and sin halves followed by the q and k norm weights (3 * dh total). dh is the geometry's swa_dh with sliding_window, else its dh.

Parameters:

Name	Type	Description	Default
sliding_window	bool	Build for the sliding-window layers' head dim.	False
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_rope(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> _RopeKernel:
    """Build a layer's q/k norm and rotary embedding, from ``flm_gemma4/decode_rope.cc``.

    Three entry points, one read-only head (``dh`` bf16) each.
    ``rope_q(q, qkv, rope_w, i)`` and ``rope_k(k, qkv, rope_w, i)``
    RMS-normalize ``qkv`` against the head's weight and rotate it into
    output slot ``i``; ``rope_v(v, qkv, i)``
    RMS-normalizes it into slot ``i``. ``q`` is ``q_heads_padded * dh`` bf16,
    ``k`` and ``v`` ``num_kv_heads * dh``, and ``rope_w`` the cos and sin
    halves followed by the q and k norm weights (``3 * dh`` total).
    ``dh`` is the geometry's ``swa_dh`` with ``sliding_window``, else its
    ``dh``.

    Args:
        sliding_window: Build for the sliding-window layers' head dim.
        geometry: The model the kernel builds for.
    """
    dh = geometry.swa_dh if sliding_window else geometry.dh
    bf16 = np.dtype[bfloat16]
    kv = geometry.num_kv_heads * dh
    qkv_ty = np.ndarray[(dh,), bf16]
    q_ty = np.ndarray[(_decode_q_heads_padded(geometry) * dh,), bf16]
    kv_ty = np.ndarray[(kv,), bf16]
    rope_ty = np.ndarray[(dh * (3 if geometry.qk_norm else 1),), bf16]
    fn = _decode_kernel(
        "rope",
        "rope_q",
        [q_ty, qkv_ty, rope_ty, np.int32],
        (Out, InOut, In, Param),
        geometry,
        flags=(f"-DFLM_GEMMA4_DECODE_ROPE_SWA={int(sliding_window)}",),
        cls=_RopeKernel,
    )
    bind = fn.object_file.bind
    fn.rope_k = bind("rope_k", [kv_ty, qkv_ty, rope_ty, np.int32])
    fn.rope_v = bind("rope_v", [kv_ty, qkv_ty, np.int32])
    return fn

flm_gemma4_rope_core_ref

flm_gemma4_rope_core_ref(x, rope_w)

Numpy reference for flm_gemma4_rope_core.

Per call, x is a q head then a k head and rope_w is [cos | sin | q weight | k weight]. Each head n = x * w / sqrt(mean(x^2) + 1e-6) is rotated by halves: y1 = n1 cos - n2 sin, y2 = n1 sin + n2 cos, with n1, n2 its first and second half.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_rope_core_ref(x, rope_w):
    """Numpy reference for [`flm_gemma4_rope_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_rope_core].

    Per call, ``x`` is a q head then a k head and ``rope_w`` is ``[cos |
    sin | q weight | k weight]``. Each head ``n = x * w / sqrt(mean(x^2) +
    1e-6)`` is rotated by halves: ``y1 = n1 cos - n2 sin``, ``y2 = n1 sin +
    n2 cos``, with ``n1``, ``n2`` its first and second half.
    """
    a, b = _rope_core_terms(x, rope_w)
    y = (a + b).reshape(len(a), -1)
    return y.astype(np.float32).astype(bfloat16)

flm_gemma4_rope_core

flm_gemma4_rope_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

_rotate_t of flm_gemma4_decode_rope on one q and one k head.

rope_head_core(x, rope_w, y): x is a q head then a k head, rope_w the base kernel's [cos | sin | q weight | k weight].

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_rope_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> ExternalFunction:
    """``_rotate_t`` of [`flm_gemma4_decode_rope`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_rope] on one q and one k head.

    ``rope_head_core(x, rope_w, y)``: ``x`` is a q head then a k head,
    ``rope_w`` the base kernel's ``[cos | sin | q weight | k weight]``.
    """
    base = flm_gemma4_decode_rope(sliding_window=sliding_window, geometry=geometry)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    heads = np.ndarray[(2 * dh,), _BF16]
    return _wrapper(
        base,
        "decode_rope_core.cc",
        "rope_head_core",
        [heads, np.ndarray[(3 * dh,), _BF16], heads],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            reference=flm_gemma4_rope_core_ref,
            sample=partial(_rope_core_sample, dh=dh),
            tolerance=Tolerance.bounded(
                _rope_core_bound,
                note="bf16 ulp <= 2**-7 |v|. The kernel floors each normalized "
                "value to bf16 (1 ulp, plus the rms scale's error, under "
                "2**-12), so each rotation product is off by that much; it "
                "floors the rotated sum to bf16 (1 ulp) and the reference "
                "rounds it to nearest (1/2 ulp)",
            ),
            # Per element: square-add and two multiplies for the norm, then
            # a multiply and a multiply-add for the rotation.
            ops_per_call=2 * 7 * dh,
        ),
    )

flm_gemma4_v_norm_core_ref

flm_gemma4_v_norm_core_ref(x)

Numpy reference for flm_gemma4_v_norm_core: x / sqrt(mean(x^2) + RMS_EPS).

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_v_norm_core_ref(x):
    """Numpy reference for [`flm_gemma4_v_norm_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_v_norm_core]: ``x / sqrt(mean(x^2) + RMS_EPS)``."""
    x = np.asarray(x, np.float64)
    return (x * _rms(x)).astype(np.float32).astype(bfloat16)

flm_gemma4_v_norm_core

flm_gemma4_v_norm_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

rms_norm_unweighted of flm_gemma4_decode_rope on one v head.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_v_norm_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> ExternalFunction:
    """``rms_norm_unweighted`` of [`flm_gemma4_decode_rope`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_rope] on one v head."""
    base = flm_gemma4_decode_rope(sliding_window=sliding_window, geometry=geometry)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    head = np.ndarray[(dh,), _BF16]
    return _wrapper(
        base,
        "decode_rope_core.cc",
        "v_norm_core",
        [head, head],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_v_norm_core_ref,
            sample=partial(_v_norm_core_sample, dh=dh),
            tolerance=Tolerance.bounded(
                _v_norm_core_bound,
                note="bf16 ulp <= 2**-7 |v|: the kernel floors to bf16 (1 ulp), "
                "the reference rounds to nearest (1/2 ulp), and the rms scale "
                "is within 2**-12",
            ),
            # Per element: square-add, then one multiply.
            ops_per_call=3 * dh,
        ),
    )

flm_gemma4_decode_rms_residual

flm_gemma4_decode_rms_residual(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _RmsKernel

Build three RMS/residual phases: rms_begin, rms_middle, rms_finish.

Each takes an input row, four norm weights, residual scratch (two rows), and an output. Begin and middle write a packet-padded projection input; finish writes a plain row. Retain the weights and scratch across calls.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_rms_residual(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _RmsKernel:
    """Build three RMS/residual phases: ``rms_begin``, ``rms_middle``, ``rms_finish``.

    Each takes an input row, four norm weights, residual scratch (two rows),
    and an output. Begin and middle write a packet-padded projection input;
    finish writes a plain row. Retain the weights and scratch across calls.
    """
    d = geometry.model_dim
    x = np.ndarray[(d,), _BF16]
    args = [
        x,
        np.ndarray[(4, d), _BF16],
        np.ndarray[(2, d), _BF16],
        np.ndarray[(d + 16,), _BF16],
    ]
    fn = _decode_kernel(
        "rms_residual",
        "rms_begin",
        args,
        (In, In, InOut, Out),
        geometry,
        cls=_RmsKernel,
    )
    fn.rms_middle = fn.object_file.bind("rms_middle", args)
    fn.rms_finish = fn.object_file.bind("rms_finish", [*args[:3], x])
    return fn

flm_gemma4_rms_residual_core_ref

flm_gemma4_rms_residual_core_ref(x, w, residual)

Numpy reference for flm_gemma4_rms_residual_core.

residual + x * w / sqrt(mean(x^2) + RMS_EPS), rounded to bf16.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_rms_residual_core_ref(x, w, residual):
    """Numpy reference for [`flm_gemma4_rms_residual_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_rms_residual_core].

    ``residual + x * w / sqrt(mean(x^2) + RMS_EPS)``, rounded to bf16.
    """
    return (
        _rms_residual_core_terms(x, w, residual)[1].astype(np.float32).astype(bfloat16)
    )

flm_gemma4_rms_residual_core

flm_gemma4_rms_residual_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

rms_norm then residual_add of flm_gemma4_decode_rms_residual.

The post-attention norm and residual add. rms_residual_core(x, w, residual, y), all model_dim bf16: y = residual + rms_norm(x) * w.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_rms_residual_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``rms_norm`` then ``residual_add`` of [`flm_gemma4_decode_rms_residual`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_rms_residual].

    The post-attention norm and residual add. ``rms_residual_core(x, w,
    residual, y)``, all ``model_dim`` bf16: ``y = residual + rms_norm(x) * w``.
    """
    base = flm_gemma4_decode_rms_residual(geometry=geometry)
    row = np.ndarray[(geometry.model_dim,), _BF16]
    return _wrapper(
        base,
        "decode_rms_residual_core.cc",
        "rms_residual_core",
        [row, row, row, row],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Out),
            reference=flm_gemma4_rms_residual_core_ref,
            tolerance=Tolerance.bounded(
                _rms_residual_core_bound,
                note="bf16 ulp <= 2**-7 |v|. The kernel floors the normalized "
                "value to bf16 (1 ulp, plus the fast inverse sqrt's and the "
                "float32 sum of squares' error, under 2**-12), floors the sum "
                "to bf16 (1 ulp, plus its float32 add's error, under 2**-12), "
                "and the reference rounds the sum to nearest (1/2 ulp). Each "
                "flush of a subnormal adds the smallest normal",
            ),
            # Per element: square-add and two multiplies, then the add.
            ops_per_call=5 * geometry.model_dim,
        ),
    )

flm_gemma4_rms_norm_core_ref

flm_gemma4_rms_norm_core_ref(x, w)

Numpy reference for flm_gemma4_rms_norm_core.

16 zeros, then x * w / sqrt(mean(x^2) + RMS_EPS) rounded to bf16.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_rms_norm_core_ref(x, w):
    """Numpy reference for [`flm_gemma4_rms_norm_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_rms_norm_core].

    16 zeros, then ``x * w / sqrt(mean(x^2) + RMS_EPS)`` rounded to bf16.
    """
    n = _rms_residual_core_terms(x, w, 0)[0]
    header = np.zeros((len(n), _PACKET_HEADER))
    return np.concatenate([header, n], axis=1).astype(np.float32).astype(bfloat16)

flm_gemma4_rms_norm_core

flm_gemma4_rms_norm_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

rms_norm of flm_gemma4_decode_rms_residual into its packet layout.

rms_norm_core(x, w, y): x and w are model_dim bf16, y is model_dim + 16, a zeroed 16-element packet header then rms_norm(x) * w.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_rms_norm_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``rms_norm`` of [`flm_gemma4_decode_rms_residual`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_rms_residual] into its packet layout.

    ``rms_norm_core(x, w, y)``: ``x`` and ``w`` are ``model_dim`` bf16, ``y``
    is ``model_dim + 16``, a zeroed 16-element packet header then
    ``rms_norm(x) * w``.
    """
    base = flm_gemma4_decode_rms_residual(geometry=geometry)
    d = geometry.model_dim
    row = np.ndarray[(d,), _BF16]
    return _wrapper(
        base,
        "decode_rms_residual_core.cc",
        "rms_norm_core",
        [row, row, np.ndarray[(d + _PACKET_HEADER,), _BF16]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            reference=flm_gemma4_rms_norm_core_ref,
            tolerance=Tolerance.bounded(
                _rms_norm_core_bound,
                note="bf16 ulp <= 2**-7 |v|. The kernel floors the normalized "
                "value to bf16 (1 ulp, plus the fast inverse sqrt's and the "
                "float32 sum of squares' error, under 2**-12) and the "
                "reference rounds it to nearest (1/2 ulp). The header is "
                "exact. Each flush of a subnormal adds the smallest normal",
            ),
            # Per element: square-add and two multiplies.
            ops_per_call=4 * d,
        ),
    )

flm_gemma4_decode_emit

flm_gemma4_decode_emit(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    attention: bool = False,
    sliding_window: bool = False
) -> ExternalFunction

Copy one 512-bf16 projection input, reordering attention's head groups.

RMS input has a 16-bf16 packet prefix. Attention input is a float32 allocation whose first half holds normalized bf16 values in [dh/8, head, 8] order, following attn_kv_finish. The output is a contiguous piece of the head-major projection vector.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_emit(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    attention: bool = False,
    sliding_window: bool = False,
) -> ExternalFunction:
    """Copy one 512-bf16 projection input, reordering attention's head groups.

    RMS input has a 16-bf16 packet prefix. Attention input is a float32
    allocation whose first half holds normalized bf16 values in
    [dh/8, head, 8] order, following ``attn_kv_finish``.
    The output is a contiguous piece of the head-major projection vector.
    """
    dh = geometry.swa_dh if sliding_window else geometry.dh
    size = geometry.num_attn_heads * dh if attention else geometry.model_dim + 16
    return _decode_kernel(
        "emit",
        "decode_emit",
        [
            np.ndarray[(size,), _F32 if attention else _BF16],
            np.ndarray[(512,), _BF16],
            np.int32,
        ],
        (In, Out, Param),
        geometry,
        flags=(f"-DEMIT_ATTENTION={int(attention)}", f"-DEMIT_DH={dh}"),
    )

flm_gemma4_decode_proj_main

flm_gemma4_decode_proj_main(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _ProjKernel

Build q4nx projection steps: zero, accumulate, store and packet assembly.

Group sums are computed for the first row block and reused for subsequent rows. proj_store writes one row block; proj_pack combines a float32 accumulator and a bf16 row block after a 16-bf16 packet header.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_proj_main(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _ProjKernel:
    """Build q4nx projection steps: zero, accumulate, store and packet assembly.

    Group sums are computed for the first row block and reused for subsequent
    rows. ``proj_store`` writes one row block; ``proj_pack`` combines a float32
    accumulator and a bf16 row block after a 16-bf16 packet header.
    """
    m, k = Q4NX_M_TILE, Q4NX_K_TILE
    acc = np.ndarray[(m,), np.dtype[np.float32]]
    row = np.ndarray[(m,), _BF16]
    sums = np.ndarray[
        (geometry.intermediate_size * (2 if geometry.double_wide_mlp else 1) // 32,),
        _BF16,
    ]
    fn = _decode_kernel(
        "proj_main",
        "proj_accumulate",
        [
            np.ndarray[(Q4NX_BLOCK_BYTES // 2,), _BF16],
            np.ndarray[(k,), _BF16],
            acc,
            sums,
            np.int32,
            np.int32,
        ],
        (In, In, InOut, InOut, Param, Param),
        geometry,
        cls=_ProjKernel,
    )
    fn.proj_zero = fn.object_file.bind("proj_zero", [acc])
    fn.proj_store = fn.object_file.bind("proj_store", [row, acc])
    fn.proj_pack = fn.object_file.bind(
        "proj_pack", [np.ndarray[(2 * m + 16,), _BF16], acc, row, np.int32]
    )
    return fn

flm_gemma4_decode_per_layer_up

flm_gemma4_decode_per_layer_up(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _PerLayerUpKernel

Build the per-layer-input up projection, from flm_gemma4/decode_per_layer_up.cc.

per_layer_up_begin(x) gates the layer's per-layer input. bf16_proj_zero, per_layer_up_mvm and per_layer_up_store implement the projection from pli_d up to model_dim using 32 by 256 bf16 weight blocks. per_layer_up_finish(x, y) RMS-normalizes the result, adds the residual and scales it by the layer scale. x is 2 * (pli_d + model_dim) + 32 bf16: the norm weight, the layer scale padded to 32, the per-layer input, the residual and the gate. y is model_dim bf16. The begin entry point gates the per-layer input in place.

y needs a 64-byte-aligned base pointer.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_per_layer_up(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _PerLayerUpKernel:
    """Build the per-layer-input up projection, from ``flm_gemma4/decode_per_layer_up.cc``.

    ``per_layer_up_begin(x)`` gates the layer's per-layer input.
    ``bf16_proj_zero``, ``per_layer_up_mvm`` and ``per_layer_up_store`` implement
    the projection from ``pli_d`` up to ``model_dim`` using 32 by 256 bf16
    weight blocks. ``per_layer_up_finish(x, y)`` RMS-normalizes
    the result, adds the residual and scales it by the layer scale. ``x`` is
    ``2 * (pli_d + model_dim) + 32`` bf16: the norm weight, the layer scale
    padded to 32, the per-layer input, the residual and the gate. ``y`` is
    ``model_dim`` bf16. The begin entry point gates the per-layer input in place.

    ``y`` needs a 64-byte-aligned base pointer.

    Args:
        geometry: The model the kernel builds for.
    """
    bf16 = np.dtype[bfloat16]
    d, pli = geometry.model_dim, geometry.pli_d
    x_ty = np.ndarray[(2 * (pli + d) + 32,), bf16]
    y_ty = np.ndarray[(d,), bf16]
    fn = _decode_kernel(
        "per_layer_up",
        "per_layer_up_begin",
        [x_ty],
        (InOut,),
        geometry,
        cls=_PerLayerUpKernel,
    )
    fn.per_layer_up_finish = fn.object_file.bind("per_layer_up_finish", [x_ty, y_ty])
    return fn

flm_gemma4_decode_bf16_proj_zero

flm_gemma4_decode_bf16_proj_zero(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

Zero one 32-element float projection accumulator, inlined into its core.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_bf16_proj_zero(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """Zero one 32-element float projection accumulator, inlined into its core."""
    return _bf16_proj_helper(
        "bf16_proj_zero",
        0,
        [np.ndarray[(BF16_PROJ_M,), _F32]],
        (Out,),
        geometry,
    )

flm_gemma4_decode_proj_layer_embedding_mvm

flm_gemma4_decode_proj_layer_embedding_mvm(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One MVM block of the per-layer-input embedding projection, inlined into its core.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_proj_layer_embedding_mvm(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """One MVM block of the per-layer-input embedding projection, inlined into its core."""
    return _bf16_proj_mvm(
        "proj_layer_embedding_mvm",
        2 * (geometry.pli_d + geometry.model_dim) + 32,
        geometry,
    )

flm_gemma4_decode_proj_layer_embedding_store

flm_gemma4_decode_proj_layer_embedding_store(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

Store one output block of the per-layer-input embedding projection, inlined into its core.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_proj_layer_embedding_store(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """Store one output block of the per-layer-input embedding projection, inlined into its core."""
    return _bf16_proj_store("proj_layer_embedding_store", geometry.pli_d, geometry)

flm_gemma4_decode_gate_layer_embedding_mvm

flm_gemma4_decode_gate_layer_embedding_mvm(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One MVM block of the per-layer-input gate projection, inlined into its core.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_gate_layer_embedding_mvm(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """One MVM block of the per-layer-input gate projection, inlined into its core."""
    return _bf16_proj_mvm(
        "gate_layer_embedding_mvm", geometry.model_dim + geometry.pli_d, geometry
    )

flm_gemma4_decode_gate_layer_embedding_store

flm_gemma4_decode_gate_layer_embedding_store(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

Store one output block of the per-layer-input gate projection, inlined into its core.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_gate_layer_embedding_store(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """Store one output block of the per-layer-input gate projection, inlined into its core."""
    return _bf16_proj_store(
        "gate_layer_embedding_store", geometry.model_dim + geometry.pli_d, geometry
    )

flm_gemma4_decode_per_layer_up_mvm

flm_gemma4_decode_per_layer_up_mvm(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One MVM block of the per-layer-input up projection, inlined into its core.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_per_layer_up_mvm(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """One MVM block of the per-layer-input up projection, inlined into its core."""
    size = 2 * (geometry.pli_d + geometry.model_dim) + 32
    return _bf16_proj_mvm("per_layer_up_mvm", size, geometry)

flm_gemma4_decode_per_layer_up_store

flm_gemma4_decode_per_layer_up_store(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

Store one output block of the per-layer-input up projection, inlined into its core.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_per_layer_up_store(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """Store one output block of the per-layer-input up projection, inlined into its core."""
    return _bf16_proj_store("per_layer_up_store", geometry.model_dim, geometry)

flm_gemma4_bf16_proj_core_ref

flm_gemma4_bf16_proj_core_ref(w, x)

Numpy reference for flm_gemma4_bf16_proj_core: one weight block times x.

The block is stored column by column: w[k * 32 + m] multiplies x[k] into out-feature m. Sums in float64.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_bf16_proj_core_ref(w, x):
    """Numpy reference for [`flm_gemma4_bf16_proj_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_bf16_proj_core]: one weight block times ``x``.

    The block is stored column by column: ``w[k * 32 + m]`` multiplies
    ``x[k]`` into out-feature ``m``. Sums in float64.
    """
    w = np.asarray(w, np.float64).reshape(-1, BF16_PROJ_K, BF16_PROJ_M)
    x = np.asarray(x, np.float64).reshape(-1, BF16_PROJ_K)
    return np.einsum("ckm,ck->cm", w, x).astype(np.float32)

flm_gemma4_bf16_proj_core

flm_gemma4_bf16_proj_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

_mvm_bf16_bf16 of the per-layer-input projections, one block.

bf16_proj_block_core(w, x, y) writes one 32 by 256 bf16 weight block times 256 inputs to float32 y. The block does not depend on geometry.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_bf16_proj_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``_mvm_bf16_bf16`` of the per-layer-input projections, one block.

    ``bf16_proj_block_core(w, x, y)`` writes one 32 by 256 bf16 weight block
    times 256 inputs to float32 ``y``. The block does not depend on ``geometry``.
    """
    base = flm_gemma4_decode_per_layer_up(geometry=geometry)
    return _wrapper(
        base,
        "decode_per_layer_up_core.cc",
        "bf16_proj_block_core",
        [
            np.ndarray[(_DECODE_BF16_BLOCK,), _BF16],
            np.ndarray[(BF16_PROJ_K,), _BF16],
            np.ndarray[(BF16_PROJ_M,), _F32],
        ],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            reference=flm_gemma4_bf16_proj_core_ref,
            sample=_bf16_proj_core_sample,
            tolerance=Tolerance.exact(
                note="the sample keeps every product and sum exact in float32"
            ),
            ops_per_call=2 * _DECODE_BF16_BLOCK,
            acc_dtype=np.float32,
            reduction=BF16_PROJ_K,
        ),
    )

flm_gemma4_decode_proj_layer_embedding

flm_gemma4_decode_proj_layer_embedding(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

Finish the embedding from a packed input object and projected scratch.

Input layout: per-layer embedding, norm weights plus model_dim + 32 values forwarded to the up projection, then the model_dim input row. The MVM entry point reads that last row using its explicit offset. The first pli_d input elements are overwritten with the residual sum.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_proj_layer_embedding(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """Finish the embedding from a packed input object and projected scratch.

    Input layout: per-layer embedding, norm weights plus ``model_dim + 32``
    values forwarded to the up projection, then the ``model_dim`` input row.
    The MVM entry point reads that last row using its explicit offset.
    The first ``pli_d`` input elements are overwritten with the residual sum.
    """
    bf16 = np.dtype[bfloat16]
    d, pli = geometry.model_dim, geometry.pli_d
    pli_ty = np.ndarray[(pli,), bf16]
    y_ty = np.ndarray[(pli + d + 32,), bf16]
    return _decode_kernel(
        "proj_layer_embedding",
        "proj_layer_embedding_finish",
        [np.ndarray[(2 * (pli + d) + 32,), bf16], pli_ty, y_ty],
        (InOut, InOut, Out),
        geometry,
    )

flm_gemma4_decode_gate_layer_embedding

flm_gemma4_decode_gate_layer_embedding(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _GateKernel

Build the per-layer-input gate, from flm_gemma4/decode_gate_layer_embedding.cc.

gate_layer_embedding_begin(x, y) copies the layer output x (model_dim bf16) into y. bf16_proj_zero, gate_layer_embedding_mvm and gate_layer_embedding_store project it to pli_d using 32 by 256 bf16 weight blocks. gate_layer_embedding_finish(y) applies the activation into y + model_dim. y is model_dim + pli_d bf16.

y needs a 64-byte-aligned base pointer; the kernel also stores 512 bits at a time at y + model_dim.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_decode_gate_layer_embedding(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> _GateKernel:
    """Build the per-layer-input gate, from ``flm_gemma4/decode_gate_layer_embedding.cc``.

    ``gate_layer_embedding_begin(x, y)`` copies the layer output ``x``
    (``model_dim`` bf16) into ``y``.
    ``bf16_proj_zero``, ``gate_layer_embedding_mvm`` and
    ``gate_layer_embedding_store`` project it to ``pli_d`` using 32 by 256
    bf16 weight blocks. ``gate_layer_embedding_finish(y)``
    applies the activation into
    ``y + model_dim``. ``y`` is ``model_dim + pli_d`` bf16.

    ``y`` needs a 64-byte-aligned base pointer; the kernel also stores
    512 bits at a time at ``y + model_dim``.

    Args:
        geometry: The model the kernel builds for.
    """
    bf16 = np.dtype[bfloat16]
    d = geometry.model_dim
    x_ty = np.ndarray[(d,), bf16]
    y_ty = np.ndarray[(d + geometry.pli_d,), bf16]
    fn = _decode_kernel(
        "gate_layer_embedding",
        "gate_layer_embedding_begin",
        [x_ty, y_ty],
        (In, Out),
        geometry,
        lut=True,
        cls=_GateKernel,
    )
    fn.gate_layer_embedding_finish = fn.object_file.bind(
        "gate_layer_embedding_finish", [y_ty]
    )
    return fn

flm_gemma4_pli_gelu_core_ref

flm_gemma4_pli_gelu_core_ref(x)

Numpy reference for flm_gemma4_pli_gelu_core: tanh-GELU.

The tolerance covers the kernel's table.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_pli_gelu_core_ref(x):
    """Numpy reference for [`flm_gemma4_pli_gelu_core`][aie.iron.kernels.flm_gemma4.flm_gemma4_pli_gelu_core]: tanh-GELU.

    The tolerance covers the kernel's table.
    """
    return _gelu_tanh(x).astype(np.float32)

flm_gemma4_pli_gelu_core

flm_gemma4_pli_gelu_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

_activate of flm_gemma4_decode_gate_layer_embedding: GELU over pli_d.

pli_gelu_core(x, y): copies x into y and applies GELU there, through aie_runtime_lib's table (getGeluBf16). Both are pli_d bf16.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_pli_gelu_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``_activate`` of [`flm_gemma4_decode_gate_layer_embedding`][aie.iron.kernels.flm_gemma4.flm_gemma4_decode_gate_layer_embedding]: GELU over ``pli_d``.

    ``pli_gelu_core(x, y)``: copies ``x`` into ``y`` and applies GELU there,
    through aie_runtime_lib's table (``getGeluBf16``). Both are ``pli_d``
    bf16.

    Args:
        geometry: The model the kernel builds for.
    """
    n = geometry.pli_d
    return _wrapper(
        flm_gemma4_decode_gate_layer_embedding(geometry=geometry),
        "decode_gate_layer_embedding_core.cc",
        "pli_gelu_core",
        [np.ndarray[(n,), _BF16], np.ndarray[(n,), _BF16]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_pli_gelu_core_ref,
            sample=partial(_pli_gelu_core_sample, n=n),
            tolerance=Tolerance.bounded(
                _pli_gelu_core_bound,
                note="the GELU table's error against tanh-GELU, measured on the "
                "host from aie_runtime_lib's table, plus a bf16 ulp of its slope "
                "and the floor bf16 store",
            ),
            uses_lut=True,
            # Per value: the table line's multiply and add.
            ops_per_call=2 * n,
        ),
    )

flm_gemma4_q4nx_lm_head_ref

flm_gemma4_q4nx_lm_head_ref(
    w,
    x,
    sums,
    *,
    m_tile=Q4NX_M_TILE,
    k_tile=Q4NX_K_TILE,
    group=Q4NX_GROUP
)

Numpy reference for flm_gemma4_q4nx_lm_head: one q4nx_lm_head_block call.

w is one q4nx block per call, as bytes or bf16 words: bf16 scales, then bf16 minima, both indexed [k // group, m], then 4-bit codes, low nibble first, indexed [m // 16, k // 32, (k % 32) // 8, k % 8, m % 16]. x is the token then its RMS weight, and sums the token's per-32 column sums, both shared by every call; slice 0 of the token is the one the call reads. Returns the float32 accumulator, zero before the call, holding sum_k (min + scale * code) * x with the minima folded in through sums.

Per 16 rows and 32 columns the kernel sums the codes times x in float32, narrows that dot product to bf16 with the floor rounding a fresh core uses, then accumulates it times the scale, and the minimum times the column sum, into float32. Every product is exact, so each multiply-accumulate rounds once.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_ref(
    w, x, sums, *, m_tile=Q4NX_M_TILE, k_tile=Q4NX_K_TILE, group=Q4NX_GROUP
):
    """Numpy reference for [`flm_gemma4_q4nx_lm_head`][aie.iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head]: one ``q4nx_lm_head_block`` call.

    ``w`` is one q4nx block per call, as bytes or bf16 words: bf16 scales,
    then bf16 minima, both indexed ``[k // group, m]``, then 4-bit codes, low
    nibble first, indexed ``[m // 16, k // 32, (k % 32) // 8, k % 8, m % 16]``.
    ``x`` is the token then its RMS weight, and ``sums`` the token's per-32
    column sums, both shared by every call; slice 0 of the token is the one
    the call reads. Returns the float32
    accumulator, zero before the call, holding ``sum_k (min + scale * code) *
    x`` with the minima folded in through ``sums``.

    Per 16 rows and 32 columns the kernel sums the codes times ``x`` in
    float32, narrows that dot product to bf16 with the floor rounding a fresh
    core uses, then accumulates it times the scale, and the minimum times the
    column sum, into float32. Every product is exact, so each
    multiply-accumulate rounds once.
    """
    w = np.ascontiguousarray(w).view(np.uint8).reshape(-1, m_tile * k_tile * 5 // 8)
    calls, n_groups = len(w), k_tile // group
    codes, scales, mins = q4nx_unpack(
        w, m_tile=m_tile, k_tile=k_tile, group=group, ct_k=group
    )
    codes = codes.reshape(calls, n_groups, group, m_tile // 16, 16).transpose(
        0, 3, 1, 2, 4
    )
    x = np.asarray(x, np.float32).reshape(-1)[:k_tile]
    x = x.astype(np.float64).reshape(1, n_groups, 32)
    sums = np.asarray(sums, np.float32).reshape(-1)[:n_groups]
    sums = sums.astype(np.float64).reshape(1, n_groups)
    acc = np.zeros((calls, m_tile // 16, 16), np.float32)
    for g in range(n_groups):
        dot = np.zeros((calls, m_tile // 16, 16), np.float32)
        for c in range(32):
            dot = f32(dot + codes[:, :, g, c, :] * x[:, g, c, None, None])
        dot = bf16_floor(dot).astype(np.float64)
        scale = scales[:, g].reshape(calls, m_tile // 16, 16).astype(np.float64)
        low = mins[:, g].reshape(calls, m_tile // 16, 16).astype(np.float64)
        acc = f32(acc + dot * scale)
        acc = f32(acc + low * sums[:, g, None, None])
    return acc.reshape(calls, m_tile)

flm_gemma4_q4nx_lm_head

flm_gemma4_q4nx_lm_head(
    *,
    dim: int = 1536,
    m_tile: int = Q4NX_M_TILE,
    k_tile: int = Q4NX_K_TILE,
    group: int = Q4NX_GROUP
) -> _LmHeadKernel

AIE2P logits from a q4nx vocabulary, from flm_gemma4/q4nx_lm_head.cc.

One core owns m_tile out-features at a time and streams their q4nx blocks past a token of dim values. The returned kernel is q4nx_lm_head_block, which accumulates one m_tile by k_tile block times one slice of the token into float32; its slice index is bound to 0. See flm_gemma4_q4nx_lm_head_ref for the block layout.

The other entry points, attributes of the kernel:

• q4nx_lm_head_rms(x, sums): once per token; normalizes the token in x (the token, then its RMS weight) in place and writes its per-32 column sums.
• q4nx_lm_head_zero(y_acc): before a tile's k loop.
• q4nx_lm_head_epilogue(y, y_acc, softcap): after the k loop, y = c * tanh(y_acc / c), with the float32 c in RTP word 0.

Parameters:

Name	Type	Description	Default
dim	int	The token's length, a multiple of k_tile; sets the RMS norm's span.	1536
m_tile	int	Out-features per block, a multiple of 16.	Q4NX_M_TILE
k_tile	int	In-features per block, a multiple of group.	Q4NX_K_TILE
group	int	In-features per scale and minimum; must be 32.	Q4NX_GROUP

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head(
    *,
    dim: int = 1536,
    m_tile: int = Q4NX_M_TILE,
    k_tile: int = Q4NX_K_TILE,
    group: int = Q4NX_GROUP,
) -> _LmHeadKernel:
    """AIE2P logits from a q4nx vocabulary, from ``flm_gemma4/q4nx_lm_head.cc``.

    One core owns ``m_tile`` out-features at a time and streams their q4nx
    blocks past a token of ``dim`` values. The returned kernel is
    ``q4nx_lm_head_block``, which accumulates one ``m_tile`` by ``k_tile``
    block times one slice of the token into float32; its slice index is bound
    to 0. See ``flm_gemma4_q4nx_lm_head_ref`` for the block layout.

    The other entry points, attributes of the kernel:

    - ``q4nx_lm_head_rms(x, sums)``: once per token; normalizes the token in
      ``x`` (the token, then its RMS weight) in place and writes its per-32
      column sums.
    - ``q4nx_lm_head_zero(y_acc)``: before a tile's k loop.
    - ``q4nx_lm_head_epilogue(y, y_acc, softcap)``: after the k loop,
      ``y = c * tanh(y_acc / c)``, with the float32 ``c`` in RTP word 0.

    Args:
        dim: The token's length, a multiple of ``k_tile``; sets the RMS
            norm's span.
        m_tile: Out-features per block, a multiple of 16.
        k_tile: In-features per block, a multiple of ``group``.
        group: In-features per scale and minimum; must be 32.
    """
    _lm_head_geometry(dim, m_tile, k_tile, group)
    if not _arch_traits().bfp16:
        raise NotImplementedError(
            "flm_gemma4_q4nx_lm_head() is only available on aie2p."
        )
    geometry = dict(m_tile=m_tile, k_tile=k_tile, group=group)
    x = np.ndarray[(2, dim), _BF16]
    sums = np.ndarray[(dim // group,), _BF16]
    y_acc = np.ndarray[(m_tile,), _F32]
    fn = _make_extern(
        "q4nx_lm_head_block",
        _kernel_source("flm_gemma4/q4nx_lm_head.cc"),
        [np.ndarray[(m_tile * k_tile * 5 // 8 // 2,), _BF16], x, y_acc, sums, np.int32],
        compile_flags=[
            f"-DQ4NX_M_TILE={m_tile}",
            f"-DQ4NX_K_TILE={k_tile}",
            f"-DQ4NX_GROUP={group}",
            f"-DFLM_GEMMA4_LM_HEAD_DIM={dim}",
            # Without it the bf16 mmul emulation runs about 8x slower.
            "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16",
        ],
        cls=_LmHeadKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            # The token and its sums are fixed while the blocks stream past.
            roles=(In, Param, InOut, Param, Param),
            parameter_bindings=((4, 0),),
            initializers=((2, _lm_head_zero),),
            reference=partial(flm_gemma4_q4nx_lm_head_ref, **geometry),
            sample=partial(_q4nx_lm_head_sample, dim=dim, **geometry),
            tolerance=Tolerance.exact(
                note="the sample keeps every sum exact in float32; the floor "
                "bf16 narrowing of each dot product is modeled"
            ),
            ops_per_call=2 * m_tile * k_tile,
            acc_dtype=np.float32,
            reduction=k_tile,
        ),
    )
    bind = fn.object_file.bind
    fn.q4nx_lm_head_block = fn
    fn.q4nx_lm_head_rms = bind("q4nx_lm_head_rms", [x, sums])
    fn.q4nx_lm_head_zero = bind("q4nx_lm_head_zero", [y_acc])
    rtp = np.ndarray[(_LM_HEAD_RTP_WORDS,), _RTP_WORD]
    fn.q4nx_lm_head_epilogue = bind(
        "q4nx_lm_head_epilogue", [np.ndarray[(m_tile,), _BF16], y_acc, rtp]
    )
    return fn

flm_gemma4_q4nx_lm_head_epilogue_ref

flm_gemma4_q4nx_lm_head_epilogue_ref(y_acc, rtp)

Numpy reference for flm_gemma4_q4nx_lm_head_epilogue: c * tanh(y_acc / c).

rtp holds the float32 bits of c in word 0. The reference computes in float64 and rounds to bf16 once; the tolerance covers the kernel's narrowings and its tanh.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_epilogue_ref(y_acc, rtp):
    """Numpy reference for [`flm_gemma4_q4nx_lm_head_epilogue`][aie.iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head_epilogue]: ``c * tanh(y_acc / c)``.

    ``rtp`` holds the float32 bits of ``c`` in word 0. The reference computes
    in float64 and rounds to bf16 once; the tolerance covers the kernel's
    narrowings and its tanh.
    """
    c = _lm_head_softcap(rtp)
    y = np.asarray(y_acc, np.float64)
    return (c * np.tanh(y / c)).astype(bfloat16)

flm_gemma4_q4nx_lm_head_epilogue

flm_gemma4_q4nx_lm_head_epilogue(
    *,
    dim: int = 1536,
    m_tile: int = Q4NX_M_TILE,
    k_tile: int = Q4NX_K_TILE,
    group: int = Q4NX_GROUP
) -> ExternalFunction

q4nx_lm_head_epilogue of the LM head: y = c * tanh(y_acc / c).

The softcap c is float32 bits in word 0 of a 32-word RTP buffer. The tolerance assumes c is a bf16 value, so the kernel's narrowing of it is exact; Gemma 4's 30 is.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_epilogue(
    *,
    dim: int = 1536,
    m_tile: int = Q4NX_M_TILE,
    k_tile: int = Q4NX_K_TILE,
    group: int = Q4NX_GROUP,
) -> ExternalFunction:
    """``q4nx_lm_head_epilogue`` of the LM head: ``y = c * tanh(y_acc / c)``.

    The softcap ``c`` is float32 bits in word 0 of a 32-word RTP buffer. The
    tolerance assumes ``c`` is a bf16 value, so the kernel's narrowing of it
    is exact; Gemma 4's 30 is.
    """
    return _lm_head_sibling(
        "q4nx_lm_head_epilogue",
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, In, Param),
            reference=flm_gemma4_q4nx_lm_head_epilogue_ref,
            sample=partial(_lm_head_epilogue_sample, m_tile=m_tile),
            tolerance=Tolerance.bounded(
                _lm_head_epilogue_bound,
                note="c times vtanh's error, measured on npu2 (see "
                "activation._vtanh_error), plus the floor bf16 roundings",
            ),
        ),
        dim=dim,
        m_tile=m_tile,
        k_tile=k_tile,
        group=group,
    )

flm_gemma4_q4nx_lm_head_rms_ref

flm_gemma4_q4nx_lm_head_rms_ref(x, *, dim=1536)

Numpy reference for flm_gemma4_q4nx_lm_head_rms: column sums of the normalized token.

x holds, per call, the token then its RMS weight. The reference narrows each normalized value to bf16 and sums each 32 of them in float64.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_rms_ref(x, *, dim=1536):
    """Numpy reference for [`flm_gemma4_q4nx_lm_head_rms`][aie.iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head_rms]: column sums of the normalized token.

    ``x`` holds, per call, the token then its RMS weight. The reference
    narrows each normalized value to bf16 and sums each 32 of them in float64.
    """
    y = _lm_head_rms_y(x, dim).astype(np.float32).astype(bfloat16)
    sums = y.astype(np.float64).reshape(len(y), dim // 32, 32).sum(axis=2)
    return sums.astype(np.float32)

flm_gemma4_q4nx_lm_head_rms

flm_gemma4_q4nx_lm_head_rms(
    *,
    dim: int = 1536,
    m_tile: int = Q4NX_M_TILE,
    k_tile: int = Q4NX_K_TILE,
    group: int = Q4NX_GROUP
) -> ExternalFunction

q4nx_lm_head_rms of flm_gemma4_q4nx_lm_head: the token's RMS norm and per-32 column sums.

The kernel normalizes the token in place. The contract declares the token In: the write lands in the core's input element, and the next DMA fill overwrites it. The column sums depend on every normalized value.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_rms(
    *,
    dim: int = 1536,
    m_tile: int = Q4NX_M_TILE,
    k_tile: int = Q4NX_K_TILE,
    group: int = Q4NX_GROUP,
) -> ExternalFunction:
    """``q4nx_lm_head_rms`` of [`flm_gemma4_q4nx_lm_head`][aie.iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head]: the token's RMS norm and per-32 column sums.

    The kernel normalizes the token in place. The contract declares the token
    ``In``: the write lands in the core's input element, and the next DMA fill
    overwrites it. The column sums depend on every normalized value.
    """
    return _lm_head_sibling(
        "q4nx_lm_head_rms",
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=partial(flm_gemma4_q4nx_lm_head_rms_ref, dim=dim),
            tolerance=Tolerance.bounded(
                partial(_lm_head_rms_bound, dim=dim),
                note="bf16 ulp <= 2**-7 |v|. Per term: the kernel floors to bf16 "
                "(1 ulp), the reference rounds to nearest (1/2 ulp), and the fast "
                "inverse sqrt after two Newton steps adds 5e-6 relative. The "
                "kernel floors the sum to bf16 (1 ulp of the sum). 2.5 ulps of "
                "the sum of |y| plus that 5e-6 stays under 3 * 2**-7 of it",
            ),
            # Per element: square and add, two multiplies, one column-sum add.
            ops_per_call=5 * dim,
        ),
        dim=dim,
        m_tile=m_tile,
        k_tile=k_tile,
        group=group,
    )

flm_gemma4_attn_prefill_ref

flm_gemma4_attn_prefill_ref(l_bf16, y)

Numpy reference for flm_gemma4_attn_prefill: attn_epilogue at chunk 0.

y is the round's (8, 512) float32 accumulator, flat; chunk 0 is its first 64 values, eight per row of l_bf16. The kernel narrows y to bf16, multiplies by the row's 1 / l exactly and narrows again, both times with the floor rounding a fresh core uses.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_prefill_ref(l_bf16, y):
    """Numpy reference for [`flm_gemma4_attn_prefill`][aie.iron.kernels.flm_gemma4.flm_gemma4_attn_prefill]: ``attn_epilogue`` at chunk 0.

    ``y`` is the round's (8, 512) float32 accumulator, flat; chunk 0 is its
    first 64 values, eight per row of ``l_bf16``. The kernel narrows ``y`` to
    bf16, multiplies by the row's ``1 / l`` exactly and narrows again, both
    times with the floor rounding a fresh core uses.
    """
    return _prefill_epilogue(
        l_bf16, y, lq=_ATTN_PREFILL_LQ, dh=_ATTN_PREFILL_DH, chunk=0
    )

flm_gemma4_swa_prefill_ref

flm_gemma4_swa_prefill_ref(l_bf16, y)

Numpy reference for flm_gemma4_swa_prefill: attn_epilogue at chunk 33.

As flm_gemma4_attn_prefill_ref, over a (16, 256) y: rows 8 to 15, columns 64 to 127, each row scaled by its own 1 / l.

Source code in utils/pyright_stubs/aie/iron/kernels/flm_gemma4.py

def flm_gemma4_swa_prefill_ref(l_bf16, y):
    """Numpy reference for [`flm_gemma4_swa_prefill`][aie.iron.kernels.flm_gemma4.flm_gemma4_swa_prefill]: ``attn_epilogue`` at chunk 33.

    As [`flm_gemma4_attn_prefill_ref`][aie.iron.kernels.flm_gemma4.flm_gemma4_attn_prefill_ref], over a
    (16, 256) ``y``: rows 8 to 15, columns 64 to 127, each row scaled by its
    own ``1 / l``.
    """
    return _prefill_epilogue(
        l_bf16, y, lq=_SWA_PREFILL_LQ, dh=_SWA_PREFILL_DH, chunk=_SWA_PREFILL_CHUNK
    )