Operators and declarations

An operator class declares what it takes in its body: its parameters (param), the choices a device resolves (tunable), its operands (In, Out, InOut) and the values each call supplies (per_call). The class body is read into a Declaration when it finishes, so the shapes, the cache keys and the per-call words all come from one place. An operator library subclasses Operator with library=True and adds building, tuning and dispatch; amd/IRON's iron.common.declare.Operator does (amd/IRON#222).

Operator and Unresolvable live in aie.utils.compile.jit.operator; param, tunable, per_call and Select in aie.utils.compile.jit.params; In, Out and InOut are exported from aie.utils.compile.jit.

from aie.iron.device import NPU2
from aie.utils.compile.jit import In, Out
from aie.utils.compile.jit.operator import Operator
from aie.utils.compile.jit.params import param, per_call, tunable


class Scale(Operator):
    M: int = param()
    factor: float = param(default=1.0)
    tile: int | None = tunable()
    x = In(M, obj_shape=tile)
    y = Out(M)
    count = per_call(derive=lambda op: op.M // 4)

    def resolve(self, dev):
        return self.replace(tile=self.tile or 16)


device = NPU2()
op = Scale(M=1024)          # Scale(M=1024, factor=1.0, tile=None)
done = op.resolved(device)  # Scale(M=1024, factor=1.0, tile=16)

replace returns a copy with the parameters it names changed. A parameter the caller never gave, computed by a derive or a function default, is computed again from the new values; one the caller gave, a derive one included, keeps its value (pass None to compute it again).

In, Out and InOut take the same arguments: the host dims (integers, param()s and Selects), then dtype, obj_shape (what one ObjectFifo element holds), per (the param() the stream is split over), depth, replicate, when (a bool param() the tensor exists under) and like (an input whose shape this one takes on each call).

A per-call value is an expression in one integer the host gives each call (PerCall and Affine in aie.helpers.taplib). The host evaluates it, and a pattern that needs it on the device stages it as arith.

Operator

Operator

Operator(*, trace: Any = None, **given: Any)

An operator. Subclass it; each subclass's body is its declaration.

Attributes:

Name	Type	Description
declaration	Declaration	What the class declares.
reference_only	bool	The operator has a reference and no kernel yet.
trace	Any	What the build traces, or None.

Source code in python/utils/compile/jit/operator.py

def __init__(self, *, trace: Any = None, **given: Any) -> None:
    params = self.declaration.params
    unknown = [n for n in given if n not in params or not params[n].init]
    if unknown:
        raise TypeError(f"{type(self).__name__}() takes no argument {unknown}")
    missing = [
        n
        for n, p in params.items()
        if p.init and p.default is MISSING and p.derive is None and n not in given
    ]
    if missing:
        raise TypeError(f"{type(self).__name__}() needs {missing}")
    self.trace = trace
    self._given = frozenset(n for n, v in given.items() if v is not None)
    self._set(given)

bound property

bound: Mapping[str, Any]

What each present member reads as on this operator.

validate

validate() -> None

Check the parameters on their own; raise ValueError. Runs at construction.

Source code in python/utils/compile/jit/operator.py

def validate(self) -> None:
    """Check the parameters on their own; raise ``ValueError``. Runs at construction."""

compatible

compatible() -> None

Check the parameters against the resolved tunables; raise ValueError.

Runs on the copy resolved returns.

Source code in python/utils/compile/jit/operator.py

def compatible(self) -> None:
    """Check the parameters against the resolved tunables; raise ``ValueError``.

    Runs on the copy ``resolved`` returns.
    """

resolve

resolve(dev: Device) -> Self

Return a copy with every tunable filled for dev.

Parameters:

Name	Type	Description	Default
dev	Device	The device the tunables are chosen for.	required

Returns:

Type	Description
Self	The resolved copy.

Raises:

Type	Description
Unresolvable	No legal value exists on dev.

Source code in python/utils/compile/jit/operator.py

def resolve(self, dev: Device) -> Self:
    """Return a copy with every tunable filled for ``dev``.

    Args:
        dev: The device the tunables are chosen for.

    Returns:
        The resolved copy.

    Raises:
        Unresolvable: No legal value exists on ``dev``.
    """
    return self.replace()

resolved

resolved(dev: Device) -> Self

Return this operator resolved for dev, which is current while it resolves and is checked.

One already resolved for dev is itself. One that came from an unresolved operator, by resolving or by replace of a resolved one, is resolved for dev from that operator.

Raises:

Type	Description
TypeError	resolve returned self or another class.
Unresolvable	resolve left a tunable None.

Source code in python/utils/compile/jit/operator.py

def resolved(self, dev: Device) -> Self:
    """Return this operator resolved for ``dev``, which is current while it resolves and is checked.

    One already resolved for ``dev`` is itself. One that came from an
    unresolved operator, by resolving or by ``replace`` of a resolved one,
    is resolved for ``dev`` from that operator.

    Raises:
        TypeError: ``resolve`` returned ``self`` or another class.
        Unresolvable: ``resolve`` left a tunable None.
    """
    if self._resolved and self._resolved_for == dev:
        return self
    if self._unresolved is not self:
        return self._unresolved.resolved(dev)
    with set_current_device(dev):
        new = self.resolve(dev)
        if new is self or not isinstance(new, type(self)):
            raise TypeError(
                f"{type(self).__name__}.resolve() must return a copy, self.replace(...)"
            )
        missing = [
            n
            for n, p in self.declaration.params.items()
            if p.tunable and vars(new)[n] is None
        ]
        if missing:
            raise Unresolvable(
                f"{type(self).__name__}.resolve() left {missing} unset for {dev}"
            )
        new.compatible()
    new._resolved = True
    new._resolved_for = dev
    new._unresolved = self
    return new

replace

replace(**changes: Any) -> Self

Return a copy with changes, which may name any parameter.

A parameter not in changes that a derive or a function default computed is computed again; one a caller gave, a derive one included, keeps its value (give it None to compute it again). The copy is not resolved, but its keys keep naming the device its tunables were resolved for, and a copy of a resolved operator resolves for another device from the unresolved one, with changes applied, as the original does.

Source code in python/utils/compile/jit/operator.py

def replace(self, **changes: Any) -> Self:
    """Return a copy with ``changes``, which may name any parameter.

    A parameter not in ``changes`` that a ``derive`` or a function default
    computed is computed again; one a caller gave, a ``derive`` one
    included, keeps its value (give it None to compute it again). The
    copy is not resolved, but its keys keep naming the device its
    tunables were resolved for, and a copy of a resolved operator
    resolves for another device from the unresolved one, with
    ``changes`` applied, as the original does.
    """
    params = self.declaration.params
    unknown = [n for n in changes if n not in params]
    if unknown:
        raise TypeError(f"{type(self).__name__} has no parameter {unknown}")
    kept = {
        n: vars(self)[n]
        for n, p in params.items()
        if n in self._given
        or (p.derive is None and not isinstance(p.default, FunctionType))
    }
    new = object.__new__(type(self))
    new.trace = self.trace
    new._given = self._given.difference(changes) | {
        n for n, v in changes.items() if v is not None
    }
    new._set(kept | changes)
    new._resolved_for = self._resolved_for
    if self._unresolved is not self:
        unresolved = self._unresolved
        new._unresolved = unresolved.replace(**changes) if changes else unresolved
    return new

copy

copy() -> Self

Return a fresh instance with every value, still resolved as this one is.

Source code in python/utils/compile/jit/operator.py

def copy(self) -> Self:
    """Return a fresh instance with every value, still resolved as this one is."""
    new = self.replace()
    new._resolved = self._resolved
    return new

layout

layout() -> Self

Return this operator as its cores see it: the key's parameters and the members.

Source code in python/utils/compile/jit/operator.py

def layout(self) -> Self:
    """Return this operator as its cores see it: the key's parameters and the members."""
    view = object.__new__(type(self))
    vars(view).update(
        {n: vars(self)[n] for n in self.declaration.key}, trace=self.trace
    )
    # Bound to the view, a member's dims read through it, so a parameter
    # outside the key raises there too.
    view._bound = {
        m.name: m.of(view)
        for m in self.declaration.members
        if m.name in self._bound
    }
    return view

array_key

array_key() -> tuple

Return the identity two operators share their cores by.

It names the device this operator is resolved for (None until resolved), since the resolved tunables depend on it.

Source code in python/utils/compile/jit/operator.py

def array_key(self) -> tuple:
    """Return the identity two operators share their cores by.

    It names the device this operator is resolved for (None until
    resolved), since the resolved tunables depend on it.
    """
    dev = self._resolved_for
    key = (
        type(self),
        ("device", None if dev is None else dev.name),
        *((n, _key_value(vars(self)[n])) for n in self.declaration.key),
    )
    return key if self.trace is None else (*key, ("trace", self.trace))

design_key

design_key() -> tuple

Return the identity two operators share a build by: the device and every parameter.

Source code in python/utils/compile/jit/operator.py

def design_key(self) -> tuple:
    """Return the identity two operators share a build by: the device and every parameter."""
    dev = self._resolved_for
    return (
        type(self),
        (
            ("device", None if dev is None else dev.name),
            ("trace", self.trace),
            *((n, _key_value(vars(self)[n])) for n in self.declaration.params),
        ),
    )

Unresolvable

Bases: ValueError

No legal resolution exists for this operator on this device.

Parameters and per-call values

The words an operator's class body declares its values with.

M: int = param() is a compile-time parameter, tile: int = tunable() one a tuner may choose, and n = per_call(np.int32) a value that may change every call. Each stays the class attribute it was written as, so In(M, K) captures M itself and GEMM.M names the parameter. On an instance, op.M is the integer, held in the instance's own __dict__.

OperatorParam

OperatorParam(
    *, default, derive, array, repr, init, tunable, domain
)

What param() and tunable() return: a parameter of its class.

Attributes:

Name	Type	Description
default		The value when the caller gives none, a function of the operator computing it, or MISSING for a required argument. A tunable's default is its starting value, None for one resolve fills.
derive		A function of the operator a given value must equal, and the value when none is given; or None.
array		The cores read it, though no operand's shape names it.
repr		It is in the operator's repr and device symbols.
init		It is a constructor argument.
tunable		A tuner may choose it.
domain		A tunable's search space, or None.

Source code in python/utils/compile/jit/params.py

def __init__(self, *, default, derive, array, repr, init, tunable, domain):
    self.default = default
    self.derive = derive
    self.array = array
    self.repr = repr
    self.init = init
    self.tunable = tunable
    self.domain = domain
    self.name = ""
    self.owner: type | None = None

Select dataclass

Select(
    flag: OperatorParam, when_true: tuple, when_false: tuple
)

A shape chosen by a flag the caller gives or defaults, never inferred.

Select(b_col_maj, (N, K), (K, N)).

branch

branch(op) -> tuple

Return the dimensions the flag picks on op.

Source code in python/utils/compile/jit/params.py

def branch(self, op) -> tuple:
    """Return the dimensions the flag picks on ``op``."""
    flag = getattr(op, self.flag.name)
    if flag is None:
        raise ValueError(
            f"{self.flag.name} is None; a Select() on it needs a resolved operator"
        )
    return self.when_true if flag else self.when_false

param

param(
    *,
    default: Any = MISSING,
    derive: Callable[[Any], Any] | None = None,
    array: bool = False,
    repr: bool = True,
    init: bool = True
) -> Any

Declare a compile-time parameter.

The caller gives it, or the operands' shapes infer it.

Parameters:

Name	Type	Description	Default
default	Any	The value when neither the caller nor an operand gives one. A lambda or def is a function of the operator computing it; anything else, a type included, is the value.	MISSING
derive	Callable[[Any], Any] | None	A function of the operator: the value when none is given, and what a given one must equal.	None
array	bool	The cores read it (a kernel scalar, a flag the build branches on) though no operand's shape, dtype or presence names it.	False
repr	bool	It is in the operator's repr and device symbols.	True
init	bool	It is a constructor argument.	True

Returns:

Type	Description
Any	The OperatorParam the class body binds to the parameter's name.

Raises:

Type	Description
TypeError	Both default and derive are given, or default is a list, dict or set, which every instance would share.

Source code in python/utils/compile/jit/params.py

def param(
    *,
    default: Any = MISSING,
    derive: Callable[[Any], Any] | None = None,
    array: bool = False,
    repr: bool = True,
    init: bool = True,
) -> Any:
    """Declare a compile-time parameter.

    The caller gives it, or the operands' shapes infer it.

    Args:
        default: The value when neither the caller nor an operand gives one.
            A ``lambda`` or ``def`` is a function of the operator computing
            it; anything else, a type included, is the value.
        derive: A function of the operator: the value when none is given,
            and what a given one must equal.
        array: The cores read it (a kernel scalar, a flag the build branches
            on) though no operand's shape, dtype or presence names it.
        repr: It is in the operator's repr and device symbols.
        init: It is a constructor argument.

    Returns:
        The ``OperatorParam`` the class body binds to the parameter's name.

    Raises:
        TypeError: Both `default` and `derive` are given, or `default` is a
            list, dict or set, which every instance would share.
    """
    if default is not MISSING and derive is not None:
        raise TypeError("param(default=, derive=): give one, not both")
    if isinstance(default, (list, dict, set)):
        raise TypeError(
            f"param(default={default!r}) would be shared by every instance; give default=lambda op: {default!r}"
        )
    return OperatorParam(
        default=default,
        derive=derive,
        array=array,
        repr=repr,
        init=init,
        tunable=False,
        domain=None,
    )

tunable

tunable(
    *,
    default: Any = None,
    domain: Any = None,
    repr: bool = True,
    init: bool = True
) -> Any

Declare a compile-time parameter a tuner may search.

resolve fills it for the device when the caller does not. The cores always read it.

With init=False it is no constructor argument: a value resolve computes, or a subclass's fixed value of an inherited tunable.

Parameters:

Name	Type	Description	Default
default	Any	The starting value; None means resolve fills it.	None
domain	Any	The values a tuner searches: a tuple, an Enum class, or a function (op, dev, start) -> tuple such as widths. None is not searched.	None
repr	bool	It is in the operator's repr and device symbols.	True
init	bool	It is a constructor argument.	True

Returns:

Type	Description
Any	The OperatorParam the class body binds to the parameter's name.

Raises:

Type	Description
TypeError	domain is none of those.

Source code in python/utils/compile/jit/params.py

def tunable(
    *,
    default: Any = None,
    domain: Any = None,
    repr: bool = True,
    init: bool = True,
) -> Any:
    """Declare a compile-time parameter a tuner may search.

    ``resolve`` fills it for the device when the caller does not. The cores
    always read it.

    With ``init=False`` it is no constructor argument: a value ``resolve``
    computes, or a subclass's fixed value of an inherited tunable.

    Args:
        default: The starting value; None means ``resolve`` fills it.
        domain: The values a tuner searches: a tuple, an ``Enum`` class, or
            a function ``(op, dev, start) -> tuple`` such as `widths`. None
            is not searched.
        repr: It is in the operator's repr and device symbols.
        init: It is a constructor argument.

    Returns:
        The ``OperatorParam`` the class body binds to the parameter's name.

    Raises:
        TypeError: `domain` is none of those.
    """
    if not (
        domain is None
        or isinstance(domain, tuple)
        or (isinstance(domain, type) and issubclass(domain, Enum))
        or isinstance(domain, (types.FunctionType, functools.partial))
    ):
        raise TypeError(
            f"tunable(domain={domain!r}): give a tuple, an Enum or a function "
            f"(op, dev, start) -> tuple"
        )
    return OperatorParam(
        default=default,
        derive=None,
        array=True,
        repr=repr,
        init=init,
        tunable=True,
        domain=domain,
    )

per_call

per_call(
    dtype: Any = int32,
    *,
    derive: Callable[[Any], Any] | None = None,
    bounds: OperatorParam | None = None
) -> Any

Declare a value the cores read that may change every call without a rebuild.

A value no build fixes (no derive, no bounds) is given every call: an operator with a dispatch library takes it as a dispatch scalar, which rebuilds the instruction stream around it.

Parameters:

Name	Type	Description	Default
dtype	Any	An integer type.	int32
derive	Callable[[Any], Any] | None	A function of the operator computing it once per build.	None
bounds	OperatorParam | None	A param() this value bounds per call: valid = per_call(bounds=M).	None

Returns:

Type	Description
Any	The PerCallMember the class body binds to the value's name.

Raises:

Type	Description
TypeError	dtype is not an integer type, or both derive and bounds are given.

Source code in python/utils/compile/jit/params.py

def per_call(
    dtype: Any = np.int32,
    *,
    derive: Callable[[Any], Any] | None = None,
    bounds: OperatorParam | None = None,
) -> Any:
    """Declare a value the cores read that may change every call without a rebuild.

    A value no build fixes (no `derive`, no `bounds`) is given every call:
    an operator with a dispatch library takes it as a dispatch scalar,
    which rebuilds the instruction stream around it.

    Args:
        dtype: An integer type.
        derive: A function of the operator computing it once per build.
        bounds: A ``param()`` this value bounds per call: ``valid =
            per_call(bounds=M)``.

    Returns:
        The ``PerCallMember`` the class body binds to the value's name.

    Raises:
        TypeError: `dtype` is not an integer type, or both `derive` and
            `bounds` are given.
    """
    if np.dtype(dtype).kind not in "iu":
        raise TypeError(
            f"per_call({np.dtype(dtype).name}): a per-call value is an integer, "
            f"since the scratchpad encoding zeroes the top two bits of a word"
        )
    if derive is not None and bounds is not None:
        raise TypeError("per_call(derive=, bounds=): give one, not both")
    return PerCallMember(dtype, derive, bounds)

Operands

Argument markers for JIT design parameters.

Five annotation categories are defined here (all exported from aie.iron):

CompileTime[T] Marks a generator function parameter as compile-time. Changing its value causes a recompile and a new cache entry. Inspired by tl.constexpr in Triton. Implemented as Annotated[T, ...] so pyright/mypy check callers against the wrapped type T itself (e.g. CompileTime[int] behaves like int for type-checking) while _introspect.py still recovers the marker at runtime via get_type_hints(..., include_extras=True).

In Marks a generator function parameter as a runtime input tensor. Data is DMA-transferred from the host to the NPU on every kernel call.

Out Marks a generator function parameter as a runtime output tensor. Data is DMA-transferred from the NPU to the host on every kernel call.

InOut Marks a generator function parameter as a runtime bidirectional tensor. Data is DMA-transferred in both directions on every kernel call.

A generator receives each tensor as an operand of its marker's class, never its data: an In/Out/InOut with the shape and dtype of the tensor the design is called with, or of a type bound ahead of time with specialize(name=np.ndarray[...]), as an operator's op.a is. Every API that takes a tensor type also takes the operand; operand.ndarray_type is its np.ndarray[shape, np.dtype[T]]. A different shape or dtype is a different compiled design.

DispatchTime[T] Marks a keyword-only integer scalar that can vary per dispatch without recompiling. Explicit specialization instead fixes it at compile time. See DispatchTime below for generator binding and usage restrictions.

Any parameter without one of these annotations is rejected at @iron.jit decoration time when the parameter has a default value: an unannotated scalar is bound at generation time, so the default would be baked into the compiled kernel and per-call overrides silently ignored. Annotate as CompileTime[T] (recompiles on change), DispatchTime[T] (runtime scalar, one compile many values), or In/Out/InOut (DMA tensor) instead.

In

In(
    *dims: Any,
    dtype: Any = bfloat16,
    obj_shape: Any = None,
    per: Any = None,
    depth: int = 2,
    replicate: bool = False,
    when: OperatorParam | None = None,
    like: _Operand | None = None
)

Bases: _Operand

Runtime input tensor (host → NPU, DMA each call).

Out and InOut take the same arguments.

Parameters:

Name	Type	Description	Default
dims	Any	The host shape: integers, param()s and Selects. A param() that is None on an instance is a dimension the shape leaves out there.	()
dtype	Any	The element type, or a param() holding it.	bfloat16
obj_shape	Any	What one ObjectFifo element holds. With one the tensor streams into the array; without one a hand-written sequence moves it.	None
per	Any	The param() (or several, multiplied) the stream is split over, one ObjectFifo each.	None
depth	int	The ObjectFifo depth.	2
replicate	bool	Every per= lane gets the whole tensor, not a share.	False
when	OperatorParam | None	A bool param(); the tensor exists only where it is true.	None
like	_Operand | None	An input whose dims these are; a caller gives this tensor the shape that input has on the call, so Out(like=x) of a flat x is (rows, cols) where x is.	None

Raises:

Type	Description
TypeError	replicate without per, or both dims and like.

Source code in python/utils/compile/jit/markers.py

def __init__(
    self,
    *dims: Any,
    dtype: Any = bfloat16,
    obj_shape: Any = None,
    per: Any = None,
    depth: int = 2,
    replicate: bool = False,
    when: OperatorParam | None = None,
    like: _Operand | None = None,
):
    if replicate and per is None:
        raise TypeError(
            "replicate=True needs per=<dim>: every lane receives the whole tensor"
        )
    if like is not None and dims:
        raise TypeError(f"give dims {dims} or like={like!r}, not both")
    self.like = like
    self.dims = dims if like is None else like.dims
    self.declared_dtype = dtype
    self.obj_dims = (
        None
        if obj_shape is None
        else (
            tuple(obj_shape)
            if isinstance(obj_shape, (tuple, list))
            else (obj_shape,)
        )
    )
    self.per = (
        None
        if per is None
        else tuple(per) if isinstance(per, (tuple, list)) else (per,)
    )
    self.depth = depth
    self.replicate = replicate
    self.when = when

Out

Out(
    *dims: Any,
    dtype: Any = bfloat16,
    obj_shape: Any = None,
    per: Any = None,
    depth: int = 2,
    replicate: bool = False,
    when: OperatorParam | None = None,
    like: _Operand | None = None
)

Bases: _Operand

Runtime output tensor (NPU → host, DMA each call).

Takes the arguments In documents.

Source code in python/utils/compile/jit/markers.py

def __init__(
    self,
    *dims: Any,
    dtype: Any = bfloat16,
    obj_shape: Any = None,
    per: Any = None,
    depth: int = 2,
    replicate: bool = False,
    when: OperatorParam | None = None,
    like: _Operand | None = None,
):
    if replicate and per is None:
        raise TypeError(
            "replicate=True needs per=<dim>: every lane receives the whole tensor"
        )
    if like is not None and dims:
        raise TypeError(f"give dims {dims} or like={like!r}, not both")
    self.like = like
    self.dims = dims if like is None else like.dims
    self.declared_dtype = dtype
    self.obj_dims = (
        None
        if obj_shape is None
        else (
            tuple(obj_shape)
            if isinstance(obj_shape, (tuple, list))
            else (obj_shape,)
        )
    )
    self.per = (
        None
        if per is None
        else tuple(per) if isinstance(per, (tuple, list)) else (per,)
    )
    self.depth = depth
    self.replicate = replicate
    self.when = when

InOut

InOut(
    *dims: Any,
    dtype: Any = bfloat16,
    obj_shape: Any = None,
    per: Any = None,
    depth: int = 2,
    replicate: bool = False,
    when: OperatorParam | None = None,
    like: _Operand | None = None
)

Bases: _Operand

Runtime bidirectional tensor (DMA in both directions each call).

Takes the arguments In documents.

Source code in python/utils/compile/jit/markers.py

def __init__(
    self,
    *dims: Any,
    dtype: Any = bfloat16,
    obj_shape: Any = None,
    per: Any = None,
    depth: int = 2,
    replicate: bool = False,
    when: OperatorParam | None = None,
    like: _Operand | None = None,
):
    if replicate and per is None:
        raise TypeError(
            "replicate=True needs per=<dim>: every lane receives the whole tensor"
        )
    if like is not None and dims:
        raise TypeError(f"give dims {dims} or like={like!r}, not both")
    self.like = like
    self.dims = dims if like is None else like.dims
    self.declared_dtype = dtype
    self.obj_dims = (
        None
        if obj_shape is None
        else (
            tuple(obj_shape)
            if isinstance(obj_shape, (tuple, list))
            else (obj_shape,)
        )
    )
    self.per = (
        None
        if per is None
        else tuple(per) if isinstance(per, (tuple, list)) else (per,)
    )
    self.depth = depth
    self.replicate = replicate
    self.when = when

Declaration

Declaration

Declaration(
    name: str,
    params: dict[str, OperatorParam],
    members: tuple[Member, ...],
    nullable: frozenset[str] = frozenset(),
    scalars: tuple[str, ...] = (),
    signature: Signature | None = None,
)

What an operator declares, read from a class body or a generator's signature.

Attributes:

Name	Type	Description
name		The class's or the generator's name.
params		Every compile-time parameter (param(), tunable(), CompileTime[T]), by name, a base's first.
members		The operands and per-call values, the most derived body's order first.
nullable		The parameters annotated T | None: a shape dimension of one is present only where it is not None.
scalars		A generator's parameters with no marker.
signature		A generator's signature, its annotations evaluated, which binds a call's arguments.
tunables		The tunables a caller or the tuner sets (init=True).
key		The parameters the cores are built from: every tunable, what a stream's obj_shape=, per=, dtype= or when= names, and every param(array=True).

Source code in python/utils/compile/jit/declaration.py

def __init__(
    self,
    name: str,
    params: dict[str, OperatorParam],
    members: tuple[Member, ...],
    nullable: frozenset[str] = frozenset(),
    scalars: tuple[str, ...] = (),
    signature: inspect.Signature | None = None,
) -> None:
    self.name = name
    self.params = params
    self.members = members
    self.nullable = nullable
    self.scalars = scalars
    self.signature = signature
    self.tunables = tuple(n for n, p in params.items() if p.tunable and p.init)
    named = {n for n, p in params.items() if p.array}
    for m in self.operands:
        if m.streamed:
            named |= _names(m.obj_dims) | _names(m.per or ())
            if isinstance(m.declared_dtype, OperatorParam):
                named.add(m.declared_dtype.name)
            if m.when is not None:
                named.add(m.when.name)
    self.key = tuple(n for n in params if n in named)

infer

infer(
    inputs: Mapping[str, Any],
    outputs: Sequence[Any] = (),
    **kwargs: Any
) -> dict[str, Any]

Return the parameters the operands' shapes give.

Each declared dimension is a param() or a literal, so this is a lookup. An optional input's when= flag is true where inputs names it.

Parameters:

Name	Type	Description	Default
inputs	Mapping[str, Any]	The input shapes, by operand name.	required
outputs	Sequence[Any]	The shapes of caller-supplied outputs, in declaration order.	()
**kwargs	Any	Construction arguments; the dimensions and flags among them pin values and are checked for agreement.	{}

Returns:

Type	Description
dict[str, Any]	{param: value}: the pinned and the inferred parameters, and
dict[str, Any]	the operand flags.

Raises:

Type	Description
TypeError	The operands are not the declared ones.
ValueError	A shape disagrees with the declaration or another operand.

Source code in python/utils/compile/jit/declaration.py

def infer(
    self,
    inputs: Mapping[str, Any],
    outputs: Sequence[Any] = (),
    **kwargs: Any,
) -> dict[str, Any]:
    """Return the parameters the operands' shapes give.

    Each declared dimension is a ``param()`` or a literal, so this is a
    lookup. An optional input's ``when=`` flag is true where ``inputs``
    names it.

    Args:
        inputs: The input shapes, by operand name.
        outputs: The shapes of caller-supplied outputs, in declaration
            order.
        **kwargs: Construction arguments; the dimensions and flags among
            them pin values and are checked for agreement.

    Returns:
        ``{param: value}``: the pinned and the inferred parameters, and
        the operand flags.

    Raises:
        TypeError: The operands are not the declared ones.
        ValueError: A shape disagrees with the declaration or another
            operand.
    """
    operands = self.operands
    names = set(self.params)
    for m in operands:
        names |= _names(m.dims)
    given = {k: v for k, v in kwargs.items() if k in names}
    flags: dict[str, bool] = {}
    for m in operands:
        if not m.fills or m.when is None:
            continue
        flag, has = m.when.name, m.name in inputs
        if flags.setdefault(flag, has) != has:
            raise TypeError(
                f"{self.name}: the operands {flag}= brings are given together"
            )
        if flag in given and bool(given[flag]) != has:
            raise TypeError(
                f"{self.name}({flag}=True) takes {m.name}="
                if not has
                else f"{self.name}: {m.name}= is an operand only where {flag} is true"
            )
    given |= flags
    # An operand's flag left out is its plain default; a computed one is unknown, so false.
    present = {
        m.name
        for m in operands
        if m.when is None
        or (
            (v := given.get(m.when.name, m.when.default)) is not MISSING
            and not isinstance(v, FunctionType)
            and bool(v)
        )
    }
    ins = [m for m in operands if m.fills and m.name in present]
    if [m.name for m in ins] != list(inputs):
        raise TypeError(
            f"{self.name} takes {len(ins)} operand(s) ({', '.join(m.name for m in ins)}), got {len(inputs)}"
        )
    for m in operands:
        if outputs and not m.fills and m.when is not None:
            v = given.get(m.when.name, m.when.default)
            if v is MISSING or isinstance(v, FunctionType):
                raise ValueError(
                    f"{self.name}: {m.when.name} decides whether {m.name} is an output and has no default; pass it explicitly"
                )
    outs = [m for m in operands if not m.fills and m.name in present]
    if outputs and len(outputs) != len(outs):
        raise TypeError(
            f"{self.name} produces {len(outs)} output(s) ({', '.join(m.name for m in outs)}), got {len(outputs)}"
        )
    bound: dict[str, Any] = dict(given)
    # Each dimension an operand's shape gives: (param, value, where).
    found: list[tuple[OperatorParam, int | None, str]] = []
    for m, shape in [(m, inputs[m.name]) for m in ins] + list(zip(outs, outputs)):
        shape = tuple(int(s) for s in shape)
        dims = list(m.dims)
        at = next(
            (
                i
                for i, d in enumerate(dims)
                if isinstance(d, OperatorParam) and d.name in self.nullable
            ),
            None,
        )
        if at is not None:
            optional = dims.pop(at)
            if len(shape) == len(dims) + 1:
                found.append((optional, shape[at], f"{m.name}.shape[{at}]"))
                shape = shape[:at] + shape[at + 1 :]
            elif len(shape) == len(dims):
                found.append((optional, None, f"{m.name} (rank {len(shape)})"))
            else:
                raise ValueError(
                    f"{self.name}: operand {m.name} has rank {len(shape)}, declared {m!r}"
                )
        expanded: list = []
        while dims:
            d = dims.pop(0)
            if not isinstance(d, Select):
                expanded.append(d)
                continue
            value = bound.get(d.flag.name, d.flag.default)
            if value is MISSING or isinstance(value, FunctionType):
                raise ValueError(
                    f"{self.name}: {d.flag.name} selects {m.name}'s shape and has no default; pass it explicitly"
                )
            dims[:0] = d.when_true if value else d.when_false
        dims = expanded
        if len(dims) == 1 and len(shape) != 1:
            # A flat operand takes a tensor of any rank: its one
            # dimension is the element count.
            shape = (math.prod(shape),)
        if len(shape) != len(dims):
            raise ValueError(
                f"{self.name}: operand {m.name} has rank {len(shape)} {shape}, declared rank {len(dims)} {m!r}"
            )
        for i, (d, n) in enumerate(zip(dims, shape)):
            if isinstance(d, OperatorParam):
                found.append((d, n, f"{m.name}.shape[{i}]"))
            elif int(d) != n:
                raise ValueError(
                    f"{self.name}: operand {m.name}.shape[{i}] is {n}, declared {d}"
                )
    for m, shape in zip(outs, outputs):
        if m.like is None:
            continue
        like = tuple(inputs[m.like.name])
        if tuple(shape) != like:
            raise ValueError(
                f"{self.name}: output {m.name} is {tuple(shape)}, but shaped like {m.like.name}, {like}"
            )
    origin = {k: "given" for k in given}
    for ref, value, where in found:
        if ref.name in bound and bound[ref.name] != value:
            raise ValueError(
                f"{self.name}: {ref.name} is {value} from {where} but {bound[ref.name]} from {origin[ref.name]}"
            )
        bound[ref.name] = value
        origin.setdefault(ref.name, where)
    return bound

Per-call expressions

PerCall

PerCall(name: str, marker: Any, scalar_type: type)

An integer scalar given per call.

Integer arithmetic on one makes an Affine. marker says how a call gives it, and only its library reads it: "aie.dispatch_time" for a DispatchTime parameter. ssa is its staged value once a runtime sequence binds it.

Source code in python/helpers/taplib/_per_call.py

def __init__(self, name: str, marker: Any, scalar_type: type):
    self.name, self.marker, self.scalar_type = name, marker, scalar_type
    self.ssa: Any = None

Affine dataclass

Affine(
    value: PerCall,
    scale: int = 1,
    bias: int = 0,
    down: int = 0,
    mul: int = 1,
    add: int = 0,
    lo: int | None = None,
    hi: int | None = None,
    mod: int | None = None,
    post: int = 0,
)

An integer expression in one per-call value v.

clamp(((v * scale + bias) >> down) * mul + add, lo, hi) % mod + post

>> floors, a bound of None is none, and mod is a power of two (a ring's row) or None. post is added after the wrap, so a run of a ring written twice over, ring[a : a + R], has its end. A linear one has down 0 and keeps its terms in scale and bias: (p + 1) * 64 is Affine(p, 64, 64).

Sums and products with integers, floor divisions by a power of two (so ceildiv), sums of two linear ones in one value, and the quotient of two that are multiples of one another ((4 * n) // n is 4) are expressions too. A clamp is carried through a sum, a product or a floor division by an integer. Anything else (a comparison, another divisor, two values) raises TypeError.

name property

name: str

An identifier for the device symbol: p, p_x64, p_x64_p64, p_m1.

affine

affine() -> Affine

Return this expression where only a linear one will do.

Raises:

Type	Description
TypeError	It floors, is clamped or wraps.

Source code in python/helpers/taplib/_per_call.py

def affine(self) -> Affine:
    """Return this expression where only a linear one will do.

    Raises:
        TypeError: It floors, is clamped or wraps.
    """
    if not self.linear:
        raise TypeError(f"{self} is not linear in {self.value.name}")
    return self

clamp

clamp(
    lo: int | None = None, hi: int | None = None
) -> Affine | int

Return this expression held to [lo, hi] too, a bound of None none.

Held to a range it never reaches, it is that range's nearer bound.

Raises:

Type	Description
TypeError	The expression wraps.
ValueError	lo exceeds hi.

Source code in python/helpers/taplib/_per_call.py

def clamp(self, lo: int | None = None, hi: int | None = None) -> Affine | int:
    """Return this expression held to ``[lo, hi]`` too, a bound of None none.

    Held to a range it never reaches, it is that range's nearer bound.

    Raises:
        TypeError: The expression wraps.
        ValueError: ``lo`` exceeds ``hi``.
    """
    if self.mod is not None:
        raise TypeError(f"no expression clamps ({self}) after it wraps")
    if lo is not None and hi is not None and lo > hi:
        raise ValueError(f"{self} held to [{lo}, {hi}] holds no value")
    if lo is not None and self.hi is not None and self.hi < lo:
        return lo
    if hi is not None and self.lo is not None and self.lo > hi:
        return hi
    if lo is None or self.lo is not None and self.lo > lo:
        lo = self.lo
    if hi is None or self.hi is not None and self.hi < hi:
        hi = self.hi
    return dataclasses.replace(self, lo=lo, hi=hi)

evaluate

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

Its value where values gives its per-call value by name.

Source code in python/helpers/taplib/_per_call.py

def evaluate(self, values: Mapping[str, int]) -> int:
    """Its value where ``values`` gives its per-call value by name."""
    v = int(values[self.value.name])
    x = ((self.scale * v + self.bias) >> self.down) * self.mul + self.add
    if self.lo is not None:
        x = max(x, self.lo)
    if self.hi is not None:
        x = min(x, self.hi)
    return x if self.mod is None else x % self.mod + self.post

stage

stage()

Stage this expression on its value's i64, as eudsl's ScalarValue.

Raises:

Type	Description
TypeError	No runtime sequence has bound the value.

Source code in python/helpers/taplib/_per_call.py

def stage(self):
    """Stage this expression on its value's ``i64``, as eudsl's ``ScalarValue``.

    Raises:
        TypeError: No runtime sequence has bound the value.
    """
    # _symbolic imports this module.
    from ._symbolic import sint

    if self.value.ssa is None:
        raise TypeError(
            f"{self.value.name} has no staged value outside the runtime "
            "sequence it is bound to"
        )
    x = ScalarValue(sint(self.value.ssa))
    x = x * self.scale + self.bias
    if self.down:
        x = x // (1 << self.down)
    x = x * self.mul + self.add
    # ScalarValue has no min or max.
    if self.lo is not None:
        x = ScalarValue(arith.maxsi(x, constant(self.lo, T.i64())))
    if self.hi is not None:
        x = ScalarValue(arith.minsi(x, constant(self.hi, T.i64())))
    if self.mod is not None:  # a power of two: the floored remainder
        x = (x & (self.mod - 1)) + self.post
    return x