Skip to content

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.

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.

 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.
Source code in python/iron/kernels/eltwise.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
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.
        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``.
    """
    bit_width = _dtype_to_bit_width(dtype, factory_name="passthrough")
    tile_ty = np.ndarray[(tile_size,), np.dtype[dtype]]
    return _make_extern(
        "passThroughLine",
        _default_source_path("passThrough.cc"),
        [tile_ty, tile_ty, np.int32],
        compile_flags=[f"-DBIT_WIDTH={bit_width}"],
    )

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 
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
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}")

    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}",
        _default_source_path("scale.cc"),
        [tile_ty, tile_ty, scalar_ty, np.int32],
        compile_flags=[f"-DBIT_WIDTH={bit_width}"],
        use_chess=use_chess,
    )

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 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
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 
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
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)

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 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
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",
        _default_source_path("relu.cc"),
        [tile_ty, tile_ty],
    )

Reduction

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

reduce_add 

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

Reduction kernel: sums all elements of a tile to a scalar.

Parameters:

Name Type Description Default
tile_size int

Number of elements in the input tile.

 1024
dtype type

Element data type (only np.int32 supported).

 int32
vectorized bool

If True use vectorized path; False selects scalar.

 True

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the reduce_add kernel.

Raises:

Type Description
ValueError

When dtype is not np.int32.
Source code in python/iron/kernels/reduce.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def reduce_add(
    tile_size: int = 1024, dtype: type = np.int32, vectorized: bool = True
) -> ExternalFunction:
    """Reduction kernel: sums all elements of a tile to a scalar.

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

    Returns:
        ExternalFunction configured for the reduce_add kernel.

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

reduce_min 

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

Reduction kernel: finds the minimum element of a tile.

Parameters:

Name Type Description Default
tile_size int

Number of elements in the input tile.

 1024
dtype type

Element data type (only np.int32 supported).

 int32
vectorized bool

If True use vectorized path; False selects scalar.

 True

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the reduce_min kernel.

Raises:

Type Description
ValueError

When dtype is not np.int32.
Source code in python/iron/kernels/reduce.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def reduce_min(
    tile_size: int = 1024, dtype: type = np.int32, vectorized: bool = True
) -> ExternalFunction:
    """Reduction kernel: finds the minimum element of a tile.

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

    Returns:
        ExternalFunction configured for the reduce_min kernel.

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

reduce_max 

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

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

Parameters:

Name Type Description Default
tile_size int

Number of elements in the input tile.

 1024
dtype type

Element data type (np.int32 or bfloat16).

 int32
vectorized bool

If True use vectorized path; False selects scalar.

 True

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the reduce_max kernel.

Raises:

Type Description
ValueError

When dtype is not np.int32 or bfloat16.
Source code in python/iron/kernels/reduce.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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
    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}",
        _default_source_path("reduce_max.cc"),
        [in_ty, out_ty, np.int32],
        shared_object_file_name=_REDUCE_MAX_OBJ,
    )

compute_max 

compute_max(dtype: type = int32) -> ExternalFunction

Pairwise scalar max — companion to reduce_max 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; sharing the output .o (via shared_object_file_name) means both factories in the same design compile the source exactly once.

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 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
def compute_max(dtype: type = np.int32) -> ExternalFunction:
    """Pairwise scalar max — companion to [`reduce_max`][iron.kernels.reduce.reduce_max] 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]; sharing the
    output ``.o`` (via ``shared_object_file_name``) means both factories
    in the same design compile the source exactly once.

    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}",
        _default_source_path("reduce_max.cc"),
        [out_ty, out_ty, out_ty],
        shared_object_file_name=_REDUCE_MAX_OBJ,
    )

Linear algebra

Linear algebra kernel factories: mm, mv, cascade_mm.

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

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

The compiled .o exports both the matmul_* and zero_* symbols. Use kernels.mm(...).zero to get a sibling Kernel binding the zero symbol against the same .o, suitable for accumulator initialization.

Parameters:

Name Type Description Default
dim_m int

Number of rows of A / C.

 64
dim_k int

Number of columns of A / rows of B.

 64
dim_n int

Number of columns of B / C.

 64
input_dtype type

Input element type (np.int8, np.int16, or bfloat16).

 int16
output_dtype type

Output element type.

 int16
vectorized bool

If True use the vectorized variant.

 True
b_col_maj bool

If True compile with -DB_COL_MAJ so the kernel consumes B laid out column-major. Must agree with the design's B dims_to_stream.

 False
c_col_maj bool

If True compile with -DC_COL_MAJ so the kernel writes C laid out column-major. Must agree with the design's C output dims_to_stream.

 False
use_chess bool

If True build with xchesscc_wrapper instead of Peano's clang++. All ExternalFunctions in a single @iron.jit design must share the same toolchain.

 False
emulate_bf16_mmul_with_bfp16 bool

AIE2P only, bf16 inputs only. When True compile with -DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16 so the kernel uses BFP16-based emulation of the bf16 MMUL. Changes the micro-kernel dims to (8, 8, 8); designs reading .mac_dims will see the new geometry automatically. Ignored for non-bf16 inputs and on AIE2.

 False

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the matmul kernel.

Raises:

Type Description
ValueError

When (input_dtype, output_dtype) is not a supported combination.
Source code in python/iron/kernels/linalg.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
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,
) -> ExternalFunction:
    """Matrix-multiply kernel: C += A * B.

    The compiled ``.o`` exports both the ``matmul_*`` and ``zero_*`` symbols.
    Use ``kernels.mm(...).zero`` to get a sibling Kernel binding the zero
    symbol against the same .o, suitable for accumulator initialization.

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

    Returns:
        ExternalFunction configured for the matmul kernel.

    Raises:
        ValueError: When ``(input_dtype, output_dtype)`` is not a supported combination.
    """
    key = (input_dtype, output_dtype)
    if key not in _MM_COMBOS:
        raise ValueError(
            f"mm(): unsupported (input_dtype, output_dtype) = {key}. "
            f"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 == "aie2p" and input_dtype is bfloat16
    )
    if bf16_emulated:
        compile_flags.append("-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16")
    extern = _make_extern(
        f"{prefix}_{suffix}",
        _default_source_path("mm.cc"),
        [a_ty, b_ty, c_ty],
        compile_flags=compile_flags,
        use_chess=use_chess,
    )
    if bf16_emulated:
        extern.mac_dims = _MM_EMULATED_BF16_MAC_DIMS_AIE2P[key]
    else:
        extern.mac_dims = _MM_MAC_DIMS[arch][key]
    # mm.cc emits both matmul_* and zero_* symbols; expose the zero binding
    # as a sibling Kernel pointing at the same .o so the design does
    # `matmul = kernels.mm(...); zero = matmul.zero` instead of a separate
    # kernels.mm_zero call (which would compile mm.cc a second time).
    zero_prefix = "zero" if vectorized else "zero_scalar"
    extern.zero = Kernel(
        f"{zero_prefix}_{_ZERO_SUFFIX[output_dtype]}",
        extern.object_file_name,
        [c_ty],
    )
    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,
) -> ExternalFunction

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

Parameters:

Name Type Description Default
dim_m int

Number of rows of A (output vector length).

 32
dim_k int

Number of columns of A (input vector length).

 32
input_dtype type

Input element type. Only np.int16 is supported.

 int16
output_dtype type

Output element type. Only np.int32 is supported.

 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

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 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
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,
) -> ExternalFunction:
    """Matrix-vector multiply kernel: c += A * b.

    Args:
        dim_m: Number of rows of A (output vector length).
        dim_k: Number of columns of A (input vector length).
        input_dtype: Input element type. Only ``np.int16`` is supported.
        output_dtype: Output element type. Only ``np.int32`` is supported.
        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).

    Returns:
        ExternalFunction configured for the matvec kernel.

    Raises:
        ValueError: When the dtype combination is not supported.
    """
    if input_dtype != np.int16 or output_dtype != np.int32:
        raise ValueError(
            f"mv(): only (np.int16, np.int32) is supported, "
            f"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]]
    extern = _make_extern(
        f"{prefix}_i16_i32",
        _default_source_path("mv.cc"),
        [a_ty, b_ty, c_ty],
        compile_flags=[f"-DDIM_M={dim_m}", f"-DDIM_K={dim_k}"],
        use_chess=use_chess,
    )
    # mv.cc emits both matvec_* and zero_* symbols; expose the zero binding
    # as a sibling Kernel pointing at the same .o.
    zero_prefix = "zero_vectorized" if vectorized else "zero_scalar"
    extern.zero = Kernel(f"{zero_prefix}_i32", extern.object_file_name, [c_ty])
    return extern

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

Cascade matrix-multiply kernel for multi-core accumulation.

cascade_mm.cc emits all three cascade variants (get_only, put_only, put_get) plus a zero companion in one .o. The returned ExternalFunction binds the get_only symbol; the other three are sibling Kernel\s available as attributes:

  • .get_only — same as the returned EF (top of the cascade chain).
  • .put_only — bottom of the chain.
  • .put_get — middle of the chain.
  • .zero — accumulator initializer.

Designs typically use all four together, one per row of compute cores.

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 
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
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,
) -> ExternalFunction:
    """Cascade matrix-multiply kernel for multi-core accumulation.

    cascade_mm.cc emits all three cascade variants (``get_only``,
    ``put_only``, ``put_get``) plus a ``zero`` companion in one .o.  The
    returned ExternalFunction binds the ``get_only`` symbol; the other
    three are sibling [`Kernel`][iron.Kernel]\\s available as attributes:

    * ``.get_only`` — same as the returned EF (top of the cascade chain).
    * ``.put_only`` — bottom of the chain.
    * ``.put_get`` — middle of the chain.
    * ``.zero`` — accumulator initializer.

    Designs typically use all four together, one per row of compute cores.

    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}. "
            f"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]]
    extern = _make_extern(
        f"matmul_scalar_cascade_get_only_{suffix}",
        _default_source_path("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,
    )
    extern.get_only = extern
    extern.put_only = Kernel(
        f"matmul_scalar_cascade_put_only_{suffix}",
        extern.object_file_name,
        [a_ty, b_ty, c_ty],
    )
    extern.put_get = Kernel(
        f"matmul_scalar_cascade_put_get_{suffix}",
        extern.object_file_name,
        [a_ty, b_ty, c_ty],
    )
    extern.zero = Kernel(
        f"zero_scalar_{_ZERO_SUFFIX[output_dtype]}",
        extern.object_file_name,
        [c_ty],
    )
    arch = _detect_arch()
    if arch not in _CASCADE_MM_MAC_DIMS:
        raise ValueError(
            f"cascade_mm(): unsupported arch {arch!r}; "
            f"cascade_mm.cc only ships for {sorted(_CASCADE_MM_MAC_DIMS)}."
        )
    extern.mac_dims = _CASCADE_MM_MAC_DIMS[arch][key]
    return extern

Convolution

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

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 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
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,
        _default_source_path("conv2dk1.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(4)],
        compile_flags=flags,
    )

conv2dk3 

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

3x3 convolution kernel.

Parameters:

Name Type Description Default
input_width int

Spatial width of the input.

 32
input_channels int

Number of input channels.

 64
output_channels int

Number of output channels produced by this call.

 64
act_dtype type

Activation data type (np.int8 or np.uint8).

 int8
weight_output_channels int | None

Total number of output channels stored in the weights buffer. Defaults to output_channels. Set higher than output_channels when the weights buffer is shared across multiple workers that each produce a slice of the output (the channel_offset runtime arg selects a worker's slice).

 None

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the conv2dk3 kernel.

Raises:

Type Description
ValueError

When act_dtype is not np.int8 or np.uint8.
Source code in python/iron/kernels/conv.py 
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def conv2dk3(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
    act_dtype: type = np.int8,
    weight_output_channels: int | None = None,
) -> ExternalFunction:
    """3x3 convolution kernel.

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

    Returns:
        ExternalFunction configured for the conv2dk3 kernel.

    Raises:
        ValueError: When ``act_dtype`` is not ``np.int8`` or ``np.uint8``.
    """
    func_name, flags = _conv_act_dtype_info(
        "conv2dk3", act_dtype, factory_name="conv2dk3"
    )
    if weight_output_channels is None:
        weight_output_channels = output_channels
    line_size = input_width * input_channels
    line_ty = np.ndarray[(line_size,), np.dtype[act_dtype]]
    wt_ty = np.ndarray[
        (3 * 3 * input_channels * weight_output_channels,), np.dtype[np.int8]
    ]
    out_ty = np.ndarray[(input_width * output_channels,), np.dtype[np.uint8]]
    return _make_extern(
        func_name,
        _default_source_path("conv2dk3.cc"),
        [line_ty, line_ty, line_ty, wt_ty, out_ty, *_i32s(8)],
        compile_flags=flags,
    )

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.
Source code in python/iron/kernels/conv.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
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``.
    """
    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,
        _default_source_path("conv2dk1_skip.cc", subdir="aie2"),
        [in0_ty, in1_ty, wt_ty, out_ty, skip_ty, *_i32s(5)],
        compile_flags=flags,
    )

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 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
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",
        _default_source_path("conv2dk1_i8.cc"),
        [in_ty, wt_ty, out_ty, *_i32s(4)],
        compile_flags=["-DINT8_ACT"],
    )

conv2dk14 

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

14x14 convolution kernel (aie2p only).

Parameters:

Name Type Description Default
input_width int

Spatial width of the input.

 224
input_channels int

Number of input channels.

 16
output_channels int

Number of output channels.

 16
kernel_width int

Width (and height) of the convolution kernel.

 14

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the conv2dk14 kernel.
Source code in python/iron/kernels/conv.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
def conv2dk14(
    input_width: int = 224,
    input_channels: int = 16,
    output_channels: int = 16,
    kernel_width: int = 14,
) -> ExternalFunction:
    """14x14 convolution kernel (aie2p only).

    Args:
        input_width: Spatial width of the input.
        input_channels: Number of input channels.
        output_channels: Number of output channels.
        kernel_width: Width (and height) of the convolution kernel.

    Returns:
        ExternalFunction configured for the conv2dk14 kernel.
    """
    tiles = input_width // kernel_width
    pixels = kernel_width * kernel_width
    _RGBA = 4
    _ACC_FACTOR = 8
    in_ty = np.ndarray[(tiles * pixels * _RGBA,), np.dtype[np.uint8]]
    wt_ty = np.ndarray[(output_channels * pixels * _RGBA,), np.dtype[np.int8]]
    out_ty = np.ndarray[(output_channels * tiles * _ACC_FACTOR,), np.dtype[np.int8]]
    return _make_extern(
        "conv2dk14_i8",
        _default_source_path("conv2dk14.cc", subdir="aie2p"),
        [in_ty, wt_ty, out_ty, *_i32s(5)],
    )

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.
Source code in python/iron/kernels/conv.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
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``.
    """
    func_name, flags = _conv_act_dtype_info(
        "conv2dk1_skip_init", act_dtype, factory_name="conv2dk1_skip_init"
    )
    if skip_input_channels is None:
        skip_input_channels = input_channels
    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,
        _default_source_path("conv2dk1_skip_init.cc", subdir="aie2"),
        [in0_ty, in1_ty, wt_ty, out_ty, skip_ty, *_i32s(7)],
        compile_flags=flags,
    )

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 
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
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",
        _default_source_path("bottleneck/bn_conv2dk1_relu.cc", subdir="aie2"),
        [in_ty, wt_ty, out_ty, *_i32s(4)],
        compile_flags=["-DREGULAR", "-DINT8_ACT"],
    )

bn_conv2dk3 

bn_conv2dk3(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
) -> 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.

 64

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the bn_conv2dk3 kernel.
Source code in python/iron/kernels/conv.py 
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
def bn_conv2dk3(
    input_width: int = 32,
    input_channels: int = 64,
    output_channels: int = 64,
) -> 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.

    Returns:
        ExternalFunction configured for the bn_conv2dk3 kernel.
    """
    line_size = input_width * input_channels
    line_ty = np.ndarray[(line_size,), np.dtype[np.int8]]
    wt_ty = np.ndarray[(3 * 3 * input_channels * 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",
        _default_source_path("bottleneck/bn_conv2dk3.cc", subdir="aie2"),
        [line_ty, line_ty, line_ty, wt_ty, out_ty, *_i32s(8)],
    )

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 
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
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",
        _default_source_path("bottleneck/bn_conv2dk1_i8.cc", subdir="aie2"),
        [in_ty, wt_ty, out_ty, *_i32s(4)],
        compile_flags=["-DREGULAR", "-DSCALAR"],
    )

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 
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
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,
        _default_source_path("bottleneck/bn_conv2dk1_skip.cc", subdir="aie2"),
        [in_ty, wt_ty, out_ty, skip_ty, *_i32s(5)],
        compile_flags=flags,
    )

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 
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
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,
        _default_source_path("bottleneck/bn_conv2dk3_dw.cc", subdir="aie2"),
        [line_ty, line_ty, line_ty, wt_ty, out_ty, *_i32s(8)],
        compile_flags=["-DREGULAR", "-DSCALAR", f"-DSTRIDE{stride}"],
    )

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 
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
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",
        _default_source_path("bottleneck/bn_conv2dk1_relu.cc", subdir="aie2"),
        [in_ty, wt_ty, out_ty, *_i32s(8)],
        compile_flags=["-DSCALAR", "-DCONV_XYPOOL_FUSED_LARGE_PADDED", "-DINT8_ACT"],
    )

bn_conv2dk1_partial_put_i8 

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

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

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

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

Parameters:

Name Type Description Default
input_width int

Spatial width of the input slice.

 7
input_channels int

Number of input channels.

 80
weight_count int

Per-call weight chunk size in elements (the design streams weights in chunks; full weight tensor is shared across multiple kernel invocations).

 4800
block_index int

13 or 14; selects the per-block C++ wrapper.

 13

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the PUT tile.

Raises:

Type Description
ValueError

When block_index is not 13 or 14.
Source code in python/iron/kernels/conv.py 
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
def bn_conv2dk1_partial_put_i8(
    input_width: int = 7,
    input_channels: int = 80,
    weight_count: int = 4800,
    *,
    block_index: int = 13,
) -> ExternalFunction:
    """Cascade-PUT half of a width-split 1x1 conv on int8 activations.

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

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

    Args:
        input_width: Spatial width of the input slice.
        input_channels: Number of input channels.
        weight_count: Per-call weight chunk size in elements (the design
            streams weights in chunks; full weight tensor is shared
            across multiple kernel invocations).
        block_index: ``13`` or ``14``; selects the per-block C++ wrapper.

    Returns:
        ExternalFunction configured for the PUT tile.

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

bn_conv2dk1_partial_get_relu_i8 

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

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

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

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

Parameters:

Name Type Description Default
input_width int

Spatial width of the input slice.

 7
input_channels int

Number of input channels.

 80
output_channels int

Number of output channels (full L1 output width).

 480
weight_count int

Per-call weight chunk size in elements.

 4800
block_index int

13 or 14; selects the per-block C++ wrapper.

 13

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the GET tile.

Raises:

Type Description
ValueError

When block_index is not 13 or 14.
Source code in python/iron/kernels/conv.py 
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
def bn_conv2dk1_partial_get_relu_i8(
    input_width: int = 7,
    input_channels: int = 80,
    output_channels: int = 480,
    weight_count: int = 4800,
    *,
    block_index: int = 13,
) -> ExternalFunction:
    """Cascade-GET half of a width-split 1x1 conv + ReLU on int8 activations.

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

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

    Args:
        input_width: Spatial width of the input slice.
        input_channels: Number of input channels.
        output_channels: Number of output channels (full L1 output width).
        weight_count: Per-call weight chunk size in elements.
        block_index: ``13`` or ``14``; selects the per-block C++ wrapper.

    Returns:
        ExternalFunction configured for the GET tile.

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

bn_conv2dk3_dw_out_split 

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

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

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

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

Parameters:

Name Type Description Default
input_width int

Spatial width of the input.

 7
input_channels int

Number of input channels (== output channels — depthwise).

 480
output_split_channels int

Channels per output slice (half of input_channels for the typical 2-way split).

 240
block_index int

13 or 14; selects the per-block C++ wrapper.

 13

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the split-output DW kernel.

Raises:

Type Description
ValueError

When block_index is not 13 or 14.
Source code in python/iron/kernels/conv.py 
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
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")
    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",
        _default_source_path("bottleneck/bn_conv2dk3_dw.cc", subdir="aie2"),
        [line_ty, line_ty, line_ty, wt_ty, out_ty, out_ty, *_i32s(8)],
        compile_flags=["-DSCALAR", f"-DBN{block_index}", "-DSTRIDE1_OUT_SPLIT"],
    )

bn_conv2dk1_input_split_partial_put_ui8 

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

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

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

Parameters:

Name Type Description Default
input_width int

Spatial width of the input slice.

 7
input_channels int

Number of input channels (one half of the full input after split).

 240
weight_count int

Per-call weight chunk size in elements.

 9600
block_index int

13 or 14; selects the per-block C++ wrapper.

 13

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the input-split PUT tile.

Raises:

Type Description
ValueError

When block_index is not 13 or 14.
Source code in python/iron/kernels/conv.py 
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
def bn_conv2dk1_input_split_partial_put_ui8(
    input_width: int = 7,
    input_channels: int = 240,
    weight_count: int = 9600,
    *,
    block_index: int = 13,
) -> ExternalFunction:
    """Input-split cascade-PUT half of a 1x1 conv on uint8 activations.

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

    Args:
        input_width: Spatial width of the input slice.
        input_channels: Number of input channels (one half of the
            full input after split).
        weight_count: Per-call weight chunk size in elements.
        block_index: ``13`` or ``14``; selects the per-block C++ wrapper.

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

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

bn_conv2dk1_input_split_partial_skip_get 

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

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

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

Parameters:

Name Type Description Default
input_width int

Spatial width of the input slice.

 7
input_channels int

Number of input channels (one half after split).

 240
output_channels int

Final output channels.

 80
weight_count int

Per-call weight chunk size in elements.

 9600
block_index int

13 or 14; selects the per-block C++ wrapper.

 13

Returns:

Type Description
ExternalFunction

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

Raises:

Type Description
ValueError

When block_index is not 13 or 14.
Source code in python/iron/kernels/conv.py 
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
def bn_conv2dk1_input_split_partial_skip_get(
    input_width: int = 7,
    input_channels: int = 240,
    output_channels: int = 80,
    weight_count: int = 9600,
    *,
    block_index: int = 13,
) -> ExternalFunction:
    """Input-split cascade-GET half of a 1x1 conv + skip-add (uint8 in, int8 out).

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

    Args:
        input_width: Spatial width of the input slice.
        input_channels: Number of input channels (one half after split).
        output_channels: Final output channels.
        weight_count: Per-call weight chunk size in elements.
        block_index: ``13`` or ``14``; selects the per-block C++ wrapper.

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

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

bn_fc_relu_ui16_pad 

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

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

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

Parameters:

Name Type Description Default
input_channels int

Number of input channels (e.g. 1280).

 1280
output_channels int

Number of output channels per call (slice width, since the full FC is split across multiple tiles).

 16
weight_chunk_count int | None

Override the weight buffer's element count when the design streams weights in chunks (cascade/ping-pong). None means use the full input_channels * output_channels tile.

 None

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the post-L2 FC kernel.
Source code in python/iron/kernels/conv.py 
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
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",
        _default_source_path("bottleneck/bn_conv2dk1_relu.cc", subdir="aie2"),
        [in_ty, wt_ty, out_ty, *_i32s(5)],
        compile_flags=["-DSCALAR", "-DPOSTL2_PAD", "-DUINT16_ACT"],
    )

Activation functions

Activation kernel factories + numpy reference implementations.

Factories (each returns an ExternalFunction): softmax, gelu, silu, swiglu, bf16_exp.

Companion numpy reference implementations for host-side verification

relu_ref, silu_ref, gelu_ref, bf16_exp_ref, softmax_ref. These compute the AIE kernel's op in float32 so designs don't each reimplement the math in their verify path. Pair with count_mismatches (rtol=0.128 is the canonical LUT-tolerance default; see each ref's docstring for per-op recommendations).

softmax 

softmax(tile_size: int = 1024) -> ExternalFunction

Softmax activation kernel for bf16 tiles (tile_size must be 1024).

Parameters:

Name Type Description Default
tile_size int

Number of elements per tile.

 1024

Returns:

Type Description
ExternalFunction

ExternalFunction configured for the softmax kernel.
Source code in python/iron/kernels/activation.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def softmax(tile_size: int = 1024) -> ExternalFunction:
    """Softmax activation kernel for bf16 tiles (tile_size must be 1024).

    Args:
        tile_size: Number of elements per tile.

    Returns:
        ExternalFunction configured for the softmax kernel.
    """
    _require_fixed_tile_size("softmax", tile_size, _LUT_FIXED_TILE)
    tile_ty = np.ndarray[(tile_size,), np.dtype[bfloat16]]
    return _create_lut_kernel(
        "softmax_bf16",
        "softmax.cc",
        [tile_ty, tile_ty, np.int32],
    )

gelu 

gelu(tile_size: int = 1024) -> ExternalFunction

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

Source code in python/iron/kernels/activation.py 
112
113
114
def gelu(tile_size: int = 1024) -> ExternalFunction:
    """GELU activation kernel (tanh approximation) for bf16 tiles (must be 1024)."""
    return _bf16_lut_factory("gelu", "gelu_bf16", "gelu.cc", tile_size, arg_arity=2)

silu 

silu(tile_size: int = 1024) -> ExternalFunction

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

Source code in python/iron/kernels/activation.py 
117
118
119
def silu(tile_size: int = 1024) -> ExternalFunction:
    """SiLU (Swish) activation kernel for bf16 tiles (must be 1024)."""
    return _bf16_lut_factory("silu", "silu_bf16", "silu.cc", tile_size, arg_arity=2)

swiglu 

swiglu(tile_size: int = 1024) -> ExternalFunction

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

Source code in python/iron/kernels/activation.py 
122
123
124
125
126
def swiglu(tile_size: int = 1024) -> ExternalFunction:
    """SwiGLU gated activation kernel for bf16 tiles (must be 1024)."""
    return _bf16_lut_factory(
        "swiglu", "swiglu_bf16", "swiglu.cc", tile_size, arg_arity=4
    )

bf16_exp 

bf16_exp(tile_size: int = 1024) -> ExternalFunction

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

Source code in python/iron/kernels/activation.py 
129
130
131
132
133
def bf16_exp(tile_size: int = 1024) -> ExternalFunction:
    """Element-wise exponential kernel for bf16 tiles (must be 1024)."""
    return _bf16_lut_factory(
        "bf16_exp", "exp_bf16_1024", "bf16_exp.cc", tile_size, arg_arity=2
    )

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 
146
147
148
149
150
151
152
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 
155
156
157
158
159
160
161
162
def silu_ref(x):
    """numpy reference for [`silu`][iron.kernels.activation.silu] (Swish) — ``x * sigmoid(x)``.

    LUT-approximation territory; pair with ``rtol=0.128`` (the default
    in `count_mismatches`) when verifying.
    """
    xf = x.astype(np.float32)
    return (xf / (1.0 + np.exp(-xf))).astype(x.dtype)

gelu_ref 

gelu_ref(x)

numpy reference for gelu — tanh approximation 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))).

Matches the C++ kernel's tanh-GELU formula; pair with rtol=0.128, atol=0.05 when verifying.

Source code in python/iron/kernels/activation.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
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; pair with ``rtol=0.128,
    atol=0.05`` when verifying.
    """
    import math as _math

    xf = x.astype(np.float32)
    return (
        0.5 * xf * (1.0 + np.tanh(_math.sqrt(2.0 / _math.pi) * (xf + 0.044715 * xf**3)))
    ).astype(x.dtype)

bf16_exp_ref 

bf16_exp_ref(x)

numpy reference for bf16_exp — element-wise exp(x).

LUT approximation territory; the AIE kernel saturates on large inputs. Pair with the canonical 12.8% relative tolerance and stop_at_ nonfinite=True (the default in count_mismatches) when verifying.

Source code in python/iron/kernels/activation.py 
180
181
182
183
184
185
186
187
188
189
190
def bf16_exp_ref(x):
    """numpy reference for [`bf16_exp`][iron.kernels.activation.bf16_exp] — element-wise ``exp(x)``.

    LUT approximation territory; the AIE kernel saturates on large inputs.
    Pair with the canonical 12.8% relative tolerance and ``stop_at_
    nonfinite=True`` (the default in
    `count_mismatches`) when verifying.
    """
    xf = x.astype(np.float32)
    with np.errstate(over="ignore", invalid="ignore"):
        return np.exp(xf).astype(x.dtype)

softmax_ref 

softmax_ref(x, *, tile_size: int = 1024)

numpy reference for softmax.

The AIE kernel computes softmax independently per tile_size-element tile (no cross-tile reduction), so the reference splits x the same way before applying the float32 softmax. x.size must be a multiple of tile_size.

Source code in python/iron/kernels/activation.py 
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
def softmax_ref(x, *, tile_size: int = 1024):
    """numpy reference for [`softmax`][iron.kernels.activation.softmax].

    The AIE kernel computes softmax independently per ``tile_size``-element
    tile (no cross-tile reduction), so the reference splits ``x`` the same
    way before applying the float32 softmax.  ``x.size`` must be a
    multiple of ``tile_size``.
    """
    xf = x.astype(np.float32)
    if xf.size % tile_size != 0:
        raise ValueError(
            f"softmax_ref: x has {xf.size} elements; not a multiple of "
            f"tile_size={tile_size}"
        )
    flat = xf.reshape(-1, tile_size)
    flat = flat - flat.max(axis=1, keepdims=True)
    exp = np.exp(flat)
    out = exp / exp.sum(axis=1, keepdims=True)
    return out.reshape(x.shape).astype(x.dtype)

Vision

Vision kernel factories: color conversion, threshold, filter2d, add_weighted.

rgba2hue 

rgba2hue(
    line_width: int = 1920, use_chess: bool = False
) -> ExternalFunction

Converts a line of RGBA pixels to hue values.

Source code in python/iron/kernels/vision.py 
48
49
50
51
52
def rgba2hue(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Converts a line of RGBA pixels to hue values."""
    return _color_convert_kernel(
        "rgba2hueLine", "rgba2hue.cc", line_width * 4, line_width, use_chess=use_chess
    )

threshold 

threshold(
    line_width: int = 1920,
    dtype: type = uint8,
    use_chess: bool = False,
) -> ExternalFunction

Applies 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, np.int16, or np.int32).

 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 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def threshold(
    line_width: int = 1920, dtype: type = np.uint8, use_chess: bool = False
) -> ExternalFunction:
    """Applies a threshold operation to a line of pixels.

    Args:
        line_width: Number of elements per line.
        dtype: Element data type (``np.uint8``, ``np.int16``, or ``np.int32``).
        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")
    scalar_ty = np.int32 if bit_width == 32 else np.int16
    line_ty = np.ndarray[(line_width,), np.dtype[dtype]]
    return _make_extern(
        "thresholdLine",
        _default_source_path("threshold.cc"),
        [line_ty, line_ty, np.int32, scalar_ty, scalar_ty, np.int8],
        compile_flags=[f"-DBIT_WIDTH={bit_width}"],
        use_chess=use_chess,
    )

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 
81
82
83
84
85
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 
88
89
90
91
92
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

Converts a grayscale line to RGBA.

Source code in python/iron/kernels/vision.py 
95
96
97
98
99
def gray2rgba(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Converts a grayscale line to RGBA."""
    return _color_convert_kernel(
        "gray2rgbaLine", "gray2rgba.cc", line_width, line_width * 4, use_chess=use_chess
    )

rgba2gray 

rgba2gray(
    line_width: int = 1920, use_chess: bool = False
) -> ExternalFunction

Converts an RGBA line to grayscale.

Source code in python/iron/kernels/vision.py 
102
103
104
105
106
def rgba2gray(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Converts an RGBA line to grayscale."""
    return _color_convert_kernel(
        "rgba2grayLine", "rgba2gray.cc", line_width * 4, line_width, use_chess=use_chess
    )

filter2d 

filter2d(
    line_width: int = 1920, use_chess: bool = False
) -> ExternalFunction

Applies a 3x3 2D convolution filter across three input lines.

Source code in python/iron/kernels/vision.py 
109
110
111
112
113
114
115
116
117
118
def filter2d(line_width: int = 1920, use_chess: bool = False) -> ExternalFunction:
    """Applies a 3x3 2D convolution filter across three input lines."""
    line_ty = np.ndarray[(line_width,), np.dtype[np.uint8]]
    kernel_ty = np.ndarray[(3, 3), np.dtype[np.int16]]
    return _make_extern(
        "filter2dLine",
        _default_source_path("filter2d.cc"),
        [line_ty, line_ty, line_ty, line_ty, np.int32, kernel_ty],
        use_chess=use_chess,
    )

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, np.int16, or np.int32).

 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 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
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``, ``np.int16``, or ``np.int32``).
        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="add_weighted")
    gamma_ty = {8: np.int8, 16: np.int16, 32: np.int32}[bit_width]
    line_ty = np.ndarray[(line_width,), np.dtype[dtype]]
    return _make_extern(
        "addWeightedLine",
        _default_source_path("addWeighted.cc"),
        [line_ty, line_ty, line_ty, np.int32, np.int16, np.int16, gamma_ty],
        compile_flags=[f"-DBIT_WIDTH={bit_width}"],
        use_chess=use_chess,
    )