diff --git a/.github/workflows/ci-lint.yml b/.github/workflows/ci-lint.yml index 29ab99193..b1d5b603c 100755 --- a/.github/workflows/ci-lint.yml +++ b/.github/workflows/ci-lint.yml @@ -21,7 +21,8 @@ jobs: uses: actions/checkout@v3 - name: Install dependencies run: | - apt-get update && apt-get install -y clang-format + sudo apt-get update + sudo apt-get install -y clang-format pip install reuse black - name: Check Licenses diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 5aa2bb523..c1ca04e35 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -705,9 +705,6 @@ class _PerBufferCallable(SequenceCallable): def _make_buffer(self, n_elements): raise NotImplementedError - def _make_subbuffer(self, parent, offset_bytes, size_bytes): - raise NotImplementedError - def _allocate_buffers(self): self._buffers = {} for name, (_, _, length) in self.op.subbuffer_layout.items(): @@ -718,8 +715,9 @@ def _resolve_buffer(self, buf_name): return self._buffers[buf_name] if buf_name in self.op.slice_info: base_name, start_bytes, end_bytes = self.op.slice_info[buf_name] - sub = self._make_subbuffer( - self._buffers[base_name], start_bytes, end_bytes - start_bytes + size_bytes = end_bytes - start_bytes + sub = self._buffers[base_name].subview( + start_bytes, (size_bytes // BF16.itemsize,), BF16 ) self._buffers[buf_name] = sub return sub @@ -756,11 +754,6 @@ def __init__(self, op, dispatch): def _make_buffer(self, n_elements): return XRTTensor((n_elements,), dtype=ml_dtypes.bfloat16) - def _make_subbuffer(self, parent, offset_bytes, size_bytes): - return parent.subview( - offset_bytes, (size_bytes // BF16.itemsize,), ml_dtypes.bfloat16 - ) - def _allocate_buffers(self): super()._allocate_buffers() dispatch = self._dispatch @@ -807,15 +800,9 @@ class SequenceReferenceCallable(_PerBufferCallable): def _make_buffer(self, n_elements): return CPUOnlyTensor((n_elements,), dtype=BF16) - def _make_subbuffer(self, parent, offset_bytes, size_bytes): - start = offset_bytes // BF16.itemsize - end = (offset_bytes + size_bytes) // BF16.itemsize - # Alias the parent's memory (numpy slice is zero-copy) so a write to - # this slice is visible when a later step reads the parent by name. - view = CPUOnlyTensor((end - start,), dtype=BF16) - view._data = parent.data[start:end] - view._shape = view._data.shape - return view + def _sync_inputs(self): + # CPU-only inputs must stay CPU-resident, including lazily created subviews. + pass def _run(self): torch = _torch() diff --git a/iron/common/test_utils.py b/iron/common/test_utils.py index bca6c81c9..afda7607f 100644 --- a/iron/common/test_utils.py +++ b/iron/common/test_utils.py @@ -6,6 +6,7 @@ import numpy as np import torch import aie.utils as aie_utils +from aie.utils.benchmark import run_iters from ml_dtypes import bfloat16 from .base import AIEOperatorBase @@ -216,16 +217,10 @@ def run_test( else: raise ValueError(f"Unsupported direction: {spec.direction}") - # Run warmup iterations - for _ in range(warmup_iters): - op_func(*args) - - # Run timed iterations and measure NPU execution time - total_npu_ns = 0 - for _ in range(timed_iters): - result = op_func(*args) - total_npu_ns += result.npu_time - latency_us = (total_npu_ns / timed_iters) / 1e3 + benchmark = run_iters(op_func, *args, warmup=warmup_iters, iters=timed_iters) + if benchmark.npu is None: + raise RuntimeError("Operator callable did not report NPU execution time") + latency_us = benchmark.npu.avg_us # Verify outputs errors = {} diff --git a/iron/common/tracing_utils.py b/iron/common/tracing_utils.py index b98396148..99dedb2e7 100644 --- a/iron/common/tracing_utils.py +++ b/iron/common/tracing_utils.py @@ -140,7 +140,7 @@ def dump_traces( mlir_path, mlir_text = lowered_mlir(run) print(f"[trace] parsing against {mlir_path}") - words = buffer.to_torch().numpy().astype(np.uint8).view(np.uint32) + words = buffer.numpy().view(np.uint32).reshape(-1) tag = _slug(tag) raw = (out_dir / tag).with_suffix(".txt") raw.write_text("\n".join(f"{w:08x}" for w in words) + "\n") diff --git a/iron/tests/infrastructure/benchmark.py b/iron/tests/infrastructure/benchmark.py new file mode 100644 index 000000000..e1fdb4897 --- /dev/null +++ b/iron/tests/infrastructure/benchmark.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The operator test adapter keeps reporting NPU, not host, latency.""" + +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + +from aie.utils.hostruntime.tensor_class import CPUOnlyTensor +from iron.common.base import AIEOperatorBase, AIERuntimeArgSpec +from iron.common import test_utils + + +class _Operator(AIEOperatorBase): + def __init__(self, results): + self.results = iter(results) + self.calls = 0 + + def set_up_artifacts(self): + pass + + def compile(self): + return self + + def get_arg_spec(self): + return [ + AIERuntimeArgSpec("in", (32,)), + AIERuntimeArgSpec("out", (32,)), + ] + + def get_callable(self): + def run(source, target): + self.calls += 1 + target[:] = source.numpy() + return next(self.results) + + return run + + +@pytest.mark.parametrize("tuple_result", [False, True]) +def test_run_test_uses_upstream_npu_timing(monkeypatch, tuple_result): + monkeypatch.setattr(test_utils.aie_utils, "DEFAULT_TENSOR_CLASS", CPUOnlyTensor) + results = [SimpleNamespace(npu_time=ns) for ns in (1000000, 2000, 4000)] + if tuple_result: + results = [(None, result) for result in results] + op = _Operator(results) + data = torch.ones(32, dtype=torch.bfloat16) + + errors, latency_us, bandwidth = test_utils.run_test( + op, {"in": data}, {"out": data}, warmup_iters=1, timed_iters=2 + ) + + assert op.calls == 3 + assert errors == {} + assert latency_us == 3.0 + assert bandwidth == pytest.approx(128 / (3e-6) / 1e9) + + +def test_missing_npu_timing_is_rejected(monkeypatch): + monkeypatch.setattr(test_utils.aie_utils, "DEFAULT_TENSOR_CLASS", CPUOnlyTensor) + op = _Operator([None]) + data = torch.ones(32, dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="NPU execution time"): + test_utils.run_test( + op, {"in": data}, {"out": data}, warmup_iters=0, timed_iters=1 + ) diff --git a/iron/tests/infrastructure/sequence.py b/iron/tests/infrastructure/sequence.py index ee877e3bf..a1399e3d9 100644 --- a/iron/tests/infrastructure/sequence.py +++ b/iron/tests/infrastructure/sequence.py @@ -203,6 +203,73 @@ def test_dispatch_modes_bit_identical(dispatch, aie_context): ) +# --------------------------------------------------------------------------- +# 3b. dispatch="reference" resolves slice-notation buffers the same way the +# NPU dispatch paths do: via NpuTensor.subview() on the CPU backend, +# rather than a hand-rolled numpy view. Not covered by +# test_dispatch_modes_bit_identical above, since reference() is a CPU +# re-implementation and only expected to match the NPU output within +# tolerance, not bit-for-bit (see CompareDispatch's rel_tol/abs_tol). +# --------------------------------------------------------------------------- + +_SLICE_SIZE = 1024 +_SLICE_BYTES = _SLICE_SIZE * 2 # bf16 + + +def _build_packed_output_sequence(context, dispatch, name): + """Two independent adds writing into disjoint halves of one explicitly + sized buffer via slice notation ("packed[start:end]"). Unlike + _build_add_relu_sequence's "temp" hand-off (a whole-buffer alias), this + exercises slice_info/explicit_buffer_sizes resolution directly.""" + add0 = ElementwiseAdd( + size=_SLICE_SIZE, tile_size=_SLICE_SIZE, num_aie_columns=1, context=context + ) + add1 = ElementwiseAdd( + size=_SLICE_SIZE, tile_size=_SLICE_SIZE, num_aie_columns=1, context=context + ) + return OperatorSequence( + name=name, + runlist=[ + (add0, "a0", "b0", f"packed[0:{_SLICE_BYTES}]"), + (add1, "a1", "b1", f"packed[{_SLICE_BYTES}:{2 * _SLICE_BYTES}]"), + ], + input_args=["a0", "b0", "a1", "b1"], + output_args=["packed"], + buffer_sizes={"packed": 2 * _SLICE_BYTES}, + dispatch=dispatch, + context=context, + ) + + +def test_reference_dispatch_resolves_sliced_buffer(aie_context): + """dispatch="reference" must resolve slice-notation buffers via + subview() on the CPU backend, matching SequenceXclbinCallable's behaviour, + and each slice's write must be visible through the parent buffer name.""" + torch.manual_seed(0) + a0 = torch.rand(_SLICE_SIZE, dtype=torch.bfloat16) + b0 = torch.rand(_SLICE_SIZE, dtype=torch.bfloat16) + a1 = torch.rand(_SLICE_SIZE, dtype=torch.bfloat16) + b1 = torch.rand(_SLICE_SIZE, dtype=torch.bfloat16) + + seq = _build_packed_output_sequence( + aie_context, "reference", "infra_reference_sliced_packed" + ) + seq.compile() + run = seq.get_callable() + _set_input(run, "a0", a0) + _set_input(run, "b0", b0) + _set_input(run, "a1", a1) + _set_input(run, "b1", b1) + run() + packed = run.get_buffer("packed").torch_view()[: 2 * _SLICE_SIZE].clone() + + expected = torch.cat([a0 + b0, a1 + b1]) + errors = verify_buffer(packed, "packed", expected, rel_tol=0.04, abs_tol=1e-6) + assert ( + not errors + ), f"reference-dispatch sliced buffer produced {len(errors)} mismatches" + + # --------------------------------------------------------------------------- # 4. Compare mode flags (and by default raises on) a per-step reference/NPU # mismatch on its own. diff --git a/iron/tests/infrastructure/sequence_subviews.py b/iron/tests/infrastructure/sequence_subviews.py new file mode 100644 index 000000000..64ef417b6 --- /dev/null +++ b/iron/tests/infrastructure/sequence_subviews.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Host-only coverage of the shared upstream tensor subview path.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +from ml_dtypes import bfloat16 + +from iron.common.sequence import SequenceReferenceCallable + + +@pytest.fixture +def run(): + op = SimpleNamespace( + subbuffer_layout={"packed": ("output", 0, 1024)}, + slice_info={ + "first": ("packed", 0, 512), + "second": ("packed", 512, 1024), + }, + ) + return SequenceReferenceCallable(op) + + +@pytest.mark.parametrize("name, start", [("first", 0), ("second", 256)]) +def test_slices_alias_the_parent_and_are_cached(run, name, start): + parent = run.get_buffer("packed") + parent.fill_(0) + view = run.get_buffer(name) + + assert view is run.get_buffer(name) + assert view is run._resolve_buffer(name) + assert view.dtype == np.dtype(bfloat16) + assert view.shape == (256,) + assert np.shares_memory(view.data, parent.data) + + view.fill_(3) + expected = np.zeros(512, dtype=bfloat16) + expected[start : start + 256] = 3 + np.testing.assert_array_equal(parent.numpy(), expected) + + parent.fill_(7) + np.testing.assert_array_equal(view.numpy(), np.full(256, 7, dtype=bfloat16)) + + +def test_unknown_buffer_is_rejected(run): + with pytest.raises(ValueError, match="Unknown buffer"): + run.get_buffer("missing") + + +def test_out_of_bounds_slice_is_rejected_by_upstream(run): + run.op.slice_info["invalid"] = ("packed", 512, 1536) + with pytest.raises(ValueError): + run.get_buffer("invalid") + + +def test_input_slices_resolve_during_reference_dispatch(run, monkeypatch): + run.op.input_args = ["packed"] + parent = run.get_buffer("packed") + + def evaluate(): + assert parent.device == "cpu" + for name in ("first", "second"): + view = run._resolve_buffer(name) + assert view.device == "cpu" + np.testing.assert_array_equal(view.numpy(), parent.numpy()[:256]) + + monkeypatch.setattr(run, "_run", evaluate) + for value in (3, 7): + parent.fill_(value) + run() diff --git a/iron/tests/infrastructure/tracing.py b/iron/tests/infrastructure/tracing.py new file mode 100644 index 000000000..cd2de3566 --- /dev/null +++ b/iron/tests/infrastructure/tracing.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Trace dumps use the upstream tensor's host interface without torch.""" + +import json +from types import SimpleNamespace + +import numpy as np +import pytest +from aie.utils.hostruntime.tensor_class import CPUOnlyTensor + +from iron.common import tracing_utils + + +@pytest.mark.parametrize("dtype", [np.int8, np.uint8]) +def test_dump_preserves_raw_trace_bits(monkeypatch, tmp_path, dtype): + words = np.array([0xFFFFFFFF, 0x80000000, 0x12345678, 0], dtype=np.uint32) + buffer = CPUOnlyTensor(words.view(dtype), dtype=dtype) + run = SimpleNamespace(trace_buffer=buffer) + monkeypatch.setattr( + tracing_utils, "lowered_mlir", lambda run: (tmp_path / "test.mlir", "mlir") + ) + events = [{"name": "event"}] + + def parse(actual, mlir_text, colshift): + np.testing.assert_array_equal(actual, words) + assert mlir_text == "mlir" + assert colshift == 2 + return [(None, events)] + + monkeypatch.setattr(tracing_utils, "parse_trace_buffer", parse) + written = tracing_utils.dump_traces( + run, "test", out_dir=tmp_path, colshift=2, summary=False + ) + + assert written == [tmp_path / "test_trace.json"] + assert json.loads(written[0].read_text()) == events + assert (tmp_path / "test.txt").read_text().splitlines() == [ + "ffffffff", + "80000000", + "12345678", + "00000000", + ] + + +def test_untraced_run_needs_no_buffer(tmp_path): + assert tracing_utils.dump_traces(SimpleNamespace(), "test", tmp_path) == [] diff --git a/requirements.txt b/requirements.txt index fdedee171..ffddf4ba8 100755 --- a/requirements.txt +++ b/requirements.txt @@ -9,13 +9,12 @@ # CUDA build served from PyPI. We therefore also pin torch to the "+cpu" local # version below, which is only available from the PyTorch CPU index. --index-url https://download.pytorch.org/whl/cpu ---find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/v1.4.3 --find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-4 --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly --extra-index-url https://pypi.org/simple -mlir_aie==1.4.3 -llvm-aie==22.0.0.2026090701+3e93bf7b +mlir_aie==1.4.4.dev4+g20a9c2f +llvm-aie==22.0.0.2026091701+773413fb black reuse