Writing an operator¶

What an operator is declared with. The declaration itself (param, tunable, per_call) is mlir-aie's, in aie.utils.compile.jit; the contributing guide walks through adding one.

Declaration¶

IRON's half of an operator: building, graphs and tuning.

The declaration itself (param, tunable, per_call, Select and the base operator) is mlir-aie's, in aie.utils.compile.jit.

Operand ¶

Operand(member: Any, op: Operator)

An operand as one operator holds it (op.x).

The resolved In, Out or InOut, read through on use (a shape may follow a tunable), and its bound lane_handles. An InOut keeps what its stream does not drain.

Source code in iron/common/declare/operand.py

def __init__(self, member: Any, op: Operator) -> None:
    self.op = op
    self.lane_handles: list[Any] = []
    self._decl = member.of(op)

flat_type `property` ¶

flat_type

The runtime-sequence argument type: this buffer flattened to 1-D.

In its own element units (blocks, for block float).

obj_type `property` ¶

obj_type

The fifo element type: np.ndarray[shape, dtype].

l1_depth `property` ¶

l1_depth: int

The depth a core-side fifo of this operand's tiles can have: two, or one.

One when a tile spans more than a hardware bank of the bound device's core memory.

batch_axes `property` ¶

batch_axes: int

1 if a leading T | None dimension is present on this instance, else 0.

bounded `property` ¶

bounded: tuple[Value, int, Value] | None

(extent, axis, word) when a bound extent sizes an axis.

word is the per-call tiles per lane.

bind ¶

bind(handles) -> None

Bind each lane, in lane order, to the fifo end build() built for it.

A one-lane operand may take its handle bare.

Raises:

Type	Description
`ValueError`	Not one handle per lane, or the operand is bound already.

Source code in iron/common/declare/operand.py

def bind(self, handles) -> None:
    """Bind each lane, in lane order, to the fifo end ``build()`` built for it.

    A one-lane operand may take its handle bare.

    Raises:
        ValueError: Not one handle per lane, or the operand is bound already.
    """
    handles = list(handles) if isinstance(handles, (list, tuple)) else [handles]
    if len(handles) != self.lanes:
        raise ValueError(
            f"{self.name} has {self.lanes} lanes; bind() was given "
            f"{len(handles)} handles"
        )
    if self.lane_handles:
        raise ValueError(f"{self.name} is already bound")
    self.lane_handles = handles

pin ¶

pin(endpoints: list[DmaEndpoint]) -> None

Pin each lane, in order, to a shim DmaEndpoint built elsewhere.

DmaEndpoint is MM2S for an input, S2MM for an output, so its transfers are shim DMA tasks; build() returns the endpoints among its objects.

Raises:

Type	Description
`ValueError`	An endpoint moves the other way or names no channel, or `bind` refuses the endpoints.

Source code in iron/common/declare/operand.py

def pin(self, endpoints: list[DmaEndpoint]) -> None:
    """Pin each lane, in order, to a shim ``DmaEndpoint`` built elsewhere.

    ``DmaEndpoint`` is MM2S for an input, S2MM for an output, so its
    transfers are shim DMA tasks; ``build()`` returns the endpoints
    among its objects.

    Raises:
        ValueError: An endpoint moves the other way or names no channel,
            or ``bind`` refuses the endpoints.
    """
    moves = DMAChannelDir.S2MM if self.drains else DMAChannelDir.MM2S
    for e in endpoints:
        if e.direction != moves or e.channel is None:
            raise ValueError(
                f"{self.name} moves {moves} on a channel it names, not {e}"
            )
    self.bind(endpoints)

fill ¶

fill(tap: TensorAccessPattern | None = None, *, lane: int | None = None, group: TaskGroup | None = None, wait: bool = False, source: Operand | None = None, offset_parameter: Value | None = None, length_parameter: Value | None = None, length_unit: int | None = None, repeat_parameter: Value | None = None, managed: bool = True)

Fill lane of this operand's stream from the host in sequence(rt).

Returns the Task. A per-call Value is patched into the descriptor on a full ELF; on an xclbin it regenerates the instruction stream every call.

Parameters:

Name	Type	Description	Default
`tap`	`TensorAccessPattern \| None`	The walk over `source`; all of it where None.	`None`
`lane`	`int \| None`	Required unless the operand has one lane.	`None`
`group`	`TaskGroup \| None`	The `TaskGroup` this transfer's `Task` joins, or None for its own wait; an unmanaged transfer joins no group.	`None`
`wait`	`bool`	End with an await rather than a free.	`False`
`source`	`Operand \| None`	Another input whose data rides this operand's lane.	`None`
`offset_parameter`	`Value \| None`	A per-call element offset of `tap`.	`None`
`length_parameter`	`Value \| None`	A per-call count of `length_unit` elements sizing `tap`'s outermost dimension below the iteration, in whole 16-byte units; `tap` must fit one descriptor.	`None`
`length_unit`	`int \| None`	The element count one `length_parameter` unit counts, matching a descriptor dimension's inner product.	`None`
`repeat_parameter`	`Value \| None`	A per-call count of extra runs of `tap`, at most its outermost size less one.	`None`
`managed`	`bool`	False hands the queue slot and descriptors to the compiler, which frees them after a later wait; joins no group.	`True`

Source code in iron/common/declare/operand.py

def fill(
    self,
    tap: TensorAccessPattern | None = None,
    *,
    lane: int | None = None,
    group: TaskGroup | None = None,
    wait: bool = False,
    source: Operand | None = None,
    offset_parameter: Value | None = None,
    length_parameter: Value | None = None,
    length_unit: int | None = None,
    repeat_parameter: Value | None = None,
    managed: bool = True,
):
    """Fill ``lane`` of this operand's stream from the host in ``sequence(rt)``.

    Returns the ``Task``. A per-call ``Value`` is patched into the
    descriptor on a full ELF; on an xclbin it regenerates the
    instruction stream every call.

    Args:
        tap: The walk over ``source``; all of it where None.
        lane: Required unless the operand has one lane.
        group: The ``TaskGroup`` this transfer's ``Task`` joins, or None
            for its own wait; an unmanaged transfer joins no group.
        wait: End with an await rather than a free.
        source: Another input whose data rides this operand's lane.
        offset_parameter: A per-call element offset of ``tap``.
        length_parameter: A per-call count of ``length_unit`` elements
            sizing ``tap``'s outermost dimension below the iteration, in
            whole 16-byte units; ``tap`` must fit one descriptor.
        length_unit: The element count one ``length_parameter`` unit
            counts, matching a descriptor dimension's inner product.
        repeat_parameter: A per-call count of extra runs of ``tap``, at
            most its outermost size less one.
        managed: False hands the queue slot and descriptors to the
            compiler, which frees them after a later wait; joins no group.
    """
    return self._runtime().transfer(
        "fill",
        self,
        tap,
        lane=lane,
        group=group,
        wait=wait,
        source=source,
        offset_parameter=offset_parameter,
        length_parameter=length_parameter,
        length_unit=length_unit,
        repeat_parameter=repeat_parameter,
        managed=managed,
    )

drain ¶

drain(tap: TensorAccessPattern | None = None, *, lane: int | None = None, group: TaskGroup | None = None, wait: bool = False, offset_parameter: Value | None = None, length_parameter: Value | None = None, length_unit: int | None = None, repeat_parameter: Value | None = None, managed: bool = True)

Drain lane lane of this operand's stream to the host.

The arguments are fill's, and a drain the sequence must see finish before it ends passes wait=True.

Source code in iron/common/declare/operand.py

def drain(
    self,
    tap: TensorAccessPattern | None = None,
    *,
    lane: int | None = None,
    group: TaskGroup | None = None,
    wait: bool = False,
    offset_parameter: Value | None = None,
    length_parameter: Value | None = None,
    length_unit: int | None = None,
    repeat_parameter: Value | None = None,
    managed: bool = True,
):
    """Drain lane ``lane`` of this operand's stream to the host.

    The arguments are ``fill``'s, and a drain the sequence must see
    finish before it ends passes ``wait=True``.
    """
    return self._runtime().transfer(
        "drain",
        self,
        tap,
        lane=lane,
        group=group,
        wait=wait,
        offset_parameter=offset_parameter,
        length_parameter=length_parameter,
        length_unit=length_unit,
        repeat_parameter=repeat_parameter,
        managed=managed,
    )

extent_axis ¶

extent_axis(extent: PerCallMember) -> int | None

Return the axis of the resolved shape that extent bounds sizes.

None where no axis matches, or extent bounds no parameter.

Source code in iron/common/declare/operand.py

def extent_axis(self, extent: PerCallMember) -> int | None:
    """Return the axis of the resolved shape that ``extent`` bounds sizes.

    None where no axis matches, or ``extent`` bounds no parameter.
    """
    if extent.bounds is None:
        return None
    return axis_of(self.dims, extent.bounds.name, self.op)

fits ¶

fits(tap: TensorAccessPattern) -> bool

Whether one buffer descriptor of a lane's shim tile holds tap.

tap is a pattern over this operand.

Source code in iron/common/declare/operand.py

def fits(self, tap: TensorAccessPattern) -> bool:
    """Whether one buffer descriptor of a lane's shim tile holds ``tap``.

    ``tap`` is a pattern over this operand.
    """
    dev = self.op.dev
    assert (
        dev is not None
    ), "fits() is read during sequence derivation, with a device bound"
    # Every shim tile's descriptor is alike; column 0's stands for a lane's.
    return dev.dma_bd_fits(0, 0, tap, self.dtype)

split ¶

split(elements: int) -> tuple[int, int] | None

Return a contiguous run of elements of this operand as (d1, d0).

d0 is the largest a lane's shim descriptor holds as its innermost dimension with d1 beside it, or None where no d0 does.

Source code in iron/common/declare/operand.py

def split(self, elements: int) -> tuple[int, int] | None:
    """Return a contiguous run of ``elements`` of this operand as ``(d1, d0)``.

    ``d0`` is the largest a lane's shim descriptor holds as its
    innermost dimension with ``d1`` beside it, or None where no ``d0`` does.
    """
    dev = self.op.dev
    assert (
        dev is not None
    ), "split() is read during sequence derivation, with a device bound"
    for d0 in range(elements, 0, -1):
        if elements % d0:
            continue
        # The run inside a strided walk, so neither dimension is exempt
        # from its wrap as a linear transfer is.
        probe = TensorAccessPattern(
            (8 * elements,),
            0,
            [1, 2, elements // d0, d0],
            [0, 4 * elements, d0, 1],
        )
        if dev.dma_bd_fits(0, 0, probe, self.dtype):
            return elements // d0, d0
    return None

extent_unit ¶

extent_unit(axis: int) -> int

Return the rows along axis of one round-robin unit under a bound.

Operator.extent_unit's, else the stream tile's.

Source code in iron/common/declare/operand.py

def extent_unit(self, axis: int) -> int:
    """Return the rows along ``axis`` of one round-robin unit under a bound.

    ``Operator.extent_unit``'s, else the stream tile's.
    """
    unit = self.op.extent_unit(self.name)
    if unit is not None:
        return unit
    obj_shape = self.obj_shape if self.streamed else ()
    k = axis - (len(self.shape) - len(obj_shape))
    return obj_shape[k] if k >= 0 else 1

transfers ¶

transfers() -> list[tuple[int, TensorAccessPattern, Value | None, int | None]]

Return the derived (lane, tap, length_parameter, length_unit) transfers.

The whole buffer (one lane or replicated) or one block of the first non-batch axis per lane; under a bound, round-robin by tile instead.

Raises:

Type	Description
`ValueError`	No `obj_shape=`, or the split does not divide, cannot bound that axis, or does not fit a descriptor.

Source code in iron/common/declare/operand.py

def transfers(
    self,
) -> list[tuple[int, TensorAccessPattern, Value | None, int | None]]:
    """Return the derived ``(lane, tap, length_parameter, length_unit)`` transfers.

    The whole buffer (one lane or replicated) or one block of the first
    non-batch axis per lane; under a bound, round-robin by tile instead.

    Raises:
        ValueError: No ``obj_shape=``, or the split does not divide, cannot
            bound that axis, or does not fit a descriptor.
    """
    name = f"{type(self.op).__name__}.{self.name}"
    if not self.streamed:
        raise ValueError(
            f"{name} has no obj_shape=, so its sequence cannot be derived; add "
            f"obj_shape= or override sequence(rt)"
        )
    count, shape = self.lanes, self.shape
    bounded = self.bounded
    if bounded is None:
        if count == 1:
            return [(0, self.tap, None, None)]
        if self.replicate:
            return [(i, self.tap, None, None) for i in range(count)]
        axis = self.batch_axes
        if axis >= len(shape) or shape[axis] % count:
            raise ValueError(
                f"{self.name} {shape} does not divide across its {count} lanes "
                f"on axis {axis}. Check {type(self.op).__name__}.compatible()"
            )
        parts = TensorAccessPattern.full(shape).partition(count, axis)
        return [(i, parts[i], None, None) for i in range(count)]
    extent, axis, word = bounded
    if axis != self.batch_axes:
        raise ValueError(
            f"{name}: {extent.name} bounds axis {axis}, but the derived "
            f"sequence splits axis {self.batch_axes}; override sequence(rt) "
            f"to bound another axis"
        )
    lanes = 1 if self.replicate else count
    inner = math.prod(shape[axis + 1 :])
    tile_rows = self.extent_unit(axis)
    if shape[axis] % (lanes * tile_rows):
        raise ValueError(
            f"{self.name} {shape}: axis {axis} does not divide into "
            f"{tile_rows}-row tiles over {lanes} lanes"
        )
    if axis > 1:
        raise ValueError(
            f"{self.name} {shape}: bounding axis {axis} needs {axis + 3} "
            f"descriptor dimensions; a descriptor has four"
        )
    tiles = shape[axis] // (lanes * tile_rows)
    run = tile_rows * inner
    halves = self.split(run)
    if halves is None:
        raise ValueError(
            f"{self.name}: a {run}-element tile does not fit one descriptor"
        )
    dev = self.op.dev
    assert (
        dev is not None
    ), "transfers() is read during sequence derivation, with a device bound"
    # A unit D1 is stripped, which would drop the tile count into D1's
    # wrap; past it, a tile split in two keeps the count on D2, which has
    # none (and a half of no whole granules fits no descriptor below).
    wrap = (1 << dev.get_dma_bd_wrap_bits(0, 0)) - 1
    if halves[0] == 1 and tiles > wrap and run % 2 == 0:
        halves = (2, run // 2)
    iterations = shape[0] if axis else 1
    tiled = TensorAccessPattern.full((iterations, tiles, lanes, run))
    out = []
    for lane in range(lanes):
        tap = tiled[:, :, lane].split(2, halves[1])
        if not self.fits(tap):
            raise ValueError(
                f"{self.name} {shape}: the round-robin split over {lanes} "
                f"lanes does not fit one descriptor per lane"
            )
        for k in range(count) if self.replicate else [lane]:
            out.append((k, tap, word, run))
    return out

Built `dataclass` ¶

Built(objects: list, rtps: list[Rtp] = list())

What build() made: the array's objects and each core's Rtp.

The objects are what the array declares (Workers, Flows, PacketFlows, Locks, TileDmas, Buffers, pinned DmaEndpoints). A sequence needing more reads a subclass via rt.built.

Operator ¶

Operator(*, bound_values: dict[str, Affine | None] | None = None, pinned: frozenset[str] | None = None, **given: Any)

Bases: Operator

An operator, whose subclass body is its declaration.

The declaration is mlir-aie's param(), tunable(), per_call(); aie.iron's In/Out/InOut.

Attributes:

Name	Type	Description
`bound_values`	`dict[str, Affine \| None]`	The graph expression bound to each per-call value, by name.
`pinned`	`frozenset[str]`	The tunables its caller gave, which the tuner holds; it may move a profile's and `resolve()`'s.

Source code in iron/common/declare/operator.py

def __init__(
    self,
    *,
    bound_values: dict[str, Affine | None] | None = None,
    pinned: frozenset[str] | None = None,
    **given: Any,
) -> None:
    if bound_values:
        self.bound_values = bound_values
    super().__init__(**given)
    if pinned is None:
        pinned = frozenset(
            n for n in self.declaration.tunables if given.get(n) is not None
        )
    self.pinned = pinned
    self._check_complete()

values `property` ¶

values: list[Value]

The per-call values: each gets a device parameter.

residents `property` ¶

residents: dict[str, Any]

What the preamble writes once per build, by name.

name `property` ¶

name: str

The label of the resolved operator: its class, shown fields and device.

step `classmethod` ¶

step(*operands: Any, **kwargs: Any) -> Handle

Record a step inside a graph's body: GEMV.step(w, h, **kw).

Raises:

Type	Description
`TypeError`	Called outside a graph's body.

Source code in iron/common/declare/operator.py

@classmethod
def step(cls, *operands: Any, **kwargs: Any) -> Handle:
    """Record a step inside a graph's body: ``GEMV.step(w, h, **kw)``.

    Raises:
        TypeError: Called outside a graph's body.
    """
    recorder = graph_recorder.get()
    if recorder is None:
        raise TypeError(f"{cls.__name__}.step() is for a graph's body")
    return recorder.call(cls, operands, kwargs)

build ¶

build() -> Built

Build the array for the bound device on layout().

Bind every operand's lanes to fifo ends and give each core reading Values an Rtp.

Source code in iron/common/declare/operator.py

def build(self) -> Built:
    """Build the array for the bound device on ``layout()``.

    Bind every operand's lanes to fifo ends and give each core reading
    ``Value``s an ``Rtp``.
    """
    raise NotImplementedError(f"{type(self).__name__}.build() is not implemented")

tolerance ¶

tolerance() -> Tolerance | None

Return what each output of the resolved operator is judged by.

None where nothing bounds it, which check mode refuses.

Source code in iron/common/declare/operator.py

def tolerance(self) -> Tolerance | None:
    """Return what each output of the resolved operator is judged by.

    None where nothing bounds it, which check mode refuses.
    """
    return None

ops ¶

ops() -> int

Return the arithmetic operations one call performs, for its throughput.

Source code in iron/common/declare/operator.py

def ops(self) -> int:
    """Return the arithmetic operations one call performs, for its throughput."""
    return sum(b.elements for b in self.outputs)

array_device ¶

array_device()

Return the device the array's Program is built for.

The bound device unless the array takes a part of it.

Source code in iron/common/declare/operator.py

def array_device(self):
    """Return the device the array's Program is built for.

    The bound device unless the array takes a part of it.
    """
    return self.dev

xclbin_operator ¶

xclbin_operator() -> Self

Return the operator whose xclbin this one runs when it runs alone.

Source code in iron/common/declare/operator.py

def xclbin_operator(self) -> Self:
    """Return the operator whose xclbin this one runs when it runs alone."""
    return self

exported_design ¶

exported_design() -> Callable[..., Any] | None

Return another tool's design for this operator; None derives it.

Source code in iron/common/declare/operator.py

def exported_design(self) -> Callable[..., Any] | None:
    """Return another tool's design for this operator; ``None`` derives it."""
    return None

shim_columns `classmethod` ¶

shim_columns(dev, num_channels: int = 1, flags: Mapping[str, Any] | None = None) -> int

How many of dev's columns fit its streams in the shim DMA budget.

A replicate stream costs once per channel, one with no per= once, and a when= flag flags leaves out is its default.

Raises:

Type	Description
`Unresolvable`	Not even one column's streams fit.

Source code in iron/common/declare/operator.py

@classmethod
def shim_columns(
    cls, dev, num_channels: int = 1, flags: Mapping[str, Any] | None = None
) -> int:
    """How many of ``dev``'s columns fit its streams in the shim DMA budget.

    A ``replicate`` stream costs once per channel, one with no ``per=``
    once, and a ``when=`` flag ``flags`` leaves out is its default.

    Raises:
        Unresolvable: Not even one column's streams fit.
    """
    flags = flags or {}
    streams = [
        m
        for m in cls.declaration.operands
        if m.streamed and (m.when is None or flags.get(m.when.name, m.when.default))
    ]
    cols = dev.cols
    for drains, budget in (
        (False, dev.shim_dma_channels_in),
        (True, dev.shim_dma_channels_out),
    ):
        ours = [m for m in streams if m.drains == drains]
        lanes = [m for m in ours if m.per is not None]
        per_core = sum(not m.replicate for m in lanes) * num_channels
        once = len(ours) - len(lanes)
        shared = sum(m.replicate for m in lanes) * num_channels + once
        if per_core:
            cols = min(cols, (budget - shared) // per_core)
    if cols < 1:
        raise Unresolvable(
            f"{cls.__name__}: one column's streams exceed this device's shim DMA budget"
        )
    return cols

call_operands `classmethod` ¶

call_operands(args, kwargs) -> tuple[dict[str, Any], list, dict[str, Any]]

Split a call into inputs by name, outputs and other keywords.

As (x, ..., [outputs...], *, weight=None, ...) would: inputs are in declaration order, and an optional input is a keyword, so never taken for an output.

Source code in iron/common/declare/operator.py

@classmethod
def call_operands(cls, args, kwargs) -> tuple[dict[str, Any], list, dict[str, Any]]:
    """Split a call into inputs by name, outputs and other keywords.

    As ``(x, ..., [outputs...], *, weight=None, ...)`` would: inputs are
    in declaration order, and an optional input is a keyword, so never
    taken for an output.
    """
    ins = [m for m in cls.declaration.operands if m.fills]
    required = [m.name for m in ins if m.when is None]
    optional = {m.name for m in ins if m.when is not None}
    given = {k: v for k, v in kwargs.items() if k in optional and v is not None}
    if len(args) < len(required):
        raise TypeError(
            f"{cls.__name__} takes {len(required)} positional operand(s) "
            f"({', '.join(required)}), got {len(args)}"
        )
    bound = {**dict(zip(required, args)), **given}
    inputs = {m.name: bound[m.name] for m in ins if m.name in bound}
    rest = {k: v for k, v in kwargs.items() if k not in optional}
    return inputs, list(args[len(required) :]), rest

resolve_columns ¶

resolve_columns(dev, given: int | None, num_channels: int = 1, *, fits: Callable[[int], bool] | None = None) -> int

given, checked against the shim budget, or the most columns that fits.

When no count fits, one column is returned and compatible names the rule.

Source code in iron/common/declare/operator.py

def resolve_columns(
    self,
    dev,
    given: int | None,
    num_channels: int = 1,
    *,
    fits: Callable[[int], bool] | None = None,
) -> int:
    """``given``, checked against the shim budget, or the most columns that ``fits``.

    When no count fits, one column is returned and ``compatible`` names the rule.
    """
    if given is not None:
        if dev is not None:
            self.check_shim_columns(dev, given, num_channels)
        return given
    if dev is None:
        raise Unresolvable(
            f"{type(self).__name__}: the column count defaults from the "
            f"device; none is bound and none was given"
        )
    budget = self.shim_columns(dev, num_channels, vars(self))
    return next((c for c in range(budget, 0, -1) if fits is None or fits(c)), 1)

addressed_inputs ¶

addressed_inputs(addresses: Mapping[str, int], rng: Generator) -> dict[str, ndarray]

Return the content of each input that names a device address.

For a run alone: addresses is every buffer's device address, by buffer name; random bytes there would steer the DMA anywhere. Other inputs are filled at random.

Source code in iron/common/declare/operator.py

def addressed_inputs(
    self, addresses: Mapping[str, int], rng: np.random.Generator
) -> dict[str, np.ndarray]:
    """Return the content of each input that names a device address.

    For a run alone: ``addresses`` is every buffer's device address, by
    buffer name; random bytes there would steer the DMA anywhere. Other
    inputs are filled at random.
    """
    return {}

judge ¶

judge(inputs: Mapping[str, ndarray], written: Mapping[str, ndarray], tolerance: Tolerance, expected: Mapping[str, ndarray] | None = None) -> dict[str, Verdict]

Each written buffer's verdict against its reference result.

Every mapping is by buffer name. expected defaults to reference on inputs, which a bound tolerance is also evaluated on.

Raises:

Type	Description
`ValueError`	`written` or `expected` name other buffers than it writes, an expected result is None, or a written buffer's element count differs from its result's.

Source code in iron/common/declare/operator.py

def judge(
    self,
    inputs: Mapping[str, np.ndarray],
    written: Mapping[str, np.ndarray],
    tolerance: Tolerance,
    expected: Mapping[str, np.ndarray] | None = None,
) -> dict[str, Verdict]:
    """Each written buffer's verdict against its reference result.

    Every mapping is by buffer name. ``expected`` defaults to
    ``reference`` on ``inputs``, which a bound tolerance is also
    evaluated on.

    Raises:
        ValueError: ``written`` or ``expected`` name other buffers than it
            writes, an expected result is None, or a written buffer's
            element count differs from its result's.
    """
    if expected is None:
        expected = self._call_reference(inputs)
    names = [b.name for b in self.outputs]
    for given, what in ((written, "written"), (expected, "expected")):
        if sorted(given) != sorted(names):
            raise ValueError(
                f"{type(self).__name__} writes {names}; the {what} "
                f"buffers are {sorted(given)}"
            )
    unknown = [n for n in names if expected[n] is None]
    if unknown:
        raise ValueError(f"{type(self).__name__}: no expected result for {unknown}")
    bound = None
    if tolerance.kind == "bound":
        assert tolerance.bound is not None, "kind == 'bound' iff bound is set"
        bound = tolerance.bound(*(inputs[b.name] for b in self.inputs))
    return {
        n: compare(
            np.reshape(written[n], np.shape(expected[n])),
            expected[n],
            tolerance,
            bound=bound,
        )
        for n in names
    }

sequence ¶

sequence(rt) -> None

Override to write the runtime sequence by hand; otherwise it is derived.

Source code in iron/common/declare/operator.py

def sequence(self, rt) -> None:
    """Override to write the runtime sequence by hand; otherwise it is derived."""
    raise NotImplementedError

design_key ¶

design_key()

Identity for sharing a build: equal keys generate byte-identical MLIR.

A graph's bound per-call values add their names.

Source code in iron/common/declare/operator.py

def design_key(self):
    """Identity for sharing a build: equal keys generate byte-identical MLIR.

    A graph's bound per-call values add their names.
    """
    cls, own = super().design_key()
    if self.bound_values:
        names = {
            k: None if v is None else v.name for k, v in self.bound_values.items()
        }
        own += (("values", tuple(sorted(names.items()))),)
    return cls, own

resolved ¶

resolved(dev: Device | None = None) -> Self

Return this operator resolved for dev, the bound device unless given.

Source code in iron/common/declare/operator.py

def resolved(self, dev: Device | None = None) -> Self:
    """Return this operator resolved for ``dev``, the bound device unless given."""
    dev = self.dev if dev is None else dev
    if dev is None:
        raise Unresolvable(
            f"{type(self).__name__}: resolved() needs a device; none is bound"
        )
    return super().resolved(dev)

replace ¶

replace(**changes: Any) -> Self

Return a copy with changes, checked against the device once complete.

changes may name any parameter, the bound_values or the pinned tunables.

Source code in iron/common/declare/operator.py

def replace(self, **changes: Any) -> Self:
    """Return a copy with ``changes``, checked against the device once complete.

    ``changes`` may name any parameter, the ``bound_values`` or the
    ``pinned`` tunables.
    """
    bound_values = changes.pop("bound_values", self.bound_values)
    pinned = changes.pop("pinned", self.pinned)
    new = super().replace(**changes)
    new.bound_values, new.pinned = bound_values, pinned
    new._check_complete()
    return new

derived_at ¶

derived_at(name: str, **extents: int | Affine) -> Any

Return the word the host writes for name.

With the extents at the given bounds (an Affine form where the caller evaluates it symbolically rather than at one number).

Source code in iron/common/declare/operator.py

def derived_at(self, name: str, **extents: int | Affine) -> Any:
    """Return the word the host writes for ``name``.

    With the extents at the given bounds (an ``Affine`` form where the
    caller evaluates it symbolically rather than at one number).
    """
    member = next((m for m in self._value_members if m.name == name), None)
    if member is None or member.derive is None:
        raise TypeError(f"{type(self).__name__}.{name} is not a derived value")
    unknown = set(extents) - {
        m.name for m in self.declaration.per_calls if m.bounds is not None
    }
    if unknown:
        raise TypeError(
            f"{type(self).__name__} declares no extent {sorted(unknown)}"
        )
    return member.derive(self._at_extents(**extents))

extent_unit ¶

extent_unit(buffer: str) -> int | None

Return the rows of buffer one lane takes at a time under a bound.

None for the stream tile's rows; 0 when a bound does not shorten it.

Source code in iron/common/declare/operator.py

def extent_unit(self, buffer: str) -> int | None:
    """Return the rows of ``buffer`` one lane takes at a time under a bound.

    ``None`` for the stream tile's rows; ``0`` when a bound does not shorten it.
    """
    return None

value ¶

value(name: str) -> Value

Return the device word of value name.

An extent attribute reads as an integer.

Source code in iron/common/declare/operator.py

def value(self, name: str) -> Value:
    """Return the device word of value ``name``.

    An extent attribute reads as an integer.
    """
    try:
        return self._values[name]
    except KeyError:
        raise TypeError(
            f"{type(self).__name__} declares no value {name!r}"
        ) from None

use_value ¶

use_value(name: str, bound_to: Affine | None = None) -> None

Record that a graph binds the per-call value name to its expression bound_to.

Source code in iron/common/declare/operator.py

def use_value(self, name: str, bound_to: Affine | None = None) -> None:
    """Record that a graph binds the per-call value ``name`` to its expression ``bound_to``."""
    if not any(m.name == name for m in self._value_members):
        raise TypeError(
            f"{type(self).__name__} declares no per-call value {name!r}"
        )
    # Rebound, not updated: replace() hands a copy the same dict.
    self.bound_values = {**self.bound_values, name: bound_to}
    vars(self).pop("_derived_under_bound", None)

compile ¶

compile() -> Artifacts

Build, once, the image this operator runs on outside a graph.

Returns its artifacts (the xclbin, the instruction stream and their cache entry).

Source code in iron/common/declare/operator.py

def compile(self) -> Artifacts:
    """Build, once, the image this operator runs on outside a graph.

    Returns its artifacts (the xclbin, the instruction stream and their
    cache entry).
    """
    if self._image is None:
        # The image layer builds on this one: imported here to break the cycle.
        from ..image.standalone import OperatorImage

        self._image = OperatorImage(self)
    return self._image.compile().artifacts

layout ¶

layout() -> Self

Return this operator as its cores see it.

Each operand is bound to the view and its values this operator's; an extent reads as its number only where the key holds its bound.

Source code in iron/common/declare/operator.py

def layout(self) -> Self:
    """Return this operator as its cores see it.

    Each operand is bound to the view and its values this operator's; an
    extent reads as its number only where the key holds its bound.
    """
    view = object.__new__(type(self))
    vars(view).update(
        {n: vars(self)[n] for n in self.declaration.key}, trace=self.trace
    )
    view._values = self._values
    view._bound = {}
    for m in self.declaration.members:
        if m.name not in self._bound:
            continue
        if isinstance(m, PerCallMember) and m.bounds is not None:
            if m.bounds.name in vars(view):
                view._bound[m.name] = vars(view)[m.bounds.name]
        elif isinstance(m, PerCallMember):
            view._bound[m.name] = self._bound[m.name]
        else:
            view._bound[m.name] = Operand(m, view)
    return view

Precision ¶

Bases: StrEnum

The roundings a matmul may make beyond its reference's.

Named as torch.set_float32_matmul_precision names them; each level allows the one above's and more, and a matmul takes the fastest kernel its level allows.

HIGHEST: f32 partial sums, rounded once to the output's type.
HIGH: partial sums rounded to the output's type per K tile.
MEDIUM: A and B rounded to bfp16 for the multiplies, where the device has bfp16 macs.

Profile ¶

Profile()

Tunable values by operator shape, as (cls, dims, tunables) entries.

An open tunable comes from the most specific match (a dimension left out matches any, a subclass its base's); equally specific ones that disagree raise.

Source code in iron/common/declare/profile.py

def __init__(self) -> None:
    self._entries: list[tuple[type, dict[str, Any], dict[str, Any]]] = []

make ¶

make(cls: type[_Op], **kwargs: Any) -> _Op

Build cls(**kwargs) with its open tunables filled from this profile.

Standalone use of what a graph's recorder applies implicitly to every step it constructs.

Source code in iron/common/declare/profile.py

def make(self, cls: type[_Op], **kwargs: Any) -> _Op:
    """Build ``cls(**kwargs)`` with its open tunables filled from this profile.

    Standalone use of what a graph's recorder applies implicitly to every
    step it constructs.
    """
    given = {name: value for name, value in kwargs.items() if value is not None}
    pinned = frozenset(given.keys() & set(cls.declaration.tunables))
    return cls(**{**self.tunables_for(cls, given), **given}, pinned=pinned)

add ¶

add(cls: type, **fields: Any) -> None

Add an entry for cls.

param() values to match (one left out matches any value) and tunable() values to give.

Source code in iron/common/declare/profile.py

def add(self, cls: type, **fields: Any) -> None:
    """Add an entry for ``cls``.

    ``param()`` values to match (one left out matches any value) and
    ``tunable()`` values to give.
    """
    dims, tunables, unknown, fixed = {}, {}, [], []
    for name, value in fields.items():
        param = cls.declaration.params.get(name)
        if name in cls.declaration.tunables:
            tunables[name] = value
        elif param is not None and param.tunable:
            fixed.append(name)
        elif param is not None:
            dims[name] = value
        else:
            unknown.append(name)
    if unknown:
        raise TypeError(f"{cls.__name__} declares no field {unknown}")
    if fixed:
        raise TypeError(
            f"{cls.__name__}.{fixed} are derived or fixed by the class, not "
            f"tunables a profile gives"
        )
    if not tunables:
        raise TypeError(f"an entry for {cls.__name__} must give a tunable")
    self._entries.append((cls, dims, tunables))

tunables_for ¶

tunables_for(cls: type, given: Mapping[str, Any]) -> dict[str, Any]

Resolve the tunables for the cls a call with given keywords makes.

Each comes from the most specific entry naming it.

Source code in iron/common/declare/profile.py

def tunables_for(self, cls: type, given: Mapping[str, Any]) -> dict[str, Any]:
    """Resolve the tunables for the ``cls`` a call with ``given`` keywords makes.

    Each comes from the most specific entry naming it.
    """
    dims = self.dims(cls, given)
    chosen: dict[str, tuple[int, Any]] = {}
    for base, matched, tunables in self._entries:
        if not issubclass(cls, base) or not self.matches(matched, dims):
            continue
        rank = len(matched)
        for name, value in tunables.items():
            if name not in chosen or rank > chosen[name][0]:
                chosen[name] = (rank, value)
            elif rank == chosen[name][0] and chosen[name][1] != value:
                raise ValueError(
                    f"profile is ambiguous for {cls.__name__} {dims}: "
                    f"{name}={chosen[name][1]} and {name}={value} from "
                    f"entries of equal specificity"
                )
    return {name: value for name, (_, value) in chosen.items()}

dims `staticmethod` ¶

dims(kind: type, given: Mapping[str, Any]) -> dict[str, Any]

Collect every dimension (param()) a kind call with given makes.

This is what an entry is matched against: one it leaves out is its default, None where a function computes it.

Source code in iron/common/declare/profile.py

@staticmethod
def dims(kind: type, given: Mapping[str, Any]) -> dict[str, Any]:
    """Collect every dimension (``param()``) a ``kind`` call with ``given`` makes.

    This is what an entry is matched against: one it leaves out is its
    default, ``None`` where a function computes it.
    """
    dims = {}
    for name, p in kind.declaration.params.items():
        if p.tunable:
            continue
        if name in given:
            dims[name] = given[name]
        elif p.derive is not None or isinstance(p.default, FunctionType):
            # Unknown until made, so an entry naming it matches only a call
            # giving it: entries key on what a call constructs with.
            dims[name] = None
        elif p.default is not MISSING:
            dims[name] = p.default
    return dims

matches `staticmethod` ¶

matches(matched: Mapping[str, Any], dims: Mapping[str, Any]) -> bool

Whether an entry's matched dimensions cover dims.

Source code in iron/common/declare/profile.py

@staticmethod
def matches(matched: Mapping[str, Any], dims: Mapping[str, Any]) -> bool:
    """Whether an entry's ``matched`` dimensions cover ``dims``."""
    return all(name in dims and dims[name] == v for name, v in matched.items())

load `classmethod` ¶

load(path: str | Path) -> Profile

Load the profile save wrote to path.

Source code in iron/common/declare/profile.py

@classmethod
def load(cls, path: str | Path) -> Profile:
    """Load the profile ``save`` wrote to ``path``."""
    from ...operators import CATALOG  # the operators import this module

    with open(path) as f:
        data = json.load(f)
    if not isinstance(data, dict) or set(data) != {"entries"}:
        raise ValueError(f"{path}: a profile is {{'entries': [...]}}")
    profile = cls()
    for entry in data["entries"]:
        fields = dict(entry)
        name = fields.pop("operator")
        if name not in CATALOG:
            raise ValueError(f"{path}: no operator {name!r} in iron.operators")
        op_cls = CATALOG[name]
        # An Enum-domain tunable is its member's value in JSON, not its ordinal.
        for key, param in op_cls.declaration.params.items():
            domain = param.domain
            if (
                key in fields
                and isinstance(domain, type)
                and issubclass(domain, Enum)
            ):
                fields[key] = domain(fields[key])
        profile.add(op_cls, **fields)
    return profile

save ¶

save(path: str | Path) -> None

Write the profile as JSON, one entry to a line.

Raises:

Type	Description
`ValueError`	An entry's operator is not in `iron.operators.CATALOG`.

Source code in iron/common/declare/profile.py

def save(self, path: str | Path) -> None:
    """Write the profile as JSON, one entry to a line.

    Raises:
        ValueError: An entry's operator is not in ``iron.operators.CATALOG``.
    """
    from ...operators import CATALOG  # the operators import this module

    names = {cls: name for name, cls in CATALOG.items()}
    lines = []
    for cls, dims, tunables in self._entries:
        if cls not in names:
            raise ValueError(f"{cls.__qualname__} is not in iron.operators.CATALOG")
        name = names[cls]
        lines.append(json.dumps({"operator": name, **dims, **tunables}))
    with open(path, "w") as f:
        f.write('{"entries": [\n  ' + ",\n  ".join(lines) + "\n]}\n")

Block ¶

Block(words: int, *, derive: Any)

Bases: PerCallMember

A block of words words a core reads, written once per build.

derive computes them from the operator. A shipped image's parameters are one.

Source code in iron/common/declare/rtp.py

def __init__(self, words: int, *, derive: Any) -> None:
    super().__init__(np.int32, derive, None)
    self.words = words

Rtp ¶

Bases: tuple

The words one core reads per call: a buffer and scratchpad words.

The sequence writes the buffer, then opens barrier on it. A tuple, so a Worker takes it as one argument whose leaves mlir-aie registers; the core calls read.

Attributes:

Name	Type	Description
`slots`	`dict[Value, int]`	Each value the buffer holds and its first word there.
`buffer`	`Buffer \| None`	None when every value is a scratchpad word.

Raises:

Type	Description
`ValueError`	A value of several words is read from the scratchpad, whose parameters are one word each.

read ¶

read() -> dict[Value, Any]

Read each value's word (a list of them for a block), keyed by the value.

In a core: wait for the sequence, read, re-arm.

Source code in iron/common/declare/rtp.py

def read(self) -> dict[Value, Any]:
    """Read each value's word (a list of them for a block), keyed by the value.

    In a core: wait for the sequence, read, re-arm.
    """
    self.barrier.wait_for_value(1)
    words: dict[Value, Any] = {}
    for value in self.values:
        if value.in_scratchpad:
            words[value] = value.param.read()
            continue
        at = self.slots[value]
        assert (
            self.buffer is not None
        ), "a slot exists, so some value is not scratchpad"
        block = [self.buffer[at + i] for i in range(value.words)]
        words[value] = block if value.words > 1 else block[0]
    self.barrier.release_with_value(1)
    return words

write ¶

write(words: Mapping[str, Any]) -> None

Write the words of each value the buffer holds, by value name.

Called in the preamble: a resident's number, or an xclbin's per-call scalar.

Raises:

Type	Description
`ValueError`	A value is given another number of words than it declares.

Source code in iron/common/declare/rtp.py

def write(self, words: Mapping[str, Any]) -> None:
    """Write the words of each value the buffer holds, by value name.

    Called in the preamble: a resident's number, or an xclbin's per-call
    scalar.

    Raises:
        ValueError: A value is given another number of words than it declares.
    """
    for value, at in self.slots.items():
        word = words[value.name]
        block = word if isinstance(word, (list, tuple)) else [word]
        if len(block) != value.words:
            raise ValueError(
                f"{value.name} holds {value.words} words; it was given "
                f"{len(block)}"
            )
        assert (
            self.buffer is not None
        ), "a slot exists, so some value is not scratchpad"
        for i, w in enumerate(block):
            self.buffer[at + i] = w

Value ¶

Value(op: Operator, name: str, *, dtype: Any = int32, derive: Callable[[Any], Any] | None = None, bounds: OperatorParam | None = None, words: int = 1)

Bases: PerCallMember

A per_call() member bound to one operator (op.value("count")).

Or op.count unless an extent. Per call when a graph binds it, else written once per build before the first DMA. A core reads it through its Rtp.

Attributes:

Name	Type	Description
`derive`		The once-per-build word, or `words` of them, from the operator; None for one every call gives.
`bounds`		For an extent, the `param()` it bounds per call; else None.
`param`	`Any`	The build's parameter: a `ScratchpadParameter` on a full ELF, a dispatch-time scalar on an xclbin.
`ssa`	`Any`	Its value in the sequence body, on an xclbin.

Source code in iron/common/declare/value.py

def __init__(
    self,
    op: Operator,
    name: str,
    *,
    dtype: Any = np.int32,
    derive: Callable[[Any], Any] | None = None,
    bounds: OperatorParam | None = None,
    words: int = 1,
) -> None:
    super().__init__(dtype, derive, bounds)
    self.op = op
    self.name = name
    self.words = words

per_call `property` ¶

per_call: bool

A graph binds this value, or derive reads a bound extent.

A value with no derive and no bounds is always per call.

dispatch `property` ¶

dispatch: bool

The instruction stream is regenerated around this value every call.

One every call gives, on an operator with an exported design. A full ELF cannot carry it.

symbol `property` ¶

symbol: str

The device symbol the host writes it through.

An exported design's own name, else <instance>_<value>[_<its graph value>]; a derived value's names all the instance's graph values, so differently bound ones differ.

in_scratchpad `property` ¶

in_scratchpad: bool

The cores read this value from the parameter scratchpad.

A build gave it a ScratchpadParameter (a per-call value on a full ELF).

Xclbin ¶

Xclbin(url: str, sha256: str, filename: str, **kernel: Any)

Bases: NPUKernel

A downloaded xclbin pinned by sha256.

Fetched over https into the JIT cache's prebuilt/<filename> on first read of xclbin_path. As an image= it has no instructions; a build runs it with_instructions.

Parameters:

Name	Type	Description	Default
`**kernel`	`Any`	`NPUKernel`'s other arguments, kept as `settings`.	`{}`

Source code in iron/common/declare/xclbin.py

def __init__(self, url: str, sha256: str, filename: str, **kernel: Any) -> None:
    super().__init__(**kernel)
    self.url, self.sha256, self.filename = url, sha256, filename
    self.settings = kernel

with_instructions ¶

with_instructions(insts_path, **kernel) -> Xclbin

Return this image running insts_path.

With kernel (as NPUKernel takes it) over the declared settings.

Source code in iron/common/declare/xclbin.py

def with_instructions(self, insts_path, **kernel) -> Xclbin:
    """Return this image running ``insts_path``.

    With ``kernel`` (as `NPUKernel` takes it) over the declared
    settings.
    """
    run = Xclbin(
        self.url,
        self.sha256,
        self.filename,
        **{**self.settings, **kernel, "insts_path": insts_path},
    )
    run._xclbin_path = self._xclbin_path
    return run

fetch ¶

fetch(directory=None) -> Path

Return the downloaded file in directory.

NPU_CACHE_HOME's prebuilt/ by default, fetched unless the pinned content is there.

Raises:

Type	Description
`ValueError`	The URL is not https.
`RuntimeError`	The download has another digest.

Source code in iron/common/declare/xclbin.py

def fetch(self, directory=None) -> Path:
    """Return the downloaded file in ``directory``.

    ``NPU_CACHE_HOME``'s ``prebuilt/`` by default, fetched unless
    the pinned content is there.

    Raises:
        ValueError: The URL is not https.
        RuntimeError: The download has another digest.
    """
    if directory is None:
        directory = Path(NPU_CACHE_HOME) / "prebuilt"
    target = Path(directory) / self.filename
    if target.exists():
        with open(target, "rb") as f:
            if hashlib.file_digest(f, "sha256").hexdigest() == self.sha256:
                return target
    if not self.url.startswith("https://"):
        raise ValueError(f"refusing to download over {self.url!r}")
    with urllib.request.urlopen(self.url, timeout=60) as response:
        content = response.read()
    if (got := hashlib.sha256(content).hexdigest()) != self.sha256:
        raise RuntimeError(f"{self.url} has SHA-256 {got}, expected {self.sha256}")
    target.parent.mkdir(parents=True, exist_ok=True)
    # A file of this fetch's own, renamed over the target whole: concurrent
    # fetches never share one, and none sees a partial target.
    partial = tempfile.NamedTemporaryFile(
        dir=target.parent, prefix=f"{target.name}.", suffix=".part", delete=False
    )
    try:
        with partial:
            partial.write(content)
            # A temporary file is created owner-only; the cache is read as any file is.
            os.fchmod(partial.fileno(), 0o644)
        os.replace(partial.name, target)
    except BaseException:
        os.unlink(partial.name)
        raise
    return target

Elementwise templates¶

The shared elementwise array.

One core per (column, channel), each streaming fixed-size lines of every operand through the kernel a subclass's kernel() names, with the scalars scalars() supplies; a parameter either reads is declared array=True unless a stream names it. A replicate=True input is one line each core holds for the whole call and every line meets (RowwiseMul's row).

Sweep ¶

Sweep(lengths: Iterable[int] = LENGTHS, *, channels: Iterable[int] | None = (1, 2), tile_cap: int | None = None, regular: int | None = 2048, rows: bool = False, bench: int | None = BENCH_ELEMENTS, lower: bool = False, **extra: Any)

The cases of an elementwise operator.

At the regular length, one core and the most its shim budget allows (the grids between run the same loop); at the other lengths, in the extensive suite, the smallest and the largest tile.

Parameters:

Name	Type	Description	Default
`channels`	`Iterable[int] \| None`	None leaves the count to the operator (a binary one).	`(1, 2)`
`tile_cap`	`int \| None`	Else the class's `tile_cap`.	`None`
`regular`	`int \| None`	The one length in the default suite; None for all.	`2048`
`rows`	`bool`	A length is that many elements in rows of `tile_size` (`Rowwise`).	`False`
`bench`	`int \| None`	The length of one CI-tracked case at the widest grid; None adds none.	`BENCH_ELEMENTS`
`lower`	`bool`	Flags its first default case `Case.lower`.	`False`
`**extra`	`Any`	Given to every case.	`{}`

Source code in iron/common/elementwise.py

def __init__(
    self,
    lengths: Iterable[int] = LENGTHS,
    *,
    channels: Iterable[int] | None = (1, 2),
    tile_cap: int | None = None,
    regular: int | None = 2048,
    rows: bool = False,
    bench: int | None = BENCH_ELEMENTS,
    lower: bool = False,
    **extra: Any,
):
    self.lengths = tuple(lengths)
    self.bench = bench
    self.lower = lower
    self.channels = None if channels is None else tuple(channels)
    self.tile_cap = tile_cap
    self.regular = regular
    self.rows = rows
    self.extra = extra

KernelCall `dataclass` ¶

KernelCall(kernel: ExternalFunction, scalars: tuple = (), operands: tuple[Operand | None, ...] = ())

One kernel call an elementwise operator's core makes over a line.

Parameters:

Name	Type	Description	Default
`scalars`	`tuple`	The scalar arguments its contract leaves unbound, in order.	`()`
`operands`	`tuple[Operand \| None, ...]`	Each input tile, in the kernel's order: an input of the operator (`self.x`), or None, the line the call before it wrote.	`()`

Elementwise ¶

Elementwise(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int = None)

Bases: Operator

The array for an elementwise kernel over lines of tile_size elements.

Subclasses declare operands with the line as their tile, most split over every core (per=CORES), and implement kernel; tile_cap is the longest line.

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`

Attributes:

Name	Type	Description
`count`		Per call.

Source code in iron/common/elementwise.py

class Elementwise(Operator):

valid_elements `property` ¶

valid_elements: int

The elements a call processes, the bounded extent's worth if bound.

chain ¶

chain() -> tuple[KernelCall, ...]

Return the kernel calls a core makes over one line, in order.

reference and judge follow them; more than kernel() only with an array of its own.

Source code in iron/common/elementwise.py

def chain(self) -> tuple[KernelCall, ...]:
    """Return the kernel calls a core makes over one line, in order.

    ``reference`` and ``judge`` follow them; more than ``kernel()``
    only with an array of its own.
    """
    config = self.layout()
    return (KernelCall(config.kernel(), config.scalars(), tuple(self.inputs)),)

over ¶

over(elements: int, line: int) -> Self

Return this operator over elements elements in lines of line, on one core.

Raises:

Type	Description
`ValueError`	Its shape is not a run of lines.

Source code in iron/common/elementwise.py

def over(self, elements: int, line: int) -> Self:
    """Return this operator over ``elements`` elements in lines of ``line``, on one core.

    Raises:
        ValueError: Its shape is not a run of lines.
    """
    raise ValueError(f"{type(self).__name__} is not a run of lines")

reference ¶

reference(*inputs: Any, **kwargs: Any)

Run the kernel contract's reference, line by line, what the cores compute.

A bounded call's elements short of a whole line are zero-padded to whole lines, as the device runs whole tiles.

Source code in iron/common/elementwise.py

def reference(self, *inputs: Any, **kwargs: Any):
    """Run the kernel contract's reference, line by line, what the cores compute.

    A bounded call's elements short of a whole line are zero-padded
    to whole lines, as the device runs whole tiles.
    """
    (out,) = self.outputs
    line = self.default_tile if self.tile_size is None else self.tile_size
    if out.elements % line:
        whole = -(-out.elements // line) * line
        padded = [
            (
                x
                if b.replicate
                else np.concatenate(
                    [np.ravel(x), np.zeros(whole - np.size(x), np.asarray(x).dtype)]
                )
            )
            for b, x in zip(self.inputs, inputs)
        ]
        y = np.asarray(self.over(whole, line).reference(*padded))
        return y.reshape(-1)[: out.elements].reshape(out.host_shape)
    op, lines, scalars = self._kernel_arguments(inputs)
    (out,) = op.outputs
    y = (
        op.layout()
        .kernel()
        .expected(lines, scalars=scalars, arch=resolve_target_arch(op.dev))
    )
    y = np.asarray(y).astype(out.host_dtype, copy=False)
    return y.reshape(out.host_shape, copy=False)

judge ¶

judge(inputs: Mapping[str, ndarray], written: Mapping[str, ndarray], tolerance: Tolerance, expected: Mapping[str, ndarray] | None = None) -> dict[str, Verdict]

Its output held to its reference by its kernel's judge, line by line.

Where a core makes one kernel call; as any operator's where it makes more.

Source code in iron/common/elementwise.py

def judge(
    self,
    inputs: Mapping[str, np.ndarray],
    written: Mapping[str, np.ndarray],
    tolerance: Tolerance,
    expected: Mapping[str, np.ndarray] | None = None,
) -> dict[str, Verdict]:
    """Its output held to its reference by its kernel's ``judge``, line by line.

    Where a core makes one kernel call; as any operator's where it
    makes more.
    """
    op = self.resolved()
    if len(op.chain()) > 1:
        return super().judge(inputs, written, tolerance, expected)
    if expected is None:
        expected = self._call_reference(inputs)
    _, lines, scalars = self._kernel_arguments([inputs[b.name] for b in op.inputs])
    (out,) = op.outputs
    verdict = (
        op.layout()
        .kernel()
        .judge(
            written[out.name],
            expected[out.name],
            calls=op.lines,
            tolerance=tolerance,
            inputs=lines,
            scalars=scalars,
            arch=resolve_target_arch(op.dev),
        )
    )
    return {out.name: verdict}

kernel ¶

kernel() -> ExternalFunction

Return the ExternalFunction each core calls over one line.

Source code in iron/common/elementwise.py

def kernel(self) -> ExternalFunction:
    """Return the ``ExternalFunction`` each core calls over one line."""
    raise NotImplementedError(f"{type(self).__qualname__} declares no kernel()")

scalars ¶

scalars() -> tuple

Return the kernel's scalar arguments its contract leaves unbound, in order.

Source code in iron/common/elementwise.py

def scalars(self) -> tuple:
    """Return the kernel's scalar arguments its contract leaves unbound, in order."""
    return ()

held_segments ¶

held_segments() -> int

Return the lines each held input is cut into.

A core holds all of them and meets, in turn, with that many lines of every other input.

Source code in iron/common/elementwise.py

def held_segments(self) -> int:
    """Return the lines each held input is cut into.

    A core holds all of them and meets, in turn, with that many
    lines of every other input.
    """
    return 1

arguments ¶

arguments(kernel: ExternalFunction, given: tuple, n_in: int, n_out: int) -> tuple[list[int], dict]

Return where each tile goes in a call of kernel, and its scalar arguments.

n_in inputs and n_out outputs; given fills the scalars the contract leaves unbound.

Returns:

Type	Description
`list[int]`	The argument positions of the inputs then the outputs, in the
`dict`	order a core acquires them, and the scalar values by position.

Source code in iron/common/elementwise.py

def arguments(
    self, kernel: ExternalFunction, given: tuple, n_in: int, n_out: int
) -> tuple[list[int], dict]:
    """Return where each tile goes in a call of ``kernel``, and its scalar arguments.

    ``n_in`` inputs and ``n_out`` outputs; ``given`` fills the
    scalars the contract leaves unbound.

    Returns:
        The argument positions of the inputs then the outputs, in the
        order a core acquires them, and the scalar values by position.
    """
    contract = kernel.contract
    if contract is None:
        return list(range(n_in + n_out)), {}
    if contract.accumulates:
        raise ValueError(
            f"{type(self).__qualname__}: an elementwise kernel writes its output; "
            f"{kernel.name} accumulates into one"
        )
    roles = contract.roles
    scalars = dict(contract.parameter_bindings)
    free = [
        i for i, r in enumerate(roles) if r is kernels.Param and i not in scalars
    ]
    if len(free) != len(given):
        raise ValueError(
            f"{type(self).__qualname__}: {kernel.name} leaves {len(free)} scalar(s) "
            f"unbound; {len(given)} given"
        )
    scalars.update(zip(free, given))
    outs = list(contract.out_indices)
    ins = [i for i in range(len(roles)) if i not in scalars and i not in outs]
    if (len(ins), len(outs)) != (n_in, n_out):
        raise ValueError(
            f"{type(self).__qualname__}: {kernel.name} takes {len(ins)} input and "
            f"{len(outs)} output tile(s); it is given {n_in} and {n_out}"
        )
    return ins + outs, scalars

UnaryElementwise ¶

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

Bases: Elementwise

A flat buffer in, a flat buffer of the same size out.

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/common/elementwise.py

class UnaryElementwise(Elementwise):

BinaryElementwise ¶

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

Bases: Elementwise

Two flat buffers in, one of the same size out.

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/common/elementwise.py

class BinaryElementwise(Elementwise):

Rowwise ¶

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

Bases: Elementwise

rows rows of tile_size elements in, the same out, one kernel call per row.

For a kernel that reduces over its line, so the line is a param() the cores read: another line would compute something else.

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(rows, tile_size, obj_shape=(tile_size,), per=CORES)`
`y`		`Out(rows, tile_size, obj_shape=(tile_size,), per=CORES)`

Source code in iron/common/elementwise.py

class Rowwise(Elementwise):

WeightedRowwise ¶

WeightedRowwise(*, num_aie_columns: int = None, num_channels: int = None, tile_size: int, rows: int, weighted: bool = False)

Bases: Rowwise

A Rowwise kernel's rows, optionally times a learned weight row.

Weighted, two cores per (column, channel) are pipelined, the second multiplying by the row, which halves the line a core holds.

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`

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/common/elementwise.py

class WeightedRowwise(Rowwise):

reference ¶

reference(x, weight=None)

Run the kernels' references in turn, the kernel's row rounded to bf16.

Source code in iron/common/elementwise.py

def reference(self, x, weight=None):
    """Run the kernels' references in turn, the kernel's row rounded to bf16."""
    normed = super().reference(x)
    if weight is None:
        return normed
    y = eltwise.mul_sized(self.tile_size).contract.reference(normed, weight)
    return y.astype(normed.dtype)

RowBroadcast ¶

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

Bases: Rowwise

rows rows of tile_size elements, each met elementwise by one row.

A core holds row whole; past tile_cap a row is cut into segments lines of line elements, each met by its own segment: the same sums, in more calls.

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/common/elementwise.py

class RowBroadcast(Rowwise):

segments_for `classmethod` ¶

segments_for(width: int) -> int | None

Return the fewest lines a width-element row is cut into for a core.

A core holds it beside two lines each of x and y, as much as a row at the cap takes; None where no cut does.

Source code in iron/common/elementwise.py

@classmethod
def segments_for(cls, width: int) -> int | None:
    """Return the fewest lines a ``width``-element row is cut into for a core.

    A core holds it beside two lines each of x and y, as much as a
    row at the cap takes; None where no cut does.
    """
    return next(
        (
            s
            for s in range(1, width + 1)
            if width % s == 0 and width + 4 * (width // s) <= 5 * cls.tile_cap
        ),
        None,
    )

Design and runtime sequence¶

The library-owned build of a declared operator.

OperatorDesign.module builds aie.iron's Program from the operator's build() and an OperatorRuntime. For a full ELF, configuration puts each BuiltArray and its sequences into one DeviceConfiguration.

BuiltArray ¶

BuiltArray(design: OperatorDesign, words: Mapping[str, str] | None = None, dispatch: Mapping[str, object] | None = None)

One design's array, built on a resolved copy of its operator (op).

Its operands' lanes are bound. Designs of one array (identities) build it once, each running its own sequence over it (runtime).

Parameters:

Name	Type	Description	Default
`words`	`Mapping[str, str] \| None`	On a full ELF, design symbol -> the scratchpad word it shares with others that always hold the same number.	`None`
`dispatch`	`Mapping[str, object] \| None`	On an xclbin, each per-call value's dispatch-time scalar, by device symbol.	`None`

Raises:

Type	Description
`ValueError`	A full ELF is asked to carry a `DispatchTime` value, an xclbin is given no scalar for a value, or a trace would reuse an output a full ELF consolidates.

Source code in iron/common/design/design.py

def __init__(
    self,
    design: OperatorDesign,
    words: Mapping[str, str] | None = None,
    dispatch: Mapping[str, object] | None = None,
):
    words = words or {}
    dispatch = dispatch or {}
    self._words, self._dispatch = words, dispatch
    op = design.op.copy()  # a build binds streams
    # Before build(), which may hand a core-read value to a worker.
    for value in op.values:
        if design.full_elf and not value.dispatch:
            value.param = ScratchpadParameter(
                words.get(value.symbol, value.symbol),
                value.dtype,
            )
        elif design.full_elf:
            raise ValueError(
                f"{type(op).__name__}.{value.name} is a DispatchTime value, "
                f"which a full ELF cannot carry (its stream is fixed at build "
                f"time); compile with full_elf=False"
            )
        elif value.symbol not in dispatch:
            raise ValueError(
                f"{type(op).__name__}.{value.name}: no dispatch parameter "
                f"{value.symbol!r} was handed to the build"
            )
        else:
            value.param = dispatch[value.symbol]
    view = op.layout()
    self.built = view.build()
    for mine, built in zip(op.buffers, view.buffers):
        mine.lane_handles = built.lane_handles
    self.design = design
    self.op = op
    self.workers = [w for w in self.built.objects if isinstance(w, Worker)]
    self.traced: list[Worker] = []
    if op.trace is not None:
        if op.trace.reuse_output_buffer:
            raise ValueError(
                f"{type(op).__name__}: its trace takes a buffer of its own "
                f"(reuse_output_buffer=False), since a full ELF consolidates "
                f"its outputs and an xclbin's build is traced the same way"
            )
        self.traced = [w for w in self.workers if w.trace is not None]
        self.traced = self.traced or self.workers[:1]

reads `cached` `property` ¶

reads: frozenset[str]

Return the scratchpad words this design's cores read.

Beside its array_key, all its array depends on. Read off the design built anew and resolved alone, since a core reads a word through the parameter its function closes over.

program ¶

program() -> Program

Return this array alone with its operator's own sequence.

Source code in iron/common/design/design.py

def program(self) -> Program:
    """Return this array alone with its operator's own sequence."""
    return Program(
        self.op.array_device(),
        OperatorRuntime(self.op, self.built),
        workers=self.workers,
    )

runtime ¶

runtime(other: BuiltArray, name: str | None) -> OperatorRuntime

other's runtime sequence over this array, named name.

Raises:

Type	Description
`ValueError`	`other` is another array.

Source code in iron/common/design/design.py

def runtime(self, other: BuiltArray, name: str | None) -> OperatorRuntime:
    """``other``'s runtime sequence over this array, named ``name``.

    Raises:
        ValueError: ``other`` is another array.
    """
    if other is self:
        return OperatorRuntime(self.op, self.built, name=name)
    if other.op.array_key() != self.op.array_key():
        raise ValueError(
            f"{type(other.op).__name__} builds another array than "
            f"{type(self.op).__name__}; its sequence cannot run over this one"
        )
    op = other.op
    for mine, theirs in zip(op.buffers, self.op.buffers):
        mine.lane_handles = theirs.lane_handles
    # A sequence fills only flows it registers, which the configuration
    # keeps once; tile DMAs and shim channels are the owner's to declare.
    # It writes the owner's runtime parameters, and reads what else the
    # owner's build kept.
    reached = (Flow, PacketFlow, Lock)
    return OperatorRuntime(
        op,
        dataclasses.replace(
            self.built,
            objects=[b for b in self.built.objects if isinstance(b, reached)],
        ),
        name=name,
    )

OperatorDesign ¶

OperatorDesign(op: Operator, full_elf: bool = True)

One operator's design (module, or Operator.exported_design).

CompilableDesign compiles and caches it; identity adds to the cache key the fields the operator's repr leaves out.

Source code in iron/common/design/design.py

def __init__(self, op: Operator, full_elf: bool = True):
    if op.reference_only:
        raise TypeError(f"{type(op).__name__} is reference-only: it has no kernel")
    # Keyed by what it builds: a replaced copy of a resolved operator
    # resolves afresh, and may resolve to other tunables than it holds.
    op = self.op = op.resolved(op.dev)
    self.full_elf = full_elf
    P = inspect.Parameter
    build = op.exported_design()
    if build is None:
        build = self.module
        params = [
            P(
                v.symbol,
                P.KEYWORD_ONLY,
                annotation=DispatchTime[v.dtype],
            )
            for v in op.values
            if not full_elf
        ]
    else:
        params = list(inspect.signature(build).parameters.values())
    self.build = build

    def generator(
        op: CompileTime[Operator],
        identity: CompileTime[str],
        **kwargs,
    ):
        return build(**kwargs)

    keys = [
        P("op", P.KEYWORD_ONLY, annotation=CompileTime[Operator]),
        P("identity", P.KEYWORD_ONLY, annotation=CompileTime[str]),
    ]
    setattr(generator, "__signature__", inspect.Signature([*params, *keys]))
    self.generator = generator

identity `cached` `property` ¶

identity: str

What the module is built from, sources aside.

name `property` ¶

name: str

The design's device symbol in a full ELF, stable across source edits.

module ¶

module(**dispatch)

Return the operator's MLIR module on the bound device.

On an xclbin, which has no parameter scratchpad, dispatch is each per-call value's scalar by symbol.

Source code in iron/common/design/design.py

def module(self, **dispatch):
    """Return the operator's MLIR module on the bound device.

    On an xclbin, which has no parameter scratchpad, ``dispatch`` is each
    per-call value's scalar by symbol.
    """
    array = BuiltArray(self, dispatch=dispatch)
    prog = array.program()
    if array.traced:
        assert (
            array.op.trace is not None
        ), "traced is non-empty only when op.trace is set"
        prog.enable_trace(array.op.trace.trace_size, workers=array.traced)
    return prog.resolve_program()

compilable ¶

compilable(**options) -> CompilableDesign

Return the CompilableDesign given options.

The operator's aiecc_flags are added to theirs.

Source code in iron/common/design/design.py

def compilable(self, **options) -> CompilableDesign:
    """Return the ``CompilableDesign`` given ``options``.

    The operator's ``aiecc_flags`` are added to theirs.
    """
    flags = [*options.pop("aiecc_flags", ()), *self.op.aiecc_flags]
    return CompilableDesign(
        self.generator,
        compile_kwargs={"op": self.op, "identity": self.identity},
        aiecc_flags=flags,
        **options,
    )

compile ¶

compile(image: CompilableDesign | None = None, **options) -> CompilableDesign

Compile the design, or find it in the cache.

image, for an insts_only stream, is the design it runs on: the compile fails, naming each resource laid out otherwise than there, rather than run against it.

Source code in iron/common/design/design.py

def compile(
    self, image: CompilableDesign | None = None, **options
) -> CompilableDesign:
    """Compile the design, or find it in the cache.

    ``image``, for an insts_only stream, is the design it runs on: the
    compile fails, naming each resource laid out otherwise than there,
    rather than run against it.
    """
    # The key reads the current device: bound first, keys before and after agree.
    aie_utils.ensure_current_device()
    design = self.compilable(**options)
    design.compile(image=image)
    return design

OperatorRuntime ¶

OperatorRuntime(op: Operator, built: Built, name: str | None = None)

Bases: Runtime

The rt an operator's sequence(rt) receives.

aie.iron's Runtime running the preamble, then that sequence or the one its declaration derives. A managed transfer on a pinned lane keeps at most the operand's depth in flight, all awaited at the sequence's end.

Parameters:

Name	Type	Description	Default
`op`	`Operator`	The resolved operator, its array built and its operands' lanes bound or pinned.	required
`name`	`str \| None`	The sequence's symbol; the Program names it when None.	`None`

Attributes:

Name	Type	Description
`built`		What the operator's `build()` made, which `sequence` may read.
`pins`		Each pinned lane this runtime declares, by its symbol.
`buffers`		Buffers no worker places: a shipped image's parameter blocks, or a pool only descriptors address.

Source code in iron/common/design/runtime.py

def __init__(self, op: Operator, built: Built, name: str | None = None) -> None:
    self.op = op
    self.built = built
    self.buffers = [b for b in built.objects if isinstance(b, Buffer)]
    self._tensors: dict[str, Any] = {}
    self._queues: dict[str, list[Task]] = {}
    # Raises if any stream of a built array is unbound.
    lanes = [
        (b, i, h) for b in op.buffers if b.streamed for i, h in enumerate(b.handles)
    ]
    # A shared array's pins are declared by the runtime of the array that built them.
    ours = {id(e) for e in built.objects if isinstance(e, DmaEndpoint)}
    self.pins = {f"{b.name}_{i}": h for b, i, h in lanes if id(h) in ours}
    if len(self.pins) != len(ours):
        raise TypeError(
            f"{type(op).__name__}.build() makes a DmaEndpoint no operand "
            f"is pinned to"
        )
    handles = [h for _, _, h in lanes if not isinstance(h, DmaEndpoint)]
    super().__init__(
        self.body,
        [
            *(b.flat_type for b in op.buffers),
            handles,
            *(v.param for v in op.values),
        ],
        name=name,
    )
    # Before resolving: the sequence body, generated last, may address any of it.
    for obj in built.objects:
        if isinstance(obj, (Flow, PacketFlow)):
            self.add_flow(obj)
        elif isinstance(obj, Lock):
            self.add_lock(obj)
        elif isinstance(obj, TileDma):
            self.add_tile_dma(obj)
        elif not isinstance(obj, (Worker, DmaEndpoint, Buffer)):
            raise TypeError(
                f"{type(op).__name__}.build() makes Workers, Flows, "
                f"PacketFlows, Locks, TileDmas, the DmaEndpoints it pinned "
                f"and Buffers; got {obj!r}"
            )

body ¶

body(*args) -> None

Record the sequence body, given the operator's runtime tensors, its lanes' handles.

Then each per-call value not in the scratchpad as its scalar.

Source code in iron/common/design/runtime.py

def body(self, *args) -> None:
    """Record the sequence body, given the operator's runtime tensors, its lanes' handles.

    Then each per-call value not in the scratchpad as its scalar.
    """
    buffers = self.op.buffers
    self._tensors = {b.name: a for b, a in zip(buffers, args)}
    for value, arg in zip(self.op.values, args[len(buffers) + 1 :]):
        if not value.in_scratchpad:
            value.ssa = arg
    token = active_runtime.set(self)
    try:
        if not self.op.own_preamble:
            self.preamble()
        self.run()
    finally:
        active_runtime.reset(token)
    for queue in self._queues.values():
        for task in queue:
            task.await_()

declare ¶

declare(device) -> None

Its pins and buffers, at device scope ahead of the sequence.

aie.iron names a shim channel by its route and places a buffer by the worker that reads it, and these have neither.

Source code in iron/common/design/runtime.py

def declare(self, device) -> None:
    """Its pins and buffers, at ``device`` scope ahead of the sequence.

    aie.iron names a shim channel by its route and places a buffer by
    the worker that reads it, and these have neither.
    """
    for symbol, end in self.pins.items():
        device.resolve_tile(end.tile)
        shim_dma_allocation(symbol, end.tile.op, end.direction, end.channel)
    for buffer in self.buffers:
        device.resolve_tile(buffer.tile)
        buffer.resolve()

transfer ¶

transfer(verb: str, operand: Operand, tap: TensorAccessPattern | None, *, lane: int | None, group, wait: bool, source: Operand | None = None, offset_parameter: Value | None = None, length_parameter: Value | None = None, length_unit: int | None = None, repeat_parameter: Value | None = None, managed: bool = True)

Operand.fill/drain (verb): source's (else operand's) tap.

Through lane of operand's stream, its per-call words a scratchpad parameter on a full ELF, the pattern's sizes and offset on an xclbin.

Source code in iron/common/design/runtime.py

def transfer(
    self,
    verb: str,
    operand: Operand,
    tap: TensorAccessPattern | None,
    *,
    lane: int | None,
    group,
    wait: bool,
    source: Operand | None = None,
    offset_parameter: Value | None = None,
    length_parameter: Value | None = None,
    length_unit: int | None = None,
    repeat_parameter: Value | None = None,
    managed: bool = True,
):
    """``Operand.fill``/``drain`` (``verb``): ``source``'s (else ``operand``'s) ``tap``.

    Through ``lane`` of ``operand``'s stream, its per-call words a
    scratchpad parameter on a full ELF, the pattern's sizes and offset
    on an xclbin.
    """
    if lane is None:
        if operand.lanes != 1:
            raise ValueError(f"{operand.name} is per-{operand.lanes}; name a lane")
        lane = 0
    buffer = operand if source is None else source
    tap = buffer.tap if tap is None else tap
    handle = operand.handles[lane]
    pinned = isinstance(handle, DmaEndpoint)
    symbol = f"{operand.name}_{lane}"
    fn = (
        functools.partial(emit_shim_transfer, symbol)
        if pinned
        else getattr(handle, verb)
    )
    for value in (offset_parameter, length_parameter, repeat_parameter):
        if value is None:
            continue
        if not isinstance(value, Value):
            raise TypeError(
                f"a per-call parameter is a value member (self.count), "
                f"got {value!r}"
            )
        if value.param is None:
            raise ValueError(
                f"{value.name} has no device parameter: it is not per call "
                f"here (Value.per_call), or the build has not created it yet"
            )
    if length_parameter is not None and length_parameter.dispatch:
        raise ValueError(
            f"{length_parameter.name} is a DispatchTime value; a per-call "
            f"length is a scratchpad word patched into the descriptor"
        )
    _, sizes, strides = shim_dims(tap)
    # Read before a per-call size or repeat is patched in below, so still concrete.
    static_sizes, static_strides = concrete(sizes, tap), concrete(strides, tap)
    dim = scale = None
    if length_parameter is not None:
        assert length_unit is not None, "length_unit accompanies length_parameter"
        # A patched descriptor's length ends at its outermost dimension
        # below the iteration; an xclbin's stream, regenerated per call,
        # sizes any dimension of length_unit-element units.
        patched = length_parameter.ssa is None
        dim = next(
            (
                d
                for d in range(1, 4)
                if (prod(static_sizes[1:d]) == 1 or not patched)
                and prod(static_sizes[d + 1 :]) == length_unit
            ),
            None,
        )
        if (
            dim is None
            and prod(static_sizes[1:3]) == 1
            and static_sizes[3] % length_unit == 0
        ):
            # One linear run, whose length counts its elements: units of them.
            dim, scale = 3, length_unit
        if dim is None:
            raise ValueError(
                f"{type(self.op).__name__}.{buffer.name}: {tap} has no "
                f"descriptor dimension of {length_unit}-element units"
                + (" below the iteration, outermost" if patched else "")
            )
    data = self._tensors[buffer.name]
    common: dict[str, Any]
    if pinned and managed:
        # The image's memtiles hold depth objects: past that, wait for the oldest.
        queue = self._queues.setdefault(symbol, [])
        if len(queue) == operand.depth:
            queue.pop(0).await_()
        common = dict(wait=True, managed=False)
    elif managed:
        common = dict(wait=wait, group=group)
    else:
        if group is not None:
            raise ValueError("an unmanaged transfer joins no group")
        common = dict(wait=wait, managed=False)
    # A per-call size lands in one descriptor's length, a repeat count in
    # its push; an offset patches every split piece.
    per_call = length_parameter is not None or repeat_parameter is not None
    if per_call and not buffer.fits(tap):
        raise ValueError(
            f"{type(self.op).__name__}.{buffer.name}: {tap} moves by a per-call "
            f"size, so it must fit one buffer descriptor, and does not"
        )
    dynamic = any(
        v is not None and v.ssa is not None
        for v in (offset_parameter, length_parameter, repeat_parameter)
    )
    if not dynamic:
        task = fn(
            data,
            tap=tap,
            offset_parameter=(
                None if offset_parameter is None else offset_parameter.param
            ),
            length_parameter=(
                None if length_parameter is None else length_parameter.param
            ),
            length_unit=None if length_parameter is None else length_unit,
            repeat_parameter=(
                None if repeat_parameter is None else repeat_parameter.param
            ),
            **common,
        )
        if pinned and managed:
            queue.append(task)
        return task
    if dim is not None:
        # dim is set only inside the length_parameter is not None branch above.
        assert length_parameter is not None
        # A unit dimension's stride is normalised to 0, which a dynamic
        # size may not have; its count never passes 1, so the stride the
        # next block would have moves the same bytes.
        if sizes[dim] == 1 and strides[dim] == 0:
            # dim's own size and stride are patched below, but nothing past
            # it, which stays concrete.
            inner = zip(static_sizes[dim + 1 :], static_strides[dim + 1 :])
            strides[dim] = max(n * stride for n, stride in inner)
        sizes[dim] = length_parameter.ssa
        if scale is not None:
            sizes[dim] = arith.muli(
                length_parameter.ssa,
                arith.constant(scale, IntegerType.get_signless(32)),
            )
    if repeat_parameter is not None:
        sizes[0] = repeat_parameter.ssa + arith.constant(
            1, IntegerType.get_signless(32)
        )
    offset: Any = tap.offset
    if offset_parameter is not None and offset_parameter.ssa is not None:
        offset = offset_parameter.ssa
        if tap.offset:
            offset = offset + arith.constant(
                concrete((tap.offset,), tap)[0], IntegerType.get_signless(32)
            )
    task = fn(
        data,
        tap=TensorAccessPattern(tap.tensor_dims, offset, sizes, strides),
        **common,
    )
    if pinned and managed:
        queue.append(task)
    return task

preamble ¶

preamble(**values) -> None

Write each core's runtime parameters, the parameter sync, then each core's barrier, before any DMA.

A shipped image's parameter blocks (buffers) get the residents' words in declaration order, for its sequence to release. An own_preamble operator's sequence calls it, values overriding residents.

Raises:

Type	Description
`ValueError`	`values` names no resident of the operator, or a parameter block holds another number of words.

Source code in iron/common/design/runtime.py

def preamble(self, **values) -> None:
    """Write each core's runtime parameters, the parameter sync, then each core's barrier, before any DMA.

    A shipped image's parameter blocks (``buffers``) get the residents'
    words in declaration order, for its sequence to release. An
    ``own_preamble`` operator's sequence calls it, ``values`` overriding
    residents.

    Raises:
        ValueError: ``values`` names no resident of the operator, or a
            parameter block holds another number of words.
    """
    op = self.op
    residents = op.residents
    unknown = set(values) - set(residents)
    if unknown:
        raise ValueError(
            f"{type(op).__name__} has no resident {sorted(unknown)} to override"
        )
    # Without a scratchpad, a per-call word a core reads is the dispatch-time scalar.
    words = {v.name: v.ssa for v in op.values if not v.in_scratchpad}
    words.update(residents, **values)
    for rtp in self.built.rtps:
        rtp.write(words)
    block = []
    for word in (words[name] for name in residents):
        block += word if isinstance(word, (list, tuple)) else [word]
    for buffer in self.buffers if op.external is not None else ():
        if len(block) != prod(buffer.shape):
            raise ValueError(
                f"{type(op).__name__}'s parameter block holds "
                f"{prod(buffer.shape)} words; its residents are {len(block)}"
            )
        for i, word in enumerate(block):
            buffer[i] = word
    if any(value.in_scratchpad for value in op.values):
        sync_parameters()
    for rtp in self.built.rtps:
        rtp.barrier.set(1)

configuration ¶

configuration(name: str, arrays: Mapping[str, BuiltArray], steps: Sequence[str] = ()) -> DeviceConfiguration

Build arrays as one device configuration name.

One array per identity, and each entry's sequence over it, sequence if it is alone, else named for its key. Given steps (keys of arrays), one sequence runs theirs in order.

Raises:

Type	Description
`ValueError`	The arrays are built for different devices, or traced with different trace sizes.

Source code in iron/common/design/design.py

def configuration(
    name: str, arrays: Mapping[str, BuiltArray], steps: Sequence[str] = ()
) -> DeviceConfiguration:
    """Build ``arrays`` as one device configuration ``name``.

    One array per identity, and each entry's sequence over it, ``sequence``
    if it is alone, else named for its key. Given ``steps`` (keys of
    ``arrays``), one ``sequence`` runs theirs in order.

    Raises:
        ValueError: The arrays are built for different devices, or traced
            with different trace sizes.
    """
    of = identities(arrays)
    owners: dict[Hashable, BuiltArray] = {}
    for key, array in arrays.items():
        owners.setdefault(of[key], array)
    kinds = {type(a.op.array_device()).__name__ for a in owners.values()}
    if len(kinds) != 1:
        raise ValueError(f"{list(arrays)} are built for different devices {kinds}")
    traced = [owner for owner in owners.values() if owner.traced]
    sizes = set()
    for owner in traced:
        assert (
            owner.op.trace is not None
        ), "traced is non-empty only when op.trace is set"
        sizes.add(owner.op.trace.trace_size)
    if len(sizes) > 1:
        raise ValueError(
            f"{name}: its traced arrays ask for trace sizes {sorted(sizes)}; one "
            f"configuration has one trace buffer, so give them one size"
        )
    runtimes = (
        [
            MergedRuntime.merge(
                [owners[of[k]].runtime(arrays[k], None) for k in steps], "sequence"
            )
        ]
        if steps
        else [
            owners[of[key]].runtime(array, "sequence" if len(arrays) == 1 else key)
            for key, array in arrays.items()
        ]
    )
    first = next(iter(owners.values()))
    configured = DeviceConfiguration(
        name,
        first.op.array_device(),
        workers=[w for owner in owners.values() for w in owner.workers],
        runtimes=runtimes,
    )
    if traced:
        configured.enable_trace(
            sizes.pop(),
            workers=[w for owner in traced for w in owner.traced],
        )
    return configured

identities ¶

identities(arrays: Mapping[str, BuiltArray]) -> dict[str, Hashable]

Return what each of arrays is, equal where two are one array.

Its array_key, and where another shares that, the words its cores read (BuiltArray.reads).

Source code in iron/common/design/design.py

def identities(arrays: Mapping[str, BuiltArray]) -> dict[str, Hashable]:
    """Return what each of ``arrays`` is, equal where two are one array.

    Its ``array_key``, and where another shares that, the words its cores
    read (``BuiltArray.reads``).
    """
    keys = [a.op.array_key() for a in arrays.values()]
    return {
        name: key if keys.count(key) == 1 else (key, array.reads)
        for (name, array), key in zip(arrays.items(), keys)
    }

Writing an operator¶

Declaration¶

Operand ¶

flat_type property ¶

obj_type property ¶

l1_depth property ¶

batch_axes property ¶

bounded property ¶

bind ¶

pin ¶

fill ¶

drain ¶

extent_axis ¶

fits ¶

split ¶

extent_unit ¶

transfers ¶

Built dataclass ¶

Operator ¶

values property ¶

residents property ¶

name property ¶

step classmethod ¶

build ¶

tolerance ¶

ops ¶

array_device ¶

xclbin_operator ¶

exported_design ¶

shim_columns classmethod ¶

call_operands classmethod ¶

resolve_columns ¶

addressed_inputs ¶

judge ¶

sequence ¶

design_key ¶

resolved ¶

replace ¶

derived_at ¶

extent_unit ¶

value ¶

use_value ¶

compile ¶

layout ¶

Precision ¶

Profile ¶

make ¶

add ¶

tunables_for ¶

dims staticmethod ¶

matches staticmethod ¶

load classmethod ¶

save ¶

Block ¶

Rtp ¶

read ¶

write ¶

Value ¶

per_call property ¶

dispatch property ¶

symbol property ¶

in_scratchpad property ¶

Xclbin ¶

with_instructions ¶

fetch ¶

Elementwise templates¶

Sweep ¶

KernelCall dataclass ¶

Elementwise ¶

valid_elements property ¶

chain ¶

over ¶

reference ¶

judge ¶

kernel ¶

scalars ¶

held_segments ¶

arguments ¶

UnaryElementwise ¶

BinaryElementwise ¶

flat_type `property` ¶

obj_type `property` ¶

l1_depth `property` ¶

batch_axes `property` ¶

bounded `property` ¶

Built `dataclass` ¶

values `property` ¶

residents `property` ¶

name `property` ¶

step `classmethod` ¶

shim_columns `classmethod` ¶

call_operands `classmethod` ¶

dims `staticmethod` ¶

matches `staticmethod` ¶

load `classmethod` ¶

per_call `property` ¶

dispatch `property` ¶

symbol `property` ¶

in_scratchpad `property` ¶

KernelCall `dataclass` ¶

valid_elements `property` ¶

segments_for `classmethod` ¶

reads `cached` `property` ¶

identity `cached` `property` ¶

name `property` ¶