Operators

The operator library. Each operator is constructed by keyword (GEMV(M=..., K=...)); outside a graph, op(*tensors) compiles once and runs on the NPU, and inside a graph's body() calling it records a step. iron.operators.CATALOG holds every operator by the name a profile gives it.

The IRON operator library.

AXPY

AXPY(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = None, size: int, scalar_factor: float = 3.0)

Bases: BinaryElementwise

AIE-accelerated aX + Y operator.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
size	int		required
scalar_factor	float		3.0

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
a		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
b		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
y		Out(like=a, obj_shape=(Elementwise.tile_size,), per=CORES)

Source code in iron/operators/axpy.py

class AXPY(BinaryElementwise):

Clamp

Clamp(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = None, size: int, low: float, high: float)

Bases: UnaryElementwise

x clamped to [low, high], elementwise, the bounds rounded to bf16.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
size	int		required
low	float		required
high	float		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
y		Out(like=x, obj_shape=(Elementwise.tile_size,), per=CORES)
low_bits		Per call.
high_bits		Per call.

Source code in iron/operators/clamp.py

class Clamp(UnaryElementwise):

bits

bits() -> tuple

Return the bounds as the kernel takes them: each bf16's bits.

Source code in iron/operators/clamp.py

def bits(self) -> tuple:
    """Return the bounds as the kernel takes them: each bf16's bits."""
    return tuple(
        int(np.array(b, bfloat16).view(np.uint16)) for b in (self.low, self.high)
    )

Copy

Copy(*, input_buffer_size: int, src: TensorAccessPattern | tuple[TensorAccessPattern, ...] = lambda op: TensorAccessPattern.full((op.input_buffer_size,)), output_buffer_size: int = lambda op: sum(prod(p.sizes) for p in ((op.dst,) if op.src_repeat else op.src if isinstance(op.src, tuple) else (op.src,))), dst: TensorAccessPattern | tuple[TensorAccessPattern, ...] = lambda op: TensorAccessPattern.full((op.output_buffer_size // op.src.sizes[-1], op.src.sizes[-1]) if isinstance(op.src, TensorAccessPattern) else (op.output_buffer_size // op.src[0].sizes[-1], op.src[0].sizes[-1])), src_bound: int | None = None, dst_bound: int | None = None, src_repeat: bool = False, tile_size: int = None, num_channels: int = None, dtype: Any = bfloat16)

Bases: Operator

A copy between two views, gathered by src and scattered by dst.

Through num_channels memtile pass-throughs, split on the innermost run every pattern holds whole. A tile_size short of a descriptor never completes it, and the drain hangs (ERT_CMD_STATE_TIMEOUT).

Parameters:

Name	Type	Description	Default
input_buffer_size	int		required
src	TensorAccessPattern | tuple[TensorAccessPattern, ...]		lambda op: TensorAccessPattern.full((op.input_buffer_size,))
output_buffer_size	int	A repeated source's walk is its capacity: the destination's rows are the copy's.	lambda op: sum(prod(p.sizes) for p in ((op.dst,) if op.src_repeat else op.src if isinstance(op.src, tuple) else (op.src,)))
dst	TensorAccessPattern | tuple[TensorAccessPattern, ...]	Rows as wide as the source's innermost axis: the channels split both alike.	lambda op: TensorAccessPattern.full((op.output_buffer_size // op.src.sizes[-1], op.src.sizes[-1]) if isinstance(op.src, TensorAccessPattern) else (op.output_buffer_size // op.src[0].sizes[-1], op.src[0].sizes[-1]))
src_bound	int | None	The axis of each pattern a graph bounds per call.	None
dst_bound	int | None		None
src_repeat	bool	The source's first axis runs a per-call count of times (Handle.at's repeat).	False
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
dtype	Any		bfloat16

Attributes:

Name	Type	Description
x		In(input_buffer_size, dtype=dtype, obj_shape=(tile_size,), per=(num_channels,), depth=1)
y		Out(output_buffer_size, dtype=dtype, obj_shape=(tile_size,), per=(num_channels,), depth=1)
in_offset		Per call. In elements.
out_offset		Per call.
src_valid		Per call.
dst_valid		Per call.
src_repeats0		Per call. Each walk's runs past the first, as a task queue's repeat count counts them.
src_repeats1		Per call.
src_repeats2		Per call.
src_repeats3		Per call.
src_repeats4		Per call.
src_repeats5		Per call.
src_repeats6		Per call.
src_repeats7		Per call.

Source code in iron/operators/copy.py

class Copy(Operator):

channel_share property

channel_share: int

What each channel's transfers split into.

walks

walks() -> list[tuple[TensorAccessPattern, ...]]

Return the patterns each side walks in turn, [src, dst].

Source code in iron/operators/copy.py

def walks(self) -> list[tuple[TensorAccessPattern, ...]]:
    """Return the patterns each side walks in turn, ``[src, dst]``."""
    return [t if isinstance(t, tuple) else (t,) for t in (self.src, self.dst)]

reference

reference(x, y=None, *, in_offset=0, out_offset=0, src_valid=None, dst_valid=None, src_repeats0=None, src_repeats1=None, src_repeats2=None, src_repeats3=None, src_repeats4=None, src_repeats5=None, src_repeats6=None, src_repeats7=None)

Gather by src, scatter by dst, into y in place if given.

Source code in iron/operators/copy.py

def reference(
    self,
    x,
    y=None,
    *,
    in_offset=0,
    out_offset=0,
    src_valid=None,
    dst_valid=None,
    src_repeats0=None,
    src_repeats1=None,
    src_repeats2=None,
    src_repeats3=None,
    src_repeats4=None,
    src_repeats5=None,
    src_repeats6=None,
    src_repeats7=None,
):
    """Gather by ``src``, scatter by ``dst``, into ``y`` in place if given."""

    # Offsets count elements: the firmware scales the word by the element size.
    def at(tap, bound, valid):
        if valid is None:
            return tap
        return tap[(*[slice(None)] * bound, slice(int(valid)))]

    src, dst = self.walks()
    if self.src_repeat:
        repeats = (
            src_repeats0,
            src_repeats1,
            src_repeats2,
            src_repeats3,
            src_repeats4,
            src_repeats5,
            src_repeats6,
            src_repeats7,
        )
        src = [w if r is None else w[: int(r) + 1] for w, r in zip(src, repeats)]
    src = tuple(at(tap, self.src_bound, src_valid) for tap in src)
    dst = tuple(at(tap, self.dst_bound, dst_valid) for tap in dst)
    gather, scatter = (
        np.concatenate(
            [p.gather(np.arange(prod(concrete(p.tensor_dims, p)))) for p in tap]
        )
        + int(offset)
        for tap, offset in ((src, in_offset), (dst, out_offset))
    )
    if len(gather) != len(scatter):
        raise ValueError(
            f"Copy: pattern element counts differ ({len(gather)} vs {len(scatter)}); "
            "src and dst must move the same number of elements"
        )
    out = (
        np.zeros(self.output_buffer_size, dtype=x.dtype)
        if y is None
        else y.reshape(-1)
    )
    out[scatter] = x.reshape(-1)[gather]
    if y is None:
        return out
    # A strided y reshapes to a copy, so the result is written back through it.
    y[...] = out.reshape(y.shape)
    return y

Gather

Gather(*, rows: int, table_rows: int, row: int, dtype: Any = bfloat16, words: int = lambda op: 20 * op.chain + 6 * op.rows + 3 * len(op.batches) + (3 if op.rows % op.chain else 0), feeds: int = None)

Bases: Operator

Copy(table[ids]) with ids per call: the rows a call names.

Moved through feeds pairs of shims; a row past a bound moves pad_bytes instead, a descriptor's minimum 16-byte unit.

Parameters:

Name	Type	Description	Default
rows	int		required
table_rows	int		required
row	int		required
dtype	Any		bfloat16
words	int	Derived where not given; a given value must equal it.	lambda op: 20 * op.chain + 6 * op.rows + 3 * len(op.batches) + (3 if op.rows % op.chain else 0)
feeds	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None

Attributes:

Name	Type	Description
table		In(table_rows, row, dtype=dtype)
control		In(words, dtype=np.uint32)
y		Out(rows, row, dtype=dtype, obj_shape=(row,), per=(feeds,))
valid_rows		Per call.
drained0		Per call. Per feed, the 16-byte units it drains under a bound.
drained1		Per call.
drained2		Per call.
drained3		Per call.

Source code in iron/operators/copy.py

class Gather(Operator):

batches property

batches: list[int]

The rows of each batch, chain at most.

runs property

runs: list[range]

Per feed, the batches it streams: a run of whole pairs.

template cached property

template: tuple[ndarray, ndarray]

Return the control words with every row's address zero.

And where each row's address-low word is (its length one before, its high word two after).

drained

drained(feed: int) -> int

Return the 16-byte units feed feed drains.

With valid_rows rows valid.

Source code in iron/operators/copy.py

def drained(self, feed: int) -> int:
    """Return the 16-byte units feed ``feed`` drains.

    With ``valid_rows`` rows valid.
    """
    valid = self.valid_rows
    if feed >= self.feeds:
        return 0
    run = self.runs[feed]
    first = run.start * self.chain
    count = min(run.stop * self.chain, self.rows) - first
    kept = min(max(valid - first, 0), count)
    row_bytes = self.row * np.dtype(self.dtype).itemsize
    return (kept * row_bytes + (count - kept) * self.pad_bytes) // 16

word_source

word_source() -> GatherWords

Return the step writing row addresses into the template.

For ids the device made.

Source code in iron/operators/copy.py

def word_source(self) -> "GatherWords":
    """Return the step writing row addresses into the template.

    For ids the device made.
    """
    return GatherWords(
        rows=self.rows,
        table_rows=self.table_rows,
        row_bytes=self.row * np.dtype(self.dtype).itemsize,
        words=self.words,
        bounded="valid_rows" in self.bound_extents,
    )

control_words

control_words(ids: Sequence[int] | ndarray, address: int, valid: int | None = None) -> ndarray

Compute the control words gathering ids from a table at address.

Parameters:

Name	Type	Description	Default
ids	Sequence[int] | ndarray	Each row to gather, by table row index.	required
address	int	The table's device address.	required
valid	int | None	Rows past it move pad_bytes of the table's first row instead of their own; None moves every row.	None

Raises:

Type	Description
IndexError	An id is past the table, either way, or valid is outside [0, rows].

Source code in iron/operators/copy.py

def control_words(
    self, ids: Sequence[int] | np.ndarray, address: int, valid: int | None = None
) -> np.ndarray:
    """Compute the control words gathering ``ids`` from a table at ``address``.

    Args:
        ids: Each row to gather, by table row index.
        address: The table's device address.
        valid: Rows past it move ``pad_bytes`` of the table's first row
            instead of their own; ``None`` moves every row.

    Raises:
        IndexError: An id is past the table, either way, or ``valid`` is
            outside ``[0, rows]``.
    """
    ids = np.asarray(ids).astype(np.int64)
    if ids.shape != (self.rows,):
        raise ValueError(
            f"Gather: a gather of {self.rows} rows takes {self.rows} ids, got {ids.shape}"
        )
    valid = self.rows if valid is None else int(valid)
    if not 0 <= valid <= self.rows:
        raise IndexError(f"Gather: {valid} valid rows of a gather of {self.rows}")
    ids = ids[:valid]
    n = self.table_rows
    if ((ids < -n) | (ids >= n)).any():
        raise IndexError(
            f"Gather: gathering {ids[(ids < -n) | (ids >= n)]} from {n} rows"
        )
    template, slots = self.template
    stride = self.row * np.dtype(self.dtype).itemsize
    rows = np.full(self.rows, address + APERTURE, dtype=np.uint64)
    rows[:valid] += ((ids % n) * stride).astype(np.uint64)
    words = template.copy()
    words[slots[valid:] - 1] = self.pad_bytes // 4
    words[slots] = (rows & np.uint64(0xFFFFFFFC)).astype(np.uint32)
    words[slots + 2] = ((rows >> np.uint64(32)) & np.uint64(0xFFFF)).astype(
        np.uint32
    )
    return words

GatherWords

GatherWords(*, rows: int, table_rows: int, row_bytes: int, words: int, pairs: int = lambda op: 2 * op.rows, bounded: bool = False, record: int = lambda op: op.pairs + (op.rows if op.bounded else 0))

Bases: Operator

A Gather's control words for ids the device made, in a graph.

Each id, clipped to the table, becomes its row's address words, split at low_bits. Where bounded, it also writes each row's length, capped past valid_rows.

Parameters:

Name	Type	Description	Default
rows	int		required
table_rows	int		required
row_bytes	int		required
words	int		required
pairs	int	Derived where not given; a given value must equal it.	lambda op: 2 * op.rows
bounded	bool		False
record	int	Derived where not given; a given value must equal it. The words a call writes: the address pairs, then if bounded the lengths.	lambda op: op.pairs + (op.rows if op.bounded else 0)

Attributes:

Name	Type	Description
ids		In(rows, dtype=np.int32, obj_shape=(rows,), depth=1)
control		Out(words, dtype=np.uint32, obj_shape=(record,), depth=1)
base_lo		Per call.
base_hi		Per call.
valid_rows		Per call.
valid		Per call.

Source code in iron/operators/copy.py

class GatherWords(Operator):

address_words

address_words(address: int) -> dict[str, int]

base_lo and base_hi for a table at address.

Source code in iron/operators/copy.py

def address_words(self, address: int) -> dict[str, int]:
    """``base_lo`` and ``base_hi`` for a table at ``address``."""
    return dict(
        base_lo=address & ((1 << self.low_bits) - 1),
        base_hi=address >> self.low_bits,
    )

taps

taps() -> list[TensorAccessPattern]

Return where the rows' address words go in control, in row order.

Then if bounded their lengths, the word before each address.

Source code in iron/operators/copy.py

def taps(self) -> list[TensorAccessPattern]:
    """Return where the rows' address words go in ``control``, in row order.

    Then if ``bounded`` their lengths, the word before each address.
    """
    chain, at = Gather.chain, 20 * Gather.chain + 2
    full, short = divmod(self.rows, chain)
    taps = []
    for start, width in [(at, 2)] + ([(at - 1, 1)] if self.bounded else []):
        if full:
            taps.append(
                TensorAccessPattern(
                    (self.words,),
                    start,
                    [1, full, chain, width],
                    [0, 6 * chain + 3, 6, 2],
                )
            )
        if short:
            taps.append(
                TensorAccessPattern(
                    (self.words,),
                    start + (6 * chain + 3) * full,
                    [1, 1, short, width],
                    [0, 0, 6, 2],
                )
            )
    return taps

DepthwiseConv1d

DepthwiseConv1d(*, rows: int, channels: int, num_aie_columns: int = None)

Bases: Operator

A causal depthwise convolution over time, five taps, channels last.

One core per column takes a run of rows, its input fifo the sliding window; every run but the first is filled from four rows before it.

Parameters:

Name	Type	Description	Default
rows	int		required
channels	int		required
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None

Attributes:

Name	Type	Description
x		In(rows, channels, obj_shape=(channels,), per=(num_aie_columns,), depth=TAPS + 1)
weight		In(TAPS, channels, obj_shape=(channels,), depth=TAPS)
y		Out(rows, channels, obj_shape=(channels,), per=(num_aie_columns,))
valid		Per call. Every row is still computed under a bound: a row reads only those before it.
steps		Per call.

Source code in iron/operators/depthwise_conv1d.py

class DepthwiseConv1d(Operator):

sequence

sequence(rt)

Write the weights to every core, then each run with the four rows before it.

Source code in iron/operators/depthwise_conv1d.py

def sequence(self, rt):
    """Write the weights to every core, then each run with the four rows before it."""
    run = self.rows // self.num_aie_columns
    weights = TaskGroup()
    self.weight.fill(group=weights)
    tg = TaskGroup()
    for k in range(self.num_aie_columns):
        start = max(0, k * run - (TAPS - 1))
        self.x.fill(self.x[start : (k + 1) * run], lane=k, group=tg)
    for k in range(self.num_aie_columns):
        self.y.drain(self.y[k * run : (k + 1) * run], lane=k, group=tg, wait=True)
    tg.finish()
    weights.finish()

ElementwiseAdd

ElementwiseAdd(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = None, size: int)

Bases: BinaryElementwise

AIE-accelerated element-wise addition.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
size	int		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
a		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
b		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
y		Out(like=a, obj_shape=(Elementwise.tile_size,), per=CORES)

Source code in iron/operators/elementwise_add.py

class ElementwiseAdd(BinaryElementwise):

RowwiseAdd

RowwiseAdd(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int, rows: int)

Bases: RowBroadcast

Each row plus one row, x + row: a projection's bias.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int		required
rows	int		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(Rowwise.rows, Rowwise.tile_size, obj_shape=(line,), per=CORES)
y		Out(Rowwise.rows, Rowwise.tile_size, obj_shape=(line,), per=CORES)
row		In(Rowwise.tile_size, obj_shape=(line,), per=(Elementwise.num_channels,), replicate=True)

Source code in iron/operators/elementwise_add.py

class RowwiseAdd(RowBroadcast):

ElementwiseMul

ElementwiseMul(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = None, size: int)

Bases: BinaryElementwise

AIE-accelerated element-wise multiplication.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
size	int		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
a		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
b		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
y		Out(like=a, obj_shape=(Elementwise.tile_size,), per=CORES)

Source code in iron/operators/elementwise_mul.py

class ElementwiseMul(BinaryElementwise):

RowwiseMul

RowwiseMul(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int, rows: int)

Bases: RowBroadcast

Each row times one row, x * row.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int		required
rows	int		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(Rowwise.rows, Rowwise.tile_size, obj_shape=(line,), per=CORES)
y		Out(Rowwise.rows, Rowwise.tile_size, obj_shape=(line,), per=CORES)
row		In(Rowwise.tile_size, obj_shape=(line,), per=(Elementwise.num_channels,), replicate=True)

Source code in iron/operators/elementwise_mul.py

class RowwiseMul(RowBroadcast):

Emit

Emit(*, slots: int, carried: int, rows: int = lambda op: op.slots + op.carried)

Bases: Operator

Evaluate an emit program: the next call's scratchpad image and state.

Parameters:

Name	Type	Description	Default
slots	int		required
carried	int		required
rows	int	Derived where not given; a given value must equal it.	lambda op: op.slots + op.carried

Attributes:

Name	Type	Description
program		In(rows, ROW, dtype=np.int32, obj_shape=(rows, ROW), depth=1)
planes		In(2, carried, dtype=np.int32, obj_shape=(2, carried), depth=1)
image		Out(slots, dtype=np.int32, obj_shape=(slots,), depth=1)
state		Out(carried, dtype=np.int32, obj_shape=(carried,), depth=1)

Source code in iron/operators/emit.py

class Emit(Operator):

GELU

GELU(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = None, size: int, approximate: str = 'tanh')

Bases: UnaryElementwise

AIE-accelerated GELU.

By its tanh approximation, or exact (approximate="none", x * Phi(x)), as torch's approximate.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
size	int		required
approximate	str		'tanh'

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
y		Out(like=x, obj_shape=(Elementwise.tile_size,), per=CORES)

Source code in iron/operators/gelu.py

class GELU(UnaryElementwise):

GEMM

GEMM(*, M: int, K: int, N: int, tile_m: int = None, tile_k: int = None, tile_n: int = None, num_aie_columns: int = None, placement: GEMMPlacement = None, b_col_maj: bool = False, c_col_maj: bool = False, precision: Precision, input_dtype: Any = bfloat16, output_dtype: Any = bfloat16)

Bases: Operator

C = A @ B, tiled m x k x n.

A broadcasts across columns and splits across cores' rows, B splits across columns and broadcasts across rows, C joins across rows and splits across columns. Used wherever flm.ProjectionGEMM doesn't apply (a weight operand, MEDIUM/HIGH, two or more rows).

Parameters:

Name	Type	Description	Default
M	int		required
K	int		required
N	int		required
tile_m	int	Tunable: the call's, else the profile's, else resolve(dev)'s. None: the widest of 64, 32, 16, 8 that splits M and the m block divides.	None
tile_k	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 64. Searched at HIGHEST alone: below it C rounds to bf16 after every K tile, so tile_k changes the result.	None
tile_n	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 64.	None
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s. None: the most columns the shim budget allows that split N into tiles.	None
placement	GEMMPlacement	Tunable: the call's, else the profile's, else resolve(dev)'s, else GEMMPlacement.SHIMS.	None
b_col_maj	bool		False
c_col_maj	bool		False
precision	Precision		required
input_dtype	Any		bfloat16
output_dtype	Any		bfloat16

Attributes:

Name	Type	Description
A		In(M, K, dtype=input_dtype, obj_shape=(a_l2,), per=(n_shim_mem_a,))
B		In(Select(b_col_maj, (N, K), (K, N)), dtype=input_dtype, obj_shape=(b_l2,), per=(num_aie_columns,))
C		Out(Select(c_col_maj, (N, M), (M, N)), dtype=output_dtype, obj_shape=(c_l2,), per=(num_aie_columns,))
valid		Per call. A bound limits the compute, not the DMAs: A's pattern fills all four BD dims.
valid_k		Per call. The last K tile zeroes A's columns and B's rows past the bound.
k_div_k		Per call. Reduction steps per output tile, and output tiles per core.
n_tiles		Per call.
row_blocks_valid		Per call. Under a bound: the row blocks it covers, and the column tiles of each.
col_tiles		Per call.
k_valid		Per call.

Source code in iron/operators/gemm.py

class GEMM(Operator):

wider_accumulator property

wider_accumulator: bool

Check whether build() and GEMM.held() build HIGHEST's f32 buffer.

For a bf16 C's partial sums, beside C's own.

mac_block

mac_block(dev) -> tuple[int, int, int]

Return what tile_m, tile_k and tile_n must each be a multiple of.

mm.cc's r x s x t blocks, two in m and n (four in m on aie2). dev=None assumes aie2p's geometry.

Source code in iron/operators/gemm.py

def mac_block(self, dev) -> tuple[int, int, int]:
    """Return what ``tile_m``, ``tile_k`` and ``tile_n`` must each be a multiple of.

    mm.cc's r x s x t blocks, two in m and n (four in m on aie2).
    ``dev=None`` assumes aie2p's geometry.
    """
    r, s, t = kernels.mm.mac_dims(  # pyright: ignore[reportFunctionMemberAccess]
        self.input_dtype,
        self.output_dtype,
        device=dev,
        arch=None if dev is not None else "aie2p",
        emulate_bf16_mmul_with_bfp16=self.precision == Precision.MEDIUM,
    )
    aie2 = dev is not None and dev.arch is AIEArch.AIE2
    return (4 if aie2 else 2) * r, s, 2 * t

held

held(dev, tile_m: int) -> int

Core memory a core of dev holds at tile_m.

Source code in iron/operators/gemm.py

def held(self, dev, tile_m: int) -> int:
    """Core memory a core of ``dev`` holds at ``tile_m``."""
    m, k, n = tile_m, self.tile_k, self.tile_n
    din = np.dtype(self.input_dtype).itemsize
    dout = np.dtype(self.output_dtype).itemsize
    wider = self.replace(tile_m=m).layout().wider_accumulator
    c = m * n * (dout + 4 if wider else FIFO_DEPTH * dout)
    return FIFO_DEPTH * (m * k + k * n) * din + c + STACK_BYTES

matmul

matmul()

Return the m x k x n kernel each core runs.

Source code in iron/operators/gemm.py

def matmul(self):
    """Return the m x k x n kernel each core runs."""
    return kernels.linalg.mm(
        self.tile_m,
        self.tile_k,
        self.tile_n,
        input_dtype=self.input_dtype,
        # HIGHEST accumulates in f32 in L1, so the matmul's C is f32.
        output_dtype=(np.float32 if self.wider_accumulator else self.output_dtype),
        b_col_maj=self.b_col_maj,
        c_col_maj=self.c_col_maj,
        emulate_bf16_mmul_with_bfp16=self.precision == Precision.MEDIUM,
        round_conv_even=True,
    )

matmul_calls

matmul_calls(values: Mapping[str, int]) -> int

Return the fewest matmul kernel calls the busiest core makes in a step.

At per-call values: each output tile under the bound takes one call per K tile under it.

Source code in iron/operators/gemm.py

def matmul_calls(self, values: Mapping[str, int]) -> int:
    """Return the fewest matmul kernel calls the busiest core makes in a step.

    At per-call ``values``: each output tile under the bound takes one
    call per K tile under it.
    """
    calls = (
        ceildiv(values.get("valid", self.M), self.tile_m)
        * ceildiv(values.get("valid_k", self.K), self.tile_k)
        * (self.N // self.tile_n)
    )
    return ceildiv(calls, self.n_aie_rows * self.num_aie_columns)

reference

reference(A, B)

Compute A @ B over the first valid_k of K, a bounded call's.

Source code in iron/operators/gemm.py

def reference(self, A, B):
    """Compute ``A @ B`` over the first ``valid_k`` of K, a bounded call's."""
    # float32, not mm_ref's float64, which doubles the weight's host copy.
    a = A.astype(np.float32)[:, : self.valid_k]
    b = (B.T if self.b_col_maj else B).astype(np.float32)[: self.valid_k]
    C = np.matmul(b.T, a.T) if self.c_col_maj else np.matmul(a, b)
    return C.astype(self.output_dtype)

tolerance

tolerance() -> Tolerance

Return C's roundings, in units of 2^-8 of what each rounds.

2 of |A| @ |B| for C's own, 2 per K tile's partial sum but the last when C accumulates in bf16, and 2 more for bfp16 inputs; measured at no more than 0.83 of this over K = 256-8192 on npu2.

Source code in iron/operators/gemm.py

def tolerance(self) -> Tolerance:
    """Return C's roundings, in units of 2^-8 of what each rounds.

    2 of ``|A| @ |B|`` for C's own, 2 per K tile's partial sum but the
    last when C accumulates in bf16, and 2 more for bfp16 inputs;
    measured at no more than 0.83 of this over K = 256-8192 on npu2.
    """
    if np.issubdtype(np.dtype(self.input_dtype), np.integer):
        return Tolerance.exact(note="integer matmul")
    units = 4.0 if self.precision == Precision.MEDIUM else 2.0
    per_tile = (
        np.dtype(self.output_dtype) == np.dtype(bfloat16)
        and self.precision != Precision.HIGHEST
    )

    def bound(A, B):
        a = A.astype(np.float32)
        b = (B.T if self.b_col_maj else B).astype(np.float32)
        err = units * (np.abs(a) @ np.abs(b))
        if per_tile:
            partial = np.zeros_like(err)
            for k0 in range(0, self.K - self.tile_k, self.tile_k):
                partial += a[:, k0 : k0 + self.tile_k] @ b[k0 : k0 + self.tile_k]
                err += 2 * np.abs(partial)
        err *= 2.0**-8
        return err.T if self.c_col_maj else err

    return Tolerance.bounded(
        bound,
        note=f"{units:g} units of |A| @ |B|"
        + (", 2 of each bf16 partial sum" if per_tile else ""),
    )

GEMV

GEMV(*, M: int, K: int, num_batches: int | None = None, repeat: int = 1, num_matrices: int | None = lambda op: None if op.num_batches is None else op.num_batches // op.repeat, num_aie_columns: int = None, num_channels: int = None, tile_size_input: int = None, tile_size_output: int = None, kernel_vector_size: int = None)

Bases: Operator

C = A @ B as row-blocks of A per lane, tile_size_input rows per mv.cc call.

B rides each lane's A stream ahead of its rows so every input channel of a column can carry A.

Parameters:

Name	Type	Description	Default
M	int		required
K	int		required
num_batches	int | None	None: a single batch, with no batch dimension.	None
repeat	int	Batches per matrix of A (grouped-query scores): no repeated copy.	1
num_matrices	int | None	Derived where not given; a given value must equal it.	lambda op: None if op.num_batches is None else op.num_batches // op.repeat
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
tile_size_input	int	Tunable: the call's, else the profile's, else resolve(dev)'s. None: two rows where they fit one L1 programmer bank, else one.	None
tile_size_output	int	Tunable: the call's, else the profile's, else resolve(dev)'s. At least two rows: the shim moves C in 4-byte granules.	None
kernel_vector_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None

Attributes:

Name	Type	Description
A		In(num_matrices, M, K, obj_shape=(tile_size_input, K), per=(num_aie_columns, num_channels), depth=2)
B		In(num_batches, K)
C		Out(num_batches, M, obj_shape=(tile_size_output,), per=(num_aie_columns, num_channels), depth=2)
valid		Per call.
tiles		Per call.
batches		Per call.

Source code in iron/operators/gemv.py

class GEMV(Operator):

sequence

sequence(rt)

Sequence per batch and lane: B's head tile, then A's rows.

C drains are issued first: an undrained C would stop A.

Source code in iron/operators/gemv.py

def sequence(self, rt):
    """Sequence per batch and lane: B's head tile, then A's rows.

    C drains are issued first: an undrained C would stop A.
    """
    if self.repeat > 1:
        for buf in (self.A, self.C):
            if buf.bounded is not None:
                raise ValueError(
                    f"GEMV: {buf.name} takes no per-call bound when repeat "
                    f"({self.repeat}) is above 1"
                )
    M, K, tsi = self.M, self.K, self.tile_size_input
    drains = [
        self.C.drain(
            tap,
            lane=k,
            wait=True,
            length_parameter=word,
            length_unit=unit,
            managed=False,
        )
        for k, tap, word, unit in self.C.transfers()
    ]
    rows = self.A.transfers()
    if (self.num_matrices or 1) > 1 and self.A.bounded is None:
        parts = TensorAccessPattern.full((M, K)).partition(len(rows), 0)
        rows = [
            (
                k,
                TensorAccessPattern(
                    self.A.shape, part.offset, part.sizes, part.strides
                ),
                None,
                None,
            )
            for (k, _, _, _), part in zip(rows, parts, strict=True)
        ]
    elif (self.num_matrices or 1) > 1:
        rows = [
            (
                k,
                TensorAccessPattern(
                    tap.tensor_dims,
                    tap.offset,
                    [1, *tap.sizes[1:]],
                    [0, *tap.strides[1:]],
                ),
                word,
                unit,
            )
            for k, tap, word, unit in rows
        ]
    for b in range(self.num_batches or 1):
        # B's row repeated over an A tile in the iteration dim (stride 0).
        head = TensorAccessPattern(
            self.B.shape, b * K, [tsi, 1, 1, K], [0, 0, 0, 1]
        )
        for k, tap, word, unit in rows:
            self.A.fill(head, lane=k, source=self.B, managed=False)
            (offset,) = concrete((tap.offset,), tap)
            matrix = TensorAccessPattern(
                tap.tensor_dims,
                offset + b // self.repeat * M * K,
                tap.sizes,
                tap.strides,
            )
            self.A.fill(
                matrix,
                lane=k,
                length_parameter=word,
                length_unit=unit,
                managed=False,
            )
    for task in drains:
        task.await_()

reference

reference(A, B)

C = A @ B, one product per batch.

Source code in iron/operators/gemv.py

def reference(self, A, B):
    """``C = A @ B``, one product per batch."""
    if self.repeat > 1:
        A = np.repeat(A.reshape(-1, *A.shape[-2:]), self.repeat, axis=0)
    # float32: numpy's bf16 matmul accumulates in bf16, the kernel in f32.
    a, b = A.astype(np.float32), B.astype(np.float32)
    if A.ndim == 3:
        b = b.reshape(A.shape[0], A.shape[2], 1)
        return np.matmul(a, b).reshape(A.shape[0], A.shape[1]).astype(A.dtype)
    return (a @ b.reshape(A.shape[-1])).astype(A.dtype)

tolerance

tolerance() -> Tolerance

C accumulates in f32 and rounds once; tighter than linalg.mv's contract.

Source code in iron/operators/gemv.py

def tolerance(self) -> Tolerance:
    """C accumulates in f32 and rounds once; tighter than linalg.mv's contract."""
    return Tolerance.relative(0.04, 1e-3, note="f32 accumulation, rounded once")

LayerNorm

LayerNorm(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int, rows: int, weighted: bool = False, epsilon: float = 1e-05)

Bases: WeightedRowwise

AIE-accelerated Layer Normalization of each row (beta 0), optionally times a learned weight row (gamma).

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int		required
rows	int		required
weighted	bool		False
epsilon	float	The epsilon under the root: torch's 1e-5 by default; a model states its own.	1e-05

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(Rowwise.rows, Rowwise.tile_size, obj_shape=(Rowwise.tile_size,), per=CORES)
y		Out(Rowwise.rows, Rowwise.tile_size, obj_shape=(Rowwise.tile_size,), per=CORES)
weight		In(Rowwise.tile_size, obj_shape=(Rowwise.tile_size,), per=(Elementwise.num_channels,), replicate=True, when=weighted) Shared by every column of a channel: counted once per channel.

Source code in iron/operators/layer_norm.py

class LayerNorm(WeightedRowwise):

Limbs

Limbs(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = lambda op: 6 * op.line, rows: int, line: int)

Bases: Rowwise

Each row of float32 as six planes of its bf16 limbs, for a float32 product by GEMM.

[hi, mid, hi, lo, mid, hi].

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Derived where not given; a given value must equal it.	lambda op: 6 * op.line
rows	int		required
line	int		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(Rowwise.rows, line, dtype=np.float32, obj_shape=(line,), per=CORES)
y		Out(rows, tile_size, obj_shape=(tile_size,), per=CORES)

Source code in iron/operators/limbs.py

class Limbs(Rowwise):

Log

Log(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = None, size: int, offset: float = 0.0)

Bases: UnaryElementwise

log(x + offset) of float32, rounded once to bf16.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
size	int		required
offset	float		0.0

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(UnaryElementwise.size, dtype=np.float32, obj_shape=(UnaryElementwise.tile_size,), per=CORES)
y		Out(like=x, obj_shape=(Elementwise.tile_size,), per=CORES)

Source code in iron/operators/log.py

class Log(UnaryElementwise):

Magnitude

Magnitude(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = lambda op: op.width // 2, rows: int, width: int)

Bases: Rowwise

The magnitude of each row of complex float32 held as [re | im].

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Derived where not given; a given value must equal it.	lambda op: op.width // 2
rows	int		required
width	int		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(Rowwise.rows, width, dtype=np.float32, obj_shape=(width,), per=CORES)
y		Out(Rowwise.rows, tile_size, dtype=np.float32, obj_shape=(tile_size,), per=CORES)

Source code in iron/operators/magnitude.py

class Magnitude(Rowwise):

Merge

Merge(*, rows: int, placeholders: tuple[tuple[int, int], ...], block: int = 64, offset: bool = False, kinds: int = lambda op: len(op.placeholders))

Bases: Operator

The row each position of ids takes from the text's embeddings and the towers' soft tokens.

For each placeholder (token, at), the j-th token takes row at + j; any other id its own position.

Given before, the count of each placeholder before these ids (a prompt chunk after others), the j-th takes row at + before + j.

Parameters:

Name	Type	Description	Default
rows	int		required
placeholders	tuple[tuple[int, int], ...]	Each placeholder's (token id, first row of its soft tokens).	required
block	int		64
offset	bool		False
kinds	int	Derived where not given; a given value must equal it.	lambda op: len(op.placeholders)

Attributes:

Name	Type	Description
ids		In(rows, dtype=np.int32, obj_shape=(block,))
before		In(kinds, dtype=np.int32, obj_shape=(kinds,), depth=1, when=offset)
merge		Out(rows, dtype=np.int32, obj_shape=(block,))
count		Per call.

Source code in iron/operators/merge.py

class Merge(Operator):

MHA

MHA(*, num_heads: int, num_KV_heads: int = lambda op: op.num_heads, seq_len: int = lambda op: op.seq_pad, seq_pad: int = lambda op: op.seq_padding(op.seq_len), kv_len: int = lambda op: op.seq_pad, heads_interleaved: bool = False, kv_interleaved: bool = False, d: int = 64, causal: bool = True, window: int | None = None, scale: float = lambda op: float(1 / np.sqrt(op.d)), B_q: int = None, num_pipelines: int = None, pv_cores: int = None)

Bases: Operator

Fused attention over (B_q, d) Q, (d, B_kv) K and (B_kv, d) V blocks.

More than six pipelines split Q and O over two shims.

Parameters:

Name	Type	Description	Default
num_heads	int		required
num_KV_heads	int		lambda op: op.num_heads
seq_len	int		lambda op: op.seq_pad
seq_pad	int	Derived where not given; a given value must equal it.	lambda op: op.seq_padding(op.seq_len)
kv_len	int	A cache longer than the queries, whose last rows the queries are.	lambda op: op.seq_pad
heads_interleaved	bool	(seq, heads, d) as a projection writes it, so no copy reorders the heads.	False
kv_interleaved	bool		False
d	int		64
causal	bool		True
window	int | None	Keys more than this many positions from a query are masked.	None
scale	float		lambda op: float(1 / np.sqrt(op.d))
B_q	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_pipelines	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
pv_cores	int	Tunable: the call's, else the profile's, else resolve(dev)'s. Two P*V cores, by halves of the head, where one's L1 holds no block.	None

Attributes:

Name	Type	Description
Q		In(Select(heads_interleaved, (seq_pad, num_heads, d), (num_heads, seq_pad, d)), obj_shape=(join_rows, d), per=(q_shims,))
K		In(Select(kv_interleaved, (kv_len, num_KV_heads, d), (num_KV_heads, kv_len, d)), obj_shape=(d, B_kv), per=(kv_lanes,))
V		In(Select(kv_interleaved, (kv_len, num_KV_heads, d), (num_KV_heads, kv_len, d)), obj_shape=(B_kv, pv_width), per=(kv_lanes, pv_cores))
O		Out(Select(heads_interleaved, (seq_pad, num_heads, d), (num_heads, seq_pad, d)), obj_shape=(join_rows, pv_width), per=(q_shims, pv_cores))
valid		Per call. A bound on the queries limits the compute; Q and O stream every row.
kv_valid		Per call. A bound on the keys limits K's and V's blocks: the work follows the context.
heads		Per call. One query packed sits past every key block: no causal skip, no diagonal.
q_blocks_per_pipeline		Per call.
q_blocks_valid		Per call.
kv_blocks		Per call.
s_q		Per call.
s_kv		Per call.
q_start		Per call.
q_offset		Per call.

Source code in iron/operators/mha.py

class MHA(Operator):

q_tokens property

q_tokens: int

The query rows attended from: the call's bound, else seq_len.

kv_tokens property

kv_tokens: int

The keys attended over: the call's bound, else kv_len less Q's padding.

first_query property

first_query: int

The position of Q's first row, which mha.cc masks by; 0 unmasked.

first_query_block property

first_query_block: int

The Q block first_query is in, q_offset rows into it.

fits classmethod

fits(dev) -> bool

Whether dev has the columns of a layout MHA is placed on.

Source code in iron/operators/mha.py

@classmethod
def fits(cls, dev) -> bool:
    """Whether ``dev`` has the columns of a layout MHA is placed on."""
    return dev.cols >= min(cls.PACKED_PIPELINES)

seq_padding

seq_padding(seq_len: int) -> int

seq_len rounded up to 64; one query, packed, is not padded.

Source code in iron/operators/mha.py

def seq_padding(self, seq_len: int) -> int:
    """``seq_len`` rounded up to 64; one query, packed, is not padded."""
    if seq_len == 1:
        return 1
    return ceildiv(seq_len, 64) * 64

key_block

key_block(B_q: int, dev: Device | None) -> int

B_kv for B_q: itself, which mha.cc's causal skip compares by.

Packed, the widest divisor of kv_len whose cores fit dev (any, where None), since mha.cc's vectorized softmax reads 64-key rows.

Source code in iron/operators/mha.py

def key_block(self, B_q: int, dev: Device | None) -> int:
    """``B_kv`` for ``B_q``: itself, which mha.cc's causal skip compares by.

    Packed, the widest divisor of ``kv_len`` whose cores fit ``dev``
    (any, where None), since mha.cc's vectorized softmax reads 64-key rows.
    """
    if not self.packed:
        return B_q
    return next(
        (
            b
            for b in (64, 32, 16)
            if B_q <= b
            and not self.kv_len % b
            and (
                dev is None
                or self.core_bytes(dev, B_q, b, 1) <= dev.core_memory_bytes
            )
        ),
        B_q,
    )

core_bytes

core_bytes(dev: Device, B_q: int, B_kv: int, pv_cores: int) -> int

Return the larger of a pipeline's QK^T and P*V cores' L1 bytes.

Kernels' data and stack included, over the device's default stack.

Parameters:

Name	Type	Description	Default
dev	Device	The device the cores' memory is measured against.	required
B_q	int	The query block size.	required
B_kv	int	The key/value block size.	required
pv_cores	int	The P*V cores the head is split over.	required

Source code in iron/operators/mha.py

def core_bytes(self, dev: Device, B_q: int, B_kv: int, pv_cores: int) -> int:
    """Return the larger of a pipeline's QK^T and P*V cores' L1 bytes.

    Kernels' data and stack included, over the device's default stack.

    Args:
        dev: The device the cores' memory is measured against.
        B_q: The query block size.
        B_kv: The key/value block size.
        pv_cores: The P*V cores the head is split over.
    """
    d, pv_width = self.d, self.d // pv_cores
    arch = resolve_target_arch(dev)
    pv = 4 * B_q * B_kv + 4 * pv_width * B_kv + 6 * pv_width * B_q + 32 * B_q
    qk = (4 if pv_cores == 1 else 2) * d * B_q + 4 * d * B_kv + 4 * B_q * B_kv
    # Only measured, never built, so it stays out of the build's registry.
    # dev is bound because the factory reads its architecture from it.
    with set_current_device(dev), compile_context(_iron_kernel_registry=set()):
        qk_mha = self.mha_kernel(B_q, d, B_kv, dev.arch is AIEArch.AIE2p)
        pv_mha = qk_mha if pv_cores == 1 else self.mha_kernel(B_q, pv_width, B_kv)
        bind = pv_mha.object_file.bind
        pv += max(
            pv_mha.frame_bytes(arch, entry=bind(name).name)
            for name in ("matmul_PV", "rescale_O")
        )
        zero = kernels.zero(tile_size=(B_q, B_kv), dtype=bfloat16)
        # Summed, the two objects' stacks bound the deeper one from above.
        qk += qk_mha.frame_bytes(arch) + zero.frame_bytes(arch)
    return dev.default_core_stack_bytes + max(pv, qk)

mha_kernel

mha_kernel(B_q: int, width: int, B_kv: int, accurate_exp2=False)

mha.cc at a B_q x width x B_kv block.

Source code in iron/operators/mha.py

def mha_kernel(self, B_q: int, width: int, B_kv: int, accurate_exp2=False):
    """mha.cc at a ``B_q`` x ``width`` x ``B_kv`` block."""
    return kernels.linalg.mha(
        B_q,
        width,
        B_kv,
        b_col_maj=True,
        accurate_exp2=accurate_exp2,
        causal=self.causal,
        window=self.window or 0,
        # Unwindowed decode runs as it has, without mha.cc's packed path.
        packed=self.packed and self.window is not None,
    )

ops

ops() -> int

Q K^T and its product with V per head, over the keys each query keeps.

Source code in iron/operators/mha.py

def ops(self) -> int:
    """Q K^T and its product with V per head, over the keys each query keeps."""
    keys = self.kv_len - self.seq_pad + self.seq_len
    position = keys - self.seq_len + np.arange(self.seq_len)
    ahead = 0 if self.causal else (self.window or keys)
    lo = np.maximum(position - self.window, 0) if self.window else 0
    hi = np.minimum(position + ahead, keys - 1)
    return 4 * self.num_heads * self.d * int((hi - lo + 1).sum())

reference

reference(Q, K, V, s_q=None, s_kv=None)

Attention per head, the queries the keys' last rows.

Query rows from s_q on are zeros and keys from s_kv on are masked.

Source code in iron/operators/mha.py

def reference(self, Q, K, V, s_q=None, s_kv=None):
    """Attention per head, the queries the keys' last rows.

    Query rows from ``s_q`` on are zeros and keys from ``s_kv`` on are
    masked.
    """
    if self.heads_interleaved:
        Q = np.swapaxes(Q, 0, 1)
    if self.kv_interleaved:
        K, V = (np.swapaxes(t, 0, 1) for t in (K, V))
    groups = self.num_heads // self.num_KV_heads
    # The keys before the first query: those not in Q's padding, or a
    # bounded call's.
    offset = self.kv_tokens - self.q_tokens
    start = offset if self.causal or self.window else 0
    s_q = start + self.q_tokens if s_q is None else int(s_q)
    s_kv = self.kv_tokens if s_kv is None else int(s_kv)
    K, V = K[:, :s_kv], V[:, :s_kv]
    rows = max(s_q - start, 0)
    q = Q[:, :rows].astype(np.float32)
    k, v = (t.astype(np.float32) for t in (K, V))
    key = np.arange(k.shape[1])
    scale = np.float32(self.scale)
    out = np.zeros(Q.shape, dtype=Q.dtype)
    # A block of query rows at a time, its scores, mask and exponentials
    # 64 MiB each at any context, where all rows' are GiBs at 128K keys.
    step = max(1, 2**24 // max(k.shape[1], 1))
    for first in range(0, rows, step):
        block = slice(first, min(first + step, rows))
        position = offset + np.arange(block.start, block.stop)[:, None]
        masked = np.zeros((len(position), k.shape[1]), dtype=bool)
        if self.causal or self.window:
            # The enclosing "if" means self.window is set whenever not causal.
            masked |= key > position + (0 if self.causal else (self.window or 0))
        if self.window:
            masked |= key < position - self.window
        mask = np.where(masked, -np.inf, 0).astype(np.float32)
        for h in range(q.shape[0]):
            scores = q[h, block] @ k[h // groups].T * scale + mask
            e = np.exp(scores - scores.max(axis=-1, keepdims=True))
            out[h, block] = (e / e.sum(axis=-1, keepdims=True)) @ v[h // groups]
    if self.heads_interleaved:
        return np.ascontiguousarray(np.swapaxes(out, 0, 1))
    return out

sequence

sequence(rt)

One descriptor set per KV group.

K and V re-read per (head, block) from the iteration slot, not six descriptors per block.

Source code in iron/operators/mha.py

def sequence(self, rt):
    """One descriptor set per KV group.

    K and V re-read per (head, block) from the iteration slot, not six
    descriptors per block.
    """
    kv_heads = self.num_KV_heads
    group = self.num_heads // kv_heads
    rows = self.join_rows  # Q rows each shim carries per block
    blocks = ceildiv(self.seq_pad, rows * self.q_shims)  # per pipeline
    B_kv = self.B_kv
    kv_blocks = ceildiv(self.kv_tokens, B_kv)
    kv_valid = self.kv_blocks if self.kv_blocks.per_call else None

    def by_head(buffer, interleaved):
        tap = TensorAccessPattern.full(buffer.shape)
        return tap.permute((1, 0, 2)) if interleaved else tap

    def q_rows(buffer, head0, shim):
        heads = by_head(buffer, self.heads_interleaved)[head0 : head0 + group]
        by_shim = heads.split(1, self.q_shims * rows)
        return by_shim[:, :, shim * rows : (shim + 1) * rows]

    def kv_rows(buffer, kv_head, reads):
        head = by_head(buffer, self.kv_interleaved)[kv_head, : kv_blocks * B_kv]
        return head.split(0, B_kv).repeat(reads)

    if self.packed:
        steps = kv_heads // self.num_pipelines
        # A group that does not divide a block is two overlapping runs.
        width = 1 << (group.bit_length() - 1)

        def packed_rows(buffer, kv_head):
            heads = by_head(buffer, self.heads_interleaved)
            if width == group:
                heads = heads[kv_head * group : (kv_head + 1) * group, 0]
                return heads.coalesce().repeat(self.B_q // group)
            first = heads[kv_head * group : kv_head * group + width, 0]
            (stride0,) = concrete(first.strides[:1], first)
            both = TensorAccessPattern(
                first.tensor_dims,
                first.offset,
                (2, *first.sizes),
                ((group - width) * stride0, *first.strides),
            )
            return both.coalesce().repeat(self.B_q // (2 * width))

        for step in range(steps):
            owned = [p * steps + step for p in range(self.num_pipelines)]
            tg = TaskGroup()
            for kv_head in owned:
                self.Q.fill(packed_rows(self.Q, kv_head), lane=0, group=tg)
            for p, kv_head in enumerate(owned):
                for x in (self.K, self.V):
                    x.fill(
                        kv_rows(x, kv_head, 1),
                        lane=p,
                        group=tg,
                        length_parameter=kv_valid,
                        length_unit=B_kv * x.shape[-1],
                    )
            for kv_head in owned:
                self.O.drain(
                    packed_rows(self.O, kv_head), lane=0, group=tg, wait=True
                )
            tg.finish()
        return

    cores, pv_width = self.pv_cores, self.pv_width
    dev = self.dev
    if dev is None:
        raise Unresolvable("MHA: sequence() needs a bound device")
    # A shim descriptor's iteration count.
    iterations = 1 << dev.get_dma_bd_iter_bits(0, 0)
    for kv_head in range(kv_heads):
        head0 = kv_head * group
        tg = TaskGroup()
        for shim in range(self.q_shims):
            self.Q.fill(q_rows(self.Q, head0, shim), lane=shim, group=tg)
        # O drains before K and V fill: past the queue's depth a fill
        # waits for a free slot, which the cores free only while O drains.
        for shim in range(self.q_shims):
            for h in range(cores):
                self.O.drain(
                    q_rows(self.O, head0, shim)[
                        ..., h * pv_width : (h + 1) * pv_width
                    ],
                    lane=shim * cores + h,
                    group=tg,
                    wait=True,
                )
        # A bounded pattern is one descriptor: re-reads at most the slot's count.
        reads = group * blocks
        for first in range(0, reads, iterations):
            count = min(iterations, reads - first)
            self.K.fill(
                kv_rows(self.K, kv_head, count),
                group=tg,
                length_parameter=kv_valid,
                length_unit=B_kv * self.K.shape[-1],
            )
            for h in range(cores):
                self.V.fill(
                    kv_rows(self.V, kv_head, count)[
                        ..., h * pv_width : (h + 1) * pv_width
                    ],
                    lane=h,
                    group=tg,
                    length_parameter=kv_valid,
                    length_unit=B_kv * pv_width,
                )
        tg.finish()

ReLU

ReLU(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = None, size: int)

Bases: UnaryElementwise

AIE-accelerated ReLU activation function.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
size	int		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
y		Out(like=x, obj_shape=(Elementwise.tile_size,), per=CORES)

Source code in iron/operators/relu.py

class ReLU(UnaryElementwise):

PatchPositions

PatchPositions(*, rows: int, out_height: int, out_width: int, pool: int = 3, positions: int = 10240, cores: int = 16, block: int = 256)

Bases: Operator

Where each of rows patch rows of the vision tower comes from.

For an image resized to out_rows by out_columns pixels.

The patches are in pooling-window order; order is each one's row of Resample's raster. Rows past the image hold positions, 2 * positions, past each table, which a gather bounded to the image's patches never reads, and the raster's zero patch.

Parameters:

Name	Type	Description	Default
rows	int		required
out_height	int		required
out_width	int		required
pool	int		3
positions	int	The rows of each axis's position-embedding table.	10240
cores	int	Resample's cores, which its raster's columns are padded past.	16
block	int		256

Attributes:

Name	Type	Description
out_rows		Per call.
out_columns		Per call.
xy		Out(rows, 2, dtype=np.int32, obj_shape=(block, 2))
position_ids		Out(2, rows, dtype=np.int32, obj_shape=(2, block))
order		Out(rows, dtype=np.int32, obj_shape=(block,))
count		Per call.
dst_rows		Per call.
dst_columns		Per call.

Source code in iron/operators/resample/op.py

class PatchPositions(Operator):

sequence

sequence(rt)

Write xy and order as the core computes them.

Writes position_ids with a block's x into the first half and its y into the second.

Source code in iron/operators/resample/op.py

def sequence(self, rt):
    """Write `xy` and `order` as the core computes them.

    Writes `position_ids` with a block's x into the first half and its
    y into the second.
    """
    T, B = self.rows, self.block
    tg = TaskGroup()
    self.xy.drain(group=tg, wait=True)
    self.position_ids.drain(
        TensorAccessPattern((2, T), 0, [1, T // B, 2, B], [0, B, T, 1]),
        group=tg,
        wait=True,
    )
    self.order.drain(group=tg, wait=True)
    tg.finish()

Resample

Resample(*, height: int, width: int, out_height: int, out_width: int, chunk: int = 4096, words: int = 356, patch_columns: int = lambda op: max(8, op.default_patch_columns()), image_chunks: int = lambda op: op.height * -(-3 * op.width // op.chunk), width_chunks: int = lambda op: 16 * -(-op.out_width // (16 * SIDE)), height_chunks: int = lambda op: 16 * -(-op.out_height // (16 * SIDE)), patch_rows: int = lambda op: op.padded_rows(), num_aie_columns: int = None, num_channels: int = None)

Bases: Operator

An image resized as torchvision's antialiased bicubic resize does it on the CPU.

Written as the vision tower's 16x16 patches of u8 / 255 in bf16 (vision.resize).

Every core receives the image in chunk-byte chunks; core k owns patch columns k, k + cores, ..., resamples each row across them, then each band of 16 rows down them. Patch (py, px) is row py * pad + px of patches, pad the patch columns padded past the cores with zeros: row pad - 1 is the zero patch padding gathers. The resized sizes are whole patches; a call's sizes that are not write zeros.

Parameters:

Name	Type	Description	Default
height	int		required
width	int		required
out_height	int		required
out_width	int		required
chunk	int	Every chunk is a handshake with every core: few and large ones.	4096
words	int	A 39-tap window: a scale of up to 9.5.	356
patch_columns	int	A floor of 8, so two resizes of one image that need no more share an array.	lambda op: max(8, op.default_patch_columns())
image_chunks	int		lambda op: op.height * -(-3 * op.width // op.chunk)
width_chunks	int	Multiples of 16, as ResampleTaps's are.	lambda op: 16 * -(-op.out_width // (16 * SIDE))
height_chunks	int		lambda op: 16 * -(-op.out_height // (16 * SIDE))
patch_rows	int		lambda op: op.padded_rows()
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 2.	None

Attributes:

Name	Type	Description
rows		Per call.
columns		Per call.
out_rows		Per call.
out_columns		Per call.
image		In(image_chunks, chunk, dtype=np.uint8, obj_shape=(chunk,))
taps_w		In(width_chunks, words, dtype=np.int32, obj_shape=(words,))
taps_h		In(height_chunks, words, dtype=np.int32) Through taps_w's stream, after it.
patches		Out(patch_rows, SIDE * SIDE * 3, dtype=bfloat16, obj_shape=(SIDE * SIDE * 3,), per=(num_aie_columns, num_channels))
src_rows		Per call.
src_columns		Per call.
dst_rows		Per call.
dst_columns		Per call.
image_count		Per call.
width_count		Per call.
height_count		Per call.
patch_count		Per call.

Source code in iron/operators/resample/op.py

class Resample(Operator):

window_cap property

window_cap: int

The widest filter window a chunk of words holds 16 of.

default_patch_columns

default_patch_columns() -> int

Return patch columns a core holds: out_width's split over the cores.

And a zero one past them.

Source code in iron/operators/resample/op.py

def default_patch_columns(self) -> int:
    """Return patch columns a core holds: `out_width`'s split over the cores.

    And a zero one past them.
    """
    dev = self.dev
    cols = self.num_aie_columns or (dev.cols if dev is not None else 8)
    return self.out_width // SIDE // (cols * self.num_channels) + 1

padded_rows

padded_rows() -> int

Every band's patch columns, padded past the cores.

Source code in iron/operators/resample/op.py

def padded_rows(self) -> int:
    """Every band's patch columns, padded past the cores."""
    dev = self.dev
    cols = self.num_aie_columns or (dev.cols if dev is not None else 8)
    cores = cols * self.num_channels
    return self.out_height // SIDE * self.default_patch_columns() * cores

held

held(dev) -> int

Core memory a core of dev holds at patch_columns.

Source code in iron/operators/resample/op.py

def held(self, dev) -> int:
    """Core memory a core of `dev` holds at `patch_columns`."""
    cols, win = self.patch_columns, self.window_cap
    plane = (5 * win + 8 + 64 + 63) // 64 * 64
    row = (cols * SIDE * 3 + 31) // 32 * 32  # LINE
    kernel = (
        14 * 4  # resize_setup's own scalars
        + cols * SIDE * 4  # hstart
        + cols * SIDE * 64 * 2  # hweight, two taps vectors a slot
        + 2 * cols * 4  # span0, span
        + cols * 3 * plane  # line
        + win * row  # mid
        + SIDE * row  # hold
        + 2 * 256  # PIXEL
    )
    buffers = (
        2 * self.chunk  # image, depth 2
        + 2 * self.words * 4  # taps, depth 2
        + 2 * SIDE * SIDE * 3 * 2  # patches, depth 2
        + 5 * 4  # counts
    )
    return kernel + buffers + dev.default_core_stack_bytes

reference

reference(image, taps_w, taps_h, *, rows=None, columns=None, out_rows=None, out_columns=None, patches=None)

Return the kernel's patches (vision.resize_ref).

Given patches, written over its leading rows, the ones a call writes, the rest kept.

Source code in iron/operators/resample/op.py

def reference(
    self,
    image,
    taps_w,
    taps_h,
    *,
    rows=None,
    columns=None,
    out_rows=None,
    out_columns=None,
    patches=None,
):
    """Return the kernel's patches (`vision.resize_ref`).

    Given `patches`, written over its leading rows, the ones a call
    writes, the rest kept.
    """
    written = vision.resize_ref(
        np.asarray(image),
        np.asarray(taps_w),
        np.asarray(taps_h),
        self.height if rows is None else int(rows),
        self.width if columns is None else int(columns),
        self.out_height if out_rows is None else int(out_rows),
        self.out_width if out_columns is None else int(out_columns),
        words=self.words,
        chunk=self.chunk,
        patch_columns=self.patch_columns,
        cores=self.cores,
    )
    if patches is None:
        return written
    patches[: len(written)] = written
    return patches

sequence

sequence(rt)

Run the image, then both tables down one stream, and each core's patches.

Source code in iron/operators/resample/op.py

def sequence(self, rt):
    """Run the image, then both tables down one stream, and each core's patches."""
    tg = TaskGroup()
    moves = [
        (self.image, None, self.image, "image_count", 0, 1),
        (self.taps_w, None, self.taps_w, "width_count", 0, 1),
        (self.taps_w, None, self.taps_h, "height_count", 0, 1),
    ] + [
        (self.patches, k, self.patches, "patch_count", k, self.cores)
        for k in range(self.cores)
    ]
    for operand, lane, buffer, name, first, step in moves:
        per_call = self.value(name).per_call
        n, width = buffer.shape
        count = n // step if per_call else self.residents[name]
        tap = TensorAccessPattern(
            (n, width), first * width, [1, count, 1, width], [0, step * width, 0, 1]
        )
        length = self.value(name) if per_call else None
        if operand.drains:
            operand.drain(
                tap,
                lane=lane,
                group=tg,
                wait=True,
                length_parameter=length,
                length_unit=width,
            )
        else:
            operand.fill(
                tap,
                lane=lane,
                group=tg,
                source=buffer,
                length_parameter=length,
                length_unit=width,
            )
    tg.finish()

ResampleTaps

ResampleTaps(*, in_size: int, out_size: int, words: int = 256, slots: int | None = None, chunks: int = lambda op: 16 * -(-op.out_size // (16 * max(1, op.per_chunk(op.in_size, op.out_size)))), num_aie_columns: int = None, num_channels: int = None)

Bases: Operator

The int16 filter table and precision torch resamples one axis with.

in_samples samples to out_samples.

Each core takes chunks k, k + cores, ... and computes their weights in float64 twice, as vision.resample_peak and vision.resample_quantize do: once for the largest weight, which a chain through the cores reduces, and once to quantize at the precision it sets.

Parameters:

Name	Type	Description	Default
in_size	int		required
out_size	int		required
words	int		256
slots	int | None	The slots of a chunk, or None for as many as the window lets fit.	None
chunks	int	A multiple of 16, so every core count up to it divides.	lambda op: 16 * -(-op.out_size // (16 * max(1, op.per_chunk(op.in_size, op.out_size))))
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 2.	None

Attributes:

Name	Type	Description
table		Out(chunks, words, dtype=np.int32, obj_shape=(words,), per=(num_aie_columns, num_channels))
in_samples		Per call.
out_samples		Per call.
count		Per call.
in_length		Per call.
out_length		Per call.

Source code in iron/operators/resample/op.py

class ResampleTaps(Operator):

per_chunk

per_chunk(n_in: int, n_out: int) -> int

Return the outputs a chunk holds of n_in samples to n_out (per_chunk).

Source code in iron/operators/resample/op.py

def per_chunk(self, n_in: int, n_out: int) -> int:
    """Return the outputs a chunk holds of `n_in` samples to `n_out` (`per_chunk`)."""
    return per_chunk(n_in, n_out, self.words, self.slots)

filled

filled(n_in: int, n_out: int) -> int

Return the chunks a table of n_in samples to n_out is written in.

Parameters:

Name	Type	Description	Default
n_in	int	The input samples on the axis.	required
n_out	int	The output samples on the axis.	required

Returns:

Type	Description
int	Those its outputs fill, in a multiple of 16 as chunks is; every
int	chunk for sizes the table cannot hold, each header saying so. The
int	cores write whole rounds, so up to a round past them.

Source code in iron/operators/resample/op.py

def filled(self, n_in: int, n_out: int) -> int:
    """Return the chunks a table of `n_in` samples to `n_out` is written in.

    Args:
        n_in: The input samples on the axis.
        n_out: The output samples on the axis.

    Returns:
        Those its outputs fill, in a multiple of 16 as `chunks` is; every
        chunk for sizes the table cannot hold, each header saying so. The
        cores write whole rounds, so up to a round past them.
    """
    per = 0
    if min(n_in, n_out) >= 1:
        per = self.per_chunk(n_in, n_out)
    if not per:
        return self.chunks
    return min(16 * -(-n_out // (16 * per)), self.chunks)

sequence

sequence(rt)

Chunk c from core c % cores, the order the cores compute them in.

Source code in iron/operators/resample/op.py

def sequence(self, rt):
    """Chunk `c` from core `c % cores`, the order the cores compute them in."""
    per_call = self.count.per_call
    n, words = self.table.shape
    count = n // self.cores if per_call else self.residents["count"]
    tg = TaskGroup()
    for k in range(self.cores):
        tap = TensorAccessPattern(
            (n, words),
            k * words,
            [1, count, 1, words],
            [0, self.cores * words, 0, 1],
        )
        self.table.drain(
            tap,
            lane=k,
            group=tg,
            wait=True,
            length_parameter=self.count if per_call else None,
            length_unit=words,
        )
    tg.finish()

RMSNorm

RMSNorm(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int, rows: int, weighted: bool = False, epsilon: float = 1e-05)

Bases: WeightedRowwise

AIE-accelerated RMS Normalization of each row, optionally times a learned weight row.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int		required
rows	int		required
weighted	bool		False
epsilon	float	The epsilon under the root: 1e-5 by default; a model states its own.	1e-05

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(Rowwise.rows, Rowwise.tile_size, obj_shape=(Rowwise.tile_size,), per=CORES)
y		Out(Rowwise.rows, Rowwise.tile_size, obj_shape=(Rowwise.tile_size,), per=CORES)
weight		In(Rowwise.tile_size, obj_shape=(Rowwise.tile_size,), per=(Elementwise.num_channels,), replicate=True, when=weighted) Shared by every column of a channel: counted once per channel.

Source code in iron/operators/rms_norm.py

class RMSNorm(WeightedRowwise):

RoPE

RoPE(*, rows: int, cols: int, angle_rows: int = lambda op: op.rows, num_aie_columns: int = None, two_halves: bool = True)

Bases: Operator

AIE-accelerated RoPE: one core per column, each rotating rows of cols.

Each angle row serves rows / angle_rows consecutive rows.

Parameters:

Name	Type	Description	Default
rows	int		required
cols	int		required
angle_rows	int		lambda op: op.rows
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
two_halves	bool		True

Attributes:

Name	Type	Description
x		In(rows, cols, obj_shape=(1, cols), per=(num_aie_columns,))
angles		In(angle_rows, cols, obj_shape=(1, cols), per=(num_aie_columns,))
y		Out(rows, cols, obj_shape=(1, cols), per=(num_aie_columns,))
valid		Per call. Either row count may be bounded per call: x[:n] and angles[:m].
valid_angles		Per call.
lut_rows		Per call.
rows_per_lut		Per call.

Source code in iron/operators/rope.py

class RoPE(Operator):

extent_unit

extent_unit(buffer: str) -> int | None

Return a position's rows at a time, so each core gets its angle row.

One core takes the rows a call bounds one at a time.

Source code in iron/operators/rope.py

def extent_unit(self, buffer: str) -> int | None:
    """Return a position's rows at a time, so each core gets its angle row.

    One core takes the rows a call bounds one at a time.
    """
    if buffer not in ("x", "y"):
        return None
    return 1 if self.num_aie_columns == 1 else self.rows // self.angle_rows

reference

reference(x, angles)

datamovement.rope_ref on the packed angles buffer.

Each angle row serving rows / angles.shape[0] consecutive rows of x.

Source code in iron/operators/rope.py

def reference(self, x, angles):
    """``datamovement.rope_ref`` on the packed ``angles`` buffer.

    Each angle row serving ``rows / angles.shape[0]`` consecutive rows
    of ``x``.
    """
    rows, cols, lut_rows = x.shape[0], x.shape[-1], angles.shape[0]
    if rows % lut_rows != 0:
        raise ValueError(
            f"RoPE: angles ({lut_rows} rows) must divide x's {rows} rows"
        )
    # The table broadcast over the rows each angle row serves: no copy.
    y = datamovement.rope_ref(
        x.reshape(lut_rows, -1, cols, copy=False),
        angles.reshape(lut_rows, 1, cols, copy=False),
        two_halves=self.two_halves,
    )
    return y.reshape(x.shape, copy=False)

Sample

Sample(*, vocab: int, k_max: int = 64, drawable: int = lambda op: op.vocab, steps: int = 1, cores: int = None, chunk: int = None)

Bases: Operator

Draw the next token from a row of logits (exact temperature/top-k).

Parameters:

Name	Type	Description	Default
vocab	int		required
k_max	int		64
drawable	int	The logits that are tokens'; the rest pad an odd vocabulary to whole words, and the last select core writes them -inf.	lambda op: op.vocab
steps	int	Positions the draws and the record hold; the sequence picks one.	1
cores	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 4.	None
chunk	int	Tunable: the call's, else the profile's, else resolve(dev)'s. Logits per select call.	None

Attributes:

Name	Type	Description
logits		In(vocab, obj_shape=(chunk,), per=(cores,))
draws		In(steps, sample.ROW_WORDS, dtype=np.int32, obj_shape=(sample.ROW_WORDS,), depth=1)
tokens		InOut(steps, dtype=np.int32, obj_shape=(1,), depth=1)
token		Out(1, dtype=np.int32, obj_shape=(1,), depth=1)
row		Per call. Element offsets of this position's draw row and record.
at		Per call.

Source code in iron/operators/sample.py

class Sample(Operator):

runs property

runs: list[tuple[int, int, int]]

Return each select core's (logits, chunk, drawable).

The last reads what is left, the padding with it.

sequence

sequence(rt)

One group: the draw and the logits in, then the record and the token.

Source code in iron/operators/sample.py

def sequence(self, rt):
    """One group: the draw and the logits in, then the record and the token."""
    row = self.row if self.row.per_call else None
    at = self.at if self.at.per_call else None
    tg = TaskGroup()
    self.draws.fill(self.draws[0], group=tg, offset_parameter=row)
    for c, (run, chunk, _) in enumerate(self.layout().runs):
        # Core c's slice, once per select stream.
        streams = sample.select_streams(run, chunk)
        start = c * self.slice_size
        piece = TensorAccessPattern.full((self.vocab,))[start : start + run]
        halves = self.logits.split(run)
        tap = halves and piece.split(0, halves[1]).repeat(streams)
        taps = [tap] if tap and self.logits.fits(tap) else [piece] * streams
        for tap in taps:
            self.logits.fill(tap, lane=c, group=tg)
    self.tokens.drain(self.tokens[:1], group=tg, wait=True, offset_parameter=at)
    self.token.drain(group=tg, wait=True)
    tg.finish()

reference

reference(logits, draws, tokens, *, row=0, at=0)

(tokens, token): the draw at draws row row // 4, recorded at at.

Source code in iron/operators/sample.py

def reference(self, logits, draws, tokens, *, row=0, at=0):
    """``(tokens, token)``: the draw at ``draws`` row ``row // 4``, recorded at ``at``."""
    token = reference(
        np.asarray(logits).reshape(-1)[: self.drawable],
        np.asarray(draws).reshape(-1)[row : row + sample.ROW_WORDS],
        self.k_max,
    )
    tokens = np.array(tokens, dtype=np.int32).reshape(-1)
    tokens[at] = token
    return tokens, np.array([token], dtype=np.int32)

SelectiveScan

SelectiveScan(*, rows: int, heads: int, head_dim: int, state_size: int, groups: int = 1, columns: int = None, inner: int = lambda op: op.heads * op.head_dim, width: int = lambda op: op.inner + 2 * op.groups * op.state_size, block: int = None, slice_state: int = lambda op: SLICE * op.state_size)

Bases: Operator

Each row t of xbc updates each head's state (head_dim, state_size) and reads it out.

The row splits into x (heads, head_dim) and groups rows each of B and C (state_size,):

dt = softplus(dt[t] + dt_bias)
state = exp(dt * A) * state + dt * outer(x, B)
y[t] = state @ C + D * x

head h reading group h // (heads // groups)'s B and C and its (dt_bias, A, D) from head. state is the state before the first row; final the state after the last valid one. Float32 throughout, y rounded once.

Each of columns cores takes whole heads, SLICE rows of a head's state at a time, and gathers a block of rows' x, B, C and dt for the slice from whole rows of xbc, which one memtile joins to each row's dt and every core reads; another joins each slice's state to its head's head, so a core reads two streams. The kernel keeps the state from block to block. Scalar.

Parameters:

Name	Type	Description	Default
rows	int		required
heads	int		required
head_dim	int		required
state_size	int		required
groups	int		1
columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
inner	int	Derived where not given; a given value must equal it.	lambda op: op.heads * op.head_dim
width	int	Derived where not given; a given value must equal it.	lambda op: op.inner + 2 * op.groups * op.state_size
block	int	Tunable: the call's, else the profile's, else resolve(dev)'s. The rows of y one kernel call writes: whole SLICE-row runs that fit L1.	None
slice_state	int	Derived where not given; a given value must equal it.	lambda op: SLICE * op.state_size

Attributes:

Name	Type	Description
xbc		In(rows, width, obj_shape=(width,))
dt		In(rows, heads, obj_shape=(heads,))
head		In(heads, 3, dtype=np.float32, obj_shape=(3,), per=(columns,))
state		In(heads, head_dim, state_size, dtype=np.float32, obj_shape=(slice_state,), per=(columns,), depth=1)
y		Out(rows, inner, obj_shape=(block,), per=(columns,))
final		Out(heads, head_dim, state_size, dtype=np.float32, obj_shape=(slice_state,), per=(columns,), depth=1)
valid		Per call.
count		Per call.

Source code in iron/operators/selective_scan.py

class SelectiveScan(Operator):

core_slices property

core_slices: int

A core's slices: its heads' SLICE-row runs of state.

slice_bytes

slice_bytes() -> tuple[int, int]

Bytes of a slice's (dt_bias, A, D) then state in, and its final out.

Source code in iron/operators/selective_scan.py

def slice_bytes(self) -> tuple[int, int]:
    """Bytes of a slice's (dt_bias, A, D) then state in, and its final out."""
    return 4 * (3 + self.slice_state), 4 * self.slice_state

held

held(dev, rows: int) -> int

Return the bytes a core of dev uses at rows rows per call.

That is the input row, the block buffer, y, the kernel's data and stack (its own copy of the state included), and the device's default stack, which covers the core's own frame. One slice's state and final state are not counted.

Source code in iron/operators/selective_scan.py

def held(self, dev, rows: int) -> int:
    """Return the bytes a core of ``dev`` uses at ``rows`` rows per call.

    That is the input row, the block buffer, y, the kernel's data and
    stack (its own copy of the state included), and the device's default
    stack, which covers the core's own frame. One slice's state and final
    state are not counted.
    """
    # Only measured, never built, so it stays out of the build's registry.
    with compile_context(_iron_kernel_registry=set()):
        kernel = self.kernel(rows)
        block = 2 * int(np.prod(kernel.arg_shape(0)))
        frame = kernel.frame_bytes(resolve_target_arch(dev))
    frame += dev.default_core_stack_bytes
    return 2 * (self.width + self.heads) + block + 4 * rows * SLICE + frame

sequence

sequence(rt)

Every row of xbc and dt once a slice, for every core.

Per core, its heads' parameters and states, and its slices of y and final.

Source code in iron/operators/selective_scan.py

def sequence(self, rt):
    """Every row of ``xbc`` and ``dt`` once a slice, for every core.

    Per core, its heads' parameters and states, and its slices of ``y``
    and ``final``.
    """
    R, H, W, d_inner = self.rows, self.heads, self.width, self.inner
    P, SN = self.head_dim, SLICE * self.state_size
    S, heads = self.core_slices, H // self.columns
    per, states, run = S // heads, H * P * self.state_size, S * SN
    # A core's heads' states, a slice of each head in turn.
    slices = [
        TensorAccessPattern(
            (states,), k * run, [per, heads, SN], [SN, P * self.state_size, 1]
        )
        for k in range(self.columns)
    ]
    tg = TaskGroup()
    self.xbc.fill(TensorAccessPattern((R, W), 0, [S, R, W], [0, W, 1]), group=tg)
    self.dt.fill(TensorAccessPattern((R, H), 0, [S, R, H], [0, H, 1]), group=tg)
    for k in range(self.columns):
        self.head.fill(
            TensorAccessPattern((H, 3), 3 * k * heads, [per, heads, 3], [0, 3, 1]),
            lane=k,
            group=tg,
        )
        self.state.fill(slices[k], lane=k, group=tg)
    for k in range(self.columns):
        self.y.drain(
            TensorAccessPattern(
                (R, d_inner),
                k * heads * P,
                [per, heads, R, SLICE],
                [SLICE, P, d_inner, 1],
            ),
            lane=k,
            group=tg,
            wait=True,
        )
        self.final.drain(slices[k], lane=k, group=tg, wait=True)
    tg.finish()

ops

ops() -> int

Count per row, per head, the state update and the C readout.

The update is two multiplies and an add over its head_dim * state_size elements; the readout a state_size-deep dot product per head_dim row.

Source code in iron/operators/selective_scan.py

def ops(self) -> int:
    """Count per row, per head, the state update and the C readout.

    The update is two multiplies and an add over its
    ``head_dim * state_size`` elements; the readout a
    ``state_size``-deep dot product per ``head_dim`` row.
    """
    return 5 * self.rows * self.heads * self.head_dim * self.state_size

Sigmoid

Sigmoid(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = None, size: int)

Bases: UnaryElementwise

AIE-accelerated Sigmoid activation function.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
size	int		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
y		Out(like=x, obj_shape=(Elementwise.tile_size,), per=CORES)

Source code in iron/operators/sigmoid.py

class Sigmoid(UnaryElementwise):

SiLU

SiLU(*, num_aie_columns: int = None, tile_size: int = None, size: int)

Bases: UnaryElementwise

AIE-accelerated SiLU activation function.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
size	int		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
y		Out(like=x, obj_shape=(Elementwise.tile_size,), per=CORES)

Source code in iron/operators/silu.py

class SiLU(UnaryElementwise):

Softmax

Softmax(*, rows: int, cols: int, num_aie_columns: int = None, num_channels: int = None, block: int = None, streamed: bool = None, input_dtype: Any = ml_dtypes.bfloat16, soft_cap: float | None = None, has_addend: bool = False, window: int | None = None, period: int | None = None)

Bases: Operator

Softmax over each row, one core per (column, channel), each taking a contiguous run of rows.

streamed, a core reads a row in block-sized chunks three times (maximum, sums, output); a per-call bound zeroes elements past it and only its blocks move.

Parameters:

Name	Type	Description	Default
rows	int		required
cols	int		required
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
block	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
streamed	bool	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
input_dtype	Any		ml_dtypes.bfloat16
soft_cap	float | None	Float input only, mlir-aie's softmax_f32_bf16: a pre-cap addend and cap * tanh(x / cap).	None
has_addend	bool		False
window	int | None	Float input only, softmax_rows_banded_f32_bf16: a band computed from these two in the core (window, and the row's own query position, wrapped every period rows), not a mask operand -- a core has two input DMA channels, and x and addend already take both.	None
period	int | None		None

Attributes:

Name	Type	Description
length		Per call. cols, or fewer per call
x		In(rows, cols, dtype=input_dtype, obj_shape=(block,), per=(num_aie_columns, num_channels))
addend		In(rows, cols, dtype=np.float32, obj_shape=(block,), per=(num_aie_columns, num_channels), when=has_addend) A float32 row added to x before the cap, such as relative-position logits.
y		Out(rows, cols, obj_shape=(block,), per=(num_aie_columns, num_channels))
valid_cols		Per call.
blocks		Per call. A row held whole is one block, which the bound does not change.

Source code in iron/operators/softmax.py

class Softmax(Operator):

sequence

sequence(rt)

Each lane's rows as (block, row, element), thrice if streamed.

Source code in iron/operators/softmax.py

def sequence(self, rt):
    """Each lane's rows as ``(block, row, element)``, thrice if streamed."""
    lanes = self.cores
    rows = self.rows // lanes
    blocks = ceildiv(self.length, self.block)
    bounded = self.streamed and self.blocks.per_call
    blocks_valid = self.blocks if bounded else None
    full = TensorAccessPattern.full((self.rows, self.cols))
    taps = [
        full[k * rows : (k + 1) * rows, : blocks * self.block]
        .split(1, self.block)
        .permute((1, 0, 2))
        for k in range(lanes)
    ]
    reads = 3 if self.streamed else 1
    tg = TaskGroup()
    for k, tap in enumerate(taps):
        self.x.fill(
            tap.repeat(reads),
            lane=k,
            group=tg,
            length_parameter=blocks_valid,
            length_unit=rows * self.block,
        )
    if self.has_addend:
        for k, tap in enumerate(taps):
            self.addend.fill(
                tap.repeat(1),
                lane=k,
                group=tg,
                length_parameter=blocks_valid,
                length_unit=rows * self.block,
            )
    for k, tap in enumerate(taps):
        self.y.drain(
            tap.repeat(1),
            lane=k,
            group=tg,
            wait=True,
            length_parameter=blocks_valid,
            length_unit=rows * self.block,
        )
    tg.finish()

reference

reference(x, *rest)

Return the softmax kernel's contract reference, a row at a time.

rest is addend where has_addend, positional; the band, where given, comes from the factory's own banded contract, row_offset each row's index.

Source code in iron/operators/softmax.py

def reference(self, x, *rest):
    """Return the softmax kernel's contract reference, a row at a time.

    ``rest`` is ``addend`` where ``has_addend``, positional; the
    band, where given, comes from the factory's own banded contract,
    ``row_offset`` each row's index.
    """
    addend = rest[0] if self.has_addend else None
    # A bounded call's first ``length`` columns, the rest left zero.
    n = self.length
    x = x[..., :n]
    addend = None if addend is None else addend[..., :n]
    if self.window is not None:
        y = _kernel(self, n).contract.reference(x, addend, np.arange(self.rows))
    else:
        y = activation.softmax_ref(
            x, addend, None, tile_size=n, soft_cap=self.soft_cap
        )
    if n == self.cols:
        return y
    out = np.zeros((*np.shape(y)[:-1], self.cols), np.asarray(y).dtype)
    out[..., :n] = y
    return out

Tanh

Tanh(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = None, size: int)

Bases: UnaryElementwise

AIE-accelerated Tanh activation function.

Parameters:

Name	Type	Description	Default
num_aie_columns	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
num_channels	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1.	None
tile_size	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
size	int		required

Attributes:

Name	Type	Description
count		Per call.
valid		Per call.
x		In(size, obj_shape=(Elementwise.tile_size,), per=CORES)
y		Out(like=x, obj_shape=(Elementwise.tile_size,), per=CORES)

Source code in iron/operators/tanh.py

class Tanh(UnaryElementwise):

FastFlowLM ports

Catalogued as "flm.<name>". flm.ProjectionGEMM is not iron.operators.GEMM.

Ports of FastFlowLM's Gemma 4 NPU overlays to IRON.

https://github.com/ROCm/FastFlowLM, MIT.

DequantBFP

DequantBFP(*, K: int, N: int, run_out_features: int | None = None, run_period_out_features: int | None = None, quantized_bytes: int = lambda op: op.quantized_size(), packed_blocks: int = lambda op: op.K * op.N // bfp.BLOCK, cols: int = None)

Bases: Operator

q4nx weights to bfp16, packed the way flm.ProjectionGEMM reads B.

A 4-row grid, as wide as the device, for one q4nx tiling. Nothing depends on K or N, since every core loops forever over identical per-block work, so one xclbin serves every shape; K, N and the interleave move only the instruction stream's offsets. See README.md for the layout and constraints.

Parameters:

Name	Type	Description	Default
K	int		required
N	int		required
run_out_features	int | None	A projection interleaved with another in the same buffer: this one occupies run_out_features of every run_period_out_features.	None
run_period_out_features	int | None		None
quantized_bytes	int	Derived where not given; a given value must equal it. The q4nx input is bytes, as a block interleaves three tables at 5 bits a weight; the output is bfp16ebs8 blocks, as the drains index them.	lambda op: op.quantized_size()
packed_blocks	int	Derived where not given; a given value must equal it.	lambda op: op.K * op.N // bfp.BLOCK
cols	int	Tunable: the call's, else the profile's, else resolve(dev)'s. cols follows the device. ROWS is baked into the split offsets and the join, so it is not a field; halves is, because a stream's replication count must be declared to be indexed, but the join bakes it too.	None

Attributes:

Name	Type	Description
qw		In(quantized_bytes, dtype=np.uint8, obj_shape=(ROWS * BLOCK_BYTES,), per=(cols,), depth=2) One q4nx block per core, delivered as one per-column object the cores split; one packed half-tile out per (column, n-half), joined from the two cores that share it.
out		Out(packed_blocks, dtype=v8bfp16ebs8, obj_shape=(HALF_BLOCKS,), per=(cols, halves), depth=2)

Source code in iron/operators/flm/dequant/op.py

class DequantBFP(Operator):

config_name property

config_name: str

Stem of the artifacts that do not depend on the shape: the xclbin's.

name property

name: str

The configuration's stem, then the shape.

Two instruction streams never share a name.

column_blocks property

column_blocks: list[int]

Where each of the weight's column blocks sits in the blob, in column blocks.

run_out_features of every run_period_out_features.

quantized_size

quantized_size() -> int

Bytes of q4nx input, counting any interleave gap it strides over.

Source code in iron/operators/flm/dequant/op.py

def quantized_size(self) -> int:
    """Bytes of q4nx input, counting any interleave gap it strides over."""
    return packed_bytes((self.column_blocks[-1] + 1) * N_TILE * self.K)

xclbin_operator

xclbin_operator()

Return this configuration at its reference shape.

Its xclbin serves every shape sharing the configuration, and only the instruction stream is per shape, as flm.ProjectionGEMM's.

Source code in iron/operators/flm/dequant/op.py

def xclbin_operator(self):
    """Return this configuration at its reference shape.

    Its xclbin serves every shape sharing the configuration, and only
    the instruction stream is per shape, as flm.ProjectionGEMM's.
    """
    K, N = self._reference_shape
    return self.replace(
        K=K,
        N=N,
        run_out_features=None,
        run_period_out_features=None,
        quantized_bytes=None,
        packed_blocks=None,
    )

grid

grid(qw)

Return the weight's q4nx blocks in qw, as q4nx.unpack's grid.

A column block is its k-tiles in order, each k-tile its two k-halves, each of those its two n-halves (README.md layers 6-9).

Source code in iron/operators/flm/dequant/op.py

def grid(self, qw):
    """Return the weight's q4nx blocks in ``qw``, as ``q4nx.unpack``'s grid.

    A column block is its k-tiles in order, each k-tile its two
    k-halves, each of those its two n-halves (README.md layers 6-9).
    """
    cb_bytes = packed_bytes(N_TILE * self.K)
    qw = np.asarray(qw, np.uint8).reshape(-1, cb_bytes)[self.column_blocks]
    qw = qw.reshape(self.N // N_TILE, self.K // K_TILE_B, 2, 2, BLOCK_BYTES)
    return qw.transpose(0, 3, 1, 2, 4).reshape(
        self.N // M_TILE, self.K // K_TILE, BLOCK_BYTES
    )

reference

reference(qw)

Return the bytes the operator must produce from the blob qw.

A flat uint8 array, bit-exact against the device, since the cores round to bf16 toward negative infinity (README.md has the layout and rounding).

Source code in iron/operators/flm/dequant/op.py

def reference(self, qw):
    """Return the bytes the operator must produce from the blob ``qw``.

    A flat uint8 array, bit-exact against the device, since the cores
    round to bf16 toward negative infinity (README.md has the layout
    and rounding).
    """
    K, N = self.K, self.N
    grid = self.grid(qw)
    enc = quant.q4nx_dequant_ref(
        grid.reshape(-1, BLOCK_BYTES),
        m_tile=M_TILE,
        k_tile=K_TILE,
        group=GROUP,
        ct_k=CT_K,
        s=S,
        t=T,
    )
    # Each block's bytes are indexed [k slice, n // T, k // S in the slice,
    # n % T, 9]. Lay the grid's blocks out as one such array for the whole
    # matrix.
    enc = enc.reshape(
        N // M_TILE, K // K_TILE, K_TILE // CT_K, M_TILE // T, CT_K // S, T, 9
    )
    out = enc.transpose(1, 2, 0, 3, 4, 5, 6)
    # pack_b's order: (cb, kb, k slice, n // T in cb, k // S in the slice,
    # n % T).
    out = out.reshape(
        K // K_TILE_B, K_TILE_B // CT_K, N // N_TILE, N_TILE // T, CT_K // S, T, 9
    )
    return out.transpose(2, 0, 1, 3, 4, 5, 6).ravel()

ProjectionGEMM

ProjectionGEMM(*, M: int, K: int, N: int, epilogue: Epilogue | str = Epilogue.NONE, packed_blocks: int = lambda op: op.K * op.N // bfp.BLOCK, tile_n: int = None, k_tile: int = None, ct_k: int = None, tile_ma: int = None, m_chunk: int = None, epilogue_modes: tuple = tuple(Epilogue), rounding: Rounding | str = Rounding.CONV_EVEN, gelu: Gelu | str = Gelu.FP32, b_col_maj: bool = False, precision: Precision, bfp16_b: bool = False)

Bases: Operator

AIE-accelerated bf16 GEMM on a 4-row grid, as wide as the device.

With a fused epilogue. M, K, N and the activation are values written to the cores, so every shape on one config_name shares an xclbin. Used wherever iron.operators.GEMM doesn't apply (a weight operand, MEDIUM/HIGH, two or more rows).

Parameters:

Name	Type	Description	Default
M	int		required
K	int		required
N	int		required
epilogue	Epilogue | str	The activation, selected at run time from the compiled-in modes; its gelu is x * sigmoid(1.702x), the port's own arithmetic.	Epilogue.NONE
packed_blocks	int	Derived where not given; a given value must equal it. B's packed block count on AIE2P: blocks, not bytes, since B's declaration counts bfp16ebs8 blocks and bfp.itemsize turns that back into the byte count pack_b returns.	lambda op: op.K * op.N // bfp.BLOCK
tile_n	int	Tunable: the call's, else the profile's, else resolve(dev)'s. 64 halves the mmul's accumulator traffic per mac; 128 halves A fetches.	None
k_tile	int	Tunable: the call's, else the profile's, else resolve(dev)'s. k tile. K must be a multiple of it, and it of the compute tile's k slice.	None
ct_k	int	Tunable: the call's, else the profile's, else resolve(dev)'s. How much of the k tile a compute tile holds at a time, the kernel's chunk_k. None resolves to the deepest that fits L1 (see l1_tiles); the search tries the slice either side of it.	None
tile_ma	int	Tunable: the call's, else the profile's, else resolve(dev)'s. A-tile rows, decoupled from the accumulator's M_TILE (asymmetric tile buffering). None resolves to whatever L1 affords.	None
m_chunk	int	Tunable: the call's, else the profile's, else resolve(dev)'s, else 1. Row-blocks a core folds into one B fetch, cutting B's DDR reads by m_chunk at the cost of that many L1 accumulators. 1 by default: it must divide m_row_blocks (M % 512 == 0), and a shape that cannot use it would fork config_name. See README.md.	None
epilogue_modes	tuple	The activations the epilogue can select between at run time. Each one compiled in costs program memory, so a deployment that dispatches two should compile two.	tuple(Epilogue)
rounding	Rounding | str		Rounding.CONV_EVEN
gelu	Gelu | str		Gelu.FP32
b_col_maj	bool	B as stored (N, K), a checkpoint's layout, read as is: the memtile lays out the k panels the kernel reads as B^T blocks, so no packed copy.	False
precision	Precision	MEDIUM multiplies on AIE2P's bfp16 macs, A and B rounded to bfp16ebs8 (the port's arithmetic); above it, or on AIE2, on bf16 macs.	required
bfp16_b	bool	B in DDR as pack_b's bfp16ebs8 blocks, which AIE2P's bfp16 macs multiply unconverted: a packed B at precision='medium' there must say so, since B's host shape depends on it and no device is known yet.	False

Attributes:

Name	Type	Description
A		In(M, K, obj_shape=(a_l2,), per=(rows,), depth=A_DEPTH)
B		In(Select(bfp16_b, (packed_blocks,), (Select(b_col_maj, (N, K), (K, N)),)), dtype=b_dtype, obj_shape=(b_l2,), per=(cols,)) For bfp16 macs B is counted in bfp16ebs8 blocks, the unit every descriptor into it counts in: in bytes, a transfer would move a ninth of what it named.
C		Out(M, N, obj_shape=(c_l2,), per=(cols, joins), depth=C_DEPTH)
valid		Per call. M, or fewer rows per call (A[:n] in a graph). A and C move, and the cores compute, only the units (rows * M_TILE rows each) the bound reaches; past them the cores drain each column-block's B replay uncomputed, since its count is in the memtile's lock values, which no per-call word reaches. B moves as it does unbounded.
n_val		Per call. The parameter words every core reads once a call. The unit counts are read only at m_chunk > 1: each word costs every dispatch (README.md, "A word is not free"). The clamp bounds are read always, the identity: the kernel's clamp is not compiled out.
m_row_blocks		Per call.
k_iters		Per call.
mode		Per call.
clamp_min		Per call.
clamp_max		Per call.
n_units		Per call.
units_valid		Per call. Under a bound, the units it reaches: the cores' trip count and C's length, in row blocks.
units_repeat		Per call. A's repeat count on a full ELF: the units less one.
b_passes_repeat		Per call. Under a bound, the passes B's memtile ring takes on a column-block, less one: a streamed ring's per call, as DDR reads B once per unit reached.
b_resident		Per call. Whether the memtile replays B to every unit, as it does where B is resident, so that the cores drain the replays past a bound.
a_half		Per call. Under a bound, where the second half of the core rows starts in A and in C: past the first half's share of the units it reaches.
c_half		Per call.

Source code in iron/operators/flm/gemm/op.py

class ProjectionGEMM(Operator):

compiled_modes property

compiled_modes: tuple[Epilogue, ...]

Return the activations compiled into the epilogue.

None and epilogue_modes.

epilogue_mask property

epilogue_mask: int

Return the bitmask of the modes compiled into the epilogue.

Mode 0 is always present; the kernel falls back to it.

config_name property

config_name: str

Stem of the artifacts that do not depend on the shape: the xclbin's.

name property

name: str

Return the configuration's stem, then every runtime parameter.

Two instruction streams never share a name.

frame

frame(dev, tile_n, k_tile, ct_k, band_m, bfp16_b, emulate) -> int

Return what a core of dev holds beside its buffers at these tiles.

Its kernel's data and deepest call's stack, over the default stack. Kept per kernel, as every resolve asks it of up to six.

Source code in iron/operators/flm/gemm/op.py

def frame(self, dev, tile_n, k_tile, ct_k, band_m, bfp16_b, emulate) -> int:
    """Return what a core of ``dev`` holds beside its buffers at these tiles.

    Its kernel's data and deepest call's stack, over the default stack.
    Kept per kernel, as every resolve asks it of up to six.
    """
    tiles = (tile_n, k_tile, ct_k, band_m, bfp16_b, emulate)
    key = (
        dev,
        tiles,
        self.b_col_maj,
        self.epilogue_modes,
        self.rounding,
        self.gelu,
    )
    if key in _FRAMES:
        return _FRAMES[key]
    arch = resolve_target_arch(dev)
    # Measured, never built: kept out of the build's registry.
    with set_current_device(dev), compile_context(_iron_kernel_registry=set()):
        k = self.kernel(*tiles)
        calls = (k.mm_fused_acc_init, k.mm_fused_k_step, k.mm_fused_epilogue_chunk)
        deepest = max(k.frame_bytes(arch, entry=c.name) for c in calls)
    _FRAMES[key] = dev.default_core_stack_bytes + deepest
    return _FRAMES[key]

kernel

kernel(tile_n, k_tile, ct_k, band_m, bfp16_b, emulate)

Return the fused tile kernel each core runs at these tiles.

Source code in iron/operators/flm/gemm/op.py

def kernel(self, tile_n, k_tile, ct_k, band_m, bfp16_b, emulate):
    """Return the fused tile kernel each core runs at these tiles."""
    modes = _compiled_modes(self.epilogue_modes)
    return fused_mm(
        dim_m=M_TILE,
        dim_k=k_tile,
        dim_n=tile_n,
        band_m=band_m,
        chunk_k=ct_k,
        out_chunk=CT_OUT_LEN,
        c_depth=C_DEPTH,
        # 8x8x8 on both architectures: AIE2P lowers it onto one bfp16 mac
        # or two bf16 ones, AIE2 onto four native bf16 macs.
        mmul_shape=(R, S, T),
        bfp16_b=bfp16_b,
        b_col_maj=self.b_col_maj,
        emulate_bf16_mmul_with_bfp16=emulate,
        epilogue_modes=tuple(str(m) for m in modes),
        rounding=str(self.rounding),
        gelu=str(self.gelu),
    )

xclbin_operator

xclbin_operator()

Return this configuration at its reference shape and activation.

The xclbin every shape sharing it runs, each with its own instruction stream.

Source code in iron/operators/flm/gemm/op.py

def xclbin_operator(self):
    """Return this configuration at its reference shape and activation.

    The xclbin every shape sharing it runs, each with its own
    instruction stream.
    """
    M, K, N = self._reference_shape
    # The tiles are named: the copy resolves again, at the reference shape.
    return self.replace(
        M=M,
        K=K,
        N=N,
        tile_n=self.tile_n,
        k_tile=self.k_tile,
        ct_k=self.ct_k,
        tile_ma=self.tile_ma,
        m_chunk=self.m_chunk,
        epilogue=Epilogue.NONE,
        epilogue_modes=self.compiled_modes,
        packed_blocks=None,
    )

pack_b

pack_b(B)

Reorder a row-major (K, N) weight matrix into consumption order.

(iron.operators.flm.packing): flat uint8 bfp16ebs8 blocks on NPU2, flat bf16 on NPU1, making both B hops linear descriptors.

Raises:

Type	Description
ValueError	Under b_col_maj, which reads B as stored.

Source code in iron/operators/flm/gemm/op.py

def pack_b(self, B):
    """Reorder a row-major ``(K, N)`` weight matrix into consumption order.

    (``iron.operators.flm.packing``): flat uint8 bfp16ebs8 blocks on
    NPU2, flat bf16 on NPU1, making both B hops linear descriptors.

    Raises:
        ValueError: Under ``b_col_maj``, which reads B as stored.
    """
    t = self.resolved()
    if t.b_col_maj:
        raise ValueError(
            "flm.ProjectionGEMM: b_col_maj (True) reads B as stored (N, K)"
        )
    return pack_b(
        B,
        k_tile=t.k_tile,
        n_tile=t.tile_n,
        s=S,
        t=T,
        ct_k=t.ct_k,
        bfp16=bool(t.bfp16_b),
        round_conv_even=t.rounding is Rounding.CONV_EVEN,
    )

reference

reference(A, B)

C = epilogue(A @ B), in the kernel's order.

The product accumulates in float32 and rounds to the output dtype before the epilogue, since mm_fused_epilogue_chunk converts the accumulator to bf16 first. Not bit-exact -- the activation is a LUT on aie2, native on aie2p, worth up to ~0.02 absolute, absorbed by the tolerance. B is (N, K) under b_col_maj.

Source code in iron/operators/flm/gemm/op.py

def reference(self, A, B):
    """``C = epilogue(A @ B)``, in the kernel's order.

    The product accumulates in float32 and rounds to the output dtype
    before the epilogue, since ``mm_fused_epilogue_chunk`` converts the
    accumulator to bf16 first. Not bit-exact -- the activation is a LUT
    on aie2, native on aie2p, worth up to ~0.02 absolute, absorbed by
    the tolerance. ``B`` is ``(N, K)`` under ``b_col_maj``.
    """
    b = B.T if self.b_col_maj else B
    C = np.matmul(A.astype(np.float32), b.astype(np.float32)).astype(A.dtype)
    return Epilogue(self.epilogue).apply(C)

tolerance

tolerance() -> Tolerance | None

Return C's roundings, in units of 2^-8 of |A| @ |B|.

Under 2 for C's conversion to bf16, plus K * 2^-16 for the f32 accumulator; bfp16 macs add 2 more units (nearest) or 8 (truncating). An activated GEMM has no tolerance, since the approximation isn't bounded here.

Source code in iron/operators/flm/gemm/op.py

def tolerance(self) -> Tolerance | None:
    """Return C's roundings, in units of 2^-8 of ``|A| @ |B|``.

    Under 2 for C's conversion to bf16, plus K * 2^-16 for the f32
    accumulator; bfp16 macs add 2 more units (nearest) or 8
    (truncating). An activated GEMM has no tolerance, since the
    approximation isn't bounded here.
    """
    if self.epilogue is not Epilogue.NONE:
        return None
    units = 2.0 + self.K * 2.0**-16
    if self.emulate_bf16_mmul_with_bfp16:
        units += 2.0 if self.rounding is Rounding.CONV_EVEN else 8.0

    def bound(A, B):
        b = B.T if self.b_col_maj else B
        return (
            units
            * 2.0**-8
            * (np.abs(A.astype(np.float32)) @ np.abs(b.astype(np.float32)))
        )

    return Tolerance.bounded(bound, note=f"{units:g} units of |A| @ |B|")

DecodeLayer

DecodeLayer(*, geometry: FlmGemma4DecodeGeometry, layer_type: LayerType | str, x_len: int = lambda op: arg_sizes(op.geometry)['x'], proj_len: int = lambda op: arg_sizes(op.geometry)['proj'], rms_len: int = lambda op: arg_sizes(op.geometry)['rms'], rope_rms_len: int = lambda op: arg_sizes(op.geometry)['rope_rms'], kv_len: int = lambda op: arg_sizes(op.geometry)['kv'])

Bases: Operator

One Gemma 4 decode layer for one token, as FastFlowLM's engine drives it.

See README.md for the parameters, the buffers and the dispatch parameters.

Parameters:

Name	Type	Description	Default
geometry	FlmGemma4DecodeGeometry		required
layer_type	LayerType | str		required
x_len	int	Derived where not given; a given value must equal it.	lambda op: arg_sizes(op.geometry)['x']
proj_len	int	Derived where not given; a given value must equal it.	lambda op: arg_sizes(op.geometry)['proj']
rms_len	int	Derived where not given; a given value must equal it.	lambda op: arg_sizes(op.geometry)['rms']
rope_rms_len	int	Derived where not given; a given value must equal it.	lambda op: arg_sizes(op.geometry)['rope_rms']
kv_len	int	Derived where not given; a given value must equal it.	lambda op: arg_sizes(op.geometry)['kv']

Attributes:

Name	Type	Description
x		InOut(x_len, dtype=bfloat16) The runtime sequence's argument order.
proj		In(proj_len, dtype=bfloat16)
rms		In(rms_len, dtype=bfloat16)
rope_rms		In(rope_rms_len, dtype=bfloat16)
kv		InOut(kv_len, dtype=bfloat16)
context_len		Per call.
max_l		Per call.

Source code in iron/operators/flm/layer/op.py

class DecodeLayer(Operator):

dims property

dims: LayerDims

This layer type's dimensions and weight blob.

reference

reference(x, proj, rms, rope_rms, kv, context_len, max_l)

X and kv as the dispatch leaves them; see reference.py.

Source code in iron/operators/flm/layer/op.py

def reference(self, x, proj, rms, rope_rms, kv, context_len, max_l):
    """X and kv as the dispatch leaves them; see reference.py."""
    return tuple(
        r.view(bfloat16)
        for r in reference.reference(
            self.dims, x, proj, rms, rope_rms, kv, context_len, max_l
        )
    )

LMHead

LMHead(*, dim: int, vocab: int, softcap: float, w_words: int = lambda op: packed_bytes(op.vocab * op.dim) // 4, x_len: int = lambda op: 2 * op.dim, cols: int = None, rows: int = None)

Bases: Operator

Softcapped logits from a q4nx vocabulary, with the RMS norm folded in.

x holds the token then its RMS weight, w the q4nx vocabulary, y the logits. See README.md.

Parameters:

Name	Type	Description	Default
dim	int		required
vocab	int		required
softcap	float		required
w_words	int	Derived where not given; a given value must equal it.	lambda op: packed_bytes(op.vocab * op.dim) // 4
x_len	int	Derived where not given; a given value must equal it.	lambda op: 2 * op.dim
cols	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None
rows	int	Tunable: the call's, else the profile's, else resolve(dev)'s.	None

Attributes:

Name	Type	Description
y		Out(vocab, dtype=bfloat16, obj_shape=(rows, M_TILE), per=(cols,)) The runtime sequence's argument order. A w or y column object holds one block per core row, which the cores split or join. One broadcast sends the token and its RMS weight to every core.
w		In(w_words, dtype=np.uint32, obj_shape=(rows, BLOCK_BYTES // 4), per=(cols,))
x		In(x_len, dtype=bfloat16, obj_shape=(x_len,))

Source code in iron/operators/flm/lm_head/op.py

class LMHead(Operator):

build

build() -> Built

Each core RMS-normalizes the token once.

Each round, it accumulates one weight block per K_TILE in-features into M_TILE out-features and applies the tanh softcap to the sums.

Source code in iron/operators/flm/lm_head/op.py

def build(self) -> Built:
    """Each core RMS-normalizes the token once.

    Each round, it accumulates one weight block per K_TILE
    in-features into M_TILE out-features and applies the tanh
    softcap to the sums.
    """
    cols, rows, dim = self.cols, self.rows, self.dim
    rounds = self.vocab // (cols * rows * M_TILE)
    kernel = flm_gemma4.flm_gemma4_q4nx_lm_head(
        dim=dim, m_tile=M_TILE, k_tile=K_TILE, group=GROUP
    )
    # The kernels' argument types are the x and w fifos' element types
    # and the core-local scratch.
    # _LmHeadKernel (flm_gemma4.py) declares these entry points; the factory's
    # return type is annotated as the base ExternalFunction.
    w_blk_ty, x_ty, y_acc_ty, sums_ty, _ = kernel.q4nx_lm_head_block.arg_types()
    y_blk_ty, _, rtp_ty = kernel.q4nx_lm_head_epilogue.arg_types()
    # The w fifos move the host's words as bf16 elements. The kernel
    # reinterprets the bytes.
    (w_blk,) = np_ndarray_type_get_shape(w_blk_ty)
    w_col_ty = np.ndarray[(rows, w_blk), np.dtype[bfloat16]]

    def core_fn(x_in, w_in, y_out, k_rms, k_zero, k_block, k_epi, y_acc, sums, rtp):
        x = x_in.acquire(1)
        k_rms(x, sums)
        for _ in range_(rounds):
            k_zero(y_acc)
            for k in range_(dim // K_TILE):
                w = w_in.acquire(1)
                k_block(w, x, y_acc, sums, k)
                w_in.release(1)
            y = y_out.acquire(1)
            k_epi(y, y_acc, rtp)
            y_out.release(1)
        x_in.release(1)

    # The host addresses the DMAs by column. The build therefore pins
    # the shim tiles. The placer places the other tiles.
    shim = [Tile(j, 0) for j in range(cols)]
    of_x = ObjectFifo(x_ty)
    self.x.bind(of_x.prod(tile=shim[0]))

    workers, ws, ys, softcaps = [], [], [], []
    for j in range(cols):
        of_w = ObjectFifo(w_col_ty)
        ws.append(of_w.prod(tile=shim[j]))
        w_cores = of_w.cons().split(
            [w_blk * i for i in range(rows)],
            obj_types=[w_blk_ty] * rows,
            names=[f"w{j}_{i}" for i in range(rows)],
        )

        of_y = ObjectFifo(self.y.obj_type)
        ys.append(of_y.cons(tile=shim[j]))
        y_cores = of_y.prod().join(
            [M_TILE * i for i in range(rows)],
            obj_types=[y_blk_ty] * rows,
            names=[f"y{j}_{i}" for i in range(rows)],
        )

        for i in range(rows):
            y_acc = Buffer(type=y_acc_ty, name=f"yacc_{i}_{j}")
            sums = Buffer(type=sums_ty, name=f"sums_{i}_{j}")
            rtp = Buffer(type=rtp_ty, name=f"rtp_{i}_{j}", use_write_rtp=True)
            softcaps.append(rtp)
            workers.append(
                Worker(
                    core_fn,
                    [
                        of_x.cons(),
                        w_cores[i].cons(),
                        y_cores[i].prod(),
                        kernel.q4nx_lm_head_rms,
                        kernel.q4nx_lm_head_zero,
                        kernel.q4nx_lm_head_block,
                        kernel.q4nx_lm_head_epilogue,
                        y_acc,
                        sums,
                        rtp,
                    ],
                )
            )
    self.w.bind(ws)
    self.y.bind(ys)
    return _Built(workers, softcaps=softcaps)

projection_error

projection_error(w, x) -> float

Return the bound on the projection's error.

A share of the largest logit before the softcap.

Source code in iron/operators/flm/lm_head/op.py

def projection_error(self, w, x) -> float:
    """Return the bound on the projection's error.

    A share of the largest logit before the softcap.
    """
    uncapped = type(self)(dim=self.dim, vocab=self.vocab, softcap=1e30)
    return PROJECTION_ERROR * float(np.abs(uncapped.reference(w, x)).max())

reference

reference(w, x)

Return the softcapped logits of X for W, computed in float64 and rounded to Y's bf16.

Source code in iron/operators/flm/lm_head/op.py

def reference(self, w, x):
    """Return the softcapped logits of X for W, computed in float64 and rounded to Y's bf16."""
    weights = reference.dequantize(w, self.dim, self.vocab, self.resolved().rows)
    logits = reference.reference(weights, np.asarray(x, np.float64), self.softcap)
    return logits.astype(bfloat16)

PrefillAttention

PrefillAttention(*, max_context: int, num_heads: int, num_kv_heads: int, head_dim: int = lambda op: Geometry.of(op.kernel()).dh, qo_len: int = lambda op: op.max_context * op.num_heads * op.head_dim, kv_len: int = lambda op: 2 * op.max_context * op.num_kv_heads * op.head_dim)

Bases: _PrefillAttentionBase

Causal prefill attention with a head dim of 512, from a KV cache.

Parameters:

Name	Type	Description	Default
max_context	int		required
num_heads	int		required
num_kv_heads	int		required
head_dim	int	Derived where not given; a given value must equal it.	lambda op: Geometry.of(op.kernel()).dh
qo_len	int	Derived where not given; a given value must equal it.	lambda op: op.max_context * op.num_heads * op.head_dim
kv_len	int	Derived where not given; a given value must equal it.	lambda op: 2 * op.max_context * op.num_kv_heads * op.head_dim

Attributes:

Name	Type	Description
o		Out(qo_len, dtype=bfloat16) The runtime sequence's argument order.
q		In(qo_len, dtype=bfloat16)
kv		In(kv_len, dtype=bfloat16)
L_begin		Per call.
L_end		Per call.
max_l		Per call.

Source code in iron/operators/flm/prefill_attn/op.py

class PrefillAttention(_PrefillAttentionBase):

PrefillSlidingAttention

PrefillSlidingAttention(*, max_context: int, num_heads: int, num_kv_heads: int, head_dim: int = lambda op: Geometry.of(op.kernel()).dh, qo_len: int = lambda op: op.max_context * op.num_heads * op.head_dim, kv_len: int = lambda op: 2 * op.max_context * op.num_kv_heads * op.head_dim, window: int = 512)

Bases: _PrefillAttentionBase

Sliding-window causal prefill attention with a head dim of 256, from a KV cache.

Parameters:

Name	Type	Description	Default
max_context	int		required
num_heads	int		required
num_kv_heads	int		required
head_dim	int	Derived where not given; a given value must equal it.	lambda op: Geometry.of(op.kernel()).dh
qo_len	int	Derived where not given; a given value must equal it.	lambda op: op.max_context * op.num_heads * op.head_dim
kv_len	int	Derived where not given; a given value must equal it.	lambda op: 2 * op.max_context * op.num_kv_heads * op.head_dim
window	int		512

Attributes:

Name	Type	Description
o		Out(qo_len, dtype=bfloat16) The runtime sequence's argument order.
q		In(qo_len, dtype=bfloat16)
kv		In(kv_len, dtype=bfloat16)
L_begin		Per call.
L_end		Per call.
max_l		Per call.

Source code in iron/operators/flm/prefill_attn/op.py

class PrefillSlidingAttention(_PrefillAttentionBase):