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 python/iron/kernels/eltwise.py

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

mul_ref

mul_ref(a, b)

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

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

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

scale_ref

scale_ref(x, factor)

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

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

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

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

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

passthrough

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for passThroughLine.

Raises:

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

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

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

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

    Returns:
        ExternalFunction configured for ``passThroughLine``.

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

scale

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the scale kernel.

Raises:

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

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

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

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

    Returns:
        ExternalFunction configured for the scale kernel.

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

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

add

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction for eltwise_add_bf16.

Raises:

Type	Description
ValueError	When dtype is not bfloat16.

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

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

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

    Returns:
        ExternalFunction for eltwise_add_bf16.

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

mul

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction for eltwise_mul_bf16.

Raises:

Type	Description
ValueError	When dtype is not bfloat16.

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

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

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

    Returns:
        ExternalFunction for eltwise_mul_bf16.

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

mul_add

mul_add(tile_size: int = 1024) -> ExternalFunction

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

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

Parameters:

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

Raises:

Type	Description
ValueError	When tile_size is not 1024.

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

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

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

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

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

mul_add_ref

mul_add_ref(a, b, is_mul)

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

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

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

relu

relu(tile_size: int = 1024) -> ExternalFunction

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction for bf16_relu.

Raises:

Type	Description
ValueError	When tile_size is not 1024.

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

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

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

    Returns:
        ExternalFunction for bf16_relu.

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

add_sized

add_sized(tile_size: int = 1024) -> ExternalFunction

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

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

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

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

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

mul_sized

mul_sized(tile_size: int = 1024) -> ExternalFunction

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

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

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

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

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

relu_sized

relu_sized(tile_size: int = 1024) -> ExternalFunction

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

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

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

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

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

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 python/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 python/iron/kernels/eltwise.py

def clamp_ref(x, low, high):
    """Numpy reference for [`clamp`][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) -> ExternalFunction

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

Takes (x, y, size): x holds tile_size real parts then as many imaginary parts, and y gets tile_size float32 magnitudes. 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	Magnitudes per call, a positive multiple of 32.	320

Returns:

Type	Description
ExternalFunction	ExternalFunction for magnitude_f32.

Raises:

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

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

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

    Takes ``(x, y, size)``: ``x`` holds ``tile_size`` real parts then as many
    imaginary parts, and ``y`` gets ``tile_size`` float32 magnitudes. 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: Magnitudes per call, a positive multiple of 32.

    Returns:
        ExternalFunction for ``magnitude_f32``.

    Raises:
        ValueError: When ``tile_size`` is not a positive multiple of 32.
    """
    _require_vector_alignment("magnitude_f32", tile_size, 32)
    return _make_extern(
        "magnitude_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}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            reference=magnitude_f32_ref,
            sample=lambda rng, calls: magnitude_f32_sample(
                rng, calls, tile_size=tile_size
            ),
            tolerance=_MAGNITUDE_F32_TOLERANCE,
            ops_per_call=tile_size,
        ),
    )

magnitude_f32_ref

magnitude_f32_ref(x)

Numpy reference for magnitude_f32.

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

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

def magnitude_f32_ref(x):
    """Numpy reference for [`magnitude_f32`][iron.kernels.eltwise.magnitude_f32].

    ``hypot`` of each call's two halves, in float64.
    """
    re, im = np.split(np.asarray(x, np.float64), 2, axis=-1)
    return np.hypot(re, 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 python/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, 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 python/iron/kernels/datamovement.py

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

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

convert_copy_ref

convert_copy_ref(x)

Numpy reference for convert_copy.

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

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

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

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

affine_cast_ref

affine_cast_ref(x, gamma_beta)

Numpy reference for affine_cast.

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

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

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

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

expand_ref

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

Numpy reference for expand.

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

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

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

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

expand_sample

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

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

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

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

transpose_ref

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

Numpy reference for transpose.

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

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

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

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

axpy

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction for the saxpy kernel.

Raises:

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

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

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

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

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

    Returns:
        ExternalFunction for the saxpy kernel.

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

convert_copy

convert_copy(tile_size: int = 1024) -> ExternalFunction

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction for cast_f32_bf16_row.

Raises:

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

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

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

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

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

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

    Returns:
        ExternalFunction for ``cast_f32_bf16_row``.

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

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 python/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 python/iron/kernels/datamovement.py

def limbs_f32_ref(x):
    """Numpy reference for [`limbs_f32`][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 python/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 python/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`][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 python/iron/kernels/datamovement.py

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

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

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

    Returns:
        ExternalFunction for ``affine_cast_f32_bf16``.

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

expand

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction for expand_uint4_to_bfloat16.

Raises:

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

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

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

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

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

    Returns:
        ExternalFunction for ``expand_uint4_to_bfloat16``.

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

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 python/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`][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 python/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 python/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 python/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`][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 python/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 python/iron/kernels/datamovement.py

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

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

rope_ref

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

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

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

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

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 python/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`][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 python/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 python/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 python/iron/kernels/datamovement.py

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

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

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

    Returns:
        ExternalFunction for the selected transpose variant.

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

Independent, target-native zero-fill kernel.

zero

zero(
    tile_size: int | tuple[int, ...] = 1024,
    dtype: type | dtype = int32,
    *,
    vectorized: bool = True,
    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 python/iron/kernels/zero.py

@dtypes(tuple({"dtype": dtype} for dtype in _TYPES) + ({"dtype": v8bfp16ebs8},))
def zero(
    tile_size: int | tuple[int, ...] = 1024,
    dtype: type | np.dtype = np.int32,
    *,
    vectorized: bool = True,
    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=32,
    k_tile=256,
    group=32,
    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 python/iron/kernels/quant.py

def q4nx_unpack(
    payload, *, m_tile=32, k_tile=256, group=32, 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=32,
    k_tile=256,
    group=32,
    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 python/iron/kernels/quant.py

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

    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=32, k_tile=256, group=32, ct_k=128, s=8, t=8
) -> ExternalFunction

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

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

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

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

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

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

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

Core state

set_rounding: the core's rounding-mode register.

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

RoundingMode

Bases: str, Enum

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

set_rounding

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

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

Parameters:

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

Returns:

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

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

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

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

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

Reduction

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

reduce_add_ref

reduce_add_ref(x)

Numpy reference for reduce_add: per-tile sum.

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

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

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

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

reduce_min_ref

reduce_min_ref(x)

Numpy reference for reduce_min: per-tile minimum.

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

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

reduce_max_ref

reduce_max_ref(x)

Numpy reference for reduce_max: per-tile maximum.

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

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

reduce_add

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the reduce_add kernel.

Raises:

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

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

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

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

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

    Returns:
        ExternalFunction configured for the reduce_add kernel.

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

reduce_min

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the reduce_min kernel.

Raises:

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

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

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

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

    Returns:
        ExternalFunction configured for the reduce_min kernel.

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

reduce_max

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the reduce_max kernel.

Raises:

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

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

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

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

    Returns:
        ExternalFunction configured for the reduce_max kernel.

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

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

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

compute_max

compute_max(dtype: type = int32) -> ExternalFunction

Pairwise scalar max — companion to reduce_max.

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

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

Parameters:

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

Returns:

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

Raises:

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

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

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

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

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

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

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

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

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

compute_max_ref

compute_max_ref(a, b)

Numpy reference for compute_max.

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

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

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

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

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 python/iron/kernels/reduce.py

def argmax_ref(x, index_offset: int = 0):
    """Numpy reference for [`argmax`][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 python/iron/kernels/reduce.py

@dtypes(({"dtype": np.int32}, {"dtype": bfloat16}))
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`][iron.kernels.reduce.reduce_max] +
    [`compute_max`][iron.kernels.reduce.compute_max]: each core runs this over
    its own slice and a tree merges the records with
    [`argmax_combine`][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 python/iron/kernels/reduce.py

def argmax_combine_ref(a, b, dtype: type = np.int32):
    """Numpy reference for [`argmax_combine`][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 python/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 python/iron/kernels/reduce.py

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

    Takes two of the records described in [`argmax`][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,
        ),
    )

Linear algebra

Linear algebra kernel factories: mm, mv, cascade_mm.

StreamDimsABC

Bases: NamedTuple

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

None for an operand a build streams untransformed.

MatrixKernel

Bases: _ZeroInitializedKernel

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

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

mac_dims property

mac_dims: tuple[int, int, int]

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

stream_dims property

stream_dims: StreamDimsABC

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

mm_acc_dtype

mm_acc_dtype(input_dtype)

Return what mm.cc accumulates in for input_dtype.

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

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

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

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

mm_ref

mm_ref(a, b)

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

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

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

mv_ref

mv_ref(a, b)

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

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

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

mm_bfp_ref

mm_bfp_ref(a, b)

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

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

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

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

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

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

mm_bfp_mixed_ref

mm_bfp_mixed_ref(a, b)

Numpy reference for mm_bfp with mixed=True.

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

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

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

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

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

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

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

mm_tile_ref

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

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

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

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

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

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

mv_tile_ref

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

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

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

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

mm_bfp_tile_ref

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

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

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

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

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

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

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

mm_stream_dims

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

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

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

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

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

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

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

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

mm

mm(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    input_dtype: type = int16,
    output_dtype: type = int16,
    vectorized: bool = True,
    b_col_maj: bool = False,
    c_col_maj: bool = False,
    use_chess: bool = False,
    emulate_bf16_mmul_with_bfp16: bool = False,
    round_conv_even: bool = False,
) -> _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.

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

@dtypes(
    tuple({"input_dtype": i, "output_dtype": o} for (i, o) in _MM_MAC_DIMS["aie2p"])
)
def mm(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    input_dtype: type = np.int16,
    output_dtype: type = np.int16,
    vectorized: bool = True,
    b_col_maj: bool = False,
    c_col_maj: bool = False,
    use_chess: bool = False,
    emulate_bf16_mmul_with_bfp16: bool = False,
    round_conv_even: bool = False,
) -> _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.
    """
    key = (input_dtype, output_dtype)
    if key not in _MM_COMBOS:
        raise ValueError(
            f"mm(): unsupported (input_dtype, output_dtype) = {key}. Supported: {list(_MM_COMBOS.keys())}"
        )

    suffix, only_flag = _MM_COMBOS[key]
    prefix = "matmul" if vectorized else "matmul_scalar"
    a_ty = np.ndarray[(dim_m * dim_k,), np.dtype[input_dtype]]
    b_ty = np.ndarray[(dim_k * dim_n,), np.dtype[input_dtype]]
    c_ty = np.ndarray[(dim_m * dim_n,), np.dtype[output_dtype]]
    compile_flags = [
        f"-DDIM_M={dim_m}",
        f"-DDIM_K={dim_k}",
        f"-DDIM_N={dim_n}",
        f"-D{only_flag}",
    ]
    if b_col_maj:
        compile_flags.append("-DB_COL_MAJ")
    if c_col_maj:
        compile_flags.append("-DC_COL_MAJ")
    arch = _detect_arch()
    bf16_emulated = (
        emulate_bf16_mmul_with_bfp16
        and ARCH_TRAITS[arch].bfp16
        and input_dtype is bfloat16
    )
    if bf16_emulated:
        compile_flags.append("-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16")
    self_rounding = round_conv_even and arch != "aie2p" and input_dtype is bfloat16
    if self_rounding:
        compile_flags.append("-DROUND_CONV_EVEN")
    # The scalar kernel walks its operands element by element in row-major
    # order: its micro-tile is 1x1x1 and nothing is streamed transformed.
    r, s, t = _MatMulFactory.mac_dims(
        input_dtype,
        output_dtype,
        arch=arch,
        emulate_bf16_mmul_with_bfp16=emulate_bf16_mmul_with_bfp16,
        vectorized=vectorized,
    )
    streams = mm_stream_dims(
        dim_m, dim_k, dim_n, (r, s, t), b_col_maj=b_col_maj, c_col_maj=c_col_maj
    )
    # Host-side layout the streams assume: with b_col_maj, B is given as
    # (n, k) tiles of B^T; with c_col_maj, C is emitted as (n, m) tiles of
    # C^T. The layouts carry that, so the builder transposes the host
    # operands without knowing which kernel it is building.
    layouts = (
        _tile_layout((dim_m, dim_k), streams.A if vectorized else None, block=(r, s)),
        _tile_layout(
            (dim_k, dim_n),
            streams.B if vectorized else None,
            axes=(1, 0) if b_col_maj else None,
            block=(s, t),
        ),
        _tile_layout(
            (dim_m, dim_n),
            streams.C if vectorized else None,
            axes=(1, 0) if c_col_maj else None,
            inverse=True,
            block=(r, t),
        ),
    )
    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 sets conv_even itself and restores it; mm_aie2.h
            # does so only under round_conv_even, and otherwise stores bf16
            # in whatever mode the core is in.
            setup=(
                conv_even
                if arch != "aie2p" and output_dtype is bfloat16 and not self_rounding
                else None
            ),
            roles=(In, In, InOut),
            reference=partial(mm_tile_ref, dim_m=dim_m, dim_k=dim_k, dim_n=dim_n),
            initializers=((2, _zero_output),),
            acc_dtype=mm_acc_dtype(input_dtype),
            reduction=dim_k,
            tolerance=_linalg_tolerance(input_dtype),
            ops_per_call=2 * dim_m * dim_k * dim_n,
        ),
    )
    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,
) -> ExternalFunction

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the matvec kernel.

Raises:

Type	Description
ValueError	When the dtype combination is not supported.

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

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

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

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

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

    Returns:
        ExternalFunction configured for the matvec kernel.

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

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

mv_bf16_ref

mv_bf16_ref(m, row_offset, a, b)

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

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

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

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

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

mm_bfp

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

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

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

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

Parameters:

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

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

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

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

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

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

mm_bfp_shuffle

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

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

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

Parameters:

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

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

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

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

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

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

mha

mha(
    dim_m: int = 64,
    dim_k: int = 64,
    dim_n: int = 64,
    pv: bool = False,
    b_col_maj: bool = False,
    emulate_bf16_mmul_with_bfp16: bool = False,
    accurate_exp2: bool = False,
    causal: bool = True,
    window: int = 0,
    f32_scores: bool = False,
) -> MatrixKernel

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

One translation unit that includes softmax_aie2p.h and mm_aie2p.h and exports the symbols an attention dataflow composes over one micro-tile. The returned kernel is one of the toolkit's two matmuls, both accumulating into C, so it is a MatrixKernel judged like mm. By default that is the QK^T product matmul_bf16_bf16_wrapper, mm_aie2p.h's bf16 product on its 4x8x8 micro-tile behind an idx_buffer gate (the call runs when key block idx[0] holds a key some row of query block idx[1] keeps) bound here to [0, 0]. With pv it is the P*V product matmul_bf16_bf16_rowmaj, mha.cc's own expansion of the native 8x8x8 micro-tile and ungated. matmul_PV is that same product into a float32 O, preceded by a row rescale, which needs the online softmax's running state and so is exercised by test_mha_e2e.py instead; 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 16).	64
dim_k	int	Depth of the micro-tile (multiple of 8).	64
dim_n	int	Columns of the micro-tile (multiple of 16).	64
pv	bool	If True return the P*V product instead of QK^T.	False
b_col_maj	bool	If True compile with -DB_COL_MAJ so the QK^T product consumes K as stored, (n, k), rather than transposed. matmul_bf16_bf16_rowmaj is row-major either way.	False
emulate_bf16_mmul_with_bfp16	bool	As for mm: both products use BFP16-based emulation, and QK^T's micro-tile becomes (8, 8, 8). Ignored on AIE2.	False
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
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. AIE2P only, and not with emulate_bf16_mmul_with_bfp16. At 64-key blocks a design places only with P V split over two cores.	False

Returns:

Type	Description
MatrixKernel	MatrixKernel for matmul_bf16_bf16_wrapper, or with pv
MatrixKernel	matmul_bf16_bf16_rowmaj.

Raises:

Type	Description
ValueError	A tile dimension that is not a positive multiple, or a window the band cannot hold (see mha_softmax).

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

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

    One translation unit that includes ``softmax_aie2p.h`` and ``mm_aie2p.h`` and
    exports the symbols an attention dataflow composes over one micro-tile.
    The returned kernel is one of the toolkit's two matmuls, both
    accumulating into ``C``, so it is a
    [`MatrixKernel`][iron.kernels.linalg.MatrixKernel] judged like
    [`mm`][iron.kernels.linalg.mm]. By default that is the ``QK^T`` product
    ``matmul_bf16_bf16_wrapper``, ``mm_aie2p.h``'s bf16 product on its 4x8x8
    micro-tile behind an ``idx_buffer`` gate (the call runs when key block
    ``idx[0]`` holds a key some row of query block ``idx[1]`` keeps) bound
    here to ``[0, 0]``. With ``pv`` it is the
    ``P*V`` product ``matmul_bf16_bf16_rowmaj``, mha.cc's own expansion of
    the native 8x8x8 micro-tile and ungated. ``matmul_PV`` is that same
    product into a float32 ``O``, preceded by a row rescale, which needs the
    online softmax's running state and so is exercised by
    ``test_mha_e2e.py`` instead; ``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 16).
        dim_k: Depth of the micro-tile (multiple of 8).
        dim_n: Columns of the micro-tile (multiple of 16).
        pv: If ``True`` return the ``P*V`` product instead of ``QK^T``.
        b_col_maj: If ``True`` compile with ``-DB_COL_MAJ`` so the ``QK^T``
            product consumes ``K`` as stored, ``(n, k)``, rather than
            transposed.  ``matmul_bf16_bf16_rowmaj`` is row-major either way.
        emulate_bf16_mmul_with_bfp16: As for [`mm`][iron.kernels.linalg.mm]:
            both products use BFP16-based emulation, and ``QK^T``'s
            micro-tile becomes (8, 8, 8).  Ignored on AIE2.
        accurate_exp2: As for [`mha_softmax`][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`][iron.kernels.linalg.mha_softmax]: a
            query keeps no key past its own position.
        window: As for [`mha_softmax`][iron.kernels.linalg.mha_softmax]: a
            query keeps no key more than ``window`` positions from it; 0 is
            no window.
        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. AIE2P
            only, and not with ``emulate_bf16_mmul_with_bfp16``. At 64-key
            blocks a design places only with ``P V`` split over two cores.

    Returns:
        MatrixKernel for ``matmul_bf16_bf16_wrapper``, or with ``pv``
        ``matmul_bf16_bf16_rowmaj``.

    Raises:
        ValueError: A tile dimension that is not a positive multiple, or a
            window the band cannot hold (see ``mha_softmax``).
    """
    for name, v, mult in (
        ("dim_m", dim_m, 16),
        ("dim_k", dim_k, 8),
        ("dim_n", dim_n, 16),
    ):
        if v <= 0 or v % mult:
            raise ValueError(
                f"mha: {name} must be a positive multiple of {mult}, got {v}"
            )
    if f32_scores and _arch_traits().name == "aie2":
        raise NotImplementedError("mha: f32_scores is written for AIE2P only")
    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[(2,), np.dtype[np.int32]]
    # mha.cc only calls mm.cc's bf16 products, so build none of the others.
    flags = [
        f"-DDIM_M={dim_m}",
        f"-DDIM_K={dim_k}",
        f"-DDIM_N={dim_n}",
        *(
            ["-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", dim_m, dim_n, causal, window)
    emulate_bf16_mmul_with_bfp16 = emulate_bf16_mmul_with_bfp16 and _arch_traits().bfp16
    if emulate_bf16_mmul_with_bfp16:
        flags.append("-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16")
    # mha.cc includes mm_aie2p.h without C_COL_MAJ, and without B_COL_MAJ unless
    # b_col_maj. matmul_bf16_bf16_rowmaj is always row-major.
    b_col_maj = b_col_maj and not pv
    r, s, t = _MhaFactory.mac_dims(
        pv=pv, emulate_bf16_mmul_with_bfp16=emulate_bf16_mmul_with_bfp16
    )
    streams = mm_stream_dims(dim_m, dim_k, dim_n, (r, s, t), b_col_maj=b_col_maj)
    return _make_extern(
        "matmul_bf16_bf16_rowmaj" if pv else "matmul_bf16_bf16_wrapper",
        _kernel_source("linalg/mha.cc"),
        [a_ty, b_ty, tile] if pv else [a_ty, b_ty, tile, idx],
        compile_flags=flags,
        cls=MatrixKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            layouts=(
                _tile_layout((dim_m, dim_k), streams.A, block=(r, s)),
                _tile_layout(
                    (dim_k, dim_n),
                    streams.B,
                    axes=(1, 0) if b_col_maj else None,
                    block=(s, t),
                ),
                _tile_layout((dim_m, dim_n), streams.C, inverse=True, block=(r, t)),
                *(() if pv else (None,)),
            ),
            roles=(In, In, InOut) if pv else (In, In, InOut, Param),
            parameter_bindings=(() if pv else ((3, np.array([0, 0], np.int32)),)),
            reference=partial(mm_tile_ref, dim_m=dim_m, dim_k=dim_k, dim_n=dim_n),
            initializers=((2, _zero_output),),
            acc_dtype=np.float32,
            reduction=dim_k,
            tolerance=(
                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_softmax

mha_softmax(
    accurate_exp2: bool = False,
    dim_m: int = 64,
    dim_n: int = 64,
    causal: bool = True,
    window: int = 0,
    scale: float = log2(e) / 8,
) -> 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), a block no query row keeps a key of is skipped, and S_q_eff/S_kv_eff are the sequence lengths whose tails pad the block. A 64-key block of a multiple of 8 rows, at least 16, takes the kernel's row-group path; any other shape the per-row one.

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 16.	64
dim_n	int	Keys of the block (B_kv), a multiple of 16.	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. A multiple of dim_m, as dim_n must then be, and for self-attention, where every query row keeps its own position.	0
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

Returns:

Type	Description
ExternalFunction	ExternalFunction for partial_softmax.

Raises:

Type	Description
ValueError	A block dimension that is not a positive multiple of 16, a scale that is not positive and finite, or a window the band cannot hold.

Source code in python/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,
    scale: float = np.log2(np.e) / 8,
) -> 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)``, a block no query row keeps a key of is
    skipped, and ``S_q_eff``/``S_kv_eff`` are the sequence lengths whose
    tails pad the block. A 64-key block of a multiple of 8 rows, at least
    16, takes the kernel's row-group path; any other shape the per-row one.

    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 16.
        dim_n: Keys of the block (``B_kv``), a multiple of 16.
        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. A multiple of
            ``dim_m``, as ``dim_n`` must then be, and for self-attention,
            where every query row keeps its own position.
        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.

    Returns:
        ExternalFunction for ``partial_softmax``.

    Raises:
        ValueError: A block dimension that is not a positive multiple of
            16, a ``scale`` that is not positive and finite, or a window
            the band cannot hold.
    """
    for name, v in (("dim_m", dim_m), ("dim_n", dim_n)):
        if v <= 0 or v % 16:
            raise ValueError(
                f"mha_softmax: {name} must be a positive multiple of 16, got {v}"
            )
    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", dim_m, dim_n, causal, window)
    band = dict(dim_m=dim_m, dim_n=dim_n, causal=causal, window=window)
    tile = np.ndarray[(dim_m * dim_n,), np.dtype[bfloat16]]
    state = np.ndarray[(4 * dim_m,), np.dtype[np.float32]]
    idx = np.ndarray[(2,), np.dtype[np.int32]]
    return _make_extern(
        "partial_softmax",
        _kernel_source("linalg/mha.cc"),
        [tile, tile, state, idx, bfloat16, *([np.int32] * 4)],
        compile_flags=[
            f"-DDIM_M={dim_m}",
            "-DDIM_K=64",
            f"-DDIM_N={dim_n}",
            *(["-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
)

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 python/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
):
    """Numpy reference for [`mha_softmax`][iron.kernels.linalg.mha_softmax]: ``(P, scale_buffer)``.

    True ``exp2`` rather than the device's interpolant, from a zeroed
    running state. ``m`` is 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 = (int(i) for i in np.asarray(idx).ravel())
    rows, cols = np.indices((dim_m, dim_n))
    key = kv * dim_n + cols - (q * dim_m + 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)
    p = np.exp2(np.where(keep, scores * s - (m * s)[..., None], -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 python/iron/kernels/linalg.py

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

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

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

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

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

    Args:
        head_dim: 512 for global attention, 256 for sliding-window.
    """
    if head_dim not in _PREFILL_GEOM:
        raise ValueError(f"prefill_fv: head_dim must be 512 or 256, got {head_dim}")
    lq, lk = _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",
                        )
                        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 python/iron/kernels/linalg.py

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

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

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

mm_bfp_shuffle_ref

mm_bfp_shuffle_ref(
    tile, tile_width, tile_height, unshuffle
)

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

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

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

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

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

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

cascade_mm

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

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

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

Parameters:

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

Raises:

Type	Description
ValueError	When the dtype combination is not supported.

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

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

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

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

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

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

cascade_mm_put

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

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

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

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

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

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

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

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

fused_mm

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

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

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

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

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

b_col_maj (with an (r, 8, 8) mmul, r 8 on aie2p) takes B as stored transposed, the (n, k) layout of a checkpoint's weights: each k chunk holds k panels of s rows, every n row of a panel s contiguous k, a layout a DMA builds from (n, k) storage in 16-byte runs. On bf16 macs the core transposes each half block in one shuffle; on bfp16 macs it converts each block as loaded.

emulate_bf16_mmul_with_bfp16 (aie2p) multiplies 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, and the reference multiplies both operands as the core sees them. The mmul defaults to 8x8x8, which b_col_maj needs. bfp16_b implies it.

epilogue_modes lists the activations compiled in, for a caller that selects one at run time through mm_fused_epilogue_chunk's mode argument. It defaults to (epilogue,). epilogue is the mode fused_mm_tile and the reference use, and must be in the list. On aie2 an activation links the tanh tables, 5 KB of L1; ("none",) links none.

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

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

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

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

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

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

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

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

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

    ``b_col_maj`` (with an ``(r, 8, 8)`` mmul, r 8 on aie2p) takes B as stored
    transposed, the ``(n, k)`` layout of a checkpoint's weights: each k chunk
    holds k panels of ``s`` rows, every n row of a panel ``s`` contiguous k,
    a layout a DMA builds from ``(n, k)`` storage in 16-byte runs. On bf16
    macs the core transposes each half block in one shuffle; on bfp16 macs
    it converts each block as loaded.

    ``emulate_bf16_mmul_with_bfp16`` (aie2p) multiplies 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, and the
    reference multiplies both operands as the core sees them. The mmul
    defaults to 8x8x8, which ``b_col_maj`` needs. ``bfp16_b`` implies it.

    ``epilogue_modes`` lists the activations compiled in, for a caller that
    selects one at run time through ``mm_fused_epilogue_chunk``'s mode
    argument. It defaults to ``(epilogue,)``. ``epilogue`` is the mode
    ``fused_mm_tile`` and the reference use, and must be in the list. On
    aie2 an activation links the tanh tables, 5 KB of L1; ``("none",)``
    links none.

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

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

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

    The returned function carries every entry point of its object as an
    attribute named by the C symbol: ``fn.mm_fused_acc_init``,
    ``fn.mm_fused_k_step``, ``fn.mm_fused_epilogue_chunk`` and
    ``fn.fused_mm_tile`` (``fn`` itself).
    """
    arch = _detect_arch()
    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 arch == "aie2p" and r != 8:
        raise ValueError("fused_mm: b_col_maj on aie2p 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 = tuple(float(v) for v in clamp)

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

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

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

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

    # 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,
    )
    acc = np.ndarray[(dim_m * dim_n,), np.dtype[np.float32]]
    b_chunk = (
        np.ndarray[(chunk_k * dim_n // 8,), np.dtype[v8bfp16ebs8]]
        if bfp16_b
        else np.ndarray[(chunk_k * dim_n,), np.dtype[bfloat16]]
    )
    bind = fn.object_file.bind
    fn.fused_mm_tile = fn
    fn.mm_fused_acc_init = bind("mm_fused_acc_init", [acc])
    # The trailing int32 is the A band.
    fn.mm_fused_k_step = bind(
        "mm_fused_k_step",
        [np.ndarray[(band_m * chunk_k,), np.dtype[bfloat16]], b_chunk, acc, np.int32],
    )
    # outer, half, mode, clamp_min_bits, clamp_max_bits
    fn.mm_fused_epilogue_chunk = bind(
        "mm_fused_epilogue_chunk",
        [np.ndarray[(out_chunk,), np.dtype[bfloat16]], acc] + [np.int32] * 5,
    )
    return fn

Convolution

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

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

conv2dk1_ref

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

Numpy reference for conv2dk1: 1x1 conv, requantized.

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

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

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

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

conv2dk1_i8_ref

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

Numpy reference for conv2dk1_i8: 1x1 conv to int8.

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

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

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

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

conv2dk1_skip_ref

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

Numpy reference for conv2dk1_skip: 1x1 conv plus residual.

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

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

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

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

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

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

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

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

conv2dk3_ref

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

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

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

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

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

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

conv2dk1_skip_init_ref

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

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

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

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

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

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

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

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

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

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

conv2dk14_ref

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

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

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

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

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

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

dwconv1d_channels_first

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

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

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

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

Parameters:

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

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

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

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

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

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

dwconv1d_channels_first_ref

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

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

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

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

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

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

dwconv1d

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

Compatibility alias for dwconv1d_channels_first.

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

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

dwconv1d_ref

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

Compatibility alias for dwconv1d_channels_first_ref.

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

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

dwconv1d_channels_last

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

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

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

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

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

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

Parameters:

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

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

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

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

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

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

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

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

dwconv1d_channels_last_ref

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

Numpy reference for dwconv1d_channels_last: one timestep over all channels.

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

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

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

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

conv2dk1

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

1x1 convolution kernel.

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk1 kernel.

Raises:

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

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

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

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

    Returns:
        ExternalFunction configured for the conv2dk1 kernel.

    Raises:
        ValueError: When ``act_dtype`` is not ``np.int8`` or ``np.uint8``.
    """
    func_name, flags = _conv_act_dtype_info(
        "conv2dk1", act_dtype, factory_name="conv2dk1"
    )
    in_ty = np.ndarray[(input_width * input_channels,), np.dtype[act_dtype]]
    wt_ty = np.ndarray[(input_channels * output_channels,), np.dtype[np.int8]]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    return _make_extern(
        func_name,
        _kernel_source("conv/conv2dk1.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(4)],
        compile_flags=flags
        + _conv_dimensions(input_width, input_channels, output_channels),
        contract=KernelContract(
            alignments=_vector_args(0, 1, 2),
            trace=Trace.whole_call(),
            roles=(In, Param, Out, Param, Param, Param, Param),
            reference=conv2dk1_ref,
            acc_dtype=np.int32,
            reduction=input_channels,
            tolerance=Tolerance.exact(
                note="measured bit-exact against the reference over every data case"
            ),
            ops_per_call=2 * input_width * input_channels * output_channels,
        ),
    )

conv2dk3

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

3x3 convolution kernel.

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk3 kernel.

Raises:

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

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

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

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

    Returns:
        ExternalFunction configured for the conv2dk3 kernel.

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

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

conv2dk1_skip

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

1x1 convolution kernel with skip (residual) connection.

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk1_skip kernel.

Raises:

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

Note

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

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

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

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

    Returns:
        ExternalFunction configured for the conv2dk1_skip kernel.

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

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

conv2dk1_i8

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk1_i8 kernel.

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

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

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

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

conv2dk14

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

14x14 convolution kernel.

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk14 kernel.

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

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

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

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

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

conv2dk1_skip_init

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

1x1 convolution kernel with skip-init connection.

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the conv2dk1_skip_init kernel.

Raises:

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

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

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

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

    Returns:
        ExternalFunction configured for the conv2dk1_skip_init kernel.

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

bn_conv2dk1_relu_ref

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

Numpy reference for bn_conv2dk1_relu.

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

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

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

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

bn_conv2dk1_i8_ref

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

Numpy reference for bn_conv2dk1_i8.

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

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

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

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

bn_conv2dk1_skip_ref

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

Numpy reference for bn_conv2dk1_skip.

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

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

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

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

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

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

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

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

bn_conv2dk3_ref

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

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

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

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

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

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

bn_conv2dk3_dw_ref

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

Numpy reference for bn_conv2dk3_dw: depthwise 3x3 + ReLU.

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

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

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

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

bn_conv2dk3_dw_out_split_ref

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

Numpy reference for bn_conv2dk3_dw_out_split.

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

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

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

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

bn_conv2dk1_relu_xy_pool_padded_ref

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

Numpy reference for bn_conv2dk1_relu_xy_pool_padded into a zeroed output.

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

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

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

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

bn_fc_relu_ui16_pad_ref

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

Numpy reference for bn_fc_relu_ui16_pad.

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

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

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

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

bn_conv2dk1_relu

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk1_relu kernel.

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

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

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

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

bn_conv2dk3

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk3 kernel.

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

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

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

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

bn_conv2dk1_i8

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk1_i8 kernel.

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

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

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

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

bn_conv2dk1_skip

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk1_skip kernel.

Raises:

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

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

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

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

    Returns:
        ExternalFunction configured for the bn_conv2dk1_skip kernel.

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

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

bn_conv2dk3_dw

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the bn_conv2dk3_dw kernel.

Raises:

Type	Description
ValueError	When stride is not 1 or 2.

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

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

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

    Returns:
        ExternalFunction configured for the bn_conv2dk3_dw kernel.

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

    func_name = f"conv2dk3_dw_stride{stride}_relu_ui8_ui8"

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

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

bn_conv2dk1_relu_xy_pool_padded

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

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

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

Parameters:

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

Returns:

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

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

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

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

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

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

bn_conv2dk1_partial_put_i8

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

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the PUT tile.

Raises:

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

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

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

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

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

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

    Returns:
        ExternalFunction configured for the PUT tile.

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

bn_conv2dk1_partial_get_relu_i8

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

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

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

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

Parameters:

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

Returns:

Type	Description
ExternalFunction	ExternalFunction configured for the GET tile.

Raises:

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

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

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

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

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

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

    Returns:
        ExternalFunction configured for the GET tile.

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

bn_conv2dk3_dw_out_split

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

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

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

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

Parameters:

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

Returns:

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

Raises:

Type	Description
ValueError	When block_index is not 13 or 14.

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

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

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

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

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

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

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

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

bn_conv2dk1_input_split_partial_put_ui8

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

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

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

Parameters:

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

Returns:

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

Raises:

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

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

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

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

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

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

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

bn_conv2dk1_input_split_partial_skip_get

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

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

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

Parameters:

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

Returns:

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

Raises:

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

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

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

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

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

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

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

bn_fc_relu_ui16_pad

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

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

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

Parameters:

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

Returns:

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

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

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

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

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

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

Activation functions

Activation kernel factories and NumPy reference implementations.

softmax

softmax(
    tile_size: int = 1024,
    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,
) -> ExternalFunction

Softmax over a row of tile_size elements, into bf16.

With bf16 input this is softmax_bf16(x, y, n). 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

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, a soft_cap outside [2^-100, 2^10], window or period without the other, or either with masked.

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

@dtypes(({"input_dtype": bfloat16}, {"input_dtype": np.float32}))
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,
) -> ExternalFunction:
    """Softmax over a row of ``tile_size`` elements, into bf16.

    With bf16 input this is ``softmax_bf16(x, y, n)``. 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).

    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, 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 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"
        )
    return _create_lut_kernel(
        "softmax_bf16",
        "softmax.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=["-DEXP2_BF16_ACCURATE=1"] if accurate_exp2 else None,
        contract=_unary_lut_contract(
            lambda x: softmax_ref(x, tile_size=tile_size),
            count=tile_size,
            # Softmax outputs sum to 1 over the tile, so a typical element is
            # about 1 / tile_size and the generic LUT floor of 0.05 would
            # accept an all-zero output. A tenth of an average element still
            # covers the exp LUT's underflow on the far tail, while an
            # unwritten tile mismatches on most elements.
            tolerance=(
                _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; "
                    "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 activation kernel (tanh approximation) for bf16 tiles.

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call (must be 1024).	1024
use_lut	bool	Compute tanh from the interpolated LUT rather than AIE2P's vtanh instruction, which 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 python/iron/kernels/activation.py

def gelu(tile_size: int = 1024, use_lut: bool = False) -> ExternalFunction:
    """GELU activation kernel (tanh approximation) for bf16 tiles.

    Args:
        tile_size: Elements per call (must be 1024).
        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`][iron.kernels.activation.gelu_lut_ref]. Moot on
            aie2, which always takes the LUT and is judged against
            [`gelu_aie2_lut_ref`][iron.kernels.activation.gelu_aie2_lut_ref].
    """
    return _bf16_lut_factory(
        "gelu",
        "gelu_bf16",
        "gelu.cc",
        tile_size,
        arg_arity=2,
        contract=_unary_lut_contract(
            gelu_ref,
            count=False,
            tolerance=_GELU_VTANH_TOLERANCE,
            use_lut=use_lut,
            elementwise=gelu_lut_ref,
            branch_models={"aie2": gelu_aie2_lut_ref},
        ),
        use_lut_tanh=use_lut,
    )

silu

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

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

Parameters:

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

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

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

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

silu_sized

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

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

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

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call.	1024
use_lut	bool	As for silu.	False

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

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

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

    Args:
        tile_size: Elements per call.
        use_lut: As for [`silu`][iron.kernels.activation.silu].
    """
    width = _arch_traits().bf16_lanes
    _require_vector_alignment("silu_sized", tile_size, width)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "silu_bf16_size",
        "silu.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DSILU_ELEMS={tile_size}"],
        contract=_unary_lut_contract(
            silu_ref,
            count=tile_size,
            use_lut=use_lut,
            elementwise=silu_table_ref,
            tolerance=_vtanh_family_tolerance("silu"),
        ),
        use_lut_tanh=use_lut,
    )

gelu_sized

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

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

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

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call.	1024
use_lut	bool	As for gelu.	False

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

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

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

    Args:
        tile_size: Elements per call.
        use_lut: As for [`gelu`][iron.kernels.activation.gelu].
    """
    _require_vector_alignment("gelu_sized", tile_size, _arch_traits().bf16_lanes)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "gelu_bf16_size",
        "gelu.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DGELU_ELEMS={tile_size}"],
        contract=_unary_lut_contract(
            gelu_ref,
            count=tile_size,
            tolerance=_GELU_VTANH_TOLERANCE,
            use_lut=use_lut,
            elementwise=gelu_lut_ref,
            branch_models={"aie2": gelu_aie2_lut_ref},
        ),
        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 python/iron/kernels/activation.py

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

    ``out = (x * w1) * silu(x * w2)``; see [`swiglu_ref`][iron.kernels.activation.swiglu_ref].
    """
    return _bf16_lut_factory(
        "swiglu",
        "swiglu_bf16",
        "swiglu.cc",
        tile_size,
        arg_arity=4,
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, In, In, Out),
            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 python/iron/kernels/activation.py

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

    Computes ``exp(clip(x, -88, 88))``: the kernel saturates rather than
    overflowing for real inputs, including infinities. On AIE2P a
    range-reduced polynomial replaces the table and rounds every result
    correctly, subnormal outputs included. See
    [`bf16_exp_ref`][iron.kernels.activation.bf16_exp_ref] for why that
    clamp matches the AIE2 table's domain.
    """
    return _bf16_lut_factory(
        "bf16_exp",
        "exp_bf16_1024",
        "bf16_exp.cc",
        tile_size,
        arg_arity=2,
        contract=_unary_lut_contract(
            bf16_exp_ref,
            count=False,
            # Only aie2's tuned branch reaches getExpBf16. The other computes a
            # range-reduced polynomial, which this model does not describe,
            # so it keeps the true-function reference and a measured bound,
            # whether or not there is a tanh instruction.
            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 python/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`]
    [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 python/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}"],
        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,
        ),
    )

tanh

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

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

The count is compiled in; retain tile_size as a trailing int argument (e.g. via transform_parallel(pass_size_to_kernel=True)).

Parameters:

Name	Type	Description	Default
tile_size	int	Elements per call (a positive multiple of 32).	1024
use_lut	bool	Compute tanh from the interpolated LUT rather than AIE2P's vtanh instruction. Moot on aie2, which only has the LUT. See tanh_lut_ref for what the LUT computes and why it is the more accurate of the two.	False

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

def tanh(tile_size: int = 1024, use_lut: bool = False) -> ExternalFunction:
    """Tanh for bf16 tiles of a positive multiple of 32 elements.

    The count is compiled in; retain
    ``tile_size`` as a trailing ``int`` argument (e.g. via
    ``transform_parallel(pass_size_to_kernel=True)``).

    Args:
        tile_size: Elements per call (a positive multiple of 32).
        use_lut: Compute tanh from the interpolated LUT rather than AIE2P's
            vtanh instruction. Moot on aie2, which only has the LUT. See
            [`tanh_lut_ref`][iron.kernels.activation.tanh_lut_ref] for what
            the LUT computes and why it is the more accurate of the two.
    """
    _require_vector_alignment("tanh", tile_size, _RUNTIME_VECTOR_WIDTH)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "tanh_bf16",
        "tanh.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DTANH_ELEMS={tile_size}"],
        contract=_unary_lut_contract(
            tanh_ref,
            count=tile_size,
            use_lut=use_lut,
            elementwise=tanh_lut_ref,
            # aie2 never reaches this: _unary_lut_contract swaps in the
            # model tolerance there, since the LUT is its only tanh.
            tolerance=_VTANH_TOLERANCE,
        ),
        use_lut_tanh=use_lut,
    )

sigmoid

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

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

The count is compiled in; retain tile_size as a trailing ABI argument.

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

def sigmoid(tile_size: int = 1024, use_lut: bool = False) -> ExternalFunction:
    """Sigmoid for bf16 tiles of a positive multiple of 32 elements.

    The count is compiled in; retain ``tile_size`` as a trailing ABI argument.
    """
    _require_vector_alignment("sigmoid", tile_size, _RUNTIME_VECTOR_WIDTH)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "sigmoid_bf16",
        "sigmoid.cc",
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DSIGMOID_ELEMS={tile_size}"],
        contract=_unary_lut_contract(
            sigmoid_ref,
            count=tile_size,
            use_lut=use_lut,
            elementwise=sigmoid_table_ref,
            tolerance=_vtanh_family_tolerance("sigmoid"),
        ),
        use_lut_tanh=use_lut,
    )

leaky_relu

leaky_relu(tile_size: int = 1024) -> ExternalFunction

Leaky ReLU for bf16 tiles of at least 64 elements, in multiples of 32.

The count is compiled in, but the ABI retains (tile_size, alpha) as trailing int/bfloat16 arguments. The slope remains runtime-valued.

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

def leaky_relu(tile_size: int = 1024) -> ExternalFunction:
    """Leaky ReLU for bf16 tiles of at least 64 elements, in multiples of 32.

    The count is compiled in, but the ABI retains ``(tile_size, alpha)`` as
    trailing ``int``/``bfloat16`` arguments. The slope remains runtime-valued.
    """
    if tile_size < 64 or tile_size % _RUNTIME_VECTOR_WIDTH:
        raise ValueError(
            "leaky_relu: tile_size must be a multiple of "
            f"{_RUNTIME_VECTOR_WIDTH} and at least 64, got {tile_size}"
        )
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "leaky_relu_bf16",
        "leaky_relu.cc",
        [tile_ty, tile_ty, np.int32, bfloat16],
        compile_flags=[f"-DLEAKY_RELU_ELEMS={tile_size}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            setup=conv_even,
            roles=(In, Out, Param, Param),
            parameter_bindings=((2, tile_size),),
            reference=leaky_relu_ref,
            acc_dtype=bfloat16,
            # max(x, alpha*x) introduces exactly one rounding, on alpha*x.
            # Measured bit-exact on npu2 over 262144 elements at alpha=0.5,
            # which is a power of two and so rounds trivially; one ulp covers
            # an alpha that does not. The 0.03/0.05 with a 2% budget this
            # replaces was three orders of magnitude looser than the
            # arithmetic allows.
            tolerance=Tolerance.bf16_ulps(
                1, note="one rounding, on alpha*x; measured bit-exact at alpha=0.5"
            ),
        ),
    )

relu_ref

relu_ref(x)

Numpy reference for a ReLU kernel — element-wise max(x, 0).

Exact; tolerance comparison is not needed. See aie.utils.verify for the relaxed bf16/LUT-style comparators most kernels here want.

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

def relu_ref(x):
    """Numpy reference for a ReLU kernel — element-wise `max(x, 0)`.

    Exact; tolerance comparison is not needed.  See `aie.utils.verify`
    for the relaxed bf16/LUT-style comparators most kernels here want.
    """
    return np.maximum(x.astype(np.float32), 0.0).astype(x.dtype)

silu_ref

silu_ref(x)

Numpy reference for silu (Swish) — x * sigmoid(x).

LUT-approximation territory; pair with rtol=0.128 (the default in count_mismatches) when verifying.

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

def silu_ref(x):
    """Numpy reference for [`silu`][iron.kernels.activation.silu] (Swish) — ``x * sigmoid(x)``.

    LUT-approximation territory; pair with ``rtol=0.128`` (the default
    in `count_mismatches`) when verifying.
    """
    xf = _no_neg_inf(_f64(x))
    with np.errstate(over="ignore"):
        return _rounded(xf / (1.0 + np.exp(-xf)), x)

gelu_ref

gelu_ref(x)

Numpy reference for gelu.

Tanh approximation 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))). Matches the C++ kernel's tanh-GELU formula. It is evaluated in float64: in float32, 1 + tanh cancels for x below about -4.5 and leaves values over ten times too large.

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

def gelu_ref(x):
    """Numpy reference for [`gelu`][iron.kernels.activation.gelu].

    Tanh approximation ``0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))``.
    Matches the C++ kernel's tanh-GELU formula. It is evaluated in float64:
    in float32, ``1 + tanh`` cancels for x below about -4.5 and leaves
    values over ten times too large.
    """
    xf = _no_neg_inf(x.astype(np.float64))
    with np.errstate(over="ignore"):
        inner = math.sqrt(2.0 / math.pi) * (xf + 0.044715 * xf**3)
    return (0.5 * xf * (1.0 + np.tanh(inner))).astype(x.dtype)

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 python/iron/kernels/activation.py

def gelu_lut_ref(x):
    """Model of AIE2P's [`gelu`][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`][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 python/iron/kernels/activation.py

def gelu_aie2_lut_ref(x):
    """Model of AIE2's [`gelu`][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`][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 python/iron/kernels/activation.py

def bf16_exp_lut_ref(x):
    """Model of ``getExpBf16``, the LUT exponential the bf16_exp kernel uses.

    The kernel clamps to ``+/-_EXP_BF16_CLAMP``, converts to Q8 with a floor
    (``bfloat16_to_int(x, 8)``), then reads the byte halves of that fixed-point
    key as two table indices and multiplies: ``exp(x) = exp(int) * exp(frac)``.
    Both tables hold exactly ``bfloat16(exp(.))``, so they are written here as
    that rule rather than as 512 opaque floats -- the unreachable middle of the
    integer table (keys the clamp cannot produce) is the only part that is not
    an exponential, and it is never read.

    The product is exact in f32 (two bf16 operands), so the model differs from
    the device only where the hardware flushes the one subnormal table entry,
    ``bfloat16(exp(-88))``.
    """
    xf = np.clip(np.asarray(x).astype(np.float32), -_EXP_BF16_CLAMP, _EXP_BF16_CLAMP)
    key = np.floor(xf * 256.0).astype(np.int32).astype(np.int16).astype(np.uint16)
    i = np.arange(256)
    ilut = np.where(
        i <= 88,
        np.exp(np.minimum(i, 88.0)),
        np.where(i >= 168, np.exp(i - 256.0), np.exp(88.0)),
    )
    ilut = np.asarray(ilut, np.float32).astype(bfloat16).astype(np.float32)
    flut = np.exp(np.arange(256) / 256.0).astype(np.float32).astype(bfloat16)
    return (ilut[key >> 8] * flut.astype(np.float32)[key & 255]).astype(np.float32)

sigmoid_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 python/iron/kernels/activation.py

def sigmoid_table_ref(x):
    """Model of [`sigmoid`][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`][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 python/iron/kernels/activation.py

def silu_table_ref(x):
    """Model of [`silu`][iron.kernels.activation.silu]'s LUT build.

    activation/silu.cc multiplies x by the bf16 sigmoid of
    [`sigmoid_table_ref`][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 python/iron/kernels/activation.py

def swiglu_table_ref(x, w1, w2):
    """Model of [`swiglu`][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`][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 python/iron/kernels/activation.py

def tanh_lut_ref(x):
    """Numpy model of ``getTanhBf16``, the interpolated-LUT tanh.

    The kernel clamps x to the table's range ``[-4, 4 - 1/64]``, then evaluates
    ``slope[e] * x + offset[e]`` for ``e = floor(4x) + 16``: 32 segments of
    width 0.25 over ``[-4, 4)``. The end segments are the constants -1 and +1,
    so the clamp changes no finite result, and +-inf gives +-1 rather than
    ``0 * inf``. The product is exact in f32 (bf16 carries 8 mantissa bits and
    8 + 8 < 24), so the only rounding is the accumulator's store back to bf16,
    which is why the build using this is judged at one ulp rather than a
    percentage.

    This is what [`tanh`][iron.kernels.activation.tanh] computes with
    ``use_lut=True``, and what it always computes on aie2. The default aie2p
    build uses the ``vtanh`` instruction instead, which is a coarser
    approximation with no published spec, so it is judged against
    [`tanh_ref`][iron.kernels.activation.tanh_ref] and a measured bound.
    """
    xf = np.clip(np.asarray(x).astype(np.float32), -4.0, 4.0 - 1.0 / 64)
    e = np.clip(np.floor(xf * 4.0).astype(np.int64), -16, 15) + 16
    slope = np.asarray(_TANH_LUT_SLOPE, np.float32)[e]
    offset = np.asarray(_TANH_LUT_OFFSET, np.float32)[e]
    return (slope * xf + offset).astype(bfloat16).astype(np.asarray(x).dtype)

tanh_ref

tanh_ref(x)

Numpy reference for tanh — element-wise tanh(x).

LUT/native-approximation territory; pair with rtol=0.128 when verifying.

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

def tanh_ref(x):
    """Numpy reference for [`tanh`][iron.kernels.activation.tanh] — element-wise ``tanh(x)``.

    LUT/native-approximation territory; pair with ``rtol=0.128`` when verifying.
    """
    return _rounded(np.tanh(_f64(x)), x)

sigmoid_ref

sigmoid_ref(x)

Numpy reference for sigmoid — 1 / (1 + exp(-x)).

LUT-approximation territory; pair with rtol=0.128 when verifying.

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

def sigmoid_ref(x):
    """Numpy reference for [`sigmoid`][iron.kernels.activation.sigmoid] — ``1 / (1 + exp(-x))``.

    LUT-approximation territory; pair with ``rtol=0.128`` when verifying.
    """
    with np.errstate(over="ignore"):
        return _rounded(1.0 / (1.0 + np.exp(-_f64(x))), x)

leaky_relu_ref

leaky_relu_ref(x, alpha=0.01)

Numpy reference for leaky_relu.

x if x > 0 else alpha * x. alpha must match the slope the design passes to the kernel at runtime; the kernel takes it as bf16, so it is rounded to bf16 here. Exact up to bf16 rounding; pair with a small rtol when verifying.

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

def leaky_relu_ref(x, alpha=0.01):
    """Numpy reference for [`leaky_relu`][iron.kernels.activation.leaky_relu].

    ``x if x > 0 else alpha * x``.  ``alpha`` must match the slope the design
    passes to the kernel at runtime; the kernel takes it as bf16, so it is
    rounded to bf16 here.  Exact up to bf16 rounding; pair with a small
    ``rtol`` when verifying.
    """
    alpha = np.float32(bfloat16(alpha))
    xf = x.astype(np.float32)
    return np.where(xf > 0.0, xf, alpha * xf).astype(x.dtype)

swiglu_ref

swiglu_ref(x, w1, w2)

Numpy reference for swiglu: (x * w1) * silu(x * w2).

swiglu.cc forms the two products in bf16, then silu of the second through the tanh LUT (0.5 * (1 + tanh(z / 2))). The reference rounds the two products to bf16 as the kernel does and computes the rest in float64; LUT-approximation territory, pair with rtol=0.128.

Where silu is 0 (x * w2 is -inf, or low enough to underflow) the output is 0, as swiglu_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 python/iron/kernels/activation.py

def swiglu_ref(x, w1, w2):
    """Numpy reference for [`swiglu`][iron.kernels.activation.swiglu]: ``(x * w1) * silu(x * w2)``.

    ``swiglu.cc`` forms the two products in bf16, then ``silu`` of the second
    through the tanh LUT (``0.5 * (1 + tanh(z / 2))``). The reference rounds
    the two products to bf16 as the kernel does and computes the rest in
    float64; LUT-approximation territory, pair with ``rtol=0.128``.

    Where silu is 0 (``x * w2`` is -inf, or low enough to underflow) the
    output is 0, as swiglu_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 python/iron/kernels/activation.py

def bf16_exp_ref(x):
    """Numpy reference for [`bf16_exp`][iron.kernels.activation.bf16_exp] — element-wise ``exp(x)``.

    LUT approximation territory; pair with the canonical 12.8% relative
    tolerance and ``stop_at_nonfinite=True`` (the default in
    `count_mismatches`) when verifying.

    ``exp(clip(x, -88, 88))``, not plain ``exp(x)``: the kernel clamps to
    ``EXP_BF16_CLAMP`` before its Q8 fixed-point table lookup (see
    ``aie_runtime_lib/AIE2/lut_based_ops.h``), so it saturates rather than
    overflowing. ``+88`` is the largest value the tables carry.
    ``exp(-88)`` is a nonzero bf16 subnormal (about ``6.06e-39``), not
    zero: AIE2P preserves it through integer exponent reconstruction,
    whereas the AIE2 LUT may flush the tail under the absolute tolerance.
    The clamp also keeps the result in range: ``exp(88) = 1.65e+38`` fits
    bf16 and float32 where ``exp(89)`` would not.
    """
    xf = np.clip(_f64(x), -_EXP_BF16_CLAMP, _EXP_BF16_CLAMP)
    # The clamp rules out overflow, so that warning stays un-suppressed and
    # would now be a real signal. A NaN input still reaches exp -- the kernel
    # declares nonfinite="unspecified" and callers do feed raw bit patterns
    # (programming_examples/basic/vector_exp sweeps all 65536 of them).
    with np.errstate(invalid="ignore"):
        return _rounded(np.exp(xf), x)

exp2f_vec_ref

exp2f_vec_ref(x, min_x: float = -111.0)

Numpy reference for exp2f_vec: exact 2**x.

Unlike the LUT-based refs above, this is float64 2**x (not a reimplementation of the on-device poly): the kernel holds 8.9e-5 relative error or better, several orders tighter than the LUT-based kernels' 12.8% default, so pair with a correspondingly tight tolerance (e.g. rtol=1e-3) rather than the LUT default.

The kernel clamps its input to min_x before evaluating (see the factory's min_x), so the reference does the same: 2**-5000 is 2**min_x on the device, not zero. Pass the factory's min_x.

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

def exp2f_vec_ref(x, min_x: float = -111.0):
    """Numpy reference for [`exp2f_vec`][iron.kernels.activation.exp2f_vec]: exact ``2**x``.

    Unlike the LUT-based refs above, this is float64 ``2**x`` (not a
    reimplementation of the on-device poly): the kernel holds 8.9e-5
    relative error or better, several orders tighter than the LUT-based
    kernels' 12.8% default, so pair with a correspondingly tight
    tolerance (e.g. ``rtol=1e-3``) rather than the LUT default.

    The kernel clamps its input to ``min_x`` before evaluating (see the
    factory's ``min_x``), so the reference does the same: ``2**-5000`` is
    ``2**min_x`` on the device, not zero. Pass the factory's ``min_x``.
    """
    xf = np.maximum(x.astype(np.float64), min_x)
    return np.exp2(xf).astype(x.dtype)

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 python/iron/kernels/activation.py

def log_f32_ref(x, offset=0.0):
    """Numpy reference for [`log_f32`][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 python/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)]

softmax_ref

softmax_ref(
    x,
    addend=None,
    mask=None,
    *,
    tile_size: int = 1024,
    soft_cap: 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; 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 python/iron/kernels/activation.py

def softmax_ref(
    x, addend=None, mask=None, *, tile_size: int = 1024, soft_cap: float | None = None
):
    """Numpy reference for [`softmax`][iron.kernels.activation.softmax].

    The AIE kernel computes softmax independently per ``tile_size``-element
    tile (no cross-tile reduction), so the reference splits ``x`` the same
    way before applying the softmax, in float64.  ``x.size`` must be a
    multiple of ``tile_size``. The logits are ``x + addend``, as
    ``cap * tanh(. / cap)`` under a ``soft_cap``, plus ``mask``; 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)
    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_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 python/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 python/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 python/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`][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 python/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`][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 python/iron/kernels/activation.py

def softmax_banded_bound(
    x, addend, row_offset, *, tile_size: int, window: int, period: int, soft_cap
):
    """[`softmax_f32_bound`][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 python/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, and layer_norm support both aie2 and aie2p. They accept tile_size (default 1024) or its compatibility alias cols. rope in Data movement has the same size API and supports both interleaved and two_halves=True layouts. The transformer module re-exports the canonical bf16 norm and RoPE factories and references; importing through either module does not select a different implementation. Norm references accept eps.

Normalization kernel factories + numpy references: rms_norm, layer_norm.

rms_norm

rms_norm(
    tile_size: int = 1024, *, cols: int | None = None
) -> ExternalFunction

RMS-norm a bf16 row on aie2/aie2p; (in, out, cols), eps=1e-5.

cols is a compatibility alias for tile_size. Positive row lengths need not be vector-aligned: the kernel handles scalar tails.

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

def rms_norm(tile_size: int = 1024, *, cols: int | None = None) -> ExternalFunction:
    """RMS-norm a bf16 row on aie2/aie2p; ``(in, out, cols)``, eps=1e-5.

    ``cols`` is a compatibility alias for ``tile_size``. Positive row lengths
    need not be vector-aligned: the kernel handles scalar tails.
    """
    tile_size = _row_size("rms_norm", tile_size, cols)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _norm_extern(
        "rms_norm",
        "rms_norm.cc",
        [tile_ty, tile_ty, np.int32],
        KernelContract(
            trace=Trace.whole_call(),
            setup=None if _tuned_arch() == "aie2" else conv_even,
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            per_build=tuple(
                (build, rms_norm_ref, _RMS_NORM_TOLERANCE.get(build.tuned, _NORM_BF16))
                for build in BUILDS
            ),
            acc_dtype=np.float32,
            reduction=tile_size,
            ops_per_call=4 * tile_size,
        ),
    )

rms_norm_eps

rms_norm_eps(
    tile_size: int = 1024, *, cols: int | None = None
) -> ExternalFunction

RMS-norm a bf16 row (gamma=1); design passes (in, out, cols, epsilon).

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

def rms_norm_eps(tile_size: int = 1024, *, cols: int | None = None) -> ExternalFunction:
    """RMS-norm a bf16 row (gamma=1); design passes ``(in, out, cols, epsilon)``."""
    tile_size = _row_size("rms_norm_eps", tile_size, cols)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _norm_extern(
        "rms_norm_eps",
        "rms_norm.cc",
        [tile_ty, tile_ty, np.int32, np.float32],
        KernelContract(
            trace=Trace.whole_call(),
            setup=None if _tuned_arch() == "aie2" else conv_even,
            roles=(In, Out, Param, Param),
            parameter_bindings=((2, tile_size),),
            per_build=tuple(
                (
                    build,
                    lambda x, epsilon: rms_norm_ref(x, eps=epsilon),
                    _RMS_NORM_TOLERANCE.get(build.tuned, _NORM_BF16),
                )
                for build in BUILDS
            ),
            acc_dtype=np.float32,
            reduction=tile_size,
            ops_per_call=4 * tile_size,
        ),
    )

layer_norm

layer_norm(
    tile_size: int = 1024, *, cols: int | None = None
) -> ExternalFunction

Layer-norm a bf16 row; (in, out, cols), gamma=1, beta=0, eps=1e-5.

cols aliases tile_size, a positive multiple of 16 (the source processes whole 16-lane halves, without a scalar tail).

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

def layer_norm(tile_size: int = 1024, *, cols: int | None = None) -> ExternalFunction:
    """Layer-norm a bf16 row; ``(in, out, cols)``, gamma=1, beta=0, eps=1e-5.

    ``cols`` aliases ``tile_size``, a positive multiple of 16 (the source
    processes whole 16-lane halves, without a scalar tail).
    """
    tile_size = _row_size("layer_norm", tile_size, cols, 16)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _norm_extern(
        "layer_norm",
        "layer_norm.cc",
        [tile_ty, tile_ty, np.int32],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, tile_size),),
            per_build=tuple(
                (
                    build,
                    layer_norm_ref,
                    _LAYER_NORM_TOLERANCE.get(build.tuned, _NORM_BF16),
                )
                for build in BUILDS
            ),
            acc_dtype=np.float32,
            reduction=tile_size,
            ops_per_call=6 * tile_size,
        ),
    )

rms_norm_ref

rms_norm_ref(x, *, eps: float = 1e-05)

Numpy reference for rms_norm: x / sqrt(mean(x**2) + eps).

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

def rms_norm_ref(x, *, eps: float = 1e-5):
    """Numpy reference for [`rms_norm`][iron.kernels.norm.rms_norm]: ``x / sqrt(mean(x**2) + eps)``."""
    xf = x.astype(np.float32)
    ms = np.mean(xf * xf, axis=-1, keepdims=True)
    return (xf / np.sqrt(ms + eps)).astype(x.dtype)

layer_norm_ref

layer_norm_ref(x, *, eps: float = 1e-05)

Numpy reference for layer_norm: (x - mean) / sqrt(var + eps).

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

def layer_norm_ref(x, *, eps: float = 1e-5):
    """Numpy reference for [`layer_norm`][iron.kernels.norm.layer_norm]: ``(x - mean) / sqrt(var + eps)``."""
    xf = x.astype(np.float32)
    mean = np.mean(xf, axis=-1, keepdims=True)
    var = np.mean((xf - mean) ** 2, axis=-1, keepdims=True)
    return ((xf - mean) / np.sqrt(var + eps)).astype(x.dtype)

Transformer blocks

Transformer building blocks: rms_norm, layer_norm (bf16, f32, affine+cast), rope, mm_activation_epilogue.

The bf16 norms and RoPE are re-exported from norm and datamovement; they support aie2 and aie2p, with cols as an alias for tile_size. The activation epilogue and the f32/affine norms take their aie2p source on both generations. Each processes one row (cols elements) per call; the row length is a scalar Param the factory binds to cols. These are the kernels programming_examples/ml/{norm,rope,mm_activation_epilogue} build.

layer_norm_f32

layer_norm_f32(cols: int = 4096) -> ExternalFunction

Row-wise LayerNorm on float32 in and out (gamma = 1, beta = 0, eps 1e-5).

A separate factory rather than a dtype of layer_norm: this one is held to atol 2e-6 (1e-3 in the portable build) instead of the bf16 tolerance, which its reference meets only by computing the variance two-pass in float64. Merging them would put that numerical difference behind a dtype switch.

Parameters:

Name	Type	Description	Default
cols	int	Elements per row (multiple of 16).	4096

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

def layer_norm_f32(cols: int = 4096) -> ExternalFunction:
    """Row-wise LayerNorm on float32 in and out (gamma = 1, beta = 0, eps 1e-5).

    A separate factory rather than a dtype of
    [`layer_norm`][iron.kernels.norm.layer_norm]: this one is held to atol 2e-6
    (1e-3 in the portable build) instead of the bf16 tolerance, which its
    reference meets only by computing the variance two-pass in float64.
    Merging them would put that numerical difference behind a dtype switch.

    Args:
        cols: Elements per row (multiple of 16).
    """
    return _row_kernel(
        "layer_norm_f32",
        "layer_norm_f32",
        "transformer/layer_norm_f32.cc",
        cols,
        np.float32,
        np.float32,
        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 python/iron/kernels/transformer.py

def layer_norm_affine_cast(cols: int = 4096) -> ExternalFunction:
    """Row-wise LayerNorm, f32 in, per-column gamma/beta, bf16 out.

    The second argument holds ``gamma`` (``cols`` values) followed by ``beta``
    (``cols`` values) as float32: a tensor ``Param``, which the generic
    builder bakes into a core buffer.

    Args:
        cols: Elements per row (multiple of 16).
    """
    _cols("layer_norm_affine_cast", cols)
    in_ty = np.ndarray[(cols,), np.dtype[np.float32]]
    gb_ty = np.ndarray[(2 * cols,), np.dtype[np.float32]]
    out_ty = np.ndarray[(cols,), np.dtype[bfloat16]]
    return _make_extern(
        "layer_norm_affine_cast",
        _kernel_source("transformer/layer_norm_f32.cc"),
        [in_ty, gb_ty, out_ty, np.int32],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Param, Out, Param),
            parameter_bindings=((3, cols),),
            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 python/iron/kernels/transformer.py

def mm_activation_epilogue(tile_size: int = 1024) -> ExternalFunction:
    """GEMM epilogue on float32 rows: identity (0), SiLU (1), tanh-GELU (2) or ReLU (3) by ``mode``.

    One resident kernel whose ``mode`` is a runtime argument, so a design can
    switch activations without recompiling
    (programming_examples/ml/mm_activation_epilogue).

    AIE2 has no tanh instruction, so there SiLU and GELU read getTanhBf16's
    table, and its tuned build is judged against a model of that arithmetic
    instead of the true functions.

    Args:
        tile_size: Elements per call (multiple of 16).
    """
    _cols("mm_activation_epilogue", tile_size)
    tile_ty = np.ndarray[(tile_size,), np.dtype[np.float32]]
    source = _kernel_source("transformer/mm_activation_epilogue.cc")
    lut = not _arch_traits().native_tanh
    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),),
            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 python/iron/kernels/transformer.py

def layer_norm_f32_ref(x):
    """Numpy reference for [`layer_norm_f32`][iron.kernels.transformer.layer_norm_f32].

    Centered two-pass variance in float64 so it stays exact on the
    non-zero-mean input the f32 kernel is exercised with.
    """
    x64 = x.astype(np.float64)
    mean = x64.mean(axis=-1, keepdims=True)
    var = ((x64 - mean) ** 2).mean(axis=-1, keepdims=True)
    return ((x64 - mean) / np.sqrt(var + _EPS)).astype(np.float32)

layer_norm_affine_cast_ref

layer_norm_affine_cast_ref(x, gamma_beta)

Numpy reference for layer_norm_affine_cast.

gamma_beta is gamma then beta, each cols float32 values.

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

def layer_norm_affine_cast_ref(x, gamma_beta):
    """Numpy reference for [`layer_norm_affine_cast`][iron.kernels.transformer.layer_norm_affine_cast].

    ``gamma_beta`` is ``gamma`` then ``beta``, each ``cols`` float32 values.
    """
    x32 = x.astype(np.float32)
    gb = np.asarray(gamma_beta, dtype=np.float32).reshape(-1)
    cols = x32.shape[-1]
    gamma, beta = gb[:cols], gb[cols : 2 * cols]
    mean = x32.mean(axis=-1, keepdims=True)
    var = ((x32 - mean) ** 2).mean(axis=-1, keepdims=True)
    return ((x32 - mean) / np.sqrt(var + _EPS) * gamma + beta).astype(bfloat16)

mm_activation_epilogue_ref

mm_activation_epilogue_ref(x, mode)

Numpy reference for mm_activation_epilogue.

mode 0 identity, 1 x * sigmoid(x), 2 tanh-approximation GELU, 3 max(x, 0).

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

def mm_activation_epilogue_ref(x, mode):
    """Numpy reference for [`mm_activation_epilogue`][iron.kernels.transformer.mm_activation_epilogue].

    ``mode`` 0 identity, 1 ``x * sigmoid(x)``, 2 tanh-approximation GELU,
    3 ``max(x, 0)``.
    """
    x32 = x.astype(np.float32)
    mode = int(mode)
    if mode == 0:
        return x32.astype(x.dtype)
    if mode == 1:
        with np.errstate(over="ignore"):
            return (x32 / (1.0 + np.exp(-x32))).astype(x.dtype)
    if mode == 2:
        # gelu's limit at -inf is 0, which the most negative float gives
        # rather than -inf * 0.
        x32 = np.maximum(x32, -np.finfo(np.float32).max)
        inner = 0.7978845608 * (x32 + 0.044715 * x32**3)
        return (0.5 * x32 * (1.0 + np.tanh(inner))).astype(x.dtype)
    if mode == 3:
        return np.maximum(x32, 0.0).astype(x.dtype)
    raise ValueError(f"mm_activation_epilogue mode must be 0, 1, 2 or 3, got {mode}")

mm_activation_epilogue_lut_ref

mm_activation_epilogue_lut_ref(x, mode)

Model of mm_activation_epilogue on aie2.

Follows mm_activation_epilogue.cc's roundings around getTanhBf16 (tanh_lut_ref). SiLU splits x into hi, its top 16 bits, and lo, bf16(x - hi), and multiplies each by the bf16 sigmoid (bf16(t + 1)) / 2, where t is the table's tanh of bf16(x) / 2 narrowed to bf16. hi is finite for any finite x, so huge inputs give about x or 0 rather than NaN; +-inf still gives NaN. GELU runs in bf16: x, x * x and the inner polynomial are each rounded before the next step, and the output is bf16(x / 2) * bf16(t + 1), all reading x clamped at -8 so -inf gives 0. Both return +0 where IEEE arithmetic gives -0, as the accumulator does. Identity and ReLU are exact.

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

def mm_activation_epilogue_lut_ref(x, mode):
    """Model of [`mm_activation_epilogue`][iron.kernels.transformer.mm_activation_epilogue] on aie2.

    Follows mm_activation_epilogue.cc's roundings around getTanhBf16
    ([`tanh_lut_ref`][iron.kernels.activation.tanh_lut_ref]). SiLU splits
    ``x`` into ``hi``, its top 16 bits, and ``lo``, ``bf16(x - hi)``, and
    multiplies each by the bf16 sigmoid ``(bf16(t + 1)) / 2``, where ``t`` is
    the table's tanh of ``bf16(x) / 2`` narrowed to bf16. ``hi`` is finite for
    any finite ``x``, so huge inputs give about ``x`` or 0 rather than NaN;
    +-inf still gives NaN. GELU runs in bf16: ``x``, ``x * x`` and the inner
    polynomial are each rounded before the next step, and the output is
    ``bf16(x / 2) * bf16(t + 1)``, all reading ``x`` clamped at -8 so -inf
    gives 0. Both return +0 where IEEE arithmetic gives -0, as the accumulator
    does. Identity and ReLU are exact.
    """
    x32 = np.asarray(x, np.float32)
    mode = int(mode)
    if mode not in (1, 2):
        return mm_activation_epilogue_ref(x, mode)
    with np.errstate(over="ignore", invalid="ignore"):
        if mode == 1:
            hi = (x32.view(np.uint32) & np.uint32(0xFFFF0000)).view(np.float32)
            lo = _bf16(x32 - hi)
            t = tanh_lut_ref(_bf16(x32) * np.float32(0.5))
            sig = _bf16(_bf16(t + np.float32(1.0)) * np.float32(0.5))
            out = hi * sig + lo * sig
        else:
            c0 = _bf16(np.float32(0.7978845608))
            c0c1 = _bf16(np.float32(0.7978845608) * np.float32(0.044715))
            xl = np.maximum(_bf16(x32), np.float32(-8.0))
            poly = _bf16(c0 + c0c1 * _bf16(xl * xl))
            t = tanh_lut_ref(_bf16(xl * poly))
            half_x = _bf16(np.float32(0.5) * xl)
            out = half_x * _bf16(t + np.float32(1.0))
    return (out + np.float32(0.0)).astype(x.dtype)

Vision

Vision kernel factories.

Color conversion, threshold, filter2d, add_weighted, 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 python/iron/kernels/vision.py

def rgba2hue(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Convert a line of RGBA pixels to hue values (full-range, 0..255)."""
    _require_vector_alignment("rgba2hue", line_width, 32, param="line_width")
    # lut_inv.h pins its gather pair with AIE_BANK_A/AIE_BANK_B.
    flags = [f"-I{_runtime_lib_include()}"]
    return _color_convert_kernel(
        "rgba2hueLine",
        "rgba2hue.cc",
        line_width * 4,
        line_width,
        use_chess=use_chess,
        compile_flags=flags,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, line_width),),
            reference=rgba2hue_ref,
            acc_dtype=np.int32,
            reduction=1,
            tolerance=Tolerance.exact(note="integer reciprocal, no rounding slack"),
            uses_lut=True,
        ),
    )

threshold

threshold(
    line_width: int = 1920,
    dtype: type = uint8,
    use_chess: bool = False,
) -> ExternalFunction

Apply a threshold operation to a line of pixels.

Parameters:

Name	Type	Description	Default
line_width	int	Number of elements per line.	1920
dtype	type	Element data type, np.uint8 or np.int16. The source's 32-bit branch multiplies int32 data by int16 coefficients, a MAC AIE2 does not have, so it does not compile and is not offered.	uint8
use_chess	bool	When True, build the .o with xchesscc_wrapper instead of Peano.	False

Raises:

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

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

@dtypes(({"dtype": np.uint8}, {"dtype": np.int16}, {"dtype": np.int32}))
def threshold(
    line_width: int = 1920, dtype: type = np.uint8, use_chess: bool = False
) -> ExternalFunction:
    """Apply a threshold operation to a line of pixels.

    Args:
        line_width: Number of elements per line.
        dtype: Element data type, ``np.uint8`` or ``np.int16``. The source's
            32-bit branch multiplies int32 data by int16 coefficients, a MAC
            AIE2 does not have, so it does not compile and is not offered.
        use_chess: When ``True``, build the .o with ``xchesscc_wrapper``
            instead of Peano.

    Raises:
        ValueError: When ``dtype`` is not ``np.uint8``, ``np.int16``, or ``np.int32``.
    """
    bit_width = _dtype_to_bit_width(dtype, factory_name="threshold")
    _require_vector_alignment(
        "threshold", line_width, 512 // bit_width, param="line_width"
    )
    scalar_ty = np.int32 if bit_width == 32 else np.int16
    line_ty = np.ndarray[(line_width,), np.dtype[dtype]]
    return _make_extern(
        "thresholdLine",
        _kernel_source("vision/threshold.cc"),
        [line_ty, line_ty, np.int32, scalar_ty, scalar_ty, np.int8],
        compile_flags=[f"-DBIT_WIDTH={bit_width}", f"-DTHRESHOLD_ELEMS={line_width}"],
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param, Param, Param, Param),
            parameter_bindings=((2, line_width),),
            reference=threshold_ref,
            tolerance=Tolerance.exact(note="selection"),
        ),
    )

bitwise_or

bitwise_or(
    line_width: int = 1920,
    dtype: type = uint8,
    use_chess: bool = False,
) -> ExternalFunction

Element-wise bitwise OR of two lines.

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

@dtypes(({"dtype": np.uint8}, {"dtype": np.int16}, {"dtype": np.int32}))
def bitwise_or(
    line_width: int = 1920, dtype: type = np.uint8, use_chess: bool = False
) -> ExternalFunction:
    """Element-wise bitwise OR of two lines."""
    return _bitwise_kernel("OR", line_width, dtype, use_chess=use_chess)

bitwise_and

bitwise_and(
    line_width: int = 1920,
    dtype: type = uint8,
    use_chess: bool = False,
) -> ExternalFunction

Element-wise bitwise AND of two lines.

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

@dtypes(({"dtype": np.uint8}, {"dtype": np.int16}, {"dtype": np.int32}))
def bitwise_and(
    line_width: int = 1920, dtype: type = np.uint8, use_chess: bool = False
) -> ExternalFunction:
    """Element-wise bitwise AND of two lines."""
    return _bitwise_kernel("AND", line_width, dtype, use_chess=use_chess)

gray2rgba

gray2rgba(
    line_width: int = 1920, use_chess: bool = False
) -> ExternalFunction

Convert a grayscale line to RGBA.

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

def gray2rgba(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Convert a grayscale line to RGBA."""
    return _color_convert_kernel(
        "gray2rgbaLine",
        "gray2rgba.cc",
        line_width,
        line_width * 4,
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Param),
            parameter_bindings=((2, line_width),),
            reference=gray2rgba_ref,
            tolerance=Tolerance.exact(note="copy with alpha = 255"),
        ),
    )

rgba2gray

rgba2gray(
    line_width: int = 1920, use_chess: bool = False
) -> ExternalFunction

Convert an RGBA line to grayscale.

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

def rgba2gray(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Convert an RGBA line to grayscale."""
    flags = []
    if not use_chess and _tuned_arch() == "aie2p":
        # The pre-RA pipeliner's schedule of the 64-pixel loop ends up at II13
        # after register allocation; the postpipeliner finds II11 (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 python/iron/kernels/vision.py

def filter2d(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Apply a 3x3 2D convolution filter across three input lines."""
    if line_width % 32 or line_width < 96:
        raise ValueError(
            f"filter2d: line_width must be a multiple of 32 and at least 96 "
            f"(the vector kernel handles the two borders in 32-pixel blocks), "
            f"got {line_width}"
        )
    line_ty = np.ndarray[(line_width,), np.dtype[np.uint8]]
    kernel_ty = np.ndarray[(3, 3), np.dtype[np.int16]]
    return _make_extern(
        "filter2dLine",
        _kernel_source("vision/filter2d.cc"),
        [line_ty, line_ty, line_ty, line_ty, np.int32, kernel_ty],
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Out, Param, Param),
            parameter_bindings=((4, line_width),),
            reference=filter2d_ref,
            acc_dtype=np.int32,
            reduction=9,
            # The one-LSB slack these pixel kernels used to share was
            # absorbing a wrong carry across the 32-pixel boundary here (see
            # filter2d.cc); an exact contract is what would have caught it.
            tolerance=Tolerance.exact(
                note="vector path matches the reference bit-for-bit"
            ),
            ops_per_call=18 * line_width,
        ),
    )

add_weighted

add_weighted(
    line_width: int = 1920,
    dtype: type = uint8,
    use_chess: bool = False,
) -> ExternalFunction

Weighted addition of two lines with a gamma offset.

Parameters:

Name	Type	Description	Default
line_width	int	Number of elements per line.	1920
dtype	type	Element data type, np.uint8 or np.int16. The source's 32-bit branch multiplies int32 data by int16 coefficients, a MAC AIE2 does not have, so it does not compile and is not offered.	uint8
use_chess	bool	When True, build the .o with xchesscc_wrapper instead of Peano.	False

Raises:

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

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

@dtypes(({"dtype": np.uint8}, {"dtype": np.int16}))
def add_weighted(
    line_width: int = 1920, dtype: type = np.uint8, use_chess: bool = False
) -> ExternalFunction:
    """Weighted addition of two lines with a gamma offset.

    Args:
        line_width: Number of elements per line.
        dtype: Element data type, ``np.uint8`` or ``np.int16``. The source's
            32-bit branch multiplies int32 data by int16 coefficients, a MAC
            AIE2 does not have, so it does not compile and is not offered.
        use_chess: When ``True``, build the .o with ``xchesscc_wrapper``
            instead of Peano.

    Raises:
        ValueError: When ``dtype`` is not ``np.uint8`` or ``np.int16``.
    """
    bit_width = _dtype_to_bit_width(dtype, factory_name="add_weighted")
    if bit_width == 32:
        raise ValueError(
            "add_weighted: no int32 build; addWeighted.cc has no int32 x int16 MAC. "
            "Use np.uint8 or np.int16."
        )
    gamma_ty = {8: np.int8, 16: np.int16}[bit_width]
    _require_vector_alignment(
        "add_weighted", line_width, 256 // bit_width, param="line_width"
    )
    line_ty = np.ndarray[(line_width,), np.dtype[dtype]]
    return _make_extern(
        "addWeightedLine",
        _kernel_source("vision/addWeighted.cc"),
        [line_ty, line_ty, line_ty, np.int32, np.int16, np.int16, gamma_ty],
        compile_flags=[
            f"-DBIT_WIDTH={bit_width}",
            f"-DADD_WEIGHTED_ELEMS={line_width}",
        ],
        use_chess=use_chess,
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out, Param, Param, Param, Param),
            parameter_bindings=((3, line_width),),
            reference=add_weighted_ref,
            acc_dtype=np.int32,
            reduction=2,
            tolerance=Tolerance.exact(
                note="measured bit-exact against the reference over every data case"
            ),
            ops_per_call=3 * line_width,
        ),
    )

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 python/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`][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 python/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`][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,
            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:

Type	Description
_ResizeKernel	ExternalFunction for resize_consume(image_chunk, height_chunk),
_ResizeKernel	with resize_setup, resize_take, resize_band,
_ResizeKernel	resize_emit and resize_finish bound on its object.

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 python/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`][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`][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`` and ``resize_finish`` bound on its object.

    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_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 python/iron/kernels/vision.py

def gray2rgba_ref(y):
    """Numpy reference for [`gray2rgba`][iron.kernels.vision.gray2rgba]: ``(y, y, y, 255)`` per pixel."""
    y = np.asarray(y, dtype=np.uint8)
    out = np.empty(y.shape[:-1] + (y.shape[-1] * 4,), dtype=np.uint8)
    px = out.reshape(*y.shape, 4)
    px[..., 0] = px[..., 1] = px[..., 2] = y
    px[..., 3] = 255
    return out

rgba2gray_ref

rgba2gray_ref(rgba)

Numpy reference for rgba2gray: fixed-point BT.470 luma.

Y = (9798 R + 19235 G + 3736 B + 2**14) >> 15, saturated to uint8; alpha is ignored. Within one LSB of the kernel (rounding of the final shift).

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

def rgba2gray_ref(rgba):
    """Numpy reference for [`rgba2gray`][iron.kernels.vision.rgba2gray]: fixed-point BT.470 luma.

    ``Y = (9798 R + 19235 G + 3736 B + 2**14) >> 15``, saturated to ``uint8``;
    alpha is ignored. Within one LSB of the kernel (rounding of the final
    shift).
    """
    rgba = np.asarray(rgba, dtype=np.uint8)
    px = rgba.reshape(*rgba.shape[:-1], -1, 4).astype(np.int64)
    wr, wg, wb = _GRAY_WEIGHTS
    acc = px[..., 0] * wr + px[..., 1] * wg + px[..., 2] * wb + (1 << 14)
    return np.clip(acc >> 15, 0, 255).astype(np.uint8)

rgba2hue_ref

rgba2hue_ref(rgba)

Numpy reference for rgba2hue: full-range hue.

rgba2hue.cc multiplies by a Q7.9 reciprocal rather than dividing, so with d = max - min of R, G, B and inv = 85 * 512 / d the hue is (offset * 512 + c * inv) >> 10 for whichever channel holds the max: c = G - B at offset 1, B - R at 171, R - G at 341. Each offset carries the + 1 that rounds the final halving, so there is one rounding step rather than two. The cast to uint8 wraps, so a negative hue (R max, G < B) comes out as 256 + h -- the right circular value. Gray pixels (d == 0) are hue 0, and a max held by both G and R goes to G, as the kernel's select order does. inv truncates, which leaves hue up to one LSB below the exact value.

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

def rgba2hue_ref(rgba):
    """Numpy reference for [`rgba2hue`][iron.kernels.vision.rgba2hue]: full-range hue.

    ``rgba2hue.cc`` multiplies by a Q7.9 reciprocal rather than dividing, so
    with ``d = max - min`` of R, G, B and ``inv = 85 * 512 / d`` the hue is
    ``(offset * 512 + c * inv) >> 10`` for whichever channel holds the max:
    ``c = G - B`` at offset 1, ``B - R`` at 171, ``R - G`` at 341.
    Each offset carries the ``+ 1`` that rounds the final halving, so there is
    one rounding step rather than two. The cast to ``uint8`` wraps, so a
    negative hue (R max, G < B) comes out as ``256 + h`` -- the right circular
    value. Gray pixels (``d == 0``) are hue 0, and a max held by both G and R
    goes to G, as the kernel's select order does. ``inv`` truncates, which
    leaves hue up to one LSB below the exact value.
    """
    rgba = np.asarray(rgba, dtype=np.uint8)
    px = rgba.reshape(*rgba.shape[:-1], -1, 4).astype(np.int64)
    r, g, b = px[..., 0], px[..., 1], px[..., 2]
    mx = np.maximum(np.maximum(r, g), b)
    d = mx - np.minimum(np.minimum(r, g), b)
    inv = (85 * 512) // np.where(d == 0, 1, d)  # avoid /0; masked to 0 below
    h = np.where(
        mx == g,
        (171 * 512 + (b - r) * inv) >> 10,
        np.where(
            mx == r,
            (1 * 512 + (g - b) * inv) >> 10,
            (341 * 512 + (r - g) * inv) >> 10,
        ),
    )
    return (np.where(d == 0, 0, h) & 0xFF).astype(np.uint8)

threshold_ref

threshold_ref(x, thresh, maxval, ttype)

Numpy reference for threshold (OpenCV semantics).

ttype: 0 binary (x > thresh ? maxval : 0), 1 binary inverted, 2 truncate (min(x, thresh)), 3 to-zero (x > thresh ? x : 0), 4 to-zero inverted. Exact.

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

def threshold_ref(x, thresh, maxval, ttype):
    """Numpy reference for [`threshold`][iron.kernels.vision.threshold] (OpenCV semantics).

    ``ttype``: 0 binary (``x > thresh ? maxval : 0``), 1 binary inverted,
    2 truncate (``min(x, thresh)``), 3 to-zero (``x > thresh ? x : 0``),
    4 to-zero inverted. Exact.
    """
    x = np.asarray(x)
    t, mv, zero = x.dtype.type(thresh), x.dtype.type(maxval), x.dtype.type(0)
    above = x > t
    ttype = int(ttype)
    if ttype == 0:
        return np.where(above, mv, zero)
    if ttype == 1:
        return np.where(above, zero, mv)
    if ttype == 2:
        return np.minimum(x, t)
    if ttype == 3:
        return np.where(above, x, zero)
    if ttype == 4:
        return np.where(above, zero, x)
    raise ValueError(f"threshold type must be 0..4, got {ttype}")

bitwise_or_ref

bitwise_or_ref(a, b)

Numpy reference for bitwise_or. Exact.

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

def bitwise_or_ref(a, b):
    """Numpy reference for [`bitwise_or`][iron.kernels.vision.bitwise_or]. Exact."""
    return np.bitwise_or(np.asarray(a), np.asarray(b))

bitwise_and_ref

bitwise_and_ref(a, b)

Numpy reference for bitwise_and. Exact.

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

def bitwise_and_ref(a, b):
    """Numpy reference for [`bitwise_and`][iron.kernels.vision.bitwise_and]. Exact."""
    return np.bitwise_and(np.asarray(a), np.asarray(b))

add_weighted_ref

add_weighted_ref(a, b, alpha, beta, gamma)

Numpy reference for add_weighted: Q2.14 blend.

out = sat(((alpha * a + beta * b) >> 14) + gamma) with alpha and beta as Q2.14 fixed point (8192 is 0.5) and gamma in output units, as OpenCV's addWeighted has it; the kernel reads gamma as the data type, so for uint8 data -56 is 200. The vector path in addWeighted.cc rounds the shift down; its scalar path rounds to nearest, so the two differ by at most one LSB.

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

def add_weighted_ref(a, b, alpha, beta, gamma):
    """Numpy reference for [`add_weighted`][iron.kernels.vision.add_weighted]: Q2.14 blend.

    ``out = sat(((alpha * a + beta * b) >> 14) + gamma)`` with ``alpha`` and
    ``beta`` as Q2.14 fixed point (``8192`` is 0.5) and ``gamma`` in output
    units, as OpenCV's ``addWeighted`` has it; the kernel reads ``gamma`` as
    the data type, so for ``uint8`` data ``-56`` is ``200``. The vector path in
    ``addWeighted.cc`` rounds the shift down; its scalar path rounds to
    nearest, so the two differ by at most one LSB.
    """
    a = np.asarray(a)
    info = np.iinfo(a.dtype)
    acc = a.astype(np.int64) * int(alpha) + np.asarray(b).astype(np.int64) * int(beta)
    acc = (acc >> 14) + int(np.array(gamma).astype(a.dtype))
    return np.clip(acc, info.min, info.max).astype(a.dtype)

filter2d_ref

filter2d_ref(line0, line1, line2, kernel)

Numpy reference for filter2d: 3x3 correlation of three lines.

Produces the middle line. The vector path keeps only the top byte of each int16 coefficient (k >> 8 as int8) and shifts the sum by 4, so a Q4.12 kernel such as 4096 * [[0, 1, 0], [1, -4, 1], [0, 1, 0]] lands on integer taps. Borders replicate the edge pixel; the result is saturated to uint8. Within one LSB of the kernel.

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

def filter2d_ref(line0, line1, line2, kernel):
    """Numpy reference for [`filter2d`][iron.kernels.vision.filter2d]: 3x3 correlation of three lines.

    Produces the middle line. The vector path keeps only the top byte of
    each ``int16`` coefficient (``k >> 8`` as ``int8``) and shifts the sum by
    4, so a Q4.12 kernel such as ``4096 * [[0, 1, 0], [1, -4, 1], [0, 1, 0]]``
    lands on integer taps. Borders replicate the edge pixel; the result is
    saturated to ``uint8``. Within one LSB of the kernel.
    """
    k8 = (np.asarray(kernel, dtype=np.int16) >> 8).astype(np.int8).astype(np.int64)
    k8 = k8.reshape(3, 3)
    terms = []
    for r, line in enumerate((line0, line1, line2)):
        v = np.asarray(line, dtype=np.uint8).astype(np.int64)
        left = np.concatenate([v[..., :1], v[..., :-1]], axis=-1)
        right = np.concatenate([v[..., 1:], v[..., -1:]], axis=-1)
        terms.append(left * k8[r, 0] + v * k8[r, 1] + right * k8[r, 2])
    acc = terms[0] + terms[1] + terms[2]
    return np.clip(acc >> 4, 0, 255).astype(np.uint8)

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 python/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`][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 python/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`][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
) -> 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

Returns:

Type	Description
ndarray	(bands * pad, 768) bf16: band b's patch column j is row
ndarray	b * pad + j, pad = (out_columns // 16 // cores + 1) * cores,
ndarray	and patch columns past the image are zeros. Sizes the kernel
ndarray	refuses zero every patch, a width chunk a core refuses zeros that
ndarray	core's patch columns, and a height chunk refused, or a band whose
ndarray	rows read past the image's, zeros every patch from that band on.

Source code in python/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,
) -> np.ndarray:
    """Numpy reference for [`resize`][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.

    Returns:
        ``(bands * pad, 768)`` bf16: band ``b``'s patch column ``j`` is row
        ``b * pad + j``, ``pad = (out_columns // 16 // 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.
    """
    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 >= side
        and out_columns >= side
        and out_rows % side == 0
        and out_columns % side == 0
        and pad // cores <= patch_columns
    )
    if not sized:
        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])
            if not (1 <= p <= 22 and 1 <= win <= win_cap) or (
                t[2] != k * side or t[3] != side
            ):
                ok[c] = False
                break
            start, count, weights = _resize_slots(t, win)
            first, end = int(start[0]), int(start[-1] + count[-1])
            if (
                (start < first).any()
                or (count < 0).any()
                or (count > win).any()
                or (start + count > 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(side):
                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])
        if not (1 <= p <= 22 and 1 <= win <= win_cap) or (
            vt[2] != b * side or vt[3] != side
        ):
            break
        start, count, weights = _resize_slots(vt, win)
        e = min(max(int(start[-1] + count[-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):
                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]
    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 as factory.dtypes, and the factories above export their references as *_ref functions (add_ref, reduce_max_ref, mm_ref, ...), so host code never reimplements the math.

Shared helpers for the kernels submodules.

KernelContract dataclass

KernelContract(
    roles: tuple[type, ...],
    reference: (
        Callable[..., ndarray | tuple[ndarray, ...]] | None
    ) = None,
    tolerance: Tolerance | None = None,
    ops_per_call: int | None = None,
    out_valid: int | None = None,
    sample: Callable[..., list] | None = None,
    acc_dtype: type | None = None,
    reduction: int | None = None,
    setup: Callable[[], object] | None = None,
    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,
    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.
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; factory(fn) returns the kernel that initializes the buffer.
out_offset	tuple[int, int] | None	(index, step) for a kernel that writes step elements of its one Out per call, at the offset it reads from bound scalar Param index. The builder then hands every call the same output tile, passes call * step as the offset and drains the tile once, so the calls must fill it exactly. None: each call writes a whole tile of its own.
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.
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 python/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 python/iron/kernels/_common.py

def reference_indices(self) -> list[int]:
    """Argument positions handed to ``reference``, in order.

    An ``InOut`` output is excluded like an ``Out`` one: the reference
    computes the result from the declared initializer's state. The
    builder initializes the buffer before each independent tile call.
    """
    bound = dict(self.parameter_bindings)
    return [
        i
        for i, r in enumerate(self.roles)
        if r not in (Out, InOut) and i not in bound
    ]

validate_types

validate_types(arg_types)

Validate contracts against NumPy tensor aliases and scalar dtypes.

Raw MLIR types remain usable by ExternalFunction, but the host contract requires NumPy declarations for sampling, layouts and references.

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

def validate_types(self, arg_types):
    """Validate contracts against NumPy tensor aliases and scalar dtypes.

    Raw MLIR types remain usable by ExternalFunction, but the host contract
    requires NumPy declarations for sampling, layouts and references.
    """
    if len(arg_types) != len(self.roles):
        raise ValueError("roles must have one entry per argument")
    bound = dict(self.parameter_bindings)
    for i, (role, arg_type) in enumerate(zip(self.roles, arg_types)):
        tensor = _is_tensor_type(arg_type)
        if tensor:
            try:
                np_ndarray_type_get_shape(arg_type)
                dtype = np_ndarray_type_get_dtype(arg_type)
            except (AssertionError, IndexError, TypeError) as exc:
                raise ValueError(
                    f"argument {i}: expected np.ndarray[shape, np.dtype[dtype]]"
                ) from exc
        else:
            dtype = arg_type
        if dtype not in get_args(NpuDType):
            raise ValueError(
                f"argument {i}: kernel contracts require NumPy tensor aliases "
                "or supported NumPy scalar dtypes"
            )
        if not tensor and role is not Param:
            raise ValueError(f"argument {i}: scalar arguments require Param")
        if self.layouts and self.layouts[i] is not None and not tensor:
            raise ValueError(f"argument {i}: layouts require tensor arguments")
        if i not in bound:
            continue
        value = bound[i]
        if tensor:
            value = np.asarray(value)
            layout = self.layouts[i] if self.layouts else None
            shape = (
                layout.shape
                if layout is not None
                else np_ndarray_type_get_shape(arg_type)
            )
            if value.ndim == 0 or value.size != int(np.prod(shape)):
                raise ValueError(
                    f"argument {i}: tensor parameter must contain {shape} elements"
                )
        elif not isinstance(value, (int, float, np.integer, np.floating)):
            raise ValueError(f"argument {i}: expected scalar parameter")

TensorLayout dataclass

TensorLayout(
    shape: tuple[int, ...],
    pack: Callable | None = None,
    unpack: Callable | None = None,
    stream: TensorAccessPattern | None = None,
    block: tuple[int, ...] | None = None,
)

How a kernel wants one tensor operand laid out.

shape is the logical tile. pack and unpack are the reversible host codec between (calls, *shape) and (calls, storage_elements); identity is the default. stream is the DMA transform (to_stream) a design applies on the hop that feeds this operand to the kernel or drains it, None when the operand streams as stored; block is the micro-tile the kernel consumes or produces, (r, s) for an MMUL operand. The codec is built from the same two facts, so the host and the design agree by construction. None of this is an algorithm or a whole-problem iteration schedule.

Param

Read-only test-fixture parameter; its ABI determines scalar or tensor.

The generic harness holds its value fixed across calls. This is not a C++ operand lifetime: direct designs may pass a new value on every kernel call.

Trace dataclass

Trace(shape: str, reason: str | None = None)

How a kernel's event0()/event1() markers bracket one call.

Trace.whole_call(): one pair brackets every call of the entry symbol and nothing it calls emits another, so each trace interval is one call. Trace.none(reason): a call emits no marker. Trace.partial(reason): markers exist but do not bracket each call exactly once (around an inner loop, or skipped on an early return), so intervals cannot be attributed to calls. test_kernel_trace_markers.py checks the declaration against the compiled IR of every library build.

aie.iron.algorithms.kernel_design turns any contract-bearing factory into a design of one Worker, built on the same single-core pipeline as transform, for_each and reduce. What a kernel can answer about itself -- its reference result, its safe input range, which arguments are parameters, how to judge a device output -- lives on ExternalFunction instead, so bringing up a kernel needs no test harness. See Kernel Library for the add-a-kernel procedure and the test tiers built on it.

Import the builder with from aie.iron.algorithms import kernel_design as kd and call kd.design(...), or import design directly from aie.iron.algorithms. The former aie.utils.kernel_harness module has been removed.

Build, sample and check independent kernel calls from their declarations.

Each call consumes one tile per In, reads constant Param values, and writes one tile per output. Layout codecs convert logical tiles to kernel storage on the host. The Worker and the runtime sequence are the shared single-core pipeline's (_pipeline). Whole-problem tiling, reductions across kernel calls and multi-core schedules belong to algorithms, not this kernel-validation harness.

CallCycles dataclass

CallCycles(
    kernel: tuple[int, ...] = (),
    initializers: dict[int, tuple[int, ...]] = dict(),
    setup: tuple[int, ...] = (),
    truncated: bool = False,
    untimed: str | None = None,
)

One traced run's intervals, split by the kernel that emitted them.

kernel holds one interval per call of the measured kernel, in call order; initializers the same for each traced initializer, keyed by the InOut argument it initializes; setup the setup kernel's one interval when it is traced. A trace that fills its buffer keeps a prefix of the stream, which the split still labels correctly, and truncated says the lists are short. untimed is the contract's reason when the kernel's markers do not bracket its calls; then nothing ran.

design

design(
    factory,
    *,
    calls=1,
    scalars=(),
    shape=None,
    params=None,
    aiecc_flags=None,
    guard=False,
    stack_bytes=None,
    arg_byte_offsets=(),
    **factory_kwargs
)

Wrap tile calls; params/scalars supply unbound tensor/scalar Params.

This harness embeds these values for every call; changing them recompiles the design. Direct designs can supply different operands on each call.

stack_bytes 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: size the outputs with output_size(..., guard=True) and split them with strip_guard. bfp outputs carry no guard.

arg_byte_offsets binds tensor Params at a byte offset: ((1, 16),) hands argument 1 a view 16 bytes into a buffer VIEW_PAD larger, poisoned around the data, as a design that packs several weights into one buffer does. The shift is loaded at run time, so the kernel sees an address no design-time check could have folded.

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

def design(
    factory,
    *,
    calls=1,
    scalars=(),
    shape=None,
    params=None,
    aiecc_flags=None,
    guard=False,
    stack_bytes=None,
    arg_byte_offsets=(),
    **factory_kwargs,
):
    """Wrap tile calls; ``params``/``scalars`` supply unbound tensor/scalar Params.

    This harness embeds these values for every call; changing them recompiles
    the design. Direct designs can supply different operands on each call.

    ``stack_bytes`` 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:
    size the outputs with ``output_size(..., guard=True)`` and split them
    with ``strip_guard``. bfp outputs carry no guard.

    ``arg_byte_offsets`` binds tensor Params at a byte offset: ``((1, 16),)``
    hands argument 1 a view 16 bytes into a buffer ``VIEW_PAD`` larger,
    poisoned around the data, as a design that packs several weights into
    one buffer does. The shift is loaded at run time, so the kernel sees an
    address no design-time check could have folded.
    """
    calls = _calls(calls, shape)
    _device()
    fn = factory(**factory_kwargs)
    c = _contract(fn)
    if c.unsupported:
        raise ValueError(
            f"{fn.name}: the generic harness cannot build this kernel: {c.unsupported}"
        )
    _fifo_plan(fn)  # the DMA channel budget is checked before anything builds
    offsets = _byte_offsets(fn, arg_byte_offsets)
    flags: list[str] = list(aiecc_flags or ())
    if any(bfp.is_bfp(shape_dtype(t)[1]) for t in fn.arg_types() if _is_tensor_type(t)):
        if "--dynamic-objFifos" not in flags:
            flags.append("--dynamic-objFifos")
    # A gather reads its two tables at once, so they have to sit in different
    # banks. Chess cannot be checked (the flag needs Peano LLVM IR), and an
    # arch whose path has no LUT just finds nothing to check.
    if c.uses_lut and not fn.use_chess:
        if not any(f.startswith("--check-lut-banks") for f in flags):
            flags.append("--check-lut-banks")
    return _stream.specialize(
        factory=factory,
        factory_kwargs=factory_kwargs,
        calls=calls,
        scalars=tuple(scalars),
        params=_encode_params(fn, params or ()),
        guard=guard,
        # Only when given, so every other design keeps its cache key.
        **({"stack_bytes": stack_bytes} if stack_bytes else {}),
        **({"arg_byte_offsets": offsets} if offsets else {}),
        **({"aiecc_flags": flags} if flags else {}),
    )

host_args

host_args(fn, *, calls=1, shape=None, guard=False)

Describe physical host buffers, not kernel arguments or constant Params.

Same-type streamed inputs share a buffer; each output has its own. Descriptors expose direction, shape, dtype and n_elements.

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

def host_args(fn, *, calls=1, shape=None, guard=False):
    """Describe physical host buffers, not kernel arguments or constant Params.

    Same-type streamed inputs share a buffer; each output has its own.
    Descriptors expose ``direction``, ``shape``, ``dtype`` and ``n_elements``.
    """
    calls = _calls(calls, shape)
    types = fn.arg_types()
    guarded = _guarded(fn, guard)
    result = []
    entries = [(In, g) for g in _fifo_plan(fn)[0]]
    entries += [(Out, [i]) for i in _contract(fn).out_indices]
    for direction, indices in entries:
        i = indices[0]
        dt = shape_dtype(types[i])[1]
        n = elems(types[i])
        if bfp.is_bfp(dt):
            n, dt = n * bfp.BLOCK_BYTES, np.uint8
        elif direction is Out and i in guarded:
            n += _guard_elems(types[i])
        rows = calls if direction is In else _out_tiles(fn, calls)
        s = (rows, n) if len(indices) == 1 else (rows, len(indices), n)
        result.append(_HostBuffer(direction, s, dt))
    return result

sample_inputs

sample_inputs(fn, *, calls=1, shape=None, rng=None)

One array per unbound In/tensor Param; only In has a call dimension.

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

def sample_inputs(fn, *, calls=1, shape=None, rng=None):
    """One array per unbound In/tensor Param; only In has a call dimension."""
    calls = _calls(calls, shape)
    rng = np.random.default_rng(0) if rng is None else rng
    c = _contract(fn)
    if c.sample is not None:
        return c.sample(rng, calls)
    result = []
    for i in _tensor_positions(fn)[0]:
        s, dt = shape_dtype(fn.arg_types()[i])
        layout = _layout(c, i)
        s = layout.shape if layout else (int(np.prod(s)),)
        dt = np.float32 if bfp.is_bfp(dt) else dt
        s = ((calls,) if c.roles[i] is In else ()) + s
        if np.issubdtype(np.dtype(dt), np.integer):
            r = fn.input_limit(dt) or _INT_RANGE.get(dt, 1 << 15)
            lo = 0 if np.dtype(dt).kind == "u" else -r
            result.append(rng.integers(lo, r, size=s).astype(dt))
        else:
            result.append(rng.standard_normal(s).astype(np.float32).astype(dt))
    return result

host_layout

host_layout(fn, inputs)

Pack declared layouts and interleave same-type In tiles; omit Param buffers.

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

def host_layout(fn, inputs):
    """Pack declared layouts and interleave same-type In tiles; omit Param buffers."""
    c = _contract(fn)
    positions = _tensor_positions(fn)[0]
    if len(inputs) != len(positions):
        raise ValueError(f"{fn.name}: expected {len(positions)} input arrays")
    values = dict(zip(positions, inputs))
    result = []
    for group in _fifo_plan(fn)[0]:
        arrays = []
        for i in group:
            value = np.asarray(values[i])
            layout = _layout(c, i)
            arrays.append(
                layout.encode(value)
                if layout
                else value.reshape(-1, elems(fn.arg_types()[i]))
            )
        result.append(
            np.ascontiguousarray(
                arrays[0] if len(arrays) == 1 else np.stack(arrays, axis=1)
            )
        )
    return result

output_size

output_size(fn, *, calls=1, shape=None, guard=False)

Storage elements per output; a tuple for multiple outputs.

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

def output_size(fn, *, calls=1, shape=None, guard=False):
    """Storage elements per output; a tuple for multiple outputs."""
    calls = _calls(calls, shape)
    types = fn.arg_types()
    guarded = _guarded(fn, guard)
    sizes = tuple(
        (elems(types[i]) + (_guard_elems(types[i]) if i in guarded else 0))
        * _out_tiles(fn, calls)
        * (bfp.BLOCK_BYTES if bfp.is_bfp(shape_dtype(types[i])[1]) else 1)
        for i in _contract(fn).out_indices
    )
    return sizes[0] if len(sizes) == 1 else sizes

upload

upload(inputs, out_size, out_dtype, *, fn, poison=False)

Return design inputs and output tensor(s); splat multiple outputs when calling.

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

def upload(inputs, out_size, out_dtype, *, fn, poison=False):
    """Return design inputs and output tensor(s); splat multiple outputs when calling."""
    ins = [
        tensor(a.reshape(-1), dtype=a.dtype, device="npu")
        for a in host_layout(fn, inputs)
    ]
    multiple = len(fn.contract.out_indices) > 1
    sizes, dtypes = (out_size, out_dtype) if multiple else ((out_size,), (out_dtype,))
    outs = [
        tensor(
            poisoned(n, dt) if poison else np.zeros(n, dtype=dt), dtype=dt, device="npu"
        )
        for n, dt in zip(sizes, dtypes)
    ]
    return ins, tuple(outs) if multiple else outs[0]

cycles_per_call

cycles_per_call(
    design_,
    inputs,
    out_size,
    out_dtype,
    *,
    fn,
    trace_size,
    workdir,
    calls=1
) -> CallCycles

Trace one run and split its intervals by kernel (CallCycles).

Only a kernel whose contract declares Trace.whole_call() is timed; one declaring none or partial returns its reason without running, and one declaring nothing raises. Traced initializers and a traced setup kernel are split off by position; a partial one raises, because its intervals cannot be labeled. Size trace_size from traced_intervals.

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

def cycles_per_call(
    design_, inputs, out_size, out_dtype, *, fn, trace_size, workdir, calls=1
) -> CallCycles:
    """Trace one run and split its intervals by kernel (``CallCycles``).

    Only a kernel whose contract declares ``Trace.whole_call()`` is timed;
    one declaring ``none`` or ``partial`` returns its reason without running,
    and one declaring nothing raises. Traced initializers and a traced setup
    kernel are split off by position; a ``partial`` one raises, because its
    intervals cannot be labeled. Size ``trace_size`` from
    ``traced_intervals``.
    """
    c = _contract(fn)
    if c.trace is None:
        raise ValueError(f"{fn.name}: the contract declares no trace")
    if c.trace.shape != "whole_call":
        return CallCycles(untimed=c.trace.reason)
    calls = _calls(calls)
    setup, inits = _traced(fn)
    workdir = Path(workdir)
    workdir.mkdir(parents=True, exist_ok=True)
    cfg = TraceConfig(trace_size=trace_size, trace_file=str(workdir / "trace.txt"))
    ins, out = upload(inputs, out_size, out_dtype, fn=fn)
    design_(*ins, *(out if isinstance(out, tuple) else (out,)), trace_config=cfg)
    if cfg.physical_mlir_path is None:
        raise RuntimeError("the traced run recorded no physical MLIR path")
    trace_json = workdir / "trace.json"
    cfg.trace_to_json(cfg.physical_mlir_path, str(trace_json))
    durations = [d for p in get_cycles_summary(str(trace_json)) for d in p[1:]]
    head, per_kernel, truncated = split_intervals(
        durations,
        calls=calls,
        per_call=len(inits) + 1,
        setup=int(setup),
        flush=TRACE_FLUSH,
    )
    if not per_kernel[-1]:
        raise RuntimeError(
            f"{fn.name}: the trace holds {len(durations)} intervals and none of "
            "the kernel's; the buffer is too small or the markers are missing"
        )
    return CallCycles(
        kernel=per_kernel[-1],
        initializers=dict(zip(inits, per_kernel)),
        setup=head,
        truncated=truncated,
    )

traced_intervals

traced_intervals(fn, *, calls=1) -> int

How many event0/event1 intervals a traced run of fn emits.

A trace_size that holds fewer truncates the run; 0 when the kernel itself is not timed.

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

def traced_intervals(fn, *, calls=1) -> int:
    """How many ``event0``/``event1`` intervals a traced run of ``fn`` emits.

    A ``trace_size`` that holds fewer truncates the run; ``0`` when the
    kernel itself is not timed.
    """
    trace = _contract(fn).trace
    if trace is None or trace.shape != "whole_call":
        return 0
    setup, inits = _traced(fn)
    return int(setup) + _calls(calls) * (len(inits) + 1)

split_intervals

split_intervals(
    durations, *, calls, per_call, setup=0, flush=0
)

Label an interval stream: setup intervals, then per_call per call.

The harness runs the setup kernel once, then each call runs its traced initializers in contract order and the kernel last, so interval j of call n sits at setup + n * per_call + j. Returns (setup intervals, one tuple per per-call kernel, truncated). Up to flush intervals of at most FLUSH_CYCLES may follow, the pairs the core emits to push the trace out. Anything else is some kernel emitting markers it does not declare; one whose extra intervals are that short and that few passes as a flush.

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

def split_intervals(durations, *, calls, per_call, setup=0, flush=0):
    """Label an interval stream: ``setup`` intervals, then ``per_call`` per call.

    The harness runs the setup kernel once, then each call runs its traced
    initializers in contract order and the kernel last, so interval ``j`` of
    call ``n`` sits at ``setup + n * per_call + j``. Returns ``(setup
    intervals, one tuple per per-call kernel, truncated)``. Up to ``flush``
    intervals of at most ``FLUSH_CYCLES`` may follow, the pairs the core
    emits to push the trace out. Anything else is some kernel emitting
    markers it does not declare; one whose extra intervals are that short
    and that few passes as a flush.
    """
    durations = [int(d) for d in durations]
    expected = setup + calls * per_call
    tail = durations[expected:]
    if len(tail) > flush or any(d > FLUSH_CYCLES for d in tail):
        raise RuntimeError(
            f"expected {expected} trace intervals and up to {flush} flush pairs, "
            f"got {len(durations)}, the extra {tail[:flush + 1]} cycles; a "
            "kernel on the core emits markers its contract's trace does not "
            "declare"
        )
    durations = durations[:expected]
    head, stream = durations[:setup], durations[setup:]
    return (
        tuple(head),
        [tuple(stream[j::per_call]) for j in range(per_call)],
        len(durations) < expected,
    )

test/python/npu/test_kernels_perf.py times the library: correctness first, then cycles, wall time and build size, gated by a device preflight and a measurement-sanity test. It is an ordinary pytest module, so -k selects cases and the session's exit status decides whether any numbers are written. --baseline-sources DIR measures every selected case a second time with its kernels from DIR and compares the two runs' raw output words and cycles in --perf-meta. The timing helpers live here:

Benchmarking helpers for NPU kernel callables.

Stats keeps the raw sample list and robust statistics (median, MAD, p95, coefficient of variation, 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.

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 python/utils/benchmark.py

def as_dict(self, prefix: str = "") -> dict:
    """Flat dict for JSON emission (samples omitted)."""
    return {
        f"{prefix}median_us": self.median_us,
        f"{prefix}mad_us": self.mad_us,
        f"{prefix}p95_us": self.p95_us,
        f"{prefix}min_us": self.min_us,
        f"{prefix}max_us": self.max_us,
        f"{prefix}avg_us": self.avg_us,
        f"{prefix}cov": self.cov,
        f"{prefix}n": self.n,
    }

BenchmarkResult dataclass

BenchmarkResult(e2e: Stats, npu: Stats | None)

Preflight dataclass

Preflight(
    npu: str, arch: str, device: str, pmode: str | None
)

What the runtime says about the device a run is about to use.

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 python/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 python/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 an npu_time (nanoseconds, captured by the runtime around kernel.wait()), it is reported separately so callers can see the host-side overhead delta.

arg_sets: an optional list of positional-argument tuples to rotate through, one per iteration, so consecutive runs do not hit the same host buffers (buffer rotation; see the CUTLASS measurement guidelines). When given, *args must be empty.

Source code in python/utils/benchmark.py

def run_iters(
    fn: Callable,
    *args,
    warmup: int = 0,
    iters: int = 1,
    arg_sets: list[tuple] | None = None,
    **kwargs,
) -> BenchmarkResult:
    """Invoke ``fn`` ``warmup + iters`` times, reporting timings.

    End-to-end latency is measured around the Python call. If the return
    value carries an ``npu_time`` (nanoseconds, captured by the runtime
    around ``kernel.wait()``), it is reported separately so callers can see
    the host-side overhead delta.

    ``arg_sets``: an optional list of positional-argument tuples to
    rotate through, one per iteration, so consecutive runs do not hit the
    same host buffers (buffer rotation; see the CUTLASS measurement
    guidelines). When given, ``*args`` must be empty.
    """
    if iters < 1:
        raise ValueError(f"iters must be >= 1 (got {iters})")
    if warmup < 0:
        raise ValueError(f"warmup must be >= 0 (got {warmup})")
    if arg_sets is not None and args:
        raise ValueError("pass either *args or arg_sets, not both")
    if arg_sets is not None and not arg_sets:
        raise ValueError("arg_sets must hold at least one argument tuple")
    rotation = arg_sets if arg_sets is not None else [args]

    e2e_samples: list[float] = []
    npu_samples: list[float] = []
    for i in range(warmup + iters):
        call_args = rotation[i % len(rotation)]
        start = time.perf_counter()
        ret = fn(*call_args, **kwargs)
        e2e_us = (time.perf_counter() - start) * 1_000_000
        if i < warmup:
            continue
        e2e_samples.append(e2e_us)
        npu_ns = _extract_npu_time_ns(ret)
        if npu_ns is not None:
            npu_samples.append(npu_ns / 1_000.0)

    return BenchmarkResult(
        e2e=_stats(e2e_samples),
        npu=_stats(npu_samples) if npu_samples else None,
    )

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

Returns:

Type	Description
list[Stats]	list[Stats]: Per callable, the statistics of its per-round medians.

Raises:

Type	Description
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 python/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,
) -> 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.

    Returns:
        list[Stats]: Per callable, the statistics of its per-round medians.

    Raises:
        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)))
    with CompletionCores(before) if completion_cores else nullcontext():
        for r in range(1, rounds + 1):
            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)
                    npu_ns = _extract_npu_time_ns(ret)
                    times.append(elapsed_us if npu_ns is None else npu_ns / 1e3)
                medians[i].append(float(np.median(times)))
                waited[i].append(float(np.median(waits)) if waits else 0.0)
            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 python/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 python/utils/benchmark.py

def preflight() -> Preflight:
    """Describe the device through the host runtime (any backend)."""
    import aie.utils as aie_utils
    from aie.utils.compile.utils import resolve_target_arch

    runtime = aie_utils.DefaultNPURuntime
    if runtime is None:
        raise RuntimeError("no NPU runtime is available (XRT, HRX or HSA)")
    device = runtime.device()
    arch = resolve_target_arch(device)
    return Preflight(
        npu="npu2" if arch == "aie2p" else "npu1",
        arch=arch,
        device=runtime.device_name() or type(device).__name__,
        pmode=runtime.power_mode(),
    )

provenance

provenance(**extra: str | None) -> str

Return a one-line description of what produced a measurement.

The git commit (GITHUB_SHA or git rev-parse HEAD), the Peano that compiles the kernels (peano_version()), the installed mlir_aie version, the kernel tree (MLIR_AIE_KERNEL_SOURCES when set, and kernel_tree_digest()), the selected runtime's provenance, the runner the job ran on (RUNNER_NAME, under GitHub Actions: nightly runs move between hosts of a generation, and a host swap moves npu_us as a regression would), and any extra fields (device="NPU Strix", pmode="performance") as key value pairs. A benchmark row records it so a number can be traced to a toolchain, the kernel sources and the machine.

A package that is not installed is left out rather than recorded as unknown. CI builds mlir_aie from source and puts it on PYTHONPATH, so there is no distribution to read a version from, and every published row would otherwise carry a word that reads like a lookup failure. The commit already identifies that build.

Source code in python/utils/benchmark.py

def provenance(**extra: str | None) -> str:
    """Return a one-line description of what produced a measurement.

    The git commit (``GITHUB_SHA`` or ``git rev-parse HEAD``), the Peano that
    compiles the kernels (``peano_version()``), the installed ``mlir_aie``
    version, the kernel tree
    (``MLIR_AIE_KERNEL_SOURCES`` when set, and ``kernel_tree_digest()``),
    the selected runtime's provenance, the runner the job ran
    on (``RUNNER_NAME``, under GitHub Actions: nightly runs move between
    hosts of a generation, and a host swap moves ``npu_us`` as a regression
    would), and any ``extra`` fields (``device="NPU Strix"``,
    ``pmode="performance"``) as ``key value`` pairs. A benchmark row records
    it so a number can be traced to a toolchain, the kernel sources and
    the machine.

    A package that is not installed is left out rather than recorded as
    unknown. CI builds ``mlir_aie`` from source and puts it on ``PYTHONPATH``,
    so there is no distribution to read a version from, and every published
    row would otherwise carry a word that reads like a lookup failure. The
    commit already identifies that build.
    """
    import aie.utils as aie_utils

    def pkg(name: str) -> str | None:
        try:
            return _pkg_version(name)
        except PackageNotFoundError:
            return None

    commit = os.environ.get("GITHUB_SHA")
    if not commit:
        try:
            commit = subprocess.run(
                ["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True
            ).stdout.strip()
        except (OSError, subprocess.CalledProcessError):
            commit = "unknown"
    runtime_fields = {}
    try:
        runtime = aie_utils.DefaultNPURuntime
        if runtime is not None:
            runtime_fields = runtime.provenance()
    except Exception:  # noqa: BLE001 - optional metadata must not fail a measurement
        pass
    fields = {
        "commit": commit[:10],
        "peano": peano_version(),
        "mlir_aie": pkg("mlir_aie"),
        "kernel_sources": os.environ.get("MLIR_AIE_KERNEL_SOURCES"),
        "kernels": kernel_tree_digest(),
        **runtime_fields,
        "host": os.environ.get("RUNNER_NAME"),
        **extra,
    }
    return " | ".join(f"{k} {v}" for k, v in fields.items() if v)

kernel_tree_digest

kernel_tree_digest() -> str | None

Return a 12-hex digest of the kernel sources the library factories compile.

Every file under aie_kernels_dir() and aie_runtime_lib_dir(), by relative path and content. The commit alone cannot say which kernels ran: MLIR_AIE_KERNEL_SOURCES can name another tree, and a checkout can carry uncommitted edits. None when neither directory exists.

Source code in python/utils/benchmark.py

def kernel_tree_digest() -> str | None:
    """Return a 12-hex digest of the kernel sources the library factories compile.

    Every file under ``aie_kernels_dir()`` and ``aie_runtime_lib_dir()``, by
    relative path and content. The commit alone cannot say which kernels
    ran: ``MLIR_AIE_KERNEL_SOURCES`` can name another tree, and a checkout
    can carry uncommitted edits. ``None`` when neither directory exists.
    """
    h = hashlib.sha256()
    found = False
    for top in (config.aie_kernels_dir(), config.aie_runtime_lib_dir()):
        root = Path(top)
        if not root.is_dir():
            continue
        found = True
        for path in sorted(p for p in root.rglob("*") if p.is_file()):
            h.update(f"{root.name}/{path.relative_to(root)}\0".encode())
            h.update(path.read_bytes())
    return h.hexdigest()[:12] if found else None

Static checks

These compiler-remark checks are available on demand through python -m aie.utils.compile.remarks; there is no static-check CI workflow. The trace-marker check (test/python/test_kernel_trace_markers.py) runs in lit on every PR, through trace_markers below.

Static kernel checks: compile every library kernel with Peano and read its remarks.

python -m aie.utils.compile.remarks --target aie2p --out static.json         --out-pm static-pm.json --meta static-meta.json
python -m aie.utils.compile.remarks --target aie2p --only '^gelu'         --out static.json --sources . --baseline-sources ../mlir-aie-base
python -m aie.utils.compile.remarks --target aie2p --only '^tanh/'         --cases test/python/npu/kernel_cases.py --out static.json
python -m aie.utils.compile.remarks --target aie2p --out static.json         --build cascade_mm:dim_m=16,dim_k=24,dim_n=32 --build cascade_mm

CPU-only. Every factory in aie.iron.kernels (at its defaults and for each entry of its .dtypes table) is compiled exactly as the JIT compiles it (aie.utils.compile.utils.cxx_core_compile_command), plus the optimization-record flags below, and the records become per-kernel series for benchmark-action: a Peano bump that changes a loop's schedule shows up here before anyone looks at device numbers. With --sources DIR (or MLIR_AIE_KERNEL_SOURCES=DIR) naming a checkout, that checkout's aie_kernels/ and aie_runtime_lib/ are compiled instead of the installed copies. With --cases, the builds are the ones a cases file's tests run (shape and options baked in), each named by its case; with --build, the factory builds named on the command line, for a shape that picks a code path no default reaches.

Record shapes, as llvm-aie 22.0.0.2026090201 emits them (they are Peano's, not LLVM's documented ones):

Pass	Kind / Name	Args	Tracked as
pipeliner	Passed / schedule	II, NS, Loop, Pipeliner, prologue/epilogue bundles	loop/<fn>/<bb>/II (rest as hover text)
pipeliner	Missed / canPipelineLoop	"Failed to pipeline loop"; located by DebugLoc only	unpipelined_loops (keyed L<line>)
pipeliner	Analysis / schedule	MII, SwpMaxMii, "Unable to find schedule"	schedule_notes in the meta file
aie-hardware-loops	Analysis / analysis	BasicBlock, Zero-Overhead-Loop	non_zol_loops, and loop/<fn>/<bb>/not_zol per loop
aie-asm-printer	Analysis / analysis	BasicBlock, BundleCount, ByteCount	pm_bytes (summed over the shipped functions)
aie-multi-slot-pseudo	Missed / missing-memory-bank	Instruction	missing_bank_loads
stderr	-Wpass-failed	a #pragma clang loop / AIE_* macro the compiler dropped	pass_failed_warnings, text kept
the object	llvm-readobj sections, symbols, relocations	what the entry symbol reaches	the shipped functions; libcalls (e.g. __divsf3)
the object	llvm-readobj --stack-sizes (-fstack-size-section)	frame bytes per function	kernel_stack_bytes, the deepest path from the entry
the IR (-emit-llvm)	opt print<scalar-evolution>	a constant backedge-taken count	loop/<fn>/<bb>/II_x_trips

The loop counts and pm_bytes cover only the functions the entry symbol reaches in the object, which are the ones the core link keeps. kernel_stack_bytes counts the kernel's frames only: the core that calls it also holds main's, which aiecc's measured stack size includes (dwconv1d_channels_last on AIE2P: 64 here, 256 for the core). --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 python/utils/compile/remarks.py

def parse_yaml(path: str | Path, report: StaticReport | None = None) -> StaticReport:
    r = report or StaticReport()
    with open(path) as f:
        docs = list(yaml.load_all(f, Loader=_Loader))
    for d in docs:
        if not isinstance(d, dict):
            continue
        p, name, kind = d.get("Pass"), d.get("Name"), d.get("_kind")
        fn = d.get("Function", "?")
        a = _args(d)
        if p == "pipeliner":
            if kind == "Passed" and name == "schedule":
                # A schedule with no Loop arg has been seen (zero_scalar_*);
                # fall back to the source line rather than a shared "?" key.
                loop = r.loop(fn, _block(a.get("Loop")) or _loc(d))
                loop.pipelined = True
                loop.ii = _int(a.get("II"))
                loop.ns = _int(a.get("NS"))
                loop.prologue_bundles = _int(a.get("PrologueBundles"))
                loop.epilogue_bundles = _int(a.get("EpilogueBundles"))
                loop.pipeliner = str(a.get("Pipeliner", "")) or None
                loop.file, loop.line = _source(d)
            elif kind == "Missed" and name == "canPipelineLoop":
                loop = r.loop(fn, _loc(d))
                loop.pipelined = False
                loop.missed_reason = str(a.get("String", "")) or None
                loop.file, loop.line = _source(d)
            elif kind == "Analysis" and name == "schedule":
                # The sentence as clang prints it: every Arg value in order,
                # since the typed values (MII, II, ...) sit between the
                # String fragments and joining only the strings would print
                # "Minimal Initiation Interval too large:  > ." without them.
                message = "".join(
                    str(v)
                    for arg in d.get("Args", []) or []
                    if isinstance(arg, dict)
                    for k, v in arg.items()
                    if k != "DebugLoc"
                )
                r.schedule_notes.append(f"{fn}@{_loc(d)}: {message}")
        elif p == "aie-hardware-loops":
            bb = str(a.get("BasicBlock", "?"))
            r.loop(fn, bb).zol = str(a.get("Zero-Overhead-Loop", "")).lower() == "true"
        elif p == "aie-asm-printer":
            bb = str(a.get("BasicBlock", "?"))
            bc, byc = _int(a.get("BundleCount")), _int(a.get("ByteCount"))
            r.pm_bytes_by_function[fn] = r.pm_bytes_by_function.get(fn, 0) + (byc or 0)
            if (fn, bb) in r.loops:  # only annotate known loop blocks
                r.loops[(fn, bb)].bundle_count = bc
                r.loops[(fn, bb)].byte_count = byc
        elif p == "aie-multi-slot-pseudo" and name == "missing-memory-bank":
            r.missing_bank_loads += 1
    return r

parse_stderr

parse_stderr(
    text: str, report: StaticReport
) -> StaticReport

Source code in python/utils/compile/remarks.py

def parse_stderr(text: str, report: StaticReport) -> StaticReport:
    dropped = [line.strip() for line in _WARN_PASS_FAILED.findall(text)]
    report.pass_failed_warnings += len(dropped)
    report.pass_failed.extend(dropped)
    # The stderr warning duplicates the YAML remark; only count it when no
    # YAML record was produced for the build (keeps the two channels consistent).
    if report.missing_bank_loads == 0:
        report.missing_bank_loads = len(_WARN_NO_BANK.findall(text))
    return report

report_rows

report_rows(
    report: StaticReport, prefix: str, extra: str
) -> list[dict]

benchmark-action rows for one kernel build. Smaller is better throughout.

Source code in python/utils/compile/remarks.py

def report_rows(report: StaticReport, prefix: str, extra: str) -> list[dict]:
    """benchmark-action rows for one kernel build. Smaller is better throughout."""
    out = [
        _row(f"{prefix}/unpipelined_loops", "loops", report.unpipelined_loops, extra),
        _row(f"{prefix}/non_zol_loops", "loops", report.non_zol_loops, extra),
        _row(f"{prefix}/missing_bank_loads", "loads", report.missing_bank_loads, extra),
        _row(
            f"{prefix}/pass_failed_warnings",
            "warnings",
            report.pass_failed_warnings,
            extra,
        ),
        _row(f"{prefix}/pm_bytes", "bytes", report.pm_bytes, extra),
        _row(
            f"{prefix}/libcalls",
            "symbols",
            len(report.libcalls),
            extra,
            " ".join(report.libcalls) or None,
        ),
    ]
    if report.kernel_stack_bytes is not None:
        out.append(
            _row(
                f"{prefix}/kernel_stack_bytes",
                "bytes",
                report.kernel_stack_bytes,
                extra,
                "the kernel's frames; the core adds main's"
                + (" and the runtime routines'" if report.libcalls else ""),
            )
        )
    for (fn, bb), loop in sorted(report.loops.items()):
        if loop.ii is not None:
            # The hover text names the source line, so a reader of an II
            # alert can open the loop without decoding a basic-block name.
            where = (
                f" at {Path(loop.file).name}:{loop.line}"
                if loop.file and loop.line is not None
                else ""
            )
            out.append(
                _row(
                    f"{prefix}/loop/{fn}/{bb}/II",
                    "cycles",
                    loop.ii,
                    extra,
                    f"NS={loop.ns} pro={loop.prologue_bundles} "
                    f"epi={loop.epilogue_bundles} zol={loop.zol} "
                    f"trips={loop.trips} via={loop.pipeliner}{where}",
                )
            )
            if loop.trips is not None:
                # II alone reads a loop that does 4 blocks per iteration as
                # slower than one that does 1; II x trips compares them.
                out.append(
                    _row(
                        f"{prefix}/loop/{fn}/{bb}/II_x_trips",
                        "cycles",
                        loop.ii * loop.trips,
                        extra,
                        f"{loop.trips} trips per entry{where}",
                    )
                )
        if loop.zol is not None:
            # 1 when the loop is not a zero-overhead loop: an inner loop that
            # falls off the hardware loop unit costs its bundle count in
            # branches, whatever its II says.
            out.append(
                _row(
                    f"{prefix}/loop/{fn}/{bb}/not_zol",
                    "loops",
                    int(not loop.zol),
                    extra,
                )
            )
    return out

workflow_annotations

workflow_annotations(
    name: str,
    report: StaticReport | None,
    detail: str,
    root: str | None = None,
    seen: set[tuple] | None = None,
) -> list[str]

Annotations for one kernel build: its dropped pragmas, or its compile failure.

One source is compiled once per factory build and per target, so the same dropped pragma comes back several times a run; seen (shared across builds) keeps each file, line and message to its first annotation.

Source code in python/utils/compile/remarks.py

def workflow_annotations(
    name: str,
    report: StaticReport | None,
    detail: str,
    root: str | None = None,
    seen: set[tuple] | None = None,
) -> list[str]:
    """Annotations for one kernel build: its dropped pragmas, or its compile failure.

    One source is compiled once per factory build and per target, so the same
    dropped pragma comes back several times a run; ``seen`` (shared across
    builds) keeps each file, line and message to its first annotation.
    """
    if report is None:
        first = next((ln for ln in detail.splitlines() if "error:" in ln), None)
        return [
            annotation(
                "error",
                first or detail.strip()[-500:],
                title=f"{name}: kernel failed to compile",
            )
        ]
    out = []
    for warning in report.pass_failed:
        m = _DIAG_RE.match(warning)
        key = (
            (m.group("file"), m.group("line"), warning[m.end() :]) if m else (warning,)
        )
        if seen is not None:
            if key in seen:
                continue
            seen.add(key)
        out.append(
            annotation(
                "warning",
                warning[m.end() :] if m else warning,
                title=f"{name}: pragma dropped by the compiler",
                file=m.group("file") if m else None,
                line=int(m.group("line")) if m else None,
                root=root,
            )
        )
    return out

compile_command

compile_command(
    ext_fn, target: str, out_dir: Path
) -> tuple[list[str], Path]

Return the exact Peano command the JIT would run for ext_fn, plus remark flags.

Inline-source kernels (the aie2 LUT activations) are written out under the kernel's symbol name first, as the JIT does. Kernels built with use_chess are rejected: the remarks are Peano's.

Source code in python/utils/compile/remarks.py

def compile_command(ext_fn, target: str, out_dir: Path) -> tuple[list[str], Path]:
    """Return the exact Peano command the JIT would run for ``ext_fn``, plus remark flags.

    Inline-source kernels (the aie2 LUT activations) are written out under the
    kernel's symbol name first, as the JIT does. Kernels built with
    ``use_chess`` are rejected: the remarks are Peano's.
    """
    src, include_dirs = _kernel_file(ext_fn, out_dir)
    yaml_out = out_dir / f"{ext_fn.name}.opt.yaml"
    cmd = cxx_core_compile_command(
        str(src),
        target,
        str(out_dir / f"{ext_fn.name}.o"),
        include_dirs=include_dirs,
        compile_args=[
            *ext_fn.compile_flags,
            *REMARK_FLAGS,
            f"-foptimization-record-file={yaml_out}",
            *_EXTRA_WARNINGS,
            # A non-allocated section of frame sizes; the code is unchanged.
            "-fstack-size-section",
        ],
    )
    return cmd, yaml_out

analyze

analyze(
    ext_fn, target: str, workdir: Path
) -> tuple[StaticReport | None, str]

Compile one kernel and parse its records; (None, reason) when it fails to compile.

Source code in python/utils/compile/remarks.py

def analyze(ext_fn, target: str, workdir: Path) -> tuple[StaticReport | None, str]:
    """Compile one kernel and parse its records; ``(None, reason)`` when it fails to compile."""
    cmd, yaml_out = compile_command(ext_fn, target, workdir)
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        return None, f"compile failed: {compile_failure(p.stderr)}"
    rep = parse_yaml(yaml_out) if yaml_out.exists() else StaticReport()
    for key, trips in trip_counts(ext_fn, target, workdir).items():
        if key in rep.loops:
            rep.loops[key].trips = trips
    parse_stderr(p.stderr, rep)
    reached = linked(workdir / f"{ext_fn.name}.o", entry_symbol(ext_fn))
    rep.shipped, rep.libcalls = reached.functions, reached.undefined
    rep.kernel_stack_bytes = reached.stack
    return rep, "ok"

kernel_builds

kernel_builds()

Yield (name, ExternalFunction) for every factory build the library offers.

Each exported factory at its defaults, plus one build per entry of its .dtypes table; a factory that refuses the current device (NotImplementedError) is skipped. Remarks depend on source and flags, not on the shape a test runs, so this is the whole surface.

Source code in python/utils/compile/remarks.py

def kernel_builds():
    """Yield ``(name, ExternalFunction)`` for every factory build the library offers.

    Each exported factory at its defaults, plus one build per entry of its
    ``.dtypes`` table; a factory that refuses the current device
    (``NotImplementedError``) is skipped. Remarks depend on source and flags,
    not on the shape a test runs, so this is the whole surface.
    """
    from aie.iron import kernels

    for name in kernels.factories():
        f = getattr(kernels, name)
        combos = [{}] + [dict(c) for c in getattr(f, "dtypes", ()) if c]
        seen: set[str] = set()
        for combo in combos:
            try:
                ef = f(**combo)
            except NotImplementedError:
                continue  # exists only for the other architecture
            if ef.object_file_name in seen:
                continue  # the default build is one of the dtypes entries
            seen.add(ef.object_file_name)
            yield _build_name(name, combo), ef

trace_markers

trace_markers(ir: str, entry: str) -> str

Classify how event0()/event1() bracket a call of entry.

"whole_call" when one pair brackets every call and nothing else emits a marker, "none" when a call emits no marker at all, and otherwise a sentence saying what the markers do instead.

Source code in python/utils/compile/remarks.py

def trace_markers(ir: str, entry: str) -> str:
    """Classify how ``event0()``/``event1()`` bracket a call of ``entry``.

    ``"whole_call"`` when one pair brackets every call and nothing else
    emits a marker, ``"none"`` when a call emits no marker at all, and
    otherwise a sentence saying what the markers do instead.
    """
    functions = _ir_functions(ir)
    if entry not in functions:
        raise ValueError(f"{entry} is not defined in the compiled IR")
    paths = _marker_paths(functions, entry, {})
    if isinstance(paths, str):
        return paths
    if paths == {"01"}:
        return "whole_call"
    if paths == {""}:
        return "none"
    shown = ", ".join(repr(p) for p in sorted(paths))
    return (
        f"a call of {entry} emits marker sequences {shown}, not one event0 then event1"
    )

trace_shape

trace_shape(ext_fn, target: str, out_dir: Path) -> str

Compile ext_fn to -O2 IR and classify its entry's markers (trace_markers).

Source code in python/utils/compile/remarks.py

def trace_shape(ext_fn, target: str, out_dir: Path) -> str:
    """Compile ``ext_fn`` to -O2 IR and classify its entry's markers (``trace_markers``)."""
    src, include_dirs = _kernel_file(ext_fn, out_dir)
    cmd = cxx_core_compile_command(
        str(src),
        target,
        str(out_dir / f"{ext_fn.name}.ll"),
        include_dirs=include_dirs,
        compile_args=list(ext_fn.compile_flags),
        inline=True,
    )
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        raise RuntimeError(f"{ext_fn.name}: {compile_failure(p.stderr)}")
    ir = (out_dir / f"{ext_fn.name}.ll").read_text()
    return trace_markers(ir, entry_symbol(ext_fn))

entry_symbol

entry_symbol(ext_fn) -> str

Return the symbol the kernel source defines, before the JIT's per-build prefix.

Source code in python/utils/compile/remarks.py

def entry_symbol(ext_fn) -> str:
    """Return the symbol the kernel source defines, before the JIT's per-build prefix."""
    return getattr(ext_fn, "_original_name", ext_fn.name)

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

linked

linked(obj: Path, entry: str) -> Linked

Functions and undefined symbols entry reaches in the object obj.

Source code in python/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 python/utils/compile/readobj.py

def parse_readobj(doc: dict, entry: str) -> Linked:
    """``linked`` on the parsed ``llvm-readobj --elf-output-style=JSON`` document."""
    alloc = {}
    for s in doc["Sections"]:
        s = s["Section"]
        alloc[s["Index"]] = any(f["Name"] == "SHF_ALLOC" for f in s["Flags"]["Flags"])
    symbols = [s["Symbol"] for s in doc["Symbols"]]
    # A symbol names its section by index; an undefined one names index 0.
    home = [s["Section"]["Value"] or None for s in symbols]
    functions = {}
    for s, sec in zip(symbols, home):
        if s["Type"]["Name"] == "Function" and sec is not None:
            functions.setdefault(sec, set()).add(s["Name"]["Name"])
    edges: dict[int, set[int]] = {}
    calls: dict[int, set[str]] = {}
    for rel in doc.get("Relocations", []):
        # A relocation section applies to the section its sh_info names.
        src = next(
            s["Section"]["Info"]
            for s in doc["Sections"]
            if s["Section"]["Index"] == rel["SectionIndex"]
        )
        for r in rel["Relocs"]:
            i = r["Relocation"]["Symbol"]["Value"]
            if home[i] is not None:
                edges.setdefault(src, set()).add(home[i])
            elif symbols[i]["Section"]["Name"] == "Undefined":
                calls.setdefault(src, set()).add(symbols[i]["Name"]["Name"])
    frames = {
        name: e["Entry"]["Size"]
        for e in doc.get("StackSizes", [])
        for name in e["Entry"]["Functions"]
    }

    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]

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

Host-side helpers

aie.utils.bfp is the host side of the block-floating-point kernels: the bfp16ebs8 codec and the tile shuffle the mm_bfp DMA layout needs. It is the Python counterpart of programming_examples/ml/block_datatypes/helper.h, which the examples' C++ hosts use; a Python host encodes and checks with this module.

bfp16ebs8 on the host: encode, decode and the mmul-block shuffle.

v8bfp16ebs8 is the AIE2P block-floating-point type: eight values share one 8-bit exponent and each carries an 8-bit two's-complement mantissa, so a block of 8 values is 9 bytes ([exponent, m0, ..., m7]). These are numpy ports of floatToBfp16, bfp16ebs8ToFloat and shuffleMatrixForBfp16ebs8 in programming_examples/ml/block_datatypes/helper.h, bit for bit (the host test compiles that header and compares), including the header's rounding: mantissas truncate toward negative infinity, and a value more than 31 binades below its block's maximum becomes 0 (positive) or -1 LSB (negative).

NumPy structured arrays describe the packed storage. Neither NumPy nor ml_dtypes provides shared-exponent arithmetic for this format; the codec below supplies that conversion, not a replacement scalar dtype.

The block-floating-point matmul kernels (aie.iron.kernels.mm_bfp) load 8x8 sub-tiles as one 72-byte block vector, which a DMA cannot gather at 9-byte granularity, so tiles are rearranged by shuffle on the host: within each (tile_height, tile_width) tile the 8-row by 8-block sub-tiles are made contiguous in raster order. quantize is what a kernel sees of a float input, and what a reference should multiply.

encode

encode(x, *, rounding: str = 'floor') -> ndarray

float32 (..., n) with n % 8 == 0 -> uint8 (..., n * 9 // 8).

Blocks are taken along the last axis. Inputs must be finite (the C++ silently drops inf and NaN, which shifts every later value).

floor is the header's truncation. conv_even rounds each mantissa to nearest, ties to even, as amd/IRON's f32_to_bfp16ebs8 packs weights, byte for byte. A mantissa that rounds to +128 saturates to 127 there, so decode(encode(x, rounding="conv_even")) differs from quantize, which models the core raising the exponent instead.

Source code in python/utils/bfp.py

def encode(x, *, rounding: str = "floor") -> np.ndarray:
    """float32 ``(..., n)`` with ``n % 8 == 0`` -> ``uint8`` ``(..., n * 9 // 8)``.

    Blocks are taken along the last axis. Inputs must be finite (the C++
    silently drops inf and NaN, which shifts every later value).

    ``floor`` is the header's truncation. ``conv_even`` rounds each mantissa
    to nearest, ties to even, as amd/IRON's ``f32_to_bfp16ebs8`` packs
    weights, byte for byte. A mantissa that rounds to +128 saturates to 127
    there, so ``decode(encode(x, rounding="conv_even"))`` differs from
    ``quantize``, which models the core raising the exponent instead.
    """
    if rounding not in ("floor", "conv_even"):
        raise ValueError(
            f"bfp.encode: rounding must be 'floor' or 'conv_even', got {rounding!r}"
        )
    x = np.ascontiguousarray(x, dtype=np.float32)
    n = x.shape[-1]
    if n % BLOCK:
        raise ValueError(f"bfp.encode: last axis {n} is not a multiple of {BLOCK}")
    if not np.isfinite(x).all():
        raise ValueError("bfp.encode: inputs must be finite")
    lead = x.shape[:-1]
    bits = x.view(np.uint32).reshape(*lead, n // BLOCK, BLOCK)
    sign = (bits >> 31).astype(bool)
    exp = (bits >> 23) & 0xFF
    mant = (bits & 0x7FFFFF) | np.where(exp != 0, np.uint32(0x800000), np.uint32(0))
    max_exp = exp.max(axis=-1, keepdims=True)
    shift = (max_exp - exp).astype(np.int64)
    far = shift >= 32
    if rounding == "conv_even":
        # The quotient of a 24-bit magnitude by a power of two is exact in
        # float64, so rint is the only rounding.
        signed = np.where(sign, -mant.astype(np.int64), mant.astype(np.int64))
        v = np.rint(signed / np.exp2(np.minimum(_MANTISSA_SHIFT + shift, 62)))
        v = np.clip(v, -128, 127)
    else:
        # Two's complement in 32 bits, logical shift, low byte: the header's
        # `(uint8_t)((sign ? ~m + 1 : m) >> 17)`.
        m32 = np.where(sign, (-mant.astype(np.int64)) & 0xFFFFFFFF, mant).astype(
            np.uint32
        )
        v = (
            ((m32 >> _MANTISSA_SHIFT) & 0xFF)
            .astype(np.uint8)
            .view(np.int8)
            .astype(np.int32)
        )
        v = np.right_shift(v, np.minimum(shift, 31).astype(np.int32))  # arithmetic
    out = np.empty(lead + (n // BLOCK,), dtype=_BLOCK_DTYPE)
    out["exponent"] = max_exp[..., 0]
    out["mantissas"] = np.where(far, np.where(sign, -1, 0), v)
    return out.view(np.uint8).reshape(*lead, n // BLOCK * BLOCK_BYTES)

decode

decode(b) -> ndarray

uint8 (..., n * 9 // 8) -> float32 (..., n); the exact inverse map of a block.

Source code in python/utils/bfp.py

def decode(b) -> np.ndarray:
    """``uint8`` ``(..., n * 9 // 8)`` -> float32 ``(..., n)``; the exact inverse map of a block."""
    b = np.ascontiguousarray(b, dtype=np.uint8)
    nb = b.shape[-1]
    if nb % BLOCK_BYTES:
        raise ValueError(
            f"bfp.decode: last axis {nb} is not a multiple of {BLOCK_BYTES}"
        )
    lead = b.shape[:-1]
    blk = b.view(_BLOCK_DTYPE)
    scale = np.ldexp(1.0, blk["exponent"].astype(np.int32) - 127 - 6)
    vals = blk["mantissas"].astype(np.float64) * scale[..., None]
    return vals.astype(np.float32).reshape(*lead, nb // BLOCK_BYTES * BLOCK)

quantize

quantize(x, *, rounding: str = 'floor') -> ndarray

Return what a kernel reads of x, for a given conversion rounding mode.

Which mode applies is a property of who converts, not of where. floor is decode(encode(x)): this module's encoder truncates toward negative infinity, and a core converting in floor mode agrees with it -- the q4nx_dequant kernel pins floor and its reference matches the device byte for byte. conv_even models a core converting with round-to-nearest-ties-to-even, which is what mm_bfp's mixed kernel pins so its K reduction does not accumulate a one-sided bias.

The mode is load-bearing, not a detail: on 64x64x64 mixed tiles of random and large inputs, the right one reproduces the kernel's bf16 output bit for bit, and floor mismatches about 3550 of each tile's 4096 outputs.

A mantissa that rounds up to +128 does not fit the 8-bit field, so the block's exponent goes up by one and the block is requantized. Clamping instead would cost a whole step to the one element the shared exponent was chosen for. The carry is one-sided, as on the core: -128 fits, so a block whose most negative value rounds to -128 keeps its exponent.

Source code in python/utils/bfp.py

def quantize(x, *, rounding: str = "floor") -> np.ndarray:
    """Return what a kernel reads of ``x``, for a given conversion rounding mode.

    Which mode applies is a property of who converts, not of where. ``floor``
    is ``decode(encode(x))``: this module's encoder truncates toward negative
    infinity, and a *core* converting in floor mode agrees with it -- the
    ``q4nx_dequant`` kernel pins floor and its reference matches the device
    byte for byte. ``conv_even`` models a core converting with
    round-to-nearest-ties-to-even, which is what ``mm_bfp``'s mixed kernel
    pins so its K reduction does not accumulate a one-sided bias.

    The mode is load-bearing, not a detail: on 64x64x64 mixed tiles of random
    and large inputs, the right one reproduces the kernel's bf16 output bit
    for bit, and floor mismatches about 3550 of each tile's 4096 outputs.

    A mantissa that rounds up to +128 does not fit the 8-bit field, so the
    block's exponent goes up by one and the block is requantized. Clamping
    instead would cost a whole step to the one element the shared exponent
    was chosen for. The carry is one-sided, as on the core: -128 fits, so a
    block whose most negative value rounds to -128 keeps its exponent.
    """
    if rounding == "floor":
        return decode(encode(x))
    if rounding != "conv_even":
        raise ValueError(
            f"bfp.quantize: rounding must be 'floor' or 'conv_even', got {rounding!r}"
        )
    x = np.ascontiguousarray(x, dtype=np.float32)
    n = x.shape[-1]
    if n % BLOCK:
        raise ValueError(f"bfp.quantize: last axis {n} is not a multiple of {BLOCK}")
    blocks = x.reshape(*x.shape[:-1], n // BLOCK, BLOCK)
    exp = (blocks.view(np.uint32) >> 23) & 0xFF
    scale = np.ldexp(1.0, exp.max(axis=-1, keepdims=True).astype(np.int32) - 127 - 6)
    mant = np.rint(blocks.astype(np.float64) / scale)
    carry = (mant.max(axis=-1, keepdims=True) > 127)[..., 0]
    if carry.any():
        scale = np.where(carry[..., None], scale * 2, scale)
        mant = np.rint(blocks.astype(np.float64) / scale)
    return (np.clip(mant, -128, 127) * scale).astype(np.float32).reshape(x.shape)

shuffle

shuffle(
    b,
    width: int,
    height: int,
    tile_width: int,
    tile_height: int,
    *,
    unshuffle=False
) -> ndarray

Reorder an encoded (height, width) matrix into (or out of) the mmul tile layout.

b is the encoded matrix (height rows of width * 9 // 8 bytes, flat or 2-D; width and tile_width count values, so both are multiples of 8 and the tile is tile_height rows, a multiple of 8). Within each tile, every 8-row by 8-value sub-tile (72 bytes) becomes contiguous, sub-tiles in raster order; the tiles themselves stay where they are, so a DMA that copies a tile row by row delivers the sub-tiles the kernel's block-vector loads expect. unshuffle is the inverse.

Source code in python/utils/bfp.py

def shuffle(
    b, width: int, height: int, tile_width: int, tile_height: int, *, unshuffle=False
) -> np.ndarray:
    """Reorder an encoded ``(height, width)`` matrix into (or out of) the mmul tile layout.

    ``b`` is the encoded matrix (``height`` rows of ``width * 9 // 8`` bytes,
    flat or 2-D; ``width`` and ``tile_width`` count values, so both are
    multiples of 8 and the tile is ``tile_height`` rows, a multiple of 8).
    Within each tile, every 8-row by 8-value sub-tile (72 bytes) becomes
    contiguous, sub-tiles in raster order; the tiles themselves stay where
    they are, so a DMA that copies a tile row by row delivers the sub-tiles
    the kernel's block-vector loads expect. ``unshuffle`` is the inverse.
    """
    for name, v, mult in (
        ("width", width, BLOCK),
        ("tile_width", tile_width, BLOCK),
        ("tile_height", tile_height, BLOCK),
    ):
        if v <= 0 or v % mult:
            raise ValueError(
                f"bfp.shuffle: {name} must be a positive multiple of {mult}"
            )
    if width % tile_width or height % tile_height:
        raise ValueError("bfp.shuffle: the tile must divide the matrix")
    W, tw = width * BLOCK_BYTES // BLOCK, tile_width * BLOCK_BYTES // BLOCK
    a = np.asarray(b, dtype=np.uint8).reshape(height, W)
    ty, tx = height // tile_height, W // tw
    sy, sx = tile_height // BLOCK, tw // BLOCK_BYTES
    if not unshuffle:
        # [tileY, subY, i, tileX, subX, j] -> [tileY, tileX, subY, subX, i, j]
        v = a.reshape(ty, sy, BLOCK, tx, sx, BLOCK_BYTES).transpose(0, 3, 1, 4, 2, 5)
        # ... flattened raster-wise into the tile's (tile_height, tw) bytes
        v = v.reshape(ty, tx, tile_height, tw).transpose(0, 2, 1, 3)
    else:
        v = a.reshape(ty, tile_height, tx, tw).transpose(0, 2, 1, 3)
        v = v.reshape(ty, tx, sy, sx, BLOCK, BLOCK_BYTES).transpose(0, 2, 4, 1, 3, 5)
    return np.ascontiguousarray(v.reshape(height, W))

Tolerance-based output verification helpers for examples and tests.

Mirrors the canonical test_utils::nearly_equal semantics used across the C++ testbenches so Python migrations of those examples behave identically:

|a - b|  <  max(atol, rtol * (|a| + |b|))

Defaults match the C++ default of rtol=0.128, which is the widely-used relative tolerance for bfloat16 / LUT-approximated kernels (exp, softmax, gelu, silu, swiglu, ...).

Tolerance dataclass

Tolerance(
    rtol: float | None = None,
    atol: float | None = None,
    ulps: int | None = None,
    range_frac: float | None = None,
    max_mismatch_frac: float = 0.0,
    note: str = "",
    bound: Callable | None = None,
)

How close a kernel's output must be to its reference.

Exactly one of four kinds, chosen by which fields are set:

• exact -- no field set: bit-equal after casting the reference to the output dtype. Integers, selections (relu, max), lossless copies.
• ulps -- ulps set: bf16 outputs within ulps units in the last place of the correctly rounded reference. atol may be set alongside as a floor, admitting an element that meets either -- what a kernel needs when the device flushes subnormals to zero, since a flushed value is a full 100% relative and dozens of ulps from the reference but absolutely negligible. rtol stays unset, or the kind is relative.
• relative -- rtol/atol set: the canonical |a - b| < max(atol, rtol * (|a| + |b|)) of nearly_equal. Integer outputs are compared with the same formula in exact integer arithmetic, so lsb (atol = n + 0.5) admits an n-LSB slack for fixed-point pixel kernels whose rounding shift is not modeled; under exact and ulps integers stay bit-equal.
• bound -- bound set: |a - b| <= bound(*args) element by element, where args are the reference's own arguments. For a floating-point kernel whose error is set by where its input falls rather than by what it produced -- an approximation exact in one range and a few ulps off in another -- so any one rtol/atol is either loose everywhere or unsound somewhere. compare takes the evaluated bound; ExternalFunction.judge evaluates it from the inputs it is given.

Non-finite values are never skipped: NaN must meet NaN, and an infinity must meet an infinity of the same sign, under every kind.

range_frac adds a floor scaled to the reference's own range: an element also passes at |a - b| <= range_frac * max|b|. It is for a kernel whose error is set by the magnitudes it worked from rather than by the magnitude it produced -- a dot product whose terms cancel to near zero is no less accurate than its neighbours, but an elementwise relative bound reads it as 100% wrong. The scale comes from expected, never from actual, so a kernel cannot widen its own tolerance by returning something large. Unlike atol it follows the data: the same fraction holds whether the outputs run to 34 or to 3.4e9, where a fixed floor would be either dead or permissive. Set it from a measured worst case, and say so in note.

max_mismatch_frac allows that fraction of elements to miss (LUT tails, saturation edges). note records where the number came from -- a docstring, a device run, a testbench default -- so a reviewer can tell an evidenced tolerance from a guessed one.

bf16_ulps classmethod

bf16_ulps(
    n: int = 1,
    *,
    atol: float | None = None,
    max_mismatch_frac: float = 0.0,
    note: str = ""
) -> "Tolerance"

bf16 outputs within n ulps of the correctly rounded reference.

atol is an optional strict bound (error < atol) an element may meet instead of the ulp bound. Use it for the device's subnormal flush to zero, set to the smallest normal bf16 so it admits the flushed values and nothing above them.

Source code in python/utils/verify.py

@classmethod
def bf16_ulps(
    cls,
    n: int = 1,
    *,
    atol: float | None = None,
    max_mismatch_frac: float = 0.0,
    note: str = "",
) -> "Tolerance":
    """bf16 outputs within ``n`` ulps of the correctly rounded reference.

    ``atol`` is an optional strict bound (``error < atol``) an element may
    meet instead of the ulp bound. Use it for the device's subnormal flush
    to zero, set to the smallest normal bf16 so it admits the flushed
    values and nothing above them.
    """
    return cls(ulps=n, atol=atol, max_mismatch_frac=max_mismatch_frac, note=note)

bounded classmethod

bounded(
    bound: Callable,
    *,
    max_mismatch_frac: float = 0.0,
    note: str = ""
) -> "Tolerance"

Each element within bound(*args) of the reference, absolutely.

bound takes the reference's arguments and returns one non-negative bound per output element, shaped like the reference's result (a tuple of them for several outputs). Derive it from the kernel's arithmetic and say in note which parts were measured.

Source code in python/utils/verify.py

@classmethod
def bounded(
    cls, bound: Callable, *, max_mismatch_frac: float = 0.0, note: str = ""
) -> "Tolerance":
    """Each element within ``bound(*args)`` of the reference, absolutely.

    ``bound`` takes the reference's arguments and returns one
    non-negative bound per output element, shaped like the reference's
    result (a tuple of them for several outputs). Derive it from the
    kernel's arithmetic and say in ``note`` which parts were measured.
    """
    return cls(bound=bound, max_mismatch_frac=max_mismatch_frac, note=note)

lsb classmethod

lsb(
    n: int = 1,
    *,
    max_mismatch_frac: float = 0.0,
    note: str = ""
) -> "Tolerance"

Integer outputs within n least-significant bits of the reference.

For fixed-point kernels whose final saturating shift may round or truncate (the AIE srs rounding mode is a core setting the kernel does not fix). rtol is zero: the slack is absolute.

Source code in python/utils/verify.py

@classmethod
def lsb(
    cls, n: int = 1, *, max_mismatch_frac: float = 0.0, note: str = ""
) -> "Tolerance":
    """Integer outputs within ``n`` least-significant bits of the reference.

    For fixed-point kernels whose final saturating shift may round or
    truncate (the AIE ``srs`` rounding mode is a core setting the kernel
    does not fix). ``rtol`` is zero: the slack is absolute.
    """
    return cls(
        rtol=0.0, atol=n + 0.5, max_mismatch_frac=max_mismatch_frac, note=note
    )

default_for classmethod

default_for(dtype) -> 'Tolerance'

Return the contract a kernel gets when it declares none.

Integer and boolean outputs are bit-exact. bfloat16 outputs get the repository's canonical rtol=0.128, the C++ testbench default that the LUT-approximated kernels document; float32 outputs are held to rtol=1e-4 (a few float32 ULPs of accumulation-order slack, far inside what a bf16 tolerance would hide) and float16 to 1e-2.

Source code in python/utils/verify.py

@classmethod
def default_for(cls, dtype) -> "Tolerance":
    """Return the contract a kernel gets when it declares none.

    Integer and boolean outputs are bit-exact. bfloat16 outputs get the
    repository's canonical ``rtol=0.128``, the C++ testbench default that
    the LUT-approximated kernels document; float32 outputs are held to
    ``rtol=1e-4`` (a few float32 ULPs of accumulation-order slack, far
    inside what a bf16 tolerance would hide) and float16 to ``1e-2``.
    """
    dt = np.dtype(dtype)
    if np.issubdtype(dt, np.integer) or dt == np.bool_:
        return cls.exact(note="default: integer output")
    if dt == np.dtype(np.float32) or dt == np.dtype(np.float64):
        return cls.relative(1e-4, note="default: float32 output")
    if dt == np.dtype(np.float16):
        return cls.relative(1e-2, note="default: float16 output")
    return cls.relative(_DEFAULT_RTOL, note="default: canonical bf16/LUT rtol")

Verdict dataclass

Verdict(
    ok: bool,
    n_checked: int,
    n_mismatch: int,
    max_abs_err: float,
    max_ulp_err: int | None,
    first_bad_index: int | None,
    detail: str,
)

Outcome of compare. Truthy when the comparison passed.

compare

compare(
    actual,
    expected,
    tol: Tolerance | None = None,
    *,
    range_axis: int | None = None,
    bound=None
) -> Verdict

Compare a kernel's actual output with a reference under tol.

expected may be higher precision than actual (a float64 sum, an int64 product); it is cast to actual.dtype, so the kernel is held to what a correctly rounded implementation would produce. With tol=None the output dtype's Tolerance.default_for applies.

range_axis selects the axis reduced to compute range_frac's reference scale. For (calls, tile) arrays, use 1 to scale each call independently. The default uses the whole reference array.

bound is a bound tolerance's per-element limit, already evaluated on the inputs and broadcastable to expected.

This measures; it does not model. What a kernel does on overflow, on a narrowing store, or with subnormal inputs belongs in the reference that produced expected -- a saturating kernel's reference clips, a denormal-flushing kernel's reference flushes. A reference that leaves the output range is reported as such when the comparison fails, since it means the reference is under-specified rather than the kernel wrong.

Source code in python/utils/verify.py

def compare(
    actual,
    expected,
    tol: Tolerance | None = None,
    *,
    range_axis: int | None = None,
    bound=None,
) -> Verdict:
    """Compare a kernel's ``actual`` output with a reference under ``tol``.

    ``expected`` may be higher precision than ``actual`` (a float64 sum, an
    int64 product); it is cast to ``actual.dtype``, so the kernel is held to
    what a correctly rounded implementation would produce. With ``tol=None``
    the output dtype's ``Tolerance.default_for`` applies.

    ``range_axis`` selects the axis reduced to compute ``range_frac``'s
    reference scale. For ``(calls, tile)`` arrays, use 1 to scale each call
    independently. The default uses the whole reference array.

    ``bound`` is a **bound** tolerance's per-element limit, already
    evaluated on the inputs and broadcastable to ``expected``.

    This measures; it does not model. What a kernel does on overflow, on a
    narrowing store, or with subnormal inputs belongs in the reference that
    produced ``expected`` -- a saturating kernel's reference clips, a
    denormal-flushing kernel's reference flushes. A reference that leaves the
    output range is reported as such when the comparison fails, since it
    means the reference is under-specified rather than the kernel wrong.
    """
    actual = np.asarray(actual)
    expected = np.asarray(expected)
    if tol is None:
        tol = Tolerance.default_for(actual.dtype)
    if actual.shape != expected.shape:
        return Verdict(
            False,
            0,
            0,
            float("inf"),
            None,
            None,
            f"shape mismatch {actual.shape} vs {expected.shape}",
        )
    if tol.kind == "bound":
        if bound is None:
            raise ValueError(
                "a bound tolerance needs its bound evaluated on the inputs; "
                "pass bound=, or judge through the kernel with inputs="
            )
        if not np.issubdtype(actual.dtype, np.floating) and actual.dtype != bfloat16:
            raise ValueError(
                f"a bound tolerance is for floating-point outputs, got {actual.dtype}"
            )
    a, e, n = actual.ravel(), expected.ravel(), actual.size

    def ref_scale(ref) -> float | np.ndarray:
        """Return max|ref| over finite entries.

        This is what ``range_frac`` is a fraction of, or zero when no range
        floor is in play.
        """
        if tol.range_frac is None or not n:
            return 0.0
        mag = np.abs(np.asarray(ref, dtype=np.float64).reshape(expected.shape))
        mag = np.where(np.isfinite(mag), mag, 0.0)
        scale = mag.max(axis=range_axis, keepdims=True, initial=0.0)
        if range_axis is None:
            return float(scale.item())
        return np.broadcast_to(scale, expected.shape).ravel()

    # Integers: bit-exact under exact / ulps; under a relative tolerance the
    # nearly_equal formula in int64 (Tolerance.lsb sets atol = n + 0.5).
    if np.issubdtype(actual.dtype, np.integer) or actual.dtype == np.bool_:
        n_over = 0
        if actual.dtype != np.bool_ and np.issubdtype(e.dtype, np.integer):
            info = np.iinfo(actual.dtype)
            e_wide = e.astype(np.int64)
            n_over = int(np.count_nonzero((e_wide < info.min) | (e_wide > info.max)))
        e_cast = e.astype(actual.dtype)
        a64, e64 = a.astype(np.int64), e_cast.astype(np.int64)
        err = np.abs(a64 - e64)
        scale = ref_scale(e64)
        if tol.kind == "relative":
            bound = np.maximum(
                tol.atol or 0.0,
                (tol.rtol or 0.0) * (np.abs(a64) + np.abs(e64)),
            )
            close = err < bound
            if tol.range_frac is not None:
                close |= err <= tol.range_frac * scale
            bad = ~close
        else:
            bad = a != e_cast
        v = _verdict(bad, err, None, tol, n, ref_range=scale)
        if n_over and not v.ok:
            # Not a policy, a diagnostic: the reference left the output range,
            # so what the device did there says nothing about the kernel.
            v = replace(
                v,
                detail=f"{v.detail}; the reference overflows "
                f"{np.dtype(actual.dtype).name} in {n_over} of {n} elements, so "
                "it does not model what the kernel does there (clip for a "
                "saturating kernel, cast for a wrapping one)",
            )
        return v

    a32, e32 = a.astype(np.float32), e.astype(np.float32)
    a_nan, e_nan = np.isnan(a32), np.isnan(e32)
    a_inf, e_inf = np.isinf(a32), np.isinf(e32)
    nonfinite_bad = (a_nan != e_nan) | (a_inf != e_inf) | (a_inf & e_inf & (a32 != e32))
    finite = ~(a_nan | e_nan | a_inf | e_inf)
    err = np.zeros(n, np.float64)

    if tol.kind == "ulps":
        if actual.dtype != bfloat16:
            raise ValueError(
                f"Tolerance in ULPs is defined for bfloat16 outputs, got {actual.dtype}"
            )
        # A cast to bfloat16 goes through float32 and can round twice;
        # round_to rounds the reference once, as a correctly rounded bf16
        # implementation would.
        e_bf = round_to(e, bfloat16)
        ulp = np.zeros(n, np.int64)
        ulp[finite] = bf16_ulp_distance(a[finite], e_bf[finite])
        err[finite] = np.abs(a32[finite] - e_bf[finite].astype(np.float32))
        max_ulps = tol.ulps if tol.ulps is not None else 0
        within = ulp <= max_ulps
        scale = ref_scale(e_bf.astype(np.float32))
        if tol.atol is not None:
            within |= err < tol.atol
        if tol.range_frac is not None:
            within |= err <= tol.range_frac * scale
        bad = nonfinite_bad | (finite & ~within)
        return _verdict(bad, err, ulp, tol, n, nonfinite_bad, ref_range=scale)

    if tol.kind == "bound":
        limit = np.broadcast_to(np.asarray(bound, np.float64), expected.shape).ravel()
        err[finite] = np.abs(a32[finite].astype(np.float64) - e32[finite])
        bad = nonfinite_bad | (finite & ~(err <= limit))
        return _verdict(bad, err, None, tol, n, nonfinite_bad, limit=limit)

    if tol.kind == "exact":
        e_cast = e32.astype(actual.dtype).astype(np.float32)
        err[finite] = np.abs(a32[finite] - e_cast[finite])
        bad = nonfinite_bad | (finite & (a32 != e_cast))
        return _verdict(bad, err, None, tol, n, nonfinite_bad)

    err[finite] = np.abs(a32[finite].astype(np.float64) - e32[finite])
    scale = ref_scale(e32)
    close = nearly_equal(a32, e32, rtol=tol.rtol or 0.0, atol=tol.atol)
    if tol.range_frac is not None:
        close |= err <= tol.range_frac * scale
    bad = nonfinite_bad | (finite & ~close)
    return _verdict(bad, err, None, tol, n, nonfinite_bad, ref_range=scale)

bf16_ulp_distance

bf16_ulp_distance(a, b) -> ndarray

Element-wise distance between two bf16 arrays in units in the last place.

Bit patterns are mapped to a monotonic integer scale (sign-magnitude to two's-complement style) so the distance is a plain subtraction. -0 and +0 map to the same point, so a kernel that produces the other zero is not penalized. aie.utils.accuracy.ulp_distance takes other dtypes.

Source code in python/utils/verify.py

def bf16_ulp_distance(a, b) -> np.ndarray:
    """Element-wise distance between two bf16 arrays in units in the last place.

    Bit patterns are mapped to a monotonic integer scale (sign-magnitude to
    two's-complement style) so the distance is a plain subtraction. -0 and +0
    map to the same point, so a kernel that produces the other zero is not
    penalized. ``aie.utils.accuracy.ulp_distance`` takes other dtypes.
    """
    return ulp_distance(a, b, bfloat16)

Sampling

sample_select and sample_combine draw the next token from a row of bf16 logits on the device: each sample_select core reduces one slice to its top k, and one sample_combine core draws from those summaries. sample_ref is the bit-exact host reference.

Top-k next-token sampling on the device, and its bit-exact host reference.

A row of bf16 logits is split into columns of slice_size. Each column's sample_select core reduces its slice to a summary (its top k, as entries above its k-th largest value and a bitmap of the ties at it); one sample_combine core draws the token from all the summaries. The draw is sample_ref:

1. Temperature 0 (or -0): the first index of the largest logit.
2. tau is the k-th largest logit, with multiplicity, compared on order_keys (-0 and +0 are one key). The candidates are the logits >= tau, in index order.
3. Each candidate value v weighs exp64_ref(float64(fl32(fl32(v / T) - fl32(max / T)))).
4. u = n53 * 2**-53, integer 0 <= n53 < 2**53. The token is the first candidate, in index order, whose exact prefix sum of the weights P exceeds u * S, S the exact total: P * 2**53 > n53 * S in integers. That is the inverse-CDF draw over the weights, with no rounding after them.

Every weight step is an IEEE float32 division or subtraction, or a float64 addition, subtraction, multiplication or comparison, never fused; the core computes the float32 steps with the soft-float builtins and float64 has only those, so numpy reproduces the device's bits. Each weight is an integer multiple of 2**-1074, so the core sums them exactly in fixed point. A temperature must satisfy check_order_preserving, which makes the threshold on bf16 keys select what a threshold on v / T would.

A draw is passed to the device as a four-word int32 row, draw_row.

summary_words

summary_words(slice_size: int, k_max: int) -> int

int32 words of one column's summary: header, entries, tie bitmap.

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

def summary_words(slice_size: int, k_max: int) -> int:
    """int32 words of one column's summary: header, entries, tie bitmap."""
    return SUMMARY_HEADER + 2 * k_max + (slice_size + 31) // 32

select_streams

select_streams(slice_size: int, chunk: int) -> int

How often sample_select takes its slice per position.

Once when a chunk is the whole slice, which one call passes over SELECT_PASSES times; SELECT_PASSES times otherwise.

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

def select_streams(slice_size: int, chunk: int) -> int:
    """How often ``sample_select`` takes its slice per position.

    Once when a chunk is the whole slice, which one call passes over
    ``SELECT_PASSES`` times; ``SELECT_PASSES`` times otherwise.
    """
    return 1 if chunk == slice_size else SELECT_PASSES

sample_select

sample_select(
    *, slice_size=32064, chunk=5344, k_max=64
) -> ExternalFunction

One column's half of sampling: its slice's summary, for sample_combine.

sample_select(x, row, state, summary) takes chunk bf16 logits per call; the slice arrives select_streams(slice_size, chunk) times, so a position is that many times slice_size // chunk calls with the same row and summary. state is SELECT_STATE_WORDS int32 of worker-local memory, zero before the first call; the last call of a position leaves it ready for the next. summary is summary_words(slice_size, k_max) int32.

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

def sample_select(*, slice_size=32064, chunk=5344, k_max=64) -> ExternalFunction:
    """One column's half of sampling: its slice's summary, for ``sample_combine``.

    ``sample_select(x, row, state, summary)`` takes ``chunk`` bf16 logits per
    call; the slice arrives ``select_streams(slice_size, chunk)`` times, so a
    position is that many times ``slice_size // chunk`` calls with the same
    ``row`` and ``summary``.
    ``state`` is ``SELECT_STATE_WORDS`` int32 of worker-local memory, zero
    before the first call; the last call of a position leaves it ready for the
    next. ``summary`` is ``summary_words(slice_size, k_max)`` int32.
    """
    _check_slice("sample_select", slice_size, k_max)
    _positive("sample_select", chunk=chunk)
    if slice_size % chunk:
        raise ValueError("sample_select: chunk must divide slice_size")
    if chunk % 2:
        raise ValueError("sample_select: chunk must be even (4-byte DMA)")
    return _make_extern(
        "sample_select",
        _kernel_source("sample/sample_select.cc"),
        [
            np.ndarray[(chunk,), np.dtype[bfloat16]],
            np.ndarray[(ROW_WORDS,), np.dtype[np.int32]],
            np.ndarray[(SELECT_STATE_WORDS,), np.dtype[np.int32]],
            np.ndarray[(summary_words(slice_size, k_max),), np.dtype[np.int32]],
        ],
        compile_flags=_STRICT_FP
        + [
            f"-DSAMPLE_SLICE={slice_size}",
            f"-DSAMPLE_CHUNK={chunk}",
            f"-DSAMPLE_K_MAX={k_max}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, InOut, InOut),
        ),
    )

sample_combine

sample_combine(
    *, columns=4, slice_size=32064, k_max=64
) -> ExternalFunction

Draw from columns summaries of sample_select, as sample_ref does.

sample_combine(summaries, row, token, record) reads the summaries column 0 first, and writes the drawn token (an index into the whole row of columns * slice_size) to both one-word outputs, so a design can send it two places.

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

def sample_combine(*, columns=4, slice_size=32064, k_max=64) -> ExternalFunction:
    """Draw from ``columns`` summaries of ``sample_select``, as ``sample_ref`` does.

    ``sample_combine(summaries, row, token, record)`` reads the summaries
    column 0 first, and writes the drawn token (an index into the whole row
    of ``columns * slice_size``) to both one-word outputs, so a design can
    send it two places.
    """
    _check_slice("sample_combine", slice_size, k_max)
    _positive("sample_combine", columns=columns)
    if columns * slice_size >= 1 << 31:
        raise ValueError("sample_combine: a token must fit int32")
    return _make_extern(
        "sample_combine",
        _kernel_source("sample/sample_combine.cc"),
        [
            np.ndarray[
                (columns * summary_words(slice_size, k_max),), np.dtype[np.int32]
            ],
            np.ndarray[(ROW_WORDS,), np.dtype[np.int32]],
            np.ndarray[(1,), np.dtype[np.int32]],
            np.ndarray[(1,), np.dtype[np.int32]],
        ],
        compile_flags=_STRICT_FP
        + [
            f"-DSAMPLE_SLICE={slice_size}",
            f"-DSAMPLE_COLUMNS={columns}",
            f"-DSAMPLE_K_MAX={k_max}",
        ],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out, Out),
        ),
    )

exp64_ref

exp64_ref(x)

Compute exp over float64, bit for bit as aie_kernels/sample/exp64.h.

numpy's elementwise float64 +, -, * and comparisons are correctly rounded, so evaluating exp64.h's sequence of them gives its bits; np.exp is never used. Every branch is evaluated for every element and the results selected with np.where, which is exact.

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

def exp64_ref(x):
    """Compute exp over float64, bit for bit as ``aie_kernels/sample/exp64.h``.

    numpy's elementwise float64 +, -, * and comparisons are correctly
    rounded, so evaluating exp64.h's sequence of them gives its bits; np.exp
    is never used. Every branch is evaluated for every element and the
    results selected with np.where, which is exact.
    """
    c, table = _exp64_constants()
    n = np.uint64(table.size // 2)
    table_bits = int(n).bit_length() - 1
    u64, f64 = np.uint64, np.float64
    x = np.ascontiguousarray(x, dtype=np.float64)
    shape = x.shape
    x = x.reshape(-1)
    bits = x.view(np.uint64)
    abstop = (bits >> u64(52)).astype(np.int64) & 0x7FF
    tiny = abstop < 0x3C9  # |x| < 2^-54
    huge = abstop >= 0x409  # |x| >= 1024, inf, nan
    special = abstop == 0x408  # 512 <= |x| < 1024
    negative = (bits >> u64(63)) == u64(1)

    with np.errstate(all="ignore"):
        z = c["INV_LN2_N"] * x
        kd = z + c["SHIFT"]
        ki = kd.view(np.uint64).copy()
        kd = kd - c["SHIFT"]
        khi = kd * c["NEG_LN2_HI_N"]
        klo = kd * c["NEG_LN2_LO_N"]
        r = x + khi
        r = r + klo
        idx = (ki % n) * u64(2)
        top = ki << u64(52 - table_bits)
        tail = table[idx].view(np.float64)
        sbits = table[idx + u64(1)] + top
        r2 = r * r
        p23 = r * c["C3"]
        p23 = c["C2"] + p23
        p23 = r2 * p23
        p45 = r * c["C5"]
        p45 = c["C4"] + p45
        r4 = r2 * r2
        p45 = r4 * p45
        tmp = tail + r
        tmp = tmp + p23
        tmp = tmp + p45

        scale = sbits.view(np.float64)
        st = scale * tmp
        main = scale + st

        # exp64_special, k > 0.
        scale_pos = (sbits - (u64(1009) << u64(52))).view(np.float64)
        st_pos = scale_pos * tmp
        y_pos = scale_pos + st_pos
        y_pos = f64(2.0**1009) * y_pos
        # exp64_special, k < 0.
        scale_neg = (sbits + (u64(1022) << u64(52))).view(np.float64)
        st_neg = scale_neg * tmp
        y = scale_neg + st_neg
        lo = scale_neg - y
        lo = lo + st_neg
        hi = f64(1.0) + y
        lo2 = f64(1.0) - hi
        lo2 = lo2 + y
        lo2 = lo2 + lo
        y_sub = hi + lo2
        y_sub = y_sub - f64(1.0)
        y_sub = np.where(y_sub == f64(0.0), f64(0.0), y_sub)
        y_neg = np.where(y < f64(1.0), y_sub, y)
        y_neg = f64(2.0**-1022) * y_neg
        k_positive = (ki & u64(0x80000000)) == u64(0)
        special_value = np.where(k_positive, y_pos, y_neg)

        one_plus_x = f64(1.0) + x
        huge_value = np.where(
            bits == np.float64(-np.inf).view(np.uint64),
            f64(0.0),
            np.where(
                abstop >= 0x7FF,
                one_plus_x,
                np.where(negative, f64(0.0), f64(np.inf)),
            ),
        )

    out = np.where(special, special_value, main)
    out = np.where(huge, huge_value, out)
    out = np.where(tiny, one_plus_x, out)
    return out.reshape(shape)

order_keys

order_keys(logits) -> ndarray

uint16 keys whose order is the bf16 logits' numeric order; -0 and +0 are one key.

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

def order_keys(logits) -> np.ndarray:
    """uint16 keys whose order is the bf16 logits' numeric order; -0 and +0 are one key."""
    bits = np.ascontiguousarray(logits, dtype=bfloat16).view(np.uint16)
    bits = np.where(bits == 0x8000, np.uint16(0), bits)
    negative = (bits & np.uint16(0x8000)) != 0
    return np.where(negative, ~bits, bits | np.uint16(0x8000)).astype(np.uint16)

check_order_preserving

check_order_preserving(temperature) -> None

Raise unless l -> fl32(l / T) strictly increases over every finite bf16 l with a finite image.

Then a threshold on bf16 keys selects what a threshold on l / T would, up to logits whose l / T is -inf and whose weight is 0 anyway.

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

def check_order_preserving(temperature) -> None:
    """Raise unless ``l -> fl32(l / T)`` strictly increases over every finite bf16 l with a finite image.

    Then a threshold on bf16 keys selects what a threshold on ``l / T`` would,
    up to logits whose ``l / T`` is -inf and whose weight is 0 anyway.
    """
    temperature = np.float32(temperature)
    if not 0 < temperature < np.inf:
        raise ValueError(f"temperature {temperature} is not finite and positive")
    bits = np.arange(1 << 16, dtype=np.uint32).astype(np.uint16)
    values = bits.view(bfloat16).astype(np.float32)
    values = np.sort(values[np.isfinite(values) & (bits != 0x8000)])
    with np.errstate(over="ignore"):
        x = values / temperature
    x = x[np.isfinite(x)]
    if not np.all(x[1:] > x[:-1]):
        where = int(np.flatnonzero(x[1:] <= x[:-1])[0])
        raise ValueError(
            f"temperature {temperature} maps two bf16 logits to one x "
            f"({x[where]!r}); the bf16-key threshold would not be exact"
        )

draw_row

draw_row(temperature, top_k: int, n53: int) -> ndarray

Pack the four int32 words sample_select and sample_combine read for one draw.

The kernels clamp top_k to their k_max.

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

def draw_row(temperature, top_k: int, n53: int) -> np.ndarray:
    """Pack the four int32 words ``sample_select`` and ``sample_combine`` read for one draw.

    The kernels clamp ``top_k`` to their ``k_max``.
    """
    _positive("draw_row", top_k=top_k)
    if top_k >= 1 << 31:
        raise ValueError(f"draw_row: top_k {top_k} does not fit int32")
    n53 = _check_n53(n53)
    row = np.empty(ROW_WORDS, dtype=np.uint32)
    row[0] = np.float32(temperature).view(np.uint32)
    row[1] = top_k
    row[2] = n53 & 0xFFFFFFFF
    row[3] = n53 >> 32
    return row.view(np.int32)

sample_weights

sample_weights(
    logits, temperature, top_k: int
) -> tuple[ndarray, ndarray]

(candidates, weights) of a draw at a positive temperature.

The candidates are the indices of every logit at or above the top_k-th largest (ties with it included), in index order; each weight is exp64(fl32(v / T) - fl32(max / T)), unnormalised.

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

def sample_weights(logits, temperature, top_k: int) -> tuple[np.ndarray, np.ndarray]:
    """``(candidates, weights)`` of a draw at a positive temperature.

    The candidates are the indices of every logit at or above the
    ``top_k``-th largest (ties with it included), in index order; each weight
    is ``exp64(fl32(v / T) - fl32(max / T))``, unnormalised.
    """
    _positive("sample_weights", top_k=top_k)
    row = _row(logits)
    keys = order_keys(row)
    values = row.astype(np.float32)
    temperature = np.float32(temperature)
    n = keys.size
    k = min(top_k, n)
    tau = np.partition(keys, n - k)[n - k]
    candidates = np.flatnonzero(keys >= tau)
    with np.errstate(over="ignore"):
        xm = values[int(np.argmax(keys))] / temperature
        xv = values[candidates] / temperature
    if not np.isfinite(xm):
        raise ValueError("the largest logit / T is not finite in float32")
    return candidates, exp64_ref((xv - xm).astype(np.float64))

sample_ref

sample_ref(
    logits,
    temperature,
    top_k: int,
    n53: int,
    *,
    k_max: int | None = None
) -> int

Return the token sample_combine draws from one row of bf16 logits.

With k_max, top_k is clamped to it, as kernels built with that k_max clamp it; without, top_k is taken as it is.

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

def sample_ref(
    logits, temperature, top_k: int, n53: int, *, k_max: int | None = None
) -> int:
    """Return the token ``sample_combine`` draws from one row of bf16 logits.

    With ``k_max``, ``top_k`` is clamped to it, as kernels built with that
    ``k_max`` clamp it; without, ``top_k`` is taken as it is.
    """
    _positive("sample_ref", top_k=top_k)
    n53 = _check_n53(n53)
    if k_max is not None:
        _positive("sample_ref", k_max=k_max)
        top_k = min(top_k, k_max)
    if np.float32(temperature).view(np.uint32) & 0x7FFFFFFF == 0:
        return int(np.argmax(order_keys(_row(logits))))
    candidates, weights = sample_weights(logits, temperature, top_k)
    prefix = list(accumulate(_units(weights)))
    # P * 2**53 > n53 * S is P > floor(n53 * S / 2**53) for an integer P. The
    # prefixes never step down, so the first above is a bisection; S >= 1
    # (the maximum weighs 1) and u < 1, so the last always is.
    target = n53 * prefix[-1] >> 53
    return int(candidates[bisect_right(prefix, target)])

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.

A mixture-of-experts router's top-k selection on the device.

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 python/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]],
        ],
        compile_flags=_STRICT_FP
        + [f"-DSAMPLE_SLICE={slice_size}", f"-DSAMPLE_K_MAX={top_k}"],
        contract=KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out, Out),
        ),
    )

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 python/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)
    max_value = row[int(np.argmax(keys))].astype(np.float32)
    values = row[ids].astype(np.float32)
    shifted = (values - max_value).astype(np.float32)
    w = exp64_ref(shifted.astype(np.float64))
    packed = (w / w.sum()).astype(np.float32).astype(bfloat16)
    weights = np.zeros(2 * top_k, dtype=bfloat16)
    weights[0::2] = packed
    return ids, weights

FastFlowLM Gemma 4

Kernels extracted from FastFlowLM's Gemma 4 implementation.

These are not general building blocks. Each source is the kernel code of one core in one of FastFlowLM's Gemma 4 designs (prefill attention, the LM head and the decode layer), called from that core's Worker body. Several kernels acquire and release core locks themselves, by the lock ids the factories take. The decode kernels build for the whole model's geometry (FlmGemma4DecodeGeometry) and for FastFlowLM's array layout. All are AIE2P only; their sources are in aie_kernels/flm_gemma4/.

A factory returns one ExternalFunction. As cascade_mm does with its put_only and put_get, every entry point of its object, its own included, is an attribute of it named by its C symbol: fn.attn_qk_round.

FlmGemma4DecodeGeometry dataclass

FlmGemma4DecodeGeometry(
    model_dim: int,
    num_attn_heads: int,
    num_kv_heads: int,
    intermediate_size: int,
    glu_slice: int,
    pli_d: int,
    dh: int,
    swa_dh: int,
    attn_scale: float,
    swa_attn_scale: float,
    pli_projection_scale: float,
    pli_input_scale: float,
    gelu: bool = True,
    qk_norm: bool = True,
    double_wide_mlp: bool = False,
    name: str = "custom",
)

The model geometry the flm_gemma4/decode_*.cc kernels build for.

The kernels are Gemma 4's fused decode layer, one core per stage, as FastFlowLM builds it. Every field reaches the kernels as a -DFLM_GEMMA4_DECODE_* flag. FLM_GEMMA4_E2B_DECODE and FLM_GEMMA4_E4B_DECODE are the two variants' values, and the factories build only those: the kernels keep assumptions other values break, such as eight query heads, the GELU table, the q/k norm and projections that divide into whole rounds.

flags

flags() -> list[str]

Return the -DFLM_GEMMA4_DECODE_* flags, floats as shortest round-trip literals.

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

def flags(self) -> list[str]:
    """Return the ``-DFLM_GEMMA4_DECODE_*`` flags, floats as shortest round-trip literals."""
    values = {
        "MODEL_DIM": self.model_dim,
        "NUM_ATTN_HEADS": self.num_attn_heads,
        "NUM_KV_HEADS": self.num_kv_heads,
        "INTERMEDIATE_SIZE": self.intermediate_size,
        "GLU_SLICE": self.glu_slice,
        "PLI_D": self.pli_d,
        "DH": self.dh,
        "SWA_DH": self.swa_dh,
        "ATTN_SCALE": f"{float(self.attn_scale)!r}f",
        "SWA_ATTN_SCALE": f"{float(self.swa_attn_scale)!r}f",
        "PLI_PROJECTION_SCALE": f"{float(self.pli_projection_scale)!r}f",
        "PLI_INPUT_SCALE": f"{float(self.pli_input_scale)!r}f",
        "GELU": int(self.gelu),
        "QK_NORM": int(self.qk_norm),
        "DOUBLE_WIDE_MLP": int(self.double_wide_mlp),
    }
    return [f"-DFLM_GEMMA4_DECODE_{k}={v}" for k, v in values.items()]

flm_gemma4_attn_prefill

flm_gemma4_attn_prefill(
    *, in_prod_lock: int = 2, in_cons_lock: int = 3
) -> ExternalFunction

Causal flash-attention prefill at a head dim of 512, from flm_gemma4/prefill.cc.

linalg.prefill_fv's algorithm with FastFlowLM's numerics and synchronization: the softmax scales by log2(e) alone, the epilogue multiplies by a bf16 1 / l, and k and v arrive in a ping-pong pair that the core's own locks guard.

One core holds 8 query rows of a 128-row round and folds the keys in one 8-row step at a time, with one entry point per step of the caller's round, block and step loops. The returned kernel is attn_epilogue, which writes one 64-element chunk of o = y / l; its chunk index is bound to 0. The entry points, all attributes of the kernel:

• attn_rounds(L_begin, L_end, n_out): rounds in the dispatch
• attn_blocks(L_begin, i, n_out): key blocks in round i
• attn_round_begin(prev_m, new_m, c, l, y)
• attn_block_begin(m, prev_m)
• attn_qk_step(s, q, k_ping, k_pong, m, L_begin, window_size, row, col, i, block, j)
• attn_block_mid(s, m, new_m, prev_m, c, l, y)
• attn_fv_step(y, s, v_ping, v_pong, j)
• attn_block_end(prev_m, new_m)
• attn_finalize(l, l_bf16)
• attn_epilogue(o, l_bf16, y, chunk)

The design must declare the k/v locks and fill the pair with a tile DMA.

Parameters:

Name	Type	Description	Default
in_prod_lock	int	Core lock the k/v DMA acquires to fill a buffer.	2
in_cons_lock	int	Core lock the steps acquire to read one.	3

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

def flm_gemma4_attn_prefill(
    *, in_prod_lock: int = 2, in_cons_lock: int = 3
) -> ExternalFunction:
    """Causal flash-attention prefill at a head dim of 512, from ``flm_gemma4/prefill.cc``.

    [`linalg.prefill_fv`][iron.kernels.linalg.prefill_fv]'s algorithm with
    FastFlowLM's numerics and synchronization: the softmax scales by log2(e)
    alone, the epilogue multiplies by a bf16 ``1 / l``, and k and v arrive in
    a ping-pong pair that the core's own locks guard.

    One core holds 8 query rows of a 128-row round and folds the keys in one
    8-row step at a time, with one entry point per step of the caller's
    round, block and step loops. The returned kernel is ``attn_epilogue``,
    which writes one 64-element chunk of ``o = y / l``; its chunk index is
    bound to 0. The entry points, all attributes of the kernel:

    - ``attn_rounds(L_begin, L_end, n_out)``: rounds in the dispatch
    - ``attn_blocks(L_begin, i, n_out)``: key blocks in round ``i``
    - ``attn_round_begin(prev_m, new_m, c, l, y)``
    - ``attn_block_begin(m, prev_m)``
    - ``attn_qk_step(s, q, k_ping, k_pong, m, L_begin, window_size, row, col, i, block, j)``
    - ``attn_block_mid(s, m, new_m, prev_m, c, l, y)``
    - ``attn_fv_step(y, s, v_ping, v_pong, j)``
    - ``attn_block_end(prev_m, new_m)``
    - ``attn_finalize(l, l_bf16)``
    - ``attn_epilogue(o, l_bf16, y, chunk)``

    The design must declare the k/v locks and fill the pair with a tile DMA.

    Args:
        in_prod_lock: Core lock the k/v DMA acquires to fill a buffer.
        in_cons_lock: Core lock the steps acquire to read one.
    """
    return _prefill(
        "flm_gemma4_attn_prefill",
        _ATTN_PREFILL_DH,
        _ATTN_PREFILL_LQ,
        0,
        flm_gemma4_attn_prefill_ref,
        in_prod_lock,
        in_cons_lock,
    )

flm_gemma4_swa_prefill

flm_gemma4_swa_prefill(
    *, in_prod_lock: int = 2, in_cons_lock: int = 3
) -> ExternalFunction

Sliding-window causal flash-attention prefill at a head dim of 256, from flm_gemma4/prefill.cc.

The sliding-window build of flm_gemma4_attn_prefill: one core holds 16 query rows of a 128-row round and folds in, 16 rows at a time, the keys inside the window its window_size RTP gives. The entry points are those of flm_gemma4_attn_prefill, with attn_blocks(L_begin, window_size, i, n_out) also reading the window. The returned kernel is attn_epilogue, bound to chunk 33: rows 8 to 15, columns 64 to 127 of the (16, 256) y.

Parameters:

Name	Type	Description	Default
in_prod_lock	int	Core lock the k/v DMA acquires to fill a buffer.	2
in_cons_lock	int	Core lock the steps acquire to read one.	3

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

def flm_gemma4_swa_prefill(
    *, in_prod_lock: int = 2, in_cons_lock: int = 3
) -> ExternalFunction:
    """Sliding-window causal flash-attention prefill at a head dim of 256, from ``flm_gemma4/prefill.cc``.

    The sliding-window build of [`flm_gemma4_attn_prefill`][iron.kernels.flm_gemma4.flm_gemma4_attn_prefill]:
    one core holds 16 query rows of a 128-row round and folds in, 16 rows at
    a time, the keys inside the window its ``window_size`` RTP gives. The
    entry points are those of ``flm_gemma4_attn_prefill``, with
    ``attn_blocks(L_begin, window_size, i, n_out)`` also reading the window.
    The returned kernel is ``attn_epilogue``, bound to chunk 33: rows 8 to 15,
    columns 64 to 127 of the (16, 256) ``y``.

    Args:
        in_prod_lock: Core lock the k/v DMA acquires to fill a buffer.
        in_cons_lock: Core lock the steps acquire to read one.
    """
    return _prefill(
        "flm_gemma4_swa_prefill",
        _SWA_PREFILL_DH,
        _SWA_PREFILL_LQ,
        _SWA_PREFILL_CHUNK,
        flm_gemma4_swa_prefill_ref,
        in_prod_lock,
        in_cons_lock,
    )

flm_gemma4_prefill_block_begin

flm_gemma4_prefill_block_begin(
    *, head_dim: int = 512
) -> ExternalFunction

attn_block_begin of the prefill kernels: row j of m is prev_m[j].

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

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_block_begin(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_block_begin`` of the prefill kernels: row ``j`` of ``m`` is ``prev_m[j]``."""
    return _prefill_sibling(
        "attn_block_begin",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, In),
            reference=flm_gemma4_prefill_block_begin_ref,
            tolerance=Tolerance.exact(note="bf16 broadcast"),
            ops_per_call=0,
        ),
    )

flm_gemma4_prefill_block_begin_ref

flm_gemma4_prefill_block_begin_ref(prev_m)

Numpy reference for flm_gemma4_prefill_block_begin.

m is (LQ, LK) with LK == LQ in both builds.

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

def flm_gemma4_prefill_block_begin_ref(prev_m):
    """Numpy reference for [`flm_gemma4_prefill_block_begin`][iron.kernels.flm_gemma4.flm_gemma4_prefill_block_begin].

    ``m`` is ``(LQ, LK)`` with ``LK == LQ`` in both builds.
    """
    prev_m = np.asarray(prev_m)
    return np.repeat(prev_m, prev_m.shape[-1], axis=-1)

flm_gemma4_prefill_block_end

flm_gemma4_prefill_block_end(
    *, head_dim: int = 512
) -> ExternalFunction

attn_block_end of the prefill kernels: prev_m = new_m.

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

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_block_end(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_block_end`` of the prefill kernels: ``prev_m = new_m``."""
    return _prefill_sibling(
        "attn_block_end",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, In),
            reference=flm_gemma4_prefill_block_end_ref,
            tolerance=Tolerance.exact(note="bf16 copy"),
            ops_per_call=0,
        ),
    )

flm_gemma4_prefill_block_end_ref

flm_gemma4_prefill_block_end_ref(new_m)

Numpy reference for flm_gemma4_prefill_block_end.

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

def flm_gemma4_prefill_block_end_ref(new_m):
    """Numpy reference for [`flm_gemma4_prefill_block_end`][iron.kernels.flm_gemma4.flm_gemma4_prefill_block_end]."""
    return np.array(new_m, copy=True)

flm_gemma4_prefill_rounds_ref

flm_gemma4_prefill_rounds_ref(l_begin, l_end)

(L_end >> 7) - (L_begin >> 7): the 128-row rounds a dispatch spans.

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

def flm_gemma4_prefill_rounds_ref(l_begin, l_end):
    """``(L_end >> 7) - (L_begin >> 7)``: the 128-row rounds a dispatch spans."""
    n = (_word0(l_end) >> 7) - (_word0(l_begin) >> 7)
    return n.reshape(-1, 1)

flm_gemma4_prefill_rounds

flm_gemma4_prefill_rounds(
    *, head_dim: int = 512
) -> ExternalFunction

attn_rounds of the prefill kernels: the rounds between L_begin and L_end.

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

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_rounds(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_rounds`` of the prefill kernels: the rounds between ``L_begin`` and ``L_end``."""
    return _prefill_sibling(
        "attn_rounds",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            sample=_prefill_rounds_sample,
            reference=flm_gemma4_prefill_rounds_ref,
            tolerance=Tolerance.exact(note="integer shifts and a subtraction"),
            out_valid=1,
            ops_per_call=1,
        ),
    )

flm_gemma4_prefill_blocks_ref

flm_gemma4_prefill_blocks_ref(
    l_begin, i, *, window_size=None
)

Key blocks round i folds in: all up to its own, or those in the window.

window_size is the sliding-window build's RTP buffer, None for the head-dim-512 build.

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

def flm_gemma4_prefill_blocks_ref(l_begin, i, *, window_size=None):
    """Key blocks round ``i`` folds in: all up to its own, or those in the window.

    ``window_size`` is the sliding-window build's RTP buffer, ``None`` for the
    head-dim-512 build.
    """
    begin = _word0(l_begin)
    if window_size is None:
        n = (begin >> 7) + np.int32(i) + 1
    else:
        q = begin + np.int32(i) * 128
        k = np.maximum(q - _word0(window_size), 0)
        n = ((q - k) >> 7) + 1
    return n.astype(np.int32).reshape(-1, 1)

flm_gemma4_prefill_blocks

flm_gemma4_prefill_blocks(
    *, head_dim: int = 512
) -> ExternalFunction

attn_blocks of the prefill kernels: the key blocks round i folds in.

At a head dim of 256, the sliding-window build, it also reads the window size: attn_blocks(L_begin, window_size, i, n_out).

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

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_blocks(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_blocks`` of the prefill kernels: the key blocks round ``i`` folds in.

    At a head dim of 256, the sliding-window build, it also reads the window
    size: ``attn_blocks(L_begin, window_size, i, n_out)``.
    """
    window = head_dim == 256
    return _prefill_sibling(
        "attn_blocks",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Param, Out) if window else (In, Param, Out),
            sample=partial(_prefill_blocks_sample, window=window),
            reference=(
                _prefill_swa_blocks_ref if window else flm_gemma4_prefill_blocks_ref
            ),
            tolerance=Tolerance.exact(note="integer shifts, adds and a clamp"),
            out_valid=1,
            ops_per_call=1,
        ),
    )

flm_gemma4_prefill_finalize

flm_gemma4_prefill_finalize(
    *, head_dim: int = 512
) -> ExternalFunction

attn_finalize of the prefill kernels: inv_l = bf16(1 / l) per query row.

Parameters:

Name	Type	Description	Default
head_dim	int	512 builds the global kernel's 8 rows; 256 builds the sliding-window kernel's 16 rows.	512

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

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_finalize(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_finalize`` of the prefill kernels: ``inv_l = bf16(1 / l)`` per query row.

    Args:
        head_dim: 512 builds the global kernel's 8 rows; 256 builds the
            sliding-window kernel's 16 rows.
    """
    return _prefill_sibling(
        "attn_finalize",
        head_dim,
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_prefill_finalize_ref,
            sample=partial(
                _prefill_finalize_sample,
                lq=_ATTN_PREFILL_LQ if head_dim == 512 else _SWA_PREFILL_LQ,
            ),
            # Floor narrowing lands at most one bf16 step from the nearest-even
            # cast of 1 / l if aie::inv errs by less than half a step (2**-9
            # relative).
            tolerance=Tolerance.bf16_ulps(
                1,
                note="floor bf16 narrowing of AIE2P's scalar inv instruction, "
                "whose error the AIE-ML v2 intrinsics guide does not document",
            ),
        ),
    )

flm_gemma4_prefill_finalize_ref

flm_gemma4_prefill_finalize_ref(row_sums)

Numpy reference for flm_gemma4_prefill_finalize: 1 / l.

The tolerance covers the kernel's hardware reciprocal and its floor narrowing to bf16.

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

def flm_gemma4_prefill_finalize_ref(row_sums):
    """Numpy reference for [`flm_gemma4_prefill_finalize`][iron.kernels.flm_gemma4.flm_gemma4_prefill_finalize]: ``1 / l``.

    The tolerance covers the kernel's hardware reciprocal and its floor
    narrowing to bf16.
    """
    return 1.0 / np.asarray(row_sums, dtype=np.float32)

flm_gemma4_prefill_round_begin_ref

flm_gemma4_prefill_round_begin_ref(*, head_dim=512)

Numpy reference for flm_gemma4_prefill_round_begin.

The bytes of y | c | l | prev_m | new_m: zeros, ones, zeros and two rows of NEG_INF.

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

def flm_gemma4_prefill_round_begin_ref(*, head_dim=512):
    """Numpy reference for [`flm_gemma4_prefill_round_begin`][iron.kernels.flm_gemma4.flm_gemma4_prefill_round_begin].

    The bytes of ``y | c | l | prev_m | new_m``: zeros, ones, zeros and two
    rows of NEG_INF.
    """
    lq = _prefill_lq(head_dim)
    return np.concatenate(
        [
            np.zeros(lq * head_dim, np.float32).view(np.uint8),
            np.ones(lq, np.float32).view(np.uint8),
            np.zeros(lq, np.float32).view(np.uint8),
            np.full(2 * lq, _PREFILL_NEG_INF, bfloat16).view(np.uint8),
        ]
    )

flm_gemma4_prefill_round_begin

flm_gemma4_prefill_round_begin(
    *, head_dim: int = 512
) -> ExternalFunction

attn_round_begin of the prefill kernels, its five outputs packed as bytes.

out is y | c | l | prev_m | new_m.

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

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_round_begin(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_round_begin`` of the prefill kernels, its five outputs packed as bytes.

    ``out`` is ``y | c | l | prev_m | new_m``.
    """
    lq = _prefill_lq(head_dim)
    n = 4 * lq * head_dim + 12 * lq
    return _prefill_core(
        "attn_round_begin_core",
        head_dim,
        [np.ndarray[(n,), np.dtype[np.uint8]]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out,),
            reference=partial(flm_gemma4_prefill_round_begin_ref, head_dim=head_dim),
            tolerance=Tolerance.exact(note="constant stores"),
            ops_per_call=0,
        ),
    )

flm_gemma4_prefill_qk_core_ref

flm_gemma4_prefill_qk_core_ref(
    q,
    k,
    m,
    inner_k,
    inner_q,
    inner_k_current,
    *,
    head_dim=512
)

Numpy reference for flm_gemma4_prefill_qk_core: s | m.

s = q @ k.T, NEG_INF where masked, and m folds in s lanewise. The tolerance covers the bfp16 operands and the bf16 narrowing.

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

def flm_gemma4_prefill_qk_core_ref(
    q, k, m, inner_k, inner_q, inner_k_current, *, head_dim=512
):
    """Numpy reference for [`flm_gemma4_prefill_qk_core`][iron.kernels.flm_gemma4.flm_gemma4_prefill_qk_core]: ``s | m``.

    ``s = q @ k.T``, NEG_INF where masked, and ``m`` folds in ``s``
    lanewise. The tolerance covers the bfp16 operands and the bf16 narrowing.
    """
    lq = lk = _prefill_lq(head_dim)
    q, k = _tiles(q, lq, head_dim), _tiles(k, lk, head_dim, col_major=True)
    s = q @ k.transpose(0, 2, 1)
    keep = _prefill_keep(lq, lk, inner_k, inner_q, inner_k_current)
    s = np.where(keep, s, _PREFILL_NEG_INF)
    m = np.maximum(np.asarray(m, np.float64).reshape(-1, lq, lk), s)
    return np.concatenate([s.reshape(len(s), -1), m.reshape(len(m), -1)], axis=1)

flm_gemma4_prefill_qk_core

flm_gemma4_prefill_qk_core(
    *, head_dim: int = 512
) -> ExternalFunction

attn_qk_step's arithmetic for one key chunk, without its k lock.

attn_qk_core(q, k, m, inner_k, inner_q, inner_k_current, out): s = q @ k.T through the bfp16-emulated mmul, masked, and m folded with it, into out = s | m. The three positions are apply_mask_and_get_max's, which attn_qk_step derives from its RTPs. q and k arrive in the mmul's 8x8 tiles, k's tiles in column-major order.

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

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_qk_core(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_qk_step``'s arithmetic for one key chunk, without its k lock.

    ``attn_qk_core(q, k, m, inner_k, inner_q, inner_k_current, out)``:
    ``s = q @ k.T`` through the bfp16-emulated mmul, masked, and ``m``
    folded with it, into ``out = s | m``. The three positions are
    ``apply_mask_and_get_max``'s, which attn_qk_step derives from its RTPs.
    q and k arrive in the mmul's 8x8 tiles, k's tiles in column-major order.
    """
    lq = lk = _prefill_lq(head_dim)
    kv = np.ndarray[(lq * head_dim,), _BF16]
    m = np.ndarray[(lq * lk,), _BF16]
    return _prefill_core(
        "attn_qk_core",
        head_dim,
        [kv, kv, m, np.int32, np.int32, np.int32, np.ndarray[(2 * lq * lk,), _BF16]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Param, Param, Param, Out),
            sample=partial(_prefill_qk_sample, head_dim=head_dim),
            reference=partial(flm_gemma4_prefill_qk_core_ref, head_dim=head_dim),
            tolerance=Tolerance.bounded(
                partial(_prefill_qk_bound, head_dim=head_dim),
                note="bfp16ebs8 operand rounding (zero on the sample's grid), "
                "float32 accumulation, and the bf16 narrowing of s",
            ),
            ops_per_call=2 * lq * lk * head_dim,
        ),
    )

flm_gemma4_prefill_fv_core_ref

flm_gemma4_prefill_fv_core_ref(y, sv, *, head_dim=512)

Numpy reference for flm_gemma4_prefill_fv_core: y + s @ v.

y and s are in the mmul's 8x8 tiles, v in column-major tile order; the tolerance covers the bfp16 operands.

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

def flm_gemma4_prefill_fv_core_ref(y, sv, *, head_dim=512):
    """Numpy reference for [`flm_gemma4_prefill_fv_core`][iron.kernels.flm_gemma4.flm_gemma4_prefill_fv_core]: ``y + s @ v``.

    ``y`` and ``s`` are in the mmul's 8x8 tiles, ``v`` in column-major tile
    order; the tolerance covers the bfp16 operands.
    """
    lq = lk = _prefill_lq(head_dim)
    s, v = _prefill_sv(sv, lq, lk, head_dim)
    return _to_tiles(_tiles(y, lq, head_dim) + s @ v)

flm_gemma4_prefill_fv_core

flm_gemma4_prefill_fv_core(
    *, head_dim: int = 512
) -> ExternalFunction

attn_fv_step's arithmetic for one key chunk, without its v lock.

attn_fv_core(y, sv, y_out): y_out = y + s @ v through flm_attn_fv's bfp16-emulated mmul, with sv = s | v. y and s are in the mmul's 8x8 tiles, v's tiles in column-major order.

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

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_fv_core(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_fv_step``'s arithmetic for one key chunk, without its v lock.

    ``attn_fv_core(y, sv, y_out)``: ``y_out = y + s @ v`` through
    ``flm_attn_fv``'s bfp16-emulated mmul, with ``sv = s | v``. y and s are in
    the mmul's 8x8 tiles, v's tiles in column-major order.
    """
    lq = lk = _prefill_lq(head_dim)
    y = np.ndarray[(lq * head_dim,), _F32]
    return _prefill_core(
        "attn_fv_core",
        head_dim,
        [y, np.ndarray[(lq * lk + lk * head_dim,), _BF16], y],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            sample=partial(_prefill_fv_sample, head_dim=head_dim),
            reference=partial(flm_gemma4_prefill_fv_core_ref, head_dim=head_dim),
            tolerance=Tolerance.bounded(
                partial(_prefill_fv_core_bound, head_dim=head_dim),
                note="bfp16ebs8 operand rounding (zero on the sample's grid) "
                "and float32 accumulation",
            ),
            ops_per_call=2 * lq * lk * head_dim,
        ),
    )

flm_gemma4_prefill_block_mid_core_ref

flm_gemma4_prefill_block_mid_core_ref(
    in_bf16, in_f32, *, head_dim=512
)

Numpy reference for flm_gemma4_prefill_block_mid_core.

(s | new_m, y | l | c): new_m the row max of m, p = 2**((s - new_m) * log2e) and c = 2**((prev_m - new_m) * log2e) with FastFlowLM's bf16 log2(e), l = sum(p) + c * l and y = c * y.

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

def flm_gemma4_prefill_block_mid_core_ref(in_bf16, in_f32, *, head_dim=512):
    """Numpy reference for [`flm_gemma4_prefill_block_mid_core`][iron.kernels.flm_gemma4.flm_gemma4_prefill_block_mid_core].

    ``(s | new_m, y | l | c)``: ``new_m`` the row max of ``m``,
    ``p = 2**((s - new_m) * log2e)`` and ``c = 2**((prev_m - new_m) * log2e)``
    with FastFlowLM's bf16 log2(e), ``l = sum(p) + c * l`` and ``y = c * y``.
    """
    lq = _prefill_lq(head_dim)
    p, new_m, y, row_sums, c, _, _ = _prefill_block_mid_math(in_bf16, in_f32, head_dim)
    return (
        np.concatenate([_prefill_block_mid_s(p, lq), new_m], axis=1),
        np.concatenate([y, row_sums, c], axis=1),
    )

flm_gemma4_prefill_block_mid_core

flm_gemma4_prefill_block_mid_core(
    *, head_dim: int = 512
) -> ExternalFunction

attn_block_mid of the prefill kernels, its state passed in and out.

attn_block_mid_core(in_bf16, in_f32, out_bf16, out_f32) with in_bf16 = s | m | prev_m, in_f32 = y | l, out_bf16 = s | new_m and out_f32 = y | l | c: the row max, FastFlowLM's softmax, the correction c and the rescaled l and y.

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

@dtypes([{"head_dim": 256}])
def flm_gemma4_prefill_block_mid_core(*, head_dim: int = 512) -> ExternalFunction:
    """``attn_block_mid`` of the prefill kernels, its state passed in and out.

    ``attn_block_mid_core(in_bf16, in_f32, out_bf16, out_f32)`` with
    ``in_bf16 = s | m | prev_m``, ``in_f32 = y | l``, ``out_bf16 = s | new_m``
    and ``out_f32 = y | l | c``: the row max, FastFlowLM's softmax, the
    correction ``c`` and the rescaled ``l`` and ``y``.
    """
    lq = lk = _prefill_lq(head_dim)
    return _prefill_core(
        "attn_block_mid_core",
        head_dim,
        [
            np.ndarray[(lq * 128 + lq * lk + lq,), _BF16],
            np.ndarray[(lq * head_dim + lq,), _F32],
            np.ndarray[(lq * 128 + lq,), _BF16],
            np.ndarray[(lq * head_dim + 2 * lq,), _F32],
        ],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out, Out),
            sample=partial(_prefill_block_mid_sample, head_dim=head_dim),
            reference=partial(flm_gemma4_prefill_block_mid_core_ref, head_dim=head_dim),
            tolerance=Tolerance.bounded(
                partial(_prefill_block_mid_bound, head_dim=head_dim),
                note="aie::exp2's 6.98% (linalg.mha_softmax), bf16 rounding of "
                "the exp2 arguments, in the kernel (zero on the sample's grid) "
                "and inside aie::exp2 (measured on npu2), bfp16ebs8 "
                "rounding of p in the row sums, the emulated float32 multiply",
            ),
            # Per score: subtract, scale, exp2, sum; per y element: one multiply.
            ops_per_call=4 * lq * 128 + lq * head_dim,
        ),
    )

flm_gemma4_decode_glu

flm_gemma4_decode_glu(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the decode layer's gated linear unit, from flm_gemma4/decode_glu.cc.

glu(y, x_ping, x_pong, y_ping, y_pong, skip): waits on the RTP lock, then, unless skip[0] is set, turns 2 * intermediate_size gate and up values, arriving glu_slice at a time in the x ping-pong pair, into intermediate_size activations in the core-local y (twice that with double_wide_mlp), and sends them out glu_slice // 2 at a time through the y pair once per down-projection repeat. skip is an int32 RTP buffer of 16 words.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults x_prod_lock=0, x_cons_lock=1, y_prod_lock=2, y_cons_lock=3 and rtp_lock=6.	{}

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_glu(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the decode layer's gated linear unit, from ``flm_gemma4/decode_glu.cc``.

    ``glu(y, x_ping, x_pong, y_ping, y_pong, skip)``: waits on the RTP lock,
    then, unless ``skip[0]`` is set, turns ``2 * intermediate_size`` gate and
    up values, arriving ``glu_slice`` at a time in the ``x`` ping-pong pair,
    into ``intermediate_size`` activations in the core-local ``y`` (twice that
    with ``double_wide_mlp``), and sends them out ``glu_slice // 2`` at a time
    through the ``y`` pair once per down-projection repeat. ``skip`` is an
    int32 RTP buffer of 16 words.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``x_prod_lock=0``,
            ``x_cons_lock=1``, ``y_prod_lock=2``, ``y_cons_lock=3`` and
            ``rtp_lock=6``.
    """
    bf16 = np.dtype[bfloat16]
    hid = geometry.intermediate_size * (2 if geometry.double_wide_mlp else 1)
    x_ty = np.ndarray[(geometry.glu_slice,), bf16]
    y_ty = np.ndarray[(geometry.glu_slice // 2,), bf16]
    return _decode_kernel(
        "glu",
        "glu",
        [
            np.ndarray[(hid,), bf16],
            x_ty,
            x_ty,
            y_ty,
            y_ty,
            np.ndarray[(16,), np.dtype[np.int32]],
        ],
        (Out, In, In, Out, Out, Param),
        dict(x_prod_lock=0, x_cons_lock=1, y_prod_lock=2, y_cons_lock=3, rtp_lock=6),
        locks,
        geometry,
        lut=True,
    )

flm_gemma4_glu_core_ref

flm_gemma4_glu_core_ref(x)

Numpy reference for flm_gemma4_glu_core: gelu(gate) * up.

x holds, per call, the up half then the gate half. GELU is the tanh form Gemma 4 uses; the tolerance covers the kernel's table.

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

def flm_gemma4_glu_core_ref(x):
    """Numpy reference for [`flm_gemma4_glu_core`][iron.kernels.flm_gemma4.flm_gemma4_glu_core]: ``gelu(gate) * up``.

    ``x`` holds, per call, the up half then the gate half. GELU is the tanh
    form Gemma 4 uses; the tolerance covers the kernel's table.
    """
    up, gate = _glu_split(x)
    return (_gelu_tanh(gate) * up).astype(np.float32)

flm_gemma4_glu_core

flm_gemma4_glu_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

pseduo_glu of flm_gemma4_decode_glu: gelu(gate) * up for one slice.

glu_core(x, y): x is glu_slice bf16, the up half then the gate half; y is glu_slice // 2 bf16. GELU goes through aie_runtime_lib's table (getGeluBf16).

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_glu_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``pseduo_glu`` of [`flm_gemma4_decode_glu`][iron.kernels.flm_gemma4.flm_gemma4_decode_glu]: ``gelu(gate) * up`` for one slice.

    ``glu_core(x, y)``: ``x`` is ``glu_slice`` bf16, the up half then the
    gate half; ``y`` is ``glu_slice // 2`` bf16. GELU goes through
    aie_runtime_lib's table (``getGeluBf16``).

    Args:
        geometry: The model the kernel builds for.
    """
    n = geometry.glu_slice
    return _wrapper(
        flm_gemma4_decode_glu(geometry=geometry),
        "decode_glu_core.cc",
        "glu_core",
        [np.ndarray[(n,), _BF16], np.ndarray[(n // 2,), _BF16]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_glu_core_ref,
            sample=partial(_glu_core_sample, n=n),
            tolerance=Tolerance.bounded(
                _glu_core_bound,
                note="the GELU table's error against tanh-GELU, measured on the "
                "host from aie_runtime_lib's table, plus a bf16 ulp of its slope "
                "and the floor bf16 stores of the GELU and the product",
            ),
            uses_lut=True,
            # Per output: the table line's multiply and add, then the product.
            ops_per_call=3 * (n // 2),
        ),
    )

flm_gemma4_decode_attn_kv

flm_gemma4_decode_attn_kv(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the global-attention score-times-value core for one KV head, from flm_gemma4/decode_attn_kv.cc.

The decode layer's global attention runs on two cores. This one folds each round's scores and running-max corrections from the qk core into the softmax denominator l and the float32 accumulator y, and at the end writes o = y / l. The returned kernel is attn_kv_begin, which zeroes y and l and waits for the RTPs. Per round the Worker calls attn_kv_round, then once attn_kv_finish. y and o need 64-byte-aligned base pointers.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; num_kv_heads must be 1.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults v_prod_lock=2, v_cons_lock=3, o_prod_lock=0, o_cons_lock=1 and l_cons_lock=8 (a lock of the qk core, on the tile below).	{}

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

def flm_gemma4_decode_attn_kv(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the global-attention score-times-value core for one KV head, from ``flm_gemma4/decode_attn_kv.cc``.

    The decode layer's global attention runs on two cores. This one folds
    each round's scores and running-max corrections from the qk core into
    the softmax denominator ``l`` and the float32 accumulator ``y``, and at
    the end writes ``o = y / l``. The returned kernel is ``attn_kv_begin``,
    which zeroes ``y`` and ``l`` and waits for the RTPs. Per round the
    Worker calls ``attn_kv_round``, then once ``attn_kv_finish``. ``y`` and
    ``o`` need 64-byte-aligned base pointers.

    Args:
        geometry: The model the kernel builds for; ``num_kv_heads`` must be 1.
        **locks: Core lock ids overriding the defaults ``v_prod_lock=2``,
            ``v_cons_lock=3``, ``o_prod_lock=0``, ``o_cons_lock=1`` and
            ``l_cons_lock=8`` (a lock of the qk core, on the tile below).
    """
    _decode_kv_heads("flm_gemma4_decode_attn_kv", geometry, 1)
    t = _decode_attn_types(geometry, geometry.dh, geometry.dh)
    fn = _decode_kernel(
        "attn_kv",
        "attn_kv_begin",
        [t["y"], t["l"]],
        (Out, Out),
        _DECODE_KV_LOCKS,
        locks,
        geometry,
        lut=True,
        cls=_AttnKvKernel,
    )
    bind = fn.object_file.bind
    fn.attn_kv_round = bind("attn_kv_round", [t["s"], t["kv"], t["kv"], t["y"], t["l"]])
    fn.attn_kv_finish = bind("attn_kv_finish", [t["y"], t["o"], t["l"]])
    return fn

flm_gemma4_attn_kv_core

flm_gemma4_attn_kv_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One attn_kv_round of flm_gemma4_decode_attn_kv without its v lock.

attn_kv_round_core(sv, ly_in, ly_out): ly_out is ly_in with l = c * l + rowsum(s) and y = c * y + s @ v per query head. sv holds s, then v. ly holds y, then l in a 16-float slot.

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

def flm_gemma4_attn_kv_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """One ``attn_kv_round`` of [`flm_gemma4_decode_attn_kv`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv] without its v lock.

    ``attn_kv_round_core(sv, ly_in, ly_out)``: ``ly_out`` is ``ly_in`` with
    ``l = c * l + rowsum(s)`` and ``y = c * y + s @ v`` per query head. ``sv``
    holds ``s``, then ``v``. ``ly`` holds ``y``, then ``l`` in a 16-float slot.
    """
    base = flm_gemma4_decode_attn_kv(geometry=geometry)
    dh = geometry.dh
    return _attn_kv_round_core(
        base,
        "decode_attn_kv_core.cc",
        "attn_kv_round_core",
        dh=dh,
        kv_width=dh,
        impl="1x8x1",
        rows=8,
        ref=partial(flm_gemma4_attn_kv_core_ref, dh=dh),
    )

flm_gemma4_attn_kv_core_ref

flm_gemma4_attn_kv_core_ref(sv, ly, *, dh=512)

Numpy reference for flm_gemma4_attn_kv_core, in float64.

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

def flm_gemma4_attn_kv_core_ref(sv, ly, *, dh=512):
    """Numpy reference for [`flm_gemma4_attn_kv_core`][iron.kernels.flm_gemma4.flm_gemma4_attn_kv_core], in float64."""
    return _attn_kv_round_ref(sv, ly, dh=dh, impl="1x8x1")

flm_gemma4_decode_attn_kv_kvh2

flm_gemma4_decode_attn_kv_kvh2(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE,
    **locks: int
) -> ExternalFunction

Build the global-attention score-times-value core for two KV heads, from flm_gemma4/decode_attn_kv_kvh2.cc.

The two-KV-head (E4B) sibling of flm_gemma4_decode_attn_kv: the Worker runs a round as one attn_kv_s_begin, which folds the scores into l and rescales y, then one attn_kv_v_half per KV head.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; num_kv_heads must be 2.	FLM_GEMMA4_E4B_DECODE
**locks	int	Core lock ids overriding the defaults v_prod_lock=2, v_cons_lock=3, o_prod_lock=0, o_cons_lock=1 and l_cons_lock=8 (a lock of the qk core, on the tile below).	{}

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

def flm_gemma4_decode_attn_kv_kvh2(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the global-attention score-times-value core for two KV heads, from ``flm_gemma4/decode_attn_kv_kvh2.cc``.

    The two-KV-head (E4B) sibling of
    [`flm_gemma4_decode_attn_kv`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv]:
    the Worker runs a round as one ``attn_kv_s_begin``, which folds the
    scores into ``l`` and rescales ``y``, then one ``attn_kv_v_half`` per
    KV head.

    Args:
        geometry: The model the kernel builds for; ``num_kv_heads`` must be 2.
        **locks: Core lock ids overriding the defaults ``v_prod_lock=2``,
            ``v_cons_lock=3``, ``o_prod_lock=0``, ``o_cons_lock=1`` and
            ``l_cons_lock=8`` (a lock of the qk core, on the tile below).
    """
    _decode_kv_heads("flm_gemma4_decode_attn_kv_kvh2", geometry, 2)
    t = _decode_attn_types(geometry, geometry.dh, geometry.dh)
    fn = _decode_kernel(
        "attn_kv_kvh2",
        "attn_kv_begin",
        [t["y"], t["l"]],
        (Out, Out),
        _DECODE_KV_LOCKS,
        locks,
        geometry,
        lut=True,
        cls=_AttnKvKvh2Kernel,
    )
    bind = fn.object_file.bind
    fn.attn_kv_s_begin = bind("attn_kv_s_begin", [t["s"], t["y"], t["l"]])
    fn.attn_kv_v_half = bind(
        "attn_kv_v_half", [t["s"], t["kv"], t["kv"], t["y"], np.int32]
    )
    fn.attn_kv_finish = bind("attn_kv_finish", [t["y"], t["o"], t["l"]])
    return fn

flm_gemma4_attn_kv_kvh2_core

flm_gemma4_attn_kv_kvh2_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction

attn_kv_s_begin, then one KV head's attn_kv_v_half, without the v lock.

The entry points are those of flm_gemma4_decode_attn_kv_kvh2. attn_kv_kvh2_round_core(sv, ly_in, ly_out, kv_head): ly_out is ly_in with l = c * l + rowsum(s) and y = c * y, plus s @ v for query heads 4 * kv_head to 4 * kv_head + 3. sv holds s, then v. ly holds y, then l in a 16-float slot. Both KV heads' v buffers do not fit in core memory beside ly, so a call reads one.

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

def flm_gemma4_attn_kv_kvh2_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction:
    """``attn_kv_s_begin``, then one KV head's ``attn_kv_v_half``, without the v lock.

    The entry points are those of
    [`flm_gemma4_decode_attn_kv_kvh2`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv_kvh2].
    ``attn_kv_kvh2_round_core(sv, ly_in, ly_out, kv_head)``: ``ly_out`` is
    ``ly_in`` with ``l = c * l + rowsum(s)`` and ``y = c * y``, plus ``s @ v``
    for query heads ``4 * kv_head`` to ``4 * kv_head + 3``. ``sv`` holds
    ``s``, then ``v``. ``ly`` holds ``y``, then ``l`` in a 16-float slot.
    Both KV heads' v buffers do not fit in core memory beside ``ly``, so a
    call reads one.
    """
    base = flm_gemma4_decode_attn_kv_kvh2(geometry=geometry)
    dh = geometry.dh
    return _attn_kv_round_core(
        base,
        "decode_attn_kv_kvh2_core.cc",
        "attn_kv_kvh2_round_core",
        dh=dh,
        kv_width=dh,
        impl="kvh2",
        rows=4,
        ref=partial(flm_gemma4_attn_kv_kvh2_core_ref, dh=dh),
    )

flm_gemma4_attn_kv_kvh2_core_ref

flm_gemma4_attn_kv_kvh2_core_ref(
    sv, ly, kv_head, *, dh=512
)

Numpy reference for flm_gemma4_attn_kv_kvh2_core, in float64.

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

def flm_gemma4_attn_kv_kvh2_core_ref(sv, ly, kv_head, *, dh=512):
    """Numpy reference for [`flm_gemma4_attn_kv_kvh2_core`][iron.kernels.flm_gemma4.flm_gemma4_attn_kv_kvh2_core], in float64."""
    return _attn_kv_round_ref(sv, ly, int(kv_head), dh=dh, impl="kvh2")

flm_gemma4_decode_attn_qk

flm_gemma4_decode_attn_qk(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build an attention query-times-key core, from flm_gemma4/decode_attn_qk.cc.

The other half of flm_gemma4_decode_attn_kv: it multiplies the query heads into each round's 16 key rows, masks the keys past the sequence length, keeps the running row maximum m and hands the kv core the exponentiated scores with their corrections. The returned kernel is attn_qk_begin, which sets m to -inf and releases the kv core; per round the Worker calls attn_qk_round. The global build serves one KV head (E2B), the sliding-window build both.

Parameters:

Name	Type	Description	Default
sliding_window	bool	Build for the sliding-window layers' head dim.	False
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; without sliding_window num_kv_heads must be 1.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults k_prod_lock=2, k_cons_lock=3 and l_cons_lock=8.	{}

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

@dtypes(
    [
        {"sliding_window": True},
        {"sliding_window": True, "geometry": FLM_GEMMA4_E4B_DECODE},
    ]
)
def flm_gemma4_decode_attn_qk(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int,
) -> ExternalFunction:
    """Build an attention query-times-key core, from ``flm_gemma4/decode_attn_qk.cc``.

    The other half of [`flm_gemma4_decode_attn_kv`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv]:
    it multiplies the query heads into each round's 16 key rows, masks the
    keys past the sequence length, keeps the running row maximum ``m`` and
    hands the kv core the exponentiated scores with their corrections. The
    returned kernel is ``attn_qk_begin``, which sets ``m`` to -inf and
    releases the kv core; per round the Worker calls ``attn_qk_round``.
    The global build serves one KV head (E2B), the sliding-window build both.

    Args:
        sliding_window: Build for the sliding-window layers' head dim.
        geometry: The model the kernel builds for; without ``sliding_window``
            ``num_kv_heads`` must be 1.
        **locks: Core lock ids overriding the defaults ``k_prod_lock=2``,
            ``k_cons_lock=3`` and ``l_cons_lock=8``.
    """
    if not sliding_window:
        _decode_kv_heads("flm_gemma4_decode_attn_qk", geometry, 1)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    t = _decode_attn_types(geometry, dh, geometry.num_kv_heads * dh)
    fn = _decode_kernel(
        "attn_qk",
        "attn_qk_begin",
        [t["m"]],
        (Out,),
        dict(k_prod_lock=2, k_cons_lock=3, l_cons_lock=8),
        locks,
        geometry,
        flags=(f"-DFLM_GEMMA4_DECODE_ATTN_QK_SWA={int(sliding_window)}",),
        cls=_AttnQkKernel,
    )
    fn.attn_qk_round = fn.object_file.bind(
        "attn_qk_round",
        [t["q"], t["kv"], t["kv"], t["s"], t["m"], t["c"], np.int32, np.int32],
    )
    return fn

flm_gemma4_decode_attn_qk_kvh2

flm_gemma4_decode_attn_qk_kvh2(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE,
    **locks: int
) -> ExternalFunction

Build the global-attention query-times-key core for two KV heads, from flm_gemma4/decode_attn_qk_kvh2.cc.

The two-KV-head (E4B) sibling of flm_gemma4_decode_attn_qk: the Worker runs a round as one attn_qk_half per KV head, then attn_qk_store_c, which copies the corrections into the scores.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for; num_kv_heads must be 2.	FLM_GEMMA4_E4B_DECODE
**locks	int	Core lock ids overriding the defaults k_prod_lock=2, k_cons_lock=3 and l_prod_lock=8.	{}

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

def flm_gemma4_decode_attn_qk_kvh2(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the global-attention query-times-key core for two KV heads, from ``flm_gemma4/decode_attn_qk_kvh2.cc``.

    The two-KV-head (E4B) sibling of
    [`flm_gemma4_decode_attn_qk`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_qk]:
    the Worker runs a round as one ``attn_qk_half`` per KV head, then
    ``attn_qk_store_c``, which copies the corrections into the scores.

    Args:
        geometry: The model the kernel builds for; ``num_kv_heads`` must be 2.
        **locks: Core lock ids overriding the defaults ``k_prod_lock=2``,
            ``k_cons_lock=3`` and ``l_prod_lock=8``.
    """
    _decode_kv_heads("flm_gemma4_decode_attn_qk_kvh2", geometry, 2)
    t = _decode_attn_types(geometry, geometry.dh, geometry.dh)
    fn = _decode_kernel(
        "attn_qk_kvh2",
        "attn_qk_begin",
        [t["m"]],
        (Out,),
        dict(k_prod_lock=2, k_cons_lock=3, l_prod_lock=8),
        locks,
        geometry,
        cls=_AttnQkKvh2Kernel,
    )
    bind = fn.object_file.bind
    fn.attn_qk_half = bind(
        "attn_qk_half",
        [
            t["q"],
            t["kv"],
            t["kv"],
            t["s"],
            t["m"],
            t["c"],
            np.int32,
            np.int32,
            np.int32,
        ],
    )
    fn.attn_qk_store_c = bind("attn_qk_store_c", [t["s"], t["c"]])
    return fn

flm_gemma4_attn_qk_core_ref

flm_gemma4_attn_qk_core_ref(qm, k, iter, L0)

Numpy reference for flm_gemma4_attn_qk_core.

Keys j with j < L0 - 16 * iter are valid. Per query row the new running max is max(m_in, valid scores); s holds exp(score - max) for valid keys and 0 for the others, then the 16 running maxima (rows 8 to 15 copied from m_in). The second output is c = exp(m_in - max). Both exponents are clamped at -87, as the kernel does.

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

def flm_gemma4_attn_qk_core_ref(qm, k, iter, L0):
    """Numpy reference for [`flm_gemma4_attn_qk_core`][iron.kernels.flm_gemma4.flm_gemma4_attn_qk_core].

    Keys ``j`` with ``j < L0 - 16 * iter`` are valid. Per query row the new
    running max is ``max(m_in, valid scores)``; ``s`` holds ``exp(score -
    max)`` for valid keys and 0 for the others, then the 16 running maxima
    (rows 8 to 15 copied from ``m_in``). The second output is ``c = exp(m_in
    - max)``. Both exponents are clamped at -87, as the kernel does.
    """
    _, _, m_in, m_new, _, e, c = _attn_qk_parts(qm, k, iter, L0)
    m_out = m_in.copy()
    m_out[:, :8] = m_new
    return np.concatenate([e.reshape(len(e), -1), m_out], axis=1), c

flm_gemma4_attn_qk_core

flm_gemma4_attn_qk_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One attn_qk_round of flm_gemma4_decode_attn_qk, without its k lock.

attn_qk_round_core(qm, k, s, c, iter, L0) takes the running max after q in qm and writes it after the scores in s. c holds the corrections the round stores in the tail of its s object.

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

@dtypes(
    [
        {"sliding_window": True},
        {"sliding_window": True, "geometry": FLM_GEMMA4_E4B_DECODE},
    ]
)
def flm_gemma4_attn_qk_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> ExternalFunction:
    """One ``attn_qk_round`` of [`flm_gemma4_decode_attn_qk`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_qk], without its k lock.

    ``attn_qk_round_core(qm, k, s, c, iter, L0)`` takes the running max
    after q in ``qm`` and writes it after the scores in ``s``. ``c`` holds the
    corrections the round stores in the tail of its s object.
    """
    base = flm_gemma4_decode_attn_qk(sliding_window=sliding_window, geometry=geometry)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    kv_heads = geometry.num_kv_heads
    return _wrapper(
        base,
        "decode_attn_qk_core.cc",
        "attn_qk_round_core",
        _attn_qk_core_types(geometry, dh, kv_heads),
        _attn_qk_core_contract(dh, kv_heads, flm_gemma4_attn_qk_core_ref),
    )

flm_gemma4_attn_qk_kvh2_core

flm_gemma4_attn_qk_kvh2_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction

One round of flm_gemma4_decode_attn_qk_kvh2, without its k lock.

attn_qk_kvh2_round_core(qm, k, s, c, iter, L0) runs attn_qk_half for both KV heads, whose k objects k holds back to back. Otherwise as flm_gemma4_attn_qk_core.

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

def flm_gemma4_attn_qk_kvh2_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E4B_DECODE
) -> ExternalFunction:
    """One round of [`flm_gemma4_decode_attn_qk_kvh2`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_qk_kvh2], without its k lock.

    ``attn_qk_kvh2_round_core(qm, k, s, c, iter, L0)`` runs
    ``attn_qk_half`` for both KV heads, whose k objects ``k`` holds back to
    back. Otherwise as
    [`flm_gemma4_attn_qk_core`][iron.kernels.flm_gemma4.flm_gemma4_attn_qk_core].
    """
    base = flm_gemma4_decode_attn_qk_kvh2(geometry=geometry)
    return _wrapper(
        base,
        "decode_attn_qk_kvh2_core.cc",
        "attn_qk_kvh2_round_core",
        _attn_qk_core_types(geometry, geometry.dh, 2),
        _attn_qk_core_contract(geometry.dh, 2, flm_gemma4_attn_qk_kvh2_core_ref),
    )

flm_gemma4_attn_qk_kvh2_core_ref

flm_gemma4_attn_qk_kvh2_core_ref(qm, k, iter, L0)

Numpy reference for flm_gemma4_attn_qk_kvh2_core; see flm_gemma4_attn_qk_core_ref.

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

def flm_gemma4_attn_qk_kvh2_core_ref(qm, k, iter, L0):
    """Numpy reference for [`flm_gemma4_attn_qk_kvh2_core`][iron.kernels.flm_gemma4.flm_gemma4_attn_qk_kvh2_core]; see ``flm_gemma4_attn_qk_core_ref``."""
    return flm_gemma4_attn_qk_core_ref(qm, k, iter, L0)

flm_gemma4_decode_swa_attn_kv

flm_gemma4_decode_swa_attn_kv(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the sliding-window-attention score-times-value core, from flm_gemma4/decode_swa_attn_kv.cc.

flm_gemma4_decode_attn_kv at head dim swa_dh, with entry points swa_attn_kv_begin, swa_attn_kv_round and swa_attn_kv_finish. One round covers every KV head, so it builds for one and for two.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults v_prod_lock=2, v_cons_lock=3, o_prod_lock=0, o_cons_lock=1 and l_cons_lock=8 (a lock of the qk core, on the tile below).	{}

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_swa_attn_kv(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the sliding-window-attention score-times-value core, from ``flm_gemma4/decode_swa_attn_kv.cc``.

    [`flm_gemma4_decode_attn_kv`][iron.kernels.flm_gemma4.flm_gemma4_decode_attn_kv]
    at head dim ``swa_dh``, with entry points ``swa_attn_kv_begin``,
    ``swa_attn_kv_round`` and ``swa_attn_kv_finish``. One round covers every
    KV head, so it builds for one and for two.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``v_prod_lock=2``,
            ``v_cons_lock=3``, ``o_prod_lock=0``, ``o_cons_lock=1`` and
            ``l_cons_lock=8`` (a lock of the qk core, on the tile below).
    """
    t = _decode_attn_types(
        geometry, geometry.swa_dh, geometry.num_kv_heads * geometry.swa_dh
    )
    fn = _decode_kernel(
        "swa_attn_kv",
        "swa_attn_kv_begin",
        [t["y"], t["l"]],
        (Out, Out),
        _DECODE_KV_LOCKS,
        locks,
        geometry,
        lut=True,
        cls=_SwaAttnKvKernel,
    )
    bind = fn.object_file.bind
    fn.swa_attn_kv_round = bind(
        "swa_attn_kv_round", [t["s"], t["kv"], t["kv"], t["y"], t["l"]]
    )
    fn.swa_attn_kv_finish = bind("swa_attn_kv_finish", [t["y"], t["o"], t["l"]])
    return fn

flm_gemma4_swa_attn_kv_core

flm_gemma4_swa_attn_kv_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

One swa_attn_kv_round of flm_gemma4_decode_swa_attn_kv without its v lock.

swa_attn_kv_round_core(sv, ly_in, ly_out): ly_out is ly_in with l = c * l + rowsum(s) and y = c * y + s @ v per query head, each head reading its KV head's half of v when there are two. sv holds s, then v. ly holds y, then l in a 16-float slot.

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_swa_attn_kv_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """One ``swa_attn_kv_round`` of [`flm_gemma4_decode_swa_attn_kv`][iron.kernels.flm_gemma4.flm_gemma4_decode_swa_attn_kv] without its v lock.

    ``swa_attn_kv_round_core(sv, ly_in, ly_out)``: ``ly_out`` is ``ly_in``
    with ``l = c * l + rowsum(s)`` and ``y = c * y + s @ v`` per query head,
    each head reading its KV head's half of ``v`` when there are two. ``sv``
    holds ``s``, then ``v``. ``ly`` holds ``y``, then ``l`` in a 16-float
    slot.
    """
    base = flm_gemma4_decode_swa_attn_kv(geometry=geometry)
    dh, kv_heads = geometry.swa_dh, geometry.num_kv_heads
    return _attn_kv_round_core(
        base,
        "decode_swa_attn_kv_core.cc",
        "swa_attn_kv_round_core",
        dh=dh,
        kv_width=kv_heads * dh,
        impl="1x8x1" if kv_heads == 1 else "2x4x1",
        rows=8,
        ref=partial(flm_gemma4_swa_attn_kv_core_ref, dh=dh, kv_heads=kv_heads),
    )

flm_gemma4_swa_attn_kv_core_ref

flm_gemma4_swa_attn_kv_core_ref(
    sv, ly, *, dh=256, kv_heads=1
)

Numpy reference for flm_gemma4_swa_attn_kv_core, in float64.

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

def flm_gemma4_swa_attn_kv_core_ref(sv, ly, *, dh=256, kv_heads=1):
    """Numpy reference for [`flm_gemma4_swa_attn_kv_core`][iron.kernels.flm_gemma4.flm_gemma4_swa_attn_kv_core], in float64."""
    impl = "1x8x1" if kv_heads == 1 else "2x4x1"
    return _attn_kv_round_ref(sv, ly, dh=dh, impl=impl)

flm_gemma4_decode_rope

flm_gemma4_decode_rope(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build a layer's q/k norm and rotary embedding, from flm_gemma4/decode_rope.cc.

rope(q, k, v, qkv_ping, qkv_pong, rope_w, skip_kv): takes the q, k and v projections one head (dh bf16) at a time through the qkv ping-pong pair, RMS-normalizes each q and k head in place against its weight and rotates it into q or k, and RMS-normalizes each v head into v. Unless skip_kv[0] is set, when the layer reuses another layer's KV cache, it hands k and v on. q is q_heads_padded * dh bf16, k and v num_kv_heads * dh, rope_w the cos and sin halves then, with qk_norm, the q and k norm weights (3 * dh), and skip_kv an int32 RTP buffer of 16 words. dh is the geometry's swa_dh with sliding_window, else its dh.

Parameters:

Name	Type	Description	Default
sliding_window	bool	Build for the sliding-window layers' head dim.	False
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults qkv_prod_lock=0, qkv_cons_lock=1, k_prod_lock=4, k_cons_lock=5, v_prod_lock=6, v_cons_lock=7, rope_prod_lock=8 and rope_cons_lock=9.	{}

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

@dtypes(
    [
        {"geometry": FLM_GEMMA4_E4B_DECODE},
        {"sliding_window": True},
        {"sliding_window": True, "geometry": FLM_GEMMA4_E4B_DECODE},
    ]
)
def flm_gemma4_decode_rope(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int,
) -> ExternalFunction:
    """Build a layer's q/k norm and rotary embedding, from ``flm_gemma4/decode_rope.cc``.

    ``rope(q, k, v, qkv_ping, qkv_pong, rope_w, skip_kv)``: takes the q, k
    and v projections one head (``dh`` bf16) at a time through the
    ``qkv`` ping-pong pair, RMS-normalizes each q and k head in place against
    its weight and rotates it into ``q`` or ``k``, and RMS-normalizes each v
    head into ``v``. Unless ``skip_kv[0]`` is set, when the layer reuses
    another layer's KV cache, it hands ``k`` and ``v`` on. ``q`` is
    ``q_heads_padded * dh`` bf16, ``k`` and ``v`` ``num_kv_heads * dh``,
    ``rope_w`` the cos and sin halves then, with ``qk_norm``, the q and k
    norm weights (``3 * dh``), and ``skip_kv`` an int32 RTP buffer of 16
    words. ``dh`` is the geometry's ``swa_dh`` with ``sliding_window``, else
    its ``dh``.

    Args:
        sliding_window: Build for the sliding-window layers' head dim.
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``qkv_prod_lock=0``,
            ``qkv_cons_lock=1``, ``k_prod_lock=4``, ``k_cons_lock=5``,
            ``v_prod_lock=6``, ``v_cons_lock=7``, ``rope_prod_lock=8`` and
            ``rope_cons_lock=9``.
    """
    dh = geometry.swa_dh if sliding_window else geometry.dh
    bf16 = np.dtype[bfloat16]
    kv = geometry.num_kv_heads * dh
    qkv_ty = np.ndarray[(dh,), bf16]
    return _decode_kernel(
        "rope",
        "rope",
        [
            np.ndarray[(_decode_q_heads_padded(geometry) * dh,), bf16],
            np.ndarray[(kv,), bf16],
            np.ndarray[(kv,), bf16],
            qkv_ty,
            qkv_ty,
            np.ndarray[(dh * (3 if geometry.qk_norm else 1),), bf16],
            _DECODE_RTP,
        ],
        (Out, Out, Out, InOut, InOut, In, Param),
        _DECODE_ROPE_LOCKS,
        locks,
        geometry,
        flags=(f"-DFLM_GEMMA4_DECODE_ROPE_SWA={int(sliding_window)}",),
    )

flm_gemma4_rope_core_ref

flm_gemma4_rope_core_ref(x, rope_w)

Numpy reference for flm_gemma4_rope_core.

Per call, x is a q head then a k head and rope_w is [cos | sin | q weight | k weight]. Each head n = x * w / sqrt(mean(x^2) + 1e-6) is rotated by halves: y1 = n1 cos - n2 sin, y2 = n1 sin + n2 cos, with n1, n2 its first and second half.

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

def flm_gemma4_rope_core_ref(x, rope_w):
    """Numpy reference for [`flm_gemma4_rope_core`][iron.kernels.flm_gemma4.flm_gemma4_rope_core].

    Per call, ``x`` is a q head then a k head and ``rope_w`` is ``[cos |
    sin | q weight | k weight]``. Each head ``n = x * w / sqrt(mean(x^2) +
    1e-6)`` is rotated by halves: ``y1 = n1 cos - n2 sin``, ``y2 = n1 sin +
    n2 cos``, with ``n1``, ``n2`` its first and second half.
    """
    a, b = _rope_core_terms(x, rope_w)
    y = (a + b).reshape(len(a), -1)
    return y.astype(np.float32).astype(bfloat16)

flm_gemma4_rope_core

flm_gemma4_rope_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

_rotate_t of flm_gemma4_decode_rope on one q and one k head.

rope_head_core(x, rope_w, y): x is a q head then a k head, rope_w the base kernel's [cos | sin | q weight | k weight]. The kernel normalizes x in place; the write lands in the core's input element, which the next DMA fill overwrites.

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

@dtypes(_ROPE_CORE_VARIANTS)
def flm_gemma4_rope_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> ExternalFunction:
    """``_rotate_t`` of [`flm_gemma4_decode_rope`][iron.kernels.flm_gemma4.flm_gemma4_decode_rope] on one q and one k head.

    ``rope_head_core(x, rope_w, y)``: ``x`` is a q head then a k head,
    ``rope_w`` the base kernel's ``[cos | sin | q weight | k weight]``.
    The kernel normalizes ``x`` in place; the write lands in the core's input
    element, which the next DMA fill overwrites.
    """
    base = flm_gemma4_decode_rope(sliding_window=sliding_window, geometry=geometry)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    heads = np.ndarray[(2 * dh,), _BF16]
    return _wrapper(
        base,
        "decode_rope_core.cc",
        "rope_head_core",
        [heads, np.ndarray[(3 * dh,), _BF16], heads],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            reference=flm_gemma4_rope_core_ref,
            sample=partial(_rope_core_sample, dh=dh),
            tolerance=Tolerance.bounded(
                _rope_core_bound,
                note="bf16 ulp <= 2**-7 |v|. The kernel floors each normalized "
                "value to bf16 (1 ulp, plus the rms scale's error, under "
                "2**-12), so each rotation product is off by that much; it "
                "floors the rotated sum to bf16 (1 ulp) and the reference "
                "rounds it to nearest (1/2 ulp)",
            ),
            # Per element: square-add and two multiplies for the norm, then
            # a multiply and a multiply-add for the rotation.
            ops_per_call=2 * 7 * dh,
        ),
    )

flm_gemma4_v_norm_core_ref

flm_gemma4_v_norm_core_ref(x)

Numpy reference for flm_gemma4_v_norm_core: x / sqrt(mean(x^2) + 1e-6).

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

def flm_gemma4_v_norm_core_ref(x):
    """Numpy reference for [`flm_gemma4_v_norm_core`][iron.kernels.flm_gemma4.flm_gemma4_v_norm_core]: ``x / sqrt(mean(x^2) + 1e-6)``."""
    x = np.asarray(x, np.float64)
    return (x * _rms(x)).astype(np.float32).astype(bfloat16)

flm_gemma4_v_norm_core

flm_gemma4_v_norm_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

rms_norm_unweighted of flm_gemma4_decode_rope on one v head.

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

@dtypes(_ROPE_CORE_VARIANTS)
def flm_gemma4_v_norm_core(
    *,
    sliding_window: bool = False,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
) -> ExternalFunction:
    """``rms_norm_unweighted`` of [`flm_gemma4_decode_rope`][iron.kernels.flm_gemma4.flm_gemma4_decode_rope] on one v head."""
    base = flm_gemma4_decode_rope(sliding_window=sliding_window, geometry=geometry)
    dh = geometry.swa_dh if sliding_window else geometry.dh
    head = np.ndarray[(dh,), _BF16]
    return _wrapper(
        base,
        "decode_rope_core.cc",
        "v_norm_core",
        [head, head],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_v_norm_core_ref,
            sample=partial(_v_norm_core_sample, dh=dh),
            tolerance=Tolerance.bounded(
                _v_norm_core_bound,
                note="bf16 ulp <= 2**-7 |v|: the kernel floors to bf16 (1 ulp), "
                "the reference rounds to nearest (1/2 ulp), and the rms scale "
                "is within 2**-12",
            ),
            # Per element: square-add, then one multiply.
            ops_per_call=3 * dh,
        ),
    )

flm_gemma4_decode_rms_residual

flm_gemma4_decode_rms_residual(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the decode layer's four RMS norms and two residual adds, from flm_gemma4/decode_rms_residual.cc.

rms_residual(y, x_ping, x_pong, y_final, w, x_temp_buf, is_swa, skip_kv): one layer's worth of norms. It normalizes the layer input into y for the QKV projection, then normalizes the attention output, adds the residual and normalizes the sum into y for the up/gate projection, and finally normalizes the MLP output and adds the second residual into y_final. The attention and MLP outputs arrive in the x ping-pong pair (model_dim bf16 each). y is model_dim + 16 bf16, a 16-element packet header then the row; w holds the four norm weights (4 * model_dim); x_temp_buf is 2 * model_dim of scratch. is_swa and skip_kv are int32 RTP buffers of 16 words that select how many projection repeats y is released for.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults w_prod_lock=0, w_cons_lock=1, y_prod_lock=2, y_cons_lock=3, x_prod_lock=4, x_cons_lock=5, rtp_available_lock=6, lm_head_out_prod_lock=7 and lm_head_out_cons_lock=8.	{}

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_rms_residual(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the decode layer's four RMS norms and two residual adds, from ``flm_gemma4/decode_rms_residual.cc``.

    ``rms_residual(y, x_ping, x_pong, y_final, w, x_temp_buf, is_swa, skip_kv)``:
    one layer's worth of norms. It normalizes the layer input into ``y`` for
    the QKV projection, then normalizes the attention output, adds the
    residual and normalizes the sum into ``y`` for the up/gate projection,
    and finally normalizes the MLP output and adds the second residual into
    ``y_final``. The attention and MLP outputs arrive in the ``x`` ping-pong
    pair (``model_dim`` bf16 each). ``y`` is ``model_dim + 16`` bf16, a
    16-element packet header then the row; ``w`` holds the four norm weights
    (``4 * model_dim``); ``x_temp_buf`` is ``2 * model_dim`` of scratch.
    ``is_swa`` and ``skip_kv`` are int32 RTP buffers of 16 words that select
    how many projection repeats ``y`` is released for.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``w_prod_lock=0``,
            ``w_cons_lock=1``, ``y_prod_lock=2``, ``y_cons_lock=3``,
            ``x_prod_lock=4``, ``x_cons_lock=5``, ``rtp_available_lock=6``,
            ``lm_head_out_prod_lock=7`` and ``lm_head_out_cons_lock=8``.
    """
    bf16 = np.dtype[bfloat16]
    d = geometry.model_dim
    x_ty = np.ndarray[(d,), bf16]
    return _decode_kernel(
        "rms_residual",
        "rms_residual",
        [
            np.ndarray[(d + 16,), bf16],
            x_ty,
            x_ty,
            x_ty,
            np.ndarray[(4, d), bf16],
            np.ndarray[(2, d), bf16],
            _DECODE_RTP,
            _DECODE_RTP,
        ],
        (Out, In, In, Out, In, Out, Param, Param),
        dict(
            w_prod_lock=0,
            w_cons_lock=1,
            y_prod_lock=2,
            y_cons_lock=3,
            x_prod_lock=4,
            x_cons_lock=5,
            rtp_available_lock=6,
            lm_head_out_prod_lock=7,
            lm_head_out_cons_lock=8,
        ),
        locks,
        geometry,
    )

flm_gemma4_rms_residual_core_ref

flm_gemma4_rms_residual_core_ref(x, w, residual)

Numpy reference for flm_gemma4_rms_residual_core.

residual + x * w / sqrt(mean(x^2) + 1e-6), rounded to bf16.

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

def flm_gemma4_rms_residual_core_ref(x, w, residual):
    """Numpy reference for [`flm_gemma4_rms_residual_core`][iron.kernels.flm_gemma4.flm_gemma4_rms_residual_core].

    ``residual + x * w / sqrt(mean(x^2) + 1e-6)``, rounded to bf16.
    """
    return (
        _rms_residual_core_terms(x, w, residual)[1].astype(np.float32).astype(bfloat16)
    )

flm_gemma4_rms_residual_core

flm_gemma4_rms_residual_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

rms_norm then residual_add of flm_gemma4_decode_rms_residual.

The post-attention norm and residual add. rms_residual_core(x, w, residual, y), all model_dim bf16: y = residual + rms_norm(x) * w.

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_rms_residual_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``rms_norm`` then ``residual_add`` of [`flm_gemma4_decode_rms_residual`][iron.kernels.flm_gemma4.flm_gemma4_decode_rms_residual].

    The post-attention norm and residual add. ``rms_residual_core(x, w,
    residual, y)``, all ``model_dim`` bf16: ``y = residual + rms_norm(x) * w``.
    """
    base = flm_gemma4_decode_rms_residual(geometry=geometry)
    row = np.ndarray[(geometry.model_dim,), _BF16]
    return _wrapper(
        base,
        "decode_rms_residual_core.cc",
        "rms_residual_core",
        [row, row, row, row],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, In, Out),
            reference=flm_gemma4_rms_residual_core_ref,
            tolerance=Tolerance.bounded(
                _rms_residual_core_bound,
                note="bf16 ulp <= 2**-7 |v|. The kernel floors the normalized "
                "value to bf16 (1 ulp, plus the fast inverse sqrt's and the "
                "float32 sum of squares' error, under 2**-12), floors the sum "
                "to bf16 (1 ulp, plus its float32 add's error, under 2**-12), "
                "and the reference rounds the sum to nearest (1/2 ulp). Each "
                "flush of a subnormal adds the smallest normal",
            ),
            # Per element: square-add and two multiplies, then the add.
            ops_per_call=5 * geometry.model_dim,
        ),
    )

flm_gemma4_decode_proj_main

flm_gemma4_decode_proj_main(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

One of the decode layer's 16 q4nx projection cores, from flm_gemma4/decode_proj_main.cc.

proj_main(y_ping, w_ping, x_ping, y_pong, w_pong, x_pong, is_swa, skip_kv, send_x_output): runs this core's share of every projection in a layer (QKV, output, up/gate and down), multiplying q4nx weight blocks from the w ping-pong pair into 256-element input slices from the x pair and writing each block's 32 outputs into the y pair. y_ping and y_pong are 2 * 32 + 16 bf16: a 16-element packet header, then two 32-output slots. With send_x_output set the core fills the first slot of its own pair and sends it; with it 0 the pair is the tile below's and the core fills the second slot. w_ping and w_pong are one q4nx block each (5120 bytes: 32 by 256 weights, their bf16 scales and mins, then the 4-bit codes). is_swa and skip_kv are int32 RTP buffers of 16 words that select the projection shapes.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults x_prod_lock=0, x_cons_lock=1, w_prod_lock=2, w_cons_lock=3, y_prod_ping_lock=4, y_prod_pong_lock=5, rtp_available_lock=6, y_cons_ping_lock=7 and y_cons_pong_lock=8. When send_x_output is 0 the four y locks are the tile below's.	{}

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_proj_main(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """One of the decode layer's 16 q4nx projection cores, from ``flm_gemma4/decode_proj_main.cc``.

    ``proj_main(y_ping, w_ping, x_ping, y_pong, w_pong, x_pong, is_swa, skip_kv, send_x_output)``:
    runs this core's share of every projection in a layer (QKV, output,
    up/gate and down), multiplying q4nx weight blocks from the ``w``
    ping-pong pair into 256-element input slices from the ``x`` pair and
    writing each block's 32 outputs into the ``y`` pair. ``y_ping`` and
    ``y_pong`` are ``2 * 32 + 16`` bf16: a 16-element packet header, then
    two 32-output slots. With ``send_x_output`` set the core fills the first
    slot of its own pair and sends it; with it 0 the pair is the tile
    below's and the core fills the second slot. ``w_ping`` and ``w_pong`` are one q4nx
    block each (5120 bytes: 32 by 256 weights, their bf16 scales and mins,
    then the 4-bit codes). ``is_swa`` and ``skip_kv`` are int32 RTP buffers
    of 16 words that select the projection shapes.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``x_prod_lock=0``,
            ``x_cons_lock=1``, ``w_prod_lock=2``, ``w_cons_lock=3``,
            ``y_prod_ping_lock=4``, ``y_prod_pong_lock=5``,
            ``rtp_available_lock=6``, ``y_cons_ping_lock=7`` and
            ``y_cons_pong_lock=8``. When ``send_x_output`` is 0 the four
            ``y`` locks are the tile below's.
    """
    bf16 = np.dtype[bfloat16]
    y_ty = np.ndarray[(2 * _DECODE_Q4NX_ROWS + 16,), bf16]
    # One q4nx block, as the bf16 words the weight streams move.
    w_ty = np.ndarray[(_DECODE_Q4NX_BLOCK_BYTES // 2,), bf16]
    x_ty = np.ndarray[(_DECODE_Q4NX_COLS,), bf16]
    return _decode_kernel(
        "proj_main",
        "proj_main",
        [
            y_ty,
            w_ty,
            x_ty,
            y_ty,
            w_ty,
            x_ty,
            _DECODE_RTP,
            _DECODE_RTP,
            np.int32,
        ],
        (Out, In, In, Out, In, In, Param, Param, Param),
        dict(
            x_prod_lock=0,
            x_cons_lock=1,
            w_prod_lock=2,
            w_cons_lock=3,
            y_prod_ping_lock=4,
            y_prod_pong_lock=5,
            rtp_available_lock=6,
            y_cons_ping_lock=7,
            y_cons_pong_lock=8,
        ),
        locks,
        geometry,
    )

flm_gemma4_decode_per_layer_up

flm_gemma4_decode_per_layer_up(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the per-layer-input up projection, from flm_gemma4/decode_per_layer_up.cc.

per_layer_up(x, proj_w_ping, proj_w_pong, y): gates the layer's per-layer input, projects it from pli_d up to model_dim with bf16 weight blocks (32 by 256 each) streamed through the proj_w ping-pong pair, RMS-normalizes the result, adds the residual and scales it by the layer scale into y (model_dim bf16). x is 2 * (pli_d + model_dim) + 32 bf16: the norm weight, the layer scale padded to 32, the per-layer input, the residual and the gate. The kernel gates the per-layer input in place.

y needs a 64-byte-aligned base pointer.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults x_prod_lock=0, x_cons_lock=1, proj_w_prod_lock=2, proj_w_cons_lock=3, y_prod_lock=4 and y_cons_lock=5.	{}

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_per_layer_up(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the per-layer-input up projection, from ``flm_gemma4/decode_per_layer_up.cc``.

    ``per_layer_up(x, proj_w_ping, proj_w_pong, y)``: gates the layer's
    per-layer input, projects it from ``pli_d`` up to ``model_dim`` with bf16
    weight blocks (32 by 256 each) streamed through the ``proj_w`` ping-pong
    pair, RMS-normalizes the result, adds the residual and scales it by the
    layer scale into ``y`` (``model_dim`` bf16). ``x`` is
    ``2 * (pli_d + model_dim) + 32`` bf16: the norm weight, the layer scale
    padded to 32, the per-layer input, the residual and the gate. The kernel
    gates the per-layer input in place.

    ``y`` needs a 64-byte-aligned base pointer.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``x_prod_lock=0``,
            ``x_cons_lock=1``, ``proj_w_prod_lock=2``, ``proj_w_cons_lock=3``,
            ``y_prod_lock=4`` and ``y_cons_lock=5``.
    """
    bf16 = np.dtype[bfloat16]
    d, pli = geometry.model_dim, geometry.pli_d
    w_ty = np.ndarray[(_DECODE_BF16_BLOCK,), bf16]
    return _decode_kernel(
        "per_layer_up",
        "per_layer_up",
        [
            np.ndarray[(2 * (pli + d) + 32,), bf16],
            w_ty,
            w_ty,
            np.ndarray[(d,), bf16],
        ],
        (InOut, In, In, Out),
        dict(
            x_prod_lock=0,
            x_cons_lock=1,
            proj_w_prod_lock=2,
            proj_w_cons_lock=3,
            y_prod_lock=4,
            y_cons_lock=5,
        ),
        locks,
        geometry,
    )

flm_gemma4_bf16_proj_core_ref

flm_gemma4_bf16_proj_core_ref(w, x)

Numpy reference for flm_gemma4_bf16_proj_core: one weight block times x.

The block is stored column by column: w[k * 32 + m] multiplies x[k] into out-feature m. Sums in float64.

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

def flm_gemma4_bf16_proj_core_ref(w, x):
    """Numpy reference for [`flm_gemma4_bf16_proj_core`][iron.kernels.flm_gemma4.flm_gemma4_bf16_proj_core]: one weight block times ``x``.

    The block is stored column by column: ``w[k * 32 + m]`` multiplies
    ``x[k]`` into out-feature ``m``. Sums in float64.
    """
    w = np.asarray(w, np.float64).reshape(-1, _BF16_PROJ_K, _BF16_PROJ_M)
    x = np.asarray(x, np.float64).reshape(-1, _BF16_PROJ_K)
    return np.einsum("ckm,ck->cm", w, x).astype(np.float32)

flm_gemma4_bf16_proj_core

flm_gemma4_bf16_proj_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

_mvm_bf16_bf16 of the per-layer-input projections, one block.

bf16_proj_block_core(w, x, y): zeroes a float32 accumulator in its own frame, accumulates one 32 by 256 bf16 weight block times 256 inputs into it, as linear_proj does, and copies it to y. The block does not depend on geometry.

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_bf16_proj_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``_mvm_bf16_bf16`` of the per-layer-input projections, one block.

    ``bf16_proj_block_core(w, x, y)``: zeroes a float32 accumulator in its own
    frame, accumulates one 32 by 256 bf16 weight block times 256 inputs into
    it, as ``linear_proj`` does, and copies it to ``y``. The block does not
    depend on ``geometry``.
    """
    base = flm_gemma4_decode_per_layer_up(geometry=geometry)
    return _wrapper(
        base,
        "decode_per_layer_up_core.cc",
        "bf16_proj_block_core",
        [
            np.ndarray[(_DECODE_BF16_BLOCK,), _BF16],
            np.ndarray[(_BF16_PROJ_K,), _BF16],
            np.ndarray[(_BF16_PROJ_M,), _F32],
        ],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, In, Out),
            reference=flm_gemma4_bf16_proj_core_ref,
            sample=_bf16_proj_core_sample,
            tolerance=Tolerance.exact(
                note="the sample keeps every product and sum exact in float32"
            ),
            ops_per_call=2 * _DECODE_BF16_BLOCK,
            acc_dtype=np.float32,
            reduction=_BF16_PROJ_K,
        ),
    )

flm_gemma4_decode_proj_layer_embedding

flm_gemma4_decode_proj_layer_embedding(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the per-layer-input embedding projection, from flm_gemma4/decode_proj_layer_embedding.cc.

proj_layer_embedding(norm_w, x0_per_layer, x0, x_proj, y, proj_w_ping, proj_w_pong): projects the layer input x0 (model_dim bf16) down to pli_d with bf16 weight blocks (32 by 256 each) streamed through the proj_w ping-pong pair, scales, RMS-normalizes and adds the token's per-layer embedding x0_per_layer (pli_d), and scales again, in the scratch x_proj (pli_d). norm_w (pli_d + model_dim + 32 bf16) holds the norm weight, then model_dim + 32 values the kernel copies to the start of y (same size) for the up projection; the result follows them.

x_proj needs a 64-byte-aligned base pointer.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults norm_w_prod_lock=0, x0_per_layer_prod_lock=1, x0_prod_lock=2, xw_cons_lock=3, proj_w_prod_lock=4, proj_w_cons_lock=5, y_prod_lock=6 and y_cons_lock=7.	{}

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_proj_layer_embedding(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the per-layer-input embedding projection, from ``flm_gemma4/decode_proj_layer_embedding.cc``.

    ``proj_layer_embedding(norm_w, x0_per_layer, x0, x_proj, y, proj_w_ping, proj_w_pong)``:
    projects the layer input ``x0`` (``model_dim`` bf16) down to ``pli_d``
    with bf16 weight blocks (32 by 256 each) streamed through the ``proj_w``
    ping-pong pair, scales, RMS-normalizes and adds the token's per-layer
    embedding ``x0_per_layer`` (``pli_d``), and scales again, in the scratch
    ``x_proj`` (``pli_d``). ``norm_w`` (``pli_d + model_dim + 32`` bf16)
    holds the norm weight, then ``model_dim + 32`` values the kernel copies
    to the start of ``y`` (same size) for the up projection; the result
    follows them.

    ``x_proj`` needs a 64-byte-aligned base pointer.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``norm_w_prod_lock=0``,
            ``x0_per_layer_prod_lock=1``, ``x0_prod_lock=2``,
            ``xw_cons_lock=3``, ``proj_w_prod_lock=4``,
            ``proj_w_cons_lock=5``, ``y_prod_lock=6`` and ``y_cons_lock=7``.
    """
    bf16 = np.dtype[bfloat16]
    d, pli = geometry.model_dim, geometry.pli_d
    pli_ty = np.ndarray[(pli,), bf16]
    y_ty = np.ndarray[(pli + d + 32,), bf16]
    w_ty = np.ndarray[(_DECODE_BF16_BLOCK,), bf16]
    return _decode_kernel(
        "proj_layer_embedding",
        "proj_layer_embedding",
        [
            y_ty,
            pli_ty,
            np.ndarray[(d,), bf16],
            pli_ty,
            y_ty,
            w_ty,
            w_ty,
        ],
        (In, In, In, Out, Out, In, In),
        dict(
            norm_w_prod_lock=0,
            x0_per_layer_prod_lock=1,
            x0_prod_lock=2,
            xw_cons_lock=3,
            proj_w_prod_lock=4,
            proj_w_cons_lock=5,
            y_prod_lock=6,
            y_cons_lock=7,
        ),
        locks,
        geometry,
    )

flm_gemma4_decode_gate_layer_embedding

flm_gemma4_decode_gate_layer_embedding(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE,
    **locks: int
) -> ExternalFunction

Build the per-layer-input gate, from flm_gemma4/decode_gate_layer_embedding.cc.

gate_layer_embedding(x, proj_w_ping, proj_w_pong, y): copies the layer output x (model_dim bf16, a buffer of the tile to the left) into y, projects it to pli_d with bf16 weight blocks (32 by 256 each) streamed through the proj_w ping-pong pair, and applies the activation into y + model_dim. y is model_dim + pli_d bf16.

y needs a 64-byte-aligned base pointer; the kernel also stores 512 bits at a time at y + model_dim.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE
**locks	int	Core lock ids overriding the defaults x_prod_lock=0, x_cons_lock=1, proj_w_prod_lock=2, proj_w_cons_lock=3, y_prod_lock=4, y_cons_lock=5, final_x_prod_lock=7 and final_x_cons_lock=8 (the last two locks of the tile to the left).	{}

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_decode_gate_layer_embedding(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE, **locks: int
) -> ExternalFunction:
    """Build the per-layer-input gate, from ``flm_gemma4/decode_gate_layer_embedding.cc``.

    ``gate_layer_embedding(x, proj_w_ping, proj_w_pong, y)``: copies the
    layer output ``x`` (``model_dim`` bf16, a buffer of the tile to the left)
    into ``y``, projects it to ``pli_d`` with bf16 weight blocks (32 by 256
    each) streamed through the ``proj_w`` ping-pong pair, and applies the
    activation into ``y + model_dim``. ``y`` is ``model_dim + pli_d`` bf16.

    ``y`` needs a 64-byte-aligned base pointer; the kernel also stores
    512 bits at a time at ``y + model_dim``.

    Args:
        geometry: The model the kernel builds for.
        **locks: Core lock ids overriding the defaults ``x_prod_lock=0``,
            ``x_cons_lock=1``, ``proj_w_prod_lock=2``, ``proj_w_cons_lock=3``,
            ``y_prod_lock=4``, ``y_cons_lock=5``, ``final_x_prod_lock=7`` and
            ``final_x_cons_lock=8`` (the last two locks of the tile to the
            left).
    """
    bf16 = np.dtype[bfloat16]
    d = geometry.model_dim
    w_ty = np.ndarray[(_DECODE_BF16_BLOCK,), bf16]
    return _decode_kernel(
        "gate_layer_embedding",
        "gate_layer_embedding",
        [
            np.ndarray[(d,), bf16],
            w_ty,
            w_ty,
            np.ndarray[(d + geometry.pli_d,), bf16],
        ],
        (In, In, In, Out),
        dict(
            x_prod_lock=0,
            x_cons_lock=1,
            proj_w_prod_lock=2,
            proj_w_cons_lock=3,
            y_prod_lock=4,
            y_cons_lock=5,
            final_x_prod_lock=7,
            final_x_cons_lock=8,
        ),
        locks,
        geometry,
        lut=True,
    )

flm_gemma4_pli_gelu_core_ref

flm_gemma4_pli_gelu_core_ref(x)

Numpy reference for flm_gemma4_pli_gelu_core: tanh-GELU.

The tolerance covers the kernel's table.

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

def flm_gemma4_pli_gelu_core_ref(x):
    """Numpy reference for [`flm_gemma4_pli_gelu_core`][iron.kernels.flm_gemma4.flm_gemma4_pli_gelu_core]: tanh-GELU.

    The tolerance covers the kernel's table.
    """
    return _gelu_tanh(x).astype(np.float32)

flm_gemma4_pli_gelu_core

flm_gemma4_pli_gelu_core(
    *,
    geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction

_activate of flm_gemma4_decode_gate_layer_embedding: GELU over pli_d values.

pli_gelu_core(x, y): copies x into y and applies GELU there, through aie_runtime_lib's table (getGeluBf16). Both are pli_d bf16.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry	The model the kernel builds for.	FLM_GEMMA4_E2B_DECODE

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

@dtypes([{"geometry": FLM_GEMMA4_E4B_DECODE}])
def flm_gemma4_pli_gelu_core(
    *, geometry: FlmGemma4DecodeGeometry = FLM_GEMMA4_E2B_DECODE
) -> ExternalFunction:
    """``_activate`` of [`flm_gemma4_decode_gate_layer_embedding`][iron.kernels.flm_gemma4.flm_gemma4_decode_gate_layer_embedding]: GELU over ``pli_d`` values.

    ``pli_gelu_core(x, y)``: copies ``x`` into ``y`` and applies GELU there,
    through aie_runtime_lib's table (``getGeluBf16``). Both are ``pli_d``
    bf16.

    Args:
        geometry: The model the kernel builds for.
    """
    n = geometry.pli_d
    return _wrapper(
        flm_gemma4_decode_gate_layer_embedding(geometry=geometry),
        "decode_gate_layer_embedding_core.cc",
        "pli_gelu_core",
        [np.ndarray[(n,), _BF16], np.ndarray[(n,), _BF16]],
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=flm_gemma4_pli_gelu_core_ref,
            sample=partial(_pli_gelu_core_sample, n=n),
            tolerance=Tolerance.bounded(
                _pli_gelu_core_bound,
                note="the GELU table's error against tanh-GELU, measured on the "
                "host from aie_runtime_lib's table, plus a bf16 ulp of its slope "
                "and the floor bf16 store",
            ),
            uses_lut=True,
            # Per value: the table line's multiply and add.
            ops_per_call=2 * n,
        ),
    )

flm_gemma4_q4nx_lm_head_ref

flm_gemma4_q4nx_lm_head_ref(
    w, x, sums, *, m_tile=32, k_tile=256, group=32
)

Numpy reference for flm_gemma4_q4nx_lm_head: one q4nx_lm_head_block call.

w is one q4nx block per call, as bytes or bf16 words: bf16 scales, then bf16 minima, both indexed [k // group, m], then 4-bit codes, low nibble first, indexed [m // 16, k // 32, (k % 32) // 8, k % 8, m % 16]. x is the token then its RMS weight, and sums the token's per-32 column sums, both shared by every call; slice 0 of the token is the one the call reads. Returns the float32 accumulator, zero before the call, holding sum_k (min + scale * code) * x with the minima folded in through sums.

Per 16 rows and 32 columns the kernel sums the codes times x in float32, narrows that dot product to bf16 with the floor rounding a fresh core uses, then accumulates it times the scale, and the minimum times the column sum, into float32. Every product is exact, so each multiply-accumulate rounds once.

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

def flm_gemma4_q4nx_lm_head_ref(w, x, sums, *, m_tile=32, k_tile=256, group=32):
    """Numpy reference for [`flm_gemma4_q4nx_lm_head`][iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head]: one ``q4nx_lm_head_block`` call.

    ``w`` is one q4nx block per call, as bytes or bf16 words: bf16 scales,
    then bf16 minima, both indexed ``[k // group, m]``, then 4-bit codes, low
    nibble first, indexed ``[m // 16, k // 32, (k % 32) // 8, k % 8, m % 16]``.
    ``x`` is the token then its RMS weight, and ``sums`` the token's per-32
    column sums, both shared by every call; slice 0 of the token is the one
    the call reads. Returns the float32
    accumulator, zero before the call, holding ``sum_k (min + scale * code) *
    x`` with the minima folded in through ``sums``.

    Per 16 rows and 32 columns the kernel sums the codes times ``x`` in
    float32, narrows that dot product to bf16 with the floor rounding a fresh
    core uses, then accumulates it times the scale, and the minimum times the
    column sum, into float32. Every product is exact, so each
    multiply-accumulate rounds once.
    """
    w = np.ascontiguousarray(w).view(np.uint8).reshape(-1, m_tile * k_tile * 5 // 8)
    calls, n_groups = len(w), k_tile // group
    params = np.ascontiguousarray(w[:, : 4 * m_tile * n_groups]).view("<u2")
    params = (params.astype(np.uint32) << 16).view(np.float32)
    scales, mins = params.reshape(calls, 2, n_groups, m_tile).transpose(1, 0, 2, 3)
    packed = w[:, 4 * m_tile * n_groups :]
    codes = np.empty((calls, m_tile * k_tile), np.float64)
    codes[:, 0::2], codes[:, 1::2] = packed & 15, packed >> 4
    codes = codes.reshape(calls, m_tile // 16, n_groups, 32, 16)
    x = np.asarray(x, np.float32).reshape(-1)[:k_tile]
    x = x.astype(np.float64).reshape(1, n_groups, 32)
    sums = np.asarray(sums, np.float32).reshape(-1)[:n_groups]
    sums = sums.astype(np.float64).reshape(1, n_groups)
    acc = np.zeros((calls, m_tile // 16, 16), np.float32)
    for g in range(n_groups):
        dot = np.zeros((calls, m_tile // 16, 16), np.float32)
        for c in range(32):
            dot = _f32(dot + codes[:, :, g, c, :] * x[:, g, c, None, None])
        dot = _bf16_floor(dot).astype(np.float64)
        scale = scales[:, g].reshape(calls, m_tile // 16, 16).astype(np.float64)
        low = mins[:, g].reshape(calls, m_tile // 16, 16).astype(np.float64)
        acc = _f32(acc + dot * scale)
        acc = _f32(acc + low * sums[:, g, None, None])
    return acc.reshape(calls, m_tile)

flm_gemma4_q4nx_lm_head

flm_gemma4_q4nx_lm_head(
    *,
    dim: int = 1536,
    m_tile: int = 32,
    k_tile: int = 256,
    group: int = 32
) -> ExternalFunction

AIE2P logits from a q4nx vocabulary, from flm_gemma4/q4nx_lm_head.cc.

One core owns m_tile out-features at a time and streams their q4nx blocks past a token of dim values. The returned kernel is q4nx_lm_head_block, which accumulates one m_tile by k_tile block times one slice of the token into float32; its slice index is bound to 0. See flm_gemma4_q4nx_lm_head_ref for the block layout.

The other entry points, attributes of the kernel:

• q4nx_lm_head_rms(x, sums): once per token; normalizes the token in x (the token, then its RMS weight) in place and writes its per-32 column sums.
• q4nx_lm_head_zero(y_acc): before a tile's k loop.
• q4nx_lm_head_epilogue(y, y_acc, softcap): after the k loop, y = c * tanh(y_acc / c), with the float32 c in RTP word 0.

Parameters:

Name	Type	Description	Default
dim	int	The token's length, a multiple of k_tile; sets the RMS norm's span.	1536
m_tile	int	Out-features per block, a multiple of 16.	32
k_tile	int	In-features per block, a multiple of group.	256
group	int	In-features per scale and minimum; must be 32.	32

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

def flm_gemma4_q4nx_lm_head(
    *, dim: int = 1536, m_tile: int = 32, k_tile: int = 256, group: int = 32
) -> ExternalFunction:
    """AIE2P logits from a q4nx vocabulary, from ``flm_gemma4/q4nx_lm_head.cc``.

    One core owns ``m_tile`` out-features at a time and streams their q4nx
    blocks past a token of ``dim`` values. The returned kernel is
    ``q4nx_lm_head_block``, which accumulates one ``m_tile`` by ``k_tile``
    block times one slice of the token into float32; its slice index is bound
    to 0. See ``flm_gemma4_q4nx_lm_head_ref`` for the block layout.

    The other entry points, attributes of the kernel:

    - ``q4nx_lm_head_rms(x, sums)``: once per token; normalizes the token in
      ``x`` (the token, then its RMS weight) in place and writes its per-32
      column sums.
    - ``q4nx_lm_head_zero(y_acc)``: before a tile's k loop.
    - ``q4nx_lm_head_epilogue(y, y_acc, softcap)``: after the k loop,
      ``y = c * tanh(y_acc / c)``, with the float32 ``c`` in RTP word 0.

    Args:
        dim: The token's length, a multiple of ``k_tile``; sets the RMS
            norm's span.
        m_tile: Out-features per block, a multiple of 16.
        k_tile: In-features per block, a multiple of ``group``.
        group: In-features per scale and minimum; must be 32.
    """
    _lm_head_geometry(dim, m_tile, k_tile, group)
    if not _arch_traits().bfp16:
        raise NotImplementedError(
            "flm_gemma4_q4nx_lm_head() is only available on aie2p."
        )
    geometry = dict(m_tile=m_tile, k_tile=k_tile, group=group)
    x = np.ndarray[(2, dim), _BF16]
    sums = np.ndarray[(dim // group,), _BF16]
    y_acc = np.ndarray[(m_tile,), _F32]
    fn = _make_extern(
        "q4nx_lm_head_block",
        _kernel_source("flm_gemma4/q4nx_lm_head.cc"),
        [np.ndarray[(m_tile * k_tile * 5 // 8 // 2,), _BF16], x, y_acc, sums, np.int32],
        compile_flags=[
            f"-DQ4NX_M_TILE={m_tile}",
            f"-DQ4NX_K_TILE={k_tile}",
            f"-DQ4NX_GROUP={group}",
            f"-DFLM_GEMMA4_LM_HEAD_DIM={dim}",
            # Without it the bf16 mmul emulation runs about 8x slower.
            "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16",
        ],
        cls=_LmHeadKernel,
        contract=KernelContract(
            trace=Trace.whole_call(),
            # The token and its sums are fixed while the blocks stream past.
            roles=(In, Param, InOut, Param, Param),
            parameter_bindings=((4, 0),),
            initializers=((2, _lm_head_zero),),
            reference=partial(flm_gemma4_q4nx_lm_head_ref, **geometry),
            sample=partial(_q4nx_lm_head_sample, dim=dim, **geometry),
            tolerance=Tolerance.exact(
                note="the sample keeps every sum exact in float32; the floor "
                "bf16 narrowing of each dot product is modeled"
            ),
            ops_per_call=2 * m_tile * k_tile,
            acc_dtype=np.float32,
            reduction=k_tile,
        ),
    )
    bind = fn.object_file.bind
    fn.q4nx_lm_head_block = fn
    fn.q4nx_lm_head_rms = bind("q4nx_lm_head_rms", [x, sums])
    fn.q4nx_lm_head_zero = bind("q4nx_lm_head_zero", [y_acc])
    rtp = np.ndarray[(_LM_HEAD_RTP_WORDS,), _RTP_WORD]
    fn.q4nx_lm_head_epilogue = bind(
        "q4nx_lm_head_epilogue", [np.ndarray[(m_tile,), _BF16], y_acc, rtp]
    )
    return fn

flm_gemma4_q4nx_lm_head_epilogue_ref

flm_gemma4_q4nx_lm_head_epilogue_ref(y_acc, rtp)

Numpy reference for flm_gemma4_q4nx_lm_head_epilogue: c * tanh(y_acc / c).

rtp holds the float32 bits of c in word 0. The reference computes in float64 and rounds to bf16 once; the tolerance covers the kernel's narrowings and its tanh.

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

def flm_gemma4_q4nx_lm_head_epilogue_ref(y_acc, rtp):
    """Numpy reference for [`flm_gemma4_q4nx_lm_head_epilogue`][iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head_epilogue]: ``c * tanh(y_acc / c)``.

    ``rtp`` holds the float32 bits of ``c`` in word 0. The reference computes
    in float64 and rounds to bf16 once; the tolerance covers the kernel's
    narrowings and its tanh.
    """
    c = _lm_head_softcap(rtp)
    y = np.asarray(y_acc, np.float64)
    return (c * np.tanh(y / c)).astype(bfloat16)

flm_gemma4_q4nx_lm_head_epilogue

flm_gemma4_q4nx_lm_head_epilogue(
    *,
    dim: int = 1536,
    m_tile: int = 32,
    k_tile: int = 256,
    group: int = 32
) -> ExternalFunction

q4nx_lm_head_epilogue of the LM head: y = c * tanh(y_acc / c).

The softcap c is float32 bits in word 0 of a 32-word RTP buffer. The tolerance assumes c is a bf16 value, so the kernel's narrowing of it is exact; Gemma 4's 30 is.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_epilogue(
    *, dim: int = 1536, m_tile: int = 32, k_tile: int = 256, group: int = 32
) -> ExternalFunction:
    """``q4nx_lm_head_epilogue`` of the LM head: ``y = c * tanh(y_acc / c)``.

    The softcap ``c`` is float32 bits in word 0 of a 32-word RTP buffer. The
    tolerance assumes ``c`` is a bf16 value, so the kernel's narrowing of it
    is exact; Gemma 4's 30 is.
    """
    return _lm_head_sibling(
        "q4nx_lm_head_epilogue",
        KernelContract(
            trace=Trace.whole_call(),
            roles=(Out, In, Param),
            reference=flm_gemma4_q4nx_lm_head_epilogue_ref,
            sample=partial(_lm_head_epilogue_sample, m_tile=m_tile),
            tolerance=Tolerance.bounded(
                _lm_head_epilogue_bound,
                note="c times vtanh's error, measured on npu2 (see "
                "activation._vtanh_error), plus the floor bf16 roundings",
            ),
        ),
        dim=dim,
        m_tile=m_tile,
        k_tile=k_tile,
        group=group,
    )

flm_gemma4_q4nx_lm_head_rms_ref

flm_gemma4_q4nx_lm_head_rms_ref(x, *, dim=1536)

Numpy reference for flm_gemma4_q4nx_lm_head_rms: column sums of the normalized token.

x holds, per call, the token then its RMS weight. The reference narrows each normalized value to bf16 and sums each 32 of them in float64.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_rms_ref(x, *, dim=1536):
    """Numpy reference for [`flm_gemma4_q4nx_lm_head_rms`][iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head_rms]: column sums of the normalized token.

    ``x`` holds, per call, the token then its RMS weight. The reference
    narrows each normalized value to bf16 and sums each 32 of them in float64.
    """
    y = _lm_head_rms_y(x, dim).astype(np.float32).astype(bfloat16)
    sums = y.astype(np.float64).reshape(len(y), dim // 32, 32).sum(axis=2)
    return sums.astype(np.float32)

flm_gemma4_q4nx_lm_head_rms

flm_gemma4_q4nx_lm_head_rms(
    *,
    dim: int = 1536,
    m_tile: int = 32,
    k_tile: int = 256,
    group: int = 32
) -> ExternalFunction

q4nx_lm_head_rms of flm_gemma4_q4nx_lm_head: the token's RMS norm and per-32 column sums.

The kernel normalizes the token in place. The contract declares the token In: the write lands in the core's input element, and the next DMA fill overwrites it. The column sums depend on every normalized value.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_q4nx_lm_head_rms(
    *, dim: int = 1536, m_tile: int = 32, k_tile: int = 256, group: int = 32
) -> ExternalFunction:
    """``q4nx_lm_head_rms`` of [`flm_gemma4_q4nx_lm_head`][iron.kernels.flm_gemma4.flm_gemma4_q4nx_lm_head]: the token's RMS norm and per-32 column sums.

    The kernel normalizes the token in place. The contract declares the token
    ``In``: the write lands in the core's input element, and the next DMA fill
    overwrites it. The column sums depend on every normalized value.
    """
    return _lm_head_sibling(
        "q4nx_lm_head_rms",
        KernelContract(
            trace=Trace.whole_call(),
            roles=(In, Out),
            reference=partial(flm_gemma4_q4nx_lm_head_rms_ref, dim=dim),
            tolerance=Tolerance.bounded(
                partial(_lm_head_rms_bound, dim=dim),
                note="bf16 ulp <= 2**-7 |v|. Per term: the kernel floors to bf16 "
                "(1 ulp), the reference rounds to nearest (1/2 ulp), and the fast "
                "inverse sqrt after two Newton steps adds 5e-6 relative. The "
                "kernel floors the sum to bf16 (1 ulp of the sum). 2.5 ulps of "
                "the sum of |y| plus that 5e-6 stays under 3 * 2**-7 of it",
            ),
            # Per element: square and add, two multiplies, one column-sum add.
            ops_per_call=5 * dim,
        ),
        dim=dim,
        m_tile=m_tile,
        k_tile=k_tile,
        group=group,
    )

flm_gemma4_attn_prefill_ref

flm_gemma4_attn_prefill_ref(l_bf16, y)

Numpy reference for flm_gemma4_attn_prefill: attn_epilogue at chunk 0.

y is the round's (8, 512) float32 accumulator, flat; chunk 0 is its first 64 values, eight per row of l_bf16. The kernel narrows y to bf16, multiplies by the row's 1 / l exactly and narrows again, both times with the floor rounding a fresh core uses.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_attn_prefill_ref(l_bf16, y):
    """Numpy reference for [`flm_gemma4_attn_prefill`][iron.kernels.flm_gemma4.flm_gemma4_attn_prefill]: ``attn_epilogue`` at chunk 0.

    ``y`` is the round's (8, 512) float32 accumulator, flat; chunk 0 is its
    first 64 values, eight per row of ``l_bf16``. The kernel narrows ``y`` to
    bf16, multiplies by the row's ``1 / l`` exactly and narrows again, both
    times with the floor rounding a fresh core uses.
    """
    return _prefill_epilogue(
        l_bf16, y, lq=_ATTN_PREFILL_LQ, dh=_ATTN_PREFILL_DH, chunk=0
    )

flm_gemma4_swa_prefill_ref

flm_gemma4_swa_prefill_ref(l_bf16, y)

Numpy reference for flm_gemma4_swa_prefill: attn_epilogue at chunk 33.

As flm_gemma4_attn_prefill_ref, over a (16, 256) y: rows 8 to 15, columns 64 to 127, each row scaled by its own 1 / l.

Source code in python/iron/kernels/flm_gemma4.py

def flm_gemma4_swa_prefill_ref(l_bf16, y):
    """Numpy reference for [`flm_gemma4_swa_prefill`][iron.kernels.flm_gemma4.flm_gemma4_swa_prefill]: ``attn_epilogue`` at chunk 33.

    As [`flm_gemma4_attn_prefill_ref`][iron.kernels.flm_gemma4.flm_gemma4_attn_prefill_ref], over a
    (16, 256) ``y``: rows 8 to 15, columns 64 to 127, each row scaled by its
    own ``1 / l``.
    """
    return _prefill_epilogue(
        l_bf16, y, lq=_SWA_PREFILL_LQ, dh=_SWA_PREFILL_DH, chunk=_SWA_PREFILL_CHUNK
    )