diff --git a/docs/README.md b/docs/README.md index 47b996a3d5..b6882ac316 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` | +| [Buffer Memory Model](buffer-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-abi.md b/docs/buffer-abi.md new file mode 100644 index 0000000000..ec997ef4a3 --- /dev/null +++ b/docs/buffer-abi.md @@ -0,0 +1,277 @@ +# Buffer / 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.h`](../src/common/task_interface/buffer.h) +— sizes, field offsets, and enum values all fail the build if they drift, and +[`tests/ut/cpp/types/test_buffer.cpp`](../tests/ut/cpp/types/test_buffer.cpp) +pins them again from the outside along with the blob codec. + +## Three types + +| Type | What it is | Where it lives | +| ---- | ---------- | -------------- | +| **`Buffer`** | 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 buffer 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 buffer 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. + +Both types have a member called `buffer`, and they are not the same thing: on the +wire it is the `BufferDescriptor` that says how to find the backing, on the device +it is the resolved `addr` + `size`. That is the whole of what materialization does. + +| `Tensor` (144 B, wire) | rel | `ChipTensor` (128 B, device) | +| ---------------------- | --- | ---------------------------- | +| `buffer.magic` | ⟂ | — | +| `buffer.identity` (32 B) | ⟂ | — | +| `buffer.backend_kind` | ⟂ | — | +| `buffer.body[32]` + `body_len` | ⟂ | — | +| `buffer.nbytes` | ≈ | `buffer.size` | +| `buffer.access` | ⟂ | — | +| `buffer.owner_worker_path_id` | ⟂ | — | +| — | ⟂ | `buffer.addr` | +| `byte_offset` (bytes) | ≈ | `start_offset` (elements) | +| `shapes[5]` / `strides[5]` / `ndims` / `dtype` | = | `shapes[5]` / `strides[5]` / `ndims` / `dtype` | +| `buffer.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: each named for what it is + +Both are called what they are, in every language. The wire type is +`simpler.task_interface.Tensor` in Python and `simpler::Tensor` in C++; the +device POD is `ChipTensor` in Python and the global `Tensor` in C++. + +The C++ split is a namespace rather than a rename because the device POD's name +is a cross-repo contract: kernels are authored against `__gm__ Tensor *`, and +several hundred files outside this repo dereference task args that way. A +namespace costs those files nothing, and it is also the stronger guard — a +translation unit that reaches for a bare `Tensor` while both are in scope fails +to compile with an ambiguity, where Python would silently bind whichever name +was imported last. + +**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 `buffer.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 `Buffer` 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 + +`buffer.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** | +| `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 | + +## 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 `buffer.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/examples/a2a3/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py b/examples/a2a3/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py index 0e58fb5783..73b9f6c6b5 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py +++ b/examples/a2a3/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py @@ -20,11 +20,11 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -134,7 +134,7 @@ def orch_fn(orch, _args, cfg): args.add_tensor(make_tensor_arg(out[rank]), TensorArgType.OUTPUT_EXISTING) args.add_tensor(make_tensor_arg(result[rank]), TensorArgType.OUTPUT_EXISTING) args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["notify_counter"], shapes=(1,), dtype=DataType.INT32, diff --git a/examples/a2a3/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py b/examples/a2a3/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py index d8039ecc3b..b7e2e4c165 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py +++ b/examples/a2a3/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py @@ -20,11 +20,11 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -138,7 +138,7 @@ def orch_fn(orch, _args, cfg): args = TaskArgs() args.add_tensor(make_tensor_arg(partial[rank]), TensorArgType.INPUT) args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["mailbox"], shapes=(N,), dtype=DataType.FLOAT32, @@ -148,7 +148,7 @@ def orch_fn(orch, _args, cfg): ) args.add_tensor(make_tensor_arg(result[rank]), TensorArgType.OUTPUT_EXISTING) args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["notify_counter"], shapes=(1,), dtype=DataType.INT32, diff --git a/examples/a2a3/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py b/examples/a2a3/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py index d8c55b8f5b..f4940a78f6 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py +++ b/examples/a2a3/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py @@ -28,11 +28,11 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -160,7 +160,7 @@ def orch_fn(orch, _args, cfg): domain = handle[rank] args = TaskArgs() args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["input_window"], shapes=(N,), dtype=DataType.FLOAT32, diff --git a/examples/a5/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py b/examples/a5/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py index 907399e568..a24a9328a2 100644 --- a/examples/a5/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py +++ b/examples/a5/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py @@ -20,11 +20,11 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -131,7 +131,7 @@ def orch_fn(orch, _args, cfg): args.add_tensor(make_tensor_arg(out[rank]), TensorArgType.OUTPUT_EXISTING) args.add_tensor(make_tensor_arg(result[rank]), TensorArgType.OUTPUT_EXISTING) args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["notify_counter"], shapes=(1,), dtype=DataType.INT32, diff --git a/examples/a5/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py b/examples/a5/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py index dfd2c7b570..bfb0da0e17 100644 --- a/examples/a5/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py +++ b/examples/a5/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py @@ -20,11 +20,11 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -138,7 +138,7 @@ def orch_fn(orch, _args, cfg): args = TaskArgs() args.add_tensor(make_tensor_arg(partial[rank]), TensorArgType.INPUT) args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["mailbox"], shapes=(N,), dtype=DataType.FLOAT32, @@ -148,7 +148,7 @@ def orch_fn(orch, _args, cfg): ) args.add_tensor(make_tensor_arg(result[rank]), TensorArgType.OUTPUT_EXISTING) args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["notify_counter"], shapes=(1,), dtype=DataType.INT32, diff --git a/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py b/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py index 0465e5a1ef..0833fd8038 100644 --- a/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py +++ b/examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py @@ -28,11 +28,11 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -159,7 +159,7 @@ def orch_fn(orch, _args, cfg): domain = handle[rank] args = TaskArgs() args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["input_window"], shapes=(N,), dtype=DataType.FLOAT32, diff --git a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py index ab3e6df166..75b9c62c7f 100644 --- a/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py +++ b/examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py @@ -28,11 +28,11 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -180,7 +180,7 @@ def orch_fn(orch, _args, cfg): domain = handle[rank] args = TaskArgs() args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["input_window"], shapes=(N,), dtype=DataType.FLOAT32, diff --git a/examples/workers/l2/per_task_runtime_env/main.py b/examples/workers/l2/per_task_runtime_env/main.py index c25ba099a4..f84eaef720 100644 --- a/examples/workers/l2/per_task_runtime_env/main.py +++ b/examples/workers/l2/per_task_runtime_env/main.py @@ -45,9 +45,9 @@ CallConfig, ChipCallable, ChipStorageTaskArgs, + ChipTensor, CoreCallable, DataType, - Tensor, ) from simpler.worker import Worker @@ -164,9 +164,9 @@ def _run_one(worker: Worker, chip_handle, label: str, ring: Optional[dict]) -> N worker.copy_to(dev_b, host_b.data_ptr(), NBYTES) args = ChipStorageTaskArgs() - args.add_tensor(Tensor.make(dev_a, (N_ROWS, N_COLS), DataType.FLOAT32)) - args.add_tensor(Tensor.make(dev_b, (N_ROWS, N_COLS), DataType.FLOAT32)) - args.add_tensor(Tensor.make(dev_out, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_a, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_b, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_out, (N_ROWS, N_COLS), DataType.FLOAT32)) config = _make_config(ring) print(f"[per_task_runtime_env] run '{label}': runtime_env={config.runtime_env!r}") diff --git a/examples/workers/l2/vector_add/main.py b/examples/workers/l2/vector_add/main.py index b7d6ec1f68..15fde1ea0a 100644 --- a/examples/workers/l2/vector_add/main.py +++ b/examples/workers/l2/vector_add/main.py @@ -44,9 +44,9 @@ CallConfig, ChipCallable, ChipStorageTaskArgs, + ChipTensor, CoreCallable, DataType, - Tensor, ) from simpler.worker import Worker @@ -145,12 +145,12 @@ def _run(worker: Worker, chip_handle: CallableHandle): worker.copy_to(dev_b, host_b.data_ptr(), NBYTES) # --- 3. Build TaskArgs describing the tensors visible to the orchestration --- - # Each tensor is a Tensor(data_ptr, shape, dtype). Order must + # Each tensor is a ChipTensor(data_ptr, shape, dtype). Order must # match the ``signature`` list in the ChipCallable (IN, IN, OUT). args = ChipStorageTaskArgs() - args.add_tensor(Tensor.make(dev_a, (N_ROWS, N_COLS), DataType.FLOAT32)) - args.add_tensor(Tensor.make(dev_b, (N_ROWS, N_COLS), DataType.FLOAT32)) - args.add_tensor(Tensor.make(dev_out, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_a, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_b, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_out, (N_ROWS, N_COLS), DataType.FLOAT32)) # --- 4. Run. CallConfig() defaults are fine for this kernel. --- config = CallConfig() diff --git a/examples/workers/l2/vector_add/test_run_timing.py b/examples/workers/l2/vector_add/test_run_timing.py index 1b190689f5..b71e126f17 100644 --- a/examples/workers/l2/vector_add/test_run_timing.py +++ b/examples/workers/l2/vector_add/test_run_timing.py @@ -27,7 +27,7 @@ import pytest from simpler._log import get_current_config -from simpler.task_interface import CallConfig, ChipStorageTaskArgs, DataType, Tensor +from simpler.task_interface import CallConfig, ChipStorageTaskArgs, ChipTensor, DataType from simpler.worker import Worker from simpler_setup.log_config import configure_logging @@ -72,9 +72,9 @@ def _drive_one_run(platform: str, device_id: int, *, enable_l2_swimlane: bool = worker.copy_to(dev_b, host_b.data_ptr(), NBYTES) args = ChipStorageTaskArgs() - args.add_tensor(Tensor.make(dev_a, (N_ROWS, N_COLS), DataType.FLOAT32)) - args.add_tensor(Tensor.make(dev_b, (N_ROWS, N_COLS), DataType.FLOAT32)) - args.add_tensor(Tensor.make(dev_out, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_a, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_b, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_out, (N_ROWS, N_COLS), DataType.FLOAT32)) config = CallConfig() config.enable_l2_swimlane = enable_l2_swimlane diff --git a/examples/workers/l3/allreduce/main.py b/examples/workers/l3/allreduce/main.py index 945f7b5f97..339553c897 100644 --- a/examples/workers/l3/allreduce/main.py +++ b/examples/workers/l3/allreduce/main.py @@ -34,11 +34,11 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker # noqa: E402 @@ -158,7 +158,7 @@ def orch_fn(orch, _args, cfg): chip_args.add_tensor(make_tensor_arg(host_inputs[i]), TensorArgType.INPUT) chip_args.add_tensor(make_tensor_arg(host_outputs[i]), TensorArgType.OUTPUT_EXISTING) chip_args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["scratch"], shapes=(float_elems,), dtype=DataType.FLOAT32, diff --git a/examples/workers/l3/child_memory/main.py b/examples/workers/l3/child_memory/main.py index 470847f8c3..346ef0c2b5 100644 --- a/examples/workers/l3/child_memory/main.py +++ b/examples/workers/l3/child_memory/main.py @@ -14,7 +14,7 @@ * ``orch.malloc(worker_id=0, nbytes)`` — allocate a buffer that lives on the chip child for as long as the chip is alive. * ``orch.copy_to(worker_id=0, dev, host, n)`` — H2D upload of the weight. - * ``Tensor.make(dev_ptr, shape, dtype, child_memory=True)`` — + * ``ChipTensor.make(dev_ptr, shape, dtype, child_memory=True)`` — wrap the worker pointer as a tensor that the runtime treats as *already on device*. ``init_runtime_impl`` skips malloc + H2D copy for these and does not record them in ``tensor_pairs``, so the buffer @@ -29,7 +29,7 @@ * ``orch.malloc / orch.copy_to`` — control-plane device-memory ops that forward to the chip child via mailbox IPC. - * ``Tensor.make(..., child_memory=True)`` — opt-out of the + * ``ChipTensor.make(..., child_memory=True)`` — opt-out of the runtime's auto-malloc + auto-free for a tensor whose lifetime you manage yourself. @@ -54,10 +54,10 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -165,7 +165,7 @@ def orch_fn(orch, _args, cfg): # on the worker; do NOT auto-malloc, do NOT auto-free at end-of-task". # Without this flag the first task's teardown would free dev_w and # the second task would read freed memory. - w_dev = Tensor.make(dev_w, (SIZE,), DataType.FLOAT32, child_memory=True) + w_dev = ChipTensor.make(dev_w, (SIZE,), DataType.FLOAT32, child_memory=True) # Two kernel invocations sharing the weight, pinned to chip 0. for out in (host_f1, host_f2): diff --git a/examples/workers/l3/domain_rank_map/main.py b/examples/workers/l3/domain_rank_map/main.py index 801b54302f..1cd0a65d9f 100644 --- a/examples/workers/l3/domain_rank_map/main.py +++ b/examples/workers/l3/domain_rank_map/main.py @@ -42,11 +42,11 @@ CallConfig, ChipCallable, ChipDomainContext, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker # noqa: E402 @@ -118,7 +118,7 @@ def build_allreduce_callable(platform: str) -> ChipCallable: def _add_domain_scratch(args: TaskArgs, domain: ChipDomainContext) -> None: args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["scratch"], shapes=(COUNT,), dtype=DataType.FLOAT32, diff --git a/examples/workers/l3/dual_domain_overlap/main.py b/examples/workers/l3/dual_domain_overlap/main.py index 98ad5504c0..f3e380b9c5 100644 --- a/examples/workers/l3/dual_domain_overlap/main.py +++ b/examples/workers/l3/dual_domain_overlap/main.py @@ -47,11 +47,11 @@ CallConfig, ChipCallable, ChipDomainContext, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker # noqa: E402 @@ -164,7 +164,7 @@ def _scratch_buffers() -> list[CommBufferSpec]: def _add_domain_scratch(args: TaskArgs, domain: ChipDomainContext) -> None: args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["scratch"], shapes=(COUNT,), dtype=DataType.FLOAT32, diff --git a/examples/workers/l3/ep_dispatch_combine/main.py b/examples/workers/l3/ep_dispatch_combine/main.py index b9e3fe0d33..cbc3b94e45 100644 --- a/examples/workers/l3/ep_dispatch_combine/main.py +++ b/examples/workers/l3/ep_dispatch_combine/main.py @@ -76,11 +76,11 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker # noqa: E402 @@ -568,7 +568,7 @@ def orch_fn(orch, _args, cfg): chip_args.add_tensor(make_tensor_arg(recv_y_outs[i]), TensorArgType.OUTPUT_EXISTING) chip_args.add_tensor(make_tensor_arg(routed_y_outs[i]), TensorArgType.OUTPUT_EXISTING) chip_args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["scratch"], shapes=(SCRATCH_NBYTES // 4,), dtype=DataType.FLOAT32, diff --git a/examples/workers/l3/ffn_tp_parallel/main.py b/examples/workers/l3/ffn_tp_parallel/main.py index c005db185e..d95d9434ec 100644 --- a/examples/workers/l3/ffn_tp_parallel/main.py +++ b/examples/workers/l3/ffn_tp_parallel/main.py @@ -42,11 +42,11 @@ ArgDirection, CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker # noqa: E402 @@ -224,7 +224,7 @@ def orch_fn(orch, _args, cfg): a2.add_tensor(make_tensor_arg(host_partial[i]), TensorArgType.INPUT) a2.add_tensor(make_tensor_arg(host_y[i]), TensorArgType.OUTPUT_EXISTING) a2.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["scratch"], shapes=(scratch_count,), dtype=DataType.FLOAT32, diff --git a/mkdocs.yml b/mkdocs.yml index c9ac8b3c72..bbb2e6d6a9 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 + - Buffer Memory Model: buffer-abi.md - Orchestrator: orchestrator.md - Scheduler: scheduler.md - A5 AICPU core selection: design/a5-fg-pg-core-selection.en.md diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 01fbc7f49b..09f11ca25e 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -50,6 +50,7 @@ #include #include "arg_direction.h" +#include "buffer.h" #include "callable.h" #include "callable_protocol.h" #include "chip_worker.h" @@ -946,10 +947,21 @@ 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)); + // Buffer / simpler::Tensor wire ABI (buffer.h). Exported so the Python mirror in + // simpler.buffer can pin its struct formats to the C++ layout and reject drift. + m.attr("BUFFER_DESCRIPTOR_MAGIC") = static_cast(simpler::BUFFER_DESCRIPTOR_MAGIC); + m.attr("TENSOR_BLOB_MAGIC") = static_cast(simpler::TENSOR_BLOB_MAGIC); + m.attr("WIRE_TENSOR_BYTES") = static_cast(sizeof(simpler::Tensor)); + m.attr("BUFFER_DESCRIPTOR_BYTES") = static_cast(sizeof(simpler::BufferDescriptor)); + m.attr("CANONICAL_IDENTITY_BYTES") = static_cast(sizeof(simpler::CanonicalIdentity)); + m.attr("OWNER_INSTANCE_ID_BYTES") = static_cast(simpler::OWNER_INSTANCE_ID_BYTES); + m.attr("DESC_MAX_BYTES") = static_cast(simpler::DESC_MAX_BYTES); + m.attr("TENSOR_BLOB_HEADER_BYTES") = static_cast(simpler::TENSOR_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. - nb::class_(m, "Tensor") + nb::class_(m, "ChipTensor") .def(nb::init<>()) .def_static( @@ -1079,7 +1091,7 @@ NB_MODULE(_task_interface, m) { .def("__repr__", [](const Tensor &self) -> std::string { std::ostringstream os; - os << "Tensor(data=0x" << std::hex << self.buffer.addr << std::dec << ", shapes=("; + os << "ChipTensor(data=0x" << std::hex << self.buffer.addr << std::dec << ", shapes=("; for (uint32_t i = 0; i < self.ndims; ++i) { if (i) os << ", "; os << self.shapes[i]; @@ -2016,6 +2028,57 @@ NB_MODULE(_task_interface, m) { "Tags are not preserved (blob wire format strips them)." ); + // simpler::Tensor blob readers. Each validates every element it extracts (simpler::read_tensor_blob bounds + // the header, simpler::TensorBlobView::ref validates the element), so a Python consumer never walks + // the layout itself. + m.def( + "tensor_blob_descriptors", + [](uint64_t blob_ptr, size_t capacity) -> nb::list { + const uint8_t *src = reinterpret_cast(blob_ptr); + simpler::TensorBlobView view = simpler::read_tensor_blob(src, capacity); + nb::list out; + for (int32_t i = 0; i < view.tensor_count; i++) { + simpler::Tensor r = view.tensor(i); + out.append(nb::bytes(reinterpret_cast(&r.buffer), sizeof(simpler::BufferDescriptor))); + } + return out; + }, + nb::arg("blob_ptr"), nb::arg("capacity"), + "Extract each simpler::Tensor's embedded simpler::BufferDescriptor (packed bytes) from a simpler::Tensor " + "blob, in ref order. A consumer materializes these lazily on receipt." + ); + + m.def( + "tensor_blob_tensors", + [](uint64_t blob_ptr, size_t capacity) -> nb::list { + const uint8_t *src = reinterpret_cast(blob_ptr); + simpler::TensorBlobView view = simpler::read_tensor_blob(src, capacity); + nb::list out; + for (int32_t i = 0; i < view.tensor_count; i++) { + simpler::Tensor r = view.tensor(i); + out.append(nb::bytes(reinterpret_cast(&r), sizeof(simpler::Tensor))); + } + return out; + }, + nb::arg("blob_ptr"), nb::arg("capacity"), + "Extract each full packed simpler::Tensor (descriptor + view) from a simpler::Tensor blob, in ref order." + ); + + m.def( + "tensor_blob_scalars", + [](uint64_t blob_ptr, size_t capacity) -> nb::list { + const uint8_t *src = reinterpret_cast(blob_ptr); + simpler::TensorBlobView view = simpler::read_tensor_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 simpler::Tensor 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.py b/python/simpler/buffer.py new file mode 100644 index 0000000000..1afcc84ad8 --- /dev/null +++ b/python/simpler/buffer.py @@ -0,0 +1,773 @@ +# 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 Buffer ABI (Python mirror of ``buffer.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.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_BYTES, + BUFFER_DESCRIPTOR_MAGIC, + CANONICAL_IDENTITY_BYTES, + DESC_MAX_BYTES, + MAX_TENSOR_DIMS, + OWNER_INSTANCE_ID_BYTES, + TENSOR_BLOB_HEADER_BYTES, + TENSOR_BLOB_MAGIC, + WIRE_TENSOR_BYTES, + DataType, + tensor_blob_descriptors, + tensor_blob_scalars, + tensor_blob_tensors, +) + + +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): + """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 + + +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.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") +# BufferDescriptor (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 BufferDescriptor: + """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) -> BufferDescriptor: + if len(raw) < BUFFER_DESCRIPTOR_BYTES: + raise ValueError(f"descriptor too small: {len(raw)} < {BUFFER_DESCRIPTOR_BYTES}") + magic, address_space, access, backend_kind = _DESC_PREFIX.unpack_from(raw, 0) + if magic != BUFFER_DESCRIPTOR_MAGIC: + raise ValueError(f"not a BufferDescriptor: 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]), + ) + + +# Wire Tensor (144 B): BufferDescriptor(88) + byte_offset u64, ndims u32, shapes[MAX] u32, +# strides[MAX] u32, dtype u8, _pad[3]. +_WIRE_TENSOR_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 ``Buffer``, carrying that handle's descriptor. + + Self-describing — the consumer materializes ``handle`` on first receipt (no prior handshake), + keyed by ``buffer.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). + """ + + buffer: BufferDescriptor + 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 = _WIRE_TENSOR_TAIL.pack(self.byte_offset, ndims, *shapes, *strides, self.dtype) + return self.buffer.pack() + tail + + @classmethod + def unpack(cls, raw: bytes) -> Tensor: + handle = BufferDescriptor.unpack(raw[:BUFFER_DESCRIPTOR_BYTES]) + vals = _WIRE_TENSOR_TAIL.unpack(raw[BUFFER_DESCRIPTOR_BYTES:WIRE_TENSOR_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(buffer=handle, byte_offset=byte_offset, shapes=shapes, strides=strides, dtype=dtype) + + +def pack_tensor_blob(tensors: list[Tensor], scalars: tuple[int, ...] = ()) -> bytes: + """Serialize tensors + scalars into the versioned wire blob (mirror of write_tensor_blob).""" + header = _TENSOR_BLOB_HEADER.pack(TENSOR_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 Buffer: + """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) -> BufferDescriptor: + """The wire descriptor for this backing — what a consumer needs to resolve it.""" + return BufferDescriptor( + 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) — ``buffer.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( + buffer=self.to_descriptor(), + byte_offset=byte_offset, + shapes=shapes, + strides=strides, + dtype=_dtype_value(dtype), + ) + + 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() + 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, +) -> Buffer: + """Allocate a POSIX-shm host backing and wrap it as an owner ``Buffer`` (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 Buffer( + 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: BufferDescriptor) -> Buffer: + """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-tensor map cost. + """ + return Buffer( + 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 = BufferDescriptor( + 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( + buffer=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, +) -> Buffer: + """Wrap a pre-fork, fork-inherited host allocation as a zero-copy ``Buffer``. + + 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 Buffer( + 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, +) -> Buffer: + """Wrap a device pointer (from a worker device malloc) as a ``DEVICE_MALLOC`` ``Buffer``. + + 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 tensor over this buffer 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 Buffer( + 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, +) -> Buffer: + """Wrap a domain-window-carved device VA as a ``VMM_WINDOW`` ``Buffer``. + + 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 tensor 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 Buffer( + 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 tensor'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: BufferDescriptor | bytes) -> ImportedBuffer: + """Map ``descriptor``'s backing into this process on first sight of its identity; reuse the + cached ImportedBuffer thereafter (map-once).""" + desc = BufferDescriptor.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 tensor 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_tensor_blob``: packed identity -> (local base, address_space).""" + for desc_bytes in tensor_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, t in enumerate(Tensor.unpack(tb) for tb in tensor_blob_tensors(blob_ptr, capacity)): + if t.buffer.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"({t.buffer.backend_kind.name}); it cannot be mapped into a host process" + ) + tensors.append(MappedArg(self.materialize(t.buffer), t.byte_offset, t.shapes, t.strides, t.dtype)) + return MappedArgs(tensors, tuple(tensor_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}") + return imported + + def materialization_map(self) -> dict[bytes, tuple[int, int]]: + """Snapshot for ``materialize_tensor_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: + """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() + self._by_identity.clear() diff --git a/python/simpler/l3_l2_message_queue.py b/python/simpler/l3_l2_message_queue.py index 656888a9f7..4da601256e 100644 --- a/python/simpler/l3_l2_message_queue.py +++ b/python/simpler/l3_l2_message_queue.py @@ -22,7 +22,7 @@ NotifyOp, WaitCmp, ) -from .task_interface import DataType, Tensor +from .task_interface import ChipTensor, DataType L3L2_QUEUE_MAGIC = 0x4C335132 L3L2_QUEUE_ABI_MAJOR = 1 @@ -173,7 +173,7 @@ def _host_byte_span(buffer: Any, nbytes: int, *, writable: bool) -> _HostByteSpa return _HostByteSpan(nbytes=nbytes, ptr=ptr, view=None) access = "writable" if writable else "readable" - raise ValueError(f"L3-L2 queue requires a registered Tensor or {access} contiguous ordinary host buffer") + raise ValueError(f"L3-L2 queue requires a registered ChipTensor or {access} contiguous ordinary host buffer") def make_l3_l2_queue_layout(depth: int, input_arena_bytes: int, output_arena_bytes: int) -> L3L2QueueLayout: @@ -258,9 +258,9 @@ def __init__( orch: Any, region: L3L2OrchRegion, layout: L3L2QueueLayout, - desc_fields: Tensor, - desc_seq: Tensor, - desc_read: Tensor, + desc_fields: ChipTensor, + desc_seq: ChipTensor, + desc_read: ChipTensor, ) -> None: self._orch = orch self._region = region @@ -331,16 +331,16 @@ def _ensure_live(self) -> None: raise RuntimeError("L3-L2 queue expired after orchestration run") self._region._ensure_live() - def _validate_registered_buffer(self, buffer: Any, nbytes: int) -> Tensor: - if not isinstance(buffer, Tensor): - raise ValueError("L3-L2 queue requires a registered Tensor returned by orch.alloc(...)") + def _validate_registered_buffer(self, buffer: Any, nbytes: int) -> ChipTensor: + if not isinstance(buffer, ChipTensor): + raise ValueError("L3-L2 queue requires a registered ChipTensor returned by orch.alloc(...)") self._region._validate_host_buffer(buffer) if int(nbytes) > int(buffer.nbytes()): - raise ValueError(f"L3-L2 queue nbytes={nbytes} exceeds registered Tensor size {int(buffer.nbytes())}") + raise ValueError(f"L3-L2 queue nbytes={nbytes} exceeds registered ChipTensor size {int(buffer.nbytes())}") return buffer - def _registered_buffer_or_none(self, buffer: Any, nbytes: int) -> Tensor | None: - if not isinstance(buffer, Tensor): + def _registered_buffer_or_none(self, buffer: Any, nbytes: int) -> ChipTensor | None: + if not isinstance(buffer, ChipTensor): return None return self._validate_registered_buffer(buffer, nbytes) diff --git a/python/simpler/l3_l2_orch_comm.py b/python/simpler/l3_l2_orch_comm.py index e4c89fa0e0..ecef464b29 100644 --- a/python/simpler/l3_l2_orch_comm.py +++ b/python/simpler/l3_l2_orch_comm.py @@ -25,7 +25,7 @@ _l3_host_mapped_region_close, ) -from .task_interface import Tensor +from .task_interface import ChipTensor class NotifyOp(IntEnum): @@ -222,7 +222,7 @@ def validate_region_create_reply( class _PinnedBuffer: def __init__(self, obj: Any, *, writable: bool = False) -> None: self._keepalive: Any = obj - if isinstance(obj, Tensor): + if isinstance(obj, ChipTensor): if obj.child_memory: raise ValueError("L3-L2 payload buffer must be host storage, not child_memory device storage") if not obj.is_contiguous: diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index e341a871d5..bab9e89858 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -46,12 +46,12 @@ def my_orch(orch, args, cfg): from .task_interface import ( CallConfig, ChipCallable, + ChipTensor, CommBufferSpec, CommDomainHandle, DataType, RemoteAddressSpace, TaskArgs, - Tensor, _empty_remote_sidecar_for, _remote_sidecar_for, _RemoteTaskArgsSidecar, @@ -667,10 +667,10 @@ def copy_from(self, worker_id: int, dst: int, src: int, size: int) -> None: self._worker._child_prov_require_live_range(wid, s, int(size), api="copy_from") self._o.copy_from(wid, int(dst), s, int(size)) - def alloc(self, shape: Sequence[int], dtype: DataType) -> Tensor: + def alloc(self, shape: Sequence[int], dtype: DataType) -> ChipTensor: """Allocate a runtime-managed intermediate buffer. - Returns a ``Tensor`` whose backing memory comes from a + Returns a ``ChipTensor`` whose backing memory comes from a per-allocation MAP_SHARED mmap (visible to forked child workers). Lifetime is bound to a synthetic task slot that the Orchestrator treats as the buffer's producer; the buffer is freed when all diff --git a/python/simpler/remote_l3_protocol.py b/python/simpler/remote_l3_protocol.py index 24ca8ab402..02507da55f 100644 --- a/python/simpler/remote_l3_protocol.py +++ b/python/simpler/remote_l3_protocol.py @@ -15,7 +15,7 @@ import struct from dataclasses import dataclass -from .task_interface import MAX_TENSOR_DIMS, CallConfig, DataType, Tensor +from .task_interface import MAX_TENSOR_DIMS, CallConfig, ChipTensor, DataType # 2: CallConfig lost its block_dim field — a run always takes the whole # device, so the payload is one int32 shorter than v1's. @@ -143,7 +143,7 @@ class RemoteTensorSidecar: @dataclass(frozen=True) class RemoteTaskArgsWire: - tensor_metadata: tuple[Tensor, ...] + tensor_metadata: tuple[ChipTensor, ...] remote_desc: tuple[RemoteTensorSidecar, ...] scalars: tuple[int, ...] inline_payload: bytes @@ -428,7 +428,7 @@ def decode_call_config(reader: _Reader) -> CallConfig: return cfg -def decode_tensor(reader: _Reader) -> Tensor: +def decode_tensor(reader: _Reader) -> ChipTensor: data = reader.u64() if data != 0: raise ValueError("remote_wire: remote TASK tensor data must be zero") @@ -442,8 +442,8 @@ def decode_tensor(reader: _Reader) -> Tensor: raise ValueError("remote_wire: tensor child_memory must be 0 or 1") for _ in range(7): if reader.u8() != 0: - raise ValueError("remote_wire: Tensor reserved bytes must be zero") - return Tensor.make(0, tuple(shapes[:ndims]), dtype, bool(child_memory)) + raise ValueError("remote_wire: ChipTensor reserved bytes must be zero") + return ChipTensor.make(0, tuple(shapes[:ndims]), dtype, bool(child_memory)) def decode_remote_tensor_desc(reader: _Reader) -> RemoteTensorDesc: diff --git a/python/simpler/remote_l3_session.py b/python/simpler/remote_l3_session.py index c62a613779..20739ced86 100644 --- a/python/simpler/remote_l3_session.py +++ b/python/simpler/remote_l3_session.py @@ -74,7 +74,7 @@ read_frame, send_frame, ) -from .task_interface import ChipCallable, TaskArgs, Tensor +from .task_interface import ChipCallable, ChipTensor, TaskArgs from .worker import Worker sys.modules.setdefault("simpler.remote_l3_session", sys.modules[__name__]) @@ -472,8 +472,8 @@ def _buffer_key(buffer_id: int, generation: int) -> tuple[int, int]: return int(buffer_id), int(generation) -def _tensor_with_data(tensor: Tensor, data: int) -> Tensor: - return Tensor.make(int(data), tuple(tensor.shapes), tensor.dtype, bool(tensor.child_memory)) +def _tensor_with_data(tensor: ChipTensor, data: int) -> ChipTensor: + return ChipTensor.make(int(data), tuple(tensor.shapes), tensor.dtype, bool(tensor.child_memory)) def _materialize_task_args( # noqa: PLR0912 diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index 5133e63ac9..63d873a2ee 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -9,13 +9,16 @@ # ruff: noqa: PLW0603, PLC0415 """Public Python API for task_interface nanobind bindings. -Re-exports the canonical C++ types (DataType, Tensor, ChipStorageTaskArgs, -TaskArgs, TensorArgType) plus ``scalar_to_uint64``. Torch-aware helpers -(``make_tensor_arg``, ``torch_dtype_to_datatype``) live in +``Tensor`` is the task argument you build and submit — a buffer handle descriptor plus a strided +view, carrying no address. ``ChipTensor`` is the device POD the L2 runtime ABI reads, which the +framework materializes from a ``Tensor`` on the consuming endpoint; you rarely name it directly. + +Also re-exports DataType, ChipStorageTaskArgs, TaskArgs, TensorArgType and ``scalar_to_uint64``. +Torch-aware helpers (``make_tensor_arg``, ``torch_dtype_to_datatype``) live in ``simpler_setup.torch_interop`` — this module has no torch dependency. Usage: - from simpler.task_interface import DataType, Tensor, ChipStorageTaskArgs + from simpler.task_interface import DataType, Tensor, TensorArgType from simpler_setup.torch_interop import make_tensor_arg """ @@ -50,12 +53,12 @@ CallConfig, ChipCallable, ChipStorageTaskArgs, + ChipTensor, CoreCallable, DataType, RuntimeEnv, TaskArgs, TaskState, - Tensor, TensorArgType, WorkerType, _ChipWorker, @@ -66,6 +69,8 @@ read_args_from_blob, ) +from .buffer import Tensor + def _assert_bindings_match_source_tree() -> None: """Refuse a `_task_interface` built from a different revision of this tree. @@ -136,6 +141,7 @@ def _assert_bindings_match_source_tree() -> None: "get_element_size", "get_dtype_name", "MAX_TENSOR_DIMS", + "ChipTensor", "Tensor", "ChipStorageTaskArgs", "TensorArgType", @@ -766,16 +772,16 @@ def _storage_for_remote_task_args(args: TaskArgs) -> _RemoteTaskArgsStorage: def _task_args_add_tensor( - self: TaskArgs, tensor: Tensor | RemoteTensorRef, tag: TensorArgType = TensorArgType.INPUT + self: TaskArgs, tensor: ChipTensor | RemoteTensorRef, tag: TensorArgType = TensorArgType.INPUT ) -> None: if isinstance(tensor, RemoteTensorRef): storage = _storage_for_remote_task_args(self) - metadata = Tensor.make(0, tensor.shape, tensor.dtype) + metadata = ChipTensor.make(0, tensor.shape, tensor.dtype) _TASK_ARGS_ADD_TENSOR(self, metadata, tag) storage.sidecars.append(_sidecar_from_ref(storage, tensor)) return - if not isinstance(tensor, Tensor): - raise TypeError("TaskArgs.add_tensor expects Tensor or RemoteTensorRef") + if not isinstance(tensor, ChipTensor): + raise TypeError("TaskArgs.add_tensor expects ChipTensor or RemoteTensorRef") _TASK_ARGS_ADD_TENSOR(self, tensor, tag) with _REMOTE_TASK_ARGS_STORAGE_LOCK: storage = _REMOTE_TASK_ARGS_STORAGE.get(self) diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 9487d92882..e6551830be 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -100,6 +100,11 @@ def my_l4_orch(orch, args, config): ) from . import _log as _simpler_log +from .buffer import ( + Buffer, + create_host_shared_buffer, + mint_owner_instance_id, +) from .callable_identity import ( CALLABLE_HASH_DIGEST_BYTES, CallableHandle, @@ -139,6 +144,7 @@ def my_l4_orch(orch, args, config): CallConfig, ChipCallable, ChipDomainContext, + ChipTensor, ChipWorker, CommBufferSpec, CommDomainHandle, @@ -146,7 +152,6 @@ def my_l4_orch(orch, args, config): RemoteBufferExport, RemoteBufferHandle, TaskArgs, - Tensor, _Worker, ) @@ -184,12 +189,12 @@ def my_l4_orch(orch, args, config): _CFG_FMT = struct.Struct("=iiiiii" + ("Q" * _RUNTIME_ENV_UINT64_FIELD_COUNT) + "1024s") # The generation-safe pipeline lease follows CONFIG. Args start after the # lease, rounded up to 8 bytes so the first -# Tensor.data (uint64_t at OFF_ARGS+8) is 8-byte aligned, avoiding +# ChipTensor.data (uint64_t at OFF_ARGS+8) is 8-byte aligned, avoiding # SIGBUS on strict-alignment platforms (aarch64 atomics, some ARM cores). _PIPELINE_LEASE_FMT = struct.Struct("=IIQ") _OFF_PIPELINE_LEASE = (_OFF_CONFIG + _CFG_FMT.size + 7) & ~7 _OFF_ARGS = (_OFF_PIPELINE_LEASE + _PIPELINE_LEASE_FMT.size + 7) & ~7 -assert _OFF_ARGS % 8 == 0, "_OFF_ARGS must be 8-aligned for Tensor.data" +assert _OFF_ARGS % 8 == 0, "_OFF_ARGS must be 8-aligned for ChipTensor.data" _OFF_TASK_CALLABLE_HASH = _OFF_ARGS _OFF_TASK_ARGS_BLOB = _OFF_TASK_CALLABLE_HASH + CALLABLE_HASH_DIGEST_BYTES # MAILBOX_ARGS_CAPACITY mirrors the C++ constexpr in worker_manager.h so the @@ -334,9 +339,9 @@ def _local_task_frame_count(platform: str, _runtime: str, pipeline_depth: int) - _HOST_BUF_MAP_HEADER = struct.Struct(" TaskArgs: to C++ run use the zero-copy `run_from_blob` path instead — see those loops for the matching comment. - Delegates to the nanobind helper so the Tensor layout is + Delegates to the nanobind helper so the ChipTensor layout is parsed by C++ `read_blob` (single source of truth) instead of being reimplemented in Python. The Python re-implementation that lived here previously dropped the `child_memory` byte (offset 33), which @@ -3710,6 +3715,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; `_buffers` 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._buffers: dict[int, Buffer] = {} 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)) @@ -6747,8 +6759,8 @@ def _poison_l3_l2_region_from_endpoint_error( return poisoned def _register_l3_l2_orch_comm_host_buffer(self, tensor) -> None: - if not isinstance(tensor, Tensor): - raise TypeError("L3-L2 host buffer registration expects a Tensor") + if not isinstance(tensor, ChipTensor): + raise TypeError("L3-L2 host buffer registration expects a ChipTensor") if tensor.child_memory: raise ValueError("L3-L2 payload buffer must be host storage, not child_memory device storage") if not tensor.is_contiguous: @@ -6765,8 +6777,8 @@ def _register_l3_l2_orch_comm_host_buffer(self, tensor) -> None: ) def _validate_l3_l2_orch_comm_host_buffer(self, tensor) -> None: - if not isinstance(tensor, Tensor): - raise ValueError("L3-L2 payload buffer must be a Tensor returned by orch.alloc(...)") + if not isinstance(tensor, ChipTensor): + raise ValueError("L3-L2 payload buffer must be a ChipTensor returned by orch.alloc(...)") if tensor.child_memory: raise ValueError("L3-L2 payload buffer must be host storage, not child_memory device storage") if not tensor.is_contiguous: @@ -6779,10 +6791,10 @@ def _validate_l3_l2_orch_comm_host_buffer(self, tensor) -> None: buffers = self._l3_l2_orch_comm_host_buffers if resources is None else resources.l3_l2_orch_comm_host_buffers registered_nbytes = buffers.get(base) if registered_nbytes is None: - raise ValueError("L3-L2 payload Tensor is not registered; use a tensor returned by orch.alloc(...)") + raise ValueError("L3-L2 payload ChipTensor is not registered; use a tensor returned by orch.alloc(...)") if nbytes > int(registered_nbytes): raise ValueError( - f"L3-L2 payload Tensor size {nbytes} exceeds registered shared storage {registered_nbytes}" + f"L3-L2 payload ChipTensor size {nbytes} exceeds registered shared storage {registered_nbytes}" ) def _consume_l3_host_mapped_cleanup_error_locked(self, api: str) -> RuntimeError | None: @@ -7976,6 +7988,67 @@ def _close_host_shm(entry: _HostBufEntry) -> str | None: pass return warn + # ------------------------------------------------------------------ + # Owner-side Buffer allocation + # ------------------------------------------------------------------ + + def create_buffer(self, nbytes: int) -> Buffer: + """Allocate a shared ``Buffer`` 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) -> Buffer: + # 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._buffers[buffer_id] = handle + return handle + + def _release_all_buffers(self) -> None: + """Close + unlink every owner Buffer (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._buffers.values()) + self._buffers.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()). @@ -8074,7 +8147,7 @@ def _find_host_buf_entry(self, addr: int, nbytes: int) -> _HostBufEntry | None: found by bisecting the snapshot's sorted keys so this stays log-time on the per-submit hot path rather than scanning every buffer. - Sub-view matching assumes the blob's ``Tensor.buffer.addr`` is the + Sub-view matching assumes the blob's ``ChipTensor.buffer.addr`` is the contiguous base of the host buffer (``make_tensor_arg`` builds tensors with ``start_offset == 0``); a non-zero ``start_offset`` would shift ``addr`` and is not modelled here. @@ -8966,6 +9039,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_buffers) # 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/simpler_setup/torch_interop.py b/simpler_setup/torch_interop.py index 2cff1018ca..7c09a35e09 100644 --- a/simpler_setup/torch_interop.py +++ b/simpler_setup/torch_interop.py @@ -10,7 +10,7 @@ """Torch integration helpers. Canonical home for torch-aware helpers that convert ``torch.Tensor`` and -``torch.dtype`` values into the runtime's ``Tensor`` / ``DataType`` +``torch.dtype`` values into the runtime's ``ChipTensor`` / ``DataType`` types. These helpers live in ``simpler_setup`` (not ``simpler``) so that the stable ``simpler`` runtime API can remain torch-free; torch integration is a setup-time/test-framework concern. @@ -33,7 +33,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from simpler.task_interface import DataType, Tensor + from simpler.task_interface import ChipTensor, DataType _TORCH_DTYPE_MAP = None @@ -81,30 +81,30 @@ def torch_dtype_to_datatype(dt) -> DataType: return _TORCH_DTYPE_MAP[dt] # pyright: ignore[reportOptionalSubscript] -def make_tensor_arg(tensor) -> Tensor: - """Create a ``Tensor`` from a torch.Tensor. +def make_tensor_arg(tensor) -> ChipTensor: + """Create a ``ChipTensor`` from a torch.Tensor. The result is always contiguous (row-major strides, ``start_offset == 0``) — - the unified ``Tensor`` can express strided views, but this construction path + the unified ``ChipTensor`` can express strided views, but this construction path is constrained to contiguous memory. The input torch tensor MUST therefore be contiguous; a non-contiguous tensor raises ``ValueError`` (call ``.contiguous()`` first). It must also be a CPU tensor: a device tensor's ``data_ptr()`` is a device pointer that requires ``child_memory=True``, which this helper does not set, so a non-CPU tensor raises ``ValueError``. Its ``data_ptr()``, shape, and dtype are read and stored in the returned - ``Tensor``. + ``ChipTensor``. """ - from simpler.task_interface import Tensor + from simpler.task_interface import ChipTensor _ensure_torch_map() dt = _TORCH_DTYPE_MAP.get(tensor.dtype) # pyright: ignore[reportOptionalMemberAccess] if dt is None: - raise ValueError(f"Unsupported tensor dtype for Tensor: {tensor.dtype}") + raise ValueError(f"Unsupported tensor dtype for ChipTensor: {tensor.dtype}") if tensor.device.type != "cpu": raise ValueError( f"make_tensor_arg requires a CPU tensor, got device={tensor.device}. " "A device pointer must be wrapped explicitly via " - "Tensor.make(..., child_memory=True)." + "ChipTensor.make(..., child_memory=True)." ) if not tensor.is_contiguous(): raise ValueError( @@ -112,4 +112,4 @@ def make_tensor_arg(tensor) -> Tensor: "contiguous); call tensor.contiguous() before passing it." ) shapes = tuple(int(s) for s in tensor.shape) - return Tensor.make(tensor.data_ptr(), shapes, dt) + return ChipTensor.make(tensor.data_ptr(), shapes, dt) diff --git a/src/common/task_interface/buffer.h b/src/common/task_interface/buffer.h new file mode 100644 index 0000000000..124757f46d --- /dev/null +++ b/src/common/task_interface/buffer.h @@ -0,0 +1,418 @@ +/* + * 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 + +/** + * Buffer / Tensor 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. + * - BufferDescriptor : the owner's self-describing wire descriptor — backing properties plus + * a length-delimited backend body. Embedded whole in every Tensor + * built over the buffer. + * - Tensor : the blob-carried wire element. Embeds the full BufferDescriptor + * 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. Not the same type as the global `Tensor` in + * tensor.h, which is the device POD; hence the namespace. + * + * 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" + +// The wire ABI lives in its own namespace: the global `Tensor` in tensor.h is the device POD the +// L2 runtime reads, and these two types are not interchangeable. A file that touches both must +// qualify — do not open this namespace with a using-directive. +namespace simpler { + +// Leading sentinel of a BufferDescriptor, 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 Tensor 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_tensor_blob` rejects any other value. +inline constexpr uint32_t TENSOR_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 Tensor 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 Tensor 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 BufferDescriptor { + 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(BufferDescriptor) == 88, "BufferDescriptor is wire ABI"); +static_assert(offsetof(BufferDescriptor, magic) == 0); +static_assert(offsetof(BufferDescriptor, address_space) == 2); +static_assert(offsetof(BufferDescriptor, access) == 3); +static_assert(offsetof(BufferDescriptor, backend_kind) == 4); +static_assert(offsetof(BufferDescriptor, identity) == 8); +static_assert(offsetof(BufferDescriptor, nbytes) == 40); +static_assert(offsetof(BufferDescriptor, owner_worker_path_id) == 48); +static_assert(offsetof(BufferDescriptor, body_len) == 52); +static_assert(offsetof(BufferDescriptor, 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 + * `buffer.identity`, and reuses the cached base for later tensors over 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 Tensor { + BufferDescriptor buffer; + 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 tensor_extent_bytes(const Tensor &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 Tensor 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_tensor(const Tensor &r) { + auto reject = [](const char *what) { + throw std::invalid_argument(what); + }; + const BufferDescriptor &h = r.buffer; + + 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"); + } + + // 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); + 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 = tensor_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 tensors_overlap(const Tensor &a, const Tensor &b) { + if (!(a.buffer.identity == b.buffer.identity)) return false; + const uint64_t a_end = a.byte_offset + tensor_extent_bytes(a); + const uint64_t b_end = b.byte_offset + tensor_extent_bytes(b); + return a.byte_offset < b_end && b.byte_offset < a_end; +} + +static_assert(std::is_trivially_copyable_v, "Tensor must be trivially copyable for blob memcpy"); +static_assert(sizeof(Tensor) == 144, "Tensor is wire ABI"); +static_assert(offsetof(Tensor, buffer) == 0); +static_assert(offsetof(Tensor, byte_offset) == 88); +static_assert(offsetof(Tensor, ndims) == 96); +static_assert(offsetof(Tensor, shapes) == 100); +static_assert(offsetof(Tensor, strides) == 120); +static_assert(offsetof(Tensor, dtype) == 140); + +// ============================================================================ +// Tensor wire blob — versioned, length-prefixed (P1-B). +// ============================================================================ +// +// Byte layout: +// offset 0: uint32 magic = TENSOR_BLOB_MAGIC +// offset 4: int32 tensor_count = R +// offset 8: int32 scalar_count = S +// offset 12: uint32 reserved (= 0) +// offset 16: Tensor tensors[R] (sizeof(Tensor) B each) +// offset 16 + R*sizeof(Tensor): uint64_t scalars[S] +// +// The element is Tensor (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 +// tensors[0] (whose first field is a u64) and is rejected when non-zero. + +inline constexpr size_t TENSOR_BLOB_HEADER_SIZE = 16; + +struct TensorBlobView { + int32_t tensor_count; + int32_t scalar_count; + const uint8_t *tensor_bytes; // R contiguous Tensor; extract element i with tensor(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. + Tensor tensor(int32_t i) const { + if (i < 0 || i >= tensor_count) { + throw std::out_of_range("TensorBlobView::tensor: index outside the blob's tensor_count"); + } + Tensor r; + std::memcpy(&r, tensor_bytes + static_cast(i) * sizeof(Tensor), sizeof(Tensor)); + validate_tensor(r); + return r; + } +}; + +inline size_t tensor_blob_size(int32_t tensor_count, int32_t scalar_count) { + return TENSOR_BLOB_HEADER_SIZE + static_cast(tensor_count) * sizeof(Tensor) + + static_cast(scalar_count) * sizeof(uint64_t); +} + +// Serialize refs + scalars into `dst` (caller ensures room for tensor_blob_size). +inline void write_tensor_blob( + uint8_t *dst, const Tensor *tensors, int32_t tensor_count, const uint64_t *scalars, int32_t scalar_count +) { + uint32_t magic = TENSOR_BLOB_MAGIC; + uint32_t reserved = 0; + std::memcpy(dst + 0, &magic, sizeof(magic)); + std::memcpy(dst + 4, &tensor_count, sizeof(tensor_count)); + std::memcpy(dst + 8, &scalar_count, sizeof(scalar_count)); + std::memcpy(dst + 12, &reserved, sizeof(reserved)); + if (tensor_count > 0) { + std::memcpy(dst + TENSOR_BLOB_HEADER_SIZE, tensors, static_cast(tensor_count) * sizeof(Tensor)); + } + if (scalar_count > 0) { + std::memcpy( + dst + TENSOR_BLOB_HEADER_SIZE + static_cast(tensor_count) * sizeof(Tensor), scalars, + static_cast(scalar_count) * sizeof(uint64_t) + ); + } +} + +// Zero-copy view into a blob written by write_tensor_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 TensorBlobView read_tensor_blob(const uint8_t *src, size_t capacity) { + if (capacity < TENSOR_BLOB_HEADER_SIZE) { + throw std::runtime_error( + "read_tensor_blob: capacity " + std::to_string(capacity) + " < header size " + + std::to_string(TENSOR_BLOB_HEADER_SIZE) + ); + } + uint32_t magic; + std::memcpy(&magic, src + 0, sizeof(magic)); + if (magic != TENSOR_BLOB_MAGIC) { + throw std::runtime_error( + "read_tensor_blob: not a Tensor blob — magic " + std::to_string(magic) + " (expected " + + std::to_string(TENSOR_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_tensor_blob: reserved header word must be zero"); + } + if (R < 0 || S < 0) { + throw std::runtime_error( + "read_tensor_blob: negative counts — tensors=" + std::to_string(R) + ", scalars=" + std::to_string(S) + ); + } + const size_t needed = tensor_blob_size(R, S); + if (needed > capacity) { + throw std::runtime_error( + "read_tensor_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 TensorBlobView{ + R, + S, + src + TENSOR_BLOB_HEADER_SIZE, + reinterpret_cast(src + TENSOR_BLOB_HEADER_SIZE + static_cast(R) * sizeof(Tensor)), + }; +} + +} // namespace simpler diff --git a/src/common/task_interface/data_type.h b/src/common/task_interface/data_type.h index 7c48db0b39..43ef6b8231 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 can describe. Both the wire `simpler::Tensor` and the device +// `Tensor` size their shapes[] / strides[] 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 Buffer 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/st/task_timing/task_timing_slots/test_task_timing_e2e.py b/tests/st/task_timing/task_timing_slots/test_task_timing_e2e.py index feebc21fd7..1d4f6dfb25 100644 --- a/tests/st/task_timing/task_timing_slots/test_task_timing_e2e.py +++ b/tests/st/task_timing/task_timing_slots/test_task_timing_e2e.py @@ -32,9 +32,9 @@ CallConfig, ChipCallable, ChipStorageTaskArgs, + ChipTensor, CoreCallable, DataType, - Tensor, ) from simpler.worker import Worker @@ -46,7 +46,7 @@ _A2A3_VECTOR_ADD = os.path.join(_PROJECT_ROOT, "examples", "workers", "l2", "vector_add") # a5's ccec/pto-isa needs the qualified `pto::Stride` spelling; the l2/vector_add # kernel above (a2a3-only example) uses unqualified `Stride` and does not compile -# under the a5 AICore toolchain. Use the a5-native add kernel (same out=a+b Tensor* +# under the a5 AICore toolchain. Use the a5-native add kernel (same out=a+b ChipTensor* # ABI) on a5. _A5_VECTOR_ADD = os.path.join(_PROJECT_ROOT, "examples", "a5", "tensormap_and_ringbuffer", "vector_example") @@ -144,9 +144,9 @@ def _drive( worker.copy_to(dev_b, host_b.data_ptr(), NBYTES) args = ChipStorageTaskArgs() - args.add_tensor(Tensor.make(dev_a, (N_ROWS, N_COLS), DataType.FLOAT32)) - args.add_tensor(Tensor.make(dev_b, (N_ROWS, N_COLS), DataType.FLOAT32)) - args.add_tensor(Tensor.make(dev_out, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_a, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_b, (N_ROWS, N_COLS), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(dev_out, (N_ROWS, N_COLS), DataType.FLOAT32)) config = CallConfig() config.enable_l2_swimlane = False # slots must work with swimlane OFF @@ -334,11 +334,11 @@ def test_mix_task_aggregates_across_subtasks(st_platform, st_device_ids, capfd): worker.copy_to(bufs[name], t.contiguous().data_ptr(), nb) args = ChipStorageTaskArgs() - args.add_tensor(Tensor.make(bufs["A"], (_MATMUL_SIZE, _MATMUL_SIZE), DataType.FLOAT32)) - args.add_tensor(Tensor.make(bufs["B"], (_MATMUL_SIZE, _MATMUL_SIZE), DataType.FLOAT32)) - args.add_tensor(Tensor.make(bufs["C"], (_TILE_ELEMS,), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(bufs["A"], (_MATMUL_SIZE, _MATMUL_SIZE), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(bufs["B"], (_MATMUL_SIZE, _MATMUL_SIZE), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(bufs["C"], (_TILE_ELEMS,), DataType.FLOAT32)) for n in "DEFGHI": - args.add_tensor(Tensor.make(bufs[n], (_TILE_ELEMS,), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(bufs[n], (_TILE_ELEMS,), DataType.FLOAT32)) config = CallConfig() config.enable_l2_swimlane = False diff --git a/tests/st/worker/collectives/_helpers.py b/tests/st/worker/collectives/_helpers.py index 6550cbb9e4..5e8e4e5375 100644 --- a/tests/st/worker/collectives/_helpers.py +++ b/tests/st/worker/collectives/_helpers.py @@ -18,7 +18,7 @@ import ctypes import torch -from simpler.task_interface import CommBufferSpec, DataType, TaskArgs, Tensor, TensorArgType +from simpler.task_interface import ChipTensor, CommBufferSpec, DataType, TaskArgs, TensorArgType from simpler_setup import Tensor as STensor from simpler_setup.scene_test import TaskArgsBuilder @@ -127,7 +127,7 @@ def allreduce_orch_fn(orch, callables, task_args, config): chip_args.add_tensor(make_tensor_arg(getattr(task_args, f"in_{i}")), TensorArgType.INPUT) chip_args.add_tensor(make_tensor_arg(getattr(task_args, f"out_{i}")), TensorArgType.OUTPUT_EXISTING) chip_args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["scratch"], shapes=(float_elems,), dtype=DataType.FLOAT32, @@ -196,7 +196,7 @@ def generic_collective_orch_fn( chip_args.add_tensor(make_tensor_arg(getattr(task_args, f"in_{i}")), TensorArgType.INPUT) chip_args.add_tensor(make_tensor_arg(getattr(task_args, f"out_{i}")), TensorArgType.OUTPUT_EXISTING) chip_args.add_tensor( - Tensor.make( + ChipTensor.make( data=domain.buffer_ptrs["scratch"], shapes=(float_elems,), dtype=DataType.FLOAT32, diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index 35216c8a64..c9956ef11c 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -395,6 +395,7 @@ set_tests_properties(test_native_run_launch_signal PROPERTIES LABELS "no_hardwar # --------------------------------------------------------------------------- # Types / task_interface tests (src/common/task_interface/) # --------------------------------------------------------------------------- +add_task_interface_test(test_buffer types/test_buffer.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.cpp b/tests/ut/cpp/types/test_buffer.cpp new file mode 100644 index 0000000000..4d7f970140 --- /dev/null +++ b/tests/ut/cpp/types/test_buffer.cpp @@ -0,0 +1,272 @@ +/* + * 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 Buffer / simpler::Tensor (buffer.h); the contract they pin is described +// in docs/buffer-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.h" +#include "task_args.h" + +namespace { + +simpler::CanonicalIdentity make_identity() { + simpler::CanonicalIdentity id{}; + for (uint32_t i = 0; i < simpler::OWNER_INSTANCE_ID_BYTES; ++i) + id.owner_instance_id[i] = static_cast(0xA0 + i); + id.buffer_id = 0x0102030405060708ULL; + id.generation = 7; + return id; +} + +simpler::Tensor make_tensor() { + simpler::Tensor r{}; + r.buffer.magic = simpler::BUFFER_DESCRIPTOR_MAGIC; + r.buffer.backend_kind = static_cast(simpler::BackendKind::POSIX_SHM); + r.buffer.identity = make_identity(); + r.buffer.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(BufferAbi, StructSizesAreFrozen) { + EXPECT_EQ(sizeof(simpler::CanonicalIdentity), 32u); + EXPECT_EQ(sizeof(simpler::Tensor), 144u); + EXPECT_EQ(sizeof(simpler::BufferDescriptor), 88u); +} + +TEST(BufferAbi, ConstantsAreFrozen) { + EXPECT_EQ(simpler::BUFFER_DESCRIPTOR_MAGIC, 0x5342); + EXPECT_EQ(simpler::TENSOR_BLOB_MAGIC, 0x424F4C42u); + EXPECT_EQ(simpler::OWNER_INSTANCE_ID_BYTES, 8u); + EXPECT_EQ(simpler::DESC_MAX_BYTES, 32u); +} + +TEST(BufferAbi, EnumValuesAreFrozen) { + EXPECT_EQ(static_cast(AddressSpace::HOST), 0); + EXPECT_EQ(static_cast(AddressSpace::DEVICE), 1); + EXPECT_EQ(static_cast(simpler::AccessMode::READ), 0); + EXPECT_EQ(static_cast(simpler::AccessMode::WRITE), 1); + EXPECT_EQ(static_cast(simpler::AccessMode::READWRITE), 2); + EXPECT_EQ(static_cast(simpler::BackendKind::FORK_SHM), 0); + EXPECT_EQ(static_cast(simpler::BackendKind::POSIX_SHM), 1); + EXPECT_EQ(static_cast(simpler::BackendKind::VMM_WINDOW), 2); + EXPECT_EQ(static_cast(simpler::BackendKind::REMOTE_SIDECAR), 3); + EXPECT_EQ(static_cast(simpler::BackendKind::DEVICE_MALLOC), 4); + EXPECT_EQ(static_cast(simpler::BackendKind::FORK_COW), 5); +} + +// --- memcpy round trip ------------------------------------------------------------------------- + +TEST(BufferAbi, TensorSurvivesByteRoundTrip) { + simpler::Tensor src = make_tensor(); + uint8_t bytes[sizeof(simpler::Tensor)]; + std::memcpy(bytes, &src, sizeof(simpler::Tensor)); + simpler::Tensor dst{}; + std::memcpy(&dst, bytes, sizeof(simpler::Tensor)); + EXPECT_EQ(std::memcmp(&src, &dst, sizeof(simpler::Tensor)), 0); + EXPECT_EQ(dst.byte_offset, 4096u); + EXPECT_EQ(dst.dtype, DataType::FLOAT16); + EXPECT_EQ(dst.strides[0], 32u); + EXPECT_EQ(dst.buffer.identity, src.buffer.identity); +} + +TEST(BufferAbi, DescriptorSurvivesByteRoundTrip) { + simpler::BufferDescriptor src{}; + src.magic = simpler::BUFFER_DESCRIPTOR_MAGIC; + src.address_space = static_cast(AddressSpace::DEVICE); + src.access = static_cast(simpler::AccessMode::READWRITE); + src.backend_kind = static_cast(simpler::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(simpler::BufferDescriptor)]; + std::memcpy(bytes, &src, sizeof(simpler::BufferDescriptor)); + simpler::BufferDescriptor dst{}; + std::memcpy(&dst, bytes, sizeof(simpler::BufferDescriptor)); + EXPECT_EQ(std::memcmp(&src, &dst, sizeof(simpler::BufferDescriptor)), 0); + EXPECT_EQ(dst.magic, simpler::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(BufferAbi, IdentityDistinguishesGenerationAndIncarnation) { + simpler::CanonicalIdentity a = make_identity(); + simpler::CanonicalIdentity b = make_identity(); + EXPECT_EQ(a, b); + + b.generation = a.generation + 1; // buffer_id reuse across generations (ABA) + EXPECT_NE(a, b); + + simpler::CanonicalIdentity c = make_identity(); + c.owner_instance_id[0] ^= 0xFF; // different owner incarnation nonce + EXPECT_NE(a, c); + + simpler::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(BufferAbi, IdentityPaddingIsExcludedFromKeyAndHash) { + simpler::CanonicalIdentityHash h; + simpler::CanonicalIdentity clean = make_identity(); + simpler::CanonicalIdentity dirty = make_identity(); + std::memset(dirty._pad, 0xA5, sizeof(dirty._pad)); + EXPECT_EQ(clean, dirty); + EXPECT_EQ(h(clean), h(dirty)); +} + +TEST(BufferAbi, IdentityHashMatchesEquality) { + simpler::CanonicalIdentityHash h; + simpler::CanonicalIdentity a = make_identity(); + simpler::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) +} + +// --- simpler::Tensor wire blob: versioned length-prefixed round trip + rejection ---------------------- + +simpler::Tensor make_tensor_b() { + simpler::Tensor r = make_tensor(); + r.buffer.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(TensorBlob, RoundTrip) { + simpler::Tensor tensors[2] = {make_tensor(), make_tensor_b()}; + uint64_t scalars[2] = {42, 0xC0FFEE}; + size_t sz = simpler::tensor_blob_size(2, 2); + EXPECT_EQ(sz, simpler::TENSOR_BLOB_HEADER_SIZE + 2 * sizeof(simpler::Tensor) + 2 * sizeof(uint64_t)); + + std::vector buf(sz); + simpler::write_tensor_blob(buf.data(), tensors, 2, scalars, 2); + + simpler::TensorBlobView v = simpler::read_tensor_blob(buf.data(), sz); + ASSERT_EQ(v.tensor_count, 2); + ASSERT_EQ(v.scalar_count, 2); + simpler::Tensor r0 = v.tensor(0); + simpler::Tensor r1 = v.tensor(1); + EXPECT_EQ(std::memcmp(&r0, &tensors[0], sizeof(simpler::Tensor)), 0); + EXPECT_EQ(std::memcmp(&r1, &tensors[1], sizeof(simpler::Tensor)), 0); + EXPECT_EQ(v.scalars[0], 42u); + EXPECT_EQ(v.scalars[1], 0xC0FFEEu); +} + +TEST(TensorBlob, EmptyBlob) { + size_t sz = simpler::tensor_blob_size(0, 0); + EXPECT_EQ(sz, simpler::TENSOR_BLOB_HEADER_SIZE); + std::vector buf(sz); + simpler::write_tensor_blob(buf.data(), nullptr, 0, nullptr, 0); + simpler::TensorBlobView v = simpler::read_tensor_blob(buf.data(), sz); + EXPECT_EQ(v.tensor_count, 0); + EXPECT_EQ(v.scalar_count, 0); +} + +TEST(TensorBlob, RejectsBadMagic) { + std::vector buf(simpler::tensor_blob_size(0, 0)); + simpler::write_tensor_blob(buf.data(), nullptr, 0, nullptr, 0); + uint32_t bad = simpler::TENSOR_BLOB_MAGIC + 1; + std::memcpy(buf.data(), &bad, sizeof(bad)); + EXPECT_THROW(simpler::read_tensor_blob(buf.data(), buf.size()), std::runtime_error); +} + +TEST(TensorBlob, RejectsTruncatedCapacity) { + simpler::Tensor tensors[1] = {make_tensor()}; + std::vector buf(simpler::tensor_blob_size(1, 0)); + simpler::write_tensor_blob(buf.data(), tensors, 1, nullptr, 0); + EXPECT_THROW(simpler::read_tensor_blob(buf.data(), simpler::TENSOR_BLOB_HEADER_SIZE), std::runtime_error); + EXPECT_THROW(simpler::read_tensor_blob(buf.data(), 4), std::runtime_error); +} + +TEST(TensorBlob, RejectsNegativeCount) { + std::vector buf(64, 0); + uint32_t ver = simpler::TENSOR_BLOB_MAGIC; + std::memcpy(buf.data() + 0, &ver, sizeof(ver)); + int32_t neg = -1; + std::memcpy(buf.data() + 4, &neg, sizeof(neg)); // tensor_count = -1 + EXPECT_THROW(simpler::read_tensor_blob(buf.data(), buf.size()), std::runtime_error); +} + +TEST(TensorBlob, RejectsNonZeroReserved) { + std::vector buf(simpler::tensor_blob_size(0, 0)); + simpler::write_tensor_blob(buf.data(), nullptr, 0, nullptr, 0); + uint32_t dirty = 1; + std::memcpy(buf.data() + 12, &dirty, sizeof(dirty)); // reserved header word + EXPECT_THROW(simpler::read_tensor_blob(buf.data(), buf.size()), std::runtime_error); +} + +// --- simpler::validate_tensor: the shared receive-side gate ------------------------------------------ + +TEST(BufferAbi, ForkCowGrantsReadOnly) { + simpler::Tensor r = make_tensor(); + r.buffer.backend_kind = static_cast(simpler::BackendKind::FORK_COW); + r.buffer.access = static_cast(simpler::AccessMode::READ); + EXPECT_NO_THROW(simpler::validate_tensor(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 : {simpler::AccessMode::WRITE, simpler::AccessMode::READWRITE}) { + r.buffer.access = static_cast(bad); + EXPECT_THROW(simpler::validate_tensor(r), std::invalid_argument); + } + + // The same grants stay legal over MAP_SHARED, where a child's write does reach the owner. + r.buffer.backend_kind = static_cast(simpler::BackendKind::FORK_SHM); + r.buffer.access = static_cast(simpler::AccessMode::READWRITE); + EXPECT_NO_THROW(simpler::validate_tensor(r)); +} + +TEST(TensorBlob, RefIndexIsBoundsChecked) { + simpler::Tensor r = make_tensor(); + std::vector buf(simpler::tensor_blob_size(1, 0)); + simpler::write_tensor_blob(buf.data(), &r, 1, nullptr, 0); + simpler::TensorBlobView view = simpler::read_tensor_blob(buf.data(), buf.size()); + + EXPECT_NO_THROW(view.tensor(0)); + EXPECT_THROW(view.tensor(1), std::out_of_range); + EXPECT_THROW(view.tensor(-1), std::out_of_range); +} + +} // namespace diff --git a/tests/ut/py/test_buffer.py b/tests/ut/py/test_buffer.py new file mode 100644 index 0000000000..50d6dc277b --- /dev/null +++ b/tests/ut/py/test_buffer.py @@ -0,0 +1,411 @@ +# 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: identity/descriptor pack-unpack + create/import round trip. + +Imports come from the defining module because most of what is exercised here — the registry, the +blob codec, the handle constructors — is internal to it. `Tensor` is the exception: it is also the +public name, so the first test pins that the two paths are one object. +""" + +import ctypes +import struct +from dataclasses import replace + +import pytest +from _task_interface import ( + BUFFER_DESCRIPTOR_BYTES, + CANONICAL_IDENTITY_BYTES, + OWNER_INSTANCE_ID_BYTES, + TENSOR_BLOB_MAGIC, + WIRE_TENSOR_BYTES, + DataType, + tensor_blob_tensors, +) +from simpler.buffer import ( + AccessMode, + AddressSpace, + BackendKind, + BufferDescriptor, + CanonicalIdentity, + ImportRegistry, + Tensor, + create_host_shared_buffer, + intern_worker_path, + mint_owner_instance_id, + re_export, + wrap_device_malloc, +) +from simpler.task_interface import ChipTensor +from simpler.task_interface import Tensor as PublicTensor + +_OID = bytes(range(0xA0, 0xA0 + OWNER_INSTANCE_ID_BYTES)) + + +def test_public_tensor_name_is_this_module_s_type(): + # Users reach the wire type as simpler.task_interface.Tensor; simpler.buffer defines it. A + # second class behind the public name would typecheck everywhere and fail only on the wire. + assert PublicTensor is Tensor + assert ChipTensor is not Tensor # the device POD is a distinct type, not an alias + + +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 = BufferDescriptor( + 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_DESCRIPTOR_BYTES + assert BufferDescriptor.unpack(raw) == host + + dev = BufferDescriptor( + 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 BufferDescriptor.unpack(dev.pack()) == dev + + +def test_descriptor_rejects_bad_magic(): + raw = bytearray( + BufferDescriptor( + 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"): + BufferDescriptor.unpack(bytes(raw)) + + +def test_descriptor_rejects_body_len_past_the_array(): + raw = bytearray( + BufferDescriptor( + 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"): + BufferDescriptor.unpack(bytes(raw)) + + +def test_worker_path_is_diagnostic_and_survives_an_unknown_id(): + h = BufferDescriptor( + 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 BufferDescriptor.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"): + BufferDescriptor( + 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 tensor 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()).buffer.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 = BufferDescriptor( + 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"): + BufferDescriptor( + 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), + ]: + BufferDescriptor(_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("". + # "Unsupported tensor dtype for ChipTensor: torch.". t = torch.zeros(2, 4, dtype=torch_dt) arg = make_tensor_arg(t) assert arg.dtype == expected @@ -233,62 +233,62 @@ def test_scalar_to_uint64_float_ctypes(self): # ============================================================================ -# Tensor +# ChipTensor # ============================================================================ class TestTensor: def test_default_constructor(self): - arg = Tensor() + arg = ChipTensor() assert arg is not None def test_max_dims_constant(self): assert MAX_TENSOR_DIMS == 5 def test_make(self): - arg = Tensor.make(0xDEAD, (4, 8), DataType.FLOAT32) + arg = ChipTensor.make(0xDEAD, (4, 8), DataType.FLOAT32) assert arg.data == 0xDEAD assert arg.shapes == (4, 8) assert arg.ndims == 2 assert arg.dtype == DataType.FLOAT32 def test_nbytes(self): - arg = Tensor.make(0, (10, 20), DataType.FLOAT32) + arg = ChipTensor.make(0, (10, 20), DataType.FLOAT32) assert arg.nbytes() == 10 * 20 * 4 def test_nbytes_int8(self): - arg = Tensor.make(0, (256,), DataType.INT8) + arg = ChipTensor.make(0, (256,), DataType.INT8) assert arg.nbytes() == 256 def test_shapes_setter(self): - arg = Tensor() + arg = ChipTensor() arg.shapes = (3, 5, 7) assert arg.ndims == 3 assert arg.shapes == (3, 5, 7) def test_max_dims(self): - arg = Tensor.make(0, (1, 2, 3, 4, 5), DataType.INT32) + arg = ChipTensor.make(0, (1, 2, 3, 4, 5), DataType.INT32) assert arg.ndims == 5 def test_exceed_max_dims(self): with pytest.raises((ValueError, RuntimeError)): - Tensor.make(0, (1, 2, 3, 4, 5, 6), DataType.INT32) + ChipTensor.make(0, (1, 2, 3, 4, 5, 6), DataType.INT32) def test_dtype_readwrite(self): - arg = Tensor.make(0, (1,), DataType.FLOAT32) + arg = ChipTensor.make(0, (1,), DataType.FLOAT32) arg.dtype = DataType.INT64 assert arg.dtype == DataType.INT64 def test_data_readwrite(self): - arg = Tensor.make(0x1000, (1,), DataType.FLOAT32) + arg = ChipTensor.make(0x1000, (1,), DataType.FLOAT32) assert arg.data == 0x1000 arg.data = 0x2000 assert arg.data == 0x2000 def test_repr(self): - arg = Tensor.make(0x1000, (4, 8), DataType.FLOAT16) + arg = ChipTensor.make(0x1000, (4, 8), DataType.FLOAT16) r = repr(arg) - assert "Tensor" in r + assert "ChipTensor" in r assert "4" in r assert "8" in r assert "FLOAT16" in r @@ -308,7 +308,7 @@ def test_empty(self): def test_add_tensor(self): args = ChipStorageTaskArgs() - t = Tensor.make(0xBEEF, (4, 8), DataType.FLOAT32) + t = ChipTensor.make(0xBEEF, (4, 8), DataType.FLOAT32) args.add_tensor(t) assert args.tensor_count() == 1 assert args.scalar_count() == 0 @@ -323,8 +323,8 @@ def test_add_scalar(self): def test_mixed(self): args = ChipStorageTaskArgs() - args.add_tensor(Tensor.make(0x1, (2,), DataType.INT32)) - args.add_tensor(Tensor.make(0x2, (3,), DataType.FLOAT16)) + args.add_tensor(ChipTensor.make(0x1, (2,), DataType.INT32)) + args.add_tensor(ChipTensor.make(0x2, (3,), DataType.FLOAT16)) args.add_scalar(99) args.add_scalar(100) assert args.tensor_count() == 2 @@ -335,12 +335,12 @@ def test_tensor_before_scalar_enforced(self): args = ChipStorageTaskArgs() args.add_scalar(42) with pytest.raises(RuntimeError): - args.add_tensor(Tensor.make(0x1, (2,), DataType.INT32)) + args.add_tensor(ChipTensor.make(0x1, (2,), DataType.INT32)) def test_tensor_access(self): args = ChipStorageTaskArgs() - args.add_tensor(Tensor.make(0xA, (4,), DataType.FLOAT32)) - args.add_tensor(Tensor.make(0xB, (8,), DataType.INT32)) + args.add_tensor(ChipTensor.make(0xA, (4,), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(0xB, (8,), DataType.INT32)) assert args.tensor(0).data == 0xA assert args.tensor(1).data == 0xB assert args.tensor(0).shapes == (4,) @@ -365,7 +365,7 @@ def test_scalar_out_of_range(self): def test_clear(self): args = ChipStorageTaskArgs() - args.add_tensor(Tensor.make(0, (1,), DataType.INT8)) + args.add_tensor(ChipTensor.make(0, (1,), DataType.INT8)) args.add_scalar(42) args.clear() assert len(args) == 0 @@ -404,24 +404,24 @@ def test_empty(self): def test_add_tensor_default_tag(self): args = TaskArgs() - t = Tensor.make(0xBEEF, (4, 8), DataType.FLOAT32) + t = ChipTensor.make(0xBEEF, (4, 8), DataType.FLOAT32) args.add_tensor(t) assert args.tensor_count() == 1 assert args.tag(0) == TensorArgType.INPUT def test_add_tensor_with_tag(self): args = TaskArgs() - t = Tensor.make(0xBEEF, (4, 8), DataType.FLOAT32) + t = ChipTensor.make(0xBEEF, (4, 8), DataType.FLOAT32) args.add_tensor(t, TensorArgType.OUTPUT) assert args.tag(0) == TensorArgType.OUTPUT def test_multiple_tensors_with_tags(self): args = TaskArgs() - args.add_tensor(Tensor.make(0x1, (2,), DataType.INT32), TensorArgType.INPUT) - args.add_tensor(Tensor.make(0x2, (3,), DataType.FLOAT16), TensorArgType.OUTPUT) - args.add_tensor(Tensor.make(0x3, (4,), DataType.INT8), TensorArgType.INOUT) - args.add_tensor(Tensor.make(0x4, (5,), DataType.FLOAT32), TensorArgType.OUTPUT_EXISTING) - args.add_tensor(Tensor.make(0x5, (6,), DataType.INT32), TensorArgType.NO_DEP) + args.add_tensor(ChipTensor.make(0x1, (2,), DataType.INT32), TensorArgType.INPUT) + args.add_tensor(ChipTensor.make(0x2, (3,), DataType.FLOAT16), TensorArgType.OUTPUT) + args.add_tensor(ChipTensor.make(0x3, (4,), DataType.INT8), TensorArgType.INOUT) + args.add_tensor(ChipTensor.make(0x4, (5,), DataType.FLOAT32), TensorArgType.OUTPUT_EXISTING) + args.add_tensor(ChipTensor.make(0x5, (6,), DataType.INT32), TensorArgType.NO_DEP) assert args.tensor_count() == 5 assert args.tag(0) == TensorArgType.INPUT assert args.tag(1) == TensorArgType.OUTPUT @@ -431,7 +431,7 @@ def test_multiple_tensors_with_tags(self): def test_set_tag(self): args = TaskArgs() - args.add_tensor(Tensor.make(0x1, (2,), DataType.INT32)) + args.add_tensor(ChipTensor.make(0x1, (2,), DataType.INT32)) assert args.tag(0) == TensorArgType.INPUT args.set_tag(0, TensorArgType.INOUT) assert args.tag(0) == TensorArgType.INOUT @@ -445,8 +445,8 @@ def test_add_scalar(self): def test_mixed_with_tags(self): args = TaskArgs() - args.add_tensor(Tensor.make(0x1, (2,), DataType.INT32), TensorArgType.INPUT) - args.add_tensor(Tensor.make(0x2, (3,), DataType.FLOAT16), TensorArgType.OUTPUT) + args.add_tensor(ChipTensor.make(0x1, (2,), DataType.INT32), TensorArgType.INPUT) + args.add_tensor(ChipTensor.make(0x2, (3,), DataType.FLOAT16), TensorArgType.OUTPUT) args.add_scalar(99) args.add_scalar(100) assert args.tensor_count() == 2 @@ -461,12 +461,12 @@ def test_tensor_before_scalar_enforced(self): args = TaskArgs() args.add_scalar(42) with pytest.raises(RuntimeError): - args.add_tensor(Tensor.make(0x1, (2,), DataType.INT32)) + args.add_tensor(ChipTensor.make(0x1, (2,), DataType.INT32)) def test_tensor_access(self): args = TaskArgs() - args.add_tensor(Tensor.make(0xA, (4,), DataType.FLOAT32)) - args.add_tensor(Tensor.make(0xB, (8,), DataType.INT32)) + args.add_tensor(ChipTensor.make(0xA, (4,), DataType.FLOAT32)) + args.add_tensor(ChipTensor.make(0xB, (8,), DataType.INT32)) assert args.tensor(0).data == 0xA assert args.tensor(1).data == 0xB assert args.tensor(0).shapes == (4,) @@ -501,7 +501,7 @@ def test_set_tag_out_of_range(self): def test_clear(self): args = TaskArgs() - args.add_tensor(Tensor.make(0, (1,), DataType.INT8), TensorArgType.OUTPUT) + args.add_tensor(ChipTensor.make(0, (1,), DataType.INT8), TensorArgType.OUTPUT) args.add_scalar(42) args.clear() assert len(args) == 0 @@ -512,7 +512,7 @@ def test_no_capacity_limit_tensors(self): """TaskArgs is vector-backed — no per-class capacity limit on tensors.""" args = TaskArgs() for i in range(20): - args.add_tensor(Tensor.make(i, (1,), DataType.INT8)) + args.add_tensor(ChipTensor.make(i, (1,), DataType.INT8)) assert args.tensor_count() == 20 def test_no_capacity_limit_scalars(self): @@ -724,7 +724,7 @@ def test_host_inline_descriptor_materializes_local_tensor_data(self): ) from simpler.remote_l3_session import _materialize_task_args - tensor = Tensor.make(0, (4,), DataType.UINT8) + tensor = ChipTensor.make(0, (4,), DataType.UINT8) desc = RemoteTensorDesc( address_space=WireRemoteAddressSpace.HOST_INLINE, owner_worker_id=0, @@ -759,7 +759,7 @@ def test_remote_buffer_descriptor_materializes_session_registry_address(self): backing = ctypes.create_string_buffer(b"01234567", 8) entry = _RemoteBufferEntry(backing, 8, 1, WireRemoteAddressSpace.REMOTE_DEVICE) - tensor = Tensor.make(0, (4,), DataType.UINT8) + tensor = ChipTensor.make(0, (4,), DataType.UINT8) desc = RemoteTensorDesc( address_space=WireRemoteAddressSpace.REMOTE_DEVICE, owner_worker_id=2, @@ -962,45 +962,45 @@ def test_repr(self): # ============================================================================ -# Tensor.child_memory +# ChipTensor.child_memory # ============================================================================ class TestChildMemory: def test_default_is_false(self): - t = Tensor() + t = ChipTensor() assert t.child_memory is False def test_make_default_is_false(self): - t = Tensor.make(0x1000, (4,), DataType.FLOAT32) + t = ChipTensor.make(0x1000, (4,), DataType.FLOAT32) assert t.child_memory is False def test_make_child_memory_true(self): - t = Tensor.make(0xDEAD, (8,), DataType.FLOAT16, child_memory=True) + t = ChipTensor.make(0xDEAD, (8,), DataType.FLOAT16, child_memory=True) assert t.child_memory is True assert t.data == 0xDEAD assert t.shapes == (8,) assert t.dtype == DataType.FLOAT16 def test_set_child_memory(self): - t = Tensor.make(0x1000, (4,), DataType.FLOAT32) + t = ChipTensor.make(0x1000, (4,), DataType.FLOAT32) assert t.child_memory is False t.child_memory = True assert t.child_memory is True def test_repr_shows_child_memory_when_set(self): - t = Tensor.make(0x1000, (4,), DataType.FLOAT32, child_memory=True) + t = ChipTensor.make(0x1000, (4,), DataType.FLOAT32, child_memory=True) r = repr(t) assert "child_memory=True" in r def test_repr_hides_child_memory_when_default(self): - t = Tensor.make(0x1000, (4,), DataType.FLOAT32) + t = ChipTensor.make(0x1000, (4,), DataType.FLOAT32) r = repr(t) assert "child_memory" not in r def test_chip_storage_preserves_child_memory(self): args = ChipStorageTaskArgs() - t = Tensor.make(0x2000, (16,), DataType.INT32, child_memory=True) + t = ChipTensor.make(0x2000, (16,), DataType.INT32, child_memory=True) args.add_tensor(t) out = args.tensor(0) assert out.child_memory is True diff --git a/tests/ut/py/test_worker/test_child_addr_guard.py b/tests/ut/py/test_worker/test_child_addr_guard.py index d514c85f91..bec1e81ec0 100644 --- a/tests/ut/py/test_worker/test_child_addr_guard.py +++ b/tests/ut/py/test_worker/test_child_addr_guard.py @@ -26,7 +26,7 @@ import pytest import simpler.orchestrator as orch_mod -from _task_interface import DataType, Tensor, TensorArgType +from _task_interface import ChipTensor, DataType, TensorArgType from simpler.orchestrator import Orchestrator from simpler.task_interface import TaskArgs from simpler.worker import Worker, _ChildProvEntry, _Lifecycle @@ -38,7 +38,7 @@ def _l3() -> Worker: def _child_args(ptr: int, *, n: int = 16) -> TaskArgs: args = TaskArgs() - args.add_tensor(Tensor.make(ptr, (n,), DataType.FLOAT32, child_memory=True), TensorArgType.OUTPUT_EXISTING) + args.add_tensor(ChipTensor.make(ptr, (n,), DataType.FLOAT32, child_memory=True), TensorArgType.OUTPUT_EXISTING) return args @@ -334,7 +334,7 @@ def test_host_only_args_are_not_guarded(self, _fake_handle): fake = MagicMock() o = Orchestrator(fake, w) args = TaskArgs() - args.add_tensor(Tensor.make(0, (16,), DataType.FLOAT32, child_memory=False), TensorArgType.INPUT) + args.add_tensor(ChipTensor.make(0, (16,), DataType.FLOAT32, child_memory=False), TensorArgType.INPUT) o.submit_next_level(object(), args, None, worker=0) fake.submit_next_level.assert_called_once() diff --git a/tests/ut/py/test_worker/test_group_task.py b/tests/ut/py/test_worker/test_group_task.py index 0671b8ccc9..739c4af5b5 100644 --- a/tests/ut/py/test_worker/test_group_task.py +++ b/tests/ut/py/test_worker/test_group_task.py @@ -28,9 +28,9 @@ from multiprocessing.shared_memory import SharedMemory from simpler.task_interface import ( + ChipTensor, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -59,7 +59,7 @@ def _sync_args(ptr: int, tag: TensorArgType) -> TaskArgs: value just needs to be a unique non-zero key. """ args = TaskArgs() - args.add_tensor(Tensor.make(ptr, (1,), DataType.UINT8), tag) + args.add_tensor(ChipTensor.make(ptr, (1,), DataType.UINT8), tag) return args diff --git a/tests/ut/py/test_worker/test_host_addr_rewrite.py b/tests/ut/py/test_worker/test_host_addr_rewrite.py index fde974cc0b..61c1bae8bf 100644 --- a/tests/ut/py/test_worker/test_host_addr_rewrite.py +++ b/tests/ut/py/test_worker/test_host_addr_rewrite.py @@ -27,7 +27,7 @@ def _make_blob(tensors: list[tuple[int, int]]) -> bytearray: - """Build a task-args blob: [int32 T][int32 S][Tensor*T]. + """Build a task-args blob: [int32 T][int32 S][ChipTensor*T]. ``tensors`` is a list of ``(addr, child_memory)``. Only the two fields the rewrite reads (buffer.addr at offset 0, child_memory at its struct offset) diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index 97b156df1c..9634d7f61a 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -6684,7 +6684,7 @@ def orch(o, args, cfg): class TestOrchAlloc: def test_alloc_returns_valid_tensor(self): - """alloc returns a Tensor whose data ptr is non-zero and writeable.""" + """alloc returns a ChipTensor whose data ptr is non-zero and writeable.""" captured = [] hw = Worker(level=3, num_sub_workers=1) @@ -6798,7 +6798,7 @@ def orch(o, args, cfg): class TestSubCallableArgs: def test_sub_callable_receives_tensor_metadata(self): """Sub callable receives TaskArgs with correct tensor count and shape.""" - from simpler.task_interface import Tensor # noqa: PLC0415 + from simpler.task_interface import ChipTensor # noqa: PLC0415 result_shm, result_buf = _make_shared_counter() try: @@ -6816,7 +6816,7 @@ def check_args(args): # Use a synthetic non-zero pointer — sub callable only checks metadata, # doesn't dereference the pointer. - ct = Tensor.make(0xCAFE0000, (4,), DataType.FLOAT32) + ct = ChipTensor.make(0xCAFE0000, (4,), DataType.FLOAT32) def orch(o, args, cfg): sub_args = TaskArgs() diff --git a/tests/ut/py/test_worker/test_l3_l2_message_queue.py b/tests/ut/py/test_worker/test_l3_l2_message_queue.py index 5e5302cceb..202bea6932 100644 --- a/tests/ut/py/test_worker/test_l3_l2_message_queue.py +++ b/tests/ut/py/test_worker/test_l3_l2_message_queue.py @@ -36,7 +36,7 @@ WaitCmp, ) from simpler.orchestrator import Orchestrator -from simpler.task_interface import DataType, Tensor, get_element_size +from simpler.task_interface import ChipTensor, DataType, get_element_size from simpler.worker import _IDLE, _OFF_STATE, Worker, _buffer_field_addr, _mailbox_store_i32 @@ -111,7 +111,7 @@ def alloc(self, shape, dtype): storage_t = ctypes.c_uint8 * nbytes storage = storage_t() self._buffers.append(storage) - return Tensor.make(ctypes.addressof(storage), tuple(int(x) for x in shape), dtype) + return ChipTensor.make(ctypes.addressof(storage), tuple(int(x) for x in shape), dtype) class _FakeClient: diff --git a/tests/ut/py/test_worker/test_l3_l2_orch_comm.py b/tests/ut/py/test_worker/test_l3_l2_orch_comm.py index 7291c70b03..f32aaa8fec 100644 --- a/tests/ut/py/test_worker/test_l3_l2_orch_comm.py +++ b/tests/ut/py/test_worker/test_l3_l2_orch_comm.py @@ -25,7 +25,7 @@ SignalTestResult, WaitCmp, ) -from simpler.task_interface import DataType, Tensor +from simpler.task_interface import ChipTensor, DataType from simpler.worker import ( _IDLE, _OFF_STATE, @@ -222,7 +222,7 @@ def test_sim_direct_region_uses_lifecycle_control_and_l3_host_metadata(monkeypat ) region = worker._create_l3_l2_region(0, 64, 128) - payload = Tensor.make(0x1234_0000, (16,), DataType.UINT8) + payload = ChipTensor.make(0x1234_0000, (16,), DataType.UINT8) region.payload_write(0, payload, nbytes=8) region.payload_read(8, payload, nbytes=8) result = region.counter(64).test(7, WaitCmp.EQ) @@ -1264,7 +1264,7 @@ def test_sim_direct_transfer_failure_poisons_only_region(monkeypatch): ) region = worker._create_l3_l2_region(0, 64, 128) - payload = Tensor.make(0x1234_0000, (16,), DataType.UINT8) + payload = ChipTensor.make(0x1234_0000, (16,), DataType.UINT8) with pytest.raises(RuntimeError, match="copy failed"): region.payload_write(0, payload, nbytes=8) with pytest.raises(RuntimeError, match="poisoned"):