Graphs

A graph is a subclass of iron.Graph whose body() is recorded on handles, once per input shape, and compiled to one image per shape. These are the names import iron gives.

IRON: operators for the NPU, and graphs over them.

Its names are imported on first use, so import iron stays light.

Carried module-attribute

Carried = Annotated[T, Kind.CARRIED.value]

A scratchpad value a graph computes for its next call and returns last with iron.carry(name=...): an expression of its values (position + 1) or a one-element handle (a sampled token); the host seeds only the first call.

Scratchpad module-attribute

Scratchpad = Annotated[T, Kind.SCRATCHPAD.value]

A per-call value patched into a DMA descriptor or read by a core: free per call, works under full ELF, cannot change a DMA size or stride; 30 bits at most.

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")

CarriedLoop

CarriedLoop(first: CompiledGraph, body: CompiledGraph, depth: int = 2)

A graph's carried values looped on the device.

first runs once from the host, then body (no tensor) with depth runs in flight, each run's Emit writing the next's scratchpad, so the host only waits and restarts.

Source code in iron/common/graph/carried.py

def __init__(self, first: CompiledGraph, body: CompiledGraph, depth: int = 2):
    for version in (first, body):
        if version.emit is None:
            raise ValueError(
                f"{version.recording.name}: a full-ELF version with carried values "
                f"loops; this one has no Emit step (compile one that takes a "
                f"tensor with feeds= the version it starts)"
            )
    if first.arena is None or first.arena is not body.arena:
        raise ValueError("first and body are versions of one graph")
    if body.recording.inputs:
        raise ValueError(
            f"{body.recording.name}: the looped version takes no tensor, got "
            f"{[h.name for h in body.recording.inputs]}"
        )
    if depth < 1:
        raise ValueError(f"depth is at least 1, got {depth}")
    self.first, self.body, self.depth = first, body, depth
    first.emit_to(body)
    body.emit_to(body)
    # Each run holds a copy of its ELF's control code in the device heap,
    # 64 MB on older amdxdna, so the loop reuses each version's own run.
    first_call, body_call = first.callable, body.callable
    # Only a full ELF has an Emit step.
    assert isinstance(first_call, FullELFCallable)
    assert isinstance(body_call, FullELFCallable)
    self._first_call, self._body_call = first_call, body_call
    # Built on the first start(): first runs once and, where it is not
    # body itself, is dead after, so the heap holds first's ctrlcode and
    # body's first copy, then body's depth copies, never 1 + depth.
    self._runs: list[FullELFRun] | None = None

run

run(steps: int, /, *tensors, **values) -> Iterator[int]

start now, then steps, each index yielded once complete.

Source code in iron/common/graph/carried.py

def run(self, steps: int, /, *tensors, **values) -> Iterator[int]:
    """``start`` now, then ``steps``, each index yielded once complete."""
    self.start(*tensors, **values)
    return self.steps(steps)

start

start(*tensors, **values) -> None

Run first on tensors and values, and wait for it.

Source code in iron/common/graph/carried.py

def start(self, *tensors, **values) -> None:
    """Run ``first`` on ``tensors`` and ``values``, and wait for it."""
    self.body.upload()
    first_run = self._first_call.run
    runs = self._runs or [self._body_call.run]
    first_run.bind_feedback(runs[0].scratchpad_alias())
    self.first.start(first_run, *tensors, **values)
    self._first_call.wait(first_run)
    if self._first_call is not self._body_call:
        # first never runs again: free its ctrlcode heap copy, as any
        # other callable's idle run is (FullELFCallable.release()), not
        # just when a later new_run() happens to need the room.
        first_run.wired = False
        self._first_call.release()
    # pyxrt frees a run's ctrlcode only once nothing references the run.
    del first_run
    if self._runs is None:
        runs += [self._body_call.new_run() for _ in range(self.depth - 1)]
        for k, run in enumerate(runs):
            run.bind_feedback(runs[(k + 1) % self.depth].scratchpad_alias())
        self._runs = runs

steps

steps(steps: int) -> Iterator[int]

Run steps steps of body, yielding each index once complete.

Closing early waits out the steps in flight.

Source code in iron/common/graph/carried.py

def steps(self, steps: int) -> Iterator[int]:
    """Run ``steps`` steps of ``body``, yielding each index once complete.

    Closing early waits out the steps in flight.
    """
    assert self._runs is not None, "start() builds body's runs first"
    in_flight: deque[int] = deque()
    for k in range(min(self.depth, steps)):
        self._runs[k].start()
        in_flight.append(k)
    try:
        for k in range(steps):
            run = self._runs[k % self.depth]
            self._body_call.wait(run)
            in_flight.popleft()
            if k + self.depth < steps:
                run.start()
                in_flight.append(k + self.depth)
            yield k
    finally:
        self._body_call.wait(*(self._runs[k % self.depth] for k in in_flight))

CompiledGraph

CompiledGraph(recording: Recording, plan: Plan, check: bool = False, arena: ScratchArena | None = None, emit: EmitSite | None = None)

A recorded graph built into an image, ready to call.

Called with the per-call values this version reads (Recording.reads); the graph's others go unused.

Source code in iron/common/graph/compiled.py

def __init__(
    self,
    recording: Recording,
    plan: Plan,
    check: bool = False,
    arena: ScratchArena | None = None,
    emit: EmitSite | None = None,
):
    self.recording = recording
    self.plan = plan
    self.arena = arena
    self.emit = emit
    self.words, self.shared = _words(
        recording, share=plan.full_elf, extents=not plan.full_elf
    )
    packaging = (
        FullELF(commands=plan.commands, shared_words=self.shared)
        if plan.full_elf
        else PerStep(check)
    )
    placement: dict[str, Any] = (
        {}
        if arena is None
        else dict(arena=arena.plan, residents=recording.residents)
    )
    self.sequence = Image(recording, packaging, **placement).compile()
    self.image = self.sequence.image
    self.artifacts = self.sequence.artifacts
    if plan.full_elf:
        # An extent read only through its derivations has no word.
        self.words = [
            w for w in self.words if w.symbol in self.artifacts.parameters
        ]
    self._callable = None
    self._addressed: dict[str, tuple[Any, int]] = {}
    # Shared with the arena when there is one: its weights are every image's.
    self._loaded: set = set() if arena is None else arena.loaded

callable property

callable: FullELFCallable | PerStepCallable

The loaded image, made on first use so a host without an NPU can compile.

release

release() -> None

Free this version's runs, image and the buffers it holds outside the arena.

Source code in iron/common/graph/compiled.py

def release(self) -> None:
    """Free this version's runs, image and the buffers it holds outside the arena."""
    if isinstance(self._callable, FullELFCallable):
        self._callable.release(force=True)
    self._callable = None
    self._addressed.clear()
    if self.arena is None:
        self._loaded.clear()

read

read(x)

x's buffer as the device left it: a view of its host memory.

Source code in iron/common/graph/compiled.py

def read(self, x):
    """``x``'s buffer as the device left it: a view of its host memory."""
    buf = self._storage(x)
    return buf.numpy()[: int(np.prod(x.shape))].reshape(tuple(x.shape))

push

push(x) -> None

Send to the device what the host wrote into read(x)'s view.

Source code in iron/common/graph/compiled.py

def push(self, x) -> None:
    """Send to the device what the host wrote into ``read(x)``'s view."""
    buf = self._storage(x)
    buf.numpy_view()  # marks the host copy the newer, so to() sends it
    buf.to("npu")

upload

upload(release: Callable[[ndarray], None] | None = None, piece_bytes: int = UPLOAD_PIECE) -> None

Copy every closed-over weight into its buffer, once per storage.

Source code in iron/common/graph/compiled.py

def upload(
    self,
    release: Callable[[np.ndarray], None] | None = None,
    piece_bytes: int = UPLOAD_PIECE,
) -> None:
    """Copy every closed-over weight into its buffer, once per storage."""
    for key, (tensor, handle) in self.recording.weights.items():
        if key not in self._loaded:
            self._copy_in(handle.name, tensor, release, piece_bytes)
            self._loaded.add(key)

load

load(release: Callable[[ndarray], None] | None = None, piece_bytes: int = UPLOAD_PIECE) -> CompiledGraph

Load the image and upload its weights now rather than on first call.

Source code in iron/common/graph/compiled.py

def load(
    self,
    release: Callable[[np.ndarray], None] | None = None,
    piece_bytes: int = UPLOAD_PIECE,
) -> CompiledGraph:
    """Load the image and upload its weights now rather than on first call."""
    self.callable
    self.upload(release, piece_bytes)
    return self

start

start(run: FullELFRun, /, *tensors, **values) -> None

Start a call of this full ELF on run without waiting for it.

Raises:

Type	Description
ValueError	A host step runs between its commands.

Source code in iron/common/graph/compiled.py

def start(self, run: FullELFRun, /, *tensors, **values) -> None:
    """Start a call of this full ELF on ``run`` without waiting for it.

    Raises:
        ValueError: A host step runs between its commands.
    """
    if self.recording.host:
        raise ValueError(
            f"{self.recording.name}: a host step runs between its commands; "
            f"call it, which waits on each"
        )
    self._stage(tensors, values, run)
    call = self.callable
    assert isinstance(call, FullELFCallable), "a run is a full ELF's"
    call.start(run, idle=self.plan.idle(self.recording.steps, values))

emit_to

emit_to(target: CompiledGraph) -> None

Program this version's Emit to start a call of target.

Source code in iron/common/graph/compiled.py

def emit_to(self, target: CompiledGraph) -> None:
    """Program this version's Emit to start a call of ``target``."""
    if self.emit is None:
        raise ValueError(f"{self.recording.name}: this version has no Emit step")
    if target.recording.addresses:
        tables = sorted(
            {buffer for buffer, _, _ in target.recording.addresses.values()}
        )
        raise ValueError(
            f"{target.recording.name}: an Emit cannot start it: its gather by ids "
            f"the device made reads the address of {tables}, which the host "
            f"writes each call and no carried value holds"
        )
    program = compose(
        self.emit, self.recording.carry, target.words, target.parameters
    )
    self.write(self.emit.program, program)

Graph

A graph: a subclass whose body is recorded on handles.

body's positional parameters are the inputs (one defaulting to None may be left out), its keyword-only ones the per-call values, its return the outputs; signature() may name them instead. Weights and states are what it holds, by attribute path (self.layers[3].q is layers.3.q); profile is a Profile or a directory of <device>.json ones. Called in another graph's body, it records there.

signature

signature() -> Signature

Return the inputs and per-call values body takes, self left out.

body's own parameters, unless a subclass names them here.

Source code in iron/common/graph/compiled.py

def signature(self) -> inspect.Signature:
    """Return the inputs and per-call values ``body`` takes, ``self`` left out.

    ``body``'s own parameters, unless a subclass names them here.
    """
    signature = inspect.signature(type(self).body, eval_str=True)
    return signature.replace(parameters=list(signature.parameters.values())[1:])

held

held() -> dict[int, tuple[str, Any]]

Each param, weight and state held, by id, with its path.

Raises:

Type	Description
TypeError	A param was never loaded, or a held array is no param.

Source code in iron/common/graph/compiled.py

def held(self) -> dict[int, tuple[str, Any]]:
    """Each ``param``, ``weight`` and state held, by ``id``, with its path.

    Raises:
        TypeError: A ``param`` was never loaded, or a held array is no ``param``.
    """
    held: dict[int, tuple[str, Any]] = {}
    for path, x, _, _ in walk(self):
        if isinstance(x, _Unbound):
            raise TypeError(f"{self.name}: {path} was never loaded")
        if isinstance(x, np.ndarray) and not isinstance(x, _Param):
            raise TypeError(
                f"{self.name}: {path} is an array, not a weight: "
                f"wrap it in iron.param"
            )
        if isinstance(x, (_Param, Weight, State)):
            held.setdefault(id(x), (path, x))
    return held

release

release() -> None

Free the device memory this graph's versions and arena hold.

A call after it loads them again.

Source code in iron/common/graph/compiled.py

def release(self) -> None:
    """Free the device memory this graph's versions and arena hold.

    A call after it loads them again.
    """
    for version in self._versions.values():
        version.release()
    if "_arena" in vars(self):
        self._arena.release()

load_state_dict

load_state_dict(tree, *, path: str = '') -> None

Bind each param held, at any depth, to its tensor in tree.

The tensor is found at path (this graph's prefix, "model.") plus its dotted path in tree, cast as torch's copy_; a param held at several paths binds once.

Raises:

Type	Description
ValueError	path lacks its trailing dot, the graph is already loaded, or a tensor is missing, extra, or of another shape or kind, one error naming every such path.
TypeError	A param is in a tuple or a frozen dataclass.

Source code in iron/common/graph/compiled.py

def load_state_dict(self, tree, *, path: str = "") -> None:
    """Bind each ``param`` held, at any depth, to its tensor in ``tree``.

    The tensor is found at ``path`` (this graph's prefix, ``"model."``)
    plus its dotted path in ``tree``, cast as torch's ``copy_``; a param
    held at several paths binds once.

    Raises:
        ValueError: ``path`` lacks its trailing dot, the graph is already
            loaded, or a tensor is missing, extra, or of another shape or
            kind, one error naming every such path.
        TypeError: A param is in a tuple or a frozen dataclass.
    """
    if path and not path.endswith("."):
        raise ValueError(
            f"{self.name}: load_state_dict's path {path!r} must end in a dot"
        )
    if self.__dict__.get("_bound"):
        raise ValueError(f"{self.name}: already loaded")
    tensors = {n: t for n, t, _, _ in walk(tree) if n.startswith(path)}
    tied: dict[int, list[str]] = {}
    params: dict[str, _Unbound] = {}
    shared: list[str] = []
    for n, p, parent, _ in walk(self, shared=shared):
        if not isinstance(p, _Unbound):
            continue
        if isinstance(parent, tuple) or (
            dataclasses.is_dataclass(parent)
            and cast(Any, parent).__dataclass_params__.frozen
        ):
            raise TypeError(
                f"{self.name}: {path}{n} is in a tuple or a frozen "
                f"dataclass, which load_state_dict cannot rebind; hold it in a list"
            )
        tied.setdefault(id(p), []).append(path + n)
        params.setdefault(tied[id(p)][0], p)
    held = {n for paths in tied.values() for n in paths}
    missing = sorted(params.keys() - tensors.keys())
    extra = sorted(
        n
        for n in tensors.keys() - held
        if not any(n.startswith(f"{path}{s}.") for s in shared)
    )
    given = [(n, p, tensors[n]) for n, p in params.items() if n in tensors]
    shapes = [
        f"{n} {tuple(t.shape)}, not {p.shape}"
        for n, p, t in given
        if tuple(t.shape) != p.shape
    ]
    kinds = [
        f"{n} {t.dtype}, not {p.dtype}"
        for n, p, t in given
        if _kind(t.dtype) != _kind(p.dtype)
    ]
    if missing or extra or shapes or kinds:
        found = {
            "missing": missing,
            "extra": extra,
            "wrong shape": shapes,
            "wrong kind": kinds,
        }
        raise ValueError(
            f"{self.name}: cannot load: "
            + "; ".join(f"{k} {', '.join(v)}" for k, v in found.items() if v)
        )
    for first, *others in tied.values():
        tensors[first] = param(
            np.asanyarray(tensors[first]).astype(params[first].dtype, copy=False)
        )
        tensors.update(dict.fromkeys(others, tensors[first]))
    for n, x, parent, key in walk(self, graphs=False):
        if isinstance(x, Graph):
            if not x.__dict__.get("_bound"):
                x.load_state_dict(tensors, path=f"{path}{n}.")
        elif not isinstance(x, _Unbound):
            continue
        elif isinstance(parent, (list, dict)):
            parent[key] = tensors[path + n]
        else:
            setattr(parent, key, tensors[path + n])
    self._bound = True

record

record(**shapes) -> Recording

Run body on handles of the given shapes.

An optional input given none is None.

Source code in iron/common/graph/compiled.py

def record(self, **shapes) -> Recording:
    """Run ``body`` on handles of the given shapes.

    An optional input given none is None.
    """
    shapes = {k: v for k, v in shapes.items() if v is not None}
    params = self._parameters
    missing = [
        p for p in params.inputs if p not in shapes and p not in params.optional
    ]
    unknown = [k for k in shapes if k not in params.inputs]
    if missing or unknown:
        raise TypeError(
            f"{self.name}: shapes for {missing} missing"
            + (f"; {unknown} are not inputs" if unknown else "")
        )
    inputs: list[Handle] = []
    args: list[Handle | None] = []
    for name in params.inputs:
        if name not in shapes:
            args.append(None)
            continue
        shape, dtype = _shape_and_dtype(shapes[name])
        inputs.append(Handle.whole(shape, dtype, name, "input"))
        args.append(inputs[-1])
    values = [PerCall(n, kind, dtype) for n, (kind, dtype) in params.values.items()]
    with Recorder(
        self.name, self._profile(), self.held(), params.inputs
    ) as recorder:
        result = self.body(*args, **{v.name: v for v in values})
    items, carry = self._split_carry(result)
    outputs = self._outputs(items, recorder)
    return recorder.finish(
        inputs, outputs, values, self._recorded_carry(carry, values, recorder)
    )

compile

compile(*, dev: Device | None = None, full_elf: bool | None = None, check: bool = False, feeds: CompiledGraph | None = None, **shapes: Sequence[int] | tuple[Sequence[int], DTypeLike]) -> CompiledGraph

Compile and return the version for these input shapes.

Compile all before the first call, as a later one grows (copies) the arena.

Parameters:

Name	Type	Description	Default
dev	Device | None	The device, for this compile alone; else the current one.	None
full_elf	bool | None	One full ELF, else an xclbin per step; None derives it.	None
check	bool	Run each step on its own dispatch, raising at the first outside its tolerance().	False
feeds	CompiledGraph | None	The version a full ELF's Emit starts; by default itself when it takes no tensor.	None
**shapes	Sequence[int] | tuple[Sequence[int], DTypeLike]	Each input's shape, or (shape, dtype).	{}

Source code in iron/common/graph/compiled.py

def compile(
    self,
    *,
    dev: Device | None = None,
    full_elf: bool | None = None,
    check: bool = False,
    feeds: CompiledGraph | None = None,
    **shapes: Sequence[int] | tuple[Sequence[int], DTypeLike],
) -> CompiledGraph:
    """Compile and return the version for these input shapes.

    Compile all before the first call, as a later one grows (copies) the
    arena.

    Args:
        dev: The device, for this compile alone; else the current one.
        full_elf: One full ELF, else an xclbin per step; ``None`` derives it.
        check: Run each step on its own dispatch, raising at the first
            outside its ``tolerance()``.
        feeds: The version a full ELF's Emit starts; by default itself
            when it takes no tensor.
        **shapes: Each input's shape, or ``(shape, dtype)``.
    """
    with (
        aie_utils.set_current_device(dev)
        if dev is not None
        else contextlib.nullcontext()
    ):
        recording = self.record(**shapes)
        chosen = plan(
            require_device(),
            recording,
            full_elf,
            check,
            step_words=_step_words(recording),
        )
        logger.info("%s", chosen.report(self.name))
        signature = self._signature(
            {h.name: (h.shape, h.dtype) for h in recording.inputs}
        )
        emit, sized = self._emit(recording, chosen, feeds)
        version = CompiledGraph(
            recording,
            chosen,
            check=check,
            arena=self._arena,
            emit=emit,
        )
        if sized is not None and len(version.parameters) != len(sized):
            read = {p.name for p in version.parameters}
            raise NotImplementedError(
                f"{self.name}: the image reads {len(version.parameters)} of the "
                f"{len(sized)} words its Emit was sized for; a word it reads only "
                f"through a derivation cannot be fed yet (sized for but not "
                f"read: {sorted(w.symbol for w in sized if w.symbol not in read)})"
            )
        self._versions[signature] = version
        return version

inline

inline(method: Callable, /, *args, **kwargs) -> Any

Record method, one of this graph's, into the graph being recorded.

Recorded under this graph's profile; that graph names its values.

Raises:

Type	Description
TypeError	No graph's body is being recorded.

Source code in iron/common/graph/compiled.py

def inline(self, method: Callable, /, *args, **kwargs) -> Any:
    """Record ``method``, one of this graph's, into the graph being recorded.

    Recorded under this graph's profile; that graph names its values.

    Raises:
        TypeError: No graph's body is being recorded.
    """
    recorder = graph_recorder.get()
    if recorder is None:
        raise TypeError(f"{self.name}: inline() records into a graph's body")
    own = self._profile()
    if own is None:
        return method(*args, **kwargs)
    previous, recorder.profile = recorder.profile, own
    try:
        return method(*args, **kwargs)
    finally:
        recorder.profile = previous

reference

reference(*tensors, dev=None, **values) -> Any

Run the body on the host, each step through its operator's reference().

Tiled for dev (else the current device) without needing an NPU; outputs are numpy arrays of their shape, cut to bounds.

Raises:

Type	Description
RuntimeError	No device is given or bound.
TypeError	An input or per-call value is missing or unknown.
IndexError	The values index or bound a view outside it.

Source code in iron/common/graph/compiled.py

def reference(self, *tensors, dev=None, **values) -> Any:
    """Run the body on the host, each step through its operator's ``reference()``.

    Tiled for ``dev`` (else the current device) without needing an NPU;
    outputs are numpy arrays of their shape, cut to bounds.

    Raises:
        RuntimeError: No device is given or bound.
        TypeError: An input or per-call value is missing or unknown.
        IndexError: The values index or bound a view outside it.
    """
    given = {n: np.asarray(t) for n, t in self._given(tensors).items()}
    values = {**self._parameters.defaults, **values}
    _check_values(self.name, self._parameters.values, values)
    shapes: dict[str, Any] = {
        n: (t.shape, np.dtype(t.dtype).type) for n, t in given.items()
    }
    with (
        aie_utils.set_current_device(dev)
        if dev is not None
        else contextlib.nullcontext()
    ):
        require_device()
        signature = self._signature(shapes)
        recording = self._references.get(signature)
        if recording is None:
            token = device_fit.set(False)
            try:
                recording = self._references[signature] = self.record(**shapes)
            finally:
                device_fit.reset(token)
        outputs, carry = recording.reference(given, values)
    return _results(outputs, Carry(**carry) if recording.carry else None)

carry

carry(**next_values: Handle | Affine | PerCall | int) -> Carry

Return with a graph's outputs the next value of each carried value.

return logits, iron.carry(token=sampled, position=position + 1).

Source code in iron/common/graph/handle.py

def carry(**next_values: Handle | Affine | PerCall | int) -> Carry:
    """Return with a graph's outputs the next value of each carried value.

    ``return logits, iron.carry(token=sampled, position=position + 1)``.
    """
    return Carry(**next_values)

clamp

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

Return a per-call value's expression x held to [lo, hi], a bound of None none.

The nearer bound itself where x never reaches the range.

Source code in iron/common/graph/handle.py

def clamp(
    x: PerCall | Affine, lo: int | None = None, hi: int | None = None
) -> Affine | int:
    """Return a per-call value's expression ``x`` held to ``[lo, hi]``, a bound of None none.

    The nearer bound itself where ``x`` never reaches the range.
    """
    return (x if isinstance(x, Affine) else x.affine()).clamp(lo, hi)

on_host

on_host(fn: Callable[..., None], *reads, writes=()) -> None

Run fn(*reads, *writes) on the host between the steps around it.

A full ELF ends a command there, so the graph has no device loop. Each is contiguous: reads cut to their bounds, writes parts of states no step writes.

Raises:

Type	Description
TypeError	No graph's body is being recorded.

Source code in iron/common/graph/handle.py

def on_host(fn: Callable[..., None], *reads, writes=()) -> None:
    """Run ``fn(*reads, *writes)`` on the host between the steps around it.

    A full ELF ends a command there, so the graph has no device loop. Each is
    contiguous: reads cut to their bounds, writes parts of states no step writes.

    Raises:
        TypeError: No graph's body is being recorded.
    """
    recorder = graph_recorder.get()
    if recorder is None:
        raise TypeError("on_host records into a graph's body")
    recorder.on_host(fn, reads, writes)

param

param(*shape, dtype=bfloat16) -> ndarray

Declare a weight a graph holds, as torch's nn.Parameter.

Declared by shape for Graph.load_state_dict to bind, or param(array) holding array. A graph refuses to hold a bare array.

Source code in iron/common/graph/compiled.py

def param(*shape, dtype=bfloat16) -> np.ndarray:
    """Declare a weight a graph holds, as torch's ``nn.Parameter``.

    Declared by shape for ``Graph.load_state_dict`` to bind, or
    ``param(array)`` holding ``array``. A graph refuses to hold a bare array.
    """
    if len(shape) == 1 and isinstance(shape[0], np.ndarray):
        a = shape[0]
        return a if isinstance(a, _Param) else a.view(_Param)
    return np.broadcast_to(np.zeros((), dtype), shape).view(_Unbound)

state

state(shape, dtype=bfloat16, name=None) -> State

Declare device-resident state a graph holds.

Source code in iron/common/graph/handle.py

def state(shape, dtype=bfloat16, name=None) -> State:
    """Declare device-resident state a graph holds."""
    return State(shape, dtype, name)

weight

weight(array) -> Weight

Return a tensor a graph's body indexes per call (iron.weight(table)[ids]).

Source code in iron/common/graph/handle.py

def weight(array) -> Weight:
    """Return a tensor a graph's body indexes per call (``iron.weight(table)[ids]``)."""
    return Weight(array)

Handles

What a graph's body() reads and writes: each operator step returns one.

Handle dataclass

Handle(buffer: Buffer, shape: tuple[int, ...], dtype: Any, valid: Mapping[int, Affine] = dict(), layout: Dense | Strided | Gather = Dense())

A recorded tensor: its place in a graph buffer, shape, dtype, and per-axis bound.

layout is where it is in the buffer; valid holds, per axis, the leading entries valid this call (x[:n]).

h[key] takes numpy's basic indexing, a per-call value on one axis (keys[:, pos]), and leading-axis rows, fixed (table[ids]) or a 1-D integer handle's per call; any view but a contiguous static slice is a copy's.

T property

T: Handle

transpose(1, 0), as numpy's .T of a matrix.

whole classmethod

whole(shape, dtype, name: str, role: str) -> Handle

Make a new buffer name of shape, whole.

Source code in iron/common/graph/handle.py

@classmethod
def whole(cls, shape, dtype, name: str, role: str) -> Handle:
    """Make a new buffer ``name`` of ``shape``, whole."""
    shape = tuple(int(s) for s in shape)
    return cls(Buffer(name, role, shape, dtype), shape, dtype)

ref

ref() -> BufferRef

Return where a step finds this handle.

A part of its buffer for a contiguous run, else all of it.

Source code in iron/common/graph/handle.py

def ref(self) -> BufferRef:
    """Return where a step finds this handle.

    A part of its buffer for a contiguous run, else all of it.
    """
    match self.layout:
        case Dense(offset=int(offset)):
            return BufferRef(
                self.name, offset * bfp.itemsize(self.dtype), self.nbytes
            )
    return BufferRef(self.name)

reshape

reshape(*shape) -> Handle

Return the same buffer seen with another shape (no data moves).

Source code in iron/common/graph/handle.py

def reshape(self, *shape) -> Handle:
    """Return the same buffer seen with another shape (no data moves)."""
    if len(shape) == 1 and isinstance(shape[0], (tuple, list)):
        shape = tuple(shape[0])
    if prod(shape) != self.elements:
        raise ValueError(f"cannot reshape {self!r} to {list(shape)}")
    if not isinstance(self.layout, Dense):
        raise ValueError(f"cannot reshape a view {self!r}; reshape what it views")
    return dataclasses.replace(
        self, shape=shape, valid=_rescale_bounds(self, shape)
    )

transpose

transpose(*axes) -> Handle

Return the same buffer walked with its axes permuted (no data moves).

Source code in iron/common/graph/handle.py

def transpose(self, *axes) -> Handle:
    """Return the same buffer walked with its axes permuted (no data moves)."""
    if len(axes) == 1 and isinstance(axes[0], (tuple, list)):
        axes = tuple(axes[0])
    if not isinstance(self.layout, Dense):
        raise ValueError(
            f"cannot transpose a view {self!r}; transpose what it views"
        )
    if sorted(axes) != list(range(len(self.shape))):
        raise ValueError(
            f"axes {axes} do not permute a shape of rank {len(self.shape)}"
        )
    shape = tuple(self.shape[a] for a in axes)
    tap = TensorAccessPattern.full(self.shape).permute(axes)
    valid = {axes.index(axis): b for axis, b in self.valid.items()}
    return Handle(
        self.buffer, shape, self.dtype, valid, Strided(self._over_buffer(tap))
    )

at

at(tap: TensorAccessPattern | tuple[TensorAccessPattern, ...], shape, *, bounds=None, repeat=None, length=None)

Return what tap walks, as shape: a slice if contiguous, else a view.

Per call, repeat runs each pattern's first dimension (a tuple, one per pattern of a tuple tap) and length takes entries of its second; bounds then gives the rows the view moves a call.

Raises:

Type	Description
ValueError	With repeat or length: a pattern is not one shim descriptor, repeats don't match patterns, or no bounds.

Source code in iron/common/graph/handle.py

def at(
    self,
    tap: TensorAccessPattern | tuple[TensorAccessPattern, ...],
    shape,
    *,
    bounds=None,
    repeat=None,
    length=None,
):
    """Return what ``tap`` walks, as ``shape``: a slice if contiguous, else a view.

    Per call, ``repeat`` runs each pattern's first dimension (a tuple, one
    per pattern of a tuple ``tap``) and ``length`` takes entries of its
    second; ``bounds`` then gives the rows the view moves a call.

    Raises:
        ValueError: With ``repeat`` or ``length``: a pattern is not one shim
            descriptor, repeats don't match patterns, or no ``bounds``.
    """
    if not isinstance(self.layout, Dense):
        raise ValueError(f"cannot walk a view {self!r}; walk what it views")
    walks = tap if isinstance(tap, tuple) else (tap,)
    each: tuple = repeat if isinstance(repeat, tuple) else (repeat,)
    if isinstance(repeat, tuple) != isinstance(tap, tuple) or len(each) != len(
        walks
    ):
        raise ValueError(
            f"{len(walks)} pattern(s) take a repeat each, not {repeat}"
        )
    if repeat is None and length is None:
        return self._view(tap, shape, None, bounds or {})
    dev = aie_utils.get_current_device()
    for walk in walks:
        if (
            walk.rank != 4
            or dev is None
            or not dev.dma_bd_fits(0, 0, walk, self.dtype)
        ):
            raise ValueError(
                f"{walk} takes a per-call repeat or length as one shim "
                f"descriptor: four dimensions, the repeated first, within "
                f"a descriptor's limits"
            )
    if not bounds:
        raise ValueError(
            f"a walk of {tap} by a per-call repeat or length bounds the "
            f"rows it moves a call: give bounds"
        )
    repeats = None
    if repeat is not None:
        # A repeat count may be clamped; a length patches a size, linear.
        repeats = tuple(r if isinstance(r, Affine) else r.affine() for r in each)
    length = None if length is None else length.affine()
    recorder = graph_recorder.get()
    if recorder is not None:
        for walk, r in zip(walks, repeats or ()):
            (size0,) = concrete(walk.sizes[:1], walk)
            recorder.limit(r, 1, size0 + 1, self)
        if length is not None:
            sizes1 = concrete([walk.sizes[1] for walk in walks], walks)
            recorder.limit(length, 1, min(sizes1) + 1, self)
    layout = Strided(self._over(tap), repeat=repeats, length=length)
    return Handle(self.buffer, shape, self.dtype, bounds, layout)