Testing and tracing

Testing

How an operator declares the shapes it is tested at on a device, and the device test harness.

Draw vectors, run it, judge and time it. iron/tests/operators/test_catalog.py runs every declaration against the operator's reference(). Nothing here imports pytest, so an operator module stays importable without it.

Case dataclass

Case(kwargs: dict = dict(), extensive: bool = False, id: str | None = None, bench: bool = False, lower: bool = False)

One construction of an operator.

Attributes:

Name	Type	Description
kwargs	dict	The constructor's keywords.
extensive	bool	Keeps the case out of the default suite.
id	str | None	Its name in test output, else the arguments.
bench	bool	CI tracks it over time (pytest.mark.bench).
lower	bool	The toolchain tests lower it to instructions on every device, as they do each operator's first default case.

Testing dataclass

Testing(cases: Iterable[Case | _Cases] | _Cases, draw: Callable[[Any], dict[str, Any]] | None = None)

How an operator is checked against its reference on a device.

Attributes:

Name	Type	Description
cases	Iterable[Case | _Cases] | _Cases	Cases, or callables of the class and the device returning them (a Sweep).
draw	Callable[[Any], dict[str, Any]] | None	A callable of the operator returning extra vectors arguments (dict(normal=(op.A,))), for inputs with preconditions.

resolve

resolve(cls: type, dev: Device) -> list[Case]

Return the cases for cls on dev, so an inherited sweep reads the subclass's caps.

Source code in iron/common/testing.py

def resolve(self, cls: type, dev: Device) -> list[Case]:
    """Return the cases for ``cls`` on ``dev``, so an inherited sweep reads the subclass's caps."""
    out = []
    for c in [self.cases] if callable(self.cases) else self.cases:
        out += list(c(cls, dev)) if callable(c) else [c]
    return out

Vectors dataclass

Vectors(inputs: dict[str, ndarray], outputs: dict[str, ndarray])

One operator's test vectors, keyed by its declared buffer names.

vectors

vectors(op: Operator, *, seed: int = 42, scale: float = 4.0, normal: Iterable[Operand] = (), centered: Iterable[Operand] = (), given: Mapping[Operand, ndarray | tuple[int, ...]] | None = None) -> Vectors

Random inputs for op's declared buffers, and op.reference()'s outputs.

Parameters:

Name	Type	Description	Default
op	Operator	The operator whose declared buffers are drawn for.	required
seed	int	Seeds the draw, for reproducible vectors.	42
scale	float	A uniform draw's scale; an integer buffer draws on [0, scale].	4.0
normal	Iterable[Operand]	Inputs drawn normal instead (op.A).	()
centered	Iterable[Operand]	Inputs shifted to centre on zero.	()
given	Mapping[Operand, ndarray | tuple[int, ...]] | None	Per input, an array as it is, or a shape to draw in place of the declared one (an operand the sequence packs).	None

Raises:

Type	Description
ValueError	An operand named is not one of op's inputs.

Source code in iron/common/testing.py

def vectors(
    op: Operator,
    *,
    seed: int = 42,
    scale: float = 4.0,
    normal: Iterable[Operand] = (),
    centered: Iterable[Operand] = (),
    given: Mapping[Operand, np.ndarray | tuple[int, ...]] | None = None,
) -> Vectors:
    """Random inputs for ``op``'s declared buffers, and ``op.reference()``'s outputs.

    Args:
        op: The operator whose declared buffers are drawn for.
        seed: Seeds the draw, for reproducible vectors.
        scale: A uniform draw's scale; an integer buffer draws on ``[0, scale]``.
        normal: Inputs drawn normal instead (``op.A``).
        centered: Inputs shifted to centre on zero.
        given: Per input, an array as it is, or a shape to draw in place of
            the declared one (an operand the sequence packs).

    Raises:
        ValueError: An operand named is not one of ``op``'s inputs.
    """
    given_names = {o.name: v for o, v in (given or {}).items()}
    normal_names = {o.name for o in normal}
    centered_names = {o.name for o in centered}
    unknown = (given_names.keys() | normal_names | centered_names) - {
        b.name for b in op.inputs
    }
    if unknown:
        raise ValueError(f"{type(op).__name__} has no input {sorted(unknown)}")
    rng = np.random.default_rng(seed)
    inputs = {}
    for b in op.inputs:
        value = given_names.get(b.name)
        if isinstance(value, np.ndarray):
            inputs[b.name] = value
            continue
        shape = b.host_shape if value is None else tuple(value)
        # A dtype that follows tuning is None until tuned; the unpacked operand is bf16.
        dtype = np.dtype(bfloat16 if b.dtype is None else b.host_dtype)
        if dtype.kind not in "fc" and dtype != bfloat16:
            t = rng.integers(0, int(scale) + 1, shape).astype(dtype)
        else:
            draw = rng.standard_normal if b.name in normal_names else rng.random
            t = (draw(shape) * scale).astype(dtype)
            if b.name in centered_names:
                t = (t.astype(np.float32) - scale / 2).astype(dtype)
        inputs[b.name] = t
    return Vectors(inputs, op._call_reference(inputs))

time_call

time_call(fn: Callable, *args: Any, bytes_moved: int = 0, ops: float = 0.0, warmup: int = 0, iters: int = 20, record: Callable[[str, float], object] | None = None, **kwargs: Any) -> None

Time any compiled callable (an Operator, a graph version, a shipped binary).

The median NPU latency of iters runs of fn(*args, **kwargs). run_test is this plus the checked run and its judgement.

Parameters:

Name	Type	Description	Default
fn	Callable	The compiled callable timed.	required
*args	Any	Passed to fn on every run.	()
bytes_moved	int	Bytes moved a call, for the bandwidth figure; 0 records none.	0
ops	float	Per call, for the throughput figure; 0 records none.	0.0
warmup	int	Untimed runs before the iters timed ones.	0
iters	int	The timed runs the median latency is taken over.	20
record	Callable[[str, float], object] | None	A test's record_property: latency, its median absolute deviation over the timed runs, bandwidth and throughput.	None
**kwargs	Any	Passed to fn on every run.	{}

Raises:

Type	Description
RuntimeError	fn did not report NPU execution time.

Source code in iron/common/testing.py

def time_call(
    fn: Callable,
    *args: Any,
    bytes_moved: int = 0,
    ops: float = 0.0,
    warmup: int = 0,
    iters: int = 20,
    record: Callable[[str, float], object] | None = None,
    **kwargs: Any,
) -> None:
    """Time any compiled callable (an ``Operator``, a graph version, a shipped binary).

    The median NPU latency of ``iters`` runs of ``fn(*args, **kwargs)``.
    ``run_test`` is this plus the checked run and its judgement.

    Args:
        fn: The compiled callable timed.
        *args: Passed to ``fn`` on every run.
        bytes_moved: Bytes moved a call, for the bandwidth figure; 0 records none.
        ops: Per call, for the throughput figure; 0 records none.
        warmup: Untimed runs before the ``iters`` timed ones.
        iters: The timed runs the median latency is taken over.
        record: A test's ``record_property``: latency, its median absolute
            deviation over the timed runs, bandwidth and throughput.
        **kwargs: Passed to ``fn`` on every run.

    Raises:
        RuntimeError: ``fn`` did not report NPU execution time.
    """
    benchmark = run_iters(fn, *args, warmup=warmup, iters=iters, **kwargs)
    if benchmark.npu is None:
        raise RuntimeError("the callable did not report NPU execution time")
    latency_us = benchmark.npu.median_us
    if record is not None:
        record("Latency", latency_us)
        record("Latency MAD", benchmark.npu.mad_us)
        if bytes_moved:
            record("Bandwidth", bytes_moved / (latency_us * 1e-6) / 1e9)
        if ops:
            record("Throughput", ops / (latency_us * 1e-6) / 1e9)
    logger.info(
        "Latency (us): %.1f ± %.1f%s",
        latency_us,
        benchmark.npu.mad_us,
        (
            f"  Effective Bandwidth: {bytes_moved / (latency_us * 1e-6) / 1e9:.6e} GB/s"
            if bytes_moved
            else ""
        ),
    )

run_test

run_test(operator: 'Operator', inputs: Vectors | Mapping[str, ndarray], outputs: Mapping[str, ndarray] | None = None, *, tolerance: Tolerance, timed_iters: int = 20, record: Callable[[str, float], object] | None = None) -> dict[str, Verdict]

Compile operator, run it on the device, check its outputs, time it.

The first run is the one checked, before a timed run rewrites an inout buffer.

Parameters:

Name	Type	Description	Default
operator	'Operator'	The declared operator compiled and run.	required
inputs	Vectors | Mapping[str, ndarray]	A Vectors, or every input by buffer name (an inout buffer among them).	required
outputs	Mapping[str, ndarray] | None	Each buffer the operator writes, its expected result by name; operator.reference on inputs where None.	None
tolerance	Tolerance	What a checked output is judged by.	required
timed_iters	int	The timed runs the latency is taken over.	20
record	Callable[[str, float], object] | None	time_call's.	None

Returns:

Type	Description
dict[str, Verdict]	Each failing output's Verdict, by name; empty where all passed.

Raises:

Type	Description
ValueError	inputs or outputs do not name the operator's buffers, or an expected result is None.

Source code in iron/common/testing.py

def run_test(
    operator: "Operator",
    inputs: Vectors | Mapping[str, np.ndarray],
    outputs: Mapping[str, np.ndarray] | None = None,
    *,
    tolerance: Tolerance,
    timed_iters: int = 20,
    record: Callable[[str, float], object] | None = None,
) -> dict[str, Verdict]:
    """Compile ``operator``, run it on the device, check its outputs, time it.

    The first run is the one checked, before a timed run rewrites an
    ``inout`` buffer.

    Args:
        operator: The declared operator compiled and run.
        inputs: A ``Vectors``, or every input by buffer name (an ``inout``
            buffer among them).
        outputs: Each buffer the operator writes, its expected result by
            name; ``operator.reference`` on ``inputs`` where None.
        tolerance: What a checked output is judged by.
        timed_iters: The timed runs the latency is taken over.
        record: ``time_call``'s.

    Returns:
        Each failing output's ``Verdict``, by name; empty where all passed.

    Raises:
        ValueError: ``inputs`` or ``outputs`` do not name the operator's
            buffers, or an expected result is None.
    """
    from .declare import Operator  # which imports this module

    if isinstance(inputs, Vectors):
        inputs, outputs = inputs.inputs, inputs.outputs
    if not isinstance(operator, Operator):
        raise TypeError(f"run_test runs one declared Operator, not {operator!r}")
    declared = [b.name for b in operator.inputs]
    if sorted(inputs) != sorted(declared):
        raise ValueError(
            f"{type(operator).__name__} takes the inputs {declared}, "
            f"not {sorted(inputs)}"
        )
    operator.compile()
    # The tensor type of the host runtime NPU_RUNTIME selects, which the operator dispatches on.
    tensor_class = aie_utils.DEFAULT_TENSOR_CLASS
    args, produced, total_bytes = [], {}, 0
    for b in operator.buffers:
        if b.fills:
            buf = tensor_class(inputs[b.name])
        else:
            buf = tensor_class(b.host_shape, dtype=b.host_dtype)
        if b.drains:
            produced[b.name] = buf
        args.append(buf)
        total_bytes += _nbytes(buf)

    operator(*args)
    written = {name: buf.numpy() for name, buf in produced.items()}
    verdicts = operator.judge(inputs, written, tolerance, outputs)
    for name, verdict in verdicts.items():
        _report(name, verdict, tolerance)
    errors = {name: verdict for name, verdict in verdicts.items() if not verdict}

    ops = operator.resolved().ops()
    time_call(
        operator,
        *args,
        bytes_moved=total_bytes,
        ops=ops,
        iters=timed_iters,
        record=record,
    )
    return errors

Tracing

NPU hardware tracing, read back after a run.

An operator built with trace=TraceConfig(...) (aie.utils.trace) traces the workers its build() marks Worker(trace=), else its first. After run(), dump_traces(run, path) writes mlir-aie's raw trace text and Perfetto JSON.

dump_traces

dump_traces(run: ImageCallable, trace_file: str | Path, *, colshift: int | None = None, mlir: str | Path | None = None, summary: bool = True) -> list[Path]

Write a completed run's trace buffer as text and Perfetto JSON.

The text goes to trace_file and the JSON beside it, one per traced design. An untraced build writes nothing, so a test can call it unconditionally.

Parameters:

Name	Type	Description	Default
run	ImageCallable	The completed run whose trace buffer is read.	required
trace_file	str | Path	The text trace's path; the JSON paths are derived from it.	required
colshift	int | None	None lets the parser align the columns, as a design may load into another column; override it when that picks the wrong ones.	None
mlir	str | Path | None	The lowered module the parser reads, else the build's own.	None
summary	bool	Also print mlir-aie's per-tile cycles summary for each.	True

Returns:

Type	Description
list[Path]	The JSON paths, empty on an untraced build.

Source code in iron/common/tracing.py

def dump_traces(
    run: ImageCallable,
    trace_file: str | Path,
    *,
    colshift: int | None = None,
    mlir: str | Path | None = None,
    summary: bool = True,
) -> list[Path]:
    """Write a completed run's trace buffer as text and Perfetto JSON.

    The text goes to ``trace_file`` and the JSON beside it, one per traced
    design. An untraced build writes nothing, so a test can call it
    unconditionally.

    Args:
        run: The completed run whose trace buffer is read.
        trace_file: The text trace's path; the JSON paths are derived from it.
        colshift: None lets the parser align the columns, as a design may load
            into another column; override it when that picks the wrong ones.
        mlir: The lowered module the parser reads, else the build's own.
        summary: Also print mlir-aie's per-tile cycles summary for each.

    Returns:
        The JSON paths, empty on an untraced build.
    """
    if not isinstance(run, FullELFCallable):
        if run.op.traced:
            raise TypeError(
                f"{type(run).__name__} was built with tracing enabled but has no "
                "trace buffer; only the full-ELF sequence callable allocates one"
            )
        return []
    buffer = run.trace_buffer
    if buffer is None:
        return []

    trace_file = Path(trace_file)
    trace_file.parent.mkdir(parents=True, exist_ok=True)
    words = buffer.numpy().view(np.uint32).reshape(-1)
    config = TraceConfig(trace_size=words.nbytes, trace_file=str(trace_file))
    config.write_trace(words)
    if not words.any():
        logger.warning("trace buffer is all zeros, no trace data captured")
        return []

    mlir = mlir or run.lowered_mlir_path
    logger.info("parsing the trace against %s", mlir)
    try:
        written = config.trace_to_json(
            str(mlir),
            str(trace_file.with_suffix(".json")),
            colshift=colshift,
            kernel=f"{run.device_name}:{run.sequence_name}",
        )
    except Exception as exc:  # a visualisation failure must not fail a run
        logger.warning("trace parse failed (%s); raw words kept at %s", exc, trace_file)
        return []

    paths = [Path(p) for p in written]
    for path in paths:
        logger.info("trace written to %s", path)
        if summary:
            print_cycles_summary(path)
    return paths