From ab85729af6ac425d05b1e264e756c5453f8f6b6b Mon Sep 17 00:00:00 2001 From: YunjiQin Date: Sun, 9 Aug 2026 20:05:00 -0700 Subject: [PATCH] Change: decode a chip child's mailbox blob once per dispatch Both chip-child dispatch paths read the same mailbox bytes twice: once through `ImportRegistry.materialize_blob` to reach each tensor's descriptor and build the identity -> local-base map, then again through `materialize_tensor_blob` to produce the `ChipStorageTaskArgs` the runtime ABI reads. Each element therefore paid two 144 B copies and two full `validate_tensor` passes, and at the 256-arg ceiling that is 36 KiB of copying plus 256 validations per dispatch. The bytes now decode once, into the wire `TaskArgs`, and both the mapping pass and the POD build read that object -- the three-step shape the L2 leaf already uses. `read_args_from_blob` and `materialize_task_args` both predate this change, so no binding grows a parameter; `materialize_tensor_blob` is deleted and `materialize_blob` collapses into `materialize_args`, whose body it duplicated except for the leading decode. Validation coverage is unchanged: `TaskArgsView::tensors` is the sole gate on a blob element and the surviving decode runs it over the same elements. `materialize_task_args` is now the only path to the chip POD, so the `resolved` map's shape is documented on it rather than by reference to the deleted function. --- python/bindings/task_interface.cpp | 33 +++++-------------- python/simpler/buffer.py | 14 ++------ python/simpler/worker.py | 28 ++++++++-------- .../test_l3_tensor_dispatch.py | 6 ++-- 4 files changed, 28 insertions(+), 53 deletions(-) diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index c26230f724..10b4d7a2b6 100644 --- a/python/bindings/task_interface.cpp +++ b/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(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( diff --git a/python/simpler/buffer.py b/python/simpler/buffer.py index 8068b3091b..5466214f24 100644 --- a/python/simpler/buffer.py +++ b/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 diff --git a/python/simpler/worker.py b/python/simpler/worker.py index d381278acb..7c393dc746 100644 --- a/python/simpler/worker.py +++ b/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, diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_tensor_dispatch.py b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_tensor_dispatch.py index d4f358a261..7406a321f5 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_tensor_dispatch.py +++ b/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