From 57a89fda16b3b063fd27c96b3b4102b3e154deaa Mon Sep 17 00:00:00 2001 From: YunjiQin Date: Thu, 30 Jul 2026 04:11:51 -0700 Subject: [PATCH 1/2] Add: the BufferHandle/Tensor wire ABI, its codec, and create_buffer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit At L3 and above a task argument is currently a raw pointer plus a `child_memory` bool, which cannot say which process the address is valid in, which backing it belongs to, or whether the holder may write it. This adds the typed ABI that replaces it, with nothing consuming it yet. Three wire types in one header, each pinned by static_assert on size and every field offset, so a layout drift fails the build rather than the run: - `CanonicalIdentity` (32 B) — owner_instance_id + buffer_id + generation. The key an owner registry and every consumer import cache share. Fixed-length with no length field, so hashing and comparison cannot read past it, and padding is excluded from equality so two decodes of one backing key alike. - `BufferHandleDescriptor` (88 B) — the self-describing backing: address space, access, backend kind, extent, and a length-delimited backend body. - `BufferRef` (144 B) — the blob element: the descriptor embedded whole plus a strided view. No materialized address; a consumer resolves it on receipt. `validate_buffer_ref` is the single validator every trust boundary runs. It bounds `body_len` and `ndims`, checks the address_space × backend matrix, requires strides > 0 and a known dtype, and rejects a view that extends past its backing. The blob codec lives beside the types it encodes, and its reader validates each element as it is extracted, so a decoder cannot hand out an element nobody checked. `Worker.create_buffer` allocates a POSIX-shm backing under a fresh identity; the Worker owns every handle it created and unlinks the backings on close(). `ImportRegistry` is the consumer half — map-once, keyed by identity. Also bounds-checks `get_element_size`, which indexed its table with a raw u8 that now reaches it from the wire, and moves `MAX_TENSOR_DIMS` to `data_type.h` so the view types can share it with `Tensor`. The ABI is frozen once it ships, which is why it lands on its own: the byte layout, the enum values, and the validator are what this change is asking to have reviewed. Naming a `Tensor` and submitting it, the per-consumer materialization split, and the submit-time guards follow separately. --- docs/README.md | 1 + docs/buffer-handle-abi.md | 266 ++++++++ mkdocs.yml | 1 + python/bindings/task_interface.cpp | 62 ++ python/simpler/buffer_handle.py | 760 ++++++++++++++++++++++ python/simpler/worker.py | 74 +++ src/common/task_interface/buffer_handle.h | 403 ++++++++++++ src/common/task_interface/data_type.h | 15 +- src/common/task_interface/tensor.h | 2 - tests/ut/cpp/CMakeLists.txt | 1 + tests/ut/cpp/types/test_buffer_handle.cpp | 239 +++++++ tests/ut/py/test_buffer_handle.py | 397 +++++++++++ 12 files changed, 2218 insertions(+), 3 deletions(-) create mode 100644 docs/buffer-handle-abi.md create mode 100644 python/simpler/buffer_handle.py create mode 100644 src/common/task_interface/buffer_handle.h create mode 100644 tests/ut/cpp/types/test_buffer_handle.cpp create mode 100644 tests/ut/py/test_buffer_handle.py diff --git a/docs/README.md b/docs/README.md index 47b996a3d5..51e018d0df 100644 --- a/docs/README.md +++ b/docs/README.md @@ -40,6 +40,7 @@ changing simpler's own internals. | [Chip-Level Architecture (L2)](chip-level-arch.md) | Three-program model (host / AICPU / AICore), API layers, handshake | | [Hierarchical Level Runtime](hierarchical-level-runtime.md) | The L0–L6 level model and component composition | | [Task Flow](task-flow.md) | Callable / TaskArgs / CallConfig pass-through, `IWorker` | +| [BufferHandle Memory Model](buffer-handle-abi.md) | How L3+ tasks name data: canonical identity, backend descriptor, strided view | | [Orchestrator](orchestrator.md) | DAG submission: TensorMap, Scope, Ring, task state machine | | [Scheduler](scheduler.md) | DAG dispatch: wiring / ready / completion queues, dispatch loop | | [Worker Manager](worker-manager.md) | Worker pool, THREAD/PROCESS modes, fork + mailbox mechanics | diff --git a/docs/buffer-handle-abi.md b/docs/buffer-handle-abi.md new file mode 100644 index 0000000000..2e5df6df10 --- /dev/null +++ b/docs/buffer-handle-abi.md @@ -0,0 +1,266 @@ +# BufferHandle / Tensor — the L3+ memory model + +At L3 and above, tasks name their data with typed, self-describing **buffer +handles and views**, not raw pointers. This replaces the legacy "raw pointer + +`child_memory` bool" mechanism with an ABI that carries a canonical identity, a +backend descriptor, and a strided view — so a buffer can be resolved exactly +across the L3→L2 (and L4→L3) boundaries without a side table. + +This page is the user-facing how-to. The byte layout itself is pinned by the +`static_assert`s in [`src/common/task_interface/buffer_handle.h`](../src/common/task_interface/buffer_handle.h) +— sizes, field offsets, and enum values all fail the build if they drift, and +[`tests/ut/cpp/types/test_buffer_handle.cpp`](../tests/ut/cpp/types/test_buffer_handle.cpp) +pins them again from the outside along with the blob codec. + +## Three types + +| Type | What it is | Where it lives | +| ---- | ---------- | -------------- | +| **`BufferHandle`** | An owned backing (POSIX shm / fork-COW / device malloc) with a canonical identity + lifecycle. Stays with the Worker that created it. | owner side (L3+) | +| **`Tensor`** | A **self-describing task argument**: the full handle descriptor embedded + a strided view `(byte_offset, shapes, strides, dtype)`. The wire element of `TaskArgs`. Carries no materialized address. | what you build and submit | +| **`ChipTensor`** | The materialized POD the device runtime ABI reads (address + strided view). Exists **only** at the L2 device-runtime boundary. | L2 leaf, internal | + +**You never build a `ChipTensor`.** Name a `Tensor` over a handle and submit it; +the address is resolved on the consuming endpoint, and the C++ orchestration on +the chip receives the resolved form. + +## Why `Tensor` and `ChipTensor` are two types + +The two sit on opposite sides of one resolution step. A `Tensor` *cannot* carry an +address: at submit time none exists — a `POSIX_SHM` backing maps to a different VA +in every process, and a `DEVICE_MALLOC` one is only valid on its owner chip. A +`ChipTensor` *must* carry one; it is what the kernel dereferences. Two ways to +collapse that into a single type were considered and dropped. + +### Rejected: merge `Tensor` into `ChipTensor` + +Drop `buffer.addr`, add the handle descriptor, and have the H2D staging step +rewrite the backend tag and body (and mint a fresh identity for the device copy). +That is self-consistent, but it charges the device for host-side fields: + +- **Device cost.** `sizeof(ChipTensor)` is pinned at **128 B** by + `static_assert(… == 128, "Tensor must be 2 cache lines")` and by + `PTO2_TASKPAYLOAD_TENSOR_STRIDE`, which the AICPU scheduler strides the payload + with. A merged struct is ~192–216 B, taking `PTO2TaskPayload` from 4864 B to + 6912–7680 B — **+42% to +58% per task slot** — because the payload embeds + `ChipTensor tensors[MAX_TENSOR_ARGS]` **by value** in the device task ring. +- **Zero return for that cost.** Most of what a merge would add has no reader on + the far side. The field-by-field split is below. + +#### What each type carries, and what the other has no use for + +Only the middle block means the same thing in both — the geometry of the view. +Everything above it answers *"where does this live?"*, everything below it is +*"what has the L2 runtime decided about it?"*, and neither question is open on +the other side of materialization. + +| `Tensor` (144 B, wire) | rel | `ChipTensor` (128 B, device) | +| ---------------------- | --- | ---------------------------- | +| `handle.magic` | ⟂ | — | +| `handle.identity` (32 B) | ⟂ | — | +| `handle.backend_kind` | ⟂ | — | +| `handle.body[32]` + `body_len` | ⟂ | — | +| `handle.nbytes` | ≈ | `buffer.size` | +| `handle.access` | ⟂ | — | +| `handle.owner_worker_path_id` | ⟂ | — | +| — | ⟂ | `buffer.addr` | +| `byte_offset` (bytes) | ≈ | `start_offset` (elements) | +| `shapes[5]` / `strides[5]` / `ndims` / `dtype` | = | `shapes[5]` / `strides[5]` / `ndims` / `dtype` | +| `handle.address_space` | = | `address_space` (still spelled `child_memory` until the wire flip) | +| — | ⟂ | `owner_task_id` | +| — | ⟂ | `version` | +| — | ⟂ | `manual_dep` | +| — | ⟂ | `is_contiguous`, `extent_elem_cache` | + +**Dead on the device** (`Tensor`-only): `magic` discriminates untrusted bytes at +a decode boundary the device does not have. `identity` / `backend_kind` / `body` +/ `body_len` are the *recipe* for finding the backing — spent once at +materialization, never consulted again. `access` has no device enforcement point +(it is checked at submit). `owner_worker_path_id` is diagnostic. + +**Meaningless before materialization** (`ChipTensor`-only): `buffer.addr` is the +resolved address, which by definition does not exist while the argument is still +crossing processes. `owner_task_id` / `version` / `manual_dep` are L2 OverlapMap +state — the producing task and its dependency treatment are decided by the chip +runtime, so an L3 builder has nothing to put there (`make_tensor_strided` fills +`PTO2TaskId::invalid()`, `0`, `false`). `is_contiguous` and `extent_elem_cache` +are derived from `shapes`/`strides`, cached for the AICore hot path so it can +skip cache line 2. + +So a merged struct would carry ~70 B that the AICore never reads, in a type whose +size is pinned at two cache lines precisely because the scheduler walks it per +task. + +> The one row with a plausible future device reader is `identity`: keying the L2 +> OverlapMap by it rather than by `buffer.addr` would make two views of one +> backing bucket together by construction. That needs 32 B — which fits the +> existing `_pad_cl2[36]` at `sizeof == 128`, i.e. **without** merging anything +> else. If it is ever done, the H2D staging step must mint a *new* identity for +> each staged copy, because the device buffer is a distinct backing from the host +> one it was copied from. + +### Rejected: keep the wire type transport-only, use `ChipTensor` in the L3 orch + +This keeps one user-visible type but pays a conversion at every hop: the orch +builds a `ChipTensor`, the mailbox needs the self-describing form, and the +consumer needs a `ChipTensor` again. Each conversion has to **rewrite the +address** — and the sender cannot know the receiver's, so the receiver ends up +guessing. That is the pre-P1-B mechanism: `_rewrite_blob_host_addrs` rewrote +addresses by **numeric range** and mis-rewrote device pointers that happened to +fall inside a registered host range, patched by adding a `child_memory` skip +whose own comment records the hazard. + +### Shipped instead: same name, same interface + +The Python type is named `Tensor` and carries the C++ type's vocabulary — +`shapes` / `strides` / `dtype` / `ndims`, with `handle.tensor(...)` and +`args.add_tensor(...)` mirroring `Tensor.make(...)` / `add_tensor(...)`. The +device POD is `ChipTensor` throughout this page; its Python binding still +carries the older name `Tensor` and is renamed with the dispatch wire. + +**So the split costs the user nothing to know.** You name a `Tensor`, submit it, +and the chip's C++ orchestration receives it resolved; the type name does not +even appear in your code — you write `handle.tensor(shapes, dtype)` and +`args.add_tensor(t, tag)`. Resolving the address in between is the framework's +job, and keeping the two forms as separate types is what makes that boundary a +type change rather than a silently wrong address. + +## Allocating buffers + +```python +h = worker.create_buffer(nbytes) # kind3: explicit shared host buffer (POSIX shm) +h = worker.alloc_shared_tensor((M, N), dtype) # kind3: shape-sized create_buffer +# inside an orch fn, for device (chip-private) memory: +d = orch.alloc_child_tensor(worker, (M, N), dtype) # kind4: DEVICE_MALLOC on that chip +``` + +`create_buffer` returns a `BufferHandle` backed by a POSIX shm the consumer maps +lazily on first receipt of a tensor over it (map-once, by identity); +`alloc_shared_tensor` returns a runtime-managed intermediate over a FORK_SHM ring +VA. `alloc_child_tensor` allocates device memory on a specific next-level worker +and wraps the pointer; its `.base` is the device pointer (the `orch.copy_to` +destination), and a tensor over it must be dispatched only to that worker. + +## Naming a view + +`handle.tensor(...)` names a view over the backing: + +```python +v = h.tensor(shapes=(M, N), dtype) # contiguous (row-major strides) +v = h.tensor(shapes=(N, M), dtype, strides=(1, M)) # transposed +v = h.tensor(shapes=(M, K), dtype, byte_offset=off) # sub-region +``` + +`strides` are **element** strides and are strictly > 0 — broadcast (stride 0) and +negative step are unsupported, and a singleton dimension's stride is never +normalized away. `byte_offset` is a **byte** offset and must be a multiple of the +dtype size. + +## Submitting a task + +```python +ta = TaskArgs() +ta.add_tensor(a_h.tensor((SIZE,), DataType.FLOAT32), TensorArgType.INPUT) +ta.add_tensor(out_h.tensor((SIZE,), DataType.FLOAT32), TensorArgType.OUTPUT_EXISTING) +orch.submit_next_level(chip_handle, ta, cfg, worker=0) +``` + +`TaskArgs` carries `Tensor`s. Tags drive dependency inference, which keys on +the **canonical identity** (buffer granularity, the successor of the former +buffer-address key); byte-range overlap between same-buffer views is refined by +the L2 OverlapMap on the materialized tensors, not by the L3 key. + +## Reading and writing data — torch only at the boundaries + +The orch fn is a pure DAG builder: computing on data there would be invisible to +dependency inference. So **torch is used only outside `run()`** (fill inputs, +read outputs) or **inside a Python sub-worker** (a compute leaf): + +```python +h = worker.create_buffer(n * 4) +torch.frombuffer(h.shm.buf, dtype=torch.float32, count=n).fill_(5.0) # before run() +worker.run(my_orch, ...) # orch names tensors only +result = torch.frombuffer(out_h.shm.buf, dtype=torch.float32, count=n) # after run() +``` + +## How a tensor reaches its consumer (three-way split) + +A `Tensor` on the wire is materialized differently by each consumer: + +| Consumer | What it does | +| -------- | ------------ | +| **Chip leaf (L2 runtime)** | Materialize each one to a `ChipTensor` (map-once, keyed by identity), including **strided** views; hand the POD blob to `run_from_blob`. | +| **Python sub-worker** (compute) | Map each one into a `MappedArg`; the callable computes with `torch.frombuffer(arg.buffer, ...)`. No `ChipTensor`. | +| **Nested L4→L3 orch** (forwarding) | **Re-export** each backing to a handle `H'` that keeps the source's canonical identity — no pass-through, no map on the forwarding hop. | + +**Re-export (no pass-through).** An upper-level tensor is forwarded on receipt as a +handle `H'` that keeps the source's **canonical identity unchanged** (invariant +across every edge — an L4 buffer forwarded L4→L3→L2 carries one identity at all +three layers; only role / materialized VA / view change), per-backing and without +mapping. A downstream compute leaf maps lazily, so pure forwarding carries no map +cost. Dependency inference keys on the invariant identity, so an alias / +retain-release does not split across layers. + +## Canonical identity + +Every backing carries a fixed-length 32-byte identity — an opaque per-incarnation +`owner_instance_id`, a `buffer_id` unique within that incarnation, and a `generation` that starts at +1 and increments whenever a `buffer_id` slot is reused (so a stale handle for a recycled slot is +rejected rather than silently resolving). It is **invariant across every edge**: an L4 buffer +forwarded L4→L3→L2 carries one identity at all three layers. + +Identity is what dependency inference and the map-once import cache key on — never a materialized +address, which means nothing different in another process. + +Nothing inside the identity bounds a read: it has no length field. Hashing and comparison are +therefore in-bounds for any bytes that arrive, structurally rather than by validation. The owning +worker's tree path is deliberately **not** part of it — a path is reused across restarts and +contributes no uniqueness, so it is interned to a diagnostic `owner_worker_path_id` whose table lives +only in the owning process. An id another process minted renders as ``; nothing routes, +gates, or keys on it. + +## Backends + +| Backend | Materializes to | Used for | +| ------- | --------------- | -------- | +| `POSIX_SHM` | a named shm mapped into the consumer | `create_buffer` / `alloc_shared_tensor` | +| `FORK_SHM` | the same VA, no map | a pre-fork `MAP_SHARED` host buffer (e.g. `share_memory_()`), writable from the child | +| `FORK_COW` | the same VA, no map | a pre-fork plain host buffer: copy-on-write, so **READ only** | +| `DEVICE_MALLOC` | the device pointer, no map (chip-local) | `alloc_child_tensor` | +| `REMOTE_SIDECAR` | (P2) resolved via the remote transport | an arg to a remote L3; the descriptor rides in the sidecar | + +## What a submit checks + +Naming an argument is not enough on its own — two things are verified where the +values are final, at submit: + +- **`access ⊆ granted`.** A tag may only request what the backing grants: `INPUT` + needs READ, `OUTPUT_EXISTING` needs WRITE, `INOUT` needs READWRITE. This is + re-checked at submit rather than trusted from `add_tensor`, because a tag stays + mutable afterwards. It is what stops a plain (copy-on-write) tensor from being + named as an output and then silently losing every write in the child. +- **No overlapping writes within one task.** Two arguments of one task that name + intersecting bytes of the same backing are rejected: they belong to one node, + so there is no order between them to express, and a device-staged copy of a + host backing does not even alias on the device for the L2 overlap map to + notice. Disjoint slices of one buffer stay legal — that is what `byte_offset` + is for. + +Group members are not compared against each other: a group is one node, and +naming one buffer as every member's output is how it publishes a single +completion token for a downstream task to depend on. + +## Scope / status + +This page describes the memory model end to end. What the tree has today is the +ABI itself — the three types, the canonical identity, the blob codec — plus +`create_buffer`. **The dispatch wire is not connected yet**: `TaskArgs` still +carries the device POD, so `handle.tensor(...)` has no consumer, the other +allocators (`alloc_shared_tensor`, `alloc_child_tensor`) do not exist, and the +submit-time checks above are not reachable. Those land with the wire flip, which +is what every section from "Submitting a task" onward describes. + +Single-machine (host + device) L3→L2 and L4→L3→L2 dispatch is implemented and +verified in `a2a3sim` and onboard `a2a3` on that branch. The remote **receive** +side and the buffer lifecycle robustness (`release_buffer`, in-flight retain / +deferred-free) are later phases (P2). diff --git a/mkdocs.yml b/mkdocs.yml index ec313c64f7..9a4cebbae1 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -101,6 +101,7 @@ nav: - Chip-Level Architecture: chip-level-arch.md - Hierarchical Level Runtime: hierarchical-level-runtime.md - Task Flow: task-flow.md + - BufferHandle Memory Model: buffer-handle-abi.md - Orchestrator: orchestrator.md - Scheduler: scheduler.md - Worker Manager: worker-manager.md diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 4d588551ad..82296d5286 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -48,6 +48,7 @@ #include #include "arg_direction.h" +#include "buffer_handle.h" #include "callable.h" #include "callable_protocol.h" #include "chip_worker.h" @@ -617,6 +618,17 @@ NB_MODULE(_task_interface, m) { m.attr("TENSOR_STRIDE_BYTES") = static_cast(sizeof(Tensor)); m.attr("TENSOR_CHILD_MEMORY_OFFSET") = static_cast(offsetof(Tensor, child_memory)); + // BufferHandle / BufferRef wire ABI (buffer_handle.h). Exported so the Python mirror in + // simpler.buffer_handle can pin its struct formats to the C++ layout and reject drift. + m.attr("BUFFER_DESCRIPTOR_MAGIC") = static_cast(BUFFER_DESCRIPTOR_MAGIC); + m.attr("BUFFERREF_BLOB_MAGIC") = static_cast(BUFFERREF_BLOB_MAGIC); + m.attr("BUFFER_REF_BYTES") = static_cast(sizeof(BufferRef)); + m.attr("BUFFER_HANDLE_DESCRIPTOR_BYTES") = static_cast(sizeof(BufferHandleDescriptor)); + m.attr("CANONICAL_IDENTITY_BYTES") = static_cast(sizeof(CanonicalIdentity)); + m.attr("OWNER_INSTANCE_ID_BYTES") = static_cast(OWNER_INSTANCE_ID_BYTES); + m.attr("DESC_MAX_BYTES") = static_cast(DESC_MAX_BYTES); + m.attr("BUFFERREF_BLOB_HEADER_BYTES") = static_cast(BUFFERREF_BLOB_HEADER_SIZE); + // --- Tensor --- // The unified strided tensor descriptor. Constructed contiguous via make() // (row-major strides, start_offset == 0); see src/common/task_interface/tensor.h. @@ -1603,6 +1615,56 @@ NB_MODULE(_task_interface, m) { "Tags are not preserved (blob wire format strips them)." ); + // BufferRef blob readers. Each validates every element it extracts (read_bufferref_blob bounds + // the header, BufferRefBlobView::ref validates the element), so a Python consumer never walks + // the layout itself. + m.def( + "bufferref_blob_descriptors", + [](uint64_t blob_ptr, size_t capacity) -> nb::list { + const uint8_t *src = reinterpret_cast(blob_ptr); + BufferRefBlobView view = read_bufferref_blob(src, capacity); + nb::list out; + for (int32_t i = 0; i < view.ref_count; i++) { + BufferRef r = view.ref(i); + out.append(nb::bytes(reinterpret_cast(&r.handle), sizeof(BufferHandleDescriptor))); + } + return out; + }, + nb::arg("blob_ptr"), nb::arg("capacity"), + "Extract each BufferRef's embedded BufferHandleDescriptor (packed bytes) from a BufferRef " + "blob, in ref order. A consumer materializes these lazily on receipt." + ); + + m.def( + "bufferref_blob_refs", + [](uint64_t blob_ptr, size_t capacity) -> nb::list { + const uint8_t *src = reinterpret_cast(blob_ptr); + BufferRefBlobView view = read_bufferref_blob(src, capacity); + nb::list out; + for (int32_t i = 0; i < view.ref_count; i++) { + BufferRef r = view.ref(i); + out.append(nb::bytes(reinterpret_cast(&r), sizeof(BufferRef))); + } + return out; + }, + nb::arg("blob_ptr"), nb::arg("capacity"), + "Extract each full packed BufferRef (descriptor + view) from a BufferRef blob, in ref order." + ); + + m.def( + "bufferref_blob_scalars", + [](uint64_t blob_ptr, size_t capacity) -> nb::list { + const uint8_t *src = reinterpret_cast(blob_ptr); + BufferRefBlobView view = read_bufferref_blob(src, capacity); + nb::list out; + for (int32_t i = 0; i < view.scalar_count; i++) { + out.append(view.scalars[i]); + } + return out; + }, + nb::arg("blob_ptr"), nb::arg("capacity"), "Extract the scalar args (uint64) from a BufferRef blob, in order." + ); + nb::class_(m, "_L2ChildOnboardRegionExport") .def_ro("device_addr", &L2ChildOnboardRegionExport::device_addr) .def_ro("mapping_bytes", &L2ChildOnboardRegionExport::mapping_bytes) diff --git a/python/simpler/buffer_handle.py b/python/simpler/buffer_handle.py new file mode 100644 index 0000000000..b78a5d99e4 --- /dev/null +++ b/python/simpler/buffer_handle.py @@ -0,0 +1,760 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Owner/consumer side of the BufferHandle ABI (Python mirror of ``buffer_handle.h``). + +``CanonicalIdentity`` is fixed-length (opaque ``owner_instance_id`` + ``buffer_id`` + ``generation``) +so hashing and comparison cannot read past it whatever arrives on the wire. The owning worker's tree +path is **not** part of the identity: it is interned to a diagnostic ``owner_worker_path_id`` whose +side table lives only in the owning process. Struct formats mirror ``buffer_handle.h`` byte for byte; +their sizes are asserted at import against the constants the ``_task_interface`` binding exports, so +a layout drift fails loudly here. +""" + +from __future__ import annotations + +import ctypes +import enum +import os +import struct +from collections.abc import Sequence +from dataclasses import dataclass +from multiprocessing.shared_memory import SharedMemory +from typing import Any + +from _task_interface import ( # pyright: ignore[reportMissingImports] + BUFFER_DESCRIPTOR_MAGIC, + BUFFER_HANDLE_DESCRIPTOR_BYTES, + BUFFER_REF_BYTES, + BUFFERREF_BLOB_HEADER_BYTES, + BUFFERREF_BLOB_MAGIC, + CANONICAL_IDENTITY_BYTES, + DESC_MAX_BYTES, + MAX_TENSOR_DIMS, + OWNER_INSTANCE_ID_BYTES, + DataType, + bufferref_blob_descriptors, + bufferref_blob_refs, + bufferref_blob_scalars, +) + + +def _dtype_value(dtype: Any) -> int: + """The int wire value of a dtype given either a ``DataType`` enum or its int value. + + The nanobind ``DataType`` enum is not directly ``int()``-able, so callers historically passed + ``dtype.value``; accept either form here to remove that footgun. + """ + return int(dtype.value) if hasattr(dtype, "value") else int(dtype) + + +class AddressSpace(enum.IntEnum): + HOST = 0 + DEVICE = 1 + + +class AccessMode(enum.IntEnum): + READ = 0 + WRITE = 1 + READWRITE = 2 + + +class BackendKind(enum.IntEnum): + """How a consumer turns a descriptor into a local address. + + ``FORK_SHM`` and ``FORK_COW`` resolve identically — the body is a base VA the child inherited + across the fork — but their write semantics are opposite, so they are separate tags rather than + one tag plus a hint. A child's write to a MAP_SHARED page reaches the owner; a write to a + copy-on-write page splits it into a private copy the owner never sees, silently. ``FORK_COW`` + therefore grants READ only. + """ + + FORK_SHM = 0 + POSIX_SHM = 1 + VMM_WINDOW = 2 + REMOTE_SIDECAR = 3 + DEVICE_MALLOC = 4 + FORK_COW = 5 + + +# (address_space, backend_kind) capability gate (§4.1). Absent ⇒ rejected. The two REMOTE_SIDECAR rows +# are allowed at the matrix layer; its P1 blanket reject is ImportRegistry.materialize's job (folding it +# here would double the semantics). +_CAPABILITY_OK: frozenset[tuple[int, int]] = frozenset( + { + (AddressSpace.HOST, BackendKind.FORK_SHM), + (AddressSpace.HOST, BackendKind.FORK_COW), + (AddressSpace.HOST, BackendKind.POSIX_SHM), + (AddressSpace.DEVICE, BackendKind.VMM_WINDOW), + (AddressSpace.DEVICE, BackendKind.DEVICE_MALLOC), + (AddressSpace.HOST, BackendKind.REMOTE_SIDECAR), + (AddressSpace.DEVICE, BackendKind.REMOTE_SIDECAR), + } +) + + +# --------------------------------------------------------------------------- +# Wire struct formats — mirror buffer_handle.h; sizes pinned to the binding. +# --------------------------------------------------------------------------- +# +# CanonicalIdentity (32 B): owner_instance_id[8] opaque, buffer_id u64, generation u32, _pad[12]. +_CANONICAL_IDENTITY = struct.Struct(f"<{OWNER_INSTANCE_ID_BYTES}sQI12x") +# BufferHandleDescriptor (88 B) = prefix(8) + CanonicalIdentity(32) + suffix(48). +_DESC_PREFIX = struct.Struct("". Nothing routes, +# gates or keys on a path, so an unresolvable id is never an error. Id 0 means "no path". +_PATH_BY_ID: dict[int, str] = {0: ""} +_ID_BY_PATH: dict[str, int] = {"": 0} + + +def intern_worker_path(path: str) -> int: + """The diagnostic id for ``path`` in this process, assigning one on first sight.""" + pid = _ID_BY_PATH.get(path) + if pid is None: + pid = len(_ID_BY_PATH) + _ID_BY_PATH[path] = pid + _PATH_BY_ID[pid] = path + return pid + + +def worker_path_for_id(path_id: int) -> str: + """``path_id`` rendered for humans; an id minted in another process has no local text.""" + return _PATH_BY_ID.get(int(path_id), f"") + + +@dataclass(frozen=True) +class CanonicalIdentity: + """Globally-unique allocation identity; the key of every import registry. + + Fixed-length: no field bounds a read, so hashing and comparison are structurally in-bounds for any + bytes that arrive. ``owner_instance_id`` is an opaque nonce (bytewise-compared) and is the sole + source of cross-incarnation uniqueness. ``generation`` starts at 1 and increments on every + ``buffer_id`` slot reuse; 0 means uninitialized and is rejected on decode. + """ + + owner_instance_id: bytes + buffer_id: int + generation: int = 1 + + def __post_init__(self) -> None: + if len(self.owner_instance_id) != OWNER_INSTANCE_ID_BYTES: + raise ValueError( + f"owner_instance_id must be {OWNER_INSTANCE_ID_BYTES} bytes, got {len(self.owner_instance_id)}" + ) + + def pack(self) -> bytes: + return _CANONICAL_IDENTITY.pack(self.owner_instance_id, self.buffer_id, self.generation) + + @classmethod + def unpack(cls, raw: bytes) -> CanonicalIdentity: + owner_instance_id, buffer_id, generation = _CANONICAL_IDENTITY.unpack(raw) + if generation == 0: + raise ValueError("generation 0 is reserved (uninitialized) — a live identity starts at 1") + return cls(owner_instance_id, buffer_id, generation) + + +@dataclass(frozen=True) +class BufferHandleDescriptor: + """The self-describing handle payload — embedded whole in every ``Tensor`` built over the handle. + + ``body`` is the per-backend materialization (POSIX/fork shm name UTF-8, VMM handle bytes, ...). + """ + + identity: CanonicalIdentity + address_space: AddressSpace + access: AccessMode + backend_kind: BackendKind + nbytes: int + body: bytes = b"" + owner_worker_path_id: int = 0 + + def __post_init__(self) -> None: + # §4.1 capability gate: reject an unsupported address_space×backend_kind. Runs on every + # construction — owner-side (wrap_* → to_descriptor) and wire decode (unpack → cls(...)) — so a + # bad combo fails before dispatch and can never ride the wire. + if (self.address_space, self.backend_kind) not in _CAPABILITY_OK: + raise ValueError( + f"unsupported address_space×backend: " + f"{self.address_space.name}×{self.backend_kind.name} (§4.1 capability matrix)" + ) + if self.backend_kind == BackendKind.FORK_COW and self.access != AccessMode.READ: + raise ValueError( + f"FORK_COW grants READ only, got {self.access.name} — a child's write to a " + f"copy-on-write page never reaches the owner" + ) + + @property + def owner_worker_path(self) -> str: + """The owning worker's tree path, for diagnostics only; ```` when minted elsewhere.""" + return worker_path_for_id(self.owner_worker_path_id) + + def pack(self) -> bytes: + if len(self.body) > DESC_MAX_BYTES: + raise ValueError(f"backend body exceeds DESC_MAX_BYTES ({DESC_MAX_BYTES})") + prefix = _DESC_PREFIX.pack( + BUFFER_DESCRIPTOR_MAGIC, + int(self.address_space), + int(self.access), + int(self.backend_kind), + ) + suffix = _DESC_SUFFIX.pack(self.nbytes, self.owner_worker_path_id, len(self.body), self.body) + return prefix + self.identity.pack() + suffix + + @classmethod + def unpack(cls, raw: bytes) -> BufferHandleDescriptor: + if len(raw) < BUFFER_HANDLE_DESCRIPTOR_BYTES: + raise ValueError(f"descriptor too small: {len(raw)} < {BUFFER_HANDLE_DESCRIPTOR_BYTES}") + magic, address_space, access, backend_kind = _DESC_PREFIX.unpack_from(raw, 0) + if magic != BUFFER_DESCRIPTOR_MAGIC: + raise ValueError(f"not a BufferHandleDescriptor: magic {magic:#06x} != {BUFFER_DESCRIPTOR_MAGIC:#06x}") + identity = CanonicalIdentity.unpack(raw[_DESC_PREFIX.size : _DESC_PREFIX.size + _CANONICAL_IDENTITY.size]) + nbytes, path_id, body_len, body = _DESC_SUFFIX.unpack_from(raw, _DESC_PREFIX.size + _CANONICAL_IDENTITY.size) + if body_len > DESC_MAX_BYTES: + raise ValueError(f"body_len {body_len} exceeds DESC_MAX_BYTES ({DESC_MAX_BYTES})") + return cls( + identity=identity, + address_space=AddressSpace(address_space), + access=AccessMode(access), + backend_kind=BackendKind(backend_kind), + owner_worker_path_id=path_id, + nbytes=nbytes, + body=bytes(body[:body_len]), + ) + + +# BufferRef (272 B): BufferHandleDescriptor(216) + byte_offset u64, ndims u32, shapes[MAX] u32, +# strides[MAX] u32, dtype u8, _pad[3]. +_BUFFER_REF_TAIL = struct.Struct(f" tuple[int, ...]: + """Contiguous (row-major) element strides for ``shapes``: strides[i] = prod(shapes[i+1:]).""" + strides = [1] * len(shapes) + for i in range(len(shapes) - 2, -1, -1): + strides[i] = strides[i + 1] * shapes[i + 1] + return tuple(strides) + + +@dataclass(frozen=True) +class Tensor: + """A task argument: a strided view over a ``BufferHandle``, carrying that handle's descriptor. + + Self-describing — the consumer materializes ``handle`` on first receipt (no prior handshake), + keyed by ``handle.identity``. **Carries no materialized address**: the consuming endpoint + resolves the descriptor to a local base and adds ``byte_offset`` there. + + ``strides`` are ELEMENT strides and are strictly > 0 (broadcast / negative step unsupported); + they are carried explicitly and a singleton dimension's stride is never normalized away. + ``byte_offset`` is a BYTE offset, a multiple of the dtype size (checked at materialization). + """ + + handle: BufferHandleDescriptor + byte_offset: int + shapes: tuple[int, ...] + strides: tuple[int, ...] + dtype: int | DataType # normalized to the int wire value in __post_init__ + + def __post_init__(self) -> None: + if not 0 < len(self.shapes) <= MAX_TENSOR_DIMS: + raise ValueError(f"Tensor ndims must be in [1, {MAX_TENSOR_DIMS}], got {len(self.shapes)}") + if len(self.strides) != len(self.shapes): + raise ValueError("Tensor shapes and strides must have equal length") + object.__setattr__(self, "dtype", _dtype_value(self.dtype)) + + @property + def ndims(self) -> int: + return len(self.shapes) + + def pack(self) -> bytes: + ndims = len(self.shapes) + shapes = list(self.shapes) + [0] * (MAX_TENSOR_DIMS - ndims) + strides = list(self.strides) + [0] * (MAX_TENSOR_DIMS - ndims) + tail = _BUFFER_REF_TAIL.pack(self.byte_offset, ndims, *shapes, *strides, self.dtype) + return self.handle.pack() + tail + + @classmethod + def unpack(cls, raw: bytes) -> Tensor: + handle = BufferHandleDescriptor.unpack(raw[:BUFFER_HANDLE_DESCRIPTOR_BYTES]) + vals = _BUFFER_REF_TAIL.unpack(raw[BUFFER_HANDLE_DESCRIPTOR_BYTES:BUFFER_REF_BYTES]) + byte_offset, ndims = vals[0], vals[1] + shapes = tuple(vals[2 : 2 + ndims]) + strides = tuple(vals[2 + MAX_TENSOR_DIMS : 2 + MAX_TENSOR_DIMS + ndims]) + dtype = vals[2 + 2 * MAX_TENSOR_DIMS] + return cls(handle=handle, byte_offset=byte_offset, shapes=shapes, strides=strides, dtype=dtype) + + +def pack_bufferref_blob(tensors: list[Tensor], scalars: tuple[int, ...] = ()) -> bytes: + """Serialize tensors + scalars into the versioned wire blob (mirror of write_bufferref_blob).""" + header = _BUFFERREF_BLOB_HEADER.pack(BUFFERREF_BLOB_MAGIC, len(tensors), len(scalars), 0) + body = b"".join(t.pack() for t in tensors) + tail = struct.pack(f"<{len(scalars)}Q", *scalars) if scalars else b"" + return header + body + tail + + +def mint_owner_instance_id() -> bytes: + """A fresh opaque nonce, unique per owner incarnation (defends identity against ABA). + + Must stay a full-width random draw. It is the only thing separating two Workers' buffer_id spaces, + so a structured value (timestamp/pid) would hand the same identity to two Workers constructed in + one process within one second — a routine pattern (an L4 and its L3 built back to back). + """ + return os.urandom(OWNER_INSTANCE_ID_BYTES) + + +def _shm_base_addr(shm: SharedMemory) -> int: + """Mapped base address of ``shm``; valid until ``shm.close()``.""" + view = shm.buf + assert view is not None + exporter = ctypes.c_char.from_buffer(view) + addr = ctypes.addressof(exporter) + del exporter + return addr + + +@dataclass +class BufferHandle: + """Owner-side registry object for one shared backing; owns the POSIX shm that backs it.""" + + identity: CanonicalIdentity + address_space: AddressSpace + access: AccessMode + backend_kind: BackendKind + nbytes: int + body: bytes = b"" + owner_worker_path_id: int = 0 + shm: SharedMemory | None = None + base: int = 0 + # Owner-side only, never serialized into the descriptor: which next-level worker a DEVICE_MALLOC + # backing lives on (0 for a host backing or an L2 own-device malloc). The device-pointer provenance + # guard and free/copy key on (owner_worker_id, base). + owner_worker_id: int = 0 + + def to_descriptor(self) -> BufferHandleDescriptor: + return BufferHandleDescriptor( + identity=self.identity, + address_space=self.address_space, + owner_worker_path_id=self.owner_worker_path_id, + access=self.access, + backend_kind=self.backend_kind, + nbytes=self.nbytes, + body=self.body, + ) + + def tensor( + self, + shapes: tuple[int, ...], + dtype: int | DataType, + strides: tuple[int, ...] | None = None, + byte_offset: int = 0, + ) -> Tensor: + """A self-describing ``Tensor`` viewing this handle: embeds the full descriptor + the view. + + ``strides`` default to contiguous (row-major) — ``handle.tensor(shape, dtype)`` names the + whole buffer as a contiguous view; pass explicit element strides for a strided view. + ``byte_offset`` must be a multiple of the dtype size (checked at materialization). + ``dtype`` accepts a ``DataType`` enum or its int value. + """ + shapes = tuple(shapes) + strides = _row_major_strides(shapes) if strides is None else tuple(strides) + return Tensor( + handle=self.to_descriptor(), + byte_offset=byte_offset, + shapes=shapes, + strides=strides, + dtype=_dtype_value(dtype), + ) + + def close(self) -> None: + if self.shm is not None: + self.shm.close() + self.shm.unlink() + self.shm = None + + +def create_host_shared_buffer( + nbytes: int, + owner_instance_id: bytes, + buffer_id: int, + owner_worker_path: str = "", + generation: int = 1, + access: AccessMode = AccessMode.READWRITE, +) -> BufferHandle: + """Allocate a POSIX-shm host backing and wrap it as an owner ``BufferHandle`` (backend POSIX_SHM). + + The backend body is the shm name (UTF-8); the consumer maps it by name in ``ImportRegistry``. + """ + if nbytes <= 0: + raise ValueError(f"create_host_shared_buffer: nbytes must be positive, got {nbytes}") + shm = SharedMemory(create=True, size=nbytes) + identity = CanonicalIdentity(owner_instance_id, buffer_id, generation) + return BufferHandle( + identity=identity, + owner_worker_path_id=intern_worker_path(owner_worker_path), + address_space=AddressSpace.HOST, + access=access, + backend_kind=BackendKind.POSIX_SHM, + nbytes=nbytes, + body=shm.name.encode("utf-8"), + shm=shm, + base=_shm_base_addr(shm), + ) + + +def re_export(source: BufferHandleDescriptor) -> BufferHandle: + """Re-export a received handle descriptor for forwarding — identity **invariant**, no mapping. + + Canonical identity is invariant across every edge (frozen model §5/§8): the re-exported ``H'`` + keeps the SOURCE ``(owner_instance_id, buffer_id, generation)`` and the SAME + backing (backend_kind / body / nbytes / address_space / access) as ``source`` — an + L4-owned buffer forwarded L4→L3→L2 carries one identity at all three layers. Only the mapping is + stripped: ``base=0``, ``shm=None`` (no mmap on the forwarding hop); a downstream compute leaf + materializes lazily. Dependency inference keys on the (invariant) identity, so an alias / + retain-release does not split across layers. Re-export is per-backing (memoize by identity), so + pure forwarding carries no per-ref map cost. + """ + return BufferHandle( + identity=source.identity, + owner_worker_path_id=source.owner_worker_path_id, + address_space=source.address_space, + access=source.access, + backend_kind=source.backend_kind, + nbytes=source.nbytes, + body=source.body, + shm=None, + base=0, + ) + + +def remote_sidecar_tensor( + shapes: tuple[int, ...], + dtype: int, + nbytes: int, + owner_worker_id: int, + buffer_id: int, + generation: int, + address_space: AddressSpace, +) -> Tensor: + """Build a ``REMOTE_SIDECAR`` ``Tensor`` for a task arg destined for a remote worker. + + An arg passed L4→remote-L3 cannot be materialized from a local backing — the data lives on another + machine and travels via the remote transport. Its descriptor therefore carries ``backend_kind = + REMOTE_SIDECAR`` (a consumer decode-rejects a local materialize; the authoritative remote + descriptor rides in the per-task RemoteTaskArgsSidecar). The identity encodes the remote buffer + (``owner_worker_id`` folded into the opaque nonce, plus ``buffer_id`` / ``generation``) so + dependency inference and routing stay stable across the hop. + """ + oid = int(owner_worker_id).to_bytes(OWNER_INSTANCE_ID_BYTES, "little") + # A HOST_INLINE placeholder has no backing and so no generation of its own; 0 is the reserved + # "uninitialized" value a decoder rejects, so the placeholder carries the initial generation. + identity = CanonicalIdentity(oid, buffer_id, int(generation) or 1) + handle = BufferHandleDescriptor( + identity=identity, + owner_worker_path_id=intern_worker_path(f"remote/{owner_worker_id}"), + address_space=address_space, + access=AccessMode.READWRITE, + backend_kind=BackendKind.REMOTE_SIDECAR, + nbytes=nbytes, + body=b"", + ) + shapes = tuple(shapes) + return Tensor( + handle=handle, + byte_offset=0, + shapes=shapes, + strides=_row_major_strides(shapes), + dtype=int(dtype), + ) + + +def wrap_fork_inherited( + data_ptr: int, + nbytes: int, + owner_instance_id: bytes, + buffer_id: int, + owner_worker_path: str = "", + generation: int = 1, + access: AccessMode = AccessMode.READ, +) -> BufferHandle: + """Wrap a pre-fork, fork-inherited host allocation as a zero-copy ``BufferHandle``. + + Memory allocated before the children were forked is present in every child at the *same* virtual + address; the backend body is that base VA (u64 LE) and the consumer materializes to the same VA + with no mapping and no copy. The backend tag follows the mmap the caller actually has, which is + what ``access`` states: + + * ``MAP_SHARED`` (e.g. a ``torch.Tensor.share_memory_()``) — a child's writes land in the pages + the parent reads, so it can serve as an OUTPUT. Pass ``access=READWRITE``; tagged FORK_SHM. + * plain ``MAP_PRIVATE`` — copy-on-write: a child's first write splits the page into a private + copy the parent never sees. Read-only, the default; tagged FORK_COW so the distinction is a + classification rather than something a reader has to infer from ``access``. + """ + identity = CanonicalIdentity(owner_instance_id, buffer_id, generation) + backend = BackendKind.FORK_SHM if access != AccessMode.READ else BackendKind.FORK_COW + return BufferHandle( + identity=identity, + owner_worker_path_id=intern_worker_path(owner_worker_path), + address_space=AddressSpace.HOST, + access=access, + backend_kind=backend, + nbytes=nbytes, + body=int(data_ptr).to_bytes(8, "little"), + shm=None, + base=int(data_ptr), + ) + + +def host_ptr_nbytes(obj: Any) -> tuple[int, int]: + """Host address + byte length of a copy_to/copy_from buffer, without importing torch. + + A torch tensor is read via its ``data_ptr`` / ``numel`` / ``element_size`` (duck-typed); any other + object goes through the buffer protocol and must be writable so its backing address is stable for + the duration of the synchronous copy. + """ + if hasattr(obj, "data_ptr") and hasattr(obj, "numel") and hasattr(obj, "element_size"): + return int(obj.data_ptr()), int(obj.numel()) * int(obj.element_size()) + mv = memoryview(obj) + if mv.readonly: + raise TypeError("copy_to/copy_from host buffer must be a torch tensor or a writable buffer") + return ctypes.addressof((ctypes.c_char * mv.nbytes).from_buffer(obj)), mv.nbytes + + +def wrap_device_malloc( + device_ptr: int, + nbytes: int, + owner_instance_id: bytes, + buffer_id: int, + owner_worker_path: str = "", + generation: int = 1, + access: AccessMode = AccessMode.READWRITE, + owner_worker_id: int = 0, +) -> BufferHandle: + """Wrap a device pointer (from a worker device malloc) as a ``DEVICE_MALLOC`` ``BufferHandle``. + + The backend body is the device pointer (u64 LE); the consumer materializes to that pointer with no + mapping. The pointer is valid only on the chip that allocated it, so a ref over this handle must be + dispatched only to that chip (a topology invariant, as for the former ``child_memory`` tensor). + ``owner_worker_id`` records which next-level worker the backing lives on for free/copy provenance. + """ + identity = CanonicalIdentity(owner_instance_id, buffer_id, generation) + return BufferHandle( + identity=identity, + owner_worker_path_id=intern_worker_path(owner_worker_path), + address_space=AddressSpace.DEVICE, + access=access, + backend_kind=BackendKind.DEVICE_MALLOC, + nbytes=nbytes, + body=int(device_ptr).to_bytes(8, "little"), + shm=None, + base=int(device_ptr), + owner_worker_id=int(owner_worker_id), + ) + + +def wrap_vmm_window( + device_ptr: int, + nbytes: int, + owner_instance_id: bytes, + buffer_id: int, + owner_worker_path: str = "", + generation: int = 1, + access: AccessMode = AccessMode.READWRITE, + owner_worker_id: int = 0, +) -> BufferHandle: + """Wrap a domain-window-carved device VA as a ``VMM_WINDOW`` ``BufferHandle``. + + A comm domain's per-rank window is device memory carved by ``allocate_domain``; each named buffer + slice is one such backing. The backend body is the device VA (u64 LE); the consumer materializes to + that VA with no mapping. The VA is valid only on the chip that owns the window, so a ref over this + handle must be dispatched only to that chip (``owner_worker_id``). Unlike ``DEVICE_MALLOC`` it is + not freed by ``worker.free`` — the domain owns its lifetime and reclaims it at ``release_domain``. + """ + identity = CanonicalIdentity(owner_instance_id, buffer_id, generation) + return BufferHandle( + identity=identity, + owner_worker_path_id=intern_worker_path(owner_worker_path), + address_space=AddressSpace.DEVICE, + access=access, + backend_kind=BackendKind.VMM_WINDOW, + nbytes=nbytes, + body=int(device_ptr).to_bytes(8, "little"), + shm=None, + base=int(device_ptr), + owner_worker_id=int(owner_worker_id), + ) + + +@dataclass +class ImportedBuffer: + """A handle materialized into the consumer's address space: identity -> local base.""" + + identity: CanonicalIdentity + base: int + nbytes: int + address_space: AddressSpace = AddressSpace.HOST + shm: SharedMemory | None = None # the consumer's own mapping for shm backends + + +@dataclass +class MappedArg: + """A Python compute (sub-worker) task arg: a ``Tensor`` materialized into this process, exposing a + writable ``buffer`` at the view origin plus the view geometry. The callable computes with e.g. + ``torch.frombuffer(arg.buffer, dtype=, count=prod(arg.shapes))`` — reads/writes + land in the shared backing the owner sees. + """ + + imported: ImportedBuffer + byte_offset: int + shapes: tuple[int, ...] + strides: tuple[int, ...] + dtype: int # DataType value + + @property + def buffer(self) -> memoryview: + """A memoryview over the mapped backing at this view's origin (``byte_offset``).""" + ib = self.imported + if ib.shm is not None: + base = ib.shm.buf + assert base is not None + else: + # FORK_SHM (COW): no shm object — wrap the inherited VA range. + base = memoryview((ctypes.c_char * ib.nbytes).from_address(ib.base)) + return base[self.byte_offset :] + + +class MappedArgs(Sequence): + """A Python sub-worker's task args: the mapped tensor args plus the scalar args. + + Indexes and iterates as the tensor ``MappedArg`` list (``args[i].buffer``, ``len(args)``) — the + common compute-leaf access — and additionally exposes the blob's scalars via ``scalar_count()`` / + ``scalar(i)`` (uint64, in submission order), mirroring the owner-side ``TaskArgs`` scalar API. + """ + + __slots__ = ("_tensors", "_scalars") + + def __init__(self, tensors: list[MappedArg], scalars: tuple[int, ...]) -> None: + self._tensors = list(tensors) + self._scalars = tuple(int(s) for s in scalars) + + def __getitem__(self, i): + return self._tensors[i] + + def __len__(self) -> int: + return len(self._tensors) + + def tensor_count(self) -> int: + return len(self._tensors) + + def scalar_count(self) -> int: + return len(self._scalars) + + def scalar(self, i: int) -> int: + return self._scalars[i] + + +class ImportRegistry: + """Per-consumer-endpoint lazy import cache: materialize a ``Tensor``'s embedded descriptor to a + local base on first receipt (map-once), keyed by canonical identity. + + A consumer calls ``materialize`` for each ref's embedded descriptor as it arrives; the first + sight of an identity maps its backing into this process, later sights reuse the cached base + (a bumped generation is a distinct identity, materialized fresh). Keyed by the packed canonical + identity so lookups are exact — never a numeric-range guess. + """ + + def __init__(self) -> None: + self._by_identity: dict[bytes, ImportedBuffer] = {} + + def materialize(self, descriptor: BufferHandleDescriptor | bytes) -> ImportedBuffer: + """Map ``descriptor``'s backing into this process on first sight of its identity; reuse the + cached ImportedBuffer thereafter (map-once).""" + desc = BufferHandleDescriptor.unpack(descriptor) if isinstance(descriptor, (bytes, bytearray)) else descriptor + key = desc.identity.pack() + cached = self._by_identity.get(key) + if cached is not None: + return cached + if desc.backend_kind in ( + BackendKind.FORK_SHM, + BackendKind.FORK_COW, + BackendKind.DEVICE_MALLOC, + BackendKind.VMM_WINDOW, + ): + # The body is the base pointer (u64 LE), already valid in this process — no mapping. + # FORK_SHM: a COW-inherited host VA. DEVICE_MALLOC / VMM_WINDOW: a device pointer valid on + # the chip that allocated / carved it (the ref must only reach that chip — a topology + # invariant). + base = int.from_bytes(desc.body, "little") + imported = ImportedBuffer(desc.identity, base, desc.nbytes, desc.address_space, None) + elif desc.backend_kind == BackendKind.POSIX_SHM: + shm = SharedMemory(name=desc.body.decode("utf-8")) + imported = ImportedBuffer(desc.identity, _shm_base_addr(shm), desc.nbytes, desc.address_space, shm) + elif desc.backend_kind == BackendKind.REMOTE_SIDECAR: + raise ValueError("ImportRegistry: REMOTE_SIDECAR backend is reserved for P2") + else: + raise NotImplementedError(f"ImportRegistry: backend {desc.backend_kind!r} not supported in P1-B") + self._by_identity[key] = imported + return imported + + def materialize_blob(self, blob_ptr: int, capacity: int) -> dict[bytes, tuple[int, int]]: + """Lazily materialize every embedded descriptor in a wire blob and return the resolved + map for ``materialize_bufferref_blob``: packed identity -> (local base, address_space).""" + for desc_bytes in bufferref_blob_descriptors(blob_ptr, capacity): + self.materialize(desc_bytes) + return self.materialization_map() + + def mapped_args_from_blob(self, blob_ptr: int, capacity: int) -> MappedArgs: + """Materialize a wire blob into a Python compute callable's args: every tensor becomes a + MappedArg (map-once, buffer at the view origin) and the blob's scalars ride alongside. This is + the compute-leaf map (a sub-worker reads/writes), distinct from pure forwarding (re-export, + which never maps). + """ + tensors = [] + for i, ref in enumerate(Tensor.unpack(rb) for rb in bufferref_blob_refs(blob_ptr, capacity)): + if ref.handle.address_space == AddressSpace.DEVICE: + # Depth behind the submit-time endpoint check: this process is a host compute leaf, so + # a device address here would be handed to torch as a host pointer. + raise ValueError( + f"sub-worker argument {i} is a DEVICE-space tensor " + f"({ref.handle.backend_kind.name}); it cannot be mapped into a host process" + ) + tensors.append(MappedArg(self.materialize(ref.handle), ref.byte_offset, ref.shapes, ref.strides, ref.dtype)) + return MappedArgs(tensors, tuple(bufferref_blob_scalars(blob_ptr, capacity))) + + def resolve(self, identity: CanonicalIdentity) -> ImportedBuffer: + imported = self._by_identity.get(identity.pack()) + if imported is None: + raise KeyError(f"ImportRegistry: no handle registered for {identity}") + return imported + + def materialization_map(self) -> dict[bytes, tuple[int, int]]: + """Snapshot for ``materialize_bufferref_blob``: packed identity -> (local base, address_space).""" + return {key: (ib.base, int(ib.address_space)) for key, ib in self._by_identity.items()} + + def unregister(self, identity: CanonicalIdentity) -> None: + imported = self._by_identity.pop(identity.pack(), None) + if imported is not None and imported.shm is not None: + imported.shm.close() + + def close(self) -> None: + for imported in self._by_identity.values(): + if imported.shm is not None: + imported.shm.close() + self._by_identity.clear() diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 5535007aa8..ad94f19df8 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -95,6 +95,11 @@ def my_l4_orch(orch, args, config): ) from . import _log as _simpler_log +from .buffer_handle import ( + BufferHandle, + create_host_shared_buffer, + mint_owner_instance_id, +) from .callable_identity import ( CALLABLE_HASH_DIGEST_BYTES, CallableHandle, @@ -2306,6 +2311,13 @@ def __init__( # dispatch) is now handled at the C++ boundary via mailbox_mu_, so # no quiescent-state guard is needed. self._registry_lock = threading.Lock() + # Owner-side buffer identity. `_owner_instance_id` is a fresh random draw per Worker + # incarnation, so a buffer_id reused by a later process can never collide with a live + # identity; `_buffer_handles` keeps every handle this Worker owns so close() can unlink + # the backings. + self._owner_instance_id: bytes = mint_owner_instance_id() + self._buffer_id_counter: int = 1 + self._buffer_handles: dict[int, BufferHandle] = {} self._pending_unregister_cids: set[int] = set() self._pending_remote_unregister_hashids: set[bytes] = set() self._py_control_timeout_s = float(config.get("py_control_timeout_s", _PY_CONTROL_TIMEOUT_S)) @@ -6007,6 +6019,67 @@ def _close_host_shm(entry: _HostBufEntry) -> str | None: pass return warn + # ------------------------------------------------------------------ + # Owner-side BufferHandle allocation + # ------------------------------------------------------------------ + + def create_buffer(self, nbytes: int) -> BufferHandle: + """Allocate a shared ``BufferHandle`` owned by this Worker. + + The backing is a POSIX shm; the handle carries a typed canonical identity and a + self-describing descriptor, so a consumer can resolve it with no prior handshake. Build a + tensor over ``handle.shm.buf`` with the buffer protocol. Not thread-safe against a concurrent + run/create/free on the same Worker. + """ + if self.level < 2: + raise TypeError("create_buffer requires a level >= 2 Worker") + with self._operation_lease("create_buffer"): + return self._create_buffer_locked(int(nbytes)) + + def _next_buffer_id(self) -> int: + with self._registry_lock: + bid = self._buffer_id_counter + self._buffer_id_counter += 1 + return bid + + def _create_buffer_locked(self, nbytes: int) -> BufferHandle: + # An L3+ buffer is consumed by a forked child that lazily maps it, so a childless L3+ buffer + # can reach no consumer. An L2 leaf has no children and materializes in-process, so it needs + # none. + if self.level >= 3 and not self._chip_shms and not self._sub_shms: + raise RuntimeError("create_buffer requires at least one forked chip or sub child (this Worker has none)") + if nbytes <= 0: + raise ValueError("create_buffer: nbytes must be positive") + buffer_id = self._next_buffer_id() + handle = create_host_shared_buffer( + nbytes, + owner_instance_id=self._owner_instance_id, + buffer_id=buffer_id, + owner_worker_path=f"L{self.level}", + generation=1, + ) + with self._registry_lock: + self._buffer_handles[buffer_id] = handle + return handle + + def _release_all_buffer_handles(self) -> None: + """Close + unlink every owner BufferHandle (called from close()). + + Best-effort per handle; the first error is raised after all are attempted so close() reports + a leak rather than swallowing it. + """ + with self._registry_lock: + handles = list(self._buffer_handles.values()) + self._buffer_handles.clear() + errors: list[BaseException] = [] + for handle in handles: + try: + handle.close() + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + if errors: + raise errors[0] + def _release_all_host_buffers(self) -> None: """Unmap + free every still-registered host buffer (called from close()). @@ -6713,6 +6786,7 @@ def _step(fn) -> None: _step(self._clear_child_prov) _step(self._release_active_remote_slot_refs) _step(self._flush_pending_remote_frees) + _step(self._release_all_buffer_handles) # Host buffers must be released while the local L3 child mailboxes are # still usable (before _worker.close()). _step(self._release_all_host_buffers) diff --git a/src/common/task_interface/buffer_handle.h b/src/common/task_interface/buffer_handle.h new file mode 100644 index 0000000000..78961da9df --- /dev/null +++ b/src/common/task_interface/buffer_handle.h @@ -0,0 +1,403 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +#pragma once + +/** + * BufferHandle / BufferRef ABI — typed, opaque cross-layer buffer identity. + * + * Three types: + * - CanonicalIdentity : owner_instance_id + buffer_id + generation. The key both the owner + * registry and every consumer import cache use, invariant across every + * edge. Fixed-length with no length field, so hashing and comparison + * cannot read past it. + * - BufferHandleDescriptor : the owner's self-describing wire descriptor — backing properties plus + * a length-delimited backend body. Embedded whole in every BufferRef + * built over the handle. + * - BufferRef : the blob-carried wire element. Embeds the full BufferHandleDescriptor + * plus a view (byte_offset, shape, strides, dtype) — self-describing, so + * a consumer materializes it lazily on receipt with no prior handshake. + * No materialized address. + * + * There is no wire version. Every endpoint of a run comes from one `pip install`, so a version field + * would guard a skew that cannot arise; the skew that CAN arise (a stale compiled extension against + * newer Python) is caught by SIMPLER_BUILD_COMMIT. The leading `magic` on the descriptor and on the + * blob envelope are discriminators against non-descriptor bytes, not versions. + * + * Endianness: all multi-byte integers little-endian. owner_instance_id is an opaque byte sequence + * (bytewise-compared, no integer/endianness meaning). An unknown backend / address_space / access + * value is rejected, never silently accepted. + */ + +#include +#include +#include +#include +#include +#include + +#include "data_type.h" + +// Leading sentinel of a BufferHandleDescriptor, NOT a version. A descriptor decoder needs a cheap +// leading discriminator because `TaskArgs.add_tensor` accepts raw bytes; the sentinel rejects most +// non-descriptor input before any field is trusted. It does not by itself prove the bytes are a +// descriptor — every other field is still validated. There is no multi-version wire: every endpoint +// of a run is built from one `pip install`, and build skew is caught by SIMPLER_BUILD_COMMIT. +inline constexpr uint16_t BUFFER_DESCRIPTOR_MAGIC = 0x5342; // 'BS' little-endian + +// Leading sentinel of the BufferRef blob ENVELOPE (the length-prefixed container in task_args.h), +// distinct from the descriptor sentinel above. Same role, same reason: not a version, just a cheap +// discriminator against a buffer that is not a blob. `read_bufferref_blob` rejects any other value. +inline constexpr uint32_t BUFFERREF_BLOB_MAGIC = 0x424F4C42; // "BLOB" in memory order + +// owner_instance_id is a fixed-width opaque nonce (compared bytewise; no integer/endianness meaning). +// It is the SOLE source of cross-incarnation uniqueness, so it must be generated with a full-width +// random draw — a structured value (timestamp/pid) collides between two Workers built in the same +// process and second. +inline constexpr uint32_t OWNER_INSTANCE_ID_BYTES = 8; + +// Backend body upper bound. Only POSIX_SHM uses more than 8 bytes (a shm name); every other backend +// stores a single u64 address. +inline constexpr uint32_t DESC_MAX_BYTES = 32; + +// AddressSpace (HOST/DEVICE) is shared with Tensor and lives in data_type.h. + +// The backing's granted permission. A per-arg TensorArgType requests read/write and is validated +// against this at submit (requested must be a subset of granted). +enum class AccessMode : uint8_t { + READ = 0, + WRITE = 1, + READWRITE = 2, +}; + +// Materialization backend of a handle. The consumer resolves a BufferRef to a local address via the +// import registry keyed by canonical identity; this tag selects how. REMOTE_SIDECAR is reserved for +// P2 and rejected on decode in P1. Values are frozen; 6.. reserved (unknown tag => reject). +// +// FORK_SHM and FORK_COW materialize identically — the body is a base VA the child already has, +// inherited across the fork — but the kernel's write semantics are opposite, so they are separate +// tags rather than one tag plus a hint. A child's write to a MAP_SHARED page lands in the physical +// page the parent reads; a write to a copy-on-write page splits it into a private copy the parent +// never sees, silently. FORK_COW therefore grants READ only, and that is enforced on decode. +enum class BackendKind : uint8_t { + FORK_SHM = 0, + POSIX_SHM = 1, + VMM_WINDOW = 2, + REMOTE_SIDECAR = 3, + DEVICE_MALLOC = 4, + FORK_COW = 5, +}; + +/** + * Canonical allocation identity — globally unique across owner incarnations, unchanged across every + * edge. `buffer_id` is unique only within one owner incarnation; `owner_instance_id` (a per-incarnation + * nonce) disambiguates it, and `generation` detects buffer_id slot reuse (ABA). The key of both the + * owner registry and every consumer import registry. + * + * FIXED-LENGTH BY DESIGN: no field here bounds a read. Hashing and comparison therefore cannot run + * off the end of the struct whatever bytes arrive on the wire — the property is structural, not + * something a validator has to enforce. + * + * `_pad` is excluded from comparison and hashing, so a decoded identity with dirty padding still + * matches the same backing (two views of one backing must never key differently). + * + * `generation` starts at 1; every reuse of a `buffer_id` slot increments it. 0 is reserved to mean + * uninitialized and is rejected on decode. + */ +struct CanonicalIdentity { + uint8_t owner_instance_id[OWNER_INSTANCE_ID_BYTES]; + uint64_t buffer_id; + uint32_t generation; + uint8_t _pad[12]; +}; + +static_assert(std::is_trivially_copyable_v); +static_assert(sizeof(CanonicalIdentity) == 32, "CanonicalIdentity is wire ABI"); +static_assert(offsetof(CanonicalIdentity, owner_instance_id) == 0); +static_assert(offsetof(CanonicalIdentity, buffer_id) == 8); +static_assert(offsetof(CanonicalIdentity, generation) == 16); + +inline bool operator==(const CanonicalIdentity &a, const CanonicalIdentity &b) { + return a.buffer_id == b.buffer_id && a.generation == b.generation && + std::memcmp(a.owner_instance_id, b.owner_instance_id, OWNER_INSTANCE_ID_BYTES) == 0; +} +inline bool operator!=(const CanonicalIdentity &a, const CanonicalIdentity &b) { return !(a == b); } + +// Hash for use as an unordered_map key (consumer import registry). Folds exactly the fields +// `operator==` compares — `_pad` is excluded so padding can never perturb the bucket. +struct CanonicalIdentityHash { + size_t operator()(const CanonicalIdentity &k) const { + auto mix = [](size_t h, uint64_t v) { + return (h ^ (v + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2))); + }; + size_t h = 0; + for (uint32_t i = 0; i < OWNER_INSTANCE_ID_BYTES; ++i) + h = mix(h, k.owner_instance_id[i]); + h = mix(h, k.buffer_id); + h = mix(h, k.generation); + return h; + } +}; + +/** + * The owner's self-describing handle descriptor — embedded whole in every BufferRef built over the + * handle. A consumer materializes it lazily on first receipt (no separate export handshake) and + * caches `canonical identity -> local base` (map-once). `backend_kind` + `body[0, body_len)` carry + * the per-backend materialization (POSIX shm name, fork-inherited VA, device VA, ...). + * `magic` leads as a cheap discriminator; `address_space` / `access` / `backend_kind` are raw u8 so + * an unknown value can be rejected without invoking undefined enum behavior. + * + * `owner_worker_path_id` is a DIAGNOSTIC id only — it names the owning worker in logs and + * post-mortems and takes part in no routing, visibility or identity decision. Its side table lives in + * the owning process; an id a consumer cannot resolve prints as `` and is never an error. + */ +struct BufferHandleDescriptor { + uint16_t magic; + uint8_t address_space; + uint8_t access; + uint8_t backend_kind; + uint8_t _pad0[3]; + CanonicalIdentity identity; + uint64_t nbytes; + uint32_t owner_worker_path_id; + uint16_t body_len; + uint8_t _pad1[2]; + char body[DESC_MAX_BYTES]; +}; + +static_assert(std::is_trivially_copyable_v); +static_assert(sizeof(BufferHandleDescriptor) == 88, "BufferHandleDescriptor is wire ABI"); +static_assert(offsetof(BufferHandleDescriptor, magic) == 0); +static_assert(offsetof(BufferHandleDescriptor, address_space) == 2); +static_assert(offsetof(BufferHandleDescriptor, access) == 3); +static_assert(offsetof(BufferHandleDescriptor, backend_kind) == 4); +static_assert(offsetof(BufferHandleDescriptor, identity) == 8); +static_assert(offsetof(BufferHandleDescriptor, nbytes) == 40); +static_assert(offsetof(BufferHandleDescriptor, owner_worker_path_id) == 48); +static_assert(offsetof(BufferHandleDescriptor, body_len) == 52); +static_assert(offsetof(BufferHandleDescriptor, body) == 56); + +/** + * The blob-carried, self-describing wire element: a full embedded handle descriptor plus a strided + * view onto it. Because the descriptor travels with the ref, a consumer needs no prior handshake — + * it materializes the embedded `handle` (backend selects how) on first receipt, keyed by + * `handle.identity`, and reuses the cached base for later refs to the same identity. + * + * Invariants: + * - Carries NO materialized address. The consumer materializes `handle` to a local base, then + * `Tensor.buffer.addr = base`, `Tensor.start_offset = byte_offset / dtype_bytes`. + * - `byte_offset` is a BYTE offset of the view origin; a multiple of the dtype size (validated at + * materialization). + * - `strides[i] > 0` strictly (broadcast / negative step unsupported), carried explicitly — a + * singleton dimension's stride is never normalized away. + */ +struct BufferRef { + BufferHandleDescriptor handle; + uint64_t byte_offset; + uint32_t ndims; + uint32_t shapes[MAX_TENSOR_DIMS]; + uint32_t strides[MAX_TENSOR_DIMS]; + DataType dtype; + uint8_t _pad[3]; +}; + +// Byte extent of a (possibly strided) view: the last addressable element, plus one element. Summed +// in u64 so a hostile shape/stride cannot wrap it. Callers that have not yet validated `r` must not +// trust the result for anything but a bound. +inline uint64_t buffer_ref_extent_bytes(const BufferRef &r) { + uint64_t last_elem = 0; + for (uint32_t i = 0; i < r.ndims && i < static_cast(MAX_TENSOR_DIMS); ++i) { + if (r.shapes[i] == 0) continue; + last_elem += static_cast(r.shapes[i] - 1) * static_cast(r.strides[i]); + } + return (last_elem + 1) * get_element_size(r.dtype); +} + +/** + * Reject any BufferRef whose fields are not self-consistent, BEFORE any of them is trusted. + * + * This is the single implementation behind all three trust boundaries — the builder + * (`TaskArgs.add_tensor`, which accepts raw bytes), blob decode on receipt, and materialization — + * so the three can never drift apart. Throws `std::invalid_argument` naming the field. + * + * Every remaining length-like field is bounded here: `body_len` against `DESC_MAX_BYTES` and `ndims` + * against `MAX_TENSOR_DIMS`, mirroring what the fixed-length `CanonicalIdentity` gets structurally. + * + * `REMOTE_SIDECAR` is a legal wire value and passes: an arg bound for a remote worker rides the wire + * with no local backing, and it is *materialization* that refuses it in P1. + */ +inline void validate_buffer_ref(const BufferRef &r) { + auto reject = [](const char *what) { + throw std::invalid_argument(what); + }; + const BufferHandleDescriptor &h = r.handle; + + if (h.magic != BUFFER_DESCRIPTOR_MAGIC) reject("invalid Tensor: descriptor magic"); + if (h.address_space > static_cast(AddressSpace::DEVICE)) + reject("invalid Tensor: address_space out of range"); + if (h.access > static_cast(AccessMode::READWRITE)) reject("invalid Tensor: access out of range"); + if (h.backend_kind > static_cast(BackendKind::FORK_COW)) + reject("invalid Tensor: backend_kind out of range"); + if (h.body_len > DESC_MAX_BYTES) reject("invalid Tensor: body_len exceeds DESC_MAX_BYTES"); + if (h.identity.generation == 0) reject("invalid Tensor: generation 0 is reserved (uninitialized)"); + + // address_space x backend_kind capability gate. REMOTE_SIDECAR is legal in either space. + const auto backend = static_cast(h.backend_kind); + const bool device_space = h.address_space == static_cast(AddressSpace::DEVICE); + if (backend != BackendKind::REMOTE_SIDECAR) { + const bool device_backend = backend == BackendKind::VMM_WINDOW || backend == BackendKind::DEVICE_MALLOC; + if (device_backend != device_space) reject("invalid Tensor: unsupported address_space x backend_kind"); + } + + if (r.ndims == 0 || r.ndims > static_cast(MAX_TENSOR_DIMS)) reject("invalid Tensor: ndims out of range"); + if (r.dtype >= DataType::DATA_TYPE_NUM) reject("invalid Tensor: unknown dtype"); + const uint64_t elem = get_element_size(r.dtype); + if (elem == 0) reject("invalid Tensor: unknown dtype"); + if (r.byte_offset % elem != 0) reject("invalid Tensor: byte_offset is not a multiple of the dtype size"); + + for (uint32_t i = 0; i < r.ndims; ++i) { + if (r.shapes[i] == 0) reject("invalid Tensor: shape dimension is zero"); + if (r.strides[i] == 0) reject("invalid Tensor: stride must be > 0 (broadcast and negative step unsupported)"); + } + const uint64_t extent_bytes = buffer_ref_extent_bytes(r); + if (r.byte_offset > h.nbytes || extent_bytes > h.nbytes - r.byte_offset) { + reject("invalid Tensor: view extends past the backing (byte_offset + extent > nbytes)"); + } +} + +// Do two views of the SAME backing touch a common byte? Compared as bounding ranges +// [byte_offset, byte_offset + extent): a strided view's gaps are treated as occupied, so this is +// conservative — it never misses a real overlap, and may report one for two interleaved views. +inline bool buffer_refs_overlap(const BufferRef &a, const BufferRef &b) { + if (!(a.handle.identity == b.handle.identity)) return false; + const uint64_t a_end = a.byte_offset + buffer_ref_extent_bytes(a); + const uint64_t b_end = b.byte_offset + buffer_ref_extent_bytes(b); + return a.byte_offset < b_end && b.byte_offset < a_end; +} + +static_assert(std::is_trivially_copyable_v, "BufferRef must be trivially copyable for blob memcpy"); +static_assert(sizeof(BufferRef) == 144, "BufferRef is wire ABI"); +static_assert(offsetof(BufferRef, handle) == 0); +static_assert(offsetof(BufferRef, byte_offset) == 88); +static_assert(offsetof(BufferRef, ndims) == 96); +static_assert(offsetof(BufferRef, shapes) == 100); +static_assert(offsetof(BufferRef, strides) == 120); +static_assert(offsetof(BufferRef, dtype) == 140); + +// ============================================================================ +// BufferRef wire blob — versioned, length-prefixed (P1-B). +// ============================================================================ +// +// Byte layout: +// offset 0: uint32 magic = BUFFERREF_BLOB_MAGIC +// offset 4: int32 ref_count = R +// offset 8: int32 scalar_count = S +// offset 12: uint32 reserved (= 0) +// offset 16: BufferRef refs[R] (sizeof(BufferRef) B each) +// offset 16 + R*sizeof(BufferRef): uint64_t scalars[S] +// +// The element is BufferRef (embedded handle descriptor + view, no materialized addr). `magic` is a +// leading discriminator, not a version — a decoder rejects a buffer that is not a blob rather than +// misreading it; every other field is still bounded by `capacity` below. The reserved word 8-aligns +// refs[0] (whose first field is a u64) and is rejected when non-zero. + +inline constexpr size_t BUFFERREF_BLOB_HEADER_SIZE = 16; + +struct BufferRefBlobView { + int32_t ref_count; + int32_t scalar_count; + const uint8_t *ref_bytes; // R contiguous BufferRef; extract element i with ref(i) + const uint64_t *scalars; + + // Every consumer — L4->L3 re-export, L2 materialization, the Python sub-worker map — extracts + // its elements through here, so validating on extraction covers all three receive boundaries at + // one point. Throws std::invalid_argument on a malformed element. + BufferRef ref(int32_t i) const { + BufferRef r; + std::memcpy(&r, ref_bytes + static_cast(i) * sizeof(BufferRef), sizeof(BufferRef)); + validate_buffer_ref(r); + return r; + } +}; + +inline size_t bufferref_blob_size(int32_t ref_count, int32_t scalar_count) { + return BUFFERREF_BLOB_HEADER_SIZE + static_cast(ref_count) * sizeof(BufferRef) + + static_cast(scalar_count) * sizeof(uint64_t); +} + +// Serialize refs + scalars into `dst` (caller ensures room for bufferref_blob_size). +inline void write_bufferref_blob( + uint8_t *dst, const BufferRef *refs, int32_t ref_count, const uint64_t *scalars, int32_t scalar_count +) { + uint32_t magic = BUFFERREF_BLOB_MAGIC; + uint32_t reserved = 0; + std::memcpy(dst + 0, &magic, sizeof(magic)); + std::memcpy(dst + 4, &ref_count, sizeof(ref_count)); + std::memcpy(dst + 8, &scalar_count, sizeof(scalar_count)); + std::memcpy(dst + 12, &reserved, sizeof(reserved)); + if (ref_count > 0) { + std::memcpy(dst + BUFFERREF_BLOB_HEADER_SIZE, refs, static_cast(ref_count) * sizeof(BufferRef)); + } + if (scalar_count > 0) { + std::memcpy( + dst + BUFFERREF_BLOB_HEADER_SIZE + static_cast(ref_count) * sizeof(BufferRef), scalars, + static_cast(scalar_count) * sizeof(uint64_t) + ); + } +} + +// Zero-copy view into a blob written by write_bufferref_blob; valid while `src` stays mapped. +// `capacity` bounds the read. Throws on a bad magic, negative counts, or a header that +// would walk past `capacity` (shared-memory corruption or a writer-side bug). +inline BufferRefBlobView read_bufferref_blob(const uint8_t *src, size_t capacity) { + if (capacity < BUFFERREF_BLOB_HEADER_SIZE) { + throw std::runtime_error( + "read_bufferref_blob: capacity " + std::to_string(capacity) + " < header size " + + std::to_string(BUFFERREF_BLOB_HEADER_SIZE) + ); + } + uint32_t magic; + std::memcpy(&magic, src + 0, sizeof(magic)); + if (magic != BUFFERREF_BLOB_MAGIC) { + throw std::runtime_error( + "read_bufferref_blob: not a BufferRef blob — magic " + std::to_string(magic) + " (expected " + + std::to_string(BUFFERREF_BLOB_MAGIC) + ")" + ); + } + int32_t R; + int32_t S; + uint32_t reserved; + std::memcpy(&R, src + 4, sizeof(R)); + std::memcpy(&S, src + 8, sizeof(S)); + std::memcpy(&reserved, src + 12, sizeof(reserved)); + if (reserved != 0) { + throw std::runtime_error("read_bufferref_blob: reserved header word must be zero"); + } + if (R < 0 || S < 0) { + throw std::runtime_error( + "read_bufferref_blob: negative counts — refs=" + std::to_string(R) + ", scalars=" + std::to_string(S) + ); + } + const size_t needed = bufferref_blob_size(R, S); + if (needed > capacity) { + throw std::runtime_error( + "read_bufferref_blob: header reports " + std::to_string(needed) + " bytes (R=" + std::to_string(R) + + ", S=" + std::to_string(S) + ") but capacity is " + std::to_string(capacity) + ); + } + return BufferRefBlobView{ + R, + S, + src + BUFFERREF_BLOB_HEADER_SIZE, + reinterpret_cast( + src + BUFFERREF_BLOB_HEADER_SIZE + static_cast(R) * sizeof(BufferRef) + ), + }; +} diff --git a/src/common/task_interface/data_type.h b/src/common/task_interface/data_type.h index 7c48db0b39..8fac337173 100644 --- a/src/common/task_interface/data_type.h +++ b/src/common/task_interface/data_type.h @@ -28,6 +28,17 @@ template inline constexpr bool is_supported_scalar_arg_v = std::is_arithmetic_v>> || std::is_enum_v>>; +// Maximum tensor rank a view (Tensor / BufferRef) can describe. Wire ABI: the shapes[] / strides[] +// arrays in both are sized by it, so it is shared here rather than owned by either. +constexpr int MAX_TENSOR_DIMS = 5; + +// Memory space of a backing. Byte-43 of Tensor and a BufferHandle field both store it, so it is +// shared here. Orthogonal to location (local/remote, derived) and visibility. +enum class AddressSpace : uint8_t { + HOST = 0, + DEVICE = 1, +}; + /** * Supported data types for tensor elements */ @@ -80,7 +91,9 @@ inline uint64_t get_element_size(DataType dtype) { 1, // DataType::FP8E8M0 (A5 only) 1, // DataType::FP4E2M1 (A5 only) }; - return data_type_size[static_cast(dtype)]; + // Bounds-checked: `dtype` is a raw u8 on the wire, so a decoder can hand this any value. + const auto index = static_cast(dtype); + return index < data_type_size.size() ? data_type_size[index] : 0; } /** diff --git a/src/common/task_interface/tensor.h b/src/common/task_interface/tensor.h index 8b62208bcd..e4f47ba5c9 100644 --- a/src/common/task_interface/tensor.h +++ b/src/common/task_interface/tensor.h @@ -23,8 +23,6 @@ #include "data_type.h" #include "pto_task_id.h" -constexpr int MAX_TENSOR_DIMS = 5; - /** * Buffer Handle * diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index cc75282a04..800f1d316a 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -384,6 +384,7 @@ set_tests_properties(test_run_stream_slots PROPERTIES LABELS "no_hardware") # --------------------------------------------------------------------------- # Types / task_interface tests (src/common/task_interface/) # --------------------------------------------------------------------------- +add_task_interface_test(test_buffer_handle types/test_buffer_handle.cpp) add_task_interface_test(test_child_memory types/test_child_memory.cpp) add_task_interface_test(test_chip_max_tensor_args types/test_chip_max_tensor_args.cpp) add_task_interface_test(test_call_config types/test_call_config.cpp) diff --git a/tests/ut/cpp/types/test_buffer_handle.cpp b/tests/ut/cpp/types/test_buffer_handle.cpp new file mode 100644 index 0000000000..c856cba101 --- /dev/null +++ b/tests/ut/cpp/types/test_buffer_handle.cpp @@ -0,0 +1,239 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +// Wire-ABI tests for BufferHandle / BufferRef (buffer_handle.h); the contract they pin is described +// in docs/buffer-handle-abi.md. Byte layout is pinned by static_assert in the header; these tests +// pin the sizes, enum values, and the blob codec from the outside. + +#include +#include +#include +#include +#include +#include + +#include + +#include "buffer_handle.h" +#include "task_args.h" + +namespace { + +CanonicalIdentity make_identity() { + CanonicalIdentity id{}; + for (uint32_t i = 0; i < OWNER_INSTANCE_ID_BYTES; ++i) + id.owner_instance_id[i] = static_cast(0xA0 + i); + id.buffer_id = 0x0102030405060708ULL; + id.generation = 7; + return id; +} + +BufferRef make_ref() { + BufferRef r{}; + r.handle.magic = BUFFER_DESCRIPTOR_MAGIC; + r.handle.backend_kind = static_cast(BackendKind::POSIX_SHM); + r.handle.identity = make_identity(); + r.handle.nbytes = 8192; // must cover byte_offset + the strided extent below + r.byte_offset = 4096; + r.ndims = 3; + r.shapes[0] = 2; + r.shapes[1] = 4; + r.shapes[2] = 8; + r.strides[0] = 32; + r.strides[1] = 8; + r.strides[2] = 1; + r.dtype = DataType::FLOAT16; + return r; +} + +// --- Layout / value contracts (frozen ABI) ----------------------------------------------------- + +TEST(BufferHandleAbi, StructSizesAreFrozen) { + EXPECT_EQ(sizeof(CanonicalIdentity), 32u); + EXPECT_EQ(sizeof(BufferRef), 144u); + EXPECT_EQ(sizeof(BufferHandleDescriptor), 88u); +} + +TEST(BufferHandleAbi, ConstantsAreFrozen) { + EXPECT_EQ(BUFFER_DESCRIPTOR_MAGIC, 0x5342); + EXPECT_EQ(BUFFERREF_BLOB_MAGIC, 0x424F4C42u); + EXPECT_EQ(OWNER_INSTANCE_ID_BYTES, 8u); + EXPECT_EQ(DESC_MAX_BYTES, 32u); +} + +TEST(BufferHandleAbi, EnumValuesAreFrozen) { + EXPECT_EQ(static_cast(AddressSpace::HOST), 0); + EXPECT_EQ(static_cast(AddressSpace::DEVICE), 1); + EXPECT_EQ(static_cast(AccessMode::READ), 0); + EXPECT_EQ(static_cast(AccessMode::WRITE), 1); + EXPECT_EQ(static_cast(AccessMode::READWRITE), 2); + EXPECT_EQ(static_cast(BackendKind::FORK_SHM), 0); + EXPECT_EQ(static_cast(BackendKind::POSIX_SHM), 1); + EXPECT_EQ(static_cast(BackendKind::VMM_WINDOW), 2); + EXPECT_EQ(static_cast(BackendKind::REMOTE_SIDECAR), 3); + EXPECT_EQ(static_cast(BackendKind::DEVICE_MALLOC), 4); +} + +// --- memcpy round trip ------------------------------------------------------------------------- + +TEST(BufferHandleAbi, BufferRefSurvivesByteRoundTrip) { + BufferRef src = make_ref(); + uint8_t bytes[sizeof(BufferRef)]; + std::memcpy(bytes, &src, sizeof(BufferRef)); + BufferRef dst{}; + std::memcpy(&dst, bytes, sizeof(BufferRef)); + EXPECT_EQ(std::memcmp(&src, &dst, sizeof(BufferRef)), 0); + EXPECT_EQ(dst.byte_offset, 4096u); + EXPECT_EQ(dst.dtype, DataType::FLOAT16); + EXPECT_EQ(dst.strides[0], 32u); + EXPECT_EQ(dst.handle.identity, src.handle.identity); +} + +TEST(BufferHandleAbi, HandleDescriptorSurvivesByteRoundTrip) { + BufferHandleDescriptor src{}; + src.magic = BUFFER_DESCRIPTOR_MAGIC; + src.address_space = static_cast(AddressSpace::DEVICE); + src.access = static_cast(AccessMode::READWRITE); + src.backend_kind = static_cast(BackendKind::POSIX_SHM); + src.owner_worker_path_id = 3; + src.identity = make_identity(); + src.nbytes = 1 << 20; + const char *body = "psm_deadbeef"; + src.body_len = static_cast(std::strlen(body)); + std::memcpy(src.body, body, src.body_len); + + uint8_t bytes[sizeof(BufferHandleDescriptor)]; + std::memcpy(bytes, &src, sizeof(BufferHandleDescriptor)); + BufferHandleDescriptor dst{}; + std::memcpy(&dst, bytes, sizeof(BufferHandleDescriptor)); + EXPECT_EQ(std::memcmp(&src, &dst, sizeof(BufferHandleDescriptor)), 0); + EXPECT_EQ(dst.magic, BUFFER_DESCRIPTOR_MAGIC); + EXPECT_EQ(dst.owner_worker_path_id, 3u); + EXPECT_EQ(dst.identity, src.identity); + EXPECT_EQ(std::string(dst.body, dst.body_len), "psm_deadbeef"); +} + +// --- canonical identity: the import-registry key ----------------------------------------------- + +TEST(BufferHandleAbi, IdentityDistinguishesGenerationAndIncarnation) { + CanonicalIdentity a = make_identity(); + CanonicalIdentity b = make_identity(); + EXPECT_EQ(a, b); + + b.generation = a.generation + 1; // buffer_id reuse across generations (ABA) + EXPECT_NE(a, b); + + CanonicalIdentity c = make_identity(); + c.owner_instance_id[0] ^= 0xFF; // different owner incarnation nonce + EXPECT_NE(a, c); + + CanonicalIdentity d = make_identity(); + d.buffer_id ^= 0xFFULL; + EXPECT_NE(a, d); +} + +// Padding is wire-visible but semantically absent: two decodes of one backing must key identically, +// or the same buffer lands in two import-registry buckets and its dependencies split. +TEST(BufferHandleAbi, IdentityPaddingIsExcludedFromKeyAndHash) { + CanonicalIdentityHash h; + CanonicalIdentity clean = make_identity(); + CanonicalIdentity dirty = make_identity(); + std::memset(dirty._pad, 0xA5, sizeof(dirty._pad)); + EXPECT_EQ(clean, dirty); + EXPECT_EQ(h(clean), h(dirty)); +} + +TEST(BufferHandleAbi, IdentityHashMatchesEquality) { + CanonicalIdentityHash h; + CanonicalIdentity a = make_identity(); + CanonicalIdentity b = make_identity(); + EXPECT_EQ(h(a), h(b)); + + b.generation = a.generation + 1; + EXPECT_NE(h(a), h(b)); // good hash separates the ABA case (not a strict requirement) +} + +// --- BufferRef wire blob: versioned length-prefixed round trip + rejection ---------------------- + +BufferRef make_ref_b() { + BufferRef r = make_ref(); + r.handle.identity.buffer_id = 99; + r.byte_offset = 0; + r.ndims = 1; + r.shapes[0] = 5; + r.strides[0] = 1; + r.dtype = DataType::INT32; + return r; +} + +TEST(BufferRefBlob, RoundTrip) { + BufferRef refs[2] = {make_ref(), make_ref_b()}; + uint64_t scalars[2] = {42, 0xC0FFEE}; + size_t sz = bufferref_blob_size(2, 2); + EXPECT_EQ(sz, BUFFERREF_BLOB_HEADER_SIZE + 2 * sizeof(BufferRef) + 2 * sizeof(uint64_t)); + + std::vector buf(sz); + write_bufferref_blob(buf.data(), refs, 2, scalars, 2); + + BufferRefBlobView v = read_bufferref_blob(buf.data(), sz); + ASSERT_EQ(v.ref_count, 2); + ASSERT_EQ(v.scalar_count, 2); + BufferRef r0 = v.ref(0); + BufferRef r1 = v.ref(1); + EXPECT_EQ(std::memcmp(&r0, &refs[0], sizeof(BufferRef)), 0); + EXPECT_EQ(std::memcmp(&r1, &refs[1], sizeof(BufferRef)), 0); + EXPECT_EQ(v.scalars[0], 42u); + EXPECT_EQ(v.scalars[1], 0xC0FFEEu); +} + +TEST(BufferRefBlob, EmptyBlob) { + size_t sz = bufferref_blob_size(0, 0); + EXPECT_EQ(sz, BUFFERREF_BLOB_HEADER_SIZE); + std::vector buf(sz); + write_bufferref_blob(buf.data(), nullptr, 0, nullptr, 0); + BufferRefBlobView v = read_bufferref_blob(buf.data(), sz); + EXPECT_EQ(v.ref_count, 0); + EXPECT_EQ(v.scalar_count, 0); +} + +TEST(BufferRefBlob, RejectsBadMagic) { + std::vector buf(bufferref_blob_size(0, 0)); + write_bufferref_blob(buf.data(), nullptr, 0, nullptr, 0); + uint32_t bad = BUFFERREF_BLOB_MAGIC + 1; + std::memcpy(buf.data(), &bad, sizeof(bad)); + EXPECT_THROW(read_bufferref_blob(buf.data(), buf.size()), std::runtime_error); +} + +TEST(BufferRefBlob, RejectsTruncatedCapacity) { + BufferRef refs[1] = {make_ref()}; + std::vector buf(bufferref_blob_size(1, 0)); + write_bufferref_blob(buf.data(), refs, 1, nullptr, 0); + EXPECT_THROW(read_bufferref_blob(buf.data(), BUFFERREF_BLOB_HEADER_SIZE), std::runtime_error); + EXPECT_THROW(read_bufferref_blob(buf.data(), 4), std::runtime_error); +} + +TEST(BufferRefBlob, RejectsNegativeCount) { + std::vector buf(64, 0); + uint32_t ver = BUFFERREF_BLOB_MAGIC; + std::memcpy(buf.data() + 0, &ver, sizeof(ver)); + int32_t neg = -1; + std::memcpy(buf.data() + 4, &neg, sizeof(neg)); // ref_count = -1 + EXPECT_THROW(read_bufferref_blob(buf.data(), buf.size()), std::runtime_error); +} + +TEST(BufferRefBlob, RejectsNonZeroReserved) { + std::vector buf(bufferref_blob_size(0, 0)); + write_bufferref_blob(buf.data(), nullptr, 0, nullptr, 0); + uint32_t dirty = 1; + std::memcpy(buf.data() + 12, &dirty, sizeof(dirty)); // reserved header word + EXPECT_THROW(read_bufferref_blob(buf.data(), buf.size()), std::runtime_error); +} + +} // namespace diff --git a/tests/ut/py/test_buffer_handle.py b/tests/ut/py/test_buffer_handle.py new file mode 100644 index 0000000000..137ef7f166 --- /dev/null +++ b/tests/ut/py/test_buffer_handle.py @@ -0,0 +1,397 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Unit tests for simpler.buffer_handle: identity/descriptor pack-unpack + create/import round trip.""" + +import ctypes +import struct +from dataclasses import replace + +import pytest +from _task_interface import ( + BUFFER_HANDLE_DESCRIPTOR_BYTES, + BUFFER_REF_BYTES, + BUFFERREF_BLOB_MAGIC, + CANONICAL_IDENTITY_BYTES, + OWNER_INSTANCE_ID_BYTES, + DataType, + bufferref_blob_refs, +) +from simpler.buffer_handle import ( + AccessMode, + AddressSpace, + BackendKind, + BufferHandleDescriptor, + CanonicalIdentity, + ImportRegistry, + Tensor, + create_host_shared_buffer, + intern_worker_path, + mint_owner_instance_id, + re_export, + wrap_device_malloc, +) + +_OID = bytes(range(0xA0, 0xA0 + OWNER_INSTANCE_ID_BYTES)) + + +def _identity(oid=_OID, buffer_id=7, generation=2): + return CanonicalIdentity(oid, buffer_id, generation) + + +def test_identity_roundtrip(): + ident = _identity() + raw = ident.pack() + assert len(raw) == CANONICAL_IDENTITY_BYTES + assert CanonicalIdentity.unpack(raw) == ident + + +def test_identity_is_fixed_length_no_length_field(): + # The structural guarantee: nothing inside the identity bounds a read, so a hostile blob cannot + # steer a hash or a compare past the struct. Any CANONICAL_IDENTITY_BYTES decode either yields an + # identity or is rejected on a value check — never an out-of-bounds read. + for raw in (b"\xff" * CANONICAL_IDENTITY_BYTES, bytes(range(CANONICAL_IDENTITY_BYTES))): + try: + ident = CanonicalIdentity.unpack(raw) + except ValueError: + continue + assert len(ident.pack()) == CANONICAL_IDENTITY_BYTES + + +def test_identity_rejects_bad_oid_width(): + with pytest.raises(ValueError): + CanonicalIdentity(b"\x00" * (OWNER_INSTANCE_ID_BYTES + 1), 1, 1) + + +def test_identity_rejects_reserved_generation_zero(): + raw = bytearray(_identity().pack()) + raw[16:20] = (0).to_bytes(4, "little") # generation field + with pytest.raises(ValueError, match="generation 0"): + CanonicalIdentity.unpack(bytes(raw)) + + +def test_identity_padding_does_not_perturb_the_key(): + # Two decodes of the same backing must key identically even if the wire padding differs, or the + # same buffer would land in two import-registry / dependency buckets. + clean = bytearray(_identity().pack()) + dirty = bytearray(clean) + dirty[20:32] = b"\xa5" * 12 # _pad + assert CanonicalIdentity.unpack(bytes(clean)) == CanonicalIdentity.unpack(bytes(dirty)) + assert hash(CanonicalIdentity.unpack(bytes(clean))) == hash(CanonicalIdentity.unpack(bytes(dirty))) + + +def test_identity_distinguishes_generation_and_incarnation(): + a = _identity() + assert a != _identity(generation=a.generation + 1) # ABA + assert a != _identity(oid=bytes(range(1, 1 + OWNER_INSTANCE_ID_BYTES))) # different incarnation + assert a != _identity(buffer_id=a.buffer_id + 1) + + +def test_descriptor_roundtrip_host_and_device(): + host = BufferHandleDescriptor( + identity=_identity(), + address_space=AddressSpace.HOST, + access=AccessMode.READWRITE, + backend_kind=BackendKind.POSIX_SHM, + nbytes=4096, + body=b"psm_deadbeef", + ) + raw = host.pack() + assert len(raw) == BUFFER_HANDLE_DESCRIPTOR_BYTES + assert BufferHandleDescriptor.unpack(raw) == host + + dev = BufferHandleDescriptor( + identity=_identity(buffer_id=99), + address_space=AddressSpace.DEVICE, + access=AccessMode.READ, + backend_kind=BackendKind.VMM_WINDOW, + nbytes=1 << 20, + body=(0x7F00ABCD).to_bytes(8, "little"), + ) + assert BufferHandleDescriptor.unpack(dev.pack()) == dev + + +def test_descriptor_rejects_bad_magic(): + raw = bytearray( + BufferHandleDescriptor( + identity=_identity(), + address_space=AddressSpace.HOST, + access=AccessMode.READWRITE, + backend_kind=BackendKind.POSIX_SHM, + nbytes=8, + ).pack() + ) + raw[0] = raw[0] + 1 # corrupt the leading sentinel (u16 @ offset 0) + with pytest.raises(ValueError, match="magic"): + BufferHandleDescriptor.unpack(bytes(raw)) + + +def test_descriptor_rejects_body_len_past_the_array(): + raw = bytearray( + BufferHandleDescriptor( + identity=_identity(), + address_space=AddressSpace.HOST, + access=AccessMode.READWRITE, + backend_kind=BackendKind.POSIX_SHM, + nbytes=8, + body=b"psm_x", + ).pack() + ) + raw[52:54] = (0xFFFF).to_bytes(2, "little") # body_len field + with pytest.raises(ValueError, match="body_len"): + BufferHandleDescriptor.unpack(bytes(raw)) + + +def test_worker_path_is_diagnostic_and_survives_an_unknown_id(): + h = BufferHandleDescriptor( + identity=_identity(), + address_space=AddressSpace.HOST, + access=AccessMode.READWRITE, + backend_kind=BackendKind.POSIX_SHM, + nbytes=8, + owner_worker_path_id=intern_worker_path("L4/L3[2]"), + ) + assert h.owner_worker_path == "L4/L3[2]" + assert BufferHandleDescriptor.unpack(h.pack()) == h + # An id minted in another process has no local text, and that is not an error. + foreign = replace(h, owner_worker_path_id=99_999) + assert foreign.owner_worker_path == "" + # The path takes no part in identity: two handles differing only by path are the same backing. + assert replace(h, owner_worker_path_id=0).identity == h.identity + + +def test_descriptor_rejects_oversized_body(): + with pytest.raises(ValueError, match="body"): + BufferHandleDescriptor( + identity=_identity(), + address_space=AddressSpace.HOST, + access=AccessMode.READWRITE, + backend_kind=BackendKind.POSIX_SHM, + nbytes=8, + body=b"x" * 200, # > DESC_MAX_BYTES + ).pack() + + +def test_create_export_import_resolve_zero_copy(): + oid = mint_owner_instance_id() + handle = create_host_shared_buffer(nbytes=256, owner_instance_id=oid, buffer_id=1, owner_worker_path="L4") + reg = ImportRegistry() + try: + assert handle.backend_kind == BackendKind.POSIX_SHM + imported = reg.materialize(handle.to_descriptor().pack()) + assert reg.resolve(handle.identity).base == imported.base + assert reg.materialize(handle.to_descriptor()).base == imported.base # map-once: same mapping + assert imported.nbytes == 256 + owner_shm = handle.shm + consumer_shm = imported.shm + assert owner_shm is not None + assert consumer_shm is not None + owner_buf = owner_shm.buf + consumer_buf = consumer_shm.buf + assert owner_buf is not None + assert consumer_buf is not None + owner_buf[:4] = b"\xde\xad\xbe\xef" + assert bytes(consumer_buf[:4]) == b"\xde\xad\xbe\xef" + finally: + reg.close() + handle.close() + + +def test_resolve_unregistered_raises(): + reg = ImportRegistry() + with pytest.raises(KeyError): + reg.resolve(_identity()) + + +def test_tensor_full_view_is_contiguous(): + oid = mint_owner_instance_id() + h = create_host_shared_buffer(nbytes=1024, owner_instance_id=oid, buffer_id=1) + try: + # handle.tensor(shape, dtype) is a contiguous full view: row-major strides, zero offset. + v = h.tensor(shapes=(4, 8), dtype=DataType.FLOAT32) + assert v.shapes == (4, 8) + assert v.strides == (8, 1) + assert v.ndims == 2 + assert v.byte_offset == 0 + # An explicit stride is carried verbatim; a singleton dim is never normalized away. + strided = h.tensor(shapes=(4, 1), dtype=DataType.FLOAT32, strides=(8, 3)) + assert strided.strides == (8, 3) + finally: + h.close() + + +def test_tensor_unpack_roundtrip(): + oid = mint_owner_instance_id() + h = create_host_shared_buffer(64, oid, buffer_id=1, owner_worker_path="L3") + try: + ref = h.tensor(shapes=(2, 4), dtype=DataType.FLOAT16, byte_offset=8) + assert Tensor.unpack(ref.pack()) == ref + finally: + h.close() + + +def test_re_export_preserves_identity_same_backing_no_map(): + # Frozen model §5/§8: canonical identity is invariant across every edge. Re-exporting an L4-owned + # backing for forwarding keeps the SOURCE identity (owner_instance_id / path / buffer_id / + # generation) and the same backing, only stripping the mapping. + l4 = mint_owner_instance_id() + src = create_host_shared_buffer(64, l4, buffer_id=7, owner_worker_path="L4") + try: + sdesc = src.to_descriptor() + hp = re_export(sdesc) + assert hp.identity.pack() == src.identity.pack() # identity invariant across the edge + assert hp.backend_kind == BackendKind.POSIX_SHM + assert hp.body == sdesc.body and hp.nbytes == 64 # same backing + assert hp.shm is None and hp.base == 0 # no map (lazy — a compute leaf maps) + # a ref built from H' carries the source identity + the same shm body, so L2 can materialize it + r = hp.tensor(shapes=(16,), dtype=DataType.FLOAT32) + assert Tensor.unpack(r.pack()).handle.identity.pack() == src.identity.pack() + finally: + src.close() + + +def test_device_malloc_wrap_materialize(): + # A device pointer (from orch.malloc) wrapped as DEVICE_MALLOC: materializes to the pointer with + # no map, address_space DEVICE (-> a child_memory Tensor). + oid = mint_owner_instance_id() + h = wrap_device_malloc(0xDEAD0000, 4096, oid, buffer_id=3, owner_worker_path="L3") + assert h.backend_kind == BackendKind.DEVICE_MALLOC + assert h.address_space == AddressSpace.DEVICE + assert h.shm is None and h.base == 0xDEAD0000 + reg = ImportRegistry() + imp = reg.materialize(h.to_descriptor()) + assert imp.base == 0xDEAD0000 + assert imp.address_space == AddressSpace.DEVICE + assert imp.shm is None + + +def test_materialize_remote_sidecar_rejected(): + desc = BufferHandleDescriptor( + identity=_identity(), + address_space=AddressSpace.HOST, + access=AccessMode.READWRITE, + backend_kind=BackendKind.REMOTE_SIDECAR, + nbytes=8, + ) + reg = ImportRegistry() + with pytest.raises(ValueError, match="REMOTE_SIDECAR"): + reg.materialize(desc) + + +def test_owner_instance_ids_are_distinct(): + ids = {mint_owner_instance_id() for _ in range(64)} + assert len(ids) == 64 + assert all(len(i) == OWNER_INSTANCE_ID_BYTES for i in ids) + + +@pytest.mark.parametrize( + "space,backend", + [ + (AddressSpace.HOST, BackendKind.VMM_WINDOW), + (AddressSpace.HOST, BackendKind.DEVICE_MALLOC), + (AddressSpace.DEVICE, BackendKind.FORK_SHM), + (AddressSpace.DEVICE, BackendKind.POSIX_SHM), + ], +) +def test_descriptor_rejects_bad_capability_combo(space, backend): + # §4.1 capability matrix: an unsupported address_space×backend_kind fails at construction (before + # dispatch, before it can ride the wire). + with pytest.raises(ValueError, match="capability"): + BufferHandleDescriptor( + identity=_identity(), + address_space=space, + access=AccessMode.READWRITE, + backend_kind=backend, + nbytes=64, + body=b"", + ) + + +def test_descriptor_accepts_legal_combos(): + for space, backend in [ + (AddressSpace.HOST, BackendKind.FORK_SHM), + (AddressSpace.HOST, BackendKind.POSIX_SHM), + (AddressSpace.DEVICE, BackendKind.VMM_WINDOW), + (AddressSpace.DEVICE, BackendKind.DEVICE_MALLOC), + (AddressSpace.HOST, BackendKind.REMOTE_SIDECAR), + (AddressSpace.DEVICE, BackendKind.REMOTE_SIDECAR), + ]: + BufferHandleDescriptor(_identity(), space, AccessMode.READWRITE, backend, 64, b"") + + +# --- the shared validator: blob decode rejects a malformed element -------------------------------- +# +# Decode is the boundary that sees untrusted bytes — shared memory a peer wrote. Every element +# extracted from a blob runs `validate_buffer_ref`, so a corrupted field is refused before any of +# it is trusted. + + +def _valid_packed_tensor() -> tuple[bytes, object]: + """A well-formed packed Tensor plus its owning handle (caller closes it).""" + h = create_host_shared_buffer(256, mint_owner_instance_id(), buffer_id=1) + return h.tensor(shapes=(8,), dtype=DataType.FLOAT32).pack(), h + + +def _decode_packed(packed: bytes) -> None: + """Decode a one-element blob built around `packed`, so the element crosses the receive boundary.""" + blob = struct.pack(" Date: Thu, 30 Jul 2026 05:48:46 -0700 Subject: [PATCH 2/2] Fix: close the FORK_COW write hole in the C++ validator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `validate_buffer_ref` is the one gate every receive boundary runs, but the FORK_COW rule lived only in the Python mirror's descriptor constructor. A descriptor arriving over the wire as FORK_COW with a write grant therefore passed the C++ decode: the consumer's first write would split the page into a private copy and the owner would observe nothing, with no error anywhere. Both mirrors of the ABI now reject it, and the test covers the WRITE and READWRITE grants as well as the FORK_SHM case that stays legal. `BufferRefBlobView::ref` also bounds-checks its index. Every caller loops to `ref_count` today, so nothing reaches past the validated region, but the method is the documented extraction point for three separate consumers and a future one has no way to see that the bound is the caller's job. Also: pin `FORK_COW = 5` alongside the other enumerators, which the frozen-value test had skipped; give the `BufferRef` layout comment the post-slimming sizes it describes (144 / 88, not 272 / 216 — the assertion below it was already right); add the missing `VMM_WINDOW` row to the backend table, which listed five of the six wire values; and document the two wire enums plus the handle and registry lifecycle calls, where the contract is not evident from the signature. --- docs/buffer-handle-abi.md | 1 + python/simpler/buffer_handle.py | 15 ++++++++++- src/common/task_interface/buffer_handle.h | 9 +++++++ tests/ut/cpp/types/test_buffer_handle.cpp | 33 +++++++++++++++++++++++ 4 files changed, 57 insertions(+), 1 deletion(-) diff --git a/docs/buffer-handle-abi.md b/docs/buffer-handle-abi.md index 2e5df6df10..097c2a72c4 100644 --- a/docs/buffer-handle-abi.md +++ b/docs/buffer-handle-abi.md @@ -226,6 +226,7 @@ gates, or keys on it. | `POSIX_SHM` | a named shm mapped into the consumer | `create_buffer` / `alloc_shared_tensor` | | `FORK_SHM` | the same VA, no map | a pre-fork `MAP_SHARED` host buffer (e.g. `share_memory_()`), writable from the child | | `FORK_COW` | the same VA, no map | a pre-fork plain host buffer: copy-on-write, so **READ only** | +| `VMM_WINDOW` | the device VA of the window carved by `allocate_domain`, no map | communication-domain window / buffer | | `DEVICE_MALLOC` | the device pointer, no map (chip-local) | `alloc_child_tensor` | | `REMOTE_SIDECAR` | (P2) resolved via the remote transport | an arg to a remote L3; the descriptor rides in the sidecar | diff --git a/python/simpler/buffer_handle.py b/python/simpler/buffer_handle.py index b78a5d99e4..ed53b30237 100644 --- a/python/simpler/buffer_handle.py +++ b/python/simpler/buffer_handle.py @@ -54,11 +54,15 @@ def _dtype_value(dtype: Any) -> int: class AddressSpace(enum.IntEnum): + """Which memory space a backing lives in. Wire values — mirrors the C++ enum.""" + HOST = 0 DEVICE = 1 class AccessMode(enum.IntEnum): + """What the holder of a handle may do with the backing. Wire values — mirrors the C++ enum.""" + READ = 0 WRITE = 1 READWRITE = 2 @@ -236,7 +240,7 @@ def unpack(cls, raw: bytes) -> BufferHandleDescriptor: ) -# BufferRef (272 B): BufferHandleDescriptor(216) + byte_offset u64, ndims u32, shapes[MAX] u32, +# BufferRef (144 B): BufferHandleDescriptor(88) + byte_offset u64, ndims u32, shapes[MAX] u32, # strides[MAX] u32, dtype u8, _pad[3]. _BUFFER_REF_TAIL = struct.Struct(f" BufferHandleDescriptor: + """The wire descriptor for this backing — what a consumer needs to resolve it.""" return BufferHandleDescriptor( identity=self.identity, address_space=self.address_space, @@ -384,6 +389,8 @@ def tensor( ) def close(self) -> None: + """Release the backing. The owner unlinks it, so a later consumer map fails rather than + resolving a name whose bytes are gone. Idempotent.""" if self.shm is not None: self.shm.close() self.shm.unlink() @@ -739,6 +746,8 @@ def mapped_args_from_blob(self, blob_ptr: int, capacity: int) -> MappedArgs: return MappedArgs(tensors, tuple(bufferref_blob_scalars(blob_ptr, capacity))) def resolve(self, identity: CanonicalIdentity) -> ImportedBuffer: + """The already-materialized import for ``identity``. Raises ``KeyError`` if this endpoint has + not materialized that backing — resolution never maps as a side effect.""" imported = self._by_identity.get(identity.pack()) if imported is None: raise KeyError(f"ImportRegistry: no handle registered for {identity}") @@ -749,11 +758,15 @@ def materialization_map(self) -> dict[bytes, tuple[int, int]]: return {key: (ib.base, int(ib.address_space)) for key, ib in self._by_identity.items()} def unregister(self, identity: CanonicalIdentity) -> None: + """Drop one import and close its mapping. The owner still holds the backing; only this + endpoint's view of it goes away. A no-op for an identity that was never materialized.""" imported = self._by_identity.pop(identity.pack(), None) if imported is not None and imported.shm is not None: imported.shm.close() def close(self) -> None: + """Close every mapping this endpoint made. Consumer-side only — unlinking belongs to the + owning Worker, so this never destroys a backing.""" for imported in self._by_identity.values(): if imported.shm is not None: imported.shm.close() diff --git a/src/common/task_interface/buffer_handle.h b/src/common/task_interface/buffer_handle.h index 78961da9df..8d7e8fae8d 100644 --- a/src/common/task_interface/buffer_handle.h +++ b/src/common/task_interface/buffer_handle.h @@ -256,6 +256,12 @@ inline void validate_buffer_ref(const BufferRef &r) { if (device_backend != device_space) reject("invalid Tensor: unsupported address_space x backend_kind"); } + // A copy-on-write page splits on the consumer's first write into a private copy the owner never + // sees, so a write grant over FORK_COW would be silently unobservable rather than an error. + if (backend == BackendKind::FORK_COW && h.access != static_cast(AccessMode::READ)) { + reject("invalid Tensor: FORK_COW grants READ only (a child's write would not reach the owner)"); + } + if (r.ndims == 0 || r.ndims > static_cast(MAX_TENSOR_DIMS)) reject("invalid Tensor: ndims out of range"); if (r.dtype >= DataType::DATA_TYPE_NUM) reject("invalid Tensor: unknown dtype"); const uint64_t elem = get_element_size(r.dtype); @@ -320,6 +326,9 @@ struct BufferRefBlobView { // its elements through here, so validating on extraction covers all three receive boundaries at // one point. Throws std::invalid_argument on a malformed element. BufferRef ref(int32_t i) const { + if (i < 0 || i >= ref_count) { + throw std::out_of_range("BufferRefBlobView::ref: index outside the blob's ref_count"); + } BufferRef r; std::memcpy(&r, ref_bytes + static_cast(i) * sizeof(BufferRef), sizeof(BufferRef)); validate_buffer_ref(r); diff --git a/tests/ut/cpp/types/test_buffer_handle.cpp b/tests/ut/cpp/types/test_buffer_handle.cpp index c856cba101..4a918f2fda 100644 --- a/tests/ut/cpp/types/test_buffer_handle.cpp +++ b/tests/ut/cpp/types/test_buffer_handle.cpp @@ -79,6 +79,7 @@ TEST(BufferHandleAbi, EnumValuesAreFrozen) { EXPECT_EQ(static_cast(BackendKind::VMM_WINDOW), 2); EXPECT_EQ(static_cast(BackendKind::REMOTE_SIDECAR), 3); EXPECT_EQ(static_cast(BackendKind::DEVICE_MALLOC), 4); + EXPECT_EQ(static_cast(BackendKind::FORK_COW), 5); } // --- memcpy round trip ------------------------------------------------------------------------- @@ -236,4 +237,36 @@ TEST(BufferRefBlob, RejectsNonZeroReserved) { EXPECT_THROW(read_bufferref_blob(buf.data(), buf.size()), std::runtime_error); } +// --- validate_buffer_ref: the shared receive-side gate ------------------------------------------ + +TEST(BufferHandleAbi, ForkCowGrantsReadOnly) { + BufferRef r = make_ref(); + r.handle.backend_kind = static_cast(BackendKind::FORK_COW); + r.handle.access = static_cast(AccessMode::READ); + EXPECT_NO_THROW(validate_buffer_ref(r)); + + // A write grant over copy-on-write would land in a private copy the owner never sees, so it is + // refused rather than silently dropped. + for (auto bad : {AccessMode::WRITE, AccessMode::READWRITE}) { + r.handle.access = static_cast(bad); + EXPECT_THROW(validate_buffer_ref(r), std::invalid_argument); + } + + // The same grants stay legal over MAP_SHARED, where a child's write does reach the owner. + r.handle.backend_kind = static_cast(BackendKind::FORK_SHM); + r.handle.access = static_cast(AccessMode::READWRITE); + EXPECT_NO_THROW(validate_buffer_ref(r)); +} + +TEST(BufferRefBlob, RefIndexIsBoundsChecked) { + BufferRef r = make_ref(); + std::vector buf(bufferref_blob_size(1, 0)); + write_bufferref_blob(buf.data(), &r, 1, nullptr, 0); + BufferRefBlobView view = read_bufferref_blob(buf.data(), buf.size()); + + EXPECT_NO_THROW(view.ref(0)); + EXPECT_THROW(view.ref(1), std::out_of_range); + EXPECT_THROW(view.ref(-1), std::out_of_range); +} + } // namespace