Update mlir-aie and use mlir-aie utils when possible

.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

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

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 = {}

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

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

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.

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

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) == []

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