Perf: reuse L2 IN args across SceneTest rounds

docs/task-flow.md

@@ -781,9 +781,10 @@ before execution.
 Tags (IN/OUT/INOUT/…) are used by `Orchestrator::submit_*` to derive TensorMap
 dependencies and nothing else. Scheduler, WorkerThread, child, runtime.so, and
 kernels do not inspect them. Keeping tags only in Layer ① simplifies the blob
-and makes the "tags are Orchestrator input" rule explicit. Matches existing
-runtime: `ChipStorageTaskArgs` (`task_args.h`) is already declared with
-`void` as the TensorTag parameter.
+and makes the "tags are Orchestrator input" rule explicit. The fixed
+`ChipStorageTaskArgs` runtime ABI therefore carries tensors and scalars without
+dependency tags. Its optional host-view sidecar is local materialization
+metadata and does not cross the L3 dispatch wire.
 
 ### Why no `WorkerPayload` wrapper
 
@@ -833,7 +834,7 @@ lives in the mailbox blob bytes on the child side — view doesn't care.
 - [chip-level-arch.md](chip-level-arch.md) — L2 single-chip: three-program
   model (host / AICPU / AICore)
 - [`../src/common/task_interface/task_args.h`](../src/common/task_interface/task_args.h)
-  — `TaskArgsTpl` template and the `ChipStorageTaskArgs` alias
+  — `TaskArgsTpl` template and the `ChipStorageTaskArgs` runtime ABI
 - [`../src/common/task_interface/task_args_wire.h`](../src/common/task_interface/task_args_wire.h)
   — the L3+ `TaskArgs`, `TaskArgsView`, and the mailbox blob codec
 - [`../src/common/task_interface/tensor.h`](../src/common/task_interface/tensor.h)

docs/testing.md

@@ -314,7 +314,8 @@ Layer 1  Level axis
     │       └─ Layer 4  Class — one ChipWorker per (runtime, device), reused
     │                   across every class assigned to that device
     │           └─ Layer 5  Case — serial within a class
-    │               └─ Layer 6  Rounds — `--rounds N` loop, reuses Worker
+    │               └─ Layer 6  Rounds — `--rounds N` loop, reuses Worker;
+    │                           L2 read-only inputs are uploaded once when N > 1
 ```
 
 ### Quick examples

python/bindings/task_interface.cpp

@@ -2229,6 +2229,10 @@ NB_MODULE(_task_interface, m) {
         .def("tensor_count", &ChipStorageTaskArgs::tensor_count)
         .def("scalar_count", &ChipStorageTaskArgs::scalar_count)
 
+        .def("_set_host_view", &ChipStorageTaskArgs::set_host_view, nb::arg("i"), nb::arg("addr"))
+
+        .def("_host_view", &ChipStorageTaskArgs::host_view, nb::arg("i"))
+
         .def("clear", &ChipStorageTaskArgs::clear)
 
         .def(
@@ -2286,6 +2290,10 @@ NB_MODULE(_task_interface, m) {
             "Add a uint64_t scalar. After this, add_tensor() is no longer allowed."
         )
 
+        .def("_set_host_view", &TaskArgs::set_host_view, nb::arg("i"), nb::arg("addr"))
+
+        .def("_host_view", &TaskArgs::host_view, nb::arg("i"))
+
         .def(
             "add_dep",
             [](TaskArgs &self, nb::args deps) {
@@ -3241,6 +3249,7 @@ NB_MODULE(_task_interface, m) {
             ChipStorageTaskArgs out;
             for (int32_t i = 0; i < args.tensor_count(); i++) {
                 out.add_tensor(materialize_one(args.tensor(i), resolved));
+                out.set_host_view(i, args.host_view(i));
             }
             for (int32_t i = 0; i < args.scalar_count(); i++) {
                 out.add_scalar(args.scalar(i));

simpler_setup/scene_test.py

@@ -555,14 +555,74 @@ def keep(self, *objs):
 # ---------------------------------------------------------------------------
 
 
-def _build_l2_ref_args(test_args: TaskArgsBuilder, orch_signature: list, worker):
+def _pin_l2_round_inputs(worker, test_args: TaskArgsBuilder, orch_signature: list):
+    """Pre-upload orchestration IN tensors for multi-round L2 SceneTest reuse.
+
+    Called only when ``rounds > 1``. Each nonempty IN is ``worker.malloc`` +
+    ``worker.copy_to`` once before the round loop; subsequent rounds pass the
+    device tensor while retaining the caller buffer as its host-orchestration
+    view. OUT/INOUT stay on the host-staging path.
+
+    Returns:
+        pinned: name → device ``Tensor`` for each IN
+        buffers: device ``Buffer`` handles to ``worker.free`` after the case
+    """
+    from simpler.task_interface import ArgDirection  # noqa: PLC0415
+
+    from simpler_setup.torch_interop import torch_dtype_to_datatype  # noqa: PLC0415
+
+    pinned: dict[str, Any] = {}
+    buffers: list[Any] = []
+    tensor_idx = 0
+    try:
+        for spec in test_args.specs:
+            if not isinstance(spec, TensorArg):
+                continue
+            if tensor_idx >= len(orch_signature):
+                raise ValueError(
+                    f"TensorArg '{spec.name}' at index {tensor_idx} has no matching entry in "
+                    f"orchestration signature (length {len(orch_signature)}). "
+                    f"Update CALLABLE['orchestration']['signature'] to match generate_args()."
+                )
+            direction = orch_signature[tensor_idx]
+            tensor_idx += 1
+            if direction != ArgDirection.IN:
+                continue
+            host = spec.value
+            if host.device.type != "cpu":
+                raise ValueError(f"L2 round input '{spec.name}' must be a CPU tensor, got device={host.device}.")
+            if not host.is_contiguous():
+                raise ValueError(
+                    f"L2 round input '{spec.name}' must be contiguous; call tensor.contiguous() before passing it."
+                )
+            nbytes = int(host.numel()) * int(host.element_size())
+            if nbytes == 0:
+                continue
+            buf = worker.malloc(nbytes)
+            buffers.append(buf)
+            worker.copy_to(buf, host)
+            dtype = int(torch_dtype_to_datatype(host.dtype).value)
+            shapes = tuple(int(s) for s in host.shape)
+            pinned[spec.name] = buf.tensor(shapes, dtype)
+    except Exception:
+        for buf in buffers:
+            worker.free(buf)
+        raise
+    return pinned, buffers
+
+
+def _build_l2_ref_args(test_args: TaskArgsBuilder, orch_signature: list, worker, pinned_inputs=None):
     """Build TensorArg `TaskArgs` from `TaskArgsBuilder` for the L2 `Worker.run` path.
 
     An L2 leaf consumes its own args: `Worker.run(handle, args, cfg)` materializes each TensorArg to a
     local base in-process. Each tensor is named via ``worker.make_tensor_arg`` (a host tensor;
     at L2 there is no fork, so any host tensor resolves in-process); the direction tag is inert at L2
     but set for parity with the L3 path.
 
+    When ``pinned_inputs`` is provided, orchestration IN tensors use the
+    pre-uploaded device views instead of host staging. Their original host
+    addresses remain attached as host-only metadata for HBG orchestration.
+
     Returns:
         args: TaskArgs (TensorArg)
         output_names: list of tensor names that are OUTPUT or INOUT
@@ -571,6 +631,7 @@ def _build_l2_ref_args(test_args: TaskArgsBuilder, orch_signature: list, worker)
 
     from simpler_setup.torch_interop import make_tensor_arg  # noqa: PLC0415
 
+    pinned_inputs = pinned_inputs or {}
     dir2tag = {
         ArgDirection.IN: TensorArgType.INPUT,
         ArgDirection.OUT: TensorArgType.OUTPUT_EXISTING,
@@ -588,7 +649,13 @@ def _build_l2_ref_args(test_args: TaskArgsBuilder, orch_signature: list, worker)
                     f"Update CALLABLE['orchestration']['signature'] to match generate_args()."
                 )
             direction = orch_signature[tensor_idx]
-            args.add_tensor(make_tensor_arg(worker, spec.value), dir2tag.get(direction, TensorArgType.INPUT))
+            if direction == ArgDirection.IN and spec.name in pinned_inputs:
+                tensor_arg = pinned_inputs[spec.name]
+            else:
+                tensor_arg = make_tensor_arg(worker, spec.value)
+            args.add_tensor(tensor_arg, dir2tag.get(direction, TensorArgType.INPUT))
+            if direction == ArgDirection.IN and spec.name in pinned_inputs:
+                args._set_host_view(tensor_idx, int(spec.value.data_ptr()))
             if direction in (ArgDirection.OUT, ArgDirection.INOUT):
                 output_names.append(spec.name)
             tensor_idx += 1
@@ -1717,54 +1784,60 @@ def _run_and_validate_l2(  # noqa: PLR0913 -- threads CLI diagnostic flags + cas
             handle = worker.register(callable_obj)
             type(self)._st_l2_handle = handle
 
-        # Build args
         test_args = self.generate_args(params)
-        chip_args, output_names = _build_l2_ref_args(test_args, orch_sig, worker)
+        pinned_inputs: dict[str, Any] = {}
+        pinned_buffers: list[Any] = []
+        try:
+            if rounds > 1:
+                pinned_inputs, pinned_buffers = _pin_l2_round_inputs(worker, test_args, orch_sig)
+            chip_args, output_names = _build_l2_ref_args(test_args, orch_sig, worker, pinned_inputs=pinned_inputs)
 
-        # Compute golden (unless skip_golden)
-        golden_args = None
-        if not skip_golden:
-            golden_args = test_args.clone()
-            with _golden_thread_cap():
-                self.compute_golden(golden_args, params)
+            golden_args = None
+            if not skip_golden:
+                golden_args = test_args.clone()
+                with _golden_thread_cap():
+                    self.compute_golden(golden_args, params)
 
-        _log_torch_backend_autoload_once()
-
-        # Save initial output tensor values for reset between rounds
-        initial_outputs = {}
-        if rounds > 1:
-            for name in output_names:
-                initial_outputs[name] = getattr(test_args, name).clone()
-
-        # Execute rounds. The platform emits `[STRACE]` host/device markers to
-        # stderr on every run; multi-round timing is obtained by teeing stderr
-        # to a file and parsing it offline with
-        # `python -m simpler_setup.tools.strace_timing <log> --rounds-table`
-        # (the scene test no longer captures/parses inline). See
-        # docs/dfx/l2-timing.md.
-        for round_idx in range(rounds):
-            if round_idx > 0:
-                for name, initial in initial_outputs.items():
-                    getattr(test_args, name).copy_(initial)
-
-            # Every diagnostic reaching this loop is already multi-round-safe:
-            # effective_diagnostic_options zeroes all of them when rounds > 1,
-            # so no per-round masking belongs here.
-            config = self._build_config(
-                config_dict,
-                enable_chip_swimlane=enable_chip_swimlane,
-                enable_dump_args=enable_dump_args,
-                enable_pmu=enable_pmu,
-                enable_dep_gen=enable_dep_gen,
-                enable_scope_stats=enable_scope_stats,
-                output_prefix=output_prefix,
-            )
+            _log_torch_backend_autoload_once()
+
+            # Save initial output tensor values for reset between rounds
+            initial_outputs = {}
+            if rounds > 1:
+                for name in output_names:
+                    initial_outputs[name] = getattr(test_args, name).clone()
+
+            # Execute rounds. The platform emits `[STRACE]` host/device markers to
+            # stderr on every run; multi-round timing is obtained by teeing stderr
+            # to a file and parsing it offline with
+            # `python -m simpler_setup.tools.strace_timing <log> --rounds-table`
+            # (the scene test no longer captures/parses inline). See
+            # docs/dfx/l2-timing.md.
+            for round_idx in range(rounds):
+                if round_idx > 0:
+                    for name, initial in initial_outputs.items():
+                        getattr(test_args, name).copy_(initial)
 
-            with _temporary_env(self._resolve_env()):
-                worker.run(handle, chip_args, config=config)
+                # Every diagnostic reaching this loop is already multi-round-safe:
+                # effective_diagnostic_options zeroes all of them when rounds > 1,
+                # so no per-round masking belongs here.
+                config = self._build_config(
+                    config_dict,
+                    enable_chip_swimlane=enable_chip_swimlane,
+                    enable_dump_args=enable_dump_args,
+                    enable_pmu=enable_pmu,
+                    enable_dep_gen=enable_dep_gen,
+                    enable_scope_stats=enable_scope_stats,
+                    output_prefix=output_prefix,
+                )
 
-            if not skip_golden:
-                self.compare_outputs(test_args, golden_args, output_names, params)
+                with _temporary_env(self._resolve_env()):
+                    worker.run(handle, chip_args, config=config)
+
+                if not skip_golden:
+                    self.compare_outputs(test_args, golden_args, output_names, params)
+        finally:
+            for buf in pinned_buffers:
+                worker.free(buf)
 
     def _run_and_validate_l3(  # noqa: PLR0913 -- threads CLI diagnostic flags + L3 ns context
         self,

src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp

@@ -1102,7 +1102,23 @@ extern "C" int bind_callable_to_runtime_impl(
         // as if it were a graph-heap allocation.
         if (t.is_device_memory()) {
             always_assert(t.buffer.addr < HEAP_VIRTUAL_BASE && "caller tensor reaches into the virtual heap window");
-            LOG_DEBUG("  ChipTensor %d: child memory, pass-through (0x%" PRIx64 ")", i, t.buffer.addr);
+            const uint64_t host_view_addr = orch_args->host_view(i);
+            if (host_view_addr != 0) {
+                if (signature == nullptr || i >= sig_count || signature[i] != ArgDirection::IN) {
+                    LOG_ERROR("host-orch: tensor %d has a host view but is not a read-only input", i);
+                    return PTO_RUNTIME_ERR_INTERNAL;
+                }
+                const size_t size = static_cast<size_t>(t.buffer.size);
+                void *host_view = reinterpret_cast<void *>(static_cast<uintptr_t>(host_view_addr));
+                if (!tensor_access.add(t.buffer.addr, size, host_view)) {
+                    LOG_ERROR(
+                        "host-orch: cannot attach host view for device tensor %d (addr 0x%" PRIx64 ", %zu bytes)", i,
+                        t.buffer.addr, size
+                    );
+                    return PTO_RUNTIME_ERR_INTERNAL;
+                }
+            }
+            LOG_DEBUG("  ChipTensor %d: device memory, pass-through (0x%" PRIx64 ")", i, t.buffer.addr);
             device_args.add_tensor(t);
             continue;
         }

src/a5/runtime/host_build_graph/host/runtime_maker.cpp

@@ -1118,7 +1118,23 @@ extern "C" int bind_callable_to_runtime_impl(
         // as if it were a graph-heap allocation.
         if (t.is_device_memory()) {
             always_assert(t.buffer.addr < HEAP_VIRTUAL_BASE && "caller tensor reaches into the virtual heap window");
-            LOG_DEBUG("  ChipTensor %d: child memory, pass-through (0x%" PRIx64 ")", i, t.buffer.addr);
+            const uint64_t host_view_addr = orch_args->host_view(i);
+            if (host_view_addr != 0) {
+                if (signature == nullptr || i >= sig_count || signature[i] != ArgDirection::IN) {
+                    LOG_ERROR("host-orch: tensor %d has a host view but is not a read-only input", i);
+                    return PTO_RUNTIME_ERR_INTERNAL;
+                }
+                const size_t size = static_cast<size_t>(t.buffer.size);
+                void *host_view = reinterpret_cast<void *>(static_cast<uintptr_t>(host_view_addr));
+                if (!tensor_access.add(t.buffer.addr, size, host_view)) {
+                    LOG_ERROR(
+                        "host-orch: cannot attach host view for device tensor %d (addr 0x%" PRIx64 ", %zu bytes)", i,
+                        t.buffer.addr, size
+                    );
+                    return PTO_RUNTIME_ERR_INTERNAL;
+                }
+            }
+            LOG_DEBUG("  ChipTensor %d: device memory, pass-through (0x%" PRIx64 ")", i, t.buffer.addr);
             device_args.add_tensor(t);
             continue;
         }

src/common/task_interface/task_args.h

@@ -193,5 +193,57 @@ struct TaskArgsTpl<T, S, 0, 0, TensorTag> : TensorTagMixin<TensorTag, 0> {
 
 // L2 runtime ABI: fixed POD matching runtime.so byte-for-byte, and the sole ChipTensor-typed args
 // container — the materialized form a chip child decodes the L3->L2 Tensor blob into, just before
-// simpler_run.
-using ChipStorageTaskArgs = TaskArgsTpl<ChipTensor, uint64_t, CHIP_MAX_TENSOR_ARGS, CHIP_MAX_SCALAR_ARGS>;
+// simpler_run. host_views_ is host-process-only metadata: a nonzero entry gives host orchestration
+// a readable view of a device-memory tensor without changing the ChipTensor device ABI.
+struct ChipStorageTaskArgs {
+    ChipTensor tensors_[CHIP_MAX_TENSOR_ARGS];
+    uint64_t scalars_[CHIP_MAX_SCALAR_ARGS];
+    int32_t tensor_count_{0};
+    int32_t scalar_count_{0};
+    uint64_t host_views_[CHIP_MAX_TENSOR_ARGS]{};
+
+    void add_tensor(const ChipTensor &t) {
+        if (scalar_count_ > 0) throw std::logic_error("TaskArgs: cannot add tensor after scalar");
+        if (tensor_count_ >= CHIP_MAX_TENSOR_ARGS) throw std::out_of_range("TaskArgs: tensor capacity exceeded");
+        host_views_[tensor_count_] = 0;
+        tensors_[tensor_count_++] = t;
+    }
+
+    void add_scalar(uint64_t s) {
+        if (scalar_count_ >= CHIP_MAX_SCALAR_ARGS) throw std::out_of_range("TaskArgs: scalar capacity exceeded");
+        scalars_[scalar_count_++] = s;
+    }
+
+    const ChipTensor &tensor(int32_t i) const { return tensors_[i]; }
+    ChipTensor &tensor(int32_t i) { return tensors_[i]; }
+
+    uint64_t scalar(int32_t i) const { return scalars_[i]; }
+    uint64_t &scalar(int32_t i) { return scalars_[i]; }
+
+    const uint64_t *scalars() const { return scalars_; }
+    const ChipTensor *tensor_data() const { return tensors_; }
+    const uint64_t *scalar_data() const { return scalars_; }
+
+    int32_t tensor_count() const { return tensor_count_; }
+    int32_t scalar_count() const { return scalar_count_; }
+
+    void set_host_view(int32_t i, uint64_t addr) {
+        if (i < 0 || i >= tensor_count_) throw std::out_of_range("TaskArgs: host-view index out of range");
+        host_views_[i] = addr;
+    }
+
+    uint64_t host_view(int32_t i) const {
+        if (i < 0 || i >= tensor_count_) throw std::out_of_range("TaskArgs: host-view index out of range");
+        return host_views_[i];
+    }
+
+    void clear() {
+        tensor_count_ = 0;
+        scalar_count_ = 0;
+    }
+};
+
+static_assert(
+    std::is_trivially_copyable_v<ChipStorageTaskArgs> && std::is_standard_layout_v<ChipStorageTaskArgs>,
+    "ChipStorageTaskArgs crosses the runtime.so ABI as raw bytes"
+);