iron: sync freshly written arena buffers to the device in the fused ELF callable

[atassis]

iron: sync freshly written arena buffers to the device in the fused ELF callable

Summary

A fused-ELF OperatorSequence dispatch silently runs on stale (init-zero) inputs, because writing a get_buffer() sub-view's .data never marks the parent buffer host-dirty, so the per-dispatch host to device sync is skipped by a residency-flag optimization. The op computes on zeros and produces zero, or partial and nondeterministic, output. This fixes it in the two places that own the arena sync.

Root cause

SequenceFullELFCallable re-authors its input and scratch arenas on every dispatch through the sub-views returned by get_buffer(). Those sub-views (XRTSubBuffer) alias the parent XRTTensor's buffer.

XRTTensor.to(target) is a residency tracker with a skip optimization: it only performs the copy when the tracked device flag differs from the target. The input arena is allocated with device="npu" and synced once at construction, so its flag is "npu". Writing a sub-view's .data afterward does not change that flag, so _sync_inputs' input_buffer.to("npu") sees "npu" == "npu" and does nothing. The freshly written inputs, and resident weights that a caller loads with scratch_buffer.to("npu"), never reach the device. The kernel then runs on the zeros left by the one-time init sync.

The failure looks intermittent rather than always-zero because the sub-view and parent alias the same pages, so freshly written host bytes sometimes become visible to the device without an explicit sync and sometimes do not. The existing XRTSubBuffer.to() FIXME already flags this sub-buffer sync ambiguity and asks for per-region dirty tracking.

Fix

Two authoritative points, iron/common/utils.py and iron/common/sequence.py:

1. XRTSubBuffer.data marks the sub-view and its parent "cpu" on access. A later parent .to("npu") then actually fires the host to device sync. numpy gives no write hook, and callers write with np.copyto(sub.data, ...), which goes through the buffer protocol rather than __setitem__, so a precise write-only hook would miss exactly those writes. Marking on access is the reliable choice: a redundant read sync is cheap, a silently skipped write sync is a correctness bug.

2. _sync_outputs asserts device residency before .to("cpu"). _run just rewrote the output arena on the device, so the device holds the authoritative copy; forcing the flag makes the device to host sync fire every dispatch, including after a prior read of the output .data marked the buffer "cpu". Without this a looped dispatch would read stale output.

19 lines across 2 files, no API change.

Alternative considered

Forcing the sync at the two sync sites only (_sync_inputs host to device, _sync_outputs device to host) and leaving .data untouched is cleaner for the input and output arenas, but it does not cover resident weights, which callers load into the scratch arena and sync with scratch_buffer.to("npu"); that user-driven .to() still needs the write to have marked the buffer dirty. Marking on .data access is the smaller change that makes the existing get_buffer().data plus .to() pattern correct for all three arenas. Happy to switch to whichever the maintainers prefer, or to a proper per-region dirty flag as the FIXME suggests.

Verification

On AIE2P / NPU2, against a host reference, before and after:

workload	before	after
single GEMV, fused dispatch	rel-L2 ~ 1.0, varies per process	rel-L2 0.0016, deterministic
weighted RMSNorm, fused dispatch	rel-L2 0.71 to 1.0, bimodal	rel-L2 0.0087, deterministic
RMSNorm, 8 looped dispatches	zeros or a fixed wrong value	all 8 identical, correct

End to end, a full 18-layer Gemma-3-270M greedy decode ("The capital of France is") went from 0 of 8 tokens matching the host oracle (garbage first token) to generating the correct answer " Paris" and the first 3 tokens exactly. The remaining decode divergence is a separate numerical matter, not this bug.

Repro

Self contained, no external weights. Prints rel-L2 and asserts correctness; fails on the current devel, passes with this change:

import numpy as np, ml_dtypes
from iron.common import AIEContext
from iron.common.sequence import OperatorSequence
from iron.operators.gemv.op import GEMV

bf16 = ml_dtypes.bfloat16
M, K = 64, 640
ctx = AIEContext()
op = GEMV(M=M, K=K, num_aie_columns=1, tile_size_input=2, context=ctx)
seq = OperatorSequence("repro", [(op, "A", "x", "out")],
                       input_args=["A", "x"], output_args=["out"],
                       buffer_sizes={"A": M*K*2, "x": K*2, "out": M*2},
                       dispatch="fused", context=ctx)
seq.compile(); c = seq.get_callable()
rng = np.random.default_rng(0)
A = rng.standard_normal((M, K)).astype(bf16); x = rng.standard_normal(K).astype(bf16)
np.copyto(c.get_buffer("A").data, A.reshape(-1)); np.copyto(c.get_buffer("x").data, x)
c()
out = c.get_buffer("out").data[:M].astype(np.float32)
ref = A.astype(np.float32) @ x.astype(np.float32)
rel = float(np.linalg.norm(out - ref) / np.linalg.norm(ref))
print("rel-L2 =", rel)
assert rel < 1e-2
print("PASS")

Runs on any AIE2P device with a single dispatch. I can fold this into the test suite in whatever form you prefer.