Tuning

Timing each operator's settings on the NPU and writing the fastest as a profile. Nothing in the graph core imports it.

measure

Timing designs on the NPU.

A design's step is the excess of repeats runs of it over one run, per extra step, so the dispatch cancels (measure). Runs are interleaved round by round, timed only while no other process holds the NPU, and a figure whose calls waited for a CPU (CPU_WAIT_US) is refused. Every figure is kept in a Cache keyed on what it follows, so a design is timed once per NPU and power mode.

Timing dataclass

Timing(rounds: int = 8, calls: int = 50, settle: int = 2, cutoff: float = 1.5)

How long to measure: rounds interleaved rounds of calls runs.

Where the runs are settings of one design, a setting whose runs take more than cutoff times the fastest setting's once settle rounds are in is timed no further, its figure those rounds'.

Standalone

Standalone(name: str, runlist: Sequence[Operator], values: Mapping[Operator, Mapping[str, int]] | None = None, seed: int = 0, distinct: bool = True, inputs: Mapping[Operator, Mapping[str, ndarray]] | None = None, full_elf: bool = True, load: bool = True)

A runlist of operators built alone into a loaded image.

Its inputs seeded random data.

Parameters:

Name	Type	Description	Default
full_elf	bool	How the runlist is packaged: one full ELF, else an xclbin per design dispatched step by step.	True
distinct	bool	Every step runs on buffers of its own, as a graph's do, so no step finds its inputs in a SoC cache.	True
values	Mapping[Operator, Mapping[str, int]] | None	Per operator, each per-call value by name: the tuner cannot know what a value means. One derived from bound extents follows from theirs.	None
inputs	Mapping[Operator, Mapping[str, ndarray]] | None	Per operator, input buffers by name, where random bytes would not be representative (a draw row's temperature and top-k).	None
load	bool	Write the per-call values now; with False the image is built and seeded, and load() writes them.	True

Source code in iron/tune/measure.py

def __init__(
    self,
    name: str,
    runlist: Sequence[Operator],
    values: Mapping[Operator, Mapping[str, int]] | None = None,
    seed: int = 0,
    distinct: bool = True,
    inputs: Mapping[Operator, Mapping[str, np.ndarray]] | None = None,
    full_elf: bool = True,
    load: bool = True,
):
    self.steps = list(runlist)
    values, inputs = values or {}, inputs or {}
    self._values = {op: dict(values.get(op, {})) for op in self.steps}
    firsts = {}
    for k, op in enumerate(self.steps):
        firsts.setdefault(id(op), k)
    self._slot = [
        k if distinct else firsts[id(op)] for k, op in enumerate(self.steps)
    ]
    self._filled = sorted(set(self._slot))
    in_names, out_names = [], []
    for k in self._filled:
        op = self.steps[k]
        for buf, name_ in zip(op.buffers, self._names(k, op)):
            (out_names if buf.role is Out else in_names).append(name_)
    steps = [
        SequenceStep(op, tuple(map(BufferRef, self._names(self._slot[k], op))))
        for k, op in enumerate(self.steps)
    ]
    # Configure the array every call, as a design beside others in a
    # model's image is, so one timed alone pays that too.
    packaging = FullELF(reconfigure=True) if full_elf else PerStep()
    self.sequence = Image(
        Recording.of_steps(name, steps, in_names, out_names), packaging
    ).compile()
    self.callable = self.sequence.get_callable()
    rng = np.random.default_rng(seed)
    self._before: dict[str, np.ndarray] = {}
    for k in self._filled:
        op = self.steps[k]
        names = self._names(k, op)
        addressed = op.addressed_inputs(
            {buf.name: self.address(n) for buf, n in zip(op.buffers, names)}, rng
        )
        for buf, name_ in zip(op.buffers, names):
            if buf.name in addressed:
                content = np.ascontiguousarray(addressed[buf.name], buf.dtype)
                self._bytes(name_)[: buf.nbytes] = content.view(np.uint8)
            elif buf.fills:
                content = self._content(buf, inputs.get(op, {}), rng)
                self._bytes(name_)[: buf.nbytes] = content
                if buf.drains:
                    self._before[name_] = content.copy()
    ops = {id(op): op for op in self.steps}.values()
    artifacts = self.sequence.artifacts
    # An extent read only through its derivations has no word in a full
    # ELF; an xclbin chain has no parameter table, its values dispatch-time.
    self._symbols = {
        v.symbol: np.int32(self._value(op, self._values[op], v))
        for op in ops
        for v in op.values
        if not artifacts.full_elf or v.symbol in artifacts.parameters
    }
    if load:
        self.load()

load

load() -> None

Write the per-call values, which loads a full ELF that holds any.

Any other image is loaded by its first run.

Source code in iron/tune/measure.py

def load(self) -> None:
    """Write the per-call values, which loads a full ELF that holds any.

    Any other image is loaded by its first run.
    """
    if self._symbols:
        self.callable.write_values(self._symbols)

address

address(name: str) -> int

Return the device address of buffer name.

Source code in iron/tune/measure.py

def address(self, name: str) -> int:
    """Return the device address of buffer ``name``."""
    view = self.callable.get_storage(name)
    root = view.storage.binding_handle(0, view.storage.nbytes)
    if root is None:
        raise ValueError(f"{name} has no device binding to address")
    return root.address() + view.storage_offset

digest

digest() -> str

Run once: the sha256 of what every step written.

Source code in iron/tune/measure.py

def digest(self) -> str:
    """Run once: the sha256 of what every step ``written``."""
    self.callable()
    return hashlib.sha256(
        b"".join(
            a.tobytes() for k in self._filled for a in self.written(k).values()
        )
    ).hexdigest()

inputs

inputs(k: int = 0) -> dict[str, ndarray]

Return what step k was given, by buffer name.

Views of the device buffers, an in-out one as it was before any run.

Source code in iron/tune/measure.py

def inputs(self, k: int = 0) -> dict[str, np.ndarray]:
    """Return what step ``k`` was given, by buffer name.

    Views of the device buffers, an in-out one as it was before any run.
    """
    op = self.steps[k]
    out = {}
    for buf, name in zip(op.buffers, self._names(self._slot[k], op)):
        if buf.fills:
            data = self._before.get(name, self._bytes(name)[: buf.nbytes])
            out[buf.name] = data.view(buf.host_dtype).reshape(buf.host_shape)
    return out

written

written(k: int = 0) -> dict[str, ndarray]

Return what step k wrote in the last run, by buffer name.

Along a bounded axis only the rows under the call's bound.

Source code in iron/tune/measure.py

def written(self, k: int = 0) -> dict[str, np.ndarray]:
    """Return what step ``k`` wrote in the last run, by buffer name.

    Along a bounded axis only the rows under the call's bound.
    """
    op = self.steps[k]
    out = {}
    for buf, name in zip(op.buffers, self._names(self._slot[k], op)):
        if not buf.drains:
            continue
        under = [slice(None)] * len(buf.shape)
        for extent in op.bound_extents:
            axis = buf.extent_axis(op.value(extent))
            if axis is not None:
                under[axis] = slice(0, self._values[op][extent])
        data = self.callable.get_storage(name).numpy().view(np.uint8)
        data = data[: buf.nbytes].reshape(*buf.shape, -1)
        # A block-float buffer stays its blocks' bytes, as its host shape is.
        data = np.ascontiguousarray(data[tuple(under)]).view(buf.host_dtype)
        out[buf.name] = data if bfp.is_bfp(buf.dtype) else data[..., 0]
    return out

Measurement dataclass

Measurement(t_step_us: float, round_us: tuple[float, ...], output: str, run: str, pmode: str, calls: int, default_round_us: tuple[float, ...] = (), measured: str = (lambda: isoformat())())

One design's step, as measure timed it.

Attributes:

Name	Type	Description
t_step_us	float	Its time per step, the median of round_us.
round_us	tuple[float, ...]	Its time per step as each round measured it. Two measurements of one run were timed round for round together, so their rounds pair.
output	str	The sha256 of what one run wrote, so two settings run on the same inputs compare without a rerun.
run	str	Which time_interleaved call timed it.
pmode	str	The NPU's power mode.
calls	int	The runs a round timed.
default_round_us	tuple[float, ...]	The rounds of the default timed beside it in its run; empty for the default itself.
measured	str	The day, as an ISO date.

against_default

against_default() -> tuple[float, float | None]

difference from the default timed beside it.

Source code in iron/tune/measure.py

def against_default(self) -> tuple[float, float | None]:
    """``difference`` from the default timed beside it."""
    default = self.default_round_us
    return self.difference(
        dataclasses.replace(
            self, t_step_us=statistics.median(default), round_us=default
        )
    )

difference

difference(other: Measurement) -> tuple[float, float | None]

How much longer a step of this one takes than one of other, and that difference's standard error.

Round by round where the two were timed in one run, else from each one's own.

Source code in iron/tune/measure.py

def difference(self, other: Measurement) -> tuple[float, float | None]:
    """How much longer a step of this one takes than one of ``other``, and that difference's standard error.

    Round by round where the two were timed in one run, else from each
    one's own.
    """
    delta = self.t_step_us - other.t_step_us
    if self.run == other.run:
        return delta, standard_error(
            [a - b for a, b in zip(self.round_us, other.round_us)]
        )
    own_noise, other_noise = (standard_error(m.round_us) for m in (self, other))
    if own_noise is None or other_noise is None:
        return delta, None
    return delta, math.hypot(own_noise, other_noise)

Cache

Cache(platform: str, mode: str, root: Path = Path(NPU_CACHE_HOME) / 'iron' / 'costs')

Every design measured on one NPU platform at one power mode.

One JSON file each under root/<platform>/<power mode>/.

An entry is keyed on what its time follows: the design's build (its recipe and the sources, tools and device it is compiled with, as mlir-aie's compile cache keys it), the per-call values it was run at, the contents of the inputs it was given and its packaging. Editing how a design is generated therefore misses rather than reusing a stale time.

Source code in iron/tune/measure.py

def __init__(
    self,
    platform: str,
    mode: str,
    root: Path = Path(NPU_CACHE_HOME) / "iron" / "costs",
):
    self.mode = mode
    slug = re.sub(r"[^a-z0-9]+", "-", platform.lower()).strip("-")
    self.directory = Path(root) / slug / mode

here classmethod

here() -> Cache

Return the cache of this NPU at its present power mode.

Source code in iron/tune/measure.py

@classmethod
def here(cls) -> Cache:
    """Return the cache of this NPU at its present power mode."""
    report = platform()
    return cls(report["Name"], report["Power Mode"])

key staticmethod

key(op: Operator, values: Mapping[str, int] | None = None, inputs: Mapping[str, ndarray] | None = None, full_elf: bool = True) -> str

Return the entry op, resolved for the current device, is kept in.

As run at values on inputs (random where not given), packaged as one full ELF or not.

Source code in iron/tune/measure.py

@staticmethod
def key(
    op: Operator,
    values: Mapping[str, int] | None = None,
    inputs: Mapping[str, np.ndarray] | None = None,
    full_elf: bool = True,
) -> str:
    """Return the entry ``op``, resolved for the current device, is kept in.

    As run at ``values`` on ``inputs`` (random where not given), packaged
    as one full ELF or not.
    """
    design = OperatorDesign(op.resolved()).compilable()
    given = sorted(
        (name, hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest())
        for name, a in (inputs or {}).items()
    )
    key = (
        design.recipe_hash,
        design.artifact_hash,
        sorted((values or {}).items()),
        given,
        full_elf,
    )
    return hashlib.sha256(repr(key).encode()).hexdigest()[:32]

get

get(key: str) -> Measurement | None

Return the measurement at key; None if there is none.

Also None if it was written with other fields than Measurement has.

Source code in iron/tune/measure.py

def get(self, key: str) -> Measurement | None:
    """Return the measurement at ``key``; None if there is none.

    Also None if it was written with other fields than ``Measurement`` has.
    """
    path = self.directory / f"{key}.json"
    if not path.exists():
        return None
    fields = json.loads(path.read_text())
    if fields.keys() != {f.name for f in dataclasses.fields(Measurement)}:
        return None
    return Measurement(
        **{
            **fields,
            "round_us": tuple(fields["round_us"]),
            "default_round_us": tuple(fields["default_round_us"]),
        }
    )

standard_error

standard_error(samples: Sequence[float]) -> float | None

Stats' standard error of the median of samples; None from fewer than two.

Source code in iron/tune/measure.py

def standard_error(samples: Sequence[float]) -> float | None:
    """``Stats``' standard error of the median of ``samples``; None from fewer than two."""
    if len(samples) < 2:
        return None
    return Stats.from_samples(list(samples)).median_se_us

platform

platform() -> dict[str, str]

Return the NPU's platform report from xrt-smi.

Its Name, its Power Mode, ...

Source code in iron/tune/measure.py

def platform() -> dict[str, str]:
    """Return the NPU's platform report from ``xrt-smi``.

    Its ``Name``, its ``Power Mode``, ...
    """
    out = subprocess.run(
        ["xrt-smi", "examine", "-r", "platform"],
        capture_output=True,
        text=True,
        check=True,
    ).stdout
    report = {}
    for line in out.splitlines():
        name, sep, value = line.partition(":")
        if sep:
            report[name.strip()] = value.strip()
    if "Name" not in report or "Power Mode" not in report:
        raise RuntimeError(f"xrt-smi reports no platform name or power mode:\n{out}")
    return report

pmode

pmode() -> str

Return the NPU's power mode, as xrt-smi reports it.

Source code in iron/tune/measure.py

def pmode() -> str:
    """Return the NPU's power mode, as ``xrt-smi`` reports it."""
    return platform()["Power Mode"]

judge

judge(default: Operator, op: Operator, inputs: Mapping[str, ndarray], written: Mapping[str, ndarray], values: Mapping[str, int] | None = None) -> str | None

Why op's output written on inputs may not replace that of default.

The setting it would replace: it must be within the default's gate and its own. None where it may.

Only a bit-identical setting is admitted where a gate is missing or exact, or values bound an extent short of its length.

Source code in iron/tune/measure.py

def judge(
    default: Operator,
    op: Operator,
    inputs: Mapping[str, np.ndarray],
    written: Mapping[str, np.ndarray],
    values: Mapping[str, int] | None = None,
) -> str | None:
    """Why ``op``'s output ``written`` on ``inputs`` may not replace that of ``default``.

    The setting it would replace: it must be within the default's gate and
    its own. None where it may.

    Only a bit-identical setting is admitted where a gate is missing or
    exact, or ``values`` bound an extent short of its length.
    """
    gates = {
        "the default's": default.resolved().tolerance(),
        "its own": op.resolved().tolerance(),
    }
    if any(g is None or g.kind == "exact" for g in gates.values()):
        return "only a bit-identical output is admitted"
    values = values or {}
    for extent in op.bound_extents:
        bounds = op.value(extent).bounds
        assert bounds is not None, "bound_extents holds only extents that bound one"
        if values.get(extent) != vars(op)[bounds.name]:
            return "only a bit-identical output is admitted under a bound"
    want = op._call_reference(inputs, values=values)
    for which, gate in gates.items():
        assert gate is not None, "a None gate already returned above"
        for name, verdict in op.judge(inputs, written, gate, want).items():
            if not verdict:
                return f"{name} under {which} gate: {verdict.detail}"
    return None

time_interleaved

time_interleaved(runs: Sequence[ImageCallable], timing: Timing, groups: Sequence[Hashable | None] = (), fastest: float = inf, log: Callable[[str], None] = print) -> list[Stats]

Each loaded image's per-round medians, microseconds (samples_us).

Each call's figure the time the image itself took, which its KernelResult carries: mlir-aie's run_interleaved, whose arguments and errors these are.

Source code in iron/tune/measure.py

def time_interleaved(
    runs: Sequence[ImageCallable],
    timing: Timing,
    groups: Sequence[Hashable | None] = (),
    fastest: float = math.inf,
    log: Callable[[str], None] = print,
) -> list[Stats]:
    """Each loaded image's per-round medians, microseconds (``samples_us``).

    Each call's figure the time the image itself took, which its
    ``KernelResult`` carries: mlir-aie's ``run_interleaved``, whose arguments
    and errors these are.
    """
    return run_interleaved(
        list(runs),
        rounds=timing.rounds,
        calls=timing.calls,
        settle=timing.settle,
        cutoff=timing.cutoff,
        groups=list(groups) or None,
        fastest_us=fastest,
        completion_cores=True,
        idle=True,
        log=log,
    )

write

write(path: Path, text: str) -> None

Write text to path whole, as another process may read it.

Source code in iron/tune/measure.py

def write(path: Path, text: str) -> None:
    """Write ``text`` to ``path`` whole, as another process may read it."""
    path.parent.mkdir(parents=True, exist_ok=True)
    partial = path.with_suffix(f".{os.getpid()}")
    partial.write_text(text)
    partial.replace(path)

measure

measure(ops: Sequence[Operator], *, full_elf: bool, timing: Timing = Timing(), repeats: int = 9, values: Mapping[str, int] | None = None, inputs: Mapping[str, ndarray] | None = None, cache: Cache | None = None, remeasure: bool = False, log: Callable[[str], None] = print) -> dict[int, Measurement]

Time each of ops, settings of one design with the default first.

As many at once as the device's contexts hold. Every batch times the default too, so each other setting pairs with it round by round.

Parameters:

Name	Type	Description	Default
ops	Sequence[Operator]	The settings to time, the default (index 0) first.	required
full_elf	bool	How each is packaged: one full ELF, else an xclbin per design dispatched step by step.	required
timing	Timing	How long to measure each.	Timing()
repeats	int	How many back-to-back runs a step's time is the excess over.	9
values	Mapping[str, int] | None	The per-call values each is run at.	None
inputs	Mapping[str, ndarray] | None	What its input buffers hold, by name, where random bytes would not be representative.	None
cache	Cache | None	Settings it holds are taken from it rather than run, unless remeasure; those run are written to it.	None
remeasure	bool	Run every setting even where cache already holds it.	False
log	Callable[[str], None]	Where progress and a setting that fails to build are reported.	print

Returns:

Type	Description
dict[int, Measurement]	Each one timed, by its index in ops. One other than the default
dict[int, Measurement]	that does not build is left out, and reported to log.

Raises:

Type	Description
RuntimeError	The default does not build.

Source code in iron/tune/measure.py

def measure(
    ops: Sequence[Operator],
    *,
    full_elf: bool,
    timing: Timing = Timing(),
    repeats: int = 9,
    values: Mapping[str, int] | None = None,
    inputs: Mapping[str, np.ndarray] | None = None,
    cache: Cache | None = None,
    remeasure: bool = False,
    log: Callable[[str], None] = print,
) -> dict[int, Measurement]:
    """Time each of ``ops``, settings of one design with the default first.

    As many at once as the device's contexts hold. Every batch times the
    default too, so each other setting pairs with it round by round.

    Args:
        ops: The settings to time, the default (index 0) first.
        full_elf: How each is packaged: one full ELF, else an xclbin per
            design dispatched step by step.
        timing: How long to measure each.
        repeats: How many back-to-back runs a step's time is the excess over.
        values: The per-call values each is run at.
        inputs: What its input buffers hold, by name, where random bytes
            would not be representative.
        cache: Settings it holds are taken from it rather than run, unless
            `remeasure`; those run are written to it.
        remeasure: Run every setting even where ``cache`` already holds it.
        log: Where progress and a setting that fails to build are reported.

    Returns:
        Each one timed, by its index in ``ops``. One other than the default
        that does not build is left out, and reported to ``log``.

    Raises:
        RuntimeError: The default does not build.
    """
    mode = pmode()
    if cache is not None and cache.mode != mode:
        raise ValueError(f"the cost cache is for power mode {cache.mode}, not {mode}")
    keys = (
        None
        if cache is None
        else [cache.key(op, values, inputs, full_elf) for op in ops]
    )
    out: dict[int, Measurement] = {}
    if cache is not None and not remeasure:
        assert keys is not None, "built above whenever cache is not None"
        for i, key in enumerate(keys):
            if (m := cache.get(key)) is not None:
                out[i] = m
    todo = [i for i in range(len(ops)) if i not in out]
    if not todo:
        return out
    if 0 not in todo:
        todo.insert(0, 0)  # the default, timed beside the others
    distinct = sum(b.nbytes for b in ops[0].resolved().buffers) <= DISTINCT_BYTES
    # Each setting is two runs, each holding a context; one pair is the default's.
    size = CONTEXTS // 2 - 1
    fastest = math.inf
    others = todo[1:]
    for begin in range(0, max(len(others), 1), size):
        batch = [todo[0]] + others[begin : begin + size]
        log(f"    timing {len(batch)} settings, {len(others) - begin} to go")
        built, short, long = [], [], []
        for i in batch:
            try:
                # Operator is hashable by identity at runtime (object.__hash__);
                # pyright infers __hash__ = None from jit.Operator's
                # dataclass_transform defaults.
                one = Standalone(
                    f"probe1_{i}",
                    [ops[i]],
                    values={ops[i]: values or {}},
                    inputs={ops[i]: inputs or {}},
                    full_elf=full_elf,
                    load=False,
                )
                many = Standalone(
                    f"probe{repeats}_{i}",
                    [ops[i]] * repeats,
                    values={ops[i]: values or {}},
                    distinct=distinct,
                    inputs={ops[i]: inputs or {}},
                    full_elf=full_elf,
                    load=False,
                )
            except RuntimeError as e:
                # The placer passed it, but the build did not.
                if i == 0:
                    raise
                log(f"    {ops[i]} does not build, not measured: {e}")
                continue
            # Not caught: a failed load is the device's state, not the design's.
            one.load()
            many.load()
            built.append(i)
            short.append(one)
            long.append(many)
        outputs = [run.digest() for run in short]
        # The default is timed in full; the others stop once far behind.
        groups = [None if i == 0 else i for i in built] * 2
        times = time_interleaved(
            [r.callable for r in short + long], timing, groups, fastest, log
        )
        run = uuid.uuid4().hex
        n = len(built)
        steps = [
            tuple(
                (many - one) / (repeats - 1)
                for one, many in zip(times[j].samples_us, times[n + j].samples_us)
            )
            for j in range(n)
        ]
        for j, i in enumerate(built):
            out[i] = Measurement(
                statistics.median(steps[j]),
                steps[j],
                outputs[j],
                run,
                mode,
                timing.calls,
                () if i == 0 else steps[0],
            )
            if cache is not None:
                assert keys is not None, "built above whenever cache is not None"
                cache.put(keys[i], out[i])
        fastest = min(fastest, *(times[n + j].median_us for j in range(n)))
        # A probe holds its context while it lives; the next batch needs them.
        del short, long
    return out

autotune

Choosing each operator's tunables by timing them, written as a profile.

autotune takes the calls a model runs (Timed), finds each operator shape in them, times every setting of its tunables (Operator.domains) the placer takes against the setting it was recorded at, and returns a Profile naming the fastest for each shape.

Point dataclass

Point(weight: float, values: Mapping[str, int])

A share of a graph's calls, run at one set of its per-call values.

Attributes:

Name	Type	Description
weight	float	The share; a call's points' weights sum to 1.
values	Mapping[str, int]	The graph's per-call values there, by name.

Timed dataclass

Timed(recording: Recording, values: Mapping[str, int] = dict(), contents: Mapping[str, ndarray] = dict(), points: Sequence[Point] = ())

One call of a recorded graph, as its operators are timed in.

The graph's per-call values, and what any buffer whose contents a step's time follows holds, by graph buffer name. points are the values its calls run at, where they differ from call to call: an operator whose work they change is timed at each (op_points).

op_values

op_values(op: Operator, values: Mapping[str, int] | None = None) -> dict[str, int]

Return the per-call values op is written in this call.

Or at the graph's values, by member name.

Source code in iron/tune/autotune.py

def op_values(
    self, op: Operator, values: Mapping[str, int] | None = None
) -> dict[str, int]:
    """Return the per-call values ``op`` is written in this call.

    Or at the graph's ``values``, by member name.
    """
    # A table's address the graph fills in; a word step alone reads none.
    values = {
        **dict.fromkeys(self.recording.addresses, 0),
        **(self.values if values is None else values),
    }
    return {
        b.member.name: b.expression.evaluate(values)
        for b in self.recording.bindings
        if b.op is op
    }

op_points

op_points(op: Operator) -> list[tuple[float, dict[str, int]]]

Return where op is timed, as (weight, values) pairs.

Points that give op's bound extents the same values are merged: their weights are summed, and the first one's per-call values are kept. When every point gives them the same values, as when op has no bound extent, that is a single pair.

Source code in iron/tune/autotune.py

def op_points(self, op: Operator) -> list[tuple[float, dict[str, int]]]:
    """Return where ``op`` is timed, as ``(weight, values)`` pairs.

    Points that give ``op``'s bound extents the same values are merged:
    their weights are summed, and the first one's per-call values are
    kept. When every point gives them the same values, as when ``op``
    has no bound extent, that is a single pair.
    """
    out: dict[tuple, tuple[float, dict[str, int]]] = {}
    for point in self.points:
        values = self.op_values(op, point.values)
        label = tuple(values[e] for e in sorted(op.bound_extents))
        weight, first = out.get(label, (0.0, values))
        out[label] = (weight + point.weight, first)
    if len(out) > 1:
        return list(out.values())
    return [(1.0, self.op_values(op))]

op_inputs

op_inputs(op: Operator) -> dict[str, ndarray]

Return what op's buffers hold in this call, by op's buffer name.

Source code in iron/tune/autotune.py

def op_inputs(self, op: Operator) -> dict[str, np.ndarray]:
    """Return what ``op``'s buffers hold in this call, by ``op``'s buffer name."""
    return {
        buf.name: self.contents[ref.root]
        for step in self.recording.steps
        if step.op is op
        for buf, ref in zip(op.buffers, step.refs())
        if ref.nbytes is None and ref.root in self.contents
    }

Trial dataclass

Trial(tunables: dict[str, Any], delta_us: float, noise_us: float | None, refused: str | None = None)

One setting of a shape, timed against the setting it was recorded at.

Attributes:

Name	Type	Description
tunables	dict[str, Any]	The setting.
delta_us	float	How much longer its steps take, summed over the shape's operators and weighted over their points; negative is faster.
noise_us	float | None	The standard error of delta_us; None from one round.
refused	str | None	Why it may not replace the recorded setting, or None.

gains property

gains: bool

It may be taken: admitted, and faster beyond the noise.

Shape

Shape(op: Operator)

The operators of the calls one profile entry covers.

One class at one set of dimensions (Profile.dims).

Source code in iron/tune/autotune.py

def __init__(self, op: Operator):
    self.cls = type(op)
    given = {n: vars(op)[n] for n in op.declaration.params}
    self.dims = Profile.dims(self.cls, given)
    # Each operator, where it is timed (``Timed.op_points``) and what its
    # inputs hold (``Timed.op_inputs``).
    self.members: list[
        tuple[Operator, list[tuple[float, dict[str, int]]], dict[str, np.ndarray]]
    ] = []
    # Tunables any of its operators' calls give: none of them moves.
    self.pinned: frozenset[str] = frozenset()

named

named(covering: Sequence[Mapping[str, Any]]) -> dict[str, Any]

Return the dimensions its entry names.

Those off their field's default, and every one an entry covering it names, so it outranks them.

Source code in iron/tune/autotune.py

def named(self, covering: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
    """Return the dimensions its entry names.

    Those off their field's default, and every one an entry
    ``covering`` it names, so it outranks them.
    """
    defaults = Profile.dims(self.cls, {})
    names = {n for matched in covering for n in matched}
    return {
        k: v
        for k, v in self.dims.items()
        if k in names or k not in defaults or v != defaults[k]
    }

kernel_cycles

kernel_cycles(dev, log: Callable[[str], None] = print) -> dict[str, int]

Return each kernel's cycles per call as mlir-aie's nightly last timed it.

On dev's NPU, by ExternalFunction.recipe_key; empty where they cannot be read.

Source code in iron/tune/autotune.py

def kernel_cycles(dev, log: Callable[[str], None] = print) -> dict[str, int]:
    """Return each kernel's cycles per call as mlir-aie's nightly last timed it.

    On ``dev``'s NPU, by ``ExternalFunction.recipe_key``; empty where they
    cannot be read.
    """
    url = KERNEL_CHECKS.format(npu=dev.name)
    try:
        with urllib.request.urlopen(url, timeout=30) as response:
            latest = json.load(response)
    except (OSError, ValueError) as e:
        log(f"no kernel timings from {url}: {e}")
        return {}
    if latest.get("schema") != 1:
        log(f"no kernel timings from {url}: schema {latest.get('schema')}, not 1")
        return {}
    rows = latest.get("rows", {})
    return {
        key: rows[case]["cycles"]["value"]
        for case, key in latest.get("builds", {}).items()
        if "cycles" in rows.get(case, {})
    }

least_us

least_us(op: GEMM, values: Mapping[str, int], cycles, core_mhz: float) -> float | None

Return a step time op cannot beat at per-call values.

Its busiest core's matmul calls at cycles per call, at core_mhz. None where cycles does not hold its matmul.

Source code in iron/tune/autotune.py

def least_us(
    op: GEMM, values: Mapping[str, int], cycles, core_mhz: float
) -> float | None:
    """Return a step time ``op`` cannot beat at per-call ``values``.

    Its busiest core's matmul calls at ``cycles`` per call, at ``core_mhz``.
    None where ``cycles`` does not hold its matmul.
    """
    per_call = cycles.get(op.layout().matmul().recipe_key)
    if per_call is None:
        return None
    return op.matmul_calls(values) * per_call / (core_mhz * CLOCK_HEADROOM)

settings

settings(op: Operator, dev: Device, pinned: frozenset[str] = frozenset()) -> list[dict[str, Any]]

Return op's tunables as resolved for dev, then every other combination.

Every other combination of their Operator.domains it resolves at and whose derived transfers fit their descriptors, each design once.

Parameters:

Name	Type	Description	Default
op	Operator	The operator whose tunables are searched.	required
dev	Device	The device each combination is resolved for.	required
pinned	frozenset[str]	Tunables another operator of its shape pins, held at the value op resolves them to.	frozenset()

Source code in iron/tune/autotune.py

def settings(
    op: Operator, dev: Device, pinned: frozenset[str] = frozenset()
) -> list[dict[str, Any]]:
    """Return ``op``'s tunables as resolved for ``dev``, then every other combination.

    Every other combination of their ``Operator.domains`` it resolves at
    and whose derived transfers fit their descriptors, each design once.

    Args:
        op: The operator whose tunables are searched.
        dev: The device each combination is resolved for.
        pinned: Tunables another operator of its shape pins, held at the
            value ``op`` resolves them to.
    """
    if pinned - op.pinned:
        resolved = op.resolved(dev)
        op = op.replace(
            pinned=op.pinned | pinned,
            **{name: vars(resolved)[name] for name in pinned - op.pinned},
        )
    resolved = op.resolved(dev)
    domains = resolved._domains(dev)
    default = {name: vars(resolved)[name] for name in domains}
    out = {_identity(resolved, dev): default}
    dropped = 0
    for combo in itertools.product(*domains.values()):
        tunables = dict(zip(domains, combo))
        if tunables == default:
            continue
        try:
            moved = op._with_tunables(**tunables).resolved(dev)
            if not moved.has_sequence_override():
                for buf in moved.buffers:
                    buf.transfers()
        except ValueError as e:  # unresolvable, incompatible or untransferable here
            logger.debug("%s at %s dropped: %s", type(op).__name__, tunables, e)
            dropped += 1
            continue
        out.setdefault(_identity(moved, dev), tunables)
    if dropped:
        logger.info(
            "%s: %d of %d settings dropped",
            type(op).__name__,
            dropped,
            math.prod(len(d) for d in domains.values()),
        )
    return list(out.values())

refusal

refusal(op: Operator) -> str | None

Return why the placer refuses op alone on a device, or None if it fits.

Kept under FIT_CACHE, keyed on its recipe and its build environment (the device and the tools), so a placer fix is seen.

Source code in iron/tune/autotune.py

def refusal(op: Operator) -> str | None:
    """Return why the placer refuses ``op`` alone on a device, or None if it fits.

    Kept under ``FIT_CACHE``, keyed on its recipe and its build environment
    (the device and the tools), so a placer fix is seen.
    """
    design = OperatorDesign(op.resolved())
    compilable = design.compilable()
    key = compilable.recipe_hash + compilable.artifact_hash
    record = FIT_CACHE / hashlib.sha256(key.encode()).hexdigest()[:24]
    if record.exists():
        verdict = record.read_text()
    else:
        diagnostic = fits({design.identity: design})
        verdict = "fits" if diagnostic is None else f"refused: {diagnostic}"
        write(record, verdict)
    return None if verdict == "fits" else verdict.removeprefix("refused: ")

line

line(found: Sequence[dict[str, Any]], through: dict[str, Any], name: str)

Return the settings of found a coordinate descent tries for tunable name.

From through: at each of its values, the setting that differs from through in the fewest other tunables, the first in found on a tie. Where tunables are coupled, a value is reached by moving the others with it.

Source code in iron/tune/autotune.py

def line(found: Sequence[dict[str, Any]], through: dict[str, Any], name: str):
    """Return the settings of ``found`` a coordinate descent tries for tunable ``name``.

    From ``through``: at each of its values, the setting that differs from
    ``through`` in the fewest other tunables, the first in ``found`` on a
    tie. Where tunables are coupled, a value is reached by moving the
    others with it.
    """
    at: dict[Hashable, list[dict[str, Any]]] = {}
    for s in found:
        at.setdefault(s[name], []).append(s)
    return [
        min(ss, key=lambda s: sum(s[n] != through[n] for n in s)) for ss in at.values()
    ]

shapes

shapes(calls: Sequence[Timed], dev) -> list[Shape]

Return the shapes of calls' operators that have a tunable to search.

In the order the calls first run them, each operator design once.

Source code in iron/tune/autotune.py

def shapes(calls: Sequence[Timed], dev) -> list[Shape]:
    """Return the shapes of ``calls``' operators that have a tunable to search.

    In the order the calls first run them, each operator design once.
    """
    found: dict[Hashable, Shape] = {}
    seen: set[str] = set()
    for call in calls:
        for step in call.recording.steps:
            op = step.op
            shape = Shape(op)
            shape = found.setdefault(shape.key, shape)
            shape.pinned |= op.pinned
            identity = _identity(op, dev)
            if identity not in seen:
                seen.add(identity)
                shape.members.append((op, call.op_points(op), call.op_inputs(op)))
    return [s for s in found.values() if s.members[0][0].resolved(dev)._domains(dev)]

time_shape

time_shape(shape: Shape, dev, *, full_elf: bool, timing: Timing = Timing(), repeats: int = 9, cache: Cache | None = None, remeasure: bool = False, exhaustive: int = EXHAUSTIVE, cycles: Mapping[str, int] | None = None, core_mhz: float | None = None, log: Callable[[str], None] = print) -> list[Trial]

Each setting of shape timed, the recorded one first.

Every setting its operators share and the placer takes where there are at most exhaustive, else those a coordinate descent tries, each tunable's line through the fastest setting taken so far, until a pass over the tunables moves it no further. A setting whose least_us at cycles and core_mhz is no faster than the fastest taken so far is not timed.

Source code in iron/tune/autotune.py

def time_shape(
    shape: Shape,
    dev,
    *,
    full_elf: bool,
    timing: Timing = Timing(),
    repeats: int = 9,
    cache: Cache | None = None,
    remeasure: bool = False,
    exhaustive: int = EXHAUSTIVE,
    cycles: Mapping[str, int] | None = None,
    core_mhz: float | None = None,
    log: Callable[[str], None] = print,
) -> list[Trial]:
    """Each setting of ``shape`` timed, the recorded one first.

    Every setting its operators share and the placer takes where there are
    at most ``exhaustive``, else those a coordinate descent tries, each
    tunable's ``line`` through the fastest setting taken so far, until a
    pass over the tunables moves it no further. A setting whose
    ``least_us`` at ``cycles`` and ``core_mhz`` is no faster than the
    fastest taken so far is not timed.
    """
    found = settings(shape.members[0][0], dev, shape.pinned)
    for op, _, _ in shape.members[1:]:
        own = settings(op, dev, shape.pinned)
        found = [s for s in found if s in own]
    found = [found[0]] + [
        s
        for s in found[1:]
        if all(refusal(op._with_tunables(**s)) is None for op, _, _ in shape.members)
    ]
    if len(found) == 1:
        return [Trial(found[0], 0.0, 0.0)]
    tried: dict[tuple, Trial] = {}
    least = (
        {_frozen(s): _least(shape, s, cycles, core_mhz, dev) for s in found[1:]}
        if cycles and core_mhz and issubclass(shape.cls, GEMM)
        else {}
    )
    recorded_us: float | None = None

    def slower(s: dict[str, Any]) -> bool:
        bound = least.get(_frozen(s))
        if bound is None or recorded_us is None:
            return False
        best = min((t.delta_us for t in tried.values() if t.gains), default=0.0)
        return bound >= recorded_us + best

    def run(batch: list[dict[str, Any]]) -> None:
        nonlocal recorded_us
        todo = [found[0]] + [
            s
            for s in batch
            if s is not found[0] and _frozen(s) not in tried and not slower(s)
        ]
        if skipped := sum(1 for s in batch if slower(s) and _frozen(s) not in tried):
            log(f"    {skipped} slower by their kernel alone, not timed")
        trials, recorded_us = _time(
            shape, todo, full_elf, timing, repeats, cache, remeasure, log
        )
        for s, t in zip(todo, trials):
            tried[_frozen(s)] = t

    if len(found) <= exhaustive:
        if any(b is not None for b in least.values()):
            run([])  # the recorded setting's own time, which a bound is held to
        run(found)
    else:
        best, moved = found[0], True
        while moved:
            moved = False
            for name in found[0]:
                run(line(found, best, name))
                taken = [t for t in tried.values() if t.gains or t.tunables == found[0]]
                fastest = min(taken, key=lambda t: t.delta_us)
                if fastest.tunables != best:
                    best, moved = fastest.tunables, True
    return [tried[_frozen(s)] for s in found if _frozen(s) in tried]

autotune

autotune(calls: Sequence[Timed], dev: Device, profile: Profile | None = None, *, full_elf: bool, timing: Timing = Timing(), repeats: int = 9, cache: Cache | None = None, remeasure: bool = False, cycles: Mapping[str, int] | None = None, core_mhz: float | None = None, log: Callable[[str], None] = print) -> Profile

Return profile with an entry for each shape of calls (shapes).

Each entry names the fastest setting taken (time_shape), the recorded one where none gains. profile's entries covering any of those shapes are replaced; the rest are kept.

Parameters:

Name	Type	Description	Default
calls	Sequence[Timed]	The calls whose graphs were recorded under profile.	required
dev	Device	The device the calls' operators are timed on.	required
profile	Profile | None	The profile whose unreplaced entries are kept.	None
full_elf	bool	How the calls' graphs are packaged.	required
timing	Timing	How each setting's steps are timed.	Timing()
repeats	int	How many times each setting's steps are timed.	9
cache	Cache | None	Where each time is kept across runs (Cache.here()).	None
remeasure	bool	Time again what cache holds.	False
cycles	Mapping[str, int] | None	Kernel cycles per call by recipe key (kernel_cycles).	None
core_mhz	float | None	The cores' clock the runtime reads (core_mhz), which cycles are timed at; None times every setting.	None
log	Callable[[str], None]	Where progress and the gains found are logged.	print

Source code in iron/tune/autotune.py

def autotune(
    calls: Sequence[Timed],
    dev: Device,
    profile: Profile | None = None,
    *,
    full_elf: bool,
    timing: Timing = Timing(),
    repeats: int = 9,
    cache: Cache | None = None,
    remeasure: bool = False,
    cycles: Mapping[str, int] | None = None,
    core_mhz: float | None = None,
    log: Callable[[str], None] = print,
) -> Profile:
    """Return ``profile`` with an entry for each shape of ``calls`` (``shapes``).

    Each entry names the fastest setting taken (``time_shape``), the
    recorded one where none gains. ``profile``'s entries covering any of
    those shapes are replaced; the rest are kept.

    Args:
        calls: The calls whose graphs were recorded under ``profile``.
        dev: The device the calls' operators are timed on.
        profile: The profile whose unreplaced entries are kept.
        full_elf: How the calls' graphs are packaged.
        timing: How each setting's steps are timed.
        repeats: How many times each setting's steps are timed.
        cache: Where each time is kept across runs (``Cache.here()``).
        remeasure: Time again what ``cache`` holds.
        cycles: Kernel cycles per call by recipe key (``kernel_cycles``).
        core_mhz: The cores' clock the runtime reads (``core_mhz``), which
            ``cycles`` are timed at; None times every setting.
        log: Where progress and the gains found are logged.
    """
    found = shapes(calls, dev)
    covering: dict[Hashable, list[Mapping[str, Any]]] = {s.key: [] for s in found}
    # What the covering entries set that the search does not move, as
    # GEMM's tile_k below HIGHEST.
    held: dict[Hashable, dict[str, Any]] = {s.key: {} for s in found}
    out = Profile()
    for cls, matched, tunables in profile or ():
        covered = [
            s
            for s in found
            if issubclass(s.cls, cls) and Profile.matches(matched, s.dims)
        ]
        for s in covered:
            covering[s.key].append(matched)
            held[s.key].update(tunables)
        if not covered:
            out.add(cls, **matched, **tunables)
    for i, shape in enumerate(found):
        log(f"[{i + 1}/{len(found)}] {shape}")
        tried = time_shape(
            shape,
            dev,
            full_elf=full_elf,
            timing=timing,
            repeats=repeats,
            cache=cache,
            remeasure=remeasure,
            cycles=cycles,
            core_mhz=core_mhz,
            log=log,
        )
        gains = [t for t in tried if t.gains]
        best = min(gains, key=lambda t: t.delta_us, default=tried[0])
        if gains:
            log(f"    {best.tunables}: {best.delta_us:+.1f} ± {best.noise_us:.1f} us")
        out.add(
            shape.cls,
            **shape.named(covering[shape.key]),
            **{**held[shape.key], **best.tunables},
        )
    return out

tune

Tune a model's profile on this NPU.

python -m iron.lm.tune <application> times its calls' operators and rewrites profiles/<device>.json with the fastest settings found; review and commit the change with git diff.

contexts

contexts(top: int, count: int, unit: int) -> list[tuple[float, int]]

At most count contexts halving down from top, multiples of unit.

Ascending, each weighted by the share of a uniform context in (0, top] nearest it (the trapezoid rule): the weights sum to 1.

Source code in iron/lm/tune.py

def contexts(top: int, count: int, unit: int) -> list[tuple[float, int]]:
    """At most ``count`` contexts halving down from ``top``, multiples of ``unit``.

    Ascending, each weighted by the share of a uniform context in
    ``(0, top]`` nearest it (the trapezoid rule): the weights sum to 1.
    """
    found = [top]
    while len(found) < count and found[-1] // 2 >= unit and found[-1] // 2 % unit == 0:
        found.append(found[-1] // 2)
    found.reverse()
    edges = [0] + [(a + b) / 2 for a, b in zip(found, found[1:])] + [top]
    return [((hi - lo) / top, c) for lo, hi, c in zip(edges, edges[1:], found)]

arguments

arguments(parser: ArgumentParser) -> None

Add the flags a decoder's calls follow from.

The checkpoint whose header gives the weights' shapes, the contexts attention over the cache is timed at, and the sampling Sample is timed at.

Source code in iron/lm/tune.py

def arguments(parser: argparse.ArgumentParser) -> None:
    """Add the flags a decoder's ``calls`` follow from.

    The checkpoint whose header gives the weights' shapes, the contexts
    attention over the cache is timed at, and the sampling Sample is
    timed at.
    """
    parser.add_argument(
        "--checkpoint",
        type=Path,
        required=True,
        help="the checkpoint directory, as Model.load takes it, of whose "
        "weights only the header is read",
    )
    parser.add_argument(
        "--position",
        type=int,
        default=256,
        help="the decode position the decode step's designs are timed at: "
        "what their per-call values follow from (default: 256)",
    )
    parser.add_argument(
        "--token",
        type=int,
        default=0,
        help="the token the step embeds (default: 0)",
    )
    parser.add_argument(
        "--context",
        type=int,
        default=8192,
        help="the longest context a design whose work follows it is timed "
        "at, a multiple of the prompt chunk (default: 8192)",
    )
    parser.add_argument(
        "--points",
        type=int,
        default=6,
        help="how many contexts up to --context it is timed at, each "
        "timed per setting: their cost grows with it (default: 6)",
    )
    sampling_arguments(parser)

calls

calls(model: CausalLM, position: int, token: int, context: int, points: int, sampler: Sampler) -> list[Timed]

Return the calls a decoder's designs are timed in.

The decode step and the prompt chunk, each over contexts up to context, points of them, and Sample on the draw rows sampler gives.

Source code in iron/lm/tune.py

def calls(
    model: CausalLM,
    position: int,
    token: int,
    context: int,
    points: int,
    sampler: Sampler,
) -> list[Timed]:
    """Return the calls a decoder's designs are timed in.

    The decode step and the prompt chunk, each over ``contexts`` up to
    ``context``, ``points`` of them, and Sample on the draw rows
    ``sampler`` gives.
    """
    L, C = model.config.max_seq_len, model.config.prefill_chunk
    decode_shapes, prompt_shapes = model.shapes()
    versions = [
        (
            decode_shapes,
            dict(position=position, token=token),
            [
                Point(w, dict(position=c - 1, token=token))
                for w, c in contexts(context, points, 1)
            ],
        ),
        (
            prompt_shapes,
            dict(
                position=C - 1,
                token=token,
                start=0,
                rows=C,
                rope_start=model.rope_row(C, 0),
            ),
            [
                Point(
                    w,
                    dict(
                        position=c - 1,
                        token=token,
                        start=c - C,
                        rows=C,
                        rope_start=model.rope_row(c, c - C),
                    ),
                )
                for w, c in contexts(context, points, C)
            ],
        ),
    ]
    out = []
    for shapes, values, timed in versions:
        recording = model.record(**shapes)
        # Sample's work follows its draw row's temperature and top-k: measure
        # it at the rows generation writes, not at random words.
        [k_max] = {s.op.k_max for s in recording.steps if isinstance(s.op, Sample)}
        _, draws = recording.states[id(model.draws)]
        rows = sampler.rows(L, k_max)
        out.append(Timed(recording, values, {draws.name: rows}, timed))
    return out

application_model

application_model(name: str) -> type[Model]

Return name's model class, name dotted under iron.applications.

Of those its package exports, the one that names its calls (tune_calls).

Raises:

Type	Description
ValueError	name names no such application, or its package exports other than one such class.

Source code in iron/lm/tune.py

def application_model(name: str) -> type[Model]:
    """Return ``name``'s model class, ``name`` dotted under ``iron.applications``.

    Of those its package exports, the one that names its calls
    (``tune_calls``).

    Raises:
        ValueError: ``name`` names no such application, or its package
            exports other than one such class.
    """
    try:
        package = importlib.import_module(f"iron.applications.{name}")
    except ImportError as e:
        raise ValueError(f"no application {name!r}: {e}") from e
    exported = [getattr(package, n) for n in vars(package).get("__all__", ())]
    found = [
        m
        for m in exported
        if isinstance(m, type) and any("tune_calls" in vars(k) for k in m.__mro__)
    ]
    if len(found) != 1:
        raise ValueError(
            f"iron.applications.{name} exports {len(found)} models to tune, not one"
        )
    return found[0]

main

main(name: str, model: Any, argv: Sequence[str] | None = None) -> None

Run the command line that tunes name's profile, into its profiles directory.

The calls model plans (tune_calls, on the flags tune_arguments adds).

Source code in iron/lm/tune.py

def main(name: str, model: Any, argv: Sequence[str] | None = None) -> None:
    """Run the command line that tunes ``name``'s profile, into its profiles directory.

    The calls ``model`` plans (``tune_calls``, on the flags
    ``tune_arguments`` adds).
    """
    parser = tune_cli.parser(f"Tune {name}'s profile")
    if model.profile is None:
        parser.error(f"{name} has no profiles directory to tune into")
    model.tune_arguments(parser)
    args = parse_args(parser, argv)
    tune_cli.tune(args, model.tune_calls(parser, args), Path(model.profile))

cli

cli() -> None

python -m iron.lm.tune <application>: tune <application>'s profile.

Tunes with main, its model found by application_model; --help after <application> lists the rest of its flags.

Source code in iron/lm/tune.py

def cli() -> None:
    """``python -m iron.lm.tune <application>``: tune ``<application>``'s profile.

    Tunes with ``main``, its model found by ``application_model``;
    ``--help`` after ``<application>`` lists the rest of its flags.
    """
    # No -h here: an unresolved --help passes through to main's parser,
    # which knows application's own flags.
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument(
        "application",
        help="dotted under iron.applications (e.g. llama, embedding_gemma2); "
        "the model its package exports is tuned",
    )
    args, rest = parser.parse_known_args()
    try:
        model = application_model(args.application)
    except ValueError as e:
        parser.error(str(e))
    main(args.application, model, rest)