Change: decode a chip child's mailbox blob once per dispatch

python/bindings/task_interface.cpp

@@ -2224,7 +2224,7 @@ NB_MODULE(_task_interface, m) {
             nb::arg("callable_id"), nb::arg("args"), nb::arg("config"), nb::arg("accepted_state_addr") = 0,
             nb::arg("accepted_value") = 0, nb::arg("pipeline_slot") = 0, nb::arg("pipeline_generation") = 0,
             "Launch a callable_id from the runtime.so-ABI POD a chip-child mailbox loop built with "
-            "materialize_tensor_blob, so no Python code re-implements the tensor/scalar layout."
+            "materialize_task_args, so no Python code re-implements the tensor/scalar layout."
         )
         .def(
             "unregister_callable",
@@ -2378,28 +2378,6 @@ NB_MODULE(_task_interface, m) {
 
     // --- Standalone blob helpers ---
 
-    m.def(
-        "materialize_tensor_blob",
-        [](uint64_t blob_ptr, size_t capacity, nb::dict resolved) -> ChipStorageTaskArgs {
-            TaskArgsView view = read_blob(reinterpret_cast<const uint8_t *>(blob_ptr), capacity);
-            ChipStorageTaskArgs args;
-            for (int32_t i = 0; i < view.tensor_count; i++) {
-                args.add_tensor(materialize_one(view.tensors(i), resolved));
-            }
-            for (int32_t i = 0; i < view.scalar_count; i++) {
-                args.add_scalar(view.scalars[i]);
-            }
-            return args;
-        },
-        nb::arg("blob_ptr"), nb::arg("capacity"), nb::arg("resolved"),
-        "Materialize a Tensor blob into the runtime.so-ABI ChipStorageTaskArgs POD. Each tensor's "
-        "embedded buffer identity is resolved via `resolved` {CanonicalIdentity: (local_base, "
-        "address_space)}; addr = base + byte_offset. The caller pre-populates `resolved` by "
-        "materializing each embedded descriptor (see read_args_from_blob) on first receipt. "
-        "Strided views (transpose / permute / step-slice) materialize to strided ChipTensors. "
-        "Rejects an unknown identity and a non-dtype-aligned byte_offset."
-    );
-
     m.def(
         "materialize_task_args",
         [](const TaskArgs &args, nb::dict resolved) -> ChipStorageTaskArgs {
@@ -2414,8 +2392,13 @@ NB_MODULE(_task_interface, m) {
         },
         nb::arg("args"), nb::arg("resolved"),
         "Materialize a TaskArgs held in this process into the runtime.so-ABI ChipStorageTaskArgs "
-        "POD. An L2 leaf consumes its own args, so it takes this path instead of the mailbox blob; "
-        "`resolved` has the same shape as for materialize_tensor_blob."
+        "POD — the sole path to that POD, whether the args are an L2 leaf's own or a chip child's "
+        "read back from its mailbox with read_args_from_blob. Each tensor's embedded buffer "
+        "identity is resolved via `resolved` {CanonicalIdentity: (local_base, address_space)}; "
+        "addr = base + byte_offset. The caller pre-populates `resolved` by materializing each "
+        "embedded descriptor on first receipt. Strided views (transpose / permute / step-slice) "
+        "materialize to strided ChipTensors. Rejects an unknown identity and a non-dtype-aligned "
+        "byte_offset."
     );
 
     m.def(

python/simpler/buffer.py

@@ -634,19 +634,9 @@ def materialize(self, desc: BufferDescriptor) -> ImportedBuffer:
         self._by_identity[key] = imported
         return imported
 
-    def materialize_blob(self, blob_ptr: int, capacity: int) -> dict[CanonicalIdentity, tuple[int, int]]:
-        """Materialize every embedded descriptor in a task-args blob and return the resolved map:
-        identity -> (local base, address_space), scoped to this call's own tensors."""
-        args = read_args_from_blob(blob_ptr, capacity)
-        resolved: dict[CanonicalIdentity, tuple[int, int]] = {}
-        for i in range(args.tensor_count()):
-            desc = args.tensor(i).buffer
-            imported = self.materialize(desc)
-            resolved[desc.identity] = (imported.base, int(imported.address_space))
-        return resolved
-
     def materialize_args(self, args) -> dict[CanonicalIdentity, tuple[int, int]]:
-        """The same, for a ``TaskArgs`` already held in this process (the L2-leaf path)."""
+        """Materialize every embedded descriptor in a ``TaskArgs`` and return the resolved map:
+        identity -> (local base, address_space), scoped to this call's own tensors."""
         resolved: dict[CanonicalIdentity, tuple[int, int]] = {}
         for i in range(args.tensor_count()):
             desc = args.tensor(i).buffer

python/simpler/worker.py

@@ -98,7 +98,6 @@ def my_l4_orch(orch, args, config):
     _worker_host_mapped_region_peek_cleanup_error,
     get_element_size,
     materialize_task_args,
-    materialize_tensor_blob,
     read_args_from_blob,
 )
 
@@ -2501,13 +2500,14 @@ def handle_task(task_buf) -> tuple[int, str]:
             )
             if pipeline_reserved != 0:
                 raise RuntimeError(f"chip_process dev={device_id}: invalid pipeline lease reserved field")
-            # Materialize the tensor args into a chip-POD blob the runtime reads: resolve
-            # each ref's embedded handle to a local base (map-once, cached by canonical
-            # identity), then build the chip blob at those bases. Replaces the former
-            # parent-VA range rewrite — identities resolve exactly, not by numeric range.
+            # The mailbox bytes decode once, into the wire TaskArgs; the mapping pass and the
+            # chip-POD build both read that object. Each tensor's embedded handle resolves to a
+            # local base by canonical identity (map-once, cached), and the POD is built at those
+            # bases — an exact resolution, not the parent-VA numeric-range rewrite it replaced.
             args_ptr = task_addr + _OFF_TASK_ARGS_BLOB
-            resolved = import_registry.materialize_blob(args_ptr, _MAILBOX_ARGS_CAPACITY)
-            chip_args = materialize_tensor_blob(args_ptr, _MAILBOX_ARGS_CAPACITY, resolved)
+            args = read_args_from_blob(args_ptr, _MAILBOX_ARGS_CAPACITY)
+            resolved = import_registry.materialize_args(args)
+            chip_args = materialize_task_args(args, resolved)
             # The acceptance flag lives in the mailbox, not in the materialized args, so
             # the fence still publishes through the address the parent polls.
             cw._impl.run_materialized(
@@ -2759,13 +2759,15 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None:
 
         def submit_frame(frame: _StagedFrame) -> None:
             _protocol, run_id, slot_id, generation, dispatch_id = frame.identity
-            # The frame carries the wire blob; the runtime reads the chip POD. Resolve each
-            # tensor's descriptor to a local base (map-once, cached by canonical identity) and
-            # rebuild at those bases, as the non-pipelined task path does. The lane copies the
-            # args into its own storage, so this POD need not outlive the call.
+            # The frame carries the wire blob; the runtime reads the chip POD. The bytes decode
+            # once into the wire TaskArgs, whose tensors resolve to local bases (map-once, cached
+            # by canonical identity) and rebuild at those bases, as the non-pipelined task path
+            # does. The lane copies the args into its own storage, so this POD need not outlive
+            # the call.
             args_ptr = frame.frame_addr + _OFF_TASK_ARGS_BLOB
-            resolved = import_registry.materialize_blob(args_ptr, _MAILBOX_ARGS_CAPACITY)
-            chip_args = materialize_tensor_blob(args_ptr, _MAILBOX_ARGS_CAPACITY, resolved)
+            args = read_args_from_blob(args_ptr, _MAILBOX_ARGS_CAPACITY)
+            resolved = import_registry.materialize_args(args)
+            chip_args = materialize_task_args(args, resolved)
             frame.chip_run = cw._impl._submit_chip_run_materialized(
                 frame.cid,
                 chip_args,

tests/st/a2a3/tensormap_and_ringbuffer/test_l3_tensor_dispatch.py

@@ -11,9 +11,9 @@
 
 The L3 orch is a pure DAG builder over views: it names task args as Tensors built from handles
 (create_buffer + handle.ref), never touching data. torch is used only OUTSIDE run() to fill inputs
-and read the output. The owner writes a Tensor blob to the chip mailbox; the chip child
-materializes it back to Tensors (ImportRegistry -> materialize_tensor_blob) and runs the kernel;
-the result lands in the shared output buffer with no per-run copy.
+and read the output. The owner writes a Tensor blob to the chip mailbox; the chip child materializes
+it back to Tensors (read_args_from_blob -> ImportRegistry -> materialize_task_args) and runs the
+kernel; the result lands in the shared output buffer with no per-run copy.
 """
 
 import torch