diff --git a/docs/buffer-handle-abi.md b/docs/buffer-handle-abi.md new file mode 100644 index 0000000000..eee4e2ef9c --- /dev/null +++ b/docs/buffer-handle-abi.md @@ -0,0 +1,119 @@ +# BufferHandle / BufferRef — 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. + +The wire ABI is frozen in +[`.docs/L3-new/worker-memory-model/bufferhandle-abi.md`](../.docs/L3-new/worker-memory-model/bufferhandle-abi.md); +this page is the user-facing how-to. + +## Three types + +| Type | What it is | Where it lives | +| ---- | ---------- | -------------- | +| **`BufferHandle`** | An owned backing (POSIX shm / fork-COW / device malloc) with a canonical identity + lifecycle. Stays with the Worker that created it. | owner side (L3+) | +| **`BufferRef`** | A **self-describing view**: the full handle descriptor embedded + a strided view `(byte_offset, shapes, strides, dtype)`. The wire element of `TaskArgs`. Carries no materialized address. | on the wire (L3→L2, L4→L3) | +| **`Tensor`** | The materialized form (address + strided view). Exists **only** at the L2 device-runtime boundary. | L2 leaf | + +**L3+ orch functions never touch a C++ `Tensor`.** They allocate handles, build +views, and submit refs. A `Tensor` only appears when a `BufferRef` is +materialized on the L2 chip child. + +## Allocating buffers + +```python +h = worker.create_buffer(nbytes) # kind3: explicit shared host buffer (POSIX shm) +h = worker.alloc_shared_tensor((M, N), dtype) # kind3: shape-sized create_buffer +# inside an orch fn, for device (chip-private) memory: +d = orch.alloc_child_tensor(worker, (M, N), dtype) # kind4: DEVICE_MALLOC on that chip +``` + +`create_buffer` returns a `BufferHandle` backed by a POSIX shm the consumer maps +lazily on first receipt of a ref 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 its ref must be dispatched only to that worker. + +## Building views (the view algebra) + +`handle.ref(...)` names a view; the view algebra mirrors the C++ `Tensor` +(pure metadata, returns a new `BufferRef`): + +```python +v = h.ref(shapes=(M, N), dtype) # contiguous full view (row-major strides) +v.slice(dim, start, end, step=1) # any (incl. strided) view +v.transpose(x, y) # swap two dims +v.permute((1, 0)) # reorder dims +v.view(shapes, offsets) # sub-region +v.reshape(new_shapes) # contiguous only +``` + +`slice` / `transpose` / `permute` / `view` are unconstrained (strided views are +supported); `reshape` requires a contiguous view — exactly the `Tensor` +constraints. Strides are element strides (> 0); `byte_offset` is a byte offset. + +## Submitting a task + +```python +ta = TaskArgs() +ta.add_ref(a_h.ref((SIZE,), f32), ArgDirection.INPUT) +ta.add_ref(out_h.ref((SIZE,), f32), ArgDirection.OUTPUT_EXISTING) +orch.submit_next_level(chip_handle, ta, cfg, worker=0) +``` + +`TaskArgs` carries `BufferRef`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 refs only +result = torch.frombuffer(out_h.shm.buf, dtype=torch.float32, count=n) # after run() +``` + +## How a ref reaches its consumer (three-way split) + +A `BufferRef` on the wire is materialized differently by each consumer: + +| Consumer | What it does | +| -------- | ------------ | +| **Chip leaf (L2 runtime)** | Materialize each ref to a `Tensor` (map-once, keyed by identity), including **strided** views; hand the Tensor blob to `run_from_blob`. | +| **Python sub-worker** (compute) | Map each ref into a `MappedArg`; the callable computes with `torch.frombuffer(arg.buffer, ...)`. No `Tensor`. | +| **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 ref 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. + +## 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 (copy-on-write inherited), no map | a pre-fork host buffer, zero-copy | +| `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 | + +## Scope / status + +Single-machine (host + device) L3→L2 and L4→L3→L2 dispatch is implemented and +verified in `a2a3sim` and onboard `a2a3`. The remote **receive** side and the +buffer lifecycle robustness (`release_buffer`, in-flight retain / deferred-free) +are later phases (P2); see [`.docs/L3/P1.md`](../.docs/L3/P1.md) for the phase +breakdown. diff --git a/docs/comm-domain.md b/docs/comm-domain.md index 2173ed707e..ce07beff37 100644 --- a/docs/comm-domain.md +++ b/docs/comm-domain.md @@ -43,10 +43,11 @@ allocation: if `sum(b.nbytes) > window_size`, `allocate_domain` raises | `device_ctx` | pointer to the device-side `CommContext` (pass as a kernel scalar) | | `local_window_base` | base device address of this rank's window | | `actual_window_size` | window size actually allocated | -| `buffer_ptrs` | `{buffer_name: device_ptr}` for each `CommBufferSpec` | +| `buffers` | `{buffer_name: BufferHandle}` for each `CommBufferSpec` (device `VMM_WINDOW`, owned by this chip) | Kernels read peer windows through `device_ctx` (which holds every rank's -window base, local + imported peer); `buffer_ptrs[name]` is the local slice. +window base, local + imported peer); `buffers[name]` is the local slice — +name it in a task arg with `buffers[name].ref(shapes, dtype)`. --- @@ -64,7 +65,7 @@ The handle is a context manager. Its lifecycle has **two distinct states**: This split exists because `submit_next_level()` only *enqueues* DAG work; `Worker.run()` does not wait for completion until the orch function returns. If `release()` freed memory immediately on `with`-exit, a still-queued task that -captured the domain's `device_ctx` / `buffer_ptrs` would read freed memory. So +captured the domain's `device_ctx` / `buffers` would read freed memory. So **release is deferred**: `release()` flips `released` and queues the backend free; the real free runs after the run fence, when every task that could reference the window has completed. @@ -151,14 +152,15 @@ To preload host data (rather than have a kernel write the window), use `orch.copy_to`: ```python -orch.copy_to(chip_idx, dst=handle[chip_idx].buffer_ptrs["input"], src=tensor.data_ptr(), size=n) +orch.copy_to(handle[chip_idx].buffers["input"], tensor) ``` -`copy_to` is **synchronous** (control-mailbox round-trip + synchronous -`rtMemcpy` H2D): when it returns, the bytes are in that rank's window. `src` -must be device-visible from the forked chip child — e.g. a `torch` tensor moved -to shared memory with `.share_memory_()` **before** `Worker.init()` forks the -chips. +`copy_to(dst_handle, src)` is **synchronous** (control-mailbox round-trip + +synchronous `rtMemcpy` H2D): when it returns, the bytes are in that rank's +window. `dst` is the window's `VMM_WINDOW` BufferHandle; `src` (a `torch` tensor +or writable buffer) must be device-visible from the forked chip child — e.g. a +tensor moved to shared memory with `.share_memory_()` **before** `Worker.init()` +forks the chips. **Cross-rank ordering:** when a kernel reads a *peer's* staged window, stage **all** ranks' windows before submitting any kernel — `copy_to` is synchronous @@ -168,7 +170,7 @@ rank's producer run before another rank has finished staging: ```python with orch.allocate_domain(...) as handle: for chip_idx in handle.workers: # stage all first - orch.copy_to(chip_idx, dst=handle[chip_idx].buffer_ptrs["input"], src=..., size=n) + orch.copy_to(handle[chip_idx].buffers["input"], tensor) for chip_idx in handle.workers: # then submit orch.submit_next_level(chip_handle, args, cfg, worker=chip_idx) ``` @@ -177,52 +179,46 @@ with orch.allocate_domain(...) as handle: ## 6. Host tensor visibility for `worker.run` -A host tensor passed to `worker.run(...)` / `orch.submit_next_level(...)` / -`orch.submit_sub(...)` is ultimately dereferenced from a forked local L3 child, -not the parent, so its memory must be backed by pages mapped into that child. -Fork-inherited MAP_SHARED mappings retain their virtual address, while post-fork -worker-allocated buffers may map at a different address and have their pointers -rewritten before decoding. Two sources are legal: +A host tensor named in a task arg is ultimately dereferenced from a forked local +L3 child, not the parent, so its backing must reach that child. Under the +BufferRef ABI an arg is a `BufferRef` carrying a self-describing descriptor; the +child materializes it lazily on first receipt (map-once, keyed by canonical +identity) — there is no eager broadcast and no host-pointer rewrite. Two sources +are legal: | Source | How | Why it works | | ------ | --- | ------------ | -| **fork-inherited** | `tensor.share_memory_()` **before `Worker.init()`** (before the local L3 children are forked) | the child inherits the MAP_SHARED page at fork | -| **worker-allocated post-fork** | `worker.create_host_buffer(nbytes)` after the children exist | born-shared memory attached into every local child, **zero-copy** | +| **fork-inherited** | `tensor.share_memory_()` **before `Worker.init()`**, named with `worker.make_ref_arg(t, shapes, dtype)` (FORK_SHM) | the child inherits the MAP_SHARED page at the fork; the ref resolves to that same VA | +| **worker-allocated post-fork** | `worker.create_buffer(nbytes)` after the children exist, named with `handle.ref(shapes, dtype)` (POSIX_SHM) | the child maps the shm by identity on first receipt of the ref, **zero-copy** | The local L3 children are forked eagerly in `Worker.init()`. A host tensor -created after that — the natural dynamic-shape serving pattern — is invisible to -the children unless it lives in a `create_host_buffer` buffer: +created after that — the natural dynamic-shape serving pattern — reaches the +children by naming a `create_buffer` handle as a `BufferRef`: ```python worker = Worker(level=3, ...); worker.register(chip); worker.init() # forks the chips -buf_h = worker.create_host_buffer(tokens * hidden_size * 4) # born-shared, post-fork -buf_o = worker.create_host_buffer(batch * vocab * 4) +buf_h = worker.create_buffer(tokens * hidden_size * 4) # POSIX shm, post-fork +buf_o = worker.create_buffer(batch * vocab * 4) try: - hidden = torch.frombuffer(buf_h.buffer, dtype=torch.float32).view(tokens, hidden_size) - out = torch.frombuffer(buf_o.buffer, dtype=torch.float32).view(batch, vocab) + hidden = torch.frombuffer(buf_h.shm.buf, dtype=torch.float32, count=tokens * hidden_size) + out = torch.frombuffer(buf_o.shm.buf, dtype=torch.float32, count=batch * vocab) for step in batches: - fill(hidden); worker.run(orch, ...) # no per-run copy — child reads/writes the same pages + fill(hidden) + # name each buffer as a ref in the task args; the child maps it once + worker.run(orch, ...) # no per-run copy — child reads/writes the same pages use(out) - del hidden, out # drop views before free + del hidden, out # drop views before close finally: - worker.free_host_buffer(buf_h) - worker.free_host_buffer(buf_o) + buf_h.close() + buf_o.close() ``` -**Create once, reuse many runs.** `create_host_buffer` maps a shm into each -child and keeps it mapped; the child reads and writes the same physical pages -the parent sees, so there is **no per-run copy**. Build the tensor over -`buf.buffer` (buffer protocol — `torch.frombuffer` / `np.frombuffer`) once and -reuse it; a sub-view (slice) that fits inside the buffer is resolved -automatically. simpler stays framework-free — torch/numpy appear only on the -caller's side. - -**Unregistered post-fork tensors are forwarded unvalidated.** A tensor that is -neither fork-inherited nor from `create_host_buffer` is passed through as-is: -the fork-inherited case is the common legitimate one, so it must keep working. -An anonymous post-fork tensor forwarded this way reads stale/unmapped memory in -the child — allocate it with `create_host_buffer` instead. +**Create once, reuse many runs.** The first ref over a `create_buffer` handle +maps its shm into the child (map-once, cached by identity); later runs reuse the +mapping, so there is **no per-run copy**. Build the tensor over `handle.shm.buf` +(buffer protocol — `torch.frombuffer` / `np.frombuffer`) once and reuse it. +simpler stays framework-free — torch/numpy appear only on the caller's side. ### Contract / limits @@ -230,24 +226,16 @@ the child — allocate it with `create_host_buffer` instead. child; during a `run` the child is reading/writing them, so the parent must not read or write the buffer until `run` returns (same contract as a fork-inherited `.share_memory_()` tensor). In-run cross-task ordering (a - producer task's output read by a consumer task) is still enforced by the - runtime's OverlapMap, keyed on the host address — no host-side copy involved. -- **Shape varies within `nbytes`.** A tensor built over the buffer may take any - shape whose bytes fit; a view that runs past the buffer raises before dispatch - (`overruns its host buffer`). To grow beyond `nbytes`, free and re-create a - larger buffer. Do not free a buffer while a `run` using it is in flight, and - drop every tensor/`memoryview` over `buffer.buffer` before `free_host_buffer` - (or `close()`) so the shm can be released promptly. -- **`orch.copy_to` is the unmanaged low-level path.** `create_host_buffer` - covers the `run` / `submit_next_level` host-tensor args. The explicit - `orch.copy_to(src=tensor.data_ptr())` staging path (§5) is *not* validated — - its `src` must be fork-inherited (`.share_memory_()` before `init()`) - or a `create_host_buffer` buffer. + producer task's output read by a consumer task) is enforced by the runtime's + dependency inference, keyed on the ref's canonical identity — no host-side copy. +- **`orch.copy_to` is the low-level device path.** It stages host bytes into a + domain window (`copy_to(dst_handle, src)`, §5); its `src` must be + fork-inherited (`.share_memory_()` before `init()`) or a `create_buffer` shm. - **Fork-inherited anonymous memory is copy-on-write, hence stale.** Even a tensor the child legitimately inherited is only useful as a *live* input if it is MAP_SHARED: anonymous (non-`share_memory_`) pages are COW, so writes the parent makes *after* fork do not reach the child. A live input must be - file-backed (`.share_memory_()` before `init()`) or a `create_host_buffer` one. + file-backed (`.share_memory_()` before `init()`) or a `create_buffer` one. --- 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 05b480e3f2..5975783c06 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 @@ -23,7 +23,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -31,7 +30,7 @@ from simpler_setup.elf_parser import extract_text_section from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.torch_interop import make_tensor_ref HERE = os.path.dirname(os.path.abspath(__file__)) N = 128 * 128 @@ -129,18 +128,10 @@ def orch_fn(orch, _args, cfg): for rank in range(nranks): domain = handle[rank] args = TaskArgs() - args.add_tensor(make_tensor_arg(inp[rank]), TensorArgType.INPUT) - 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( - data=domain.buffer_ptrs["notify_counter"], - shapes=(1,), - dtype=DataType.INT32, - child_memory=True, - ), - TensorArgType.INPUT, - ) + args.add_ref(make_tensor_ref(worker, inp[rank]), TensorArgType.INPUT) + args.add_ref(make_tensor_ref(worker, out[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(make_tensor_ref(worker, result[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(domain.buffers["notify_counter"].ref((1,), DataType.INT32.value), TensorArgType.INPUT) args.add_scalar(domain.device_ctx) orch.submit_next_level(chip_handle, args, cfg, worker=rank) 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 e1336eec7f..16fe0ff6ed 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 @@ -23,7 +23,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -31,7 +30,7 @@ from simpler_setup.elf_parser import extract_text_section from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.torch_interop import make_tensor_ref HERE = os.path.dirname(os.path.abspath(__file__)) N = 128 * 128 @@ -135,26 +134,10 @@ def orch_fn(orch, _args, cfg): for rank in range(nranks): domain = handle[rank] args = TaskArgs() - args.add_tensor(make_tensor_arg(partial[rank]), TensorArgType.INPUT) - args.add_tensor( - Tensor.make( - data=domain.buffer_ptrs["mailbox"], - shapes=(N,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INOUT, - ) - args.add_tensor(make_tensor_arg(result[rank]), TensorArgType.OUTPUT_EXISTING) - args.add_tensor( - Tensor.make( - data=domain.buffer_ptrs["notify_counter"], - shapes=(1,), - dtype=DataType.INT32, - child_memory=True, - ), - TensorArgType.INPUT, - ) + args.add_ref(make_tensor_ref(worker, partial[rank]), TensorArgType.INPUT) + args.add_ref(domain.buffers["mailbox"].ref((N,), DataType.FLOAT32.value), TensorArgType.INOUT) + args.add_ref(make_tensor_ref(worker, result[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(domain.buffers["notify_counter"].ref((1,), DataType.INT32.value), TensorArgType.INPUT) args.add_scalar(domain.device_ctx) orch.submit_next_level(chip_handle, args, cfg, worker=rank) diff --git a/examples/a2a3/tensormap_and_ringbuffer/prefetch_async_demo/test_prefetch_async_demo.py b/examples/a2a3/tensormap_and_ringbuffer/prefetch_async_demo/test_prefetch_async_demo.py index d5db8f7ed6..bec381a014 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/prefetch_async_demo/test_prefetch_async_demo.py +++ b/examples/a2a3/tensormap_and_ringbuffer/prefetch_async_demo/test_prefetch_async_demo.py @@ -38,15 +38,16 @@ ArgDirection, CallConfig, ChipCallable, - ChipStorageTaskArgs, CoreCallable, + TaskArgs, + TensorArgType, ) from simpler.worker import Worker from simpler_setup.elf_parser import extract_text_section from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.torch_interop import make_tensor_ref HERE = os.path.dirname(os.path.abspath(__file__)) RUNTIME = "tensormap_and_ringbuffer" @@ -105,9 +106,9 @@ def run(platform: str = "a2a3", device_id: int = 0) -> int: worker.init() try: handle = worker.register(chip_callable) - args = ChipStorageTaskArgs() - args.add_tensor(make_tensor_arg(src)) - args.add_tensor(make_tensor_arg(out)) + args = TaskArgs() + args.add_ref(make_tensor_ref(worker, src), TensorArgType.INPUT) + args.add_ref(make_tensor_ref(worker, out), TensorArgType.OUTPUT_EXISTING) worker.run(handle, args, CallConfig()) finally: worker.close() 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 3f6a58d4f1..907f082d35 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 @@ -32,7 +32,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -40,7 +39,7 @@ from simpler_setup.elf_parser import extract_text_section from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.torch_interop import make_tensor_ref HERE = os.path.dirname(os.path.abspath(__file__)) N = 128 * 128 @@ -150,26 +149,13 @@ def orch_fn(orch, _args, cfg): # each producer TGET_ASYNCs the *peer* rank's window, so all # windows must hold real data before execution begins. for rank in range(nranks): - orch.copy_to( - rank, - dst=handle[rank].buffer_ptrs["input_window"], - src=inputs[rank].data_ptr(), - size=input_nbytes, - ) + orch.copy_to(handle[rank].buffers["input_window"], inputs[rank]) for rank in range(nranks): domain = handle[rank] args = TaskArgs() - args.add_tensor( - Tensor.make( - data=domain.buffer_ptrs["input_window"], - shapes=(N,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INPUT, - ) - args.add_tensor(make_tensor_arg(out[rank]), TensorArgType.OUTPUT_EXISTING) - args.add_tensor(make_tensor_arg(result[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(domain.buffers["input_window"].ref((N,), DataType.FLOAT32.value), TensorArgType.INPUT) + args.add_ref(make_tensor_ref(worker, out[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(make_tensor_ref(worker, result[rank]), TensorArgType.OUTPUT_EXISTING) args.add_scalar(domain.device_ctx) orch.submit_next_level(chip_handle, args, cfg, worker=rank) 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 9dccbb70bf..1d03b7e530 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 @@ -23,7 +23,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -31,7 +30,7 @@ from simpler_setup.elf_parser import extract_text_section from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.torch_interop import make_tensor_ref HERE = os.path.dirname(os.path.abspath(__file__)) N = 128 * 128 @@ -126,18 +125,10 @@ def orch_fn(orch, _args, cfg): for rank in range(nranks): domain = handle[rank] args = TaskArgs() - args.add_tensor(make_tensor_arg(inp[rank]), TensorArgType.INPUT) - 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( - data=domain.buffer_ptrs["notify_counter"], - shapes=(1,), - dtype=DataType.INT32, - child_memory=True, - ), - TensorArgType.INPUT, - ) + args.add_ref(make_tensor_ref(worker, inp[rank]), TensorArgType.INPUT) + args.add_ref(make_tensor_ref(worker, out[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(make_tensor_ref(worker, result[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(domain.buffers["notify_counter"].ref((1,), DataType.INT32.value), TensorArgType.INPUT) args.add_scalar(domain.device_ctx) orch.submit_next_level(chip_handle, args, cfg, worker=rank) 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 14674000b1..2673b6fd7f 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 @@ -23,7 +23,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -31,7 +30,7 @@ from simpler_setup.elf_parser import extract_text_section from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.torch_interop import make_tensor_ref HERE = os.path.dirname(os.path.abspath(__file__)) N = 128 * 128 @@ -135,26 +134,10 @@ def orch_fn(orch, _args, cfg): for rank in range(nranks): domain = handle[rank] args = TaskArgs() - args.add_tensor(make_tensor_arg(partial[rank]), TensorArgType.INPUT) - args.add_tensor( - Tensor.make( - data=domain.buffer_ptrs["mailbox"], - shapes=(N,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INOUT, - ) - args.add_tensor(make_tensor_arg(result[rank]), TensorArgType.OUTPUT_EXISTING) - args.add_tensor( - Tensor.make( - data=domain.buffer_ptrs["notify_counter"], - shapes=(1,), - dtype=DataType.INT32, - child_memory=True, - ), - TensorArgType.INPUT, - ) + args.add_ref(make_tensor_ref(worker, partial[rank]), TensorArgType.INPUT) + args.add_ref(domain.buffers["mailbox"].ref((N,), DataType.FLOAT32.value), TensorArgType.INOUT) + args.add_ref(make_tensor_ref(worker, result[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(domain.buffers["notify_counter"].ref((1,), DataType.INT32.value), TensorArgType.INPUT) args.add_scalar(domain.device_ctx) orch.submit_next_level(chip_handle, args, cfg, worker=rank) 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..84ef994ae6 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 @@ -32,7 +32,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -40,7 +39,7 @@ from simpler_setup.elf_parser import extract_text_section from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.torch_interop import make_tensor_ref HERE = os.path.dirname(os.path.abspath(__file__)) N = 128 * 128 @@ -149,26 +148,13 @@ def orch_fn(orch, _args, cfg): # each producer TGET_ASYNCs the *peer* rank's window, so all # windows must hold real data before execution begins. for rank in range(nranks): - orch.copy_to( - rank, - dst=handle[rank].buffer_ptrs["input_window"], - src=inputs[rank].data_ptr(), - size=input_nbytes, - ) + orch.copy_to(handle[rank].buffers["input_window"], inputs[rank]) for rank in range(nranks): domain = handle[rank] args = TaskArgs() - args.add_tensor( - Tensor.make( - data=domain.buffer_ptrs["input_window"], - shapes=(N,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INPUT, - ) - args.add_tensor(make_tensor_arg(out[rank]), TensorArgType.OUTPUT_EXISTING) - args.add_tensor(make_tensor_arg(result[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(domain.buffers["input_window"].ref((N,), DataType.FLOAT32.value), TensorArgType.INPUT) + args.add_ref(make_tensor_ref(worker, out[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(make_tensor_ref(worker, result[rank]), TensorArgType.OUTPUT_EXISTING) args.add_scalar(domain.device_ctx) orch.submit_next_level(chip_cid, args, cfg, worker=rank) 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..1ecb06cbfe 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 @@ -32,7 +32,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker @@ -40,7 +39,7 @@ from simpler_setup.elf_parser import extract_text_section from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.torch_interop import make_tensor_ref HERE = os.path.dirname(os.path.abspath(__file__)) N = 128 * 128 @@ -170,26 +169,13 @@ def orch_fn(orch, _args, cfg): # each producer TGET_ASYNCs the *peer* rank's window, so all # windows must hold real data before execution begins. for rank in range(nranks): - orch.copy_to( - rank, - dst=handle[rank].buffer_ptrs["input_window"], - src=inputs[rank].data_ptr(), - size=input_nbytes, - ) + orch.copy_to(handle[rank].buffers["input_window"], inputs[rank]) for rank in range(nranks): domain = handle[rank] args = TaskArgs() - args.add_tensor( - Tensor.make( - data=domain.buffer_ptrs["input_window"], - shapes=(N,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INPUT, - ) - args.add_tensor(make_tensor_arg(out[rank]), TensorArgType.OUTPUT_EXISTING) - args.add_tensor(make_tensor_arg(result[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(domain.buffers["input_window"].ref((N,), DataType.FLOAT32.value), TensorArgType.INPUT) + args.add_ref(make_tensor_ref(worker, out[rank]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(make_tensor_ref(worker, result[rank]), TensorArgType.OUTPUT_EXISTING) args.add_scalar(domain.device_ctx) orch.submit_next_level(chip_cid, args, cfg, worker=rank) diff --git a/examples/workers/l2/hello_worker/main.py b/examples/workers/l2/hello_worker/main.py index bf85efbf42..dfc3fb038f 100644 --- a/examples/workers/l2/hello_worker/main.py +++ b/examples/workers/l2/hello_worker/main.py @@ -75,10 +75,10 @@ def run(platform: str, device_id: int) -> int: # # We at least exercise the memory-control path: one malloc/free round # trip confirms the host<->device mailbox works. - ptr = worker.malloc(4096) - assert ptr != 0, "malloc returned NULL — device memory is exhausted or not mapped" - worker.free(ptr) - print(f"[hello_worker] malloc/free round-trip OK (ptr=0x{ptr:x})") + buf = worker.malloc(4096) + assert buf.base != 0, "malloc returned NULL — device memory is exhausted or not mapped" + worker.free(buf) + print(f"[hello_worker] malloc/free round-trip OK (ptr=0x{buf.base:x})") finally: # close() releases the ACL device, unloads the runtime libraries, and # frees all per-worker C++ state. ALWAYS call it in a finally — a diff --git a/examples/workers/l2/per_task_runtime_env/main.py b/examples/workers/l2/per_task_runtime_env/main.py index c25ba099a4..15b9ba8656 100644 --- a/examples/workers/l2/per_task_runtime_env/main.py +++ b/examples/workers/l2/per_task_runtime_env/main.py @@ -44,10 +44,10 @@ ArgDirection, CallConfig, ChipCallable, - ChipStorageTaskArgs, CoreCallable, DataType, - Tensor, + TaskArgs, + TensorArgType, ) from simpler.worker import Worker @@ -157,25 +157,25 @@ def _run_one(worker: Worker, chip_handle, label: str, ring: Optional[dict]) -> N expected = host_a + host_b host_out = torch.zeros(N_ROWS, N_COLS, dtype=torch.float32) - dev_a = worker.malloc(NBYTES) - dev_b = worker.malloc(NBYTES) - dev_out = worker.malloc(NBYTES) - worker.copy_to(dev_a, host_a.data_ptr(), NBYTES) - worker.copy_to(dev_b, host_b.data_ptr(), NBYTES) + a_h = worker.malloc(NBYTES) + b_h = worker.malloc(NBYTES) + out_h = worker.malloc(NBYTES) + worker.copy_to(a_h, host_a) + worker.copy_to(b_h, host_b) - 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 = TaskArgs() + args.add_ref(a_h.ref(shapes=(N_ROWS, N_COLS), dtype=DataType.FLOAT32.value), TensorArgType.INPUT) + args.add_ref(b_h.ref(shapes=(N_ROWS, N_COLS), dtype=DataType.FLOAT32.value), TensorArgType.INPUT) + args.add_ref(out_h.ref(shapes=(N_ROWS, N_COLS), dtype=DataType.FLOAT32.value), TensorArgType.OUTPUT_EXISTING) config = _make_config(ring) print(f"[per_task_runtime_env] run '{label}': runtime_env={config.runtime_env!r}") worker.run(chip_handle, args, config) - worker.copy_from(host_out.data_ptr(), dev_out, NBYTES) - worker.free(dev_a) - worker.free(dev_b) - worker.free(dev_out) + worker.copy_from(host_out, out_h) + worker.free(a_h) + worker.free(b_h) + worker.free(out_h) assert torch.allclose(host_out, expected, rtol=1e-5, atol=1e-5), f"{label} result mismatch" print(f"[per_task_runtime_env] '{label}' golden check PASSED") diff --git a/examples/workers/l2/vector_add/README.md b/examples/workers/l2/vector_add/README.md index e434b80a38..a2839933e8 100644 --- a/examples/workers/l2/vector_add/README.md +++ b/examples/workers/l2/vector_add/README.md @@ -81,31 +81,32 @@ In our orch `func_id=0` is the only value, so one child. dev_a = worker.malloc(NBYTES) dev_b = worker.malloc(NBYTES) dev_out = worker.malloc(NBYTES) -worker.copy_to(dev_a, host_a.ctypes.data, NBYTES) -worker.copy_to(dev_b, host_b.ctypes.data, NBYTES) +worker.copy_to(dev_a, host_a) +worker.copy_to(dev_b, host_b) ``` -`malloc` returns a device pointer (uint64). `copy_to(dst_dev, src_host, n)` -does host→device DMA. +`malloc` returns a device `BufferHandle`. `copy_to(dst_handle, src)` does +host→device DMA (`src` is a torch tensor or any writable buffer). -### 5. Build `ChipStorageTaskArgs`, run +### 5. Build `TaskArgs` (BufferRefs), run ```python -args = ChipStorageTaskArgs() -args.add_tensor(Tensor.make(dev_a, shape, DataType.FLOAT32)) -args.add_tensor(Tensor.make(dev_b, shape, DataType.FLOAT32)) -args.add_tensor(Tensor.make(dev_out, shape, DataType.FLOAT32)) +args = TaskArgs() +args.add_ref(dev_a.ref(shapes=shape, dtype=DataType.FLOAT32.value), TensorArgType.INPUT) +args.add_ref(dev_b.ref(shapes=shape, dtype=DataType.FLOAT32.value), TensorArgType.INPUT) +args.add_ref(dev_out.ref(shapes=shape, dtype=DataType.FLOAT32.value), TensorArgType.OUTPUT_EXISTING) worker.run(chip_handle, args, CallConfig()) # chip_handle = worker.register(chip_callable) before init() ``` -The tensor order must match `signature` order on the `ChipCallable`. `run()` -blocks until the kernel completes. +Each `BufferRef` is a view over its handle; `.ref(shapes, dtype)` takes the +`DataType` int value. The ref order must match `signature` order on the +`ChipCallable`. `run()` blocks until the kernel completes. ### 6. Pull result back, verify, free ```python -worker.copy_from(host_out.ctypes.data, dev_out, NBYTES) +worker.copy_from(host_out, dev_out) worker.free(dev_a); worker.free(dev_b); worker.free(dev_out) np.testing.assert_allclose(host_out, expected, rtol=1e-5, atol=1e-5) ``` diff --git a/examples/workers/l2/vector_add/main.py b/examples/workers/l2/vector_add/main.py index b7d6ec1f68..90880007fc 100644 --- a/examples/workers/l2/vector_add/main.py +++ b/examples/workers/l2/vector_add/main.py @@ -43,10 +43,10 @@ ArgDirection, CallConfig, ChipCallable, - ChipStorageTaskArgs, CoreCallable, DataType, - Tensor, + TaskArgs, + TensorArgType, ) from simpler.worker import Worker @@ -137,35 +137,34 @@ def _run(worker: Worker, chip_handle: CallableHandle): host_out = torch.zeros(N_ROWS, N_COLS, dtype=torch.float32) # --- 2. Allocate device buffers + H2D copy --- - # malloc returns a uint64 device pointer. copy_to takes (dst_dev, src_host, nbytes). - dev_a = worker.malloc(NBYTES) - dev_b = worker.malloc(NBYTES) - dev_out = worker.malloc(NBYTES) - worker.copy_to(dev_a, host_a.data_ptr(), NBYTES) - 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 - # 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)) + # malloc returns a DEVICE_MALLOC BufferHandle. copy_to takes (dst_handle, host_tensor). + a_h = worker.malloc(NBYTES) + b_h = worker.malloc(NBYTES) + out_h = worker.malloc(NBYTES) + worker.copy_to(a_h, host_a) + worker.copy_to(b_h, host_b) + + # --- 3. Build TaskArgs naming each buffer as a BufferRef view. Order must + # match the ``signature`` list in the ChipCallable (IN, IN, OUT). --- + args = TaskArgs() + args.add_ref(a_h.ref(shapes=(N_ROWS, N_COLS), dtype=DataType.FLOAT32.value), TensorArgType.INPUT) + args.add_ref(b_h.ref(shapes=(N_ROWS, N_COLS), dtype=DataType.FLOAT32.value), TensorArgType.INPUT) + args.add_ref(out_h.ref(shapes=(N_ROWS, N_COLS), dtype=DataType.FLOAT32.value), TensorArgType.OUTPUT_EXISTING) # --- 4. Run. CallConfig() defaults are fine for this kernel. --- config = CallConfig() print("[vector_add] running on device...") - # run() returns None; per-stage timing is emitted as [STRACE] log markers - # (see docs/dfx/host-trace.md) — parse with simpler_setup.tools.strace_timing. + # run() materializes the BufferRefs to device Tensors in-process, then dispatches. Returns None; + # per-stage timing is emitted as [STRACE] log markers (see docs/dfx/host-trace.md). worker.run(chip_handle, args, config) # --- 5. D2H copy back + verify --- - worker.copy_from(host_out.data_ptr(), dev_out, NBYTES) + worker.copy_from(host_out, out_h) # --- 6. Free device buffers. Order doesn't matter, but leaking is bad. --- - worker.free(dev_a) - worker.free(dev_b) - worker.free(dev_out) + worker.free(a_h) + worker.free(b_h) + worker.free(out_h) max_diff = float(torch.max(torch.abs(host_out - expected))) print(f"[vector_add] max |host_out - expected| = {max_diff:.3e}") diff --git a/examples/workers/l2/vector_add/test_run_timing.py b/examples/workers/l2/vector_add/test_run_timing.py index 5c3e3c0691..166ce1b296 100644 --- a/examples/workers/l2/vector_add/test_run_timing.py +++ b/examples/workers/l2/vector_add/test_run_timing.py @@ -26,11 +26,12 @@ import re import pytest -from simpler.task_interface import CallConfig, ChipStorageTaskArgs, DataType, Tensor +from simpler.task_interface import CallConfig, DataType, TaskArgs, TensorArgType from simpler.worker import Worker from .main import N_COLS, N_ELEMS, N_ROWS, NBYTES, build_chip_callable +_F32 = DataType.FLOAT32.value _STRACE_RE = re.compile(r"\[STRACE\] .*\bname=(?P\S+)\b.*\bdur=(?P\d+)") @@ -65,13 +66,13 @@ def _drive_one_run(platform: str, device_id: int, *, enable_l2_swimlane: bool = dev_a = worker.malloc(NBYTES) dev_b = worker.malloc(NBYTES) dev_out = worker.malloc(NBYTES) - worker.copy_to(dev_a, host_a.data_ptr(), NBYTES) - worker.copy_to(dev_b, host_b.data_ptr(), NBYTES) + worker.copy_to(dev_a, host_a) + worker.copy_to(dev_b, host_b) - 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 = TaskArgs() + args.add_ref(dev_a.ref((N_ROWS, N_COLS), _F32), TensorArgType.INPUT) + args.add_ref(dev_b.ref((N_ROWS, N_COLS), _F32), TensorArgType.INPUT) + args.add_ref(dev_out.ref((N_ROWS, N_COLS), _F32), TensorArgType.OUTPUT_EXISTING) config = CallConfig() config.enable_l2_swimlane = enable_l2_swimlane @@ -82,7 +83,7 @@ def _drive_one_run(platform: str, device_id: int, *, enable_l2_swimlane: bool = # Verify the output is sane (so we know the kernel actually ran and # the markers aren't from an early-error path). host_out = torch.zeros(N_ROWS, N_COLS, dtype=torch.float32) - worker.copy_from(host_out.data_ptr(), dev_out, NBYTES) + worker.copy_from(host_out, dev_out) worker.free(dev_a) worker.free(dev_b) worker.free(dev_out) diff --git a/examples/workers/l2/worker_malloc/main.py b/examples/workers/l2/worker_malloc/main.py index e6d7ec2cd3..3563c20d99 100644 --- a/examples/workers/l2/worker_malloc/main.py +++ b/examples/workers/l2/worker_malloc/main.py @@ -12,10 +12,10 @@ This example exercises the four host<->device memory primitives on the ``Worker`` API in isolation, without ever calling ``worker.run()``: - * ``worker.malloc(nbytes)`` — allocate device memory, returns uint64 ptr - * ``worker.copy_to(dst, src, n)`` — H2D byte copy - * ``worker.copy_from(dst, src, n)`` — D2H byte copy - * ``worker.free(ptr)`` — release device memory + * ``worker.malloc(nbytes)`` — allocate device memory, returns a BufferHandle + * ``worker.copy_to(dst_handle, src)`` — H2D byte copy (src = host tensor/buffer) + * ``worker.copy_from(dst, src_handle)`` — D2H byte copy (dst = host tensor/buffer) + * ``worker.free(handle)`` — release device memory Why a standalone example for these? On real hardware (a2a3 / a5 onboard) the CANN device context is per-thread, so ``rtMalloc`` only succeeds on a thread @@ -74,13 +74,13 @@ def _round_trip(worker: Worker, nbytes: int, seed: int) -> None: src_buf = (ctypes.c_uint8 * nbytes).from_buffer_copy(src) dst_buf = (ctypes.c_uint8 * nbytes).from_buffer(dst) - dev_ptr = worker.malloc(nbytes) - assert dev_ptr != 0, f"malloc({nbytes}) returned NULL" + dev_h = worker.malloc(nbytes) + assert dev_h.base != 0, f"malloc({nbytes}) returned NULL" try: - worker.copy_to(dev_ptr, ctypes.addressof(src_buf), nbytes) - worker.copy_from(ctypes.addressof(dst_buf), dev_ptr, nbytes) + worker.copy_to(dev_h, src_buf) + worker.copy_from(dst_buf, dev_h) finally: - worker.free(dev_ptr) + worker.free(dev_h) if bytes(dst) != src: # Find first diverging byte to make the failure actionable. @@ -89,7 +89,7 @@ def _round_trip(worker: Worker, nbytes: int, seed: int) -> None: f"round-trip mismatch at byte {first_diff}/{nbytes}: " f"sent 0x{src[first_diff]:02x}, got 0x{dst[first_diff]:02x}" ) - print(f"[worker_malloc] {nbytes:>6} bytes round-trip OK (ptr=0x{dev_ptr:x})") + print(f"[worker_malloc] {nbytes:>6} bytes round-trip OK (ptr=0x{dev_h.base:x})") def run(platform: str, device_id: int) -> int: @@ -112,19 +112,20 @@ def run(platform: str, device_id: int) -> int: # 2. Hold several allocations live concurrently — exercises the # allocator's free-list bookkeeping (no buffer aliasing). print("[worker_malloc] concurrent live allocations:") - ptrs = [worker.malloc(n) for n in SIZES] - assert all(p != 0 for p in ptrs), "concurrent malloc returned NULL" - assert len(set(ptrs)) == len(ptrs), f"allocator handed out duplicate pointers: {ptrs}" - for p in ptrs: - worker.free(p) - print(f"[worker_malloc] {len(ptrs)} concurrent buffers, all distinct, freed cleanly") + handles = [worker.malloc(n) for n in SIZES] + assert all(h.base != 0 for h in handles), "concurrent malloc returned NULL" + bases = [h.base for h in handles] + assert len(set(bases)) == len(bases), f"allocator handed out duplicate pointers: {[hex(b) for b in bases]}" + for h in handles: + worker.free(h) + print(f"[worker_malloc] {len(handles)} concurrent buffers, all distinct, freed cleanly") # 3. Repeat the cycle — confirms free actually returns the slab so # we don't OOM after a handful of iterations. print("[worker_malloc] alloc/free churn:") for _ in range(8): - p = worker.malloc(SIZES[0]) - worker.free(p) + h = worker.malloc(SIZES[0]) + worker.free(h) print(f"[worker_malloc] 8x alloc/free of {SIZES[0]} bytes OK") finally: worker.close() diff --git a/examples/workers/l3/allreduce/main.py b/examples/workers/l3/allreduce/main.py index 8c25f37920..3b13d57e6a 100644 --- a/examples/workers/l3/allreduce/main.py +++ b/examples/workers/l3/allreduce/main.py @@ -38,7 +38,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker # noqa: E402 @@ -46,9 +45,9 @@ from simpler_setup.elf_parser import extract_text_section # noqa: E402 from simpler_setup.kernel_compiler import KernelCompiler # noqa: E402 from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: E402 -from simpler_setup.torch_interop import make_tensor_arg # noqa: E402 HERE = os.path.dirname(os.path.abspath(__file__)) +_F32 = DataType.FLOAT32.value _KERNEL_AIV = os.path.join(HERE, "kernels", "aiv", "allreduce_onephase_kernel.cpp") _KERNEL_ORCH = os.path.join(HERE, "kernels", "orchestration", "allreduce_onephase_orch.cpp") @@ -154,17 +153,15 @@ def orch_fn(orch, _args, cfg): for i in range(nranks): domain = handle[i] chip_args = TaskArgs() - 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( - data=domain.buffer_ptrs["scratch"], - shapes=(float_elems,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INOUT, + chip_args.add_ref( + worker.make_ref_arg(host_inputs[i], shapes=(ALLREDUCE_COUNT,), dtype=_F32), + TensorArgType.INPUT, ) + chip_args.add_ref( + worker.make_ref_arg(host_outputs[i], shapes=(ALLREDUCE_COUNT,), dtype=_F32), + TensorArgType.OUTPUT_EXISTING, + ) + chip_args.add_ref(domain.buffers["scratch"].ref((float_elems,), _F32), TensorArgType.INOUT) chip_args.add_scalar(domain.domain_size) chip_args.add_scalar(domain.device_ctx) orch.submit_next_level(chip_handle, chip_args, cfg, worker=i) diff --git a/examples/workers/l3/child_memory/main.py b/examples/workers/l3/child_memory/main.py index 470847f8c3..c2f6c7dbfc 100644 --- a/examples/workers/l3/child_memory/main.py +++ b/examples/workers/l3/child_memory/main.py @@ -57,14 +57,12 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg HERE = os.path.dirname(os.path.abspath(__file__)) KERNELS_DIR = os.path.normpath(os.path.join(HERE, "../../../a2a3/tensormap_and_ringbuffer/vector_example/kernels")) @@ -124,19 +122,14 @@ def run(platform: str, device_id: int) -> int: """Core logic — callable from both CLI and pytest.""" print(f"[child_memory] platform={platform} device={device_id}") - # --- 1. Host tensors. share_memory_() makes data_ptr() valid in the chip - # child process — the orchestration's copy_to reads it from there. + # --- 1. The H2D source weight. share_memory_() makes data_ptr() valid in the chip child, where + # orch.copy_to reads it. Inputs/outputs are handle-backed shared buffers (create_buffer, below). torch.manual_seed(0) - host_a = torch.full((SIZE,), 2.0, dtype=torch.float32).share_memory_() host_w = torch.full((SIZE,), 3.0, dtype=torch.float32).share_memory_() - host_f1 = torch.zeros(SIZE, dtype=torch.float32).share_memory_() - host_f2 = torch.zeros(SIZE, dtype=torch.float32).share_memory_() - s = host_a + host_w - expected = (s + 1) * (s + 2) + s # --- 2. Worker(level=3, ...) — single chip, no Python sub-workers. We - # still need level=3 (not level=2) because orch.malloc + child_memory - # require the L3 Orchestrator to reach into the chip child via mailbox. + # still need level=3 (not level=2) because orch.malloc + a device (DEVICE_MALLOC) + # task arg require the L3 Orchestrator to reach into the chip child via mailbox. worker = Worker( level=3, platform=platform, @@ -152,45 +145,56 @@ def run(platform: str, device_id: int) -> int: print("[child_memory] init worker...") worker.init() + f32 = DataType.FLOAT32.value + a = f1 = f2 = None try: + # Born-shared input/output buffers (post-fork, attached into the chip child). torch is used + # only here at the run boundary to fill the input and read the outputs. + def _view(handle): + shm = handle.shm + assert shm is not None + return torch.frombuffer(shm.buf, dtype=torch.float32, count=SIZE) + + a_h = worker.create_buffer(NBYTES) + f1_h = worker.create_buffer(NBYTES) + f2_h = worker.create_buffer(NBYTES) + a = _view(a_h) + f1 = _view(f1_h) + f2 = _view(f2_h) + a.fill_(2.0) + f1.zero_() + f2.zero_() + s = a + host_w + expected = (s + 1) * (s + 2) + s def orch_fn(orch, _args, cfg): - # Allocate weight on chip 0. orch.malloc forwards CTRL_MALLOC over - # the mailbox to the chip child, which calls ChipWorker::malloc -> - # device_malloc_ctx -> rtMalloc inside its own bound device context. - dev_w = orch.malloc(worker_id=0, size=NBYTES) - orch.copy_to(worker_id=0, dst=dev_w, src=host_w.data_ptr(), size=NBYTES) - - # child_memory=True: tell the runtime "this pointer already lives - # 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) - - # Two kernel invocations sharing the weight, pinned to chip 0. - for out in (host_f1, host_f2): - a = TaskArgs() - a.add_tensor(make_tensor_arg(host_a), TensorArgType.INPUT) - a.add_tensor(w_dev, TensorArgType.INPUT) - a.add_tensor(make_tensor_arg(out), TensorArgType.OUTPUT_EXISTING) - orch.submit_next_level(chip_handle, a, cfg, worker=0) - - # dev_w is reclaimed by DeviceRunner::finalize on worker.close() — - # we don't orch.free it here, that's the whole point of child_memory. + # Allocate the weight on chip 0 as a DEVICE_MALLOC handle (kind4; successor of the former + # child_memory=True Tensor). alloc_child_tensor rtMallocs inside the chip child's bound + # device context and wraps the pointer as a handle owned by this worker — NOT auto-freed at + # end-of-task, so both kernel calls share the same live weight (reclaimed at worker.close()). + w_h = orch.alloc_child_tensor(worker_id=0, shapes=(SIZE,), dtype=DataType.FLOAT32) + orch.copy_to(w_h, host_w) # H2D: host weight -> the device handle + + for out_h in (f1_h, f2_h): + ta = TaskArgs() + ta.add_ref(a_h.ref(shapes=(SIZE,), dtype=f32), TensorArgType.INPUT) + ta.add_ref(w_h.ref(shapes=(SIZE,), dtype=f32), TensorArgType.INPUT) + ta.add_ref(out_h.ref(shapes=(SIZE,), dtype=f32), TensorArgType.OUTPUT_EXISTING) + orch.submit_next_level(chip_handle, ta, cfg, worker=0) print("[child_memory] running DAG (1 malloc + 1 copy_to + 2 kernel tasks)...") worker.run(orch_fn, args=None, config=CallConfig()) - # --- 3. Verify — both outputs must equal the same golden, and they - # must agree with each other (proves the second invocation read the - # same weight, not freed/garbage memory). - for tag, got in (("f1", host_f1), ("f2", host_f2)): + # --- 3. Verify — both outputs must equal the same golden, and they must agree with each other + # (proves the second invocation read the same weight, not freed/garbage memory). + for tag, got in (("f1", f1), ("f2", f2)): max_diff = float(torch.max(torch.abs(got - expected))) print(f"[child_memory] {tag}: max |got - expected| = {max_diff:.3e}") assert torch.allclose(got, expected, rtol=1e-5, atol=1e-5), f"{tag} mismatch" - assert torch.equal(host_f1, host_f2), "f1 and f2 diverged — weight buffer was not preserved" + assert torch.equal(f1, f2), "f1 and f2 diverged — weight buffer was not preserved" print("[child_memory] golden + cross-invocation checks PASSED") finally: + a = f1 = f2 = None # drop views before close unlinks the shm worker.close() return 0 diff --git a/examples/workers/l3/domain_rank_map/main.py b/examples/workers/l3/domain_rank_map/main.py index ce049e305a..402a5892d9 100644 --- a/examples/workers/l3/domain_rank_map/main.py +++ b/examples/workers/l3/domain_rank_map/main.py @@ -46,7 +46,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker # noqa: E402 @@ -54,9 +53,9 @@ from simpler_setup.elf_parser import extract_text_section # noqa: E402 from simpler_setup.kernel_compiler import KernelCompiler # noqa: E402 from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: E402 -from simpler_setup.torch_interop import make_tensor_arg # noqa: E402 HERE = os.path.dirname(os.path.abspath(__file__)) +_F32 = DataType.FLOAT32.value OVERLAP_DIR = os.path.join(HERE, "../dual_domain_overlap") COUNT = 256 DTYPE_NBYTES = 4 @@ -117,15 +116,7 @@ def build_allreduce_callable(platform: str) -> ChipCallable: def _add_domain_scratch(args: TaskArgs, domain: ChipDomainContext) -> None: - args.add_tensor( - Tensor.make( - data=domain.buffer_ptrs["scratch"], - shapes=(COUNT,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INOUT, - ) + args.add_ref(domain.buffers["scratch"].ref((COUNT,), _F32), TensorArgType.INOUT) args.add_scalar(domain.domain_size) args.add_scalar(domain.device_ctx) @@ -181,11 +172,11 @@ def verify_orch_fn(orch, _args, _cfg): domain = handle[chip_idx] print( f"[domain_rank_map] {name} chip {chip_idx}: domain_rank={domain.domain_rank} " - f"domain_size={domain.domain_size} scratch=0x{domain.buffer_ptrs['scratch']:x}" + f"domain_size={domain.domain_size} scratch=0x{domain.buffers['scratch'].base:x}" ) if domain.domain_rank != expected_rank[name][chip_idx] or domain.domain_size != 2: state["ok"] = False - if domain.device_ctx == 0 or domain.buffer_ptrs["scratch"] == 0: + if domain.device_ctx == 0 or domain.buffers["scratch"].base == 0: state["ok"] = False # A non-member chip is absent from the domain handle. try: @@ -195,7 +186,7 @@ def verify_orch_fn(orch, _args, _cfg): except KeyError: pass # Chip 2's two domains carve distinct scratch slices. - if even[2].buffer_ptrs["scratch"] == tail[2].buffer_ptrs["scratch"]: + if even[2].buffers["scratch"].base == tail[2].buffers["scratch"].base: print("[domain_rank_map] chip 2 domains share a scratch pointer") state["ok"] = False finally: @@ -213,8 +204,14 @@ def _orch_fn(orch, _args, cfg): for worker_idx in DOMAINS[domain_name]: domain = handle[worker_idx] args = TaskArgs() - args.add_tensor(make_tensor_arg(host_inputs[worker_idx]), TensorArgType.INPUT) - args.add_tensor(make_tensor_arg(outputs[domain_name][worker_idx]), TensorArgType.OUTPUT_EXISTING) + args.add_ref( + worker.make_ref_arg(host_inputs[worker_idx], shapes=(COUNT,), dtype=_F32), + TensorArgType.INPUT, + ) + args.add_ref( + worker.make_ref_arg(outputs[domain_name][worker_idx], shapes=(COUNT,), dtype=_F32), + TensorArgType.OUTPUT_EXISTING, + ) _add_domain_scratch(args, domain) orch.submit_next_level(allreduce_handle, args, cfg, worker=worker_idx) diff --git a/examples/workers/l3/dual_domain_overlap/main.py b/examples/workers/l3/dual_domain_overlap/main.py index 16eb0e1736..fc4c2be909 100644 --- a/examples/workers/l3/dual_domain_overlap/main.py +++ b/examples/workers/l3/dual_domain_overlap/main.py @@ -42,7 +42,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker # noqa: E402 @@ -50,7 +49,9 @@ from simpler_setup.elf_parser import extract_text_section # noqa: E402 from simpler_setup.kernel_compiler import KernelCompiler # noqa: E402 from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: E402 -from simpler_setup.torch_interop import make_tensor_arg # noqa: E402 +from simpler_setup.torch_interop import make_tensor_ref # noqa: E402 + +_F32 = DataType.FLOAT32.value HERE = os.path.dirname(os.path.abspath(__file__)) @@ -154,15 +155,7 @@ def _scratch_buffers() -> list[CommBufferSpec]: def _add_domain_scratch(args: TaskArgs, domain: ChipDomainContext) -> None: - args.add_tensor( - Tensor.make( - data=domain.buffer_ptrs["scratch"], - shapes=(COUNT,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INOUT, - ) + args.add_ref(domain.buffers["scratch"].ref((COUNT,), _F32), TensorArgType.INOUT) args.add_scalar(domain.domain_size) args.add_scalar(domain.device_ctx) @@ -213,12 +206,12 @@ def _orch_fn(orch, _args, cfg): print( f"[dual_domain_overlap] {domain_name} chip {worker_idx}: " f"rank={domain.domain_rank}/{domain.domain_size} " - f"scratch=0x{domain.buffer_ptrs['scratch']:x} ctx=0x{domain.device_ctx:x}" + f"scratch=0x{domain.buffers['scratch'].base:x} ctx=0x{domain.device_ctx:x}" ) args = TaskArgs() - args.add_tensor(make_tensor_arg(host_x[worker_idx]), TensorArgType.INPUT) - args.add_tensor( - make_tensor_arg(reduce_out[domain_name][worker_idx]), TensorArgType.OUTPUT_EXISTING + args.add_ref(make_tensor_ref(worker, host_x[worker_idx]), TensorArgType.INPUT) + args.add_ref( + make_tensor_ref(worker, reduce_out[domain_name][worker_idx]), TensorArgType.OUTPUT_EXISTING ) _add_domain_scratch(args, domain) orch.submit_next_level(allreduce_handle, args, cfg, worker=worker_idx) @@ -230,10 +223,12 @@ def affine_orch_fn(orch, _args, cfg): args_list = [] for worker_idx in worker_indices: args = TaskArgs() - args.add_tensor(make_tensor_arg(reduce_out[domain_name][worker_idx]), TensorArgType.INPUT) - args.add_tensor(make_tensor_arg(scale[worker_idx]), TensorArgType.INPUT) - args.add_tensor(make_tensor_arg(bias[worker_idx]), TensorArgType.INPUT) - args.add_tensor(make_tensor_arg(affine_out[domain_name][worker_idx]), TensorArgType.OUTPUT_EXISTING) + args.add_ref(make_tensor_ref(worker, reduce_out[domain_name][worker_idx]), TensorArgType.INPUT) + args.add_ref(make_tensor_ref(worker, scale[worker_idx]), TensorArgType.INPUT) + args.add_ref(make_tensor_ref(worker, bias[worker_idx]), TensorArgType.INPUT) + args.add_ref( + make_tensor_ref(worker, affine_out[domain_name][worker_idx]), TensorArgType.OUTPUT_EXISTING + ) args_list.append(args) orch.submit_next_level_group(affine_handle, args_list, cfg, workers=worker_indices) diff --git a/examples/workers/l3/ep_dispatch_combine/main.py b/examples/workers/l3/ep_dispatch_combine/main.py index 767fa3cf02..d3d2bf2fd4 100644 --- a/examples/workers/l3/ep_dispatch_combine/main.py +++ b/examples/workers/l3/ep_dispatch_combine/main.py @@ -80,7 +80,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker # noqa: E402 @@ -88,7 +87,9 @@ from simpler_setup.elf_parser import extract_text_section # noqa: E402 from simpler_setup.kernel_compiler import KernelCompiler # noqa: E402 from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: E402 -from simpler_setup.torch_interop import make_tensor_arg # noqa: E402 +from simpler_setup.torch_interop import make_tensor_ref # noqa: E402 + +_F32 = DataType.FLOAT32.value HERE = os.path.dirname(os.path.abspath(__file__)) @@ -553,28 +554,20 @@ def orch_fn(orch, _args, cfg): print( f"[ep_dispatch] chip {i}: rank={domain.domain_rank}/{domain.domain_size} " f"window=[0x{domain.local_window_base:x} +{domain.actual_window_size}B] " - f"scratch=0x{domain.buffer_ptrs['scratch']:x}" + f"scratch=0x{domain.buffers['scratch'].base:x}" ) chip_args = TaskArgs() - chip_args.add_tensor(make_tensor_arg(indices_per_rank[i]), TensorArgType.INPUT) - chip_args.add_tensor(make_tensor_arg(x_norms[i]), TensorArgType.INPUT) - chip_args.add_tensor(make_tensor_arg(w_padded_list[i]), TensorArgType.INPUT) - chip_args.add_tensor(make_tensor_arg(idx_padded_list[i]), TensorArgType.INPUT) - chip_args.add_tensor(make_tensor_arg(recv_x_outs[i]), TensorArgType.OUTPUT_EXISTING) - chip_args.add_tensor(make_tensor_arg(recv_w_outs[i]), TensorArgType.OUTPUT_EXISTING) - chip_args.add_tensor(make_tensor_arg(recv_idx_outs[i]), TensorArgType.OUTPUT_EXISTING) - chip_args.add_tensor(make_tensor_arg(recv_count_outs[i]), TensorArgType.OUTPUT_EXISTING) - 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( - data=domain.buffer_ptrs["scratch"], - shapes=(SCRATCH_NBYTES // 4,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INOUT, - ) + chip_args.add_ref(make_tensor_ref(worker, indices_per_rank[i]), TensorArgType.INPUT) + chip_args.add_ref(make_tensor_ref(worker, x_norms[i]), TensorArgType.INPUT) + chip_args.add_ref(make_tensor_ref(worker, w_padded_list[i]), TensorArgType.INPUT) + chip_args.add_ref(make_tensor_ref(worker, idx_padded_list[i]), TensorArgType.INPUT) + chip_args.add_ref(make_tensor_ref(worker, recv_x_outs[i]), TensorArgType.OUTPUT_EXISTING) + chip_args.add_ref(make_tensor_ref(worker, recv_w_outs[i]), TensorArgType.OUTPUT_EXISTING) + chip_args.add_ref(make_tensor_ref(worker, recv_idx_outs[i]), TensorArgType.OUTPUT_EXISTING) + chip_args.add_ref(make_tensor_ref(worker, recv_count_outs[i]), TensorArgType.OUTPUT_EXISTING) + chip_args.add_ref(make_tensor_ref(worker, recv_y_outs[i]), TensorArgType.OUTPUT_EXISTING) + chip_args.add_ref(make_tensor_ref(worker, routed_y_outs[i]), TensorArgType.OUTPUT_EXISTING) + chip_args.add_ref(domain.buffers["scratch"].ref((SCRATCH_NBYTES // 4,), _F32), TensorArgType.INOUT) chip_args.add_scalar(domain.domain_size) chip_args.add_scalar(domain.device_ctx) orch.submit_next_level(chip_handle, chip_args, cfg, worker=i) diff --git a/examples/workers/l3/ffn_tp_parallel/main.py b/examples/workers/l3/ffn_tp_parallel/main.py index 7379dbf28d..d1684f9439 100644 --- a/examples/workers/l3/ffn_tp_parallel/main.py +++ b/examples/workers/l3/ffn_tp_parallel/main.py @@ -46,7 +46,6 @@ CoreCallable, DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker # noqa: E402 @@ -54,7 +53,9 @@ from simpler_setup.elf_parser import extract_text_section # noqa: E402 from simpler_setup.kernel_compiler import KernelCompiler # noqa: E402 from simpler_setup.pto_isa import ensure_pto_isa_root # noqa: E402 -from simpler_setup.torch_interop import make_tensor_arg # noqa: E402 +from simpler_setup.torch_interop import make_tensor_ref # noqa: E402 + +_F32 = DataType.FLOAT32.value HERE = os.path.dirname(os.path.abspath(__file__)) @@ -206,31 +207,23 @@ def orch_fn(orch, _args, cfg): print( f"[ffn_tp_parallel] chip {i}: rank={domain.domain_rank}/{domain.domain_size} " f"window=[0x{domain.local_window_base:x} +{domain.actual_window_size}B] " - f"scratch=0x{domain.buffer_ptrs['scratch']:x}" + f"scratch=0x{domain.buffers['scratch'].base:x}" ) # Stage 1: AIC matmul. partial_local is OUTPUT_EXISTING here; - # the framework records its buffer.addr as a producer. + # the framework records its ref identity as a producer. a1 = TaskArgs() - a1.add_tensor(make_tensor_arg(host_x_shards[i]), TensorArgType.INPUT) - a1.add_tensor(make_tensor_arg(host_w_shards[i]), TensorArgType.INPUT) - a1.add_tensor(make_tensor_arg(host_partial[i]), TensorArgType.OUTPUT_EXISTING) + a1.add_ref(make_tensor_ref(worker, host_x_shards[i]), TensorArgType.INPUT) + a1.add_ref(make_tensor_ref(worker, host_w_shards[i]), TensorArgType.INPUT) + a1.add_ref(make_tensor_ref(worker, host_partial[i]), TensorArgType.OUTPUT_EXISTING) orch.submit_next_level(ffn_handle, a1, cfg, worker=i) - # Stage 2: AIV cross-rank sum. Tagging partial_local INPUT - # with the same buffer.addr makes TensorMap auto-link this - # task as a consumer of stage 1, no explicit barrier needed. + # Stage 2: AIV cross-rank sum. Tagging partial_local INPUT resolves to the same + # memoized ref identity as stage 1's output, so TensorMap auto-links this task as a + # consumer of stage 1, no explicit barrier needed. a2 = TaskArgs() - 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( - data=domain.buffer_ptrs["scratch"], - shapes=(scratch_count,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INOUT, - ) + a2.add_ref(make_tensor_ref(worker, host_partial[i]), TensorArgType.INPUT) + a2.add_ref(make_tensor_ref(worker, host_y[i]), TensorArgType.OUTPUT_EXISTING) + a2.add_ref(domain.buffers["scratch"].ref((scratch_count,), _F32), TensorArgType.INOUT) a2.add_scalar(domain.domain_size) a2.add_scalar(domain.device_ctx) orch.submit_next_level(allreduce_handle, a2, cfg, worker=i) diff --git a/examples/workers/l3/l3_l2_orch_comm_stream/test_l3_l2_orch_comm_stream.py b/examples/workers/l3/l3_l2_orch_comm_stream/test_l3_l2_orch_comm_stream.py index 815aab2429..5f0d24a93d 100644 --- a/examples/workers/l3/l3_l2_orch_comm_stream/test_l3_l2_orch_comm_stream.py +++ b/examples/workers/l3/l3_l2_orch_comm_stream/test_l3_l2_orch_comm_stream.py @@ -74,12 +74,12 @@ def _build_chip_callable(platform: str) -> ChipCallable: ) -def _float_view(tensor): - return (ctypes.c_float * _NUMEL).from_address(int(tensor.data)) +def _float_view(handle): + return (ctypes.c_float * _NUMEL).from_address(int(handle.base)) -def _byte_view(tensor): - return (ctypes.c_uint8 * int(tensor.nbytes())).from_address(int(tensor.data)) +def _byte_view(handle): + return (ctypes.c_uint8 * int(handle.nbytes)).from_address(int(handle.base)) def _write_header(header_tensor, seq: int, opcode: int) -> None: diff --git a/examples/workers/l3/multi_chip_dispatch/main.py b/examples/workers/l3/multi_chip_dispatch/main.py index 56d3690182..349b1fb6eb 100644 --- a/examples/workers/l3/multi_chip_dispatch/main.py +++ b/examples/workers/l3/multi_chip_dispatch/main.py @@ -47,7 +47,7 @@ from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.torch_interop import make_tensor_ref HERE = os.path.dirname(os.path.abspath(__file__)) @@ -164,9 +164,9 @@ def orch_fn(orch, _args, cfg): # OUTPUT_EXISTING = "task writes to this pre-allocated tensor" for i in range(len(device_ids)): chip_args = TaskArgs() - chip_args.add_tensor(make_tensor_arg(host_a[i]), TensorArgType.INPUT) - chip_args.add_tensor(make_tensor_arg(host_b[i]), TensorArgType.INPUT) - chip_args.add_tensor(make_tensor_arg(host_out[i]), TensorArgType.OUTPUT_EXISTING) + chip_args.add_ref(make_tensor_ref(worker, host_a[i]), TensorArgType.INPUT) + chip_args.add_ref(make_tensor_ref(worker, host_b[i]), TensorArgType.INPUT) + chip_args.add_ref(make_tensor_ref(worker, host_out[i]), TensorArgType.OUTPUT_EXISTING) orch.submit_next_level(chip_handle, chip_args, cfg, worker=i) # Sub task that depends on both chip outputs. Tagging the two @@ -174,7 +174,7 @@ def orch_fn(orch, _args, cfg): # both chip tasks to finish before running the sub. sub_args = TaskArgs() for i in range(len(device_ids)): - sub_args.add_tensor(make_tensor_arg(host_out[i]), TensorArgType.INPUT) + sub_args.add_ref(make_tensor_ref(worker, host_out[i]), TensorArgType.INPUT) orch.submit_sub(sub_handle, sub_args) # --- 7. Run the DAG. Worker.run() opens a scope, invokes orch_fn, diff --git a/examples/workers/l3/per_task_runtime_env/main.py b/examples/workers/l3/per_task_runtime_env/main.py index c164133b5f..f097677316 100644 --- a/examples/workers/l3/per_task_runtime_env/main.py +++ b/examples/workers/l3/per_task_runtime_env/main.py @@ -52,7 +52,7 @@ from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.pto_isa import ensure_pto_isa_root -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.torch_interop import make_tensor_ref HERE = os.path.dirname(os.path.abspath(__file__)) # Reuse the L2 vector_add kernel verbatim — this example only varies ring sizing. @@ -199,9 +199,9 @@ def orch_fn(orch, _args, _cfg): # launch, each binding its own rings. for i, spec in enumerate(L2_TASKS): chip_args = TaskArgs() - chip_args.add_tensor(make_tensor_arg(host_a[i]), TensorArgType.INPUT) - chip_args.add_tensor(make_tensor_arg(host_b[i]), TensorArgType.INPUT) - chip_args.add_tensor(make_tensor_arg(host_out[i]), TensorArgType.OUTPUT_EXISTING) + chip_args.add_ref(make_tensor_ref(worker, host_a[i]), TensorArgType.INPUT) + chip_args.add_ref(make_tensor_ref(worker, host_b[i]), TensorArgType.INPUT) + chip_args.add_ref(make_tensor_ref(worker, host_out[i]), TensorArgType.OUTPUT_EXISTING) cfg = _l2_config(_cfg, spec) print(f"[per_task_runtime_env] submit '{spec['label']}': runtime_env={cfg.runtime_env!r}") orch.submit_next_level(chip_handle, chip_args, cfg, worker=0) diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 16c109d92d..80215da9a8 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -556,6 +556,36 @@ void append_cleanup_error(std::string &cleanup_error, const std::string &message cleanup_error += message; } +// access⊆granted (§3 access): an arg's TensorArgType may only request access the handle grants. +// INPUT→READ, OUTPUT_EXISTING→WRITE, INOUT→READWRITE; READWRITE grants everything. NO_DEP / OUTPUT +// and any other tag are unconstrained. Catches e.g. a READ-only FORK_SHM COW buffer tagged +// OUTPUT_EXISTING (a forked child's writes would silently not reach the parent). +void check_access_subset(uint8_t granted, TensorArgType tag) { + auto granted_has = [&](AccessMode need) { + return granted == static_cast(AccessMode::READWRITE) || granted == static_cast(need); + }; + bool ok = true; + switch (tag) { + case TensorArgType::INPUT: + ok = granted_has(AccessMode::READ); + break; + case TensorArgType::OUTPUT_EXISTING: + ok = granted_has(AccessMode::WRITE); + break; + case TensorArgType::INOUT: + ok = (granted == static_cast(AccessMode::READWRITE)); + break; + default: + ok = true; // NO_DEP / OUTPUT / others: unconstrained + break; + } + if (!ok) { + throw std::invalid_argument( + "TaskArgs.add_ref: arg TensorArgType requires access not granted by the buffer (access⊆granted, §3)" + ); + } +} + } // namespace // ============================================================================ @@ -612,7 +642,20 @@ NB_MODULE(_task_interface, m) { // blob walker locates tensor i's fields at i * TENSOR_STRIDE_BYTES without // reimplementing the struct layout. m.attr("TENSOR_STRIDE_BYTES") = static_cast(sizeof(Tensor)); - m.attr("TENSOR_CHILD_MEMORY_OFFSET") = static_cast(offsetof(Tensor, child_memory)); + m.attr("TENSOR_ADDRESS_SPACE_OFFSET") = static_cast(offsetof(Tensor, address_space)); + + // BufferHandle / BufferRef wire ABI (buffer_handle.h). Exported so the Python mirror in + // simpler.buffer_handle can pin its struct formats to the C++ layout and reject drift. + m.attr("BUFFER_ABI_VERSION") = static_cast(BUFFER_ABI_VERSION); + m.attr("BUFFER_DESCRIPTOR_VERSION") = static_cast(BUFFER_DESCRIPTOR_VERSION); + m.attr("MAX_TENSOR_DIMS") = static_cast(MAX_TENSOR_DIMS); + m.attr("BUFFER_REF_BYTES") = static_cast(sizeof(BufferRef)); + m.attr("BUFFER_HANDLE_DESCRIPTOR_BYTES") = static_cast(sizeof(BufferHandleDescriptor)); + m.attr("CANONICAL_IDENTITY_BYTES") = static_cast(sizeof(CanonicalIdentity)); + m.attr("OWNER_INSTANCE_ID_BYTES") = static_cast(OWNER_INSTANCE_ID_BYTES); + m.attr("PATH_MAX_BYTES") = static_cast(PATH_MAX_BYTES); + m.attr("DESC_MAX_BYTES") = static_cast(DESC_MAX_BYTES); + m.attr("BUFFERREF_BLOB_HEADER_BYTES") = static_cast(BUFFERREF_BLOB_HEADER_SIZE); // --- Tensor --- // The unified strided tensor descriptor. Constructed contiguous via make() @@ -633,7 +676,7 @@ NB_MODULE(_task_interface, m) { // start_offset == 0, buffer.size == numel * element_size. return make_tensor_external( reinterpret_cast(static_cast(data)), shp, static_cast(n), dtype, - /*manual_dep=*/false, /*version=*/0, child_memory ? 1 : 0 + /*manual_dep=*/false, /*version=*/0, child_memory ? AddressSpace::DEVICE : AddressSpace::HOST ); }, nb::arg("data"), nb::arg("shapes"), nb::arg("dtype"), nb::arg("child_memory") = false, @@ -678,7 +721,7 @@ NB_MODULE(_task_interface, m) { // Re-establish a contiguous layout over the same buffer base. self.init_external( reinterpret_cast(self.buffer.addr), numel * get_element_size(self.dtype), shp, - static_cast(n), self.dtype, self.version, self.manual_dep, self.child_memory + static_cast(n), self.dtype, self.version, self.manual_dep, self.address_space ); } ) @@ -707,10 +750,10 @@ NB_MODULE(_task_interface, m) { .def_prop_rw( "child_memory", [](const Tensor &self) -> bool { - return self.is_child_memory(); + return self.is_device_memory(); }, [](Tensor &self, bool v) { - self.child_memory = v ? 1 : 0; + self.address_space = v ? AddressSpace::DEVICE : AddressSpace::HOST; } ) @@ -753,7 +796,7 @@ NB_MODULE(_task_interface, m) { os << self.shapes[i]; } os << "), dtype=" << get_dtype_name(self.dtype); - if (self.is_child_memory()) os << ", child_memory=True"; + if (self.is_device_memory()) os << ", child_memory=True"; os << ")"; return os.str(); }); @@ -834,26 +877,33 @@ NB_MODULE(_task_interface, m) { .def(nb::init<>()) .def( - "add_tensor", - [](TaskArgs &self, const Tensor &t, TensorArgType tag) { - self.add_tensor(t, tag); + "add_ref", + [](TaskArgs &self, nb::bytes packed, TensorArgType tag) { + if (packed.size() != sizeof(BufferRef)) + throw std::invalid_argument("TaskArgs.add_ref: packed BufferRef must be sizeof(BufferRef) bytes"); + BufferRef r; + std::memcpy(&r, packed.c_str(), sizeof(BufferRef)); + check_access_subset(r.handle.access, tag); + self.add_tensor(r, tag); }, - nb::arg("t"), nb::arg("tag") = TensorArgType::INPUT, - "Add a Tensor with an optional TensorArgType tag (default INPUT)." + nb::arg("packed"), nb::arg("tag") = TensorArgType::INPUT, + "Add a BufferRef (packed bytes, e.g. buffer_handle.BufferRef.pack()) with an optional " + "TensorArgType tag (default INPUT)." ) .def( "add_scalar", &TaskArgs::add_scalar, nb::arg("s"), - "Add a uint64_t scalar. After this, add_tensor() is no longer allowed." + "Add a uint64_t scalar. After this, add_ref() is no longer allowed." ) .def( - "tensor", - [](const TaskArgs &self, int32_t i) -> const Tensor & { - if (i < 0 || i >= self.tensor_count()) throw std::out_of_range("TaskArgs tensor index out of range"); - return self.tensor(i); + "ref", + [](const TaskArgs &self, int32_t i) -> nb::bytes { + if (i < 0 || i >= self.tensor_count()) throw std::out_of_range("TaskArgs ref index out of range"); + const BufferRef &r = self.tensor(i); + return nb::bytes(reinterpret_cast(&r), sizeof(BufferRef)); }, - nb::arg("i"), nb::rv_policy::reference_internal, "Return the Tensor at index i." + nb::arg("i"), "Return the packed BufferRef bytes at index i." ) .def( @@ -1420,16 +1470,6 @@ NB_MODULE(_task_interface, m) { "Launch a callable_id previously staged via register_callable. Returns " "None; per-stage timing is emitted as `[STRACE]` log markers." ) - .def( - "run", - [](ChipWorker &self, int32_t callable_id, TaskArgs &args, const CallConfig &config) { - TaskArgsView view = make_view(args); - self.run(callable_id, view, config); - }, - nb::arg("callable_id"), nb::arg("args"), nb::arg("config"), - "Launch a callable_id from a TaskArgs (used for in-process callers). " - "Returns None; timing is emitted as `[STRACE]` log markers." - ) .def( "run_from_blob", [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity, @@ -1539,7 +1579,7 @@ NB_MODULE(_task_interface, m) { "read_args_from_blob", [](uint64_t blob_ptr) { TaskArgsView view = read_blob(reinterpret_cast(blob_ptr), MAILBOX_ARGS_CAPACITY); - TaskArgs args; + ChipStorageTaskArgs args; for (int32_t i = 0; i < view.tensor_count; i++) { args.add_tensor(view.tensors(i)); } @@ -1549,10 +1589,110 @@ NB_MODULE(_task_interface, m) { return args; }, nb::arg("blob_ptr"), - "Reconstruct a TaskArgs from a length-prefixed blob at blob_ptr. " + "Reconstruct a ChipStorageTaskArgs from a length-prefixed blob at blob_ptr. " "Tags are not preserved (blob wire format strips them)." ); + m.def( + "materialize_bufferref_blob", + [](uint64_t blob_ptr, size_t capacity, nb::dict resolved) -> nb::bytes { + const uint8_t *src = reinterpret_cast(blob_ptr); + BufferRefBlobView view = read_bufferref_blob(src, capacity); + ChipStorageTaskArgs args; + for (int32_t i = 0; i < view.ref_count; i++) { + BufferRef r = view.ref(i); + uint64_t elem = get_element_size(r.dtype); + if (elem == 0) { + throw std::runtime_error("materialize_bufferref_blob: unknown dtype"); + } + if (r.byte_offset % elem != 0) { + throw std::runtime_error("materialize_bufferref_blob: byte_offset is not a multiple of dtype size"); + } + nb::bytes key(reinterpret_cast(&r.handle.identity), sizeof(CanonicalIdentity)); + if (!resolved.contains(key)) { + throw std::runtime_error( + "materialize_bufferref_blob: canonical identity not in the import registry" + ); + } + nb::tuple val = nb::cast(resolved[key]); + auto base = nb::cast(val[0]); + auto addr_space = nb::cast(val[1]); + // The view origin is base + byte_offset (start_offset folded into addr); strides carry + // any non-row-major layout (transpose / permute / step-slice), matching the mainline + // Tensor wire which is natively strided. + Tensor t = make_tensor_strided( + reinterpret_cast(static_cast(base + r.byte_offset)), r.shapes, r.strides, + r.ndims, r.dtype, /*manual_dep=*/false, /*version=*/0, static_cast(addr_space) + ); + args.add_tensor(t); + } + for (int32_t i = 0; i < view.scalar_count; i++) { + args.add_scalar(view.scalars[i]); + } + size_t sz = task_args_blob_size(args); + std::string out(sz, '\0'); + write_blob(reinterpret_cast(out.data()), args); + return nb::bytes(out.data(), sz); + }, + nb::arg("blob_ptr"), nb::arg("capacity"), nb::arg("resolved"), + "Materialize a BufferRef blob into a Tensor blob (write_blob format). Each ref's embedded " + "handle identity is resolved via `resolved` {identity_bytes: (local_base, address_space)}; " + "addr = base + byte_offset. The caller pre-populates `resolved` by materializing each embedded " + "descriptor (see bufferref_blob_descriptors) on first receipt. Strided views (transpose / " + "permute / step-slice) materialize to strided Tensors. Rejects unknown identity, " + "non-dtype-aligned byte_offset." + ); + + m.def( + "bufferref_blob_descriptors", + [](uint64_t blob_ptr, size_t capacity) -> nb::list { + const uint8_t *src = reinterpret_cast(blob_ptr); + BufferRefBlobView view = read_bufferref_blob(src, capacity); + nb::list out; + for (int32_t i = 0; i < view.ref_count; i++) { + BufferRef r = view.ref(i); + out.append(nb::bytes(reinterpret_cast(&r.handle), sizeof(BufferHandleDescriptor))); + } + return out; + }, + nb::arg("blob_ptr"), nb::arg("capacity"), + "Extract each BufferRef's embedded BufferHandleDescriptor (packed bytes) from a BufferRef " + "blob, in ref order. A consumer materializes these lazily on receipt to build the `resolved` " + "map passed to materialize_bufferref_blob." + ); + + m.def( + "bufferref_blob_refs", + [](uint64_t blob_ptr, size_t capacity) -> nb::list { + const uint8_t *src = reinterpret_cast(blob_ptr); + BufferRefBlobView view = read_bufferref_blob(src, capacity); + nb::list out; + for (int32_t i = 0; i < view.ref_count; i++) { + BufferRef r = view.ref(i); + out.append(nb::bytes(reinterpret_cast(&r), sizeof(BufferRef))); + } + return out; + }, + nb::arg("blob_ptr"), nb::arg("capacity"), + "Extract each full packed BufferRef (descriptor + view) from a BufferRef blob, in ref order. " + "Used by an L3 orch endpoint to re-export each backing under a local identity before " + "forwarding to L2 (no BufferRef pass-through)." + ); + + m.def( + "bufferref_blob_scalars", + [](uint64_t blob_ptr, size_t capacity) -> nb::list { + const uint8_t *src = reinterpret_cast(blob_ptr); + BufferRefBlobView view = read_bufferref_blob(src, capacity); + nb::list out; + for (int32_t i = 0; i < view.scalar_count; i++) { + out.append(view.scalars[i]); + } + return out; + }, + nb::arg("blob_ptr"), nb::arg("capacity"), "Extract the scalar args (uint64) from a BufferRef blob, in order." + ); + nb::class_(m, "_L2ChildOnboardRegionExport") .def_ro("device_addr", &L2ChildOnboardRegionExport::device_addr) .def_ro("mapping_bytes", &L2ChildOnboardRegionExport::mapping_bytes) diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 63b0f7e43f..979848dfa1 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -307,42 +307,21 @@ inline void bind_worker(nb::module_ &m) { nb::arg("digest"), nb::arg("kind"), nb::arg("target_namespace"), nb::arg("args_list"), "Submit a group of SUB tasks: N args -> N workers, 1 DAG node." ) - .def( - "malloc", - [](Orchestrator &self, int worker_id, size_t size) { - return self.malloc(worker_id, size); - }, - nb::arg("worker_id"), nb::arg("size"), "Allocate memory on next-level worker." - ) - .def( - "free", - [](Orchestrator &self, int worker_id, uint64_t ptr) { - self.free(worker_id, ptr); - }, - nb::arg("worker_id"), nb::arg("ptr"), "Free memory on next-level worker." - ) - .def( - "copy_to", - [](Orchestrator &self, int worker_id, uint64_t dst, uint64_t src, size_t size) { - self.copy_to(worker_id, dst, src, size); - }, - nb::arg("worker_id"), nb::arg("dst"), nb::arg("src"), nb::arg("size"), "Copy host src to worker dst." - ) - .def( - "copy_from", - [](Orchestrator &self, int worker_id, uint64_t dst, uint64_t src, size_t size) { - self.copy_from(worker_id, dst, src, size); - }, - nb::arg("worker_id"), nb::arg("dst"), nb::arg("src"), nb::arg("size"), "Copy worker src to host dst." - ) .def( "alloc", - [](Orchestrator &self, const std::vector &shape, DataType dtype) { - return self.alloc(shape, dtype); + [](Orchestrator &self, const std::vector &shape, DataType dtype, nb::bytes identity) { + if (identity.size() != sizeof(CanonicalIdentity)) { + throw std::invalid_argument("alloc: identity must be sizeof(CanonicalIdentity) bytes"); + } + CanonicalIdentity id; + std::memcpy(&id, identity.c_str(), sizeof(CanonicalIdentity)); + return self.alloc(shape, dtype, id); }, - nb::arg("shape"), nb::arg("dtype"), - "Allocate an intermediate Tensor from the orchestrator's MAP_SHARED " - "pool (visible to forked child workers). Lifetime: until the next Worker.run() call." + nb::arg("shape"), nb::arg("dtype"), nb::arg("identity"), + "Managed HeapRing intermediate (MAP_SHARED, visible to forked children) registered under the " + "identity's canonical hash; returns its heap VA. The caller wraps the VA as a FORK_SHM " + "BufferHandle carrying `identity` so a ref dependency-wires to this slot (the alloc->BufferRef " + "bridge). Backs Worker.alloc_shared_tensor / Orchestrator.alloc (Python)." ) .def( "scope_begin", &Orchestrator::scope_begin, "Open a nested scope. Max nesting depth = MAX_SCOPE_DEPTH (64)." @@ -415,6 +394,34 @@ inline void bind_worker(nb::module_ &m) { "Python-managed (fork + _sub_worker_loop). `child_pid` is that " "forked child, used to detect an exit before mailbox completion." ) + .def( + "malloc", + [](Worker &self, int worker_id, size_t size) { + return self.malloc(worker_id, size); + }, + nb::arg("worker_id"), nb::arg("size"), "Allocate device memory on next-level worker." + ) + .def( + "free", + [](Worker &self, int worker_id, uint64_t ptr) { + self.free(worker_id, ptr); + }, + nb::arg("worker_id"), nb::arg("ptr"), "Free device memory on next-level worker." + ) + .def( + "copy_to", + [](Worker &self, int worker_id, uint64_t dst, uint64_t src, size_t size) { + self.copy_to(worker_id, dst, src, size); + }, + nb::arg("worker_id"), nb::arg("dst"), nb::arg("src"), nb::arg("size"), "H2D copy: host src to worker dst." + ) + .def( + "copy_from", + [](Worker &self, int worker_id, uint64_t dst, uint64_t src, size_t size) { + self.copy_from(worker_id, dst, src, size); + }, + nb::arg("worker_id"), nb::arg("dst"), nb::arg("src"), nb::arg("size"), "D2H copy: worker src to host dst." + ) .def( "add_remote_l3_socket", [](Worker &self, int32_t worker_id, uint64_t session_id, const std::string &transport_name, diff --git a/python/simpler/buffer_handle.py b/python/simpler/buffer_handle.py new file mode 100644 index 0000000000..d2f15d4968 --- /dev/null +++ b/python/simpler/buffer_handle.py @@ -0,0 +1,788 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Owner/consumer side of the BufferHandle ABI (Python mirror of ``buffer_handle.h``). + +Implements the frozen logical schema in ``.docs/L3-new/worker-memory-model/bufferhandle-abi.md``: +16-byte opaque ``owner_instance_id``, length-delimited UTF-8 ``owner_worker_path`` ("L4/L3[2]/L2[5]"), +and a versioned length-delimited ``backend`` body. Struct formats mirror ``buffer_handle.h`` byte for +byte; their sizes are asserted at import against the constants the ``_task_interface`` binding +exports, so a layout drift fails loudly here. +""" + +from __future__ import annotations + +import ctypes +import enum +import os +import struct +from collections.abc import Sequence +from dataclasses import dataclass, replace +from multiprocessing.shared_memory import SharedMemory +from typing import Any + +from _task_interface import ( # pyright: ignore[reportMissingImports] + BUFFER_ABI_VERSION, + BUFFER_DESCRIPTOR_VERSION, + BUFFER_HANDLE_DESCRIPTOR_BYTES, + BUFFER_REF_BYTES, + BUFFERREF_BLOB_HEADER_BYTES, + CANONICAL_IDENTITY_BYTES, + DESC_MAX_BYTES, + MAX_TENSOR_DIMS, + OWNER_INSTANCE_ID_BYTES, + PATH_MAX_BYTES, + DataType, + bufferref_blob_descriptors, + bufferref_blob_refs, + bufferref_blob_scalars, + get_element_size, +) + + +def _dtype_value(dtype: Any) -> int: + """The int wire value of a dtype given either a ``DataType`` enum or its int value. + + The nanobind ``DataType`` enum is not directly ``int()``-able, so callers historically passed + ``dtype.value``; accept either form here to remove that footgun. + """ + return int(dtype.value) if hasattr(dtype, "value") else int(dtype) + + +class AddressSpace(enum.IntEnum): + HOST = 0 + DEVICE = 1 + + +class Visibility(enum.IntEnum): + PRIVATE = 0 + SHARED = 1 + + +class AccessMode(enum.IntEnum): + READ = 0 + WRITE = 1 + READWRITE = 2 + + +class BackendKind(enum.IntEnum): + FORK_SHM = 0 + POSIX_SHM = 1 + VMM_WINDOW = 2 + REMOTE_SIDECAR = 3 + DEVICE_MALLOC = 4 + + +# (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.POSIX_SHM), + (AddressSpace.DEVICE, BackendKind.VMM_WINDOW), + (AddressSpace.DEVICE, BackendKind.DEVICE_MALLOC), + (AddressSpace.HOST, BackendKind.REMOTE_SIDECAR), + (AddressSpace.DEVICE, BackendKind.REMOTE_SIDECAR), + } +) + + +# --------------------------------------------------------------------------- +# Wire struct formats — mirror buffer_handle.h; sizes pinned to the binding. +# --------------------------------------------------------------------------- +# +# CanonicalIdentity (96 B): owner_instance_id[16] opaque, buffer_id u64, generation u32, +# path_len u16, _pad[2], owner_worker_path[PATH_MAX_BYTES] UTF-8. +_CANONICAL_IDENTITY = struct.Struct(f"<{OWNER_INSTANCE_ID_BYTES}sQIH2x{PATH_MAX_BYTES}s") +# BufferHandleDescriptor (216 B) = prefix(8) + CanonicalIdentity(96) + suffix. +_DESC_PREFIX = struct.Struct(" 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)}" + ) + if len(self.owner_worker_path.encode("utf-8")) > PATH_MAX_BYTES: + raise ValueError(f"owner_worker_path exceeds PATH_MAX_BYTES ({PATH_MAX_BYTES})") + + def pack(self) -> bytes: + path = self.owner_worker_path.encode("utf-8") + return _CANONICAL_IDENTITY.pack(self.owner_instance_id, self.buffer_id, self.generation, len(path), path) + + @classmethod + def unpack(cls, raw: bytes) -> CanonicalIdentity: + owner_instance_id, buffer_id, generation, path_len, path = _CANONICAL_IDENTITY.unpack(raw) + if path_len > PATH_MAX_BYTES: + raise ValueError(f"owner_worker_path length {path_len} exceeds PATH_MAX_BYTES ({PATH_MAX_BYTES})") + return cls(owner_instance_id, buffer_id, path[:path_len].decode("utf-8"), generation) + + +@dataclass(frozen=True) +class BufferHandleDescriptor: + """The self-describing handle payload — embedded whole in every BufferRef built over the handle. + + ``body`` is the per-backend materialization (POSIX/fork shm name UTF-8, VMM handle bytes, ...). + """ + + identity: CanonicalIdentity + address_space: AddressSpace + visibility: Visibility + access: AccessMode + backend_kind: BackendKind + nbytes: int + body: bytes = b"" + + 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)" + ) + + 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_ABI_VERSION, + int(self.address_space), + int(self.visibility), + int(self.access), + int(self.backend_kind), + BUFFER_DESCRIPTOR_VERSION, + ) + suffix = _DESC_SUFFIX.pack(self.nbytes, len(self.body), self.body) + return prefix + self.identity.pack() + suffix + + @classmethod + def unpack(cls, raw: bytes) -> BufferHandleDescriptor: + if len(raw) < BUFFER_HANDLE_DESCRIPTOR_BYTES: + raise ValueError(f"descriptor too small: {len(raw)} < {BUFFER_HANDLE_DESCRIPTOR_BYTES}") + abi_version, address_space, visibility, access, backend_kind, descriptor_version = _DESC_PREFIX.unpack_from( + raw, 0 + ) + if abi_version != BUFFER_ABI_VERSION: + raise ValueError(f"unknown BufferHandle abi_version {abi_version} (expected {BUFFER_ABI_VERSION})") + if descriptor_version != BUFFER_DESCRIPTOR_VERSION: + raise ValueError(f"unknown backend descriptor_version {descriptor_version}") + identity = CanonicalIdentity.unpack(raw[_DESC_PREFIX.size : _DESC_PREFIX.size + _CANONICAL_IDENTITY.size]) + nbytes, body_len, body = _DESC_SUFFIX.unpack_from(raw, _DESC_PREFIX.size + _CANONICAL_IDENTITY.size) + return cls( + identity=identity, + address_space=AddressSpace(address_space), + visibility=Visibility(visibility), + access=AccessMode(access), + backend_kind=BackendKind(backend_kind), + nbytes=nbytes, + body=bytes(body[:body_len]), + ) + + +# BufferRef (272 B): BufferHandleDescriptor(216) + byte_offset u64, ndims u32, shapes[MAX] u32, +# strides[MAX] u32, dtype u8, _pad[3]. +_BUFFER_REF_TAIL = struct.Struct(f" tuple[int, ...]: + """Contiguous (row-major) element strides for ``shapes``: strides[i] = prod(shapes[i+1:]).""" + strides = [1] * len(shapes) + for i in range(len(shapes) - 2, -1, -1): + strides[i] = strides[i + 1] * shapes[i + 1] + return tuple(strides) + + +@dataclass(frozen=True) +class BufferRef: + """The blob-carried wire element: a full embedded handle descriptor + a strided view onto it. + + Self-describing — the consumer materializes ``handle`` on first receipt (no prior handshake), + keyed by ``handle.identity``. Carries no materialized address. + + View algebra mirrors the C++ ``Tensor`` (``slice`` / ``transpose`` / ``permute`` / ``view`` / + ``reshape``): pure metadata rewrites of ``(byte_offset, shapes, strides)`` returning a new + ``BufferRef``. ``strides`` are ELEMENT strides (> 0); ``byte_offset`` is a BYTE offset. Matching + ``Tensor``: ``slice`` / ``transpose`` / ``permute`` / ``view`` work on any (incl. strided) view; + ``reshape`` requires a contiguous view (no allocating copy — reach contiguous first). + """ + + handle: BufferHandleDescriptor + byte_offset: int + shapes: tuple[int, ...] + strides: tuple[int, ...] + dtype: int # DataType enum value + + def __post_init__(self) -> None: + if not 0 < len(self.shapes) <= MAX_TENSOR_DIMS: + raise ValueError(f"BufferRef ndims must be in [1, {MAX_TENSOR_DIMS}], got {len(self.shapes)}") + if len(self.strides) != len(self.shapes): + raise ValueError("BufferRef shapes and strides must have equal length") + + def pack(self) -> bytes: + ndims = len(self.shapes) + shapes = list(self.shapes) + [0] * (MAX_TENSOR_DIMS - ndims) + strides = list(self.strides) + [0] * (MAX_TENSOR_DIMS - ndims) + tail = _BUFFER_REF_TAIL.pack(self.byte_offset, ndims, *shapes, *strides, self.dtype) + return self.handle.pack() + tail + + @classmethod + def unpack(cls, raw: bytes) -> BufferRef: + handle = BufferHandleDescriptor.unpack(raw[:BUFFER_HANDLE_DESCRIPTOR_BYTES]) + vals = _BUFFER_REF_TAIL.unpack(raw[BUFFER_HANDLE_DESCRIPTOR_BYTES:BUFFER_REF_BYTES]) + byte_offset, ndims = vals[0], vals[1] + shapes = tuple(vals[2 : 2 + ndims]) + strides = tuple(vals[2 + MAX_TENSOR_DIMS : 2 + MAX_TENSOR_DIMS + ndims]) + dtype = vals[2 + 2 * MAX_TENSOR_DIMS] + return cls(handle=handle, byte_offset=byte_offset, shapes=shapes, strides=strides, dtype=dtype) + + # --- view algebra (zero-copy metadata rewrites, mirroring Tensor) ------------------------------ + + @property + def ndims(self) -> int: + return len(self.shapes) + + def numel(self) -> int: + n = 1 + for s in self.shapes: + n *= s + return n + + @property + def is_contiguous(self) -> bool: + return self.strides == _row_major_strides(self.shapes) + + def _elem_bytes(self) -> int: + return get_element_size(DataType(self.dtype)) + + def slice(self, dim: int, start: int, end: int, step: int = 1) -> BufferRef: + """View ``[start, end)`` with positive ``step`` along ``dim``. Works on any (strided) view.""" + if not 0 <= dim < self.ndims: + raise ValueError(f"slice dim {dim} out of range [0, {self.ndims})") + if step < 1: + raise ValueError(f"slice step must be >= 1, got {step}") + if not 0 <= start < end <= self.shapes[dim]: + raise ValueError(f"slice [{start}, {end}) out of range [0, {self.shapes[dim]}] on dim {dim}") + old_stride = self.strides[dim] + new_shapes = list(self.shapes) + new_strides = list(self.strides) + new_shapes[dim] = (end - start + step - 1) // step + new_strides[dim] = old_stride * step + byte_offset = self.byte_offset + start * old_stride * self._elem_bytes() + return replace(self, byte_offset=byte_offset, shapes=tuple(new_shapes), strides=tuple(new_strides)) + + def transpose(self, x: int, y: int) -> BufferRef: + """Swap dims ``x`` and ``y`` (shapes + strides). Works on any (strided) view.""" + if not (0 <= x < self.ndims and 0 <= y < self.ndims): + raise ValueError(f"transpose dims ({x}, {y}) out of range [0, {self.ndims})") + new_shapes = list(self.shapes) + new_strides = list(self.strides) + new_shapes[x], new_shapes[y] = new_shapes[y], new_shapes[x] + new_strides[x], new_strides[y] = new_strides[y], new_strides[x] + return replace(self, shapes=tuple(new_shapes), strides=tuple(new_strides)) + + def permute(self, order: tuple[int, ...]) -> BufferRef: + """Reorder dims by ``order`` (a permutation of range(ndims)). Works on any (strided) view.""" + if sorted(order) != list(range(self.ndims)): + raise ValueError(f"permute order {order} must be a permutation of range({self.ndims})") + new_shapes = tuple(self.shapes[o] for o in order) + new_strides = tuple(self.strides[o] for o in order) + return replace(self, shapes=new_shapes, strides=new_strides) + + def view(self, view_shapes: tuple[int, ...], view_offsets: tuple[int, ...]) -> BufferRef: + """Sub-view: origin shifts by ``view_offsets`` (per dim), shape becomes ``view_shapes``, + strides unchanged. Each ``view_offsets[i] + view_shapes[i]`` must stay within ``shapes[i]``. + """ + if len(view_shapes) != self.ndims or len(view_offsets) != self.ndims: + raise ValueError(f"view shapes/offsets must have ndims={self.ndims}") + elem = self._elem_bytes() + byte_offset = self.byte_offset + for i in range(self.ndims): + if view_offsets[i] + view_shapes[i] > self.shapes[i]: + raise ValueError( + f"view dim {i}: offset {view_offsets[i]} + shape {view_shapes[i]} exceeds {self.shapes[i]}" + ) + byte_offset += view_offsets[i] * self.strides[i] * elem + return replace(self, byte_offset=byte_offset, shapes=tuple(view_shapes)) + + def reshape(self, new_shapes: tuple[int, ...]) -> BufferRef: + """Contiguous-only reshape (mirrors Tensor::reshape's ``always_assert(is_contiguous)``).""" + if not self.is_contiguous: + raise ValueError("reshape requires a contiguous view; reach contiguous (via a copy) first") + n = 1 + for s in new_shapes: + n *= s + if n != self.numel(): + raise ValueError(f"reshape {new_shapes} (numel {n}) does not match current numel {self.numel()}") + return replace(self, shapes=tuple(new_shapes), strides=_row_major_strides(tuple(new_shapes))) + + +def pack_bufferref_blob(refs: list[BufferRef], scalars: tuple[int, ...] = ()) -> bytes: + """Serialize refs + scalars into the versioned BufferRef blob (mirror of write_bufferref_blob).""" + header = _BUFFERREF_BLOB_HEADER.pack(BUFFER_ABI_VERSION, len(refs), len(scalars), 0) + body = b"".join(ref.pack() for ref in refs) + tail = struct.pack(f"<{len(scalars)}Q", *scalars) if scalars else b"" + return header + body + tail + + +def mint_owner_instance_id() -> bytes: + """A fresh 16-byte opaque nonce, unique per owner incarnation (defends identity against ABA).""" + return os.urandom(OWNER_INSTANCE_ID_BYTES) + + +def _shm_base_addr(shm: SharedMemory) -> int: + """Mapped base address of ``shm``; valid until ``shm.close()``.""" + view = shm.buf + assert view is not None + exporter = ctypes.c_char.from_buffer(view) + addr = ctypes.addressof(exporter) + del exporter + return addr + + +@dataclass +class BufferHandle: + """Owner-side registry object for one shared backing; owns the POSIX shm that backs it.""" + + identity: CanonicalIdentity + address_space: AddressSpace + visibility: Visibility + access: AccessMode + backend_kind: BackendKind + nbytes: int + body: bytes = b"" + shm: SharedMemory | None = None + base: int = 0 + # Owner-side only, never serialized into the descriptor: which next-level worker a DEVICE_MALLOC + # backing lives on (0 for a host backing or an L2 own-device malloc). The device-pointer provenance + # guard and free/copy key on (owner_worker_id, base). + owner_worker_id: int = 0 + + def to_descriptor(self) -> BufferHandleDescriptor: + return BufferHandleDescriptor( + identity=self.identity, + address_space=self.address_space, + visibility=self.visibility, + access=self.access, + backend_kind=self.backend_kind, + nbytes=self.nbytes, + body=self.body, + ) + + def ref( + self, + shapes: tuple[int, ...], + dtype: int | DataType, + strides: tuple[int, ...] | None = None, + byte_offset: int = 0, + ) -> BufferRef: + """A self-describing BufferRef viewing this handle: embeds the full descriptor + the view. + + ``strides`` default to contiguous (row-major) — ``handle.ref(shape, dtype)`` names the whole + buffer as a contiguous view; pass explicit element strides for a strided view (or reach one + via the BufferRef view algebra). ``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 BufferRef( + handle=self.to_descriptor(), + byte_offset=byte_offset, + shapes=shapes, + strides=strides, + dtype=_dtype_value(dtype), + ) + + def close(self) -> None: + if self.shm is not None: + self.shm.close() + self.shm.unlink() + self.shm = None + + +def create_host_shared_buffer( + nbytes: int, + owner_instance_id: bytes, + buffer_id: int, + owner_worker_path: str = "", + generation: int = 0, + visibility: Visibility = Visibility.SHARED, + access: AccessMode = AccessMode.READWRITE, +) -> BufferHandle: + """Allocate a POSIX-shm host backing and wrap it as an owner ``BufferHandle`` (backend POSIX_SHM). + + The backend body is the shm name (UTF-8); the consumer maps it by name in ``ImportRegistry``. + """ + if nbytes <= 0: + raise ValueError(f"create_host_shared_buffer: nbytes must be positive, got {nbytes}") + shm = SharedMemory(create=True, size=nbytes) + identity = CanonicalIdentity(owner_instance_id, buffer_id, owner_worker_path, generation) + return BufferHandle( + identity=identity, + address_space=AddressSpace.HOST, + visibility=visibility, + access=access, + backend_kind=BackendKind.POSIX_SHM, + nbytes=nbytes, + body=shm.name.encode("utf-8"), + shm=shm, + base=_shm_base_addr(shm), + ) + + +def re_export(source: BufferHandleDescriptor) -> BufferHandle: + """Re-export a received handle descriptor for forwarding — identity **invariant**, no mapping. + + Canonical identity is invariant across every edge (frozen model §5/§8): the re-exported ``H'`` + keeps the SOURCE ``(owner_instance_id, owner_worker_path, buffer_id, generation)`` and the SAME + backing (backend_kind / body / nbytes / address_space / visibility / access) as ``source`` — an + L4-owned buffer forwarded L4→L3→L2 carries one identity at all three layers. Only the mapping is + stripped: ``base=0``, ``shm=None`` (no mmap on the forwarding hop); a downstream compute leaf + materializes lazily. Dependency inference keys on the (invariant) identity, so an alias / + retain-release does not split across layers. Re-export is per-backing (memoize by identity), so + pure forwarding carries no per-ref map cost. + """ + return BufferHandle( + identity=source.identity, + address_space=source.address_space, + visibility=source.visibility, + access=source.access, + backend_kind=source.backend_kind, + nbytes=source.nbytes, + body=source.body, + shm=None, + base=0, + ) + + +def remote_sidecar_ref( + shapes: tuple[int, ...], + dtype: int, + nbytes: int, + owner_worker_id: int, + buffer_id: int, + generation: int, + address_space: AddressSpace, +) -> BufferRef: + """Build a ``REMOTE_SIDECAR`` BufferRef 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 BufferRef 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") + identity = CanonicalIdentity(oid, buffer_id, f"remote/{owner_worker_id}", generation) + handle = BufferHandleDescriptor( + identity=identity, + address_space=address_space, + visibility=Visibility.SHARED, + access=AccessMode.READWRITE, + backend_kind=BackendKind.REMOTE_SIDECAR, + nbytes=nbytes, + body=b"", + ) + shapes = tuple(shapes) + return BufferRef( + handle=handle, + byte_offset=0, + shapes=shapes, + strides=_row_major_strides(shapes), + dtype=int(dtype), + ) + + +def wrap_fork_inherited( + data_ptr: int, + nbytes: int, + owner_instance_id: bytes, + buffer_id: int, + owner_worker_path: str = "", + generation: int = 0, + access: AccessMode = AccessMode.READ, +) -> BufferHandle: + """Wrap a pre-fork, fork-inherited host allocation as a zero-copy ``FORK_SHM`` ``BufferHandle``. + + Memory allocated before the children were forked is present in every child at the *same* virtual + address; the backend body is that base VA (u64 LE) and the consumer materializes to the same VA + with no mapping and no copy. Read/write sharing depends on the underlying mmap: + + * a plain allocation is ``MAP_PRIVATE`` — copy-on-write, so a child's first write splits the page + into a private copy the parent never sees. Read-only from the child (input); ``access`` defaults + to READ. + * a ``MAP_SHARED`` allocation (e.g. a ``torch.Tensor.share_memory_()``) is truly shared across the + fork — the child's writes land in the same physical pages the parent reads, so it is usable as an + OUTPUT the parent reads back. Pass ``access=READWRITE`` for that case. + """ + identity = CanonicalIdentity(owner_instance_id, buffer_id, owner_worker_path, generation) + return BufferHandle( + identity=identity, + address_space=AddressSpace.HOST, + visibility=Visibility.SHARED, + access=access, + backend_kind=BackendKind.FORK_SHM, + 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 = 0, + access: AccessMode = AccessMode.READWRITE, + owner_worker_id: int = 0, +) -> BufferHandle: + """Wrap a device pointer (from a worker device malloc) as a ``DEVICE_MALLOC`` ``BufferHandle``. + + The backend body is the device pointer (u64 LE); the consumer materializes to that pointer with no + mapping. The pointer is valid only on the chip that allocated it, so a ref over this handle must be + dispatched only to that chip (a topology invariant, as for the former ``child_memory`` tensor). + ``owner_worker_id`` records which next-level worker the backing lives on for free/copy provenance. + """ + identity = CanonicalIdentity(owner_instance_id, buffer_id, owner_worker_path, generation) + return BufferHandle( + identity=identity, + address_space=AddressSpace.DEVICE, + visibility=Visibility.PRIVATE, + 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 = 0, + access: AccessMode = AccessMode.READWRITE, + owner_worker_id: int = 0, +) -> BufferHandle: + """Wrap a domain-window-carved device VA as a ``VMM_WINDOW`` ``BufferHandle``. + + A comm domain's per-rank window is device memory carved by ``allocate_domain``; each named buffer + slice is one such backing. The backend body is the device VA (u64 LE); the consumer materializes to + that VA with no mapping. The VA is valid only on the chip that owns the window, so a ref over this + handle must be dispatched only to that chip (``owner_worker_id``). Unlike ``DEVICE_MALLOC`` it is + not freed by ``worker.free`` — the domain owns its lifetime and reclaims it at ``release_domain``. + """ + identity = CanonicalIdentity(owner_instance_id, buffer_id, owner_worker_path, generation) + return BufferHandle( + identity=identity, + address_space=AddressSpace.DEVICE, + visibility=Visibility.SHARED, + 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 BufferRef 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 BufferRef's embedded descriptor to a + local base on first receipt (map-once), keyed by canonical identity. + + A consumer calls ``materialize`` for each ref's embedded descriptor as it arrives; the first + sight of an identity maps its backing into this process, later sights reuse the cached base + (a bumped generation is a distinct identity, materialized fresh). Keyed by the packed canonical + identity so lookups are exact — never a numeric-range guess. + """ + + def __init__(self) -> None: + self._by_identity: dict[bytes, ImportedBuffer] = {} + + def materialize(self, descriptor: BufferHandleDescriptor | bytes) -> ImportedBuffer: + """Map ``descriptor``'s backing into this process on first sight of its identity; reuse the + cached ImportedBuffer thereafter (map-once).""" + desc = BufferHandleDescriptor.unpack(descriptor) if isinstance(descriptor, (bytes, bytearray)) else descriptor + key = desc.identity.pack() + cached = self._by_identity.get(key) + if cached is not None: + return cached + if desc.backend_kind in (BackendKind.FORK_SHM, BackendKind.DEVICE_MALLOC, BackendKind.VMM_WINDOW): + # The body is the base pointer (u64 LE), already valid in this process — no mapping. + # FORK_SHM: a COW-inherited host VA. DEVICE_MALLOC / VMM_WINDOW: a device pointer valid on + # the chip that allocated / carved it (the ref must only reach that chip — a topology + # invariant). + base = int.from_bytes(desc.body, "little") + imported = ImportedBuffer(desc.identity, base, desc.nbytes, desc.address_space, None) + elif desc.backend_kind == BackendKind.POSIX_SHM: + shm = SharedMemory(name=desc.body.decode("utf-8")) + imported = ImportedBuffer(desc.identity, _shm_base_addr(shm), desc.nbytes, desc.address_space, shm) + elif desc.backend_kind == BackendKind.REMOTE_SIDECAR: + raise ValueError("ImportRegistry: REMOTE_SIDECAR backend is reserved for P2") + else: + raise NotImplementedError(f"ImportRegistry: backend {desc.backend_kind!r} not supported in P1-B") + self._by_identity[key] = imported + return imported + + def materialize_blob(self, blob_ptr: int, capacity: int) -> dict[bytes, tuple[int, int]]: + """Lazily materialize every embedded descriptor in a BufferRef blob and return the resolved + map for ``materialize_bufferref_blob``: packed identity -> (local base, address_space).""" + for desc_bytes in bufferref_blob_descriptors(blob_ptr, capacity): + self.materialize(desc_bytes) + return self.materialization_map() + + def mapped_args_from_blob(self, blob_ptr: int, capacity: int) -> MappedArgs: + """Materialize a BufferRef blob into a Python compute callable's args: every ref 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 = [ + MappedArg(self.materialize(ref.handle), ref.byte_offset, ref.shapes, ref.strides, ref.dtype) + for ref in (BufferRef.unpack(rb) for rb in bufferref_blob_refs(blob_ptr, capacity)) + ] + return MappedArgs(tensors, tuple(bufferref_blob_scalars(blob_ptr, capacity))) + + def resolve(self, identity: CanonicalIdentity) -> ImportedBuffer: + imported = self._by_identity.get(identity.pack()) + if imported is None: + raise KeyError(f"ImportRegistry: no handle registered for {identity}") + return imported + + def materialization_map(self) -> dict[bytes, tuple[int, int]]: + """Snapshot for ``materialize_bufferref_blob``: packed identity -> (local base, address_space).""" + return {key: (ib.base, int(ib.address_space)) for key, ib in self._by_identity.items()} + + def unregister(self, identity: CanonicalIdentity) -> None: + imported = self._by_identity.pop(identity.pack(), None) + if imported is not None and imported.shm is not None: + imported.shm.close() + + def close(self) -> None: + for imported in self._by_identity.values(): + if imported.shm is not None: + imported.shm.close() + self._by_identity.clear() diff --git a/python/simpler/l3_l2_message_queue.py b/python/simpler/l3_l2_message_queue.py index 656888a9f7..8ed4dc0f59 100644 --- a/python/simpler/l3_l2_message_queue.py +++ b/python/simpler/l3_l2_message_queue.py @@ -17,12 +17,13 @@ from enum import IntEnum from typing import Any +from .buffer_handle import BufferHandle from .l3_l2_orch_comm import ( L3L2OrchRegion, NotifyOp, WaitCmp, ) -from .task_interface import DataType, Tensor +from .task_interface import DataType L3L2_QUEUE_MAGIC = 0x4C335132 L3L2_QUEUE_ABI_MAJOR = 1 @@ -173,7 +174,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 BufferHandle 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 +259,9 @@ def __init__( orch: Any, region: L3L2OrchRegion, layout: L3L2QueueLayout, - desc_fields: Tensor, - desc_seq: Tensor, - desc_read: Tensor, + desc_fields: BufferHandle, + desc_seq: BufferHandle, + desc_read: BufferHandle, ) -> None: self._orch = orch self._region = region @@ -331,16 +332,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) -> BufferHandle: + if not isinstance(buffer, BufferHandle): + raise ValueError("L3-L2 queue requires a registered BufferHandle 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())}") + if int(nbytes) > int(buffer.nbytes): + raise ValueError(f"L3-L2 queue nbytes={nbytes} exceeds registered buffer 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) -> BufferHandle | None: + if not isinstance(buffer, BufferHandle): return None return self._validate_registered_buffer(buffer, nbytes) @@ -388,16 +389,16 @@ def _signal_notify(self, offset: int, value: int) -> None: def _write_descriptor( self, offset: int, seq: int, opcode: L3L2QueueOpcode, payload_offset: int, nbytes: int ) -> None: - fields_buf = (ctypes.c_uint8 * 24).from_address(int(self._desc_fields.data)) + fields_buf = (ctypes.c_uint8 * 24).from_address(int(self._desc_fields.base)) fields_buf[:] = _DESC.pack(0, int(opcode), int(payload_offset), int(nbytes))[8:] - seq_buf = (ctypes.c_uint8 * 8).from_address(int(self._desc_seq.data)) + seq_buf = (ctypes.c_uint8 * 8).from_address(int(self._desc_seq.base)) seq_buf[:] = struct.pack(" L3L2QueueMessage: self._run_primitive(self._region.payload_read, offset, self._desc_read, nbytes=L3L2_QUEUE_DESC_SLOT_BYTES) - raw = ctypes.string_at(int(self._desc_read.data), L3L2_QUEUE_DESC_SLOT_BYTES) + raw = ctypes.string_at(int(self._desc_read.base), L3L2_QUEUE_DESC_SLOT_BYTES) seq, opcode_value, payload_offset, payload_nbytes = _DESC.unpack(raw) try: opcode = L3L2QueueOpcode(opcode_value) diff --git a/python/simpler/l3_l2_orch_comm.py b/python/simpler/l3_l2_orch_comm.py index 4cb5566066..7fc37247d4 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 .buffer_handle import AddressSpace, BufferHandle class NotifyOp(IntEnum): @@ -222,13 +222,11 @@ 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 obj.child_memory: - raise ValueError("L3-L2 payload buffer must be host storage, not child_memory device storage") - if not obj.is_contiguous: - raise ValueError("L3-L2 payload buffer must be contiguous") - self.addr = int(obj.data) - self.nbytes = int(obj.nbytes()) + if isinstance(obj, BufferHandle): + if obj.address_space != AddressSpace.HOST: + raise ValueError("L3-L2 payload buffer must be host storage, not device storage") + self.addr = int(obj.base) + self.nbytes = int(obj.nbytes) return try: diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index 789918ec26..293d88ea2f 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -14,14 +14,14 @@ def my_orch(orch, args, cfg): # chip_handle/sub_handle come from Worker.register(...) - # build the args object yourself; tags drive dependency inference + # build the args object yourself as BufferRefs; tags drive dependency inference a = TaskArgs() - a.add_tensor(make_tensor_arg(input_tensor), TensorArgType.INPUT) - a.add_tensor(make_tensor_arg(output_tensor), TensorArgType.OUTPUT) - orch.submit_next_level(chip_handle, a, cfg, worker=0) + a.add_ref(input_handle.ref(shape, dtype), TensorArgType.INPUT) + a.add_ref(output_handle.ref(shape, dtype), TensorArgType.OUTPUT) + orch.submit_next_level(chip_handle, a, cfg, worker=0) # handle from Worker.register(chip_callable) sub_args = TaskArgs() - sub_args.add_tensor(make_tensor_arg(output_tensor), TensorArgType.INPUT) + sub_args.add_ref(output_handle.ref(shape, dtype), TensorArgType.INPUT) orch.submit_sub(sub_handle, sub_args) handle = w.submit(my_orch, my_args, my_config) @@ -41,6 +41,7 @@ def my_orch(orch, args, cfg): from _task_interface import _Orchestrator as _COrchestrator # pyright: ignore[reportMissingImports] +from .buffer_handle import AccessMode, BufferHandle, CanonicalIdentity, wrap_fork_inherited from .callable_identity import CallableHandle from .task_interface import ( CallConfig, @@ -50,11 +51,11 @@ def my_orch(orch, args, cfg): DataType, RemoteAddressSpace, TaskArgs, - Tensor, _empty_remote_sidecar_for, _remote_sidecar_for, _RemoteTaskArgsSidecar, _validate_remote_sidecar_access, + get_element_size, ) @@ -203,11 +204,6 @@ def submit_next_level(self, callable_handle: Any, args: TaskArgs, config: CallCo raise TypeError("RemoteTensorRef is only supported for RemoteCallable NEXT_LEVEL submits") remote_sidecar = None _validate_remote_sidecar_access(c_args, remote_sidecar) - # Validate the post-fork host buffers of this submit (issue #1027). Only - # the LOCAL_CHIP path dereferences raw host pointers in the forked child; - # zero-copy buffers need no per-run mirror, just an in-range fit check. - if target_namespace == "LOCAL_CHIP" and self._worker is not None: - self._worker._stage_host_buffers_for_chip_submit(c_args) final_worker_ids = _remote_data_eligible_worker_ids(remote_sidecar, eligible_worker_ids) worker = self._worker # Do the (fallible) kind4 provenance analysis BEFORE capturing remote slot @@ -287,11 +283,6 @@ def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eli if remote_sidecars is not None: for c_args, remote_sidecar in zip(c_args_list, remote_sidecars): _validate_remote_sidecar_access(c_args, remote_sidecar) - # Validate post-fork host buffers for chip dispatch (issue #1027), same as - # the single submit path. - if target_namespace == "LOCAL_CHIP" and self._worker is not None: - for c_args in c_args_list: - self._worker._stage_host_buffers_for_chip_submit(c_args) worker_id_sets = ( [ _remote_data_eligible_worker_ids(remote_sidecar, eligible_worker_ids) @@ -384,7 +375,7 @@ def allocate_domain( the IPC handshake (HCCL: aclrtMalloc + IPC import; sim: shm + ftruncate). Returns a ``CommDomainHandle`` whose ``contexts[chip_idx]`` exposes the per-chip ``ChipDomainContext`` (``device_ctx``, ``local_window_base``, - ``buffer_ptrs`` by name). + ``buffers`` by name — each a device ``VMM_WINDOW`` BufferHandle). ``name`` is a local identifier (uniqueness checked against currently-live handles); peers do not need to agree on the string. ``workers`` must be @@ -472,77 +463,66 @@ def scope(self) -> Iterator[Orchestrator]: finally: self._o.scope_end() - def malloc(self, worker_id: int, size: int) -> int: - """Allocate memory on next-level worker *worker_id*. Returns a pointer. + # A Worker is the only allocator. The Orchestrator exposes thin wrappers that delegate to the + # bound Worker's implementation so an orchestration fn can allocate / copy / free without reaching + # for the Worker — each forwards to the Worker's no-lease in-run path (the run already holds it). - This is the single L3 choke for kind4 device memory: ``Worker.malloc`` - also funnels through here, as does a user's direct ``orch.malloc``. The - returned pointer's ``(worker_id, ptr)`` provenance is recorded so a later - free / copy / kind4 dispatch to the wrong worker is rejected. - """ - wid, sz = int(worker_id), int(size) + def alloc_child_tensor(self, worker_id: int, shapes: tuple[int, ...], dtype: DataType) -> BufferHandle: + """Allocate device memory on next-level ``worker_id`` sized for ``shapes`` × ``dtype``; returns a + DEVICE_MALLOC ``BufferHandle``. Delegates to ``Worker.alloc_child_tensor``.""" if self._worker is None: - return int(self._o.malloc(wid, sz)) - with self._worker._child_prov_lock: - ptr = int(self._o.malloc(wid, sz)) - self._worker._child_prov_record_malloc(wid, ptr) - return ptr - - def free(self, worker_id: int, ptr: int) -> None: - """Free memory on next-level worker *worker_id*.""" - wid, p = int(worker_id), int(ptr) + raise RuntimeError("orch.alloc_child_tensor requires a Worker context") + return self._worker.alloc_child_tensor(int(worker_id), tuple(shapes), dtype) + + def free(self, handle: BufferHandle) -> None: + """Free a device ``BufferHandle`` (from ``alloc_child_tensor``). Delegates to ``Worker.free``.""" if self._worker is None: - self._o.free(wid, p) - return - with self._worker._child_prov_lock: - # Safety-first commit barrier: revoke provenance BEFORE the native - # free. If the native free succeeds and an async unwind (e.g. a - # KeyboardInterrupt delivered after the binding returns) fires before - # a post-free clear could run, a freed address would stay live and a - # later copy/dispatch would re-authorize it — a UAF. Revoking first - # turns a native-free failure into a terminal leak (recoverable) but - # never re-authorizes a maybe-freed address. - self._worker._child_prov_require_malloc_base(wid, p, api="free") - self._worker._child_prov_clear_malloc(wid, p) - self._o.free(wid, p) - - def copy_to(self, worker_id: int, dst: int, src: int, size: int) -> None: - """Copy *size* bytes from host *src* to worker *dst*.""" - wid, d = int(worker_id), int(dst) + raise RuntimeError("orch.free requires a Worker context") + self._worker.free(handle) + + def copy_to(self, dst: BufferHandle, src) -> None: + """H2D: copy host ``src`` into device handle ``dst``. Delegates to ``Worker.copy_to``.""" if self._worker is None: - self._o.copy_to(wid, d, int(src), int(size)) - return - with self._worker._child_prov_lock: - self._worker._child_prov_require_live(wid, d, api="copy_to") - self._o.copy_to(wid, d, int(src), int(size)) - - def copy_from(self, worker_id: int, dst: int, src: int, size: int) -> None: - """Copy *size* bytes from worker *src* to host *dst*.""" - wid, s = int(worker_id), int(src) + raise RuntimeError("orch.copy_to requires a Worker context") + self._worker.copy_to(dst, src) + + def copy_from(self, dst, src: BufferHandle) -> None: + """D2H: copy device handle ``src`` into host ``dst``. Delegates to ``Worker.copy_from``.""" if self._worker is None: - self._o.copy_from(wid, int(dst), s, int(size)) - return - with self._worker._child_prov_lock: - self._worker._child_prov_require_live(wid, s, api="copy_from") - self._o.copy_from(wid, int(dst), s, int(size)) - - def alloc(self, shape: Sequence[int], dtype: DataType) -> Tensor: - """Allocate a runtime-managed intermediate buffer. - - Returns a ``Tensor`` 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 - downstream consumers have completed and the run's scope ends. - - Use this for chip-A → chip-B intermediate buffers instead of - pre-allocating with ``torch.share_memory_()`` — the runtime owns - the lifecycle. + raise RuntimeError("orch.copy_from requires a Worker context") + self._worker.copy_from(dst, src) + + def alloc(self, shape: Sequence[int], dtype: DataType) -> BufferHandle: + """Allocate a runtime-managed intermediate buffer; returns a ``BufferHandle``. + + The backing is a MAP_SHARED slab (visible to forked child workers), auto-reclaimed once every + downstream consumer has completed and the run's scope ends — no manual free. Name it in a task + arg with ``handle.ref(shape, dtype)``: the ref's canonical identity dependency-wires to this + alloc's synthetic producer slot (tag it OUTPUT/INOUT on the producer, INPUT on the consumer). + + Use this for chip-A → chip-B intermediate buffers instead of pre-allocating with + ``torch.share_memory_()`` — the runtime owns the lifecycle. Equivalent to + ``worker.alloc_shared_tensor``, additionally registered as an L3-L2 orch-comm host buffer so it + may back an L3-L2 message-queue payload. """ - tensor = self._o.alloc(list(shape), dtype) - if self._worker is not None: - self._worker._register_l3_l2_orch_comm_host_buffer(tensor) - return tensor + assert self._worker is not None, "orch.alloc requires an L3+ orchestration context" + shape_t = tuple(int(s) for s in shape) + nbytes = get_element_size(dtype) + for s in shape_t: + nbytes *= s + oid, buffer_id, path = ( + self._worker._owner_instance_id, + self._worker._next_buffer_id(), + f"L{self._worker.level}", + ) + identity = CanonicalIdentity(oid, buffer_id, path, 0) + # alloc keys the synthetic producer slot by the ref's canonical identity (not a raw VA), so a + # consumer named via handle.ref(...) dependency-wires to it. Same managed backing as + # worker.alloc_shared_tensor; additionally registered as an L3-L2 orch-comm host buffer. + va = int(self._o.alloc(list(shape_t), dtype, identity.pack())) + handle = wrap_fork_inherited(va, int(nbytes), oid, buffer_id, path, generation=0, access=AccessMode.READWRITE) + self._worker._register_l3_l2_orch_comm_host_buffer(handle) + return handle # ------------------------------------------------------------------ # Internal (called by Worker.submit) diff --git a/python/simpler/remote_l3_session.py b/python/simpler/remote_l3_session.py index c62a613779..d7dbff4bbc 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, ChipStorageTaskArgs, Tensor from .worker import Worker sys.modules.setdefault("simpler.remote_l3_session", sys.modules[__name__]) @@ -478,10 +478,10 @@ def _tensor_with_data(tensor: Tensor, data: int) -> Tensor: def _materialize_task_args( # noqa: PLR0912 args: RemoteTaskArgsWire, buffers: dict[tuple[int, ...], _RemoteBufferEntry], worker_id: int -) -> tuple[TaskArgs, list[Any]]: +) -> tuple[ChipStorageTaskArgs, list[Any]]: if len(args.remote_desc) != len(args.tensor_metadata): raise ValueError("remote TASK descriptor count does not match tensor metadata count") - task_args = TaskArgs() + task_args = ChipStorageTaskArgs() keepalive: list[Any] = [] for tensor, sidecar in zip(args.tensor_metadata, args.remote_desc): diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index b18142329c..2090dc1261 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -59,6 +59,8 @@ read_args_from_blob, ) +from .buffer_handle import BufferHandle + def _assert_bindings_match_source_tree() -> None: """Refuse a `_task_interface` built from a different revision of this tree. @@ -645,7 +647,7 @@ class _RemoteTaskArgsStorage: inline_payload: bytearray -_TASK_ARGS_ADD_TENSOR = TaskArgs.add_tensor +_TASK_ARGS_ADD_REF = TaskArgs.add_ref _TASK_ARGS_CLEAR = TaskArgs.clear _REMOTE_TASK_ARGS_STORAGE: weakref.WeakKeyDictionary[TaskArgs, _RemoteTaskArgsStorage] = weakref.WeakKeyDictionary() _REMOTE_TASK_ARGS_STORAGE_LOCK = threading.Lock() @@ -688,18 +690,34 @@ def _storage_for_remote_task_args(args: TaskArgs) -> _RemoteTaskArgsStorage: return storage -def _task_args_add_tensor( - self: TaskArgs, tensor: Tensor | RemoteTensorRef, tag: TensorArgType = TensorArgType.INPUT -) -> None: - if isinstance(tensor, RemoteTensorRef): +def _task_args_add_ref(self: TaskArgs, ref, tag: TensorArgType = TensorArgType.INPUT) -> None: + """Add a BufferRef task arg. ``ref`` is a ``simpler.buffer_handle.BufferRef`` (packable) or its + packed bytes. A RemoteTensorRef (arg destined for a remote worker) is rewritten to a + REMOTE_SIDECAR BufferRef (no local backing) with its remote descriptor tracked in the sidecar.""" + if isinstance(ref, RemoteTensorRef): + from .buffer_handle import AddressSpace, remote_sidecar_ref + storage = _storage_for_remote_task_args(self) - metadata = Tensor.make(0, tensor.shape, tensor.dtype) - _TASK_ARGS_ADD_TENSOR(self, metadata, tag) - storage.sidecars.append(_sidecar_from_ref(storage, tensor)) + handle = ref.handle + inline = handle.address_space == RemoteAddressSpace.HOST_INLINE + nbytes = ref.nbytes + assert nbytes is not None + placeholder = remote_sidecar_ref( + shapes=tuple(int(s) for s in ref.shape), + dtype=int(ref.dtype.value), + nbytes=int(nbytes), + owner_worker_id=0 if inline else int(handle.owner_worker_id), + buffer_id=0 if inline else int(handle._buffer_id), + generation=0 if inline else int(handle._generation), + address_space=( + AddressSpace.DEVICE if handle.address_space == RemoteAddressSpace.REMOTE_DEVICE else AddressSpace.HOST + ), + ) + _TASK_ARGS_ADD_REF(self, placeholder.pack(), tag) + storage.sidecars.append(_sidecar_from_ref(storage, ref)) return - if not isinstance(tensor, Tensor): - raise TypeError("TaskArgs.add_tensor expects Tensor or RemoteTensorRef") - _TASK_ARGS_ADD_TENSOR(self, tensor, tag) + packed = ref.pack() if hasattr(ref, "pack") else ref + _TASK_ARGS_ADD_REF(self, packed, tag) with _REMOTE_TASK_ARGS_STORAGE_LOCK: storage = _REMOTE_TASK_ARGS_STORAGE.get(self) if storage is not None: @@ -712,7 +730,7 @@ def _task_args_clear(self: TaskArgs) -> None: _REMOTE_TASK_ARGS_STORAGE.pop(self, None) -TaskArgs.add_tensor = _task_args_add_tensor +TaskArgs.add_ref = _task_args_add_ref TaskArgs.clear = _task_args_clear @@ -831,9 +849,8 @@ def scalar_to_uint64(value) -> int: class CommBufferSpec: """A named slice of the per-rank communicator window. - Buffers are placed sequentially inside the window in declaration order — Buffers are placed sequentially inside the window in declaration order. - The ``CommDomainHandle.contexts[chip_idx].buffer_ptrs`` dict returned by + The ``CommDomainHandle.contexts[chip_idx].buffers`` dict returned by ``Orchestrator.allocate_domain`` is keyed by ``CommBufferSpec.name``. """ @@ -853,7 +870,9 @@ class ChipDomainContext: device_ctx: int local_window_base: int actual_window_size: int - buffer_ptrs: dict[str, int] + # Each named window slice as a device ``VMM_WINDOW`` BufferHandle owned by this chip. Name a task + # arg with ``buffers[name].ref(shapes, dtype)`` and dispatch it only to this chip (``domain_rank``). + buffers: dict[str, BufferHandle] class CommDomainHandle: @@ -902,7 +921,7 @@ def __getitem__(self, chip_idx: int) -> ChipDomainContext: if self._released: raise RuntimeError( f"CommDomainHandle({self.name!r}) already released; do not pass it to submit_* " - "after release(). Submitted tasks that captured device_ctx / buffer_ptrs before " + "after release(). Submitted tasks that captured device_ctx / buffers before" "release will still see live memory until Worker.run drains." ) return self.contexts[chip_idx] diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 2c3fa7bb6c..5c48140025 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -59,7 +59,6 @@ def my_l4_orch(orch, args, config): from __future__ import annotations -import bisect import contextlib import ctypes import enum @@ -83,7 +82,6 @@ def my_l4_orch(orch, args, config): from _task_interface import ( # pyright: ignore[reportMissingImports] MAX_REGISTERED_CALLABLE_IDS, RUNTIME_ENV_RING_COUNT, - TENSOR_CHILD_MEMORY_OFFSET, WorkerType, _l3_child_onboard_region_close, _l3_child_onboard_region_create, @@ -91,10 +89,30 @@ def my_l4_orch(orch, args, config): _l3_host_mapped_region_import_sim, _mailbox_load_i32, _mailbox_store_i32, - read_args_from_blob, + bufferref_blob_refs, + get_element_size, + materialize_bufferref_blob, ) from . import _log as _simpler_log +from .buffer_handle import ( + AccessMode, + AddressSpace, + BackendKind, + BufferHandle, + BufferHandleDescriptor, + BufferRef, + CanonicalIdentity, + ImportRegistry, + create_host_shared_buffer, + host_ptr_nbytes, + mint_owner_instance_id, + pack_bufferref_blob, + re_export, + wrap_device_malloc, + wrap_fork_inherited, + wrap_vmm_window, +) from .callable_identity import ( CALLABLE_HASH_DIGEST_BYTES, CallableHandle, @@ -139,8 +157,6 @@ def my_l4_orch(orch, args, config): RemoteAddressSpace, RemoteBufferExport, RemoteBufferHandle, - TaskArgs, - Tensor, _Worker, ) @@ -269,16 +285,6 @@ def my_l4_orch(orch, args, config): _CTRL_PY_REGISTER = 10 _CTRL_PY_UNREGISTER = 11 _CTRL_PY_IMPORT_REGISTER = 12 -# Host-buffer registration. MAP_HOST maps a named host-buffer shm -# into every local L3 child *post-fork* and keeps it mapped so later runs can copy -# through it; UNMAP_HOST drops one. The child also records the parent VA range -# the shm stands in for, so the per-task blob's host pointers (raw parent VAs) -# can be rewritten to the child's own mapping before the runtime dereferences -# them. Unlike _CTRL_REGISTER (one-shot H2D then close), these mappings persist -# for the buffer's registered lifetime — see docs/comm-domain.md. -_CTRL_MAP_HOST = 14 -_CTRL_UNMAP_HOST = 15 - # Operation names a child puts in its error message when a control command # fails, so the parent's re-raised text names the operation and not just a # numeric sub-command. Absent entries fall back to the raw number. @@ -290,20 +296,7 @@ def my_l4_orch(orch, args, config): _CTRL_PY_UNREGISTER: "py_unregister", } -# MAP_HOST payload: token (u64), parent_va (u64), nbytes (u64), then the -# NUL-free host-buffer shm name as the trailing bytes. UNMAP_HOST payload is the -# token alone. -_HOST_BUF_MAP_HEADER = struct.Struct(" bool: return self.malloc_owned or bool(self.domain_allocation_ids) -@dataclass -class _HostBufEntry: - """Parent-side record for a born-shared post-fork host buffer. - - The worker owns ``shm`` — a named buffer the local L3 children attach and - read/write through. The user builds a tensor over it (via the buffer - protocol on :class:`HostBuffer`), so the buffer *is* the shm: ``data_ptr == - shm_base`` and no per-run copy is needed (the child reads and writes the same - physical pages the parent sees). ``shm_base`` caches the mapped address. - """ - - token: int - data_ptr: int - nbytes: int - shm: SharedMemory - shm_name: str - shm_base: int - - -@dataclass(frozen=True, eq=False) -class HostBuffer: - """Handle for a worker-allocated, born-shared host buffer (zero-copy). - - Returned by ``Worker.create_host_buffer``. ``buffer`` is a ``memoryview`` - over shared memory already attached into every local L3 child; wrap it with - ``torch.frombuffer`` / ``np.frombuffer`` to get a real tensor whose writes - land directly in the child-visible pages — no per-run copy. ``token`` / - ``data_ptr`` / ``nbytes`` identify the mapping; pass this handle back to - ``free_host_buffer`` to release it. - - ``eq=False`` keeps object-identity equality/hash so the (unhashable) - ``memoryview`` field never blocks using the handle as a dict key or set - member. - """ - - token: int - data_ptr: int - nbytes: int - buffer: memoryview - - -def _rewrite_blob_host_addrs(buf: memoryview, blob_off: int, ranges: list[tuple[int, int, int]]) -> None: - """Redirect registered host pointers in a task-args blob to child mappings. - - ``ranges`` is ``(parent_lo, parent_hi, child_base)`` for each host buffer the - child has mapped via _CTRL_MAP_HOST. For every host tensor whose - ``buffer.addr`` (a parent VA) lands in a registered range, rewrite it in - place to ``child_base + (addr - parent_lo)`` so the runtime dereferences the - child's own mapping. Tensors outside every range (fork-inherited or - child-allocated) are left untouched. A ``child_memory`` tensor carries a - child-owned device pointer, never a host VA, so it is skipped even when its - address numerically falls inside a registered host range — rewriting it would - corrupt the device pointer. See _BLOB_TENSOR_STRIDE for the wire layout. - """ - tensor_count = struct.unpack_from(" str: """Decode the staged-payload shm name a broadcast_control_all left at _OFF_ARGS.""" raw = bytes(buf[_OFF_ARGS : _OFF_ARGS + _CTRL_SHM_NAME_BYTES]) @@ -537,78 +461,6 @@ def _read_ctrl_staged_shm_name(buf: memoryview) -> str: return raw[: nul if nul >= 0 else _CTRL_SHM_NAME_BYTES].decode("utf-8", "replace") -def _shm_base_addr(shm: SharedMemory) -> int: - """Mapped base address of ``shm``. The mapping outlives the temporary buffer - view, so the address stays 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 - - -def _rebuild_host_buf_ranges( - host_buf_table: dict[int, tuple[SharedMemory, int, int, int]], host_buf_ranges: list[tuple[int, int, int]] -) -> None: - host_buf_ranges.clear() - for _shm, lo, hi, base in host_buf_table.values(): - host_buf_ranges.append((lo, hi, base)) - - -def _handle_ctrl_map_host( - buf: memoryview, - host_buf_table: dict[int, tuple[SharedMemory, int, int, int]], - host_buf_ranges: list[tuple[int, int, int]], -) -> None: - """Child handler for _CTRL_MAP_HOST: persist a host-buffer mapping. - - The staged payload is ``token, parent_va, nbytes`` followed by the host - buffer's shm name. Map that shm and remember the parent VA range it stands - in for so the per-task blob rewrite can redirect host pointers to this base. - """ - payload_size = struct.unpack_from("Q", buf, _CTRL_OFF_ARG0)[0] - staged = SharedMemory(name=_read_ctrl_staged_shm_name(buf)) - try: - staged_buf = staged.buf - assert staged_buf is not None - payload = bytes(staged_buf[:payload_size]) - finally: - staged.close() - token, parent_va, nbytes = _HOST_BUF_MAP_HEADER.unpack_from(payload, 0) - host_shm_name = payload[_HOST_BUF_MAP_HEADER.size :].decode("utf-8") - prior = host_buf_table.pop(token, None) - if prior is not None: - prior[0].close() - # Rebuild ranges in a finally so a raise from SharedMemory / _shm_base_addr - # cannot leave the just-popped prior mapping's stale range in host_buf_ranges. - try: - host_shm = SharedMemory(name=host_shm_name) - host_buf_table[token] = (host_shm, parent_va, parent_va + nbytes, _shm_base_addr(host_shm)) - finally: - _rebuild_host_buf_ranges(host_buf_table, host_buf_ranges) - - -def _handle_ctrl_unmap_host( - buf: memoryview, - host_buf_table: dict[int, tuple[SharedMemory, int, int, int]], - host_buf_ranges: list[tuple[int, int, int]], -) -> None: - """Child handler for _CTRL_UNMAP_HOST: drop a host-buffer mapping by token.""" - payload_size = struct.unpack_from("Q", buf, _CTRL_OFF_ARG0)[0] - staged = SharedMemory(name=_read_ctrl_staged_shm_name(buf)) - try: - staged_buf = staged.buf - assert staged_buf is not None - token = _HOST_BUF_UNMAP.unpack_from(bytes(staged_buf[:payload_size]), 0)[0] - finally: - staged.close() - entry = host_buf_table.pop(token, None) - if entry is not None: - entry[0].close() - _rebuild_host_buf_ranges(host_buf_table, host_buf_ranges) - - def _allocate_local_slot(registry: dict[int, Any]) -> int: for i in range(MAX_REGISTERED_CALLABLE_IDS): if i not in registry: @@ -994,24 +846,21 @@ def _format_exc(prefix: str, exc: BaseException) -> str: return f"{prefix}: {type(exc).__name__}: {exc}" -def _read_args_from_mailbox(buf) -> TaskArgs: - """Decode the TaskArgs blob written by C++ write_blob from the mailbox. - - Used by the Python-targeted child loops (sub_worker, nested L4+ child) - where the destination of `args` is a Python callable that needs a - typed TaskArgs object. The chip-child loops that immediately forward - to C++ run use the zero-copy `run_from_blob` path - instead — see those loops for the matching comment. +def _reexport_args_from_mailbox(buf, worker: Worker) -> list[BufferRef]: + """Re-export the mailbox BufferRef args for an orchestrator (nested L4→L3) child. - Delegates to the nanobind helper so the Tensor 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 - silently broke any tensor carrying a chip-owned device pointer - (HCCL window slots etc.) — now structurally impossible. + Each received ref's backing is re-exported (per-backing, no map, canonical identity preserved), and + a new ref carrying the original view (byte_offset / shapes / strides / dtype) is built over it. The + inner orch fn forwards these to L2 with no map cost (no BufferRef pass-through); dependency + inference keys on the invariant identity. The compute leaf downstream maps lazily. """ - mailbox_addr = ctypes.addressof(ctypes.c_char.from_buffer(buf)) - return read_args_from_blob(mailbox_addr + _OFF_TASK_ARGS_BLOB) + args_ptr = _buffer_field_addr(buf, _OFF_TASK_ARGS_BLOB) + out: list[BufferRef] = [] + for ref_bytes in bufferref_blob_refs(args_ptr, _MAILBOX_ARGS_CAPACITY): + ref = BufferRef.unpack(ref_bytes) + h_prime = worker._reexport(ref.handle) + out.append(h_prime.ref(shapes=ref.shapes, dtype=ref.dtype, strides=ref.strides, byte_offset=ref.byte_offset)) + return out # Idle mailbox polls between `getppid()` samples in a forked child. One poll @@ -1097,8 +946,7 @@ def _sub_worker_loop( rethrows it as ``std::runtime_error``. """ state_addr = _buffer_field_addr(buf, _OFF_STATE) - host_buf_table: dict[int, tuple[SharedMemory, int, int, int]] = {} - host_buf_ranges: list[tuple[int, int, int]] = [] + import_registry = ImportRegistry() # lazy per-endpoint import cache: canonical identity -> local base def handle_task() -> tuple[int, str]: digest = _read_task_digest(buf) @@ -1107,9 +955,10 @@ def handle_task() -> tuple[int, str]: if fn is None: return 1, f"sub_worker: callable hash {_format_digest(digest)} not registered" try: - if host_buf_ranges: - _rewrite_blob_host_addrs(buf, _OFF_TASK_ARGS_BLOB, host_buf_ranges) - args = _read_args_from_mailbox(buf) + # Compute leaf: materialize each arg (map-once) into a MappedArg the Python + # callable computes on via torch.frombuffer(arg.buffer, ...). + args_ptr = _buffer_field_addr(buf, _OFF_TASK_ARGS_BLOB) + args = import_registry.mapped_args_from_blob(args_ptr, _MAILBOX_ARGS_CAPACITY) fn(args) except Exception as e: # noqa: BLE001 return 1, _format_exc("sub_worker", e) @@ -1117,19 +966,14 @@ def handle_task() -> tuple[int, str]: def handle_control(sub_cmd: int) -> tuple[int, str]: try: - if sub_cmd == _CTRL_MAP_HOST: - _handle_ctrl_map_host(buf, host_buf_table, host_buf_ranges) - elif sub_cmd == _CTRL_UNMAP_HOST: - _handle_ctrl_unmap_host(buf, host_buf_table, host_buf_ranges) - else: - _handle_py_callable_control( - buf, - registry, - identity_table, - identity_refs, - sub_cmd, - context="sub_worker", - ) + _handle_py_callable_control( + buf, + registry, + identity_table, + identity_refs, + sub_cmd, + context="sub_worker", + ) except Exception as e: # noqa: BLE001 return 1, _format_exc("sub_worker control", e) return 0, "" @@ -1137,11 +981,7 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: try: _run_mailbox_loop(buf, state_addr, handle_task=handle_task, handle_control=handle_control) finally: - for host_shm, _lo, _hi, _base in host_buf_table.values(): - try: - host_shm.close() - except Exception: # noqa: BLE001 - pass + import_registry.close() def _read_shm_name(buf, offset: int) -> str: @@ -1497,12 +1337,7 @@ def _run_chip_main_loop( # noqa: PLR0913, PLR0915 -- fork-child entry: every de """ prepared = prepared if prepared is not None else set() l3_l2_region_store = _L2HostL3L2RegionStore() - # Post-fork host buffers mapped into this child. `host_buf_table` - # owns the mmap per token (for unmap + teardown); `host_buf_ranges` is the - # parent-VA → child-VA translation table the per-task blob rewrite consults, - # rebuilt from the table on every map/unmap. - host_buf_table: dict[int, tuple[SharedMemory, int, int, int]] = {} # token -> (shm, lo, hi, child_base) - host_buf_ranges: list[tuple[int, int, int]] = [] # (parent_lo, parent_hi, child_base) + import_registry = ImportRegistry() # lazy per-endpoint import cache: canonical identity -> local base def handle_task() -> tuple[int, str]: digest = _read_task_digest(buf) @@ -1524,17 +1359,15 @@ def handle_task() -> tuple[int, str]: f"chip_process dev={device_id}: cid {cid} not prepared before TASK_READY " f"(register via _CTRL_PREPARE first)" ) - # Redirect any registered host pointer (a parent VA) in the - # blob to this child's own mapping before the runtime reads it. - # No-op when nothing is registered. - if host_buf_ranges: - _rewrite_blob_host_addrs(buf, _OFF_TASK_ARGS_BLOB, host_buf_ranges) - # Hand the mailbox bytes straight to C++ (zero-copy zero-decode): - # the blob layout is what `write_blob` already wrote, so re-parsing - # it in Python is N×40B of avoidable work and a permanent - # opportunity to drop a field. C++ reinterpret_cast - # is the source of truth. - cw._impl.run_from_blob(cid, mailbox_addr + _OFF_TASK_ARGS_BLOB, _MAILBOX_ARGS_CAPACITY, cfg) + # Materialize the BufferRef args into a Tensor blob the runtime reads: resolve + # each ref's embedded handle to a local base (map-once, cached by canonical + # identity), then build the Tensor blob at those bases. Replaces the former + # parent-VA range rewrite — identities resolve exactly, not by numeric range. + args_ptr = mailbox_addr + _OFF_TASK_ARGS_BLOB + resolved = import_registry.materialize_blob(args_ptr, _MAILBOX_ARGS_CAPACITY) + tensor_blob = materialize_bufferref_blob(args_ptr, _MAILBOX_ARGS_CAPACITY, resolved) + scratch = ctypes.create_string_buffer(tensor_blob, len(tensor_blob)) + cw._impl.run_from_blob(cid, ctypes.addressof(scratch), len(tensor_blob), cfg) except Exception as e: # noqa: BLE001 code = 1 msg = _format_exc(f"chip_process dev={device_id}", e) @@ -1635,10 +1468,6 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: # noqa: PLR0912 -- one bra _handle_ctrl_release_domain(cw, buf) elif sub_cmd == _CTRL_COMM_INIT: _handle_ctrl_comm_init(cw, buf) - elif sub_cmd == _CTRL_MAP_HOST: - _handle_ctrl_map_host(buf, host_buf_table, host_buf_ranges) - elif sub_cmd == _CTRL_UNMAP_HOST: - _handle_ctrl_unmap_host(buf, host_buf_table, host_buf_ranges) elif sub_cmd == _CTRL_L3_L2_REGION_CREATE: _handle_ctrl_l3_l2_region_create(cw, buf, chip_platform, l3_l2_region_store) elif sub_cmd == _CTRL_L3_L2_REGION_RELEASE: @@ -1657,12 +1486,8 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: # noqa: PLR0912 -- one bra try: _run_mailbox_loop(buf, state_addr, handle_task=handle_task, handle_control=handle_control) finally: + import_registry.close() _sweep_l2_host_l3_l2_regions(l3_l2_region_store) - for host_shm, _lo, _hi, _base in host_buf_table.values(): - try: - host_shm.close() - except Exception: # noqa: BLE001 - pass def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, identity tables, log config, prewarm sizing) must cross the fork as explicit COW args; the child cannot read parent state after os.fork @@ -1810,7 +1635,10 @@ def handle_task() -> tuple[int, str]: if orch_fn is None: return 1, f"child_worker: callable hash {_format_digest(digest)} not registered" try: - args = _read_args_from_mailbox(buf) + # Orchestrator (not a compute leaf): re-export each received backing to a local + # handle H' (per-backing, no map) so the inner orch sees only its own handles; + # pure forwarding to L2 carries no map cost. + args = _reexport_args_from_mailbox(buf, inner_worker) cfg = _read_config_from_mailbox(buf) inner_worker.run(orch_fn, args, cfg) except Exception as e: # noqa: BLE001 @@ -2346,28 +2174,32 @@ def __init__( self._child_alloc_prov: dict[tuple[int, int], _ChildProvEntry] = {} self._child_prov_lock = threading.Lock() - # Post-fork zero-copy host buffers (``create_host_buffer``). Keyed by the - # born-shared shm's mapped base (== the buffer's data_ptr); each entry maps - # a named shm into every local L3 child so memory created after the children - # were forked is still reachable by a later run — with no per-run copy. - self._host_buf_registry: dict[int, _HostBufEntry] = {} - # Immutable read snapshot for the lock-free per-submit lookup - # (``_find_host_buf_entry``): a ``(sorted_ptrs_tuple, registry_copy)`` pair - # rebuilt under ``_registry_lock`` on every create/free and rebound - # atomically. The reader loads it once, so the sorted keys and the dict it - # bisects into never mutate mid-lookup — no lock, no torn read, no - # IndexError from a concurrent free shrinking the list. Host buffers are - # distinct, non-overlapping allocations, so the unique candidate for an - # address is the entry with the greatest base <= addr. - self._host_buf_snapshot: tuple[tuple[int, ...], dict[int, _HostBufEntry]] = ((), {}) - self._host_buf_token_counter: int = 0 + # Owner-side BufferHandle state (P1-B): a per-incarnation opaque nonce, a monotonic buffer_id + # (0 reserved), and the live handles this Worker owns. create_buffer allocates a handle whose + # self-describing descriptor rides embedded in every BufferRef built over it (no export + # handshake); consumers materialize it lazily on receipt. + self._owner_instance_id: bytes = mint_owner_instance_id() + self._buffer_id_counter: int = 1 + self._buffer_handles: dict[int, BufferHandle] = {} + # Re-export table (points 1-4): an upper-level ref received by this worker's orch is re-exported + # to a local handle H' under this worker's identity, per-backing (keyed by source identity), + # so each level's orch sees only its own handles. No map here — H' relabels the backing; + # a compute leaf maps lazily. Lifetime is worker-scoped for now. + self._reexport_by_source: dict[bytes, BufferHandle] = {} + # make_ref_arg memo: a pre-fork host tensor's storage base -> its FORK_SHM handle, so every ref + # over the same storage shares one canonical identity (dependencies key on it). Worker-scoped. + self._fork_tensor_handles: dict[int, BufferHandle] = {} + # L2 leaf only: the in-process consumer import cache. An L2 Worker materializes its own BufferRef + # args itself (no forked child, no mailbox), resolving each ref's descriptor to a local base + # map-once — the chip-child path minus the mailbox hop. Lazily created on first L2 run. + self._l2_import_registry: ImportRegistry | None = None @property def _initialized(self) -> bool: """True only in READY — the worker's tree is live and dispatchable. False once CLOSED (the moment close() claims the epoch), so a dispatch / - register / create_host_buffer that races an in-progress close() is + register / create_buffer that races an in-progress close() is rejected rather than entering the teardown window. """ return self._lifecycle is _Lifecycle.READY @@ -4103,7 +3935,7 @@ def init(self, prewarm_config=None, *, _startup_deadline: float | None = None) - raises after a bounded rollback that reaps the children it forked best-effort (a child wedged in native code past the deadline may be left behind — see the deferred un-reaped-child / nested-shm items). - ``run`` / ``create_host_buffer`` / the remote register/memory APIs never + ``run`` / ``create_buffer`` / the remote register/memory APIs never trigger startup. Args: @@ -4849,15 +4681,13 @@ def _poison_l3_l2_region_from_endpoint_error( poisoned = True 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 tensor.child_memory: - raise ValueError("L3-L2 payload buffer must be host storage, not child_memory device storage") - if not tensor.is_contiguous: - raise ValueError("L3-L2 payload buffer must be contiguous") - base = int(tensor.data) - nbytes = int(tensor.nbytes()) + def _register_l3_l2_orch_comm_host_buffer(self, handle) -> None: + if not isinstance(handle, BufferHandle): + raise TypeError("L3-L2 host buffer registration expects a BufferHandle") + if handle.address_space != AddressSpace.HOST: + raise ValueError("L3-L2 payload buffer must be host storage, not device storage") + base = int(handle.base) + nbytes = int(handle.nbytes) if base <= 0 or nbytes <= 0: return resources = self._building_run_resources @@ -4867,25 +4697,23 @@ def _register_l3_l2_orch_comm_host_buffer(self, tensor) -> None: nbytes, ) - 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 tensor.child_memory: - raise ValueError("L3-L2 payload buffer must be host storage, not child_memory device storage") - if not tensor.is_contiguous: - raise ValueError("L3-L2 payload buffer must be contiguous") - base = int(tensor.data) - nbytes = int(tensor.nbytes()) + def _validate_l3_l2_orch_comm_host_buffer(self, handle) -> None: + if not isinstance(handle, BufferHandle): + raise ValueError("L3-L2 payload buffer must be a BufferHandle returned by orch.alloc(...)") + if handle.address_space != AddressSpace.HOST: + raise ValueError("L3-L2 payload buffer must be host storage, not device storage") + base = int(handle.base) + nbytes = int(handle.nbytes) if base <= 0 or nbytes <= 0: raise ValueError("L3-L2 payload buffer must have a nonzero address and size") resources = self._building_run_resources 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 buffer is not registered; use a handle 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 buffer size {nbytes} exceeds registered shared storage {registered_nbytes}" ) def _create_l3_l2_region(self, worker_id: int, payload_bytes: int, counter_bytes: int): # noqa: PLR0912 @@ -5189,7 +5017,17 @@ def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm device_ctx=int(device_ctx), local_window_base=int(local_window_base), actual_window_size=int(window_size), - buffer_ptrs={b.name: ptrs[i] for i, b in enumerate(buffers)}, + buffers={ + b.name: wrap_vmm_window( + ptrs[i], + int(b.nbytes), + self._owner_instance_id, + self._next_buffer_id(), + f"L{self.level}", + owner_worker_id=int(chip_idx), + ) + for i, b in enumerate(buffers) + }, ) finally: # Close + unlink local copies regardless of outcome. Children @@ -5225,8 +5063,8 @@ def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm with self._child_prov_lock: for chip_idx, ctx in contexts.items(): self._child_prov_record_domain(chip_idx, int(ctx.local_window_base), allocation_id) - for buf_ptr in ctx.buffer_ptrs.values(): - self._child_prov_record_domain(chip_idx, int(buf_ptr), allocation_id) + for buf in ctx.buffers.values(): + self._child_prov_record_domain(chip_idx, int(buf.base), allocation_id) return handle def _release_domain_handle(self, handle: CommDomainHandle, resources: _RunResources) -> None: @@ -5235,7 +5073,7 @@ def _release_domain_handle(self, handle: CommDomainHandle, resources: _RunResour Called by ``CommDomainHandle.release()``. We do NOT drive ``CTRL_RELEASE_DOMAIN`` here because the orch function is allowed to have already submitted DAG tasks that capture the handle's - ``device_ctx`` / ``buffer_ptrs``. Those tasks must see live + ``device_ctx`` / ``buffers``. Those tasks must see live memory through execution; the queue is drained by ``_execute_pending_domain_releases`` once the owning run's fence fires. @@ -5559,12 +5397,17 @@ def _child_prov_drop_domain(self, allocation_id: int) -> None: @staticmethod def _child_ptrs_in_args(args: Any) -> list[tuple[int, int]]: - """Extract ``(device_ptr, arg_index)`` for every child_memory tensor in ``args``.""" + """``(device_ptr, arg_index)`` for every device arg — used for kind4 device-pointer provenance. + + A DEVICE_MALLOC (worker device malloc) or VMM_WINDOW (domain-carved) ref carries the device + pointer in its backend body (u64 LE); that pointer is the provenance key the guard validates + against ``_child_alloc_prov``. Host-backed refs (POSIX/fork shm) contribute nothing. + """ out: list[tuple[int, int]] = [] for i in range(args.tensor_count()): - tensor = args.tensor(i) - if tensor.child_memory: - out.append((int(tensor.data), i)) + desc = BufferRef.unpack(args.ref(i)).handle + if desc.backend_kind in (BackendKind.DEVICE_MALLOC, BackendKind.VMM_WINDOW): + out.append((int.from_bytes(desc.body[:8], "little"), i)) return out def _child_prov_check_dispatch(self, child_ptrs: list[tuple[int, int]], target_worker_id: int, *, api: str) -> None: @@ -5613,363 +5456,243 @@ def _check_chip_worker_id(self, worker_id: int) -> None: if worker_id < 0 or worker_id >= len(self._chip_shms): raise IndexError(f"worker_id {worker_id} out of range (have {len(self._chip_shms)} chips)") - def malloc(self, size: int, worker_id: int = 0) -> int: - """Allocate memory on next-level chip worker *worker_id*. Returns a pointer.""" + def malloc(self, size: int) -> BufferHandle: + """Allocate device memory on this L2 worker's own chip; returns a DEVICE_MALLOC ``BufferHandle``. + + Name a task arg with ``handle.ref(shapes, dtype)`` and release with ``worker.free(handle)``. L3+ + allocates child device memory with ``alloc_child_tensor(worker_id, ...)`` instead — a Worker is + the only allocator, the Orchestrator never allocates. + """ + if self.level != 2: + raise TypeError("worker.malloc is L2-only; at L3+ use worker.alloc_child_tensor(worker_id, ...)") with self._operation_lease("malloc"): - if self.level == 2: - assert self._chip_worker is not None - # L2 is a single chip; worker_id is meaningless there, so the - # provenance is keyed on the canonical worker 0. - with self._child_prov_lock: - ptr = self._chip_worker.malloc(size) - self._child_prov_record_malloc(0, int(ptr)) - return ptr - self._check_chip_worker_id(worker_id) - assert self._orch is not None - return self._orch.malloc(worker_id, size) - - def free(self, ptr: int, worker_id: int = 0) -> None: - """Free memory allocated by ``malloc()``.""" + assert self._chip_worker is not None + with self._child_prov_lock: + ptr = int(self._chip_worker.malloc(int(size))) + self._child_prov_record_malloc(0, ptr) + return wrap_device_malloc( + ptr, int(size), self._owner_instance_id, self._next_buffer_id(), f"L{self.level}", owner_worker_id=0 + ) + + def alloc_child_tensor(self, worker_id: int, shapes: tuple[int, ...], dtype) -> BufferHandle: + """Allocate device memory on next-level ``worker_id`` sized for ``shapes`` × ``dtype``; returns a + DEVICE_MALLOC ``BufferHandle`` (successor of ``orch.malloc`` + ``child_memory``). + + Called from within an orchestration fn (capture the Worker in the closure). The pointer is + private to ``worker_id``; name the arg with ``handle.ref(shapes, dtype)``, dispatch it only to + that worker, and load host data with ``copy_to``. Not auto-freed at end-of-task. + """ + nbytes = get_element_size(dtype) + for s in shapes: + nbytes *= int(s) + self._check_chip_worker_id(int(worker_id)) + assert self._worker is not None + # The lease is re-entrant, so calling this inside the orch fn (the run already holds it) nests + # safely, and calling it outside a run acquires it fresh. + with self._operation_lease("alloc_child_tensor"), self._child_prov_lock: + ptr = int(self._worker.malloc(int(worker_id), int(nbytes))) + self._child_prov_record_malloc(int(worker_id), ptr) + return wrap_device_malloc( + ptr, + int(nbytes), + self._owner_instance_id, + self._next_buffer_id(), + f"L{self.level}", + owner_worker_id=int(worker_id), + ) + + def free(self, handle: BufferHandle) -> None: + """Free a device ``BufferHandle`` allocated by ``malloc`` / ``alloc_child_tensor``. + + The operation lease is re-entrant, so an in-run ``orch.free`` that delegates here nests safely. + """ + wid, ptr = int(handle.owner_worker_id), int(handle.base) with self._operation_lease("free"): - if self.level == 2: - assert self._chip_worker is not None - # Safety-first commit barrier (mirrors Orchestrator.free): revoke - # provenance BEFORE the native free so an async unwind after a - # successful free can never leave a freed address live. - with self._child_prov_lock: - self._child_prov_require_malloc_base(0, int(ptr), api="free") - self._child_prov_clear_malloc(0, int(ptr)) + # Reject a non-chip target (L4+, or a bad id) before touching provenance: a device op is only + # meaningful on a next-level chip. + if self.level != 2: + self._check_chip_worker_id(wid) + with self._child_prov_lock: + # Safety-first commit barrier: revoke provenance BEFORE the native free so an async unwind + # after a successful free can never leave a freed address live. + self._child_prov_require_malloc_base(wid, ptr, api="free") + self._child_prov_clear_malloc(wid, ptr) + if self.level == 2: + assert self._chip_worker is not None self._chip_worker.free(ptr) - return - self._check_chip_worker_id(worker_id) - assert self._orch is not None - self._orch.free(worker_id, ptr) + else: + assert self._worker is not None + self._worker.free(wid, ptr) - def copy_to(self, dst: int, src: int, size: int, worker_id: int = 0) -> None: - """Copy *size* bytes from host *src* to chip worker *dst*.""" + def copy_to(self, dst: BufferHandle, src) -> None: + """H2D: copy host ``src`` (a torch tensor or writable buffer) into device handle ``dst``.""" + src_addr, nbytes = host_ptr_nbytes(src) + wid, dptr = int(dst.owner_worker_id), int(dst.base) with self._operation_lease("copy_to"): - if self.level == 2: - assert self._chip_worker is not None - with self._child_prov_lock: - self._child_prov_require_live(0, int(dst), api="copy_to") - self._chip_worker.copy_to(dst, src, size) - return - self._check_chip_worker_id(worker_id) - assert self._orch is not None - self._orch.copy_to(worker_id, dst, src, size) + if self.level != 2: + self._check_chip_worker_id(wid) + with self._child_prov_lock: + self._child_prov_require_live(wid, dptr, api="copy_to") + if self.level == 2: + assert self._chip_worker is not None + self._chip_worker.copy_to(dptr, src_addr, nbytes) + else: + assert self._worker is not None + self._worker.copy_to(wid, dptr, src_addr, nbytes) - def copy_from(self, dst: int, src: int, size: int, worker_id: int = 0) -> None: - """Copy *size* bytes from chip worker *src* to host *dst*.""" + def copy_from(self, dst, src: BufferHandle) -> None: + """D2H: copy device handle ``src`` into host ``dst`` (a torch tensor or writable buffer).""" + dst_addr, nbytes = host_ptr_nbytes(dst) + wid, sptr = int(src.owner_worker_id), int(src.base) with self._operation_lease("copy_from"): - if self.level == 2: - assert self._chip_worker is not None - with self._child_prov_lock: - self._child_prov_require_live(0, int(src), api="copy_from") - self._chip_worker.copy_from(dst, src, size) - return - self._check_chip_worker_id(worker_id) - assert self._orch is not None - self._orch.copy_from(worker_id, dst, src, size) + if self.level != 2: + self._check_chip_worker_id(wid) + with self._child_prov_lock: + self._child_prov_require_live(wid, sptr, api="copy_from") + if self.level == 2: + assert self._chip_worker is not None + self._chip_worker.copy_from(dst_addr, sptr, nbytes) + else: + assert self._worker is not None + self._worker.copy_from(wid, dst_addr, sptr, nbytes) # ------------------------------------------------------------------ # Post-fork zero-copy host buffers # ------------------------------------------------------------------ - def create_host_buffer(self, nbytes: int) -> HostBuffer: - """Allocate a born-shared host buffer, attached into every local L3 child, - that a later ``run()`` reads/writes with **no per-run copy**. - - Local L3 children are forked during ``init()``; host memory allocated - afterwards is not in their address space. This hands you memory that is - *born* in a shm already attached into every child, so there is nothing to - copy: the child reads and writes the same physical pages the parent sees. - - Returns a :class:`HostBuffer` whose ``buffer`` is a ``memoryview`` over - that shm. Build a tensor over it with the buffer protocol, framework of - your choice, and pass it to ``run()`` as usual:: - - buf = worker.create_host_buffer(n * 4) - t = torch.frombuffer(buf.buffer, dtype=torch.float32, count=n) - t.uniform_(0, 1) # in place → lands in the shm - worker.run(orch(chip, t, out), args=None, config=CallConfig()) - worker.free_host_buffer(buf) # drop the tensor first - - simpler stays framework-free: torch/numpy appear only on the user's side - (``frombuffer``). Blocks until every local L3 child has attached the buffer; - not thread-safe against a concurrent ``run`` / ``create`` / ``free`` on - the same Worker — drive them from one thread, as the L3 worker is - otherwise. + def create_buffer(self, nbytes: int) -> BufferHandle: + """Allocate a shared ``BufferHandle`` owned by this Worker (P1-B). + + The backing is a POSIX shm; the handle carries a typed canonical identity and a self-describing + descriptor. No eager export handshake: the descriptor travels **embedded in every ``BufferRef``** + built over this handle, and a consumer materializes it lazily on first receipt (map-once, keyed + by canonical identity). At L3+ that consumer is a forked child; at L2 (a leaf, no children) the + Worker itself materializes the ref in-process on ``run``. 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 < 3: - raise TypeError("create_host_buffer requires a level >= 3 Worker") - with self._operation_lease("create_host_buffer"): - return self._create_host_buffer_locked(int(nbytes)) - - def _create_host_buffer_locked(self, nbytes: int) -> HostBuffer: - # A born-shared buffer is mapped into every direct process child (chip - # and sub alike, via _broadcast_host_control). Only a truly childless L3 - # has nowhere to attach it. - if not self._chip_shms and not self._sub_shms: - raise RuntimeError( - "create_host_buffer requires at least one forked chip or sub child (this Worker has none)" + 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 alloc_shared_tensor(self, shapes: tuple[int, ...], dtype) -> BufferHandle: + """Allocate a runtime-managed intermediate buffer (the BufferRef form of ``orch.alloc``). + + Called inside an orchestration fn. The backing comes from the orchestrator's HeapRing + (MAP_SHARED, visible to forked children) and is **auto-reclaimed** once every downstream consumer + has completed and the scope ends — no manual free. Returns a ``FORK_SHM`` ``BufferHandle`` whose + canonical identity is registered in the tensormap so a ref over it (``handle.ref(shapes, dtype)``) + dependency-wires to this producer slot. Chip-A→chip-B intermediates: name it as an OUTPUT of the + producing task and an INPUT of the consumer. + """ + assert self._orch is not None, "alloc_shared_tensor requires an L3+ orchestration context" + nbytes = get_element_size(dtype) + for s in shapes: + nbytes *= int(s) + oid, buffer_id, path = self._owner_instance_id, self._next_buffer_id(), f"L{self.level}" + identity = CanonicalIdentity(oid, buffer_id, path, 0) + va = int(self._orch._o.alloc(list(int(s) for s in shapes), dtype, identity.pack())) + # Wrap the ring VA under the SAME identity: the child materializes to that VA (fork-inherited, + # MAP_SHARED read-write) and infer_deps keys the ref to the slot registered above. + return wrap_fork_inherited(va, int(nbytes), oid, buffer_id, path, generation=0, access=AccessMode.READWRITE) + + def make_ref_arg(self, tensor, shapes: tuple[int, ...], dtype: int, *, strides: tuple[int, ...] | None = None): + """Name a **pre-fork** host tensor as a ``BufferRef`` over a memoized ``FORK_SHM`` handle. + + The torch (or buffer-protocol) tensor MUST be allocated before ``init()`` so its VA is + fork-inherited by the children (the mainline "fork-inherited" contract). A ``share_memory_()`` + tensor is MAP_SHARED — read-write across the fork, so usable as an OUTPUT the parent reads back; + a plain tensor is COW read-only (input only). The handle is memoized by the tensor's storage + base, so every ref over the same storage shares one canonical identity and dependencies key on + it. At L2 (no fork) any host tensor works. ``dtype`` is the ``DataType`` int value. + """ + untyped_storage = getattr(tensor, "untyped_storage", None) + if callable(untyped_storage): + st = untyped_storage() + base, nbytes = int(st.data_ptr()), int(st.nbytes()) + byte_offset = int(tensor.data_ptr()) - base # the view's start within its storage + else: + base, nbytes = host_ptr_nbytes(tensor) + byte_offset = 0 + handle = self._fork_tensor_handles.get(base) + if handle is None: + handle = wrap_fork_inherited( + base, + nbytes, + self._owner_instance_id, + self._next_buffer_id(), + f"L{self.level}", + access=AccessMode.READWRITE, ) - assert self._worker is not None - - if nbytes <= 0: - raise ValueError("create_host_buffer: nbytes must be positive") - - # Create the shm up front, then guard everything after it — mapping the - # address (``_shm_base_addr``), reserving the registry slot, and the - # broadcast — under one ``try`` so any failure closes and unlinks the shm - # instead of leaking a /dev/shm segment. The registry mutation stays under - # ``_registry_lock`` (mirrors Worker.register's discipline); the slow - # broadcast runs *outside* the lock — wire-level concurrency is serialized - # at the C++ mailbox, not here. The born-shared shm's own mapped base is - # the buffer's data_ptr, so a tensor built over buffer.buffer resolves to - # this registered range. - shm = SharedMemory(create=True, size=nbytes) - token: int | None = None - data_ptr: int | None = None - try: - data_ptr = _shm_base_addr(shm) - with self._registry_lock: - token = self._host_buf_token_counter - self._host_buf_token_counter += 1 - entry = _HostBufEntry( - token=token, - data_ptr=data_ptr, - nbytes=nbytes, - shm=shm, - shm_name=shm.name, - shm_base=data_ptr, - ) - self._host_buf_registry[data_ptr] = entry - self._rebuild_host_buf_snapshot() + self._fork_tensor_handles[base] = handle + return handle.ref(shapes=tuple(shapes), dtype=dtype, strides=strides, byte_offset=byte_offset) - payload = _HOST_BUF_MAP_HEADER.pack(token, data_ptr, nbytes) + shm.name.encode("utf-8") - errors = self._broadcast_host_control(_CTRL_MAP_HOST, payload) - if errors: - raise RuntimeError( - f"create_host_buffer: MAP_HOST failed on {len(errors)} local L3 children; first error: {errors[0]}" - ) - except BaseException: - # Roll back on any failure — a staging error before the map, a partial - # map, or an exception from the broadcast itself (any of which would - # otherwise leak the shm): unmap any child that took it, drop the - # reservation, free the shm. No user view exists yet, so close() cannot - # be blocked by an exported buffer. - try: - if token is not None: - self._broadcast_host_unmap(token) - except Exception as unmap_exc: # noqa: BLE001 -- must not mask the original failure below - sys.stderr.write( - f"[worker pid={os.getpid()}] WARN: create_host_buffer rollback UNMAP_HOST " - f"failed (continuing best-effort): {unmap_exc}\n" - ) - sys.stderr.flush() - finally: - with self._registry_lock: - if data_ptr is not None and self._host_buf_registry.pop(data_ptr, None) is not None: - self._rebuild_host_buf_snapshot() - shm.close() - try: - shm.unlink() - except FileNotFoundError: - pass - raise - - buf_view = shm.buf - assert buf_view is not None - return HostBuffer(token=token, data_ptr=data_ptr, nbytes=nbytes, buffer=buf_view) - - def free_host_buffer(self, handle: HostBuffer) -> None: - """Release a born-shared buffer created by ``create_host_buffer``. + def _next_buffer_id(self) -> int: + with self._registry_lock: + bid = self._buffer_id_counter + self._buffer_id_counter += 1 + return bid - Unmaps it from every local L3 child and frees the parent shm. Drop every - tensor / ``memoryview`` you built over ``handle.buffer`` *first*: a live - view keeps the shm's pages exported, so ``close()`` cannot release them - and the buffer only warns (and is reclaimed once the last view is gone). + def _reexport(self, source: BufferHandleDescriptor) -> BufferHandle: + """Re-export a received backing for forwarding (per-backing, memoized, no map). - Best-effort and idempotent: a stale handle whose token no longer matches - (e.g. freed twice) is a silent no-op. + An upper-level ref reaching this worker's orch is forwarded as a handle H' that keeps the + source's canonical identity unchanged (invariant across every edge, frozen model §5/§8) — H' + is never mapped here (a downstream compute leaf maps it lazily), and is built once per source + backing (keyed by identity). Worker-scoped lifetime for now. """ - if not isinstance(handle, HostBuffer): - raise TypeError("free_host_buffer expects a HostBuffer from create_host_buffer") - with self._operation_lease("free_host_buffer"): - self._free_host_buffer_locked(handle) + key = source.identity.pack() + handle = self._reexport_by_source.get(key) + if handle is None: + handle = re_export(source) + self._reexport_by_source[key] = handle + return handle - def _free_host_buffer_locked(self, handle: HostBuffer) -> None: + def _create_buffer_locked(self, nbytes: int) -> BufferHandle: + # An L3+ buffer is consumed by a forked child that lazily maps it, so a childless L3+ buffer + # can reach no consumer. An L2 leaf has no children and materializes the ref in-process itself, + # 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: - entry = self._host_buf_registry.get(handle.data_ptr) - if entry is None or entry.token != handle.token: - return - self._host_buf_registry.pop(handle.data_ptr, None) - self._rebuild_host_buf_snapshot() - errors: list[str] = [] - try: - # Gate on resource presence, not lifecycle: the child mailboxes are - # driveable whenever the C++ _worker is up — including during close() - # teardown (CLOSED), when the children are still alive to unmap. - if self._worker is not None: - errors = self._broadcast_host_unmap(entry.token) - except Exception as exc: # noqa: BLE001 - errors = [str(exc)] - finally: - close_warn = self._close_host_shm(entry) - if close_warn: - errors.append(close_warn) - if errors: - sys.stderr.write( - f"[worker pid={os.getpid()}] WARN: free_host_buffer token={entry.token} " - f"failed on {len(errors)} local L3 children; first error: {errors[0]}\n" - ) - sys.stderr.flush() + self._buffer_handles[buffer_id] = handle + return handle - @staticmethod - def _close_host_shm(entry: _HostBufEntry) -> str | None: - """Close + unlink a host-buffer's parent shm. - - Returns a warning string (else ``None``) when a still-live view over a - zero-copy buffer blocks ``close()``: ``memoryview.release()`` raises - ``BufferError`` while a tensor built via ``frombuffer`` still holds the - pages exported. The name is unlinked regardless, so the OS reclaims the - segment once the user drops that last view. - """ - warn: str | None = None - try: - entry.shm.close() - except BufferError: - warn = ( - f"host buffer token={entry.token} still has a live view (a tensor/memoryview over " - f"buffer.buffer); drop it before free_host_buffer/close() to release the shm promptly" - ) - try: - entry.shm.unlink() - except FileNotFoundError: - pass - return warn + def _close_l2_import_registry(self) -> None: + """Close the L2 in-process consumer import cache (drops its mapped shm imports).""" + if self._l2_import_registry is not None: + self._l2_import_registry.close() + self._l2_import_registry = None - def _release_all_host_buffers(self) -> None: - """Unmap + free every still-registered host buffer (called from close()). + def _release_all_buffer_handles(self) -> None: + """Close + unlink every owner BufferHandle (called from close()). - Per-buffer best-effort: every buffer's shm is closed even if its unmap - broadcast (or a prior buffer) fails, so one failure never strands the - rest; the first error is raised after all are attempted so close() - reports the leak rather than swallowing it to stderr.""" + Children drop their own lazily-mapped imports when their loops exit; the owner unlinks the + backing shm here. 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: - entries = list(self._host_buf_registry.values()) - self._host_buf_registry.clear() - self._rebuild_host_buf_snapshot() + handles = list(self._buffer_handles.values()) + self._buffer_handles.clear() errors: list[BaseException] = [] - for entry in entries: + for handle in handles: try: - try: - if self._worker is not None: # resource presence, not lifecycle (see _close_host_shm) - self._broadcast_host_unmap(entry.token) - finally: - # Tolerates a still-live view over a zero-copy buffer at close(): - # unlinks the name regardless so the OS reclaims it once dropped. - self._close_host_shm(entry) + handle.close() except BaseException as exc: # noqa: BLE001 errors.append(exc) if errors: raise errors[0] - def _broadcast_host_unmap(self, token: int) -> list[str]: - """Broadcast _CTRL_UNMAP_HOST for ``token`` to every local L3 child.""" - return self._broadcast_host_control(_CTRL_UNMAP_HOST, _HOST_BUF_UNMAP.pack(token)) - - def _broadcast_host_control(self, sub_cmd: int, payload: bytes) -> list[str]: - if self._worker is None: - return [] - results = [] - for worker_type in (WorkerType.NEXT_LEVEL, WorkerType.SUB): - results.extend( - self._worker.broadcast_control_all( - worker_type, - int(sub_cmd), - payload, - None, - timeout_s=self._py_control_timeout_s, - ) - ) - return self._control_errors(results) - - def _stage_host_buffers_for_chip_submit(self, args: Any) -> None: - """Validate the host tensors of one chip submit before dispatch. - - Called from ``Orchestrator.submit_next_level`` on the LOCAL_CHIP path — - only there does the forked child dereference raw host pointers. Each host - tensor is either: - - * inside a buffer from ``create_host_buffer`` → born-shared, so its bytes - already live in the child-visible shm and the child writes results back - into the same physical pages: nothing to copy. ``_find_host_buf_entry`` - still validates the view fits inside the buffer (else the child would - read past its shm mapping); - * unregistered → forwarded unvalidated. A fork-inherited ``share_memory_`` - tensor is the legitimate case; an unregistered post-fork tensor reads - stale/unmapped memory in the child — allocate it with - ``create_host_buffer`` instead. - - The child rewrites in-range host pointers to its own mapping; see - _rewrite_blob_host_addrs. - """ - for i in range(args.tensor_count()): - tensor = args.tensor(i) - if tensor.child_memory: - continue - addr = int(tensor.data) - if addr == 0: - continue - # Raises if an in-range view overruns its buffer; otherwise there is - # nothing to do — the born-shared bytes are already child-visible. - self._find_host_buf_entry(addr, int(tensor.nbytes())) - - def _rebuild_host_buf_snapshot(self) -> None: - """Rebuild the lock-free read snapshot from the registry. - - Caller must hold ``_registry_lock``. Rebinds ``_host_buf_snapshot`` to a - fresh ``(sorted_ptrs_tuple, registry_copy)`` pair so an in-flight - ``_find_host_buf_entry`` keeps reading the prior immutable snapshot until - the single atomic rebind swaps it — see ``_find_host_buf_entry``. - """ - registry = dict(self._host_buf_registry) - self._host_buf_snapshot = (tuple(sorted(registry)), registry) - - def _find_host_buf_entry(self, addr: int, nbytes: int) -> _HostBufEntry | None: - """Host buffer whose ``[data_ptr, data_ptr+nbytes)`` contains the whole - ``[addr, addr+nbytes)`` view, or None. Raises if a view starts inside a - buffer but runs past its end (would read past the shm in the child). - - Host buffers are distinct, non-overlapping allocations, so the only - candidate for ``addr`` is the entry with the greatest base ``<= addr`` — - 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 - 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. - """ - # Load the immutable snapshot once: sorted keys and the dict they index - # into are captured together, so a concurrent create/free (which rebinds a - # fresh snapshot) cannot mutate what we bisect and index here — no lock, no - # IndexError from a shrinking list, no torn key/dict pairing. - sorted_ptrs, registry = self._host_buf_snapshot - idx = bisect.bisect_right(sorted_ptrs, addr) - 1 - if idx < 0: - return None - entry = registry.get(sorted_ptrs[idx]) - if entry is None or addr >= entry.data_ptr + entry.nbytes: - return None - if addr + nbytes > entry.data_ptr + entry.nbytes: - raise RuntimeError( - f"Host tensor 0x{addr:x} (+{nbytes} B) overruns its host buffer " - f"0x{entry.data_ptr:x} (+{entry.nbytes} B); create a buffer at least as large." - ) - return entry - # ------------------------------------------------------------------ # run — uniform entry point # ------------------------------------------------------------------ @@ -6013,7 +5736,7 @@ def _submit_locked(self, callable, args, config) -> RunHandle: if self.level == 2: assert self._chip_worker is not None state = self._resolve_handle(callable, expected_namespace="LOCAL_CHIP") - self._chip_worker._run_slot(state.slot_id, args, cfg) + self._run_l2_materialized(state.slot_id, args, cfg) return RunHandle._completed(self) with self._submit_mu: @@ -6130,6 +5853,34 @@ def _step(fn) -> None: self._hierarchical_start_cv.notify_all() return result + def _run_l2_materialized(self, callable_id: int, args, cfg) -> None: + """Materialize an L2 leaf's BufferRef args to a Tensor blob in-process and run the kernel. + + The user builds args as ``TaskArgs`` (BufferRef) at every level; an L2 leaf is the consumer of + its own args, so it does exactly what a chip child does — resolve each ref's embedded descriptor + to a local base (map-once, cached in ``_l2_import_registry``) and build the Tensor blob the + runtime reads — only without a mailbox, since the args are already in this process. + + ``args`` is a BufferRef ``TaskArgs`` or ``None`` (no args). The chip-only POD + ``ChipStorageTaskArgs`` is not accepted here — submit it through ``ChipWorker._run_slot``. + """ + if self._l2_import_registry is None: + self._l2_import_registry = ImportRegistry() + if args is None: + refs: list[BufferRef] = [] + scalars: tuple[int, ...] = () + else: + refs = [BufferRef.unpack(args.ref(i)) for i in range(args.tensor_count())] + scalars = tuple(args.scalar(i) for i in range(args.scalar_count())) + blob = pack_bufferref_blob(refs, scalars) + in_scratch = ctypes.create_string_buffer(blob, len(blob)) + in_addr = ctypes.addressof(in_scratch) + resolved = self._l2_import_registry.materialize_blob(in_addr, len(blob)) + tensor_blob = materialize_bufferref_blob(in_addr, len(blob), resolved) + out_scratch = ctypes.create_string_buffer(tensor_blob, len(tensor_blob)) + assert self._chip_worker is not None + self._chip_worker._impl.run_from_blob(callable_id, ctypes.addressof(out_scratch), len(tensor_blob), cfg) + @property def aicpu_dlopen_count(self) -> int: """L2 only: number of distinct callable identities the AICPU has dlopened for. @@ -6186,7 +5937,6 @@ def _has_live_resources(self) -> bool: or bool(self._sub_shms or self._chip_shms or self._next_level_shms) or bool(self._live_l3_l2_regions) or bool(self._live_domains) - or bool(self._host_buf_registry) or bool(self._pending_remote_buffer_frees or self._pending_remote_import_releases) ) @@ -6206,8 +5956,6 @@ def _describe_live_resources(self) -> str: parts.append(f"{len(self._live_l3_l2_regions)} L3-L2 region(s)") if self._live_domains: parts.append(f"{len(self._live_domains)} comm domain(s)") - if self._host_buf_registry: - parts.append(f"{len(self._host_buf_registry)} host buffer(s)") n_remote = len(self._pending_remote_buffer_frees) + len(self._pending_remote_import_releases) if n_remote: parts.append(f"{n_remote} pending remote free(s)") @@ -6249,7 +5997,7 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r with self._hierarchical_start_cv: if threading.get_ident() in self._lease_depth: raise RuntimeError( - "Worker.close(): cannot be called from within a run() / submit() / create_host_buffer() " + "Worker.close(): cannot be called from within a run() / submit() / create_buffer() " "operation on this thread" ) if self._lifecycle is _Lifecycle.INITIALIZING: @@ -6528,9 +6276,9 @@ def _step(fn) -> None: _step(self._clear_child_prov) _step(self._release_active_remote_slot_refs) _step(self._flush_pending_remote_frees) - # Host buffers must be released while the local L3 child mailboxes are - # still usable (before _worker.close()). - _step(self._release_all_host_buffers) + _step(self._release_all_buffer_handles) + _step(self._fork_tensor_handles.clear) + _step(self._close_l2_import_registry) if self.level == 2: diff --git a/simpler_setup/scene_test.py b/simpler_setup/scene_test.py index a51c0b7f0d..cdebcc182c 100644 --- a/simpler_setup/scene_test.py +++ b/simpler_setup/scene_test.py @@ -230,8 +230,9 @@ def __init__(self, worker, test_args: TaskArgsBuilder): self._worker = worker self._torch = torch - self._buffers: list = [] # (HostBuffer, view) in creation order + self._buffers: list = [] # (BufferHandle, view) in creation order self._originals: dict[str, Any] = {} # name -> pre-rehost tensor + self._handles: dict[str, Any] = {} # name -> owning BufferHandle (for _build_l3_task_args refs) self._test_args = test_args self._reject_aliased_tensors(test_args) try: @@ -241,13 +242,17 @@ def __init__(self, worker, test_args: TaskArgsBuilder): # an empty tensor is never dereferenced by the child, so it is # left untouched rather than allocating a zero-length buffer. if isinstance(spec, Tensor) and isinstance(spec.value, torch.Tensor) and spec.value.numel() > 0: - view = self._rehost_one(spec.value) + handle, view = self._rehost_one(spec.value) self._originals[spec.name] = test_args._data[spec.name] + self._handles[spec.name] = handle test_args._data[spec.name] = view new_specs.append(Tensor(spec.name, view)) else: new_specs.append(spec) test_args._specs = new_specs + # Expose the owning handles so the L3 arg builder can name each rehosted tensor as a + # BufferRef (handle.ref); the L2 chip builder keeps using the raw views. + test_args._rehost_handles = self._handles except BaseException: self.release() raise @@ -289,15 +294,17 @@ def _rehost_one(self, t): "representable — build it contiguous in generate_args()" ) nbytes = t.numel() * t.element_size() - buf = self._worker.create_host_buffer(nbytes) + handle = self._worker.create_buffer(nbytes) try: - view = torch.frombuffer(buf.buffer, dtype=t.dtype, count=t.numel()).view(t.shape) + shm = handle.shm + assert shm is not None + view = torch.frombuffer(shm.buf, dtype=t.dtype, count=t.numel()).view(t.shape) view.copy_(t) except BaseException: - self._worker.free_host_buffer(buf) + handle.close() raise - self._buffers.append((buf, view)) - return view + self._buffers.append((handle, view)) + return handle, view def release(self) -> None: # Restore the builder's original entries (dropping the born-shared view @@ -309,17 +316,16 @@ def release(self) -> None: for s in self._test_args._specs ] self._originals.clear() + self._handles.clear() + if hasattr(self._test_args, "_rehost_handles"): + self._test_args._rehost_handles = {} while self._buffers: - buf, view = self._buffers.pop() + handle, view = self._buffers.pop() del view try: - buf.buffer.release() - except (ValueError, BufferError): - pass - try: - self._worker.free_host_buffer(buf) + handle.close() # drops the view above, then closes + unlinks the POSIX shm except Exception as exc: # noqa: BLE001 -- best-effort cleanup; a leak here must not mask the test result, but process-control exceptions still propagate - logger.warning("SceneTest rehost cleanup: free_host_buffer failed: %s", exc) + logger.warning("SceneTest rehost cleanup: handle.close failed: %s", exc) # --------------------------------------------------------------------------- @@ -337,12 +343,13 @@ class CallableNamespace: callables.verify # → CallableHandle Also provides ``keep()`` for lifetime management: L3 orch functions - that build transient Python objects (e.g. ChipStorageTaskArgs) whose + that build transient Python objects (e.g. a ``TaskArgs``) whose raw pointers are submitted to the C++ scheduler must register them via ``keep()`` so they outlive the scheduler drain:: def run_dag(w, callables, task_args, config): - chip_args, _ = _build_chip_task_args(task_args, callables.vector_kernel_sig) + chip_args = TaskArgs() + chip_args.add_ref(_rehosted_ref(task_args, "a"), TensorArgType.INPUT) callables.keep(chip_args) # survive until drain finishes ... """ @@ -370,11 +377,55 @@ def keep(self, *objs): # --------------------------------------------------------------------------- +def _build_l2_ref_args(test_args: TaskArgsBuilder, orch_signature: list, worker): + """Build BufferRef `TaskArgs` from `TaskArgsBuilder` for the L2 `Worker.run` path. + + An L2 leaf consumes its own args: `Worker.run(handle, args, cfg)` materializes each BufferRef to a + local base in-process. Each tensor is named as a ref over ``worker.make_ref_arg`` (a host tensor; + at L2 there is no fork, so any host tensor resolves in-process); the direction tag is inert at L2 + but set for parity with the L3 path. + + Returns: + args: TaskArgs (BufferRef) + output_names: list of tensor names that are OUTPUT or INOUT + """ + from simpler.task_interface import ArgDirection, TaskArgs, TensorArgType, scalar_to_uint64 # noqa: PLC0415 + + from simpler_setup.torch_interop import make_tensor_ref # noqa: PLC0415 + + dir2tag = { + ArgDirection.IN: TensorArgType.INPUT, + ArgDirection.OUT: TensorArgType.OUTPUT_EXISTING, + ArgDirection.INOUT: TensorArgType.INOUT, + } + args = TaskArgs() + output_names: list[str] = [] + tensor_idx = 0 + for spec in test_args.specs: + if isinstance(spec, Tensor): + if tensor_idx >= len(orch_signature): + raise ValueError( + f"Tensor '{spec.name}' at index {tensor_idx} has no matching entry in " + f"orchestration signature (length {len(orch_signature)}). " + f"Update CALLABLE['orchestration']['signature'] to match generate_args()." + ) + direction = orch_signature[tensor_idx] + args.add_ref(make_tensor_ref(worker, spec.value), dir2tag.get(direction, TensorArgType.INPUT)) + if direction in (ArgDirection.OUT, ArgDirection.INOUT): + output_names.append(spec.name) + tensor_idx += 1 + elif isinstance(spec, Scalar): + args.add_scalar(scalar_to_uint64(spec.value)) + + return args, output_names + + def _build_chip_task_args(test_args: TaskArgsBuilder, orch_signature: list): """Build `ChipStorageTaskArgs` (POD) from `TaskArgsBuilder`. - Used by the L2 path (`ChipWorker.run(callable, chip_args, config)`): the - chip worker expects the runtime.so ABI-shaped POD directly (no tags). + Used by the direct chip API (`ChipWorker._run_slot(slot, chip_args, config)`, e.g. the + prepared-callable tests): the chip worker expects the runtime.so ABI-shaped POD directly (no tags). + The `Worker.run` L2 path instead builds BufferRef args via `_build_l2_ref_args`. Returns: chip_args: ChipStorageTaskArgs (POD) @@ -411,12 +462,53 @@ def _build_chip_task_args(test_args: TaskArgsBuilder, orch_signature: list): return chip_args, output_names -def _build_l3_task_args(test_args: TaskArgsBuilder, orch_signature: list): - """Build a tagged `TaskArgs` (vector-backed, with `TensorArgType` tags) from - `TaskArgsBuilder`. +def _rehosted_ref_for(test_args: TaskArgsBuilder, tensor): + """A ``BufferRef`` over the rehosted ``tensor``'s owning handle, matched by its backing address. + + Like ``_rehosted_ref`` but keyed by the (rehosted) tensor object a hand-written orch already holds, + rather than its name — the rehosted view's ``data_ptr`` is its create_buffer handle's base. + """ + from simpler_setup.torch_interop import torch_dtype_to_datatype # noqa: PLC0415 + + base = int(tensor.data_ptr()) + for handle in getattr(test_args, "_rehost_handles", {}).values(): + if handle.base == base: + return handle.ref(shapes=tuple(tensor.shape), dtype=torch_dtype_to_datatype(tensor.dtype).value) + raise ValueError("tensor is not a rehosted arg (no create_buffer handle with a matching base)") + + +def _l3_ref(test_args: TaskArgsBuilder, name: str, worker=None): + """A ``BufferRef`` naming the L3 tensor arg ``name``. + + Two ways a scene test's host tensor becomes child-visible: the framework rehosts it into a + create_buffer (POSIX shm) — use that handle; or a test pre-allocates a ``share_memory_()`` tensor + before ``init()`` (fork-inherited MAP_SHARED) — name it via ``worker.make_ref_arg`` (FORK_SHM). + """ + from simpler_setup.torch_interop import torch_dtype_to_datatype # noqa: PLC0415 + + t = getattr(test_args, name) + dtype = torch_dtype_to_datatype(t.dtype).value + handle = getattr(test_args, "_rehost_handles", {}).get(name) + if handle is not None: + return handle.ref(shapes=tuple(t.shape), dtype=dtype) + if worker is not None: + return worker.make_ref_arg(t, shapes=tuple(t.shape), dtype=dtype) + raise ValueError(f"L3 tensor arg '{name}' has no rehosted handle; rehost the args or pass worker= for make_ref_arg") + - Used by the L3 path (`orch.submit_next_level(callable, args, config, worker=chip_id)`): - the orchestrator reads the tags to drive dependency inference. +def _rehosted_ref(test_args: TaskArgsBuilder, name: str): + """A ``BufferRef`` over the rehosted tensor ``name`` (framework rehost only). For a hand-written + L3 orch naming one rehosted tensor as a task arg (e.g. the chip output as a sub-task INPUT).""" + return _l3_ref(test_args, name, worker=None) + + +def _build_l3_task_args(test_args: TaskArgsBuilder, orch_signature: list, worker=None): + """Build a tagged `TaskArgs` (BufferRefs) from a `TaskArgsBuilder` for the L3 path + (`orch.submit_next_level`); the tags drive dependency inference. + + Names each tensor as a BufferRef over its framework-rehosted create_buffer handle, or — for a test + that pre-allocates ``share_memory_()`` tensors before ``init()`` — via ``worker.make_ref_arg`` when + ``worker`` is passed. Returns: chip_args: TaskArgs (tagged) @@ -429,14 +521,14 @@ def _build_l3_task_args(test_args: TaskArgsBuilder, orch_signature: list): scalar_to_uint64, ) - from simpler_setup.torch_interop import make_tensor_arg # noqa: PLC0415 - _DIR_TO_TAG = { ArgDirection.IN: TensorArgType.INPUT, ArgDirection.OUT: TensorArgType.OUTPUT_EXISTING, ArgDirection.INOUT: TensorArgType.INOUT, } + # Each rehosted tensor is named as a BufferRef over its owning create_buffer handle (the child + # maps it by canonical identity); tags drive dependency inference. chip_args = TaskArgs() output_names: list[str] = [] @@ -451,7 +543,7 @@ def _build_l3_task_args(test_args: TaskArgsBuilder, orch_signature: list): ) direction = orch_signature[tensor_idx] tag = _DIR_TO_TAG.get(direction, TensorArgType.INPUT) - chip_args.add_tensor(make_tensor_arg(spec.value), tag) + chip_args.add_ref(_l3_ref(test_args, spec.name, worker), tag) if direction in (ArgDirection.OUT, ArgDirection.INOUT): output_names.append(spec.name) tensor_idx += 1 @@ -1239,7 +1331,7 @@ def _run_and_validate_l2( # noqa: PLR0913 -- threads CLI diagnostic flags + cas # Build args test_args = self.generate_args(params) - chip_args, output_names = _build_chip_task_args(test_args, orch_sig) + chip_args, output_names = _build_l2_ref_args(test_args, orch_sig, worker) # Compute golden (unless skip_golden) golden_args = None diff --git a/simpler_setup/torch_interop.py b/simpler_setup/torch_interop.py index 23db6aa415..cc57ea1810 100644 --- a/simpler_setup/torch_interop.py +++ b/simpler_setup/torch_interop.py @@ -100,3 +100,23 @@ def make_tensor_arg(tensor) -> Tensor: ) shapes = tuple(int(s) for s in tensor.shape) return Tensor.make(tensor.data_ptr(), shapes, dt) + + +def make_tensor_ref(worker, tensor): + """A ``BufferRef`` over a **pre-fork** host torch tensor — the ref counterpart of ``make_tensor_arg``. + + Names ``tensor`` as a memoized ``FORK_SHM`` handle on ``worker`` (``worker.make_ref_arg``), inferring + shapes + dtype from the tensor. Use for standalone L3 examples whose host inputs/outputs are + ``share_memory_()`` tensors allocated before ``worker.init()`` (fork-inherited). A contiguous CPU + tensor is required (as for ``make_tensor_arg``). + """ + _ensure_torch_map() + dt = _TORCH_DTYPE_MAP.get(tensor.dtype) # pyright: ignore[reportOptionalMemberAccess] + if dt is None: + raise ValueError(f"Unsupported tensor dtype for BufferRef: {tensor.dtype}") + if tensor.device.type != "cpu": + raise ValueError(f"make_tensor_ref requires a CPU tensor, got device={tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError("make_tensor_ref requires a contiguous tensor; call tensor.contiguous() first.") + shapes = tuple(int(s) for s in tensor.shape) + return worker.make_ref_arg(tensor, shapes=shapes, dtype=int(dt.value)) diff --git a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp index f8de55082c..188f80e37b 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -689,7 +689,7 @@ extern "C" int bind_callable_to_runtime_impl( for (int i = 0; i < tensor_count; i++) { Tensor t = orch_args->tensor(i); - if (t.is_child_memory()) { + if (t.is_device_memory()) { LOG_DEBUG(" Tensor %d: child memory, pass-through (0x%" PRIx64 ")", i, t.buffer.addr); device_args.add_tensor(t); continue; diff --git a/src/a2a3/runtime/host_build_graph/runtime/tensor_create_info.h b/src/a2a3/runtime/host_build_graph/runtime/tensor_create_info.h index 912839a340..756f706360 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/tensor_create_info.h +++ b/src/a2a3/runtime/host_build_graph/runtime/tensor_create_info.h @@ -99,7 +99,7 @@ static_assert(offsetof(TensorCreateInfo, ndims) == offsetof(Tensor, ndims)); static_assert(offsetof(TensorCreateInfo, dtype) == offsetof(Tensor, dtype)); static_assert(offsetof(TensorCreateInfo, manual_dep) == offsetof(Tensor, manual_dep)); static_assert(offsetof(TensorCreateInfo, is_contiguous) == offsetof(Tensor, is_contiguous)); -static_assert(offsetof(TensorCreateInfo, __pad_flags__) == offsetof(Tensor, child_memory)); +static_assert(offsetof(TensorCreateInfo, __pad_flags__) == offsetof(Tensor, address_space)); static_assert(offsetof(TensorCreateInfo, shapes) == offsetof(Tensor, shapes)); // ============================================================================ diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index c7ce87c6c4..6bd7e774bc 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -341,7 +341,7 @@ class RetainedTempBump { size_t required = 0; for (int i = 0; i < orch_args->tensor_count(); i++) { Tensor t = orch_args->tensor(i); - if (t.is_child_memory() || t.nbytes() == 0) { + if (t.is_device_memory() || t.nbytes() == 0) { continue; } required += align_up(static_cast(t.nbytes())); @@ -557,7 +557,7 @@ static bool stage_device_args( for (int i = 0; i < tensor_count; i++) { Tensor t = orch_args->tensor(i); - if (t.is_child_memory()) { + if (t.is_device_memory()) { LOG_DEBUG(" Tensor %d: child memory, pass-through (0x%" PRIx64 ")", i, t.buffer.addr); out->add_tensor(t); continue; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/tensor_create_info.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/tensor_create_info.h index 912839a340..756f706360 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/tensor_create_info.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/tensor_create_info.h @@ -99,7 +99,7 @@ static_assert(offsetof(TensorCreateInfo, ndims) == offsetof(Tensor, ndims)); static_assert(offsetof(TensorCreateInfo, dtype) == offsetof(Tensor, dtype)); static_assert(offsetof(TensorCreateInfo, manual_dep) == offsetof(Tensor, manual_dep)); static_assert(offsetof(TensorCreateInfo, is_contiguous) == offsetof(Tensor, is_contiguous)); -static_assert(offsetof(TensorCreateInfo, __pad_flags__) == offsetof(Tensor, child_memory)); +static_assert(offsetof(TensorCreateInfo, __pad_flags__) == offsetof(Tensor, address_space)); static_assert(offsetof(TensorCreateInfo, shapes) == offsetof(Tensor, shapes)); // ============================================================================ diff --git a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index 187a66d91f..9af029b8a3 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -318,7 +318,7 @@ class RetainedTempBump { size_t required = 0; for (int i = 0; i < orch_args->tensor_count(); i++) { Tensor t = orch_args->tensor(i); - if (t.is_child_memory() || t.nbytes() == 0) { + if (t.is_device_memory() || t.nbytes() == 0) { continue; } required += align_up(static_cast(t.nbytes())); @@ -534,7 +534,7 @@ static bool stage_device_args( for (int i = 0; i < tensor_count; i++) { Tensor t = orch_args->tensor(i); - if (t.is_child_memory()) { + if (t.is_device_memory()) { LOG_DEBUG(" Tensor %d: child memory, pass-through (0x%" PRIx64 ")", i, t.buffer.addr); out->add_tensor(t); continue; diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/tensor_create_info.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/tensor_create_info.h index 912839a340..756f706360 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/tensor_create_info.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/tensor_create_info.h @@ -99,7 +99,7 @@ static_assert(offsetof(TensorCreateInfo, ndims) == offsetof(Tensor, ndims)); static_assert(offsetof(TensorCreateInfo, dtype) == offsetof(Tensor, dtype)); static_assert(offsetof(TensorCreateInfo, manual_dep) == offsetof(Tensor, manual_dep)); static_assert(offsetof(TensorCreateInfo, is_contiguous) == offsetof(Tensor, is_contiguous)); -static_assert(offsetof(TensorCreateInfo, __pad_flags__) == offsetof(Tensor, child_memory)); +static_assert(offsetof(TensorCreateInfo, __pad_flags__) == offsetof(Tensor, address_space)); static_assert(offsetof(TensorCreateInfo, shapes) == offsetof(Tensor, shapes)); // ============================================================================ diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index 930bac6e19..9789b152d7 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -229,30 +229,6 @@ void Orchestrator::report_task_error(TaskSlot slot, const std::string &message) record_run_error(task.run_id, std::make_exception_ptr(std::runtime_error(message))); } -uint64_t Orchestrator::malloc(int worker_id, size_t size) { - auto *wt = manager_->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); - if (!wt) throw std::runtime_error("Orchestrator::malloc: invalid worker_id"); - return wt->control_malloc(size); -} - -void Orchestrator::free(int worker_id, uint64_t ptr) { - auto *wt = manager_->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); - if (!wt) throw std::runtime_error("Orchestrator::free: invalid worker_id"); - wt->control_free(ptr); -} - -void Orchestrator::copy_to(int worker_id, uint64_t dst, uint64_t src, size_t size) { - auto *wt = manager_->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); - if (!wt) throw std::runtime_error("Orchestrator::copy_to: invalid worker_id"); - wt->control_copy_to(dst, src, size); -} - -void Orchestrator::copy_from(int worker_id, uint64_t dst, uint64_t src, size_t size) { - auto *wt = manager_->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); - if (!wt) throw std::runtime_error("Orchestrator::copy_from: invalid worker_id"); - wt->control_copy_from(dst, src, size); -} - TaskSlotState &Orchestrator::slot_state(TaskSlot s) { TaskSlotState *p = allocator_->slot_state(s); if (!p) throw std::runtime_error("Orchestrator::slot_state: invalid slot id"); @@ -265,7 +241,7 @@ TaskSlotState &Orchestrator::slot_state(TaskSlot s) { uint64_t Orchestrator::output_alloc_bytes(const Tensor &t) { return align_up(t.nbytes(), HEAP_ALIGN); } -Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { +Orchestrator::ManagedAlloc Orchestrator::alloc_managed_slot(const std::vector &shape, DataType dtype) { auto run = current_building_run(); if (shape.empty()) { // Rank-0 tensors are not supported across the ABI (Tensor enforces @@ -283,11 +259,6 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { uint64_t bytes = numel * get_element_size(dtype); uint64_t aligned = align_up(bytes, HEAP_ALIGN); - // 0-byte request (e.g. shape with a zero dim) flows straight through the - // allocator as a slot-only claim — matches reserve_outputs_and_slot. - // Skip tensormap registration when the returned heap_ptr is nullptr, - // since 0 is the sentinel for "no tensor" in infer_deps. - // // Inherit the caller's scope depth so alloc buffers land in the same // ring as any tasks submitted inside that scope — an alloc inside a // nested `with orch.scope():` uses the nested ring and reclaims @@ -301,13 +272,6 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { s.reset(); s.run_id = run->id; - uint64_t ptr = reinterpret_cast(ar.heap_ptr); - if (ptr != 0) { - TensorKey key = TensorKey::local_host(ptr); - tensormap_->insert(run->id, key, ar.slot); - s.output_keys.push_back(key); - } - // No fanin — alloc has no work to wait on. s.fanin_count = 0; s.fanin_released.store(0, std::memory_order_relaxed); @@ -326,21 +290,27 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { s.fanout_released.store(1, std::memory_order_relaxed); if (scope_ref > 0) scope_->register_task(ar.slot); - s.state.store(TaskState::COMPLETED, std::memory_order_release); + return {ar.slot, reinterpret_cast(ar.heap_ptr), bytes}; +} - increment_run_tasks(run->id); +void Orchestrator::finalize_managed_slot(const ManagedAlloc &m, const TensorKey &key) { + TaskSlotState &s = slot_state(m.slot); + // A 0-byte request (shape with a zero dim) has heap_ptr == 0 — skip tensormap registration since + // 0 is the "no tensor" sentinel in infer_deps. + if (m.va != 0) { + tensormap_->insert(s.run_id, key, m.slot); + s.output_keys.push_back(key); + } + s.state.store(TaskState::COMPLETED, std::memory_order_release); + increment_run_tasks(s.run_id); +} - // Build a contiguous external Tensor over the allocated buffer. ptr may be - // 0 for a 0-byte request (a shape with a zero dim), in which case - // init_external sets buffer.addr == 0 — the "no tensor" sentinel honored by - // infer_deps; buffer.size carries numel*elem. shape is non-empty (rejected - // at entry), so ndims >= 1 holds for init_external's assertion. - Tensor t{}; - t.init_external( - reinterpret_cast(ptr), bytes, shape.data(), static_cast(shape.size()), dtype, - /*version=*/0 - ); - return t; +uint64_t Orchestrator::alloc(const std::vector &shape, DataType dtype, const CanonicalIdentity &identity) { + ManagedAlloc m = alloc_managed_slot(shape, dtype); + // Key by the identity's canonical hash so a BufferRef over this VA (carrying the same identity) + // resolves to this slot in infer_deps — the alloc→BufferRef bridge. + finalize_managed_slot(m, TensorKey::local_host(CanonicalIdentityHash{}(identity))); + return m.va; } // ============================================================================= @@ -634,22 +604,23 @@ void Orchestrator::validate_remote_sidecars( } } for (int32_t i = 0; i < args.tensor_count(); ++i) { - const Tensor &tensor = args.tensor(i); + const BufferRef &ref = args.tensor(i); const RemoteTensorSidecar &tensor_sidecar = sidecar.tensors[static_cast(i)]; - if (tensor_sidecar.present && tensor.buffer.addr != 0) { - throw std::invalid_argument("Orchestrator: remote tensor metadata data field must be zero"); + // A remote arg carries no local backing: its BufferRef is a placeholder (nbytes 0 or a + // REMOTE_SIDECAR backend) and the real descriptor lives in the sidecar. + bool has_local_backing = + ref.handle.nbytes != 0 && ref.handle.backend_kind != static_cast(BackendKind::REMOTE_SIDECAR); + if (tensor_sidecar.present && has_local_backing) { + throw std::invalid_argument("Orchestrator: remote tensor metadata must not carry a local backing"); } - if (!tensor_sidecar.present && tensor.buffer.addr != 0) { - throw std::invalid_argument("Orchestrator: remote tensor uses a bare host pointer without sidecar"); + if (!tensor_sidecar.present && has_local_backing) { + throw std::invalid_argument("Orchestrator: remote tensor uses a local backing without sidecar"); } - if (args.tag(i) == TensorArgType::OUTPUT && tensor.buffer.addr == 0 && !tensor_sidecar.present) { + if (args.tag(i) == TensorArgType::OUTPUT && !has_local_backing && !tensor_sidecar.present) { throw std::invalid_argument("Orchestrator: remote OUTPUT tensor requires a RemoteTensorRef sidecar"); } - if (!tensor_sidecar.present && tensor.nbytes() != 0) { - throw std::invalid_argument("Orchestrator: remote tensor payload requires a RemoteTensorRef sidecar"); - } - if (tensor.is_child_memory() && !tensor_sidecar.present) { - throw std::invalid_argument("Orchestrator: remote child-memory tensor requires a sidecar"); + if (ref.handle.address_space == static_cast(AddressSpace::DEVICE) && !tensor_sidecar.present) { + throw std::invalid_argument("Orchestrator: remote device-memory tensor requires a sidecar"); } if (tensor_sidecar.present && tensor_sidecar.desc.address_space != RemoteAddressSpace::HOST_INLINE) { if (tensor_sidecar.desc.owner_worker_id < 0) { @@ -675,55 +646,20 @@ void Orchestrator::validate_remote_sidecars( } // ============================================================================= -// reserve_outputs_and_slot — atomic slot + heap carve-up for this submit +// reserve_slot — claim this submit's task slot // ============================================================================= // -// Walks every OUTPUT-tagged tensor that arrived with `data == 0` and reserves -// aligned slabs out of a single contiguous HeapRing allocation. OUTPUT tensors -// with a user-supplied data pointer are left untouched (that's the -// OUTPUT_EXISTING-equivalent back-compat path for callers that pre-fill -// OUTPUT.data themselves). The single allocator call owns both the slot and -// the heap range, so there is no partial-failure rollback. +// A slot-only allocation (0 heap bytes). Args are BufferRefs backed by handles the +// caller already owns (create_buffer / create_host_buffer), so the orchestrator no +// longer auto-allocates OUTPUT memory — an OUTPUT is a handle-backed ref like any +// other, tracked by canonical identity in infer_deps. AllocResult Orchestrator::reserve_outputs_and_slot( std::vector &args_list, const std::vector &remote_sidecars ) { - uint64_t total_bytes = 0; - for (size_t g = 0; g < args_list.size(); ++g) { - const TaskArgs &a = args_list[g]; - for (int32_t i = 0; i < a.tensor_count(); ++i) { - if (a.tag(i) != TensorArgType::OUTPUT) continue; - if (a.tensor(i).buffer.addr != 0) continue; // user supplied a pointer — leave alone - bool remote_output = !remote_sidecars.empty() && - static_cast(i) < remote_sidecars[g].tensors.size() && - remote_sidecars[g].tensors[static_cast(i)].present; - if (remote_output) continue; - total_bytes += output_alloc_bytes(a.tensor(i)); - } - } - - AllocResult ar = allocator_->alloc(total_bytes, scope_->current_depth()); - if (ar.slot == INVALID_SLOT) return ar; - - // Hand slabs out in the same order we counted them. - uint64_t off = 0; - char *base = static_cast(ar.heap_ptr); - for (size_t g = 0; g < args_list.size(); ++g) { - TaskArgs &a = args_list[g]; - for (int32_t i = 0; i < a.tensor_count(); ++i) { - if (a.tag(i) != TensorArgType::OUTPUT) continue; - Tensor &t = a.tensor(i); - if (t.buffer.addr != 0) continue; - bool remote_output = !remote_sidecars.empty() && - static_cast(i) < remote_sidecars[g].tensors.size() && - remote_sidecars[g].tensors[static_cast(i)].present; - if (remote_output) continue; - uint64_t slab = output_alloc_bytes(t); - t.buffer.addr = reinterpret_cast(base + off); - off += slab; - } - } - return ar; + (void)args_list; + (void)remote_sidecars; + return allocator_->alloc(0, scope_->current_depth()); } // ============================================================================= @@ -769,7 +705,7 @@ void Orchestrator::infer_deps( int32_t worker_id = (g < target_worker_ids.size()) ? target_worker_ids[g] : -1; const TaskArgs &a = args_list[g]; for (int32_t i = 0; i < a.tensor_count(); ++i) { - const Tensor &t = a.tensor(i); + const BufferRef &r = a.tensor(i); TensorKey key{}; bool has_key = false; if (!remote_sidecars.empty()) { @@ -787,9 +723,18 @@ void Orchestrator::infer_deps( } } if (!has_key) { - if (t.buffer.addr == 0) continue; // null tensor — nothing to track - key = t.is_child_memory() ? TensorKey::local_child(t.buffer.addr, worker_id) : - TensorKey::local_host(t.buffer.addr); + if (r.handle.nbytes == 0) continue; // placeholder / null ref — nothing to track + // Key a local BufferRef by its canonical identity alone (buffer granularity) — the + // successor of the former buffer-address key now that BufferRef carries identity, not + // an address. Any two refs to the same buffer collide (a candidate dependency); the + // byte_offset/footprint overlap that would refine this to only *conflicting* sub-views + // is a future precision pass, not part of the key (folding it in would split + // same-buffer refs into distinct keys and miss real dependencies). + CanonicalIdentityHash idh; + uint64_t k = idh(r.handle.identity); + key = r.handle.address_space == static_cast(AddressSpace::DEVICE) ? + TensorKey::local_child(k, worker_id) : + TensorKey::local_host(k); has_key = true; } TensorArgType tag = a.tag(i); diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index 5273def13c..c7ac9cf2e1 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -17,7 +17,7 @@ * - submit_next_level_group(CallableIdentity, vector, CallConfig, worker_ids) * - submit_sub(CallableIdentity, TaskArgs) * - submit_sub_group(CallableIdentity, vector) - * - alloc(shape, dtype) — runtime-owned intermediate buffer + * - alloc(shape, dtype, identity) — runtime-owned intermediate buffer (returns its VA) * * Each TaskArgs carries per-tensor TensorArgType tags. The Orchestrator * walks those tags to drive dependency inference and — for OUTPUT tags with @@ -75,22 +75,16 @@ class Orchestrator { std::function ready_notify_cb = {} ); - // Allocate an intermediate buffer from the Worker's HeapRing (MAP_SHARED, - // visible to forked child workers). Returns a contiguous Tensor whose - // `.buffer.addr` points into the ring. + // Allocate an intermediate buffer from the Worker's HeapRing (MAP_SHARED, visible to forked child + // workers) and return its VA. Registered in the tensormap under the identity's canonical hash + // (matching how infer_deps keys a BufferRef), so the caller can wrap the VA as a FORK_SHM + // BufferHandle carrying `identity` and a ref over it dependency-wires to this slot. Backs + // Worker.alloc_shared_tensor / Orchestrator.alloc (Python). // - // Lifetime: aligned with a synthetic task slot. The buffer is reclaimed - // (FIFO, via last_alive) once every downstream consumer tagging the - // pointer has reached CONSUMED and scope_end has released the scope ref. - Tensor alloc(const std::vector &shape, DataType dtype); - - // Memory management on a specific next-level worker. Thread-safe: - // can be called from the orch thread while the target worker is - // running a task (MemoryAllocator is mutex-protected). - uint64_t malloc(int worker_id, size_t size); - void free(int worker_id, uint64_t ptr); - void copy_to(int worker_id, uint64_t dst, uint64_t src, size_t size); - void copy_from(int worker_id, uint64_t dst, uint64_t src, size_t size); + // Lifetime: aligned with a synthetic task slot. The buffer is reclaimed (FIFO, via last_alive) once + // every downstream consumer tagging the ref has reached CONSUMED and scope_end has released the + // scope ref. + uint64_t alloc(const std::vector &shape, DataType dtype, const CanonicalIdentity &identity); // Submit a NEXT_LEVEL task. `callable` is the stable identity returned // by Worker.register(); the child resolves its digest to a private slot. @@ -199,6 +193,17 @@ class Orchestrator { // hold a recently-allocated slot id should always get a valid pointer. TaskSlotState &slot_state(TaskSlot s); + // Managed HeapRing intermediate backing alloc(): claims a synthetic auto-free slot + // (no fanin; fanout = scope ref) and returns (slot, heap VA, byte size). The caller registers the + // tensormap key (the identity's canonical hash) and marks the slot COMPLETED. + struct ManagedAlloc { + TaskSlot slot; + uint64_t va; + uint64_t bytes; + }; + ManagedAlloc alloc_managed_slot(const std::vector &shape, DataType dtype); + void finalize_managed_slot(const ManagedAlloc &m, const TensorKey &key); + // Shared submit machinery. Takes `args_list` by value so the Orchestrator // can patch `tensor.data` on OUTPUT tensors flagged for auto-allocation. SubmitResult submit_impl( diff --git a/src/common/hierarchical/remote_endpoint.cpp b/src/common/hierarchical/remote_endpoint.cpp index fb96e3e6f4..0148ebc09c 100644 --- a/src/common/hierarchical/remote_endpoint.cpp +++ b/src/common/hierarchical/remote_endpoint.cpp @@ -633,35 +633,39 @@ remote_l3::TaskPayloadWire RemoteL3Endpoint::build_task_payload(const TaskSlotSt payload.callable_digest = slot.callable.digest; payload.config = slot.config; - TaskArgsView view = slot.args_view(group_index); + const TaskArgs &a = slot.args(group_index); const RemoteTaskArgsSidecar &sidecar = slot.remote_sidecar_for(group_index); - if (!sidecar.tensors.empty() && sidecar.tensors.size() != static_cast(view.tensor_count)) { + if (!sidecar.tensors.empty() && sidecar.tensors.size() != static_cast(a.tensor_count())) { throw std::runtime_error("RemoteL3Endpoint::run: remote sidecar tensor count does not match TaskArgs"); } payload.args.inline_payload = sidecar.inline_payload; - payload.args.tensor_metadata.reserve(static_cast(view.tensor_count)); - payload.args.remote_desc.reserve(static_cast(view.tensor_count)); + payload.args.tensor_metadata.reserve(static_cast(a.tensor_count())); + payload.args.remote_desc.reserve(static_cast(a.tensor_count())); - for (int32_t i = 0; i < view.tensor_count; ++i) { - Tensor tensor = view.tensors(i); + for (int32_t i = 0; i < a.tensor_count(); ++i) { + const BufferRef &ref = a.tensor(i); RemoteTensorSidecar tensor_sidecar{}; if (!sidecar.tensors.empty()) tensor_sidecar = sidecar.tensors[static_cast(i)]; - if (tensor.buffer.addr != 0 && !tensor_sidecar.present) { - throw std::runtime_error("RemoteL3Endpoint::run: bare host pointer submitted without remote sidecar"); + bool has_local_backing = + ref.handle.nbytes != 0 && ref.handle.backend_kind != static_cast(BackendKind::REMOTE_SIDECAR); + if (has_local_backing && !tensor_sidecar.present) { + throw std::runtime_error("RemoteL3Endpoint::run: local backing submitted without remote sidecar"); } - if (tensor.is_child_memory() && !tensor_sidecar.present) { - throw std::runtime_error("RemoteL3Endpoint::run: child-memory tensor submitted without remote sidecar"); + if (ref.handle.address_space == static_cast(AddressSpace::DEVICE) && !tensor_sidecar.present) { + throw std::runtime_error("RemoteL3Endpoint::run: device-memory tensor submitted without remote sidecar"); } - if (!tensor_sidecar.present && tensor.nbytes() != 0) { - throw std::runtime_error("RemoteL3Endpoint::run: tensor payload submitted without remote sidecar"); - } - tensor.buffer.addr = 0; - payload.args.tensor_metadata.push_back(tensor); + // Metadata is the ref's view geometry with a null address (the remote side re-materializes + // from remote_desc); shape/dtype travel, the backing does not. + Tensor meta = make_tensor_external( + nullptr, ref.shapes, ref.ndims, ref.dtype, /*manual_dep=*/false, /*version=*/0, + static_cast(ref.handle.address_space) + ); + payload.args.tensor_metadata.push_back(meta); payload.args.remote_desc.push_back(tensor_sidecar); } - payload.args.scalars.reserve(static_cast(view.scalar_count)); - for (int32_t i = 0; i < view.scalar_count; ++i) - payload.args.scalars.push_back(view.scalars[i]); + payload.args.scalars.reserve(static_cast(a.scalar_count())); + for (int32_t i = 0; i < a.scalar_count(); ++i) + payload.args.scalars.push_back(a.scalar(i)); return payload; } diff --git a/src/common/hierarchical/remote_wire.cpp b/src/common/hierarchical/remote_wire.cpp index 8e3dbd820e..ca57ca3990 100644 --- a/src/common/hierarchical/remote_wire.cpp +++ b/src/common/hierarchical/remote_wire.cpp @@ -352,7 +352,7 @@ std::vector encode_tensor(const Tensor &tensor) { put_u32(out, shape); put_u32(out, tensor.ndims); put_u32(out, static_cast(tensor.dtype)); - put_u8(out, tensor.child_memory); + put_u8(out, static_cast(tensor.address_space)); for (int i = 0; i < 7; ++i) put_u8(out, 0); return out; @@ -368,15 +368,15 @@ Tensor decode_tensor(const uint8_t *data, size_t size, size_t &offset, bool remo ensure(ndims > 0 && ndims <= MAX_TENSOR_DIMS, "remote_wire: tensor ndims out of range"); uint32_t dtype = get_u32(data, size, offset); ensure(valid_dtype(dtype), "remote_wire: unknown tensor dtype"); - uint8_t child_memory = get_u8(data, size, offset); - ensure(child_memory == 0 || child_memory == 1, "remote_wire: tensor child_memory must be 0 or 1"); + uint8_t address_space = get_u8(data, size, offset); + ensure(address_space == 0 || address_space == 1, "remote_wire: tensor address_space must be 0 or 1"); for (int i = 0; i < 7; ++i) { ensure(get_u8(data, size, offset) == 0, "remote_wire: Tensor reserved bytes must be zero"); } // Reconstruct a contiguous external Tensor (owner=invalid, row-major strides). return make_tensor_external( reinterpret_cast(addr), shapes, ndims, static_cast(dtype), /*manual_dep=*/false, - /*version=*/0, child_memory + /*version=*/0, static_cast(address_space) ); } diff --git a/src/common/hierarchical/types.h b/src/common/hierarchical/types.h index 5819f88ae3..920d9011a5 100644 --- a/src/common/hierarchical/types.h +++ b/src/common/hierarchical/types.h @@ -402,11 +402,9 @@ struct TaskSlotState { return remote_sidecar; } - // Zero-copy view over the i-th worker's args (THREAD-mode dispatch). + // The i-th worker's BufferRef args (the L3→L2 wire element). // `i` must be 0 for non-group slots; 0..group_size()-1 for groups. - TaskArgsView args_view(int32_t i) const { - return is_group_ ? make_view(task_args_list[static_cast(i)]) : make_view(task_args); - } + const TaskArgs &args(int32_t i) const { return is_group_ ? task_args_list[static_cast(i)] : task_args; } TaskSlotState() = default; TaskSlotState(const TaskSlotState &) = delete; diff --git a/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index f8e2e82b09..5b1a2d18e5 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -176,6 +176,30 @@ class Worker { } void remote_release_import(const RemoteBufferHandle &handle) { manager_.control_remote_release_import(handle); } + // Device memory on a next-level worker. The Worker is the sole owner of buffer lifecycle; the + // Python Orchestrator's malloc/free/copy are thin wrappers that route back here. Each resolves + // the target worker thread and forwards to its control-plane op (CTRL_MALLOC / FREE / COPY). + uint64_t malloc(int worker_id, size_t size) { + auto *wt = manager_.get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (!wt) throw std::runtime_error("Worker::malloc: invalid worker_id"); + return wt->control_malloc(size); + } + void free(int worker_id, uint64_t ptr) { + auto *wt = manager_.get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (!wt) throw std::runtime_error("Worker::free: invalid worker_id"); + wt->control_free(ptr); + } + void copy_to(int worker_id, uint64_t dst, uint64_t src, size_t size) { + auto *wt = manager_.get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (!wt) throw std::runtime_error("Worker::copy_to: invalid worker_id"); + wt->control_copy_to(dst, src, size); + } + void copy_from(int worker_id, uint64_t dst, uint64_t src, size_t size) { + auto *wt = manager_.get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (!wt) throw std::runtime_error("Worker::copy_from: invalid worker_id"); + wt->control_copy_from(dst, src, size); + } + // Broadcast CTRL_REGISTER / CTRL_UNREGISTER for a ChipCallable digest to // every NEXT_LEVEL child in parallel. `blob_ptr`/`blob_size` describe // the contiguous ChipCallable bytes (see PyChipCallable::buffer_ptr / diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 979c42aa41..c4ac6582d9 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -296,7 +296,7 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis if (ring == nullptr) throw std::invalid_argument("LocalMailboxEndpoint::run: null ring"); TaskSlotState &s = *ring->slot_state(dispatch.task_slot); int32_t group_index = dispatch.group_index; - TaskArgsView view = s.args_view(group_index); + const TaskArgs &a = s.args(group_index); if (!s.remote_sidecar_for(group_index).empty()) { WorkerCompletion completion; completion.task_slot = dispatch.task_slot; @@ -332,34 +332,22 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis // Write config as a single packed POD block (see call_config.h). std::memcpy(mbox() + MAILBOX_OFF_CONFIG, &s.config, sizeof(CallConfig)); - // Write length-prefixed TaskArgs blob: [T][S][tensors][scalars]. - size_t blob_bytes = TASK_ARGS_BLOB_HEADER_SIZE + static_cast(view.tensor_count) * sizeof(Tensor) + - static_cast(view.scalar_count) * sizeof(uint64_t); + // Write the versioned, length-prefixed BufferRef blob (the L3→L2 wire; the child materializes it + // to a Tensor blob before run). + size_t blob_bytes = bufferref_blob_size(a.tensor_count(), a.scalar_count()); if (blob_bytes > MAILBOX_ARGS_CAPACITY) { completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; completion.error_message = "LocalMailboxEndpoint::run: args blob exceeds mailbox capacity: need " + std::to_string(blob_bytes) + " bytes, capacity " + std::to_string(MAILBOX_ARGS_CAPACITY) + - " bytes, tensors=" + std::to_string(view.tensor_count) + ", scalars=" + std::to_string(view.scalar_count); + " bytes, refs=" + std::to_string(a.tensor_count()) + ", scalars=" + std::to_string(a.scalar_count()); return completion; } uint8_t *hash_dst = reinterpret_cast(mbox() + MAILBOX_OFF_TASK_CALLABLE_HASH); std::memcpy(hash_dst, s.callable.digest.data(), CALLABLE_HASH_DIGEST_SIZE); uint8_t *d = reinterpret_cast(mbox() + MAILBOX_OFF_TASK_ARGS_BLOB); - std::memcpy(d + 0, &view.tensor_count, sizeof(int32_t)); - std::memcpy(d + 4, &view.scalar_count, sizeof(int32_t)); - if (view.tensor_count > 0) { - std::memcpy( - d + TASK_ARGS_BLOB_HEADER_SIZE, view.tensor_bytes, static_cast(view.tensor_count) * sizeof(Tensor) - ); - } - if (view.scalar_count > 0) { - std::memcpy( - d + TASK_ARGS_BLOB_HEADER_SIZE + static_cast(view.tensor_count) * sizeof(Tensor), view.scalars, - static_cast(view.scalar_count) * sizeof(uint64_t) - ); - } + write_bufferref_blob(d, a.tensor_data(), a.tensor_count(), a.scalar_data(), a.scalar_count()); // Signal child process. write_mailbox_state(MailboxState::TASK_READY); diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index c8a73f6393..3145be1136 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -74,15 +74,12 @@ enum class MailboxState : int32_t { INIT_FAILED = 7, }; -// Sized so the args region can hold any TaskArgs the runtime itself accepts -// (CHIP_MAX_TENSOR_ARGS tensors + CHIP_MAX_SCALAR_ARGS scalars; see the -// static_assert after MAILBOX_ARGS_CAPACITY). 4096 was too tight for composed -// child kernels with many tensor args (issue #1024). -// Bumped 16384 -> 32768 when TaskArgs moved from the former 40 B compact tensor -// to the unified 128 B Tensor: the worst-case blob (CHIP_MAX_TENSOR_ARGS tensors) -// grew ~3x, and 128*128 B = 16 KB alone exceeded the old mailbox (see the -// capacity static_assert after MAILBOX_ARGS_CAPACITY). -static constexpr size_t MAILBOX_SIZE = 32768; +// Sized so the args region holds the largest args blob the runtime accepts: +// CHIP_MAX_TENSOR_ARGS elements at sizeof(BufferRef) (the self-describing wire +// element, 272 B) + CHIP_MAX_SCALAR_ARGS scalars ~= 35.9 KB (see the capacity +// static_assert after MAILBOX_ARGS_CAPACITY). 40960 = 10 * 4 KB pages. The 272 B +// BufferRef dominates the 128 B Tensor blob, so a Tensor-wire dispatch also fits. +static constexpr size_t MAILBOX_SIZE = 40960; // Error message region lives at the mailbox tail. 256 B of headroom is // enough for `: ` produced by the child-side @@ -120,9 +117,11 @@ static constexpr ptrdiff_t MAILBOX_OFF_CONTROL_CALLABLE_HASH = static constexpr size_t MAILBOX_ARGS_CAPACITY = MAILBOX_SIZE - static_cast(MAILBOX_OFF_TASK_ARGS_BLOB) - MAILBOX_ERROR_MSG_SIZE; static_assert( - MAILBOX_ARGS_CAPACITY >= TASK_ARGS_BLOB_HEADER_SIZE + static_cast(CHIP_MAX_TENSOR_ARGS) * sizeof(Tensor) + + MAILBOX_ARGS_CAPACITY >= BUFFERREF_BLOB_HEADER_SIZE + + static_cast(CHIP_MAX_TENSOR_ARGS) * sizeof(BufferRef) + static_cast(CHIP_MAX_SCALAR_ARGS) * sizeof(uint64_t), - "mailbox args region must hold the largest TaskArgs blob the runtime accepts (issue #1024)" + "mailbox args region must hold the largest args blob the runtime accepts: the self-describing " + "BufferRef wire (272 B/arg) dominates the 128 B Tensor blob (issue #1024)" ); // Control sub-commands (written at MAILBOX_OFF_CALLABLE when state == CONTROL_*) diff --git a/src/common/task_interface/buffer_handle.h b/src/common/task_interface/buffer_handle.h new file mode 100644 index 0000000000..20e66a2499 --- /dev/null +++ b/src/common/task_interface/buffer_handle.h @@ -0,0 +1,207 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +#pragma once + +/** + * BufferHandle / BufferRef ABI — typed, versioned, opaque cross-layer buffer identity. + * + * Layout implements the frozen logical schema in + * .docs/L3-new/worker-memory-model/bufferhandle-abi.md (that doc is authoritative for field set / + * widths / enum values / endianness / evolution; exact byte offsets here are this P1 wire's choice). + * + * Three types: + * - BufferHandleDescriptor : the owner's self-describing wire descriptor. Carries backing + * properties + a versioned length-delimited backend body. abi_version + * (u16) leads so a decoder rejects an unknown version before trusting the + * rest. Embedded whole in every BufferRef built over the handle. + * - BufferRef : the blob-carried wire element. Embeds the full BufferHandleDescriptor + * plus a view (byte_offset, shape, strides, dtype) — self-describing, so + * a consumer materializes it lazily on receipt with no prior handshake. + * No materialized address. + * - CanonicalIdentity : owner_instance_id + owner_worker_path + buffer_id + generation. The + * key both the owner registry and every consumer import cache use. + * + * Endianness: all multi-byte integers little-endian. owner_instance_id is an opaque byte sequence + * (bytewise-compared, no integer/endianness meaning). Unknown version / backend / descriptor_version + * is rejected, never silently accepted. + */ + +#include +#include +#include +#include + +#include "data_type.h" + +// Wire version of the handle schema (u16). A decoder rejects an unknown abi_version rather than +// misreading a future layout. 0 is reserved (illegal). +inline constexpr uint16_t BUFFER_ABI_VERSION = 1; + +// Version of a backend_descriptor body. Unknown descriptor_version is rejected like abi_version. +inline constexpr uint8_t BUFFER_DESCRIPTOR_VERSION = 1; + +// owner_instance_id is a fixed-width opaque nonce (compared bytewise; no integer/endianness meaning). +inline constexpr uint32_t OWNER_INSTANCE_ID_BYTES = 16; + +// Bounded length-delimited limits (single constants; revisit before the final ABI freeze). +// owner_worker_path is a UTF-8 tree path like "L4/L3[2]/L2[5]"; backend body is per-backend. +inline constexpr uint32_t PATH_MAX_BYTES = 64; +inline constexpr uint32_t DESC_MAX_BYTES = 96; + +// AddressSpace (HOST/DEVICE) is shared with Tensor and lives in data_type.h. + +// Which workers may see a backing. Single-hop/multi-hop visibility is explicit, not by convention. +enum class Visibility : uint8_t { + PRIVATE = 0, + SHARED = 1, +}; + +// The backing's granted permission. A per-arg TensorArgType requests read/write and is validated +// against this at submit (requested must be a subset of granted). +enum class AccessMode : uint8_t { + READ = 0, + WRITE = 1, + READWRITE = 2, +}; + +// Materialization backend of a handle. The consumer resolves a BufferRef to a local address via the +// import registry keyed by canonical identity; this tag selects how. REMOTE_SIDECAR is reserved for +// P2 and rejected on decode in P1. Values are frozen; 5.. reserved (unknown tag => reject). +enum class BackendKind : uint8_t { + FORK_SHM = 0, + POSIX_SHM = 1, + VMM_WINDOW = 2, + REMOTE_SIDECAR = 3, + DEVICE_MALLOC = 4, +}; + +/** + * 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 16-byte + * per-incarnation nonce) and `owner_worker_path` disambiguate it, and `generation` detects buffer_id + * reuse (ABA). The key of both the owner registry and every consumer import registry. + * + * `owner_worker_path` is a bounded length-delimited UTF-8 tree path ("L4/L3[2]/L2[5]"): `path_len` + * valid bytes in `owner_worker_path[0, path_len)`; bytes beyond `path_len` and `_pad` are zero. + */ +struct CanonicalIdentity { + uint8_t owner_instance_id[OWNER_INSTANCE_ID_BYTES]; + uint64_t buffer_id; + uint32_t generation; + uint16_t path_len; + uint8_t _pad[2]; + char owner_worker_path[PATH_MAX_BYTES]; +}; + +static_assert(std::is_trivially_copyable_v); +static_assert(sizeof(CanonicalIdentity) == 96, "CanonicalIdentity is wire ABI"); +static_assert(offsetof(CanonicalIdentity, owner_instance_id) == 0); +static_assert(offsetof(CanonicalIdentity, buffer_id) == 16); +static_assert(offsetof(CanonicalIdentity, generation) == 24); +static_assert(offsetof(CanonicalIdentity, path_len) == 28); +static_assert(offsetof(CanonicalIdentity, owner_worker_path) == 32); + +inline bool operator==(const CanonicalIdentity &a, const CanonicalIdentity &b) { + return a.buffer_id == b.buffer_id && a.generation == b.generation && a.path_len == b.path_len && + std::memcmp(a.owner_instance_id, b.owner_instance_id, OWNER_INSTANCE_ID_BYTES) == 0 && + std::memcmp(a.owner_worker_path, b.owner_worker_path, a.path_len) == 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 the significant identity +// bytes; unused path bytes past path_len are excluded so they never perturb the hash. +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); + h = mix(h, k.path_len); + for (uint16_t i = 0; i < k.path_len; ++i) + h = mix(h, static_cast(k.owner_worker_path[i])); + return h; + } +}; + +/** + * The owner's self-describing handle descriptor — embedded whole in every BufferRef built over the + * handle. A consumer materializes it lazily on first receipt (no separate export handshake) and + * caches `canonical identity -> local base` (map-once). `backend_kind` + `descriptor_version` + + * `body[0, body_len)` carry the per-backend materialization (POSIX/fork shm name, VMM + * shareable-handle, device VA, ...). Version-prefixed (`abi_version` u16 leads); unknown abi_version / + * backend_kind / descriptor_version is rejected before trusting the rest. `address_space` / + * `visibility` / `access` / `backend_kind` are raw u8 so an unknown value can be rejected without + * invoking undefined enum behavior. + */ +struct BufferHandleDescriptor { + uint16_t abi_version; + uint8_t address_space; + uint8_t visibility; + uint8_t access; + uint8_t backend_kind; + uint8_t descriptor_version; + uint8_t _pad0; + CanonicalIdentity identity; + uint64_t nbytes; + uint16_t body_len; + uint8_t _pad1[6]; + char body[DESC_MAX_BYTES]; +}; + +static_assert(std::is_trivially_copyable_v); +static_assert(sizeof(BufferHandleDescriptor) == 216, "BufferHandleDescriptor is wire ABI"); +static_assert(offsetof(BufferHandleDescriptor, abi_version) == 0); +static_assert(offsetof(BufferHandleDescriptor, address_space) == 2); +static_assert(offsetof(BufferHandleDescriptor, visibility) == 3); +static_assert(offsetof(BufferHandleDescriptor, access) == 4); +static_assert(offsetof(BufferHandleDescriptor, backend_kind) == 5); +static_assert(offsetof(BufferHandleDescriptor, descriptor_version) == 6); +static_assert(offsetof(BufferHandleDescriptor, identity) == 8); +static_assert(offsetof(BufferHandleDescriptor, nbytes) == 104); +static_assert(offsetof(BufferHandleDescriptor, body_len) == 112); +static_assert(offsetof(BufferHandleDescriptor, body) == 120); + +/** + * The blob-carried, self-describing wire element: a full embedded handle descriptor plus a strided + * view onto it. Because the descriptor travels with the ref, a consumer needs no prior handshake — + * it materializes the embedded `handle` (backend selects how) on first receipt, keyed by + * `handle.identity`, and reuses the cached base for later refs to the same identity. + * + * Invariants: + * - Carries NO materialized address. The consumer materializes `handle` to a local base, then + * `Tensor.buffer.addr = base`, `Tensor.start_offset = byte_offset / dtype_bytes`. + * - `byte_offset` is a BYTE offset of the view origin; a multiple of the dtype size (validated at + * materialization). + * - `strides[i] > 0` strictly (broadcast / negative step unsupported), carried explicitly — a + * singleton dimension's stride is never normalized away. + */ +struct BufferRef { + BufferHandleDescriptor handle; + uint64_t byte_offset; + uint32_t ndims; + uint32_t shapes[MAX_TENSOR_DIMS]; + uint32_t strides[MAX_TENSOR_DIMS]; + DataType dtype; + uint8_t _pad[3]; +}; + +static_assert(std::is_trivially_copyable_v, "BufferRef must be trivially copyable for blob memcpy"); +static_assert(sizeof(BufferRef) == 272, "BufferRef is wire ABI"); +static_assert(offsetof(BufferRef, handle) == 0); +static_assert(offsetof(BufferRef, byte_offset) == 216); +static_assert(offsetof(BufferRef, ndims) == 224); +static_assert(offsetof(BufferRef, shapes) == 228); +static_assert(offsetof(BufferRef, strides) == 248); +static_assert(offsetof(BufferRef, dtype) == 268); diff --git a/src/common/task_interface/data_type.h b/src/common/task_interface/data_type.h index 17ad8f573a..9de43fb237 100644 --- a/src/common/task_interface/data_type.h +++ b/src/common/task_interface/data_type.h @@ -28,6 +28,17 @@ template inline constexpr bool is_supported_scalar_arg_v = std::is_arithmetic_v>> || std::is_enum_v>>; +// Maximum tensor rank a view (Tensor / BufferRef) can describe. Wire ABI: the shapes[] / strides[] +// arrays in both are sized by it, so it is shared here rather than owned by either. +constexpr int MAX_TENSOR_DIMS = 5; + +// Memory space of a backing. Byte-43 of Tensor and a BufferHandle field both store it, so it is +// shared here. Orthogonal to location (local/remote, derived) and visibility. +enum class AddressSpace : uint8_t { + HOST = 0, + DEVICE = 1, +}; + /** * Supported data types for tensor elements */ diff --git a/src/common/task_interface/task_args.h b/src/common/task_interface/task_args.h index ebfd4fc681..c311750c98 100644 --- a/src/common/task_interface/task_args.h +++ b/src/common/task_interface/task_args.h @@ -43,7 +43,8 @@ #include #include "arg_direction.h" -#include "tensor.h" // unified Tensor (strided) + TensorArgType, carried by TaskArgs and on the wire +#include "buffer_handle.h" // BufferRef wire type + BUFFER_ABI_VERSION for the versioned BufferRef blob +#include "tensor.h" // unified Tensor (strided) + TensorArgType, carried by TaskArgs and on the wire // ============================================================================ // TensorTagMixin — conditionally provides per-tensor tag storage @@ -173,11 +174,16 @@ struct TaskArgsTpl : TensorTagMixin { // Unified user-facing builder: vector-backed with TensorArgType tags. // Used by Orchestrator.submit_*; tags drive dependency inference at submit -// time and are stripped before the args cross the dispatch boundary. -using TaskArgs = TaskArgsTpl; - -// L2 runtime ABI: fixed POD matching runtime.so byte-for-byte. -// Assembled from a TaskArgsView on the child side just before pto2_run_runtime. +// time and are stripped before the args cross the dispatch boundary. The element +// is BufferRef (self-describing view; L3+ holds no C++ Tensor) — the L3→L2 wire +// carries BufferRefs, materialized to ChipStorageTaskArgs (Tensor) on the L2 child. +using TaskArgs = TaskArgsTpl; + +// L2 runtime ABI: fixed POD matching runtime.so byte-for-byte, and the sole +// Tensor-typed args container — the materialized form the L3→L2 BufferRef blob +// is decoded into on the L2 child, and the source/target of the Tensor-blob +// round trip (write_blob / read_blob / make_view). Assembled from a +// TaskArgsView on the child side just before pto2_run_runtime. using ChipStorageTaskArgs = TaskArgsTpl; // ============================================================================ @@ -212,8 +218,8 @@ struct TaskArgsView { } }; -// Build a view directly over a TaskArgs's vectors (THREAD-mode dispatch). -inline TaskArgsView make_view(const TaskArgs &a) { +// Build a view directly over a ChipStorageTaskArgs's arrays (THREAD-mode dispatch). +inline TaskArgsView make_view(const ChipStorageTaskArgs &a) { return TaskArgsView{ a.tensor_count(), a.scalar_count(), reinterpret_cast(a.tensor_data()), a.scalar_data() }; @@ -237,14 +243,14 @@ inline TaskArgsView make_view(const TaskArgs &a) { inline constexpr size_t TASK_ARGS_BLOB_HEADER_SIZE = 8; -inline size_t task_args_blob_size(const TaskArgs &a) { +inline size_t task_args_blob_size(const ChipStorageTaskArgs &a) { return TASK_ARGS_BLOB_HEADER_SIZE + static_cast(a.tensor_count()) * sizeof(Tensor) + static_cast(a.scalar_count()) * sizeof(uint64_t); } // Serialize a TaskArgs into `dst`. Caller must ensure `dst` has room for // task_args_blob_size(a) bytes. Tags are not written. -inline void write_blob(uint8_t *dst, const TaskArgs &a) { +inline void write_blob(uint8_t *dst, const ChipStorageTaskArgs &a) { int32_t T = a.tensor_count(); int32_t S = a.scalar_count(); std::memcpy(dst + 0, &T, sizeof(T)); @@ -301,6 +307,114 @@ inline TaskArgsView read_blob(const uint8_t *src, size_t capacity) { }; } +// ============================================================================ +// BufferRef wire blob — versioned, length-prefixed (P1-B). +// ============================================================================ +// +// Byte layout: +// offset 0: uint32 abi_version = BUFFER_ABI_VERSION +// offset 4: int32 ref_count = R +// offset 8: int32 scalar_count = S +// offset 12: uint32 reserved (= 0) +// offset 16: BufferRef refs[R] (sizeof(BufferRef) B each) +// offset 16 + R*sizeof(BufferRef): uint64_t scalars[S] +// +// The element is BufferRef (embedded handle descriptor + view, no materialized addr) and the envelope +// carries abi_version so a decoder rejects an unknown layout rather than misreading it. The +// reserved word 8-aligns refs[0] (whose first field is a u64) and gates a future layout bump: a +// non-zero reserved is rejected. + +inline constexpr size_t BUFFERREF_BLOB_HEADER_SIZE = 16; + +struct BufferRefBlobView { + int32_t ref_count; + int32_t scalar_count; + const uint8_t *ref_bytes; // R contiguous BufferRef; extract element i with ref(i) + const uint64_t *scalars; + + BufferRef ref(int32_t i) const { + BufferRef r; + std::memcpy(&r, ref_bytes + static_cast(i) * sizeof(BufferRef), sizeof(BufferRef)); + return r; + } +}; + +inline size_t bufferref_blob_size(int32_t ref_count, int32_t scalar_count) { + return BUFFERREF_BLOB_HEADER_SIZE + static_cast(ref_count) * sizeof(BufferRef) + + static_cast(scalar_count) * sizeof(uint64_t); +} + +// Serialize refs + scalars into `dst` (caller ensures room for bufferref_blob_size). Writes the +// current abi_version into the envelope. +inline void write_bufferref_blob( + uint8_t *dst, const BufferRef *refs, int32_t ref_count, const uint64_t *scalars, int32_t scalar_count +) { + uint32_t version = BUFFER_ABI_VERSION; + uint32_t reserved = 0; + std::memcpy(dst + 0, &version, sizeof(version)); + std::memcpy(dst + 4, &ref_count, sizeof(ref_count)); + std::memcpy(dst + 8, &scalar_count, sizeof(scalar_count)); + std::memcpy(dst + 12, &reserved, sizeof(reserved)); + if (ref_count > 0) { + std::memcpy(dst + BUFFERREF_BLOB_HEADER_SIZE, refs, static_cast(ref_count) * sizeof(BufferRef)); + } + if (scalar_count > 0) { + std::memcpy( + dst + BUFFERREF_BLOB_HEADER_SIZE + static_cast(ref_count) * sizeof(BufferRef), scalars, + static_cast(scalar_count) * sizeof(uint64_t) + ); + } +} + +// Zero-copy view into a blob written by write_bufferref_blob; valid while `src` stays mapped. +// `capacity` bounds the read. Throws on an unknown abi_version, negative counts, or a header that +// would walk past `capacity` (shared-memory corruption or a writer-side bug). +inline BufferRefBlobView read_bufferref_blob(const uint8_t *src, size_t capacity) { + if (capacity < BUFFERREF_BLOB_HEADER_SIZE) { + throw std::runtime_error( + "read_bufferref_blob: capacity " + std::to_string(capacity) + " < header size " + + std::to_string(BUFFERREF_BLOB_HEADER_SIZE) + ); + } + uint32_t version; + std::memcpy(&version, src + 0, sizeof(version)); + if (version != BUFFER_ABI_VERSION) { + throw std::runtime_error( + "read_bufferref_blob: unknown abi_version " + std::to_string(version) + " (expected " + + std::to_string(BUFFER_ABI_VERSION) + ")" + ); + } + int32_t R; + int32_t S; + uint32_t reserved; + std::memcpy(&R, src + 4, sizeof(R)); + std::memcpy(&S, src + 8, sizeof(S)); + std::memcpy(&reserved, src + 12, sizeof(reserved)); + if (reserved != 0) { + throw std::runtime_error("read_bufferref_blob: reserved header word must be zero"); + } + if (R < 0 || S < 0) { + throw std::runtime_error( + "read_bufferref_blob: negative counts — refs=" + std::to_string(R) + ", scalars=" + std::to_string(S) + ); + } + const size_t needed = bufferref_blob_size(R, S); + if (needed > capacity) { + throw std::runtime_error( + "read_bufferref_blob: header reports " + std::to_string(needed) + " bytes (R=" + std::to_string(R) + + ", S=" + std::to_string(S) + ") but capacity is " + std::to_string(capacity) + ); + } + return BufferRefBlobView{ + R, + S, + src + BUFFERREF_BLOB_HEADER_SIZE, + reinterpret_cast( + src + BUFFERREF_BLOB_HEADER_SIZE + static_cast(R) * sizeof(BufferRef) + ), + }; +} + // ============================================================================ // L2 ABI helper: build ChipStorageTaskArgs POD from a view (memcpy'd). // Runs on the child side immediately before crossing into runtime.so. diff --git a/src/common/task_interface/tensor.h b/src/common/task_interface/tensor.h index 8b62208bcd..785a24378c 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 * @@ -104,7 +102,7 @@ struct alignas(64) Tensor { DataType dtype; // Data type of tensor elements bool manual_dep; // True when dependency tracking is creator-only (skip OverlapMap lookup/insert) bool is_contiguous; // Cached: strides[] == row_major_stride(shapes) - uint8_t child_memory; // 0 = host memory (default), 1 = child-managed device memory (skips H2D copy) + AddressSpace address_space; // HOST (default) or DEVICE (child-managed device memory; skips H2D copy) uint32_t shapes[MAX_TENSOR_DIMS]; // Current view shape per dimension (elements) // === Cache line 2 (64B) — warm path (view metadata) === @@ -157,7 +155,7 @@ struct alignas(64) Tensor { /// True when `buffer.addr` is a device pointer allocated by the child process /// (host skips the H2D copy in init_runtime_impl). Host-side concept carried /// across the wire; runtime views inherit it via the cache-line-1 copy. - [[nodiscard]] bool is_child_memory() const { return child_memory != 0; } + [[nodiscard]] bool is_device_memory() const { return address_space == AddressSpace::DEVICE; } /// Logical byte size of the view (numel * element size). For a contiguous /// host-constructed tensor this equals buffer.size; provided for parity with @@ -181,7 +179,7 @@ struct alignas(64) Tensor { /// Enforces the ndims > 0 invariant relied upon by every downstream op. void init_external( void *addr, uint64_t buffer_size_bytes, const uint32_t in_shapes[], uint32_t in_ndims, DataType in_dtype, - int32_t in_version, bool in_manual_dep = false, uint8_t in_child_memory = 0 + int32_t in_version, bool in_manual_dep = false, AddressSpace in_address_space = AddressSpace::HOST ) { always_assert(in_ndims > 0 && in_ndims <= MAX_TENSOR_DIMS); buffer = {reinterpret_cast(addr), buffer_size_bytes}; @@ -190,7 +188,7 @@ struct alignas(64) Tensor { version = in_version; manual_dep = in_manual_dep; is_contiguous = true; - child_memory = in_child_memory; + address_space = in_address_space; start_offset = 0; owner_task_id = PTO2TaskId::invalid(); // Single reverse pass: write shapes, accumulate row-major stride, and @@ -434,10 +432,10 @@ struct alignas(64) Tensor { // default constructor is public — see above — for POD/array storage.) Tensor( void *addr, uint64_t buffer_size_bytes, const uint32_t in_shapes[], uint32_t in_ndims, DataType in_dtype, - int32_t in_version, bool in_manual_dep = false, uint8_t in_child_memory = 0 + int32_t in_version, bool in_manual_dep = false, AddressSpace in_address_space = AddressSpace::HOST ) { init_external( - addr, buffer_size_bytes, in_shapes, in_ndims, in_dtype, in_version, in_manual_dep, in_child_memory + addr, buffer_size_bytes, in_shapes, in_ndims, in_dtype, in_version, in_manual_dep, in_address_space ); } @@ -475,7 +473,11 @@ struct alignas(64) Tensor { friend struct PTO2TaskPayload; friend inline Tensor make_tensor_external( void *addr, const uint32_t shapes[], uint32_t ndims, DataType dtype, bool manual_dep, int32_t version, - uint8_t child_memory + AddressSpace address_space + ); + friend inline Tensor make_tensor_strided( + void *addr, const uint32_t shapes[], const uint32_t strides[], uint32_t ndims, DataType dtype, bool manual_dep, + int32_t version, AddressSpace address_space ); }; @@ -488,7 +490,9 @@ static_assert(offsetof(Tensor, ndims) == 36); static_assert(offsetof(Tensor, dtype) == 40); static_assert(offsetof(Tensor, manual_dep) == 41); static_assert(offsetof(Tensor, is_contiguous) == 42); -static_assert(offsetof(Tensor, child_memory) == 43, "child_memory must be at byte 43 (cacheline 1, former _pad_cl1)"); +static_assert( + offsetof(Tensor, address_space) == 43, "address_space must be at byte 43 (cacheline 1, former child_memory)" +); static_assert(offsetof(Tensor, shapes) == 44, "shapes must start at byte 44 (cacheline 1)"); static_assert(offsetof(Tensor, extent_elem_cache) == 64, "extent_elem_cache must start at byte 64 (cacheline 2)"); static_assert(offsetof(Tensor, strides) == 72); @@ -502,11 +506,41 @@ static_assert(offsetof(Tensor, strides) == 72); // ============================================================================= inline Tensor make_tensor_external( void *addr, const uint32_t shapes[], uint32_t ndims, DataType dtype = DataType::FLOAT32, bool manual_dep = false, - int32_t version = 0, uint8_t child_memory = 0 + int32_t version = 0, AddressSpace address_space = AddressSpace::HOST ) { uint64_t total = 1; for (uint32_t i = 0; i < ndims; i++) { total *= shapes[i]; } - return {addr, total * get_element_size(dtype), shapes, ndims, dtype, version, manual_dep, child_memory}; + return {addr, total * get_element_size(dtype), shapes, ndims, dtype, version, manual_dep, address_space}; +} + +// ============================================================================= +// Tensor factory — a possibly-strided external view. `addr` is the view origin +// (start_offset == 0); `strides[]` are element strides (may be non-row-major, as +// from a transpose / permute / step-sliced BufferRef). Derived fields +// (extent_elem_cache, is_contiguous) are recomputed from shapes + strides; +// buffer.size is the element extent in bytes. +// ============================================================================= +inline Tensor make_tensor_strided( + void *addr, const uint32_t shapes[], const uint32_t strides[], uint32_t ndims, DataType dtype = DataType::FLOAT32, + bool manual_dep = false, int32_t version = 0, AddressSpace address_space = AddressSpace::HOST +) { + always_assert(ndims > 0 && ndims <= MAX_TENSOR_DIMS); + Tensor t{}; + t.buffer.addr = reinterpret_cast(addr); + t.ndims = ndims; + t.dtype = dtype; + t.version = version; + t.manual_dep = manual_dep; + t.address_space = address_space; + t.start_offset = 0; + t.owner_task_id = PTO2TaskId::invalid(); + for (uint32_t i = 0; i < ndims; i++) { + t.shapes[i] = shapes[i]; + t.strides[i] = strides[i]; + } + t.refresh_derived(); // extent_elem_cache + is_contiguous from shapes/strides + t.buffer.size = t.extent_elem_cache * get_element_size(dtype); + return t; } diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py b/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py index f97378b2ef..7dd28fdc8a 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py @@ -148,8 +148,8 @@ def test_prepare_new_identity_after_start_then_run(st_platform, st_device_ids): expected = _golden(a, b) args_pre = _make_args(a, b) args_post = _make_args(a, b) - chip_args_pre, output_names_pre = _build_l3_task_args(args_pre, _ORCH_SIG) - chip_args_post, output_names_post = _build_l3_task_args(args_post, _ORCH_SIG) + chip_args_pre, output_names_pre = _build_l3_task_args(args_pre, _ORCH_SIG, worker) + chip_args_post, output_names_post = _build_l3_task_args(args_post, _ORCH_SIG, worker) assert output_names_pre == ["f"] and output_names_post == ["f"] worker.init() @@ -221,8 +221,8 @@ def test_prepare_new_identity_after_start_parallel_broadcast(st_platform, st_dev # by both broadcast targets. args_pre = _make_args(a, b) args_post = [_make_args(a, b) for _ in device_ids] - chip_args_pre, _ = _build_l3_task_args(args_pre, _ORCH_SIG) - chip_args_post = [_build_l3_task_args(args, _ORCH_SIG)[0] for args in args_post] + chip_args_pre, _ = _build_l3_task_args(args_pre, _ORCH_SIG, worker) + chip_args_post = [_build_l3_task_args(args, _ORCH_SIG, worker)[0] for args in args_post] worker.init() try: @@ -278,7 +278,7 @@ def test_prepare_capacity_overflow_post_start(st_platform, st_device_ids): a, b = 2.0, 3.0 args_pre = _make_args(a, b) - chip_args_pre, _ = _build_l3_task_args(args_pre, _ORCH_SIG) + chip_args_pre, _ = _build_l3_task_args(args_pre, _ORCH_SIG, worker) worker.init() try: @@ -326,8 +326,8 @@ def test_duplicate_prepare_same_hashid_survives_one_unregister(st_platform, st_d # the share_memory_ mappings are inherited by the forked chip child. args_one = _make_args(a, b) args_two = _make_args(a, b) - chip_args_one, _ = _build_l3_task_args(args_one, _ORCH_SIG) - chip_args_two, _ = _build_l3_task_args(args_two, _ORCH_SIG) + chip_args_one, _ = _build_l3_task_args(args_one, _ORCH_SIG, worker) + chip_args_two, _ = _build_l3_task_args(args_two, _ORCH_SIG, worker) worker.init() try: @@ -395,9 +395,9 @@ def test_unregister_last_handle_allows_reprepare_same_hashid(st_platform, st_dev args_one = _make_args(a, b) args_two = _make_args(a, b) args_three = _make_args(a, b) - chip_args_one, _ = _build_l3_task_args(args_one, _ORCH_SIG) - chip_args_two, _ = _build_l3_task_args(args_two, _ORCH_SIG) - chip_args_three, _ = _build_l3_task_args(args_three, _ORCH_SIG) + chip_args_one, _ = _build_l3_task_args(args_one, _ORCH_SIG, worker) + chip_args_two, _ = _build_l3_task_args(args_two, _ORCH_SIG, worker) + chip_args_three, _ = _build_l3_task_args(args_three, _ORCH_SIG, worker) worker.init() try: diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_bufferref_dispatch.py b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_bufferref_dispatch.py new file mode 100644 index 0000000000..659b23c74e --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_bufferref_dispatch.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""L3->L2 chip dispatch over the BufferRef wire (P1-B B3, minimal end-to-end). + +The L3 orch is a pure DAG builder over views: it names task args as BufferRefs built from handles +(create_buffer + handle.ref), never touching data. torch is used only OUTSIDE run() to fill inputs +and read the output. The owner writes a BufferRef blob to the chip mailbox; the chip child +materializes it back to Tensors (ImportRegistry -> materialize_bufferref_blob) and runs the kernel; +the result lands in the shared output buffer with no per-run copy. +""" + +import torch +from simpler.task_interface import ArgDirection as D +from simpler.task_interface import CallConfig, TaskArgs, TensorArgType + +from simpler_setup import SceneTestCase, scene_test + +KERNELS_BASE = "../../../../examples/a2a3/tensormap_and_ringbuffer/vector_example/kernels" + +SIZE = 128 * 128 +DTYPE = torch.float32 +_F32 = 0 # DataType.FLOAT32 value + + +def _golden(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + s = a + b + return (s + 1) * (s + 2) + s + + +def _one_task_orch(chip_handle, a_ref, b_ref, out_ref): + def orch_fn(orch, _args, cfg): + ta = TaskArgs() + ta.add_ref(a_ref, TensorArgType.INPUT) + ta.add_ref(b_ref, TensorArgType.INPUT) + ta.add_ref(out_ref, TensorArgType.OUTPUT_EXISTING) + orch.submit_next_level(chip_handle, ta, cfg, worker=0) + + return orch_fn + + +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestL3BufferRefDispatch(SceneTestCase): + """A single chip task dispatched over the BufferRef wire on one L3 worker.""" + + CALLABLE = { + "callables": [ + { + "name": "vector", + "orchestration": { + "source": f"{KERNELS_BASE}/orchestration/example_orchestration.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{KERNELS_BASE}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{KERNELS_BASE}/aiv/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 2, + "source": f"{KERNELS_BASE}/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + }, + ], + } + + CASES = [ + {"name": "bufferref_dispatch", "platforms": ["a2a3sim", "a2a3"]}, + ] + + def test_run(self, st_worker): + worker = st_worker + chip_handle = type(self)._st_chip_handles["vector"] + + nbytes = SIZE * DTYPE.itemsize + a_h = worker.create_buffer(nbytes) + b_h = worker.create_buffer(nbytes) + out_h = worker.create_buffer(nbytes) + a = b = out = None + result = False + try: + # torch is used only outside run(): fill inputs before, read output after. + a = torch.frombuffer(a_h.shm.buf, dtype=DTYPE, count=SIZE) + b = torch.frombuffer(b_h.shm.buf, dtype=DTYPE, count=SIZE) + out = torch.frombuffer(out_h.shm.buf, dtype=DTYPE, count=SIZE) + a.fill_(5.0) + b.fill_(7.0) + out.zero_() + # The orch names args purely as BufferRef views built from the handles. + a_ref = a_h.ref(shapes=(SIZE,), dtype=_F32) + b_ref = b_h.ref(shapes=(SIZE,), dtype=_F32) + out_ref = out_h.ref(shapes=(SIZE,), dtype=_F32) + worker.run(_one_task_orch(chip_handle, a_ref, b_ref, out_ref), args=None, config=CallConfig()) + result = torch.allclose(out, _golden(a, b), rtol=self.RTOL, atol=self.ATOL) + finally: + # Drop views before the framework-driven close unlinks the shm. + a = b = out = None + assert result, "BufferRef-dispatched chip task did not produce the golden result" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_create_buffer.py b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_create_buffer.py new file mode 100644 index 0000000000..a0cb290f09 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_create_buffer.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""L3 create_buffer owner-side allocation (P1-B). + +``Worker.create_buffer`` allocates a born-shared BufferHandle carrying a typed canonical identity. +There is **no eager export handshake**: the handle's self-describing descriptor rides embedded in +every BufferRef built over it, and a consumer materializes it lazily on first receipt. This test +covers the owner-side contract — a live L3 Worker allocates handles that carry a stable identity, a +monotonic buffer_id, and usable born-shared backing, and releases them cleanly on close. + +a2a3sim: create_buffer is pure host-side (POSIX shm), no platform branching. The vector callable +exists only to give the L3 worker a chip child to fork (create_buffer requires one). +""" + +from simpler.buffer_handle import AddressSpace, BackendKind, Visibility +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, scene_test + +KERNELS_BASE = "../../../../examples/a2a3/tensormap_and_ringbuffer/vector_example/kernels" + + +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestL3CreateBuffer(SceneTestCase): + """create_buffer owner-side allocation contract on a single L3 worker.""" + + CALLABLE = { + "callables": [ + { + "name": "vector", + "orchestration": { + "source": f"{KERNELS_BASE}/orchestration/example_orchestration.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{KERNELS_BASE}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + }, + ], + } + + CASES = [ + {"name": "create_buffer_owner_contract", "platforms": ["a2a3sim"]}, + ] + + def test_run(self, st_worker): + worker = st_worker + + h0 = worker.create_buffer(256) + h1 = worker.create_buffer(512) + + # Owner-side handle contract. + assert h0.backend_kind == BackendKind.POSIX_SHM + assert h0.address_space == AddressSpace.HOST + assert h0.visibility == Visibility.SHARED + assert len(h0.identity.owner_instance_id) == 16 + assert h0.identity.owner_instance_id == h1.identity.owner_instance_id # same incarnation + assert h0.identity.buffer_id != h1.identity.buffer_id # monotonic + assert h0.identity.generation == 1 + assert h0.nbytes == 256 and h1.nbytes == 512 + + # The backing is usable in place (born-shared shm). + assert h0.shm is not None + buf = h0.shm.buf + assert buf is not None + buf[:4] = b"\xde\xad\xbe\xef" + assert bytes(buf[:4]) == b"\xde\xad\xbe\xef" + + # Worker close (framework-driven) releases both handles via _release_all_buffer_handles. diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py index f5e8046b00..d38c45311e 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py @@ -18,8 +18,8 @@ from simpler.task_interface import ArgDirection as D from simpler.task_interface import TaskArgs, TensorArgType -from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, make_tensor_arg, scene_test -from simpler_setup.scene_test import _build_l3_task_args +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.scene_test import _build_l3_task_args, _rehosted_ref KERNELS_BASE = "../../../../examples/a2a3/tensormap_and_ringbuffer/vector_example/kernels" @@ -38,7 +38,7 @@ def run_dag(orch, callables, task_args, config): # SubTask: tag the chip output as INPUT — Orchestrator wires the dep via TensorMap. sub_args = TaskArgs() - sub_args.add_tensor(make_tensor_arg(task_args.f), TensorArgType.INPUT) + sub_args.add_ref(_rehosted_ref(task_args, "f"), TensorArgType.INPUT) orch.submit_sub(callables.verify, sub_args) diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_group.py b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_group.py index 686fc5d56c..08685dbbd6 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_group.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_group.py @@ -19,7 +19,8 @@ from simpler.task_interface import ArgDirection as D from simpler.task_interface import TaskArgs, TensorArgType -from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, make_tensor_arg, scene_test +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.scene_test import _rehosted_ref_for KERNELS_BASE = "../../../../examples/a2a3/tensormap_and_ringbuffer/vector_example/kernels" @@ -28,27 +29,27 @@ def verify(args): """SubCallable — runs after group completes.""" -def _chip_args(in_a, in_b, out_f) -> TaskArgs: - """Build per-chip TaskArgs with INPUT/INPUT/OUTPUT_EXISTING tags.""" +def _chip_args(task_args, in_a, in_b, out_f) -> TaskArgs: + """Build per-chip TaskArgs (BufferRefs) with INPUT/INPUT/OUTPUT_EXISTING tags.""" a = TaskArgs() - a.add_tensor(make_tensor_arg(in_a), TensorArgType.INPUT) - a.add_tensor(make_tensor_arg(in_b), TensorArgType.INPUT) - a.add_tensor(make_tensor_arg(out_f), TensorArgType.OUTPUT_EXISTING) + a.add_ref(_rehosted_ref_for(task_args, in_a), TensorArgType.INPUT) + a.add_ref(_rehosted_ref_for(task_args, in_b), TensorArgType.INPUT) + a.add_ref(_rehosted_ref_for(task_args, out_f), TensorArgType.OUTPUT_EXISTING) return a def run_dag(orch, callables, task_args, config): """L3 orchestration: group of 2 chips → SubTask dependency.""" - args0 = _chip_args(task_args.a0, task_args.b0, task_args.f0) - args1 = _chip_args(task_args.a1, task_args.b1, task_args.f1) + args0 = _chip_args(task_args, task_args.a0, task_args.b0, task_args.f0) + args1 = _chip_args(task_args, task_args.a1, task_args.b1, task_args.f1) callables.keep(args0, args1) # prevent GC before drain orch.submit_next_level_group(callables.vector_kernel, [args0, args1], config, workers=[0, 1]) # SubTask depends on both group outputs (f0, f1) — tag both as INPUT. sub_args = TaskArgs() - sub_args.add_tensor(make_tensor_arg(task_args.f0), TensorArgType.INPUT) - sub_args.add_tensor(make_tensor_arg(task_args.f1), TensorArgType.INPUT) + sub_args.add_ref(_rehosted_ref_for(task_args, task_args.f0), TensorArgType.INPUT) + sub_args.add_ref(_rehosted_ref_for(task_args, task_args.f1), TensorArgType.INPUT) orch.submit_sub(callables.verify, sub_args) diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py index 4b8fbca969..f962d303d3 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py @@ -12,32 +12,33 @@ A host tensor created *after* the chip children are forked (lazily on the first ``Worker.run()``) is not visible to those children: the orch fn runs in the parent and carries a raw parent VA that is unmapped (or stale) in the child. -``Worker.create_host_buffer`` hands back born-shared memory already attached into -every chip child, so a tensor built over it with ``torch.frombuffer`` round-trips +``Worker.create_buffer`` allocates a POSIX-shm backing and hands back a +``BufferHandle``; naming it in a task arg with ``handle.ref(...)`` carries a +self-describing descriptor that the forked chip child materializes lazily on +first receipt (map-once), so a tensor built over ``handle.shm.buf`` round-trips with **no per-run copy** — the child reads and writes the same physical pages the parent sees. Covers the mechanism end-to-end (allocate a post-fork buffer, fill it in place, -run, read the result back). The host-side staging (born-shared bytes need no -copy; an in-range view is validated to fit) is unit-tested in -``tests/ut/py/test_worker/test_host_buffer_registration.py``. - -a2a3sim: ``create_host_buffer`` is pure host-side (POSIX shm + a control -broadcast to the forked chip children) with no platform branching, so the sim -backend exercises the full mechanism without needing a device. The -vector_example orchestration kernels exist only for a2a3. +run, read the result back). + +a2a3sim: ``create_buffer`` is pure host-side (POSIX shm; the descriptor rides in +the mailbox, no eager broadcast) with no platform branching, so the sim backend +exercises the full mechanism without needing a device. The vector_example +orchestration kernels exist only for a2a3. """ import torch from simpler.task_interface import ArgDirection as D -from simpler.task_interface import CallConfig, TaskArgs, TensorArgType +from simpler.task_interface import CallConfig, DataType, TaskArgs, TensorArgType -from simpler_setup import SceneTestCase, make_tensor_arg, scene_test +from simpler_setup import SceneTestCase, scene_test KERNELS_BASE = "../../../../examples/a2a3/tensormap_and_ringbuffer/vector_example/kernels" SIZE = 128 * 128 DTYPE = torch.float32 +_F32 = DataType.FLOAT32.value def _golden(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: @@ -45,12 +46,12 @@ def _golden(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return (s + 1) * (s + 2) + s -def _one_task_orch(chip_handle, a, b, out): +def _one_task_orch(chip_handle, ba, bb, bout): def orch_fn(orch, _args, cfg): ta = TaskArgs() - ta.add_tensor(make_tensor_arg(a), TensorArgType.INPUT) - ta.add_tensor(make_tensor_arg(b), TensorArgType.INPUT) - ta.add_tensor(make_tensor_arg(out), TensorArgType.OUTPUT_EXISTING) + ta.add_ref(ba.ref((SIZE,), _F32), TensorArgType.INPUT) + ta.add_ref(bb.ref((SIZE,), _F32), TensorArgType.INPUT) + ta.add_ref(bout.ref((SIZE,), _F32), TensorArgType.OUTPUT_EXISTING) orch.submit_next_level(chip_handle, ta, cfg, worker=0) return orch_fn @@ -98,39 +99,39 @@ class TestPostForkHostBufferZeroCopy(SceneTestCase): ] def test_run(self, st_worker): - """Zero-copy: buffers allocated AFTER the fork via ``create_host_buffer``, + """Zero-copy: buffers allocated AFTER the fork via ``create_buffer``, filled in place, run, and read back — all without a per-run copy. ``Worker.init()`` is eager, so the chip child is already forked when the - buffers are created; a born-shared ``create_host_buffer`` is the mapped - path a post-init host tensor takes to reach that child. + buffers are created; a post-init ``create_buffer`` reaches that child by + naming the handle as a ``BufferRef`` the child materializes lazily. """ worker = st_worker chip_handle = type(self)._st_chip_handles["vector"] nbytes = SIZE * DTYPE.itemsize # element count × dtype size, not a magic 4 - ba = worker.create_host_buffer(nbytes) - bb = worker.create_host_buffer(nbytes) - bout = worker.create_host_buffer(nbytes) + ba = worker.create_buffer(nbytes) + bb = worker.create_buffer(nbytes) + bout = worker.create_buffer(nbytes) a = b = out = None result = False try: - a = torch.frombuffer(ba.buffer, dtype=DTYPE, count=SIZE) - b = torch.frombuffer(bb.buffer, dtype=DTYPE, count=SIZE) - out = torch.frombuffer(bout.buffer, dtype=DTYPE, count=SIZE) + a = torch.frombuffer(ba.shm.buf, dtype=DTYPE, count=SIZE) + b = torch.frombuffer(bb.shm.buf, dtype=DTYPE, count=SIZE) + out = torch.frombuffer(bout.shm.buf, dtype=DTYPE, count=SIZE) a.fill_(5.0) # in place → lands directly in the child-visible shm b.fill_(7.0) out.zero_() - worker.run(_one_task_orch(chip_handle, a, b, out), args=None, config=CallConfig()) + worker.run(_one_task_orch(chip_handle, ba, bb, bout), args=None, config=CallConfig()) result = torch.allclose(out, _golden(a, b), rtol=self.RTOL, atol=self.ATOL) finally: # Drop the views before freeing so the shm releases promptly, and do it # in finally so a run()/assert failure above still cleans up all three # buffers instead of leaving live views warning on the first free. del a, b, out - worker.free_host_buffer(ba) - worker.free_host_buffer(bb) - worker.free_host_buffer(bout) + ba.close() + bb.close() + bout.close() assert result diff --git a/tests/st/aicore_op_timeout/test_aicore_op_timeout.py b/tests/st/aicore_op_timeout/test_aicore_op_timeout.py index f13b31a375..3718ac9269 100644 --- a/tests/st/aicore_op_timeout/test_aicore_op_timeout.py +++ b/tests/st/aicore_op_timeout/test_aicore_op_timeout.py @@ -21,7 +21,7 @@ import time import pytest -from simpler.task_interface import CallConfig, ChipCallable, ChipStorageTaskArgs, CoreCallable +from simpler.task_interface import CallConfig, ChipCallable, CoreCallable from simpler.worker import Worker from simpler_setup.elf_parser import extract_text_section @@ -95,7 +95,7 @@ def test_aicore_op_timeout_surfaces_as_runtime_error(st_platform, st_device_ids, # single-digit seconds and surfaces either the device classification or # a valid host fallback rather than deadlocking. with pytest.raises(RuntimeError, match=r"run failed with code (-100|507(046|018|000))"): - worker.run(handle, ChipStorageTaskArgs(), config) + worker.run(handle, None, config) elapsed = time.monotonic() - t0 # CI-tight env keeps the timeout chain short; default local values are diff --git a/tests/st/runtime_fatal_codes/test_runtime_fatal_codes.py b/tests/st/runtime_fatal_codes/test_runtime_fatal_codes.py index a71e7760f2..9fafc44963 100644 --- a/tests/st/runtime_fatal_codes/test_runtime_fatal_codes.py +++ b/tests/st/runtime_fatal_codes/test_runtime_fatal_codes.py @@ -31,7 +31,7 @@ import os import pytest -from simpler.task_interface import ArgDirection, CallConfig, ChipCallable, ChipStorageTaskArgs, CoreCallable +from simpler.task_interface import ArgDirection, CallConfig, ChipCallable, CoreCallable from simpler.worker import Worker from simpler_setup.elf_parser import extract_text_section @@ -298,7 +298,7 @@ def test_fatal_code_surfaces_on_sim(st_platform, st_device_ids, case_name, monke worker, handle, config = _make_worker(st_platform, int(st_device_ids[0]), case_name, monkeypatch) try: with pytest.raises(RuntimeError, match=rf"(run_runtime|run) failed with code -{case['code']}\b"): - worker.run(handle, ChipStorageTaskArgs(), config) + worker.run(handle, None, config) captured = capfd.readouterr() log = captured.err + captured.out assert case["marker"] in log, f"missing '{case['marker']}' in host log" @@ -322,7 +322,7 @@ def test_device_error_class_reaches_host_log(st_platform, st_device_ids, case_na # CANN 507xxx instead of the runtime's own -N; we only require that the # run fails. The point of the test is the device-classified host LOG. with pytest.raises(RuntimeError): - worker.run(handle, ChipStorageTaskArgs(), config) + worker.run(handle, None, config) captured = capfd.readouterr() log = captured.err + captured.out assert case["marker"] in log, f"device error class '{case['marker']}' not in host log" diff --git a/tests/st/task_timing/task_timing_a5hbg/test_task_timing_a5hbg.py b/tests/st/task_timing/task_timing_a5hbg/test_task_timing_a5hbg.py index c528c39e1b..b8e799637c 100644 --- a/tests/st/task_timing/task_timing_a5hbg/test_task_timing_a5hbg.py +++ b/tests/st/task_timing/task_timing_a5hbg/test_task_timing_a5hbg.py @@ -24,7 +24,15 @@ import re import pytest -from simpler.task_interface import ArgDirection, CallConfig, ChipCallable, ChipStorageTaskArgs, CoreCallable +from simpler.task_interface import ( + ArgDirection, + CallConfig, + ChipCallable, + CoreCallable, + DataType, + TaskArgs, + TensorArgType, +) from simpler.worker import Worker from simpler_setup.kernel_compiler import KernelCompiler @@ -92,12 +100,14 @@ def test_a5hbg_task_timing_slots_emit_markers(st_platform, st_device_ids, capfd) out = torch.zeros(_SIZE, dtype=torch.float32) # orch copies dev_out back here expected = a + 2 * b - args = ChipStorageTaskArgs() - from simpler_setup.torch_interop import make_tensor_arg # noqa: PLC0415 - - args.add_tensor(make_tensor_arg(a)) - args.add_tensor(make_tensor_arg(b)) - args.add_tensor(make_tensor_arg(out)) + # L2 runs in-process (no fork), so any host tensor's VA is valid to the materializer; name each + # as a FORK_SHM BufferRef via make_ref_arg. a, b are inputs; out is where the hbg runtime's D2H + # lands (in-process, so the parent sees it — no share_memory_ needed at L2). + args = TaskArgs() + _f32 = DataType.FLOAT32.value + args.add_ref(worker.make_ref_arg(a, shapes=(_SIZE,), dtype=_f32), TensorArgType.INPUT) + args.add_ref(worker.make_ref_arg(b, shapes=(_SIZE,), dtype=_f32), TensorArgType.INPUT) + args.add_ref(worker.make_ref_arg(out, shapes=(_SIZE,), dtype=_f32), TensorArgType.OUTPUT_EXISTING) config = CallConfig() config.enable_l2_swimlane = False # slots must work with swimlane OFF 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 fa0f3de089..d8451fa985 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 @@ -31,10 +31,10 @@ ArgDirection, CallConfig, ChipCallable, - ChipStorageTaskArgs, CoreCallable, DataType, - Tensor, + TaskArgs, + TensorArgType, ) from simpler.worker import Worker @@ -137,16 +137,16 @@ def _drive( host_b = torch.full((N_ROWS, N_COLS), 2.0, dtype=torch.float32) expected = host_a + b_multiplier * host_b - dev_a = worker.malloc(NBYTES) - dev_b = worker.malloc(NBYTES) - dev_out = worker.malloc(NBYTES) - worker.copy_to(dev_a, host_a.data_ptr(), NBYTES) - worker.copy_to(dev_b, host_b.data_ptr(), NBYTES) + a_h = worker.malloc(NBYTES) + b_h = worker.malloc(NBYTES) + out_h = worker.malloc(NBYTES) + worker.copy_to(a_h, host_a) + worker.copy_to(b_h, host_b) - 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 = TaskArgs() + args.add_ref(a_h.ref(shapes=(N_ROWS, N_COLS), dtype=DataType.FLOAT32.value), TensorArgType.INPUT) + args.add_ref(b_h.ref(shapes=(N_ROWS, N_COLS), dtype=DataType.FLOAT32.value), TensorArgType.INPUT) + args.add_ref(out_h.ref(shapes=(N_ROWS, N_COLS), dtype=DataType.FLOAT32.value), TensorArgType.OUTPUT_EXISTING) config = CallConfig() config.enable_l2_swimlane = False # slots must work with swimlane OFF @@ -154,10 +154,10 @@ def _drive( assert worker.run(chip_handle, args, config) is None host_out = torch.zeros(N_ROWS, N_COLS, dtype=torch.float32) - worker.copy_from(host_out.data_ptr(), dev_out, NBYTES) - worker.free(dev_a) - worker.free(dev_b) - worker.free(dev_out) + worker.copy_from(host_out, out_h) + worker.free(a_h) + worker.free(b_h) + worker.free(out_h) assert torch.allclose(host_out, expected, rtol=1e-5, atol=1e-5), ( f"{func_name} output diverged; max |expected - out| = " f"{float(torch.max(torch.abs(host_out - expected))):.3e}" @@ -332,14 +332,15 @@ def test_mix_task_aggregates_across_subtasks(st_platform, st_device_ids, capfd): nb = _TILE_ELEMS * 4 bufs = {n: worker.malloc(nb) for n in "ABCDEFGHI"} for name, t in {"A": A.flatten(), "B": B.flatten(), "D": D, "E": E, "G": G, "H": H}.items(): - worker.copy_to(bufs[name], t.contiguous().data_ptr(), nb) + worker.copy_to(bufs[name], t.contiguous()) - 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)) - for n in "DEFGHI": - args.add_tensor(Tensor.make(bufs[n], (_TILE_ELEMS,), DataType.FLOAT32)) + # Arg order matches the ChipCallable signature: A,B->C (matmul), D,E->F (add), G,H->I (mul). + _shapes = {"A": (_MATMUL_SIZE, _MATMUL_SIZE), "B": (_MATMUL_SIZE, _MATMUL_SIZE)} + _outputs = {"C", "F", "I"} + args = TaskArgs() + for n in "ABCDEFGHI": + tag = TensorArgType.OUTPUT_EXISTING if n in _outputs else TensorArgType.INPUT + args.add_ref(bufs[n].ref(shapes=_shapes.get(n, (_TILE_ELEMS,)), dtype=DataType.FLOAT32.value), tag) config = CallConfig() config.enable_l2_swimlane = False @@ -348,9 +349,9 @@ def test_mix_task_aggregates_across_subtasks(st_platform, st_device_ids, capfd): out_C = torch.zeros(_TILE_ELEMS, dtype=torch.float32) out_F = torch.zeros(_TILE_ELEMS, dtype=torch.float32) out_I = torch.zeros(_TILE_ELEMS, dtype=torch.float32) - worker.copy_from(out_C.data_ptr(), bufs["C"], nb) - worker.copy_from(out_F.data_ptr(), bufs["F"], nb) - worker.copy_from(out_I.data_ptr(), bufs["I"], nb) + worker.copy_from(out_C, bufs["C"]) + worker.copy_from(out_F, bufs["F"]) + worker.copy_from(out_I, bufs["I"]) for b in bufs.values(): worker.free(b) # The AIV0/AIV1 subtasks are element-wise (layout-agnostic): their correct diff --git a/tests/st/worker/collectives/_helpers.py b/tests/st/worker/collectives/_helpers.py index 4a768340b8..cfc5e77ee9 100644 --- a/tests/st/worker/collectives/_helpers.py +++ b/tests/st/worker/collectives/_helpers.py @@ -18,11 +18,12 @@ import ctypes import torch -from simpler.task_interface import CommBufferSpec, DataType, TaskArgs, Tensor, TensorArgType +from simpler.task_interface import CommBufferSpec, DataType, TaskArgs, TensorArgType from simpler_setup import Tensor as STensor -from simpler_setup.scene_test import TaskArgsBuilder -from simpler_setup.torch_interop import make_tensor_arg +from simpler_setup.scene_test import TaskArgsBuilder, _rehosted_ref_for + +_F32 = DataType.FLOAT32.value # --------------------------------------------------------------------------- # Allreduce constants (must match kernel COUNT) @@ -123,17 +124,11 @@ def allreduce_orch_fn(orch, callables, task_args, config): for i in range(nranks): domain = handle[i] chip_args = TaskArgs() - 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( - data=domain.buffer_ptrs["scratch"], - shapes=(float_elems,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INOUT, + chip_args.add_ref(_rehosted_ref_for(task_args, getattr(task_args, f"in_{i}")), TensorArgType.INPUT) + chip_args.add_ref( + _rehosted_ref_for(task_args, getattr(task_args, f"out_{i}")), TensorArgType.OUTPUT_EXISTING ) + chip_args.add_ref(domain.buffers["scratch"].ref((float_elems,), _F32), TensorArgType.INOUT) chip_args.add_scalar(domain.domain_size) chip_args.add_scalar(domain.device_ctx) orch.submit_next_level(chip, chip_args, config, worker=i) @@ -190,17 +185,11 @@ def generic_collective_orch_fn( for i in range(nranks): domain = handle[i] chip_args = TaskArgs() - 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( - data=domain.buffer_ptrs["scratch"], - shapes=(float_elems,), - dtype=DataType.FLOAT32, - child_memory=True, - ), - TensorArgType.INOUT, + chip_args.add_ref(_rehosted_ref_for(task_args, getattr(task_args, f"in_{i}")), TensorArgType.INPUT) + chip_args.add_ref( + _rehosted_ref_for(task_args, getattr(task_args, f"out_{i}")), TensorArgType.OUTPUT_EXISTING ) + chip_args.add_ref(domain.buffers["scratch"].ref((float_elems,), _F32), TensorArgType.INOUT) chip_args.add_scalar(domain.domain_size) for s in extras: chip_args.add_scalar(s) diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index ea906b8fc7..c29a6fce69 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -344,6 +344,7 @@ add_hierarchical_test(test_pipeline_contract hierarchical/test_pipeline_contract # --------------------------------------------------------------------------- # Types / task_interface tests (src/common/task_interface/) # --------------------------------------------------------------------------- +add_task_interface_test(test_buffer_handle types/test_buffer_handle.cpp) add_task_interface_test(test_child_memory types/test_child_memory.cpp) add_task_interface_test(test_chip_max_tensor_args types/test_chip_max_tensor_args.cpp) add_task_interface_test(test_call_config types/test_call_config.cpp) diff --git a/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp b/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp index 30dba36523..3dd5a3243b 100644 --- a/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp +++ b/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp @@ -201,7 +201,10 @@ HostApi make_host_api(bool with_temporary_buffer = true) { Tensor make_tensor(std::vector &storage, bool child_memory = false) { Tensor tensor; uint32_t shape[1] = {static_cast(storage.size())}; - tensor.init_external(storage.data(), storage.size(), shape, 1, DataType::UINT8, 0, false, child_memory ? 1 : 0); + tensor.init_external( + storage.data(), storage.size(), shape, 1, DataType::UINT8, 0, false, + child_memory ? AddressSpace::DEVICE : AddressSpace::HOST + ); return tensor; } diff --git a/tests/ut/cpp/hierarchical/test_orchestrator.cpp b/tests/ut/cpp/hierarchical/test_orchestrator.cpp index 519e170dbd..1136f04eb4 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -62,15 +62,45 @@ struct OrchestratorFixture : public ::testing::Test { return c; } - // Helper: build a TaskArgs whose only tensor has the given (data, tag). - static TaskArgs single_tensor_args(uint64_t data_ptr, TensorArgType tag) { + // A canonical identity keyed by `buffer_id` (fixed owner nonce/path), so two args with the same + // buffer_id collide into one dependency — the successor of the former buffer-address key. + static CanonicalIdentity identity_for(uint64_t buffer_id) { + CanonicalIdentity id{}; + id.buffer_id = buffer_id; + return id; + } + + // The local dependency key the orchestrator derives for a ref over `buffer_id`. + static TensorKey ref_key(uint64_t buffer_id) { + return TensorKey::local_host(static_cast(CanonicalIdentityHash{}(identity_for(buffer_id)))); + } + + // Helper: a TaskArgs whose only ref is a local POSIX_SHM backing keyed by `buffer_id`. + static TaskArgs single_ref_args(uint64_t buffer_id, TensorArgType tag) { + TaskArgs a; + BufferRef r{}; + r.handle.backend_kind = static_cast(BackendKind::POSIX_SHM); + r.handle.nbytes = 1; // has a local backing (not a placeholder) -> tracked by infer_deps + r.handle.identity = identity_for(buffer_id); + r.ndims = 1; + r.shapes[0] = 1; + r.strides[0] = 1; + r.dtype = DataType::UINT8; + a.add_tensor(r, tag); + return a; + } + + // Helper: a TaskArgs whose only ref is a REMOTE_SIDECAR placeholder (no local backing; the remote + // descriptor rides in the paired RemoteTaskArgsSidecar). + static TaskArgs remote_placeholder_args(TensorArgType tag) { TaskArgs a; - Tensor t{}; - t.buffer.addr = data_ptr; - t.ndims = 1; - t.shapes[0] = 1; - t.dtype = DataType::UINT8; - a.add_tensor(t, tag); + BufferRef r{}; + r.handle.backend_kind = static_cast(BackendKind::REMOTE_SIDECAR); + r.ndims = 1; + r.shapes[0] = 1; + r.strides[0] = 1; + r.dtype = DataType::UINT8; + a.add_tensor(r, tag); return a; } }; @@ -80,7 +110,7 @@ struct OrchestratorFixture : public ::testing::Test { // --------------------------------------------------------------------------- TEST_F(OrchestratorFixture, IndependentTaskIsImmediatelyReady) { - auto a = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); + auto a = single_ref_args(0xCAFE, TensorArgType::OUTPUT); auto res = orch.submit_next_level(C(42), a, cfg, 0); EXPECT_NE(res.task_slot, INVALID_SLOT); @@ -91,7 +121,7 @@ TEST_F(OrchestratorFixture, IndependentTaskIsImmediatelyReady) { } TEST_F(OrchestratorFixture, NextLevelTargetsMustBeValidAndDistinct) { - auto args = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); + auto args = single_ref_args(0xCAFE, TensorArgType::OUTPUT); EXPECT_THROW((void)orch.submit_next_level(C(42), args, cfg, -1), std::invalid_argument); EXPECT_THROW((void)orch.submit_next_level_group(C(42), {args, args}, cfg, {0}), std::invalid_argument); EXPECT_THROW((void)orch.submit_next_level_group(C(42), {args, args}, cfg, {0, 0}), std::invalid_argument); @@ -99,13 +129,13 @@ TEST_F(OrchestratorFixture, NextLevelTargetsMustBeValidAndDistinct) { TEST_F(OrchestratorFixture, DependentTaskIsPending) { // Task A produces an OUTPUT at key 0xBEEF - auto args_a = single_tensor_args(0xBEEF, TensorArgType::OUTPUT); + auto args_a = single_ref_args(0xBEEF, TensorArgType::OUTPUT); auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot a_slot; rq.try_pop(a_slot); // Task B reads INPUT at the same key -- depends on A - auto args_b = single_tensor_args(0xBEEF, TensorArgType::INPUT); + auto args_b = single_ref_args(0xBEEF, TensorArgType::INPUT); auto b = orch.submit_next_level(C(42), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::PENDING); EXPECT_EQ(S(b.task_slot).fanin_count, 1); @@ -115,14 +145,14 @@ TEST_F(OrchestratorFixture, DependentTaskIsPending) { } TEST_F(OrchestratorFixture, GroupDuplicateInputsKeepOneProducer) { - auto producer_args = single_tensor_args(0xBEEF, TensorArgType::OUTPUT); + auto producer_args = single_ref_args(0xBEEF, TensorArgType::OUTPUT); auto producer = orch.submit_next_level(C(42), producer_args, cfg, 0); TaskSlot ready; ASSERT_TRUE(rq.try_pop(ready)); ASSERT_EQ(ready, producer.task_slot); - TaskArgs input0 = single_tensor_args(0xBEEF, TensorArgType::INPUT); - TaskArgs input1 = single_tensor_args(0xBEEF, TensorArgType::INPUT); + TaskArgs input0 = single_ref_args(0xBEEF, TensorArgType::INPUT); + TaskArgs input1 = single_ref_args(0xBEEF, TensorArgType::INPUT); auto consumer = orch.submit_next_level_group(C(43), {input0, input1}, cfg, {0, 1}); EXPECT_EQ(S(consumer.task_slot).state.load(), TaskState::PENDING); @@ -134,7 +164,7 @@ TEST_F(OrchestratorFixture, GroupDuplicateInputsKeepOneProducer) { TEST_F(OrchestratorFixture, SubmitAfterFailedProducerPoisonsConsumer) { orch.scope_begin(); - auto args_a = single_tensor_args(0xD00D, TensorArgType::OUTPUT); + auto args_a = single_ref_args(0xD00D, TensorArgType::OUTPUT); auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot ready; ASSERT_TRUE(rq.try_pop(ready)); @@ -145,7 +175,7 @@ TEST_F(OrchestratorFixture, SubmitAfterFailedProducerPoisonsConsumer) { producer.state.store(TaskState::FAILED, std::memory_order_release); producer.fanout_released.store(1, std::memory_order_release); - auto args_b = single_tensor_args(0xD00D, TensorArgType::INPUT); + auto args_b = single_ref_args(0xD00D, TensorArgType::INPUT); auto b = orch.submit_next_level(C(43), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::FAILED); @@ -163,26 +193,26 @@ TEST_F(OrchestratorFixture, SubmitAfterFailedProducerPoisonsConsumer) { } TEST_F(OrchestratorFixture, TensorMapTracksProducer) { - auto args_a = single_tensor_args(0x1234, TensorArgType::OUTPUT); + auto args_a = single_ref_args(0x1234, TensorArgType::OUTPUT); auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot drain_slot; rq.try_pop(drain_slot); - EXPECT_EQ(tm.lookup(run_id, TensorKey{0x1234, -1}), a.task_slot); + EXPECT_EQ(tm.lookup(run_id, ref_key(0x1234)), a.task_slot); } TEST_F(OrchestratorFixture, OnConsumedCleansUpTensorMap) { - auto args_a = single_tensor_args(0x42, TensorArgType::OUTPUT); + auto args_a = single_ref_args(0x42, TensorArgType::OUTPUT); auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot slot; rq.try_pop(slot); - EXPECT_EQ(tm.lookup(run_id, TensorKey{0x42, -1}), slot); + EXPECT_EQ(tm.lookup(run_id, ref_key(0x42)), slot); S(slot).state.store(TaskState::COMPLETED, std::memory_order_relaxed); orch.on_consumed(slot); - EXPECT_EQ(tm.lookup(run_id, TensorKey{0x42, -1}), INVALID_SLOT); + EXPECT_EQ(tm.lookup(run_id, ref_key(0x42)), INVALID_SLOT); EXPECT_EQ(S(slot).state.load(), TaskState::CONSUMED); } @@ -193,22 +223,22 @@ TEST_F(OrchestratorFixture, ConsumingASupersededProducerKeepsTheNewerMapping) { // second's mapping intact -- otherwise a later INPUT reader looks up an // empty slot, infers no dependency, and can be dispatched before the // second writer has written the buffer. - auto args_a = single_tensor_args(0x5150, TensorArgType::OUTPUT); + auto args_a = single_ref_args(0x5150, TensorArgType::OUTPUT); auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot drain; ASSERT_TRUE(rq.try_pop(drain)); - auto args_b = single_tensor_args(0x5150, TensorArgType::OUTPUT); + auto args_b = single_ref_args(0x5150, TensorArgType::OUTPUT); auto b = orch.submit_next_level(C(43), args_b, cfg, 0); ASSERT_TRUE(rq.try_pop(drain)); - ASSERT_EQ(tm.lookup(run_id, TensorKey{0x5150, -1}), b.task_slot); + ASSERT_EQ(tm.lookup(run_id, ref_key(0x5150)), b.task_slot); S(a.task_slot).state.store(TaskState::COMPLETED, std::memory_order_relaxed); ASSERT_TRUE(orch.on_consumed(a.task_slot)); - EXPECT_EQ(tm.lookup(run_id, TensorKey{0x5150, -1}), b.task_slot); + EXPECT_EQ(tm.lookup(run_id, ref_key(0x5150)), b.task_slot); - auto args_c = single_tensor_args(0x5150, TensorArgType::INPUT); + auto args_c = single_ref_args(0x5150, TensorArgType::INPUT); auto c = orch.submit_next_level(C(44), args_c, cfg, 0); EXPECT_EQ(S(c.task_slot).state.load(), TaskState::PENDING); EXPECT_EQ(S(c.task_slot).fanin_count, 1); @@ -218,7 +248,7 @@ TEST_F(OrchestratorFixture, ConsumingASupersededProducerKeepsTheNewerMapping) { TEST_F(OrchestratorFixture, ScopeRegistersAndReleasesRef) { orch.scope_begin(); - auto args_a = single_tensor_args(0x77, TensorArgType::OUTPUT); + auto args_a = single_ref_args(0x77, TensorArgType::OUTPUT); auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot slot; rq.try_pop(slot); @@ -242,21 +272,21 @@ TEST_F(OrchestratorFixture, ScopeRegistersAndReleasesRef) { TEST_F(OrchestratorFixture, NoDepTagSkipsDependencyTracking) { // OUTPUT-tagged input registers a producer - auto args_a = single_tensor_args(0xAAAA, TensorArgType::OUTPUT); + auto args_a = single_ref_args(0xAAAA, TensorArgType::OUTPUT); auto a = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot drain_slot; rq.try_pop(drain_slot); // Second task references same key but tagged NO_DEP -- should be independent - auto args_b = single_tensor_args(0xAAAA, TensorArgType::NO_DEP); + auto args_b = single_ref_args(0xAAAA, TensorArgType::NO_DEP); auto b = orch.submit_next_level(C(42), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::READY); EXPECT_EQ(S(b.task_slot).fanin_count, 0); } TEST_F(OrchestratorFixture, GroupTaskStoresArgsListPerMember) { - TaskArgs a0 = single_tensor_args(0xA0, TensorArgType::OUTPUT); - TaskArgs a1 = single_tensor_args(0xA1, TensorArgType::OUTPUT); + TaskArgs a0 = single_ref_args(0xA0, TensorArgType::OUTPUT); + TaskArgs a1 = single_ref_args(0xA1, TensorArgType::OUTPUT); auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); EXPECT_NE(res.task_slot, INVALID_SLOT); @@ -265,58 +295,26 @@ TEST_F(OrchestratorFixture, GroupTaskStoresArgsListPerMember) { EXPECT_EQ(S(res.task_slot).task_args_list.size(), 2u); // args_view(i) yields each member's distinct tensor list. - EXPECT_EQ(S(res.task_slot).args_view(0).tensors(0).buffer.addr, 0xA0u); - EXPECT_EQ(S(res.task_slot).args_view(1).tensors(0).buffer.addr, 0xA1u); + EXPECT_EQ(S(res.task_slot).args(0).tensor(0).handle.identity.buffer_id, 0xA0u); + EXPECT_EQ(S(res.task_slot).args(1).tensor(0).handle.identity.buffer_id, 0xA1u); // Both keys registered as producers for the group slot. - EXPECT_EQ(tm.lookup(run_id, TensorKey{0xA0, -1}), res.task_slot); - EXPECT_EQ(tm.lookup(run_id, TensorKey{0xA1, -1}), res.task_slot); + EXPECT_EQ(tm.lookup(run_id, ref_key(0xA0)), res.task_slot); + EXPECT_EQ(tm.lookup(run_id, ref_key(0xA1)), res.task_slot); } TEST_F(OrchestratorFixture, SingleTaskStoresTaskArgsDirectly) { - TaskArgs a0 = single_tensor_args(0xC0, TensorArgType::OUTPUT); + TaskArgs a0 = single_ref_args(0xC0, TensorArgType::OUTPUT); auto res = orch.submit_next_level(C(42), a0, cfg, 0); ASSERT_NE(res.task_slot, INVALID_SLOT); EXPECT_FALSE(S(res.task_slot).is_group()); EXPECT_EQ(S(res.task_slot).group_size(), 1); EXPECT_EQ(S(res.task_slot).task_args.tensor_count(), 1); - EXPECT_EQ(S(res.task_slot).args_view(0).tensors(0).buffer.addr, 0xC0u); + EXPECT_EQ(S(res.task_slot).args(0).tensor(0).handle.identity.buffer_id, 0xC0u); } -TEST_F(OrchestratorFixture, OutputAutoAllocsFromHeapRing) { - // An OUTPUT tensor submitted with `data == 0` is auto-allocated from - // the HeapRing: the slot's task_args tensor ends up with a non-zero - // data pointer that falls inside the allocator's mmap'd region, and - // the TensorMap routes that pointer to the slot. - TaskArgs args; - Tensor t{}; - t.buffer.addr = 0; - t.ndims = 1; - t.shapes[0] = 1024; // 1024 * 1 byte = 1024, one aligned slab - t.dtype = DataType::UINT8; - args.add_tensor(t, TensorArgType::OUTPUT); - - auto res = orch.submit_next_level(C(42), args, cfg, 0); - ASSERT_NE(res.task_slot, INVALID_SLOT); - - uint64_t data = S(res.task_slot).task_args.tensor(0).buffer.addr; - ASSERT_NE(data, 0u); - uintptr_t base = reinterpret_cast(allocator.heap_base(0)); - EXPECT_GE(data, base); - EXPECT_LT(data, base + allocator.heap_size(0)); - EXPECT_EQ(data % HEAP_ALIGN, 0u); - - EXPECT_EQ(tm.lookup(run_id, TensorKey{data, -1}), res.task_slot); -} - -TEST_F(OrchestratorFixture, RemoteOutputSidecarSkipsLocalAutoAllocAndRegistersRemoteKey) { - TaskArgs args; - Tensor t{}; - t.buffer.addr = 0; - t.ndims = 1; - t.shapes[0] = 1024; - t.dtype = DataType::UINT8; - args.add_tensor(t, TensorArgType::OUTPUT); +TEST_F(OrchestratorFixture, RemoteOutputSidecarRegistersRemoteKey) { + TaskArgs args = remote_placeholder_args(TensorArgType::OUTPUT); RemoteTaskArgsSidecar sidecar; sidecar.tensors.resize(1); @@ -330,14 +328,17 @@ TEST_F(OrchestratorFixture, RemoteOutputSidecarSkipsLocalAutoAllocAndRegistersRe auto res = orch.submit_next_level(C(42), args, cfg, 3, {3}, sidecar); ASSERT_NE(res.task_slot, INVALID_SLOT); - EXPECT_EQ(S(res.task_slot).task_args.tensor(0).buffer.addr, 0u); + // No local backing: the arg is a REMOTE_SIDECAR placeholder, routed by the sidecar's remote key. + EXPECT_EQ( + S(res.task_slot).task_args.tensor(0).handle.backend_kind, static_cast(BackendKind::REMOTE_SIDECAR) + ); TensorKey key = TensorKey::remote_buffer(TensorAddressKind::REMOTE_BUFFER, 3, 9, 2, 64); EXPECT_EQ(tm.lookup(run_id, key), res.task_slot); } TEST_F(OrchestratorFixture, RemoteBarePayloadFailsBeforeSlotCommit) { - TaskArgs args = single_tensor_args(0x1234, TensorArgType::INPUT); + TaskArgs args = single_ref_args(0x1234, TensorArgType::INPUT); RemoteTaskArgsSidecar sidecar; sidecar.tensors.resize(1); @@ -346,13 +347,7 @@ TEST_F(OrchestratorFixture, RemoteBarePayloadFailsBeforeSlotCommit) { } TEST_F(OrchestratorFixture, RemoteSidecarRejectsNonOwnerEligibleEndpointWithoutImport) { - TaskArgs args; - Tensor t{}; - t.buffer.addr = 0; - t.ndims = 1; - t.shapes[0] = 1; - t.dtype = DataType::UINT8; - args.add_tensor(t, TensorArgType::INPUT); + TaskArgs args = remote_placeholder_args(TensorArgType::INPUT); RemoteTaskArgsSidecar sidecar; sidecar.tensors.resize(1); @@ -367,13 +362,7 @@ TEST_F(OrchestratorFixture, RemoteSidecarRejectsNonOwnerEligibleEndpointWithoutI } TEST_F(OrchestratorFixture, RemoteInputSidecarUsesRemoteTensorMapKey) { - TaskArgs output_args; - Tensor out{}; - out.buffer.addr = 0; - out.ndims = 1; - out.shapes[0] = 1; - out.dtype = DataType::UINT8; - output_args.add_tensor(out, TensorArgType::OUTPUT); + TaskArgs output_args = remote_placeholder_args(TensorArgType::OUTPUT); RemoteTaskArgsSidecar output_sidecar; output_sidecar.tensors.resize(1); @@ -389,9 +378,7 @@ TEST_F(OrchestratorFixture, RemoteInputSidecarUsesRemoteTensorMapKey) { ASSERT_TRUE(rq_next_level.try_pop_single(3, ready)); ASSERT_EQ(ready, producer.task_slot); - TaskArgs input_args; - Tensor in = out; - input_args.add_tensor(in, TensorArgType::INPUT); + TaskArgs input_args = remote_placeholder_args(TensorArgType::INPUT); auto consumer = orch.submit_next_level(C(43), input_args, cfg, 3, {3}, output_sidecar); EXPECT_EQ(S(consumer.task_slot).state.load(), TaskState::PENDING); @@ -406,7 +393,7 @@ TEST_F(OrchestratorFixture, InoutWiresCreatorAsFanin) { // INPUT / INOUT do tensor_map.lookup. Users who want a WaW dep on // the alloc-slot (so its HeapRing slab stays live while they write) // must tag the buffer INOUT. - auto creator_args = single_tensor_args(0xFEED, TensorArgType::OUTPUT); + auto creator_args = single_ref_args(0xFEED, TensorArgType::OUTPUT); auto creator = orch.submit_next_level(C(42), creator_args, cfg, 0); TaskSlot drain; rq.try_pop(drain); @@ -414,13 +401,13 @@ TEST_F(OrchestratorFixture, InoutWiresCreatorAsFanin) { // path (COMPLETED producer with non-zero fanout). S(creator.task_slot).state.store(TaskState::COMPLETED, std::memory_order_relaxed); - auto writer_args = single_tensor_args(0xFEED, TensorArgType::INOUT); + auto writer_args = single_ref_args(0xFEED, TensorArgType::INOUT); auto writer = orch.submit_next_level(C(42), writer_args, cfg, 0); TaskSlot writer_slot; rq.try_pop(writer_slot); // TensorMap now points at the new writer. - EXPECT_EQ(tm.lookup(run_id, TensorKey{0xFEED, -1}), writer.task_slot); + EXPECT_EQ(tm.lookup(run_id, ref_key(0xFEED)), writer.task_slot); // Writer has the creator recorded as a fanin producer (via INOUT // lookup) but no *live* fanin since the creator is already COMPLETED. EXPECT_EQ(S(writer.task_slot).fanin_count, 0); @@ -445,16 +432,16 @@ TEST_F(OrchestratorFixture, OutputAndOutputExistingAreInsertOnly) { TensorArgType tag; }; for (Case c : {Case{0xABCD, TensorArgType::OUTPUT}, Case{0xBEEF, TensorArgType::OUTPUT_EXISTING}}) { - auto prior_args = single_tensor_args(c.key, TensorArgType::OUTPUT); + auto prior_args = single_ref_args(c.key, TensorArgType::OUTPUT); auto prior = orch.submit_next_level(C(42), prior_args, cfg, 0); TaskSlot drain; rq.try_pop(drain); S(prior.task_slot).state.store(TaskState::COMPLETED, std::memory_order_relaxed); - auto writer_args = single_tensor_args(c.key, c.tag); + auto writer_args = single_ref_args(c.key, c.tag); auto writer = orch.submit_next_level(C(42), writer_args, cfg, 0); - EXPECT_EQ(tm.lookup(run_id, TensorKey{c.key, -1}), writer.task_slot); + EXPECT_EQ(tm.lookup(run_id, ref_key(c.key)), writer.task_slot); EXPECT_EQ(S(writer.task_slot).fanin_count, 0); EXPECT_TRUE(S(writer.task_slot).fanin_producers.empty()) << "tag=" << static_cast(c.tag); { @@ -483,7 +470,7 @@ TEST_F(OrchestratorFixture, TimedWaitCanRetryAfterTimeout) { } TEST_F(OrchestratorFixture, OneTaskRunCompletesAfterConsumption) { - auto result = orch.submit_next_level(C(80), single_tensor_args(0x8000, TensorArgType::OUTPUT), cfg, 0); + auto result = orch.submit_next_level(C(80), single_ref_args(0x8000, TensorArgType::OUTPUT), cfg, 0); EXPECT_EQ(S(result.task_slot).run_id, run_id); orch.close_run_submission(run_id); @@ -497,7 +484,7 @@ TEST_F(OrchestratorFixture, OneTaskRunCompletesAfterConsumption) { } TEST_F(OrchestratorFixture, SequentialRunsHaveDistinctIdsAndErrorsDoNotLeak) { - auto failed_task = orch.submit_next_level(C(81), single_tensor_args(0x8100, TensorArgType::OUTPUT), cfg, 0); + auto failed_task = orch.submit_next_level(C(81), single_ref_args(0x8100, TensorArgType::OUTPUT), cfg, 0); orch.report_task_error(failed_task.task_slot, "run one failed"); S(failed_task.task_slot).failure_message = "run one failed"; S(failed_task.task_slot).state.store(TaskState::FAILED, std::memory_order_release); @@ -524,7 +511,7 @@ TEST_F(OrchestratorFixture, FailedSubmissionCompletesWithoutEnteringDeviceWork) } TEST_F(OrchestratorFixture, ReleasingRunDoesNotResetSlotsOwnedByAnotherRegisteredRun) { - auto first_task = orch.submit_next_level(C(82), single_tensor_args(0x8200, TensorArgType::OUTPUT), cfg, 0); + auto first_task = orch.submit_next_level(C(82), single_ref_args(0x8200, TensorArgType::OUTPUT), cfg, 0); S(first_task.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); EXPECT_TRUE(orch.on_consumed(first_task.task_slot)); orch.close_run_submission(run_id); @@ -550,7 +537,7 @@ TEST_F(OrchestratorFixture, ReportingErrorForAReleasedRunIsIgnored) { orch.release_run(released); run_id = orch.begin_run(); - auto live = orch.submit_next_level(C(83), single_tensor_args(0x8300, TensorArgType::OUTPUT), cfg, 0); + auto live = orch.submit_next_level(C(83), single_ref_args(0x8300, TensorArgType::OUTPUT), cfg, 0); S(live.task_slot).run_id = released; EXPECT_NO_THROW(orch.report_task_error(live.task_slot, "late endpoint failure")); @@ -559,7 +546,7 @@ TEST_F(OrchestratorFixture, ReportingErrorForAReleasedRunIsIgnored) { // on_consumed runs on the scheduler thread and reaches the same accounting. TEST_F(OrchestratorFixture, ConsumingASlotOwnedByAReleasedRunIsIgnored) { - auto task = orch.submit_next_level(C(84), single_tensor_args(0x8400, TensorArgType::OUTPUT), cfg, 0); + auto task = orch.submit_next_level(C(84), single_ref_args(0x8400, TensorArgType::OUTPUT), cfg, 0); S(task.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); EXPECT_TRUE(orch.on_consumed(task.task_slot)); orch.close_run_submission(run_id); @@ -568,7 +555,7 @@ TEST_F(OrchestratorFixture, ConsumingASlotOwnedByAReleasedRunIsIgnored) { orch.release_run(released); run_id = orch.begin_run(); - auto live = orch.submit_next_level(C(85), single_tensor_args(0x8500, TensorArgType::OUTPUT), cfg, 0); + auto live = orch.submit_next_level(C(85), single_ref_args(0x8500, TensorArgType::OUTPUT), cfg, 0); orch.close_run_submission(run_id); S(live.task_slot).run_id = released; @@ -582,7 +569,7 @@ TEST_F(OrchestratorFixture, ConsumingASlotOwnedByAReleasedRunIsIgnored) { // Failing a run is the recovery path: refusing a run whose submission already // closed would leave the caller's fence wait blocked forever. TEST_F(OrchestratorFixture, FailingAnAlreadyClosedRunStillResolvesTheFence) { - auto task = orch.submit_next_level(C(86), single_tensor_args(0x8600, TensorArgType::OUTPUT), cfg, 0); + auto task = orch.submit_next_level(C(86), single_ref_args(0x8600, TensorArgType::OUTPUT), cfg, 0); orch.close_run_submission(run_id); EXPECT_FALSE(orch.run_done(run_id)); diff --git a/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp b/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp index e5eddeb752..a3af6a7425 100644 --- a/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp +++ b/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp @@ -287,12 +287,15 @@ TaskArgs scalar_args() { TaskArgs bare_pointer_args() { TaskArgs args; - Tensor tensor{}; - tensor.buffer.addr = 0x1234; - tensor.ndims = 1; - tensor.shapes[0] = 1; - tensor.dtype = DataType::UINT8; - args.add_tensor(tensor, TensorArgType::INPUT); + BufferRef ref{}; + ref.handle.backend_kind = static_cast(BackendKind::POSIX_SHM); + ref.handle.nbytes = 1; // a local backing without a remote sidecar -> rejected by the remote endpoint + ref.handle.identity.buffer_id = 0x1234; + ref.ndims = 1; + ref.shapes[0] = 1; + ref.strides[0] = 1; + ref.dtype = DataType::UINT8; + args.add_tensor(ref, TensorArgType::INPUT); return args; } @@ -536,7 +539,7 @@ TEST(RemoteEndpoint, BareHostPointerWithoutSidecarIsEndpointFailure) { WorkerCompletion completion = endpoint.run(&ring, dispatch); EXPECT_EQ(completion.outcome, EndpointOutcome::ENDPOINT_FAILURE); - EXPECT_NE(completion.error_message.find("bare host pointer"), std::string::npos); + EXPECT_NE(completion.error_message.find("local backing"), std::string::npos); EXPECT_TRUE(transport->last_frame.empty()); ring.shutdown(); } diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 395d0137cc..f3ab7bb809 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -150,15 +150,17 @@ struct MockMailboxWorker { MailboxState s = read_state(); if (s == MailboxState::TASK_READY) { uint8_t callable_hash0 = static_cast(mailbox[MAILBOX_OFF_TASK_CALLABLE_HASH]); + // BufferRef blob header: [abi_version u32][ref_count i32][scalar_count i32][reserved]. int32_t t_count = 0; - std::memcpy(&t_count, mailbox.data() + MAILBOX_OFF_TASK_ARGS_BLOB, sizeof(int32_t)); + std::memcpy(&t_count, mailbox.data() + MAILBOX_OFF_TASK_ARGS_BLOB + sizeof(int32_t), sizeof(int32_t)); uint64_t tensor_key = 0; if (t_count > 0) { - Tensor first{}; + BufferRef first{}; std::memcpy( - &first, mailbox.data() + MAILBOX_OFF_TASK_ARGS_BLOB + TASK_ARGS_BLOB_HEADER_SIZE, sizeof(Tensor) + &first, mailbox.data() + MAILBOX_OFF_TASK_ARGS_BLOB + BUFFERREF_BLOB_HEADER_SIZE, + sizeof(BufferRef) ); - tensor_key = first.buffer.addr; + tensor_key = first.handle.identity.buffer_id; } { std::lock_guard lk(dispatched_mu); @@ -244,14 +246,27 @@ class FakeEndpoint final : public WorkerEndpoint { // Helper: build a TaskArgs whose only tensor has the given (data, tag). // --------------------------------------------------------------------------- -static TaskArgs single_tensor_args(uint64_t data_ptr, TensorArgType tag) { +static CanonicalIdentity identity_for(uint64_t buffer_id) { + CanonicalIdentity id{}; + id.buffer_id = buffer_id; + return id; +} + +static BufferRef make_local_ref(uint64_t buffer_id) { + BufferRef r{}; + r.handle.backend_kind = static_cast(BackendKind::POSIX_SHM); + r.handle.nbytes = 1; + r.handle.identity = identity_for(buffer_id); + r.ndims = 1; + r.shapes[0] = 1; + r.strides[0] = 1; + r.dtype = DataType::UINT8; + return r; +} + +static TaskArgs single_ref_args(uint64_t buffer_id, TensorArgType tag) { TaskArgs a; - Tensor t{}; - t.buffer.addr = data_ptr; - t.ndims = 1; - t.shapes[0] = 1; - t.dtype = DataType::UINT8; - a.add_tensor(t, tag); + a.add_tensor(make_local_ref(buffer_id), tag); return a; } @@ -428,7 +443,7 @@ TEST(WorkerManagerTest, ControlPrepareUsesStableNextLevelWorkerId) { } TEST_F(SchedulerFixture, IndependentTaskDispatchedAndConsumed) { - auto args_a = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); + auto args_a = single_ref_args(0xCAFE, TensorArgType::OUTPUT); auto res = orch.submit_next_level(C(42), args_a, cfg, 0); TaskSlot slot = res.task_slot; @@ -442,10 +457,10 @@ TEST_F(SchedulerFixture, IndependentTaskDispatchedAndConsumed) { } TEST_F(SchedulerFixture, DependentTaskDispatchedAfterProducerCompletes) { - auto args_a = single_tensor_args(0xBEEF, TensorArgType::OUTPUT); + auto args_a = single_ref_args(0xBEEF, TensorArgType::OUTPUT); auto a = orch.submit_next_level(C(10), args_a, cfg, 0); - auto args_b = single_tensor_args(0xBEEF, TensorArgType::INPUT); + auto args_b = single_ref_args(0xBEEF, TensorArgType::INPUT); auto b = orch.submit_next_level(C(11), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::PENDING); @@ -470,19 +485,14 @@ TEST_F(SchedulerFixture, DependentTaskDispatchedAfterProducerCompletes) { // mailbox must hold any blob the runtime itself accepts, i.e. up to // CHIP_MAX_TENSOR_ARGS / CHIP_MAX_SCALAR_ARGS. TEST_F(SchedulerFixture, ComposedKernelArgsBlobFitsMailbox) { - constexpr size_t max_blob = TASK_ARGS_BLOB_HEADER_SIZE + - static_cast(CHIP_MAX_TENSOR_ARGS) * sizeof(Tensor) + + constexpr size_t max_blob = BUFFERREF_BLOB_HEADER_SIZE + + static_cast(CHIP_MAX_TENSOR_ARGS) * sizeof(BufferRef) + static_cast(CHIP_MAX_SCALAR_ARGS) * sizeof(uint64_t); EXPECT_GE(MAILBOX_ARGS_CAPACITY, max_blob); TaskArgs args; for (int i = 0; i < 76; ++i) { - Tensor t{}; - t.buffer.addr = 0x1000u + static_cast(i) * 0x100u; - t.ndims = 1; - t.shapes[0] = 1; - t.dtype = DataType::UINT8; - args.add_tensor(t, TensorArgType::OUTPUT); + args.add_tensor(make_local_ref(0x1000u + static_cast(i) * 0x100u), TensorArgType::OUTPUT); } args.add_scalar(1); args.add_scalar(2); @@ -499,10 +509,10 @@ TEST_F(SchedulerFixture, ComposedKernelArgsBlobFitsMailbox) { } TEST_F(SchedulerFixture, FailedProducerPoisonsDependentTask) { - auto args_a = single_tensor_args(0xD00D, TensorArgType::OUTPUT); + auto args_a = single_ref_args(0xD00D, TensorArgType::OUTPUT); auto a = orch.submit_next_level(C(21), args_a, cfg, 0); - auto args_b = single_tensor_args(0xD00D, TensorArgType::INPUT); + auto args_b = single_ref_args(0xD00D, TensorArgType::INPUT); auto b = orch.submit_next_level(C(22), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::PENDING); @@ -599,8 +609,8 @@ struct GroupSchedulerFixture : public ::testing::Test { }; TEST_F(GroupSchedulerFixture, GroupDispatchesToNWorkers) { - TaskArgs a0 = single_tensor_args(0xA0, TensorArgType::OUTPUT); - TaskArgs a1 = single_tensor_args(0xA1, TensorArgType::OUTPUT); + TaskArgs a0 = single_ref_args(0xA0, TensorArgType::OUTPUT); + TaskArgs a1 = single_ref_args(0xA1, TensorArgType::OUTPUT); auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); TaskSlot slot = res.task_slot; @@ -624,8 +634,8 @@ TEST_F(GroupSchedulerFixture, GroupMapsEachMemberToItsTargetWorkerIdNotIndex) { // Reversed target order: member 0 -> worker id 1 (worker_b), member 1 -> // worker id 0 (worker_a). A map-by-registration-index bug would instead // send member 0 (a0) to worker_a; the reversed keys catch it. - TaskArgs a0 = single_tensor_args(0xA0, TensorArgType::OUTPUT); - TaskArgs a1 = single_tensor_args(0xA1, TensorArgType::OUTPUT); + TaskArgs a0 = single_ref_args(0xA0, TensorArgType::OUTPUT); + TaskArgs a1 = single_ref_args(0xA1, TensorArgType::OUTPUT); auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {1, 0}); TaskSlot slot = res.task_slot; @@ -646,8 +656,8 @@ TEST_F(GroupSchedulerFixture, GroupMapsEachMemberToItsTargetWorkerIdNotIndex) { } TEST_F(GroupSchedulerFixture, GroupCompletesOnlyWhenAllDone) { - TaskArgs a0 = single_tensor_args(0xB0, TensorArgType::OUTPUT); - TaskArgs a1 = single_tensor_args(0xB1, TensorArgType::OUTPUT); + TaskArgs a0 = single_ref_args(0xB0, TensorArgType::OUTPUT); + TaskArgs a1 = single_ref_args(0xB1, TensorArgType::OUTPUT); auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); TaskSlot slot = res.task_slot; @@ -663,12 +673,12 @@ TEST_F(GroupSchedulerFixture, GroupCompletesOnlyWhenAllDone) { } TEST_F(GroupSchedulerFixture, BlockedGroupDoesNotDispatchPartiallyOrReserveIdleWorker) { - auto running = orch.submit_next_level(C(70), single_tensor_args(0xF0, TensorArgType::OUTPUT), cfg, 0); + auto running = orch.submit_next_level(C(70), single_ref_args(0xF0, TensorArgType::OUTPUT), cfg, 0); worker_a.wait_running(); ASSERT_TRUE(worker_a.is_running.load()); - TaskArgs group_a = single_tensor_args(0xF1, TensorArgType::OUTPUT); - TaskArgs group_b = single_tensor_args(0xF2, TensorArgType::OUTPUT); + TaskArgs group_a = single_ref_args(0xF1, TensorArgType::OUTPUT); + TaskArgs group_b = single_ref_args(0xF2, TensorArgType::OUTPUT); auto group = orch.submit_next_level_group(C(71), {group_a, group_b}, cfg, {0, 1}); std::this_thread::sleep_for(std::chrono::milliseconds(20)); @@ -676,7 +686,7 @@ TEST_F(GroupSchedulerFixture, BlockedGroupDoesNotDispatchPartiallyOrReserveIdleW EXPECT_FALSE(worker_b.is_running.load()); EXPECT_EQ(worker_b.dispatched_count(), 0); - auto independent = orch.submit_next_level(C(72), single_tensor_args(0xF3, TensorArgType::OUTPUT), cfg, 1); + auto independent = orch.submit_next_level(C(72), single_ref_args(0xF3, TensorArgType::OUTPUT), cfg, 1); worker_b.wait_running(); ASSERT_TRUE(worker_b.is_running.load()); EXPECT_EQ(worker_b.dispatched[0].callable_hash0, 72u); @@ -703,23 +713,23 @@ TEST_F(GroupSchedulerFixture, BlockedGroupDoesNotDispatchPartiallyOrReserveIdleW } TEST_F(GroupSchedulerFixture, LaunchableGroupPrecedesConflictingSingles) { - auto running_a = orch.submit_next_level(C(73), single_tensor_args(0xF4, TensorArgType::OUTPUT), cfg, 0); - auto running_b = orch.submit_next_level(C(74), single_tensor_args(0xF5, TensorArgType::OUTPUT), cfg, 1); + auto running_a = orch.submit_next_level(C(73), single_ref_args(0xF4, TensorArgType::OUTPUT), cfg, 0); + auto running_b = orch.submit_next_level(C(74), single_ref_args(0xF5, TensorArgType::OUTPUT), cfg, 1); worker_a.wait_running(); worker_b.wait_running(); ASSERT_TRUE(worker_a.is_running.load()); ASSERT_TRUE(worker_b.is_running.load()); - TaskArgs group_a = single_tensor_args(0xF6, TensorArgType::OUTPUT); - TaskArgs group_b = single_tensor_args(0xF7, TensorArgType::OUTPUT); + TaskArgs group_a = single_ref_args(0xF6, TensorArgType::OUTPUT); + TaskArgs group_b = single_ref_args(0xF7, TensorArgType::OUTPUT); SubmitResult group; SubmitResult single_a; SubmitResult single_b; { std::lock_guard scheduler_pause(sched.loop_mutex()); group = orch.submit_next_level_group(C(75), {group_a, group_b}, cfg, {0, 1}); - single_a = orch.submit_next_level(C(76), single_tensor_args(0xF8, TensorArgType::OUTPUT), cfg, 0); - single_b = orch.submit_next_level(C(77), single_tensor_args(0xF9, TensorArgType::OUTPUT), cfg, 1); + single_a = orch.submit_next_level(C(76), single_ref_args(0xF8, TensorArgType::OUTPUT), cfg, 0); + single_b = orch.submit_next_level(C(77), single_ref_args(0xF9, TensorArgType::OUTPUT), cfg, 1); worker_a.complete(); worker_b.complete(); @@ -760,8 +770,8 @@ TEST_F(GroupSchedulerFixture, LaunchableGroupPrecedesConflictingSingles) { } TEST_F(GroupSchedulerFixture, GroupFailureWaitsForRunningMembersThenConsumes) { - TaskArgs a0 = single_tensor_args(0xC0, TensorArgType::OUTPUT); - TaskArgs a1 = single_tensor_args(0xC1, TensorArgType::OUTPUT); + TaskArgs a0 = single_ref_args(0xC0, TensorArgType::OUTPUT); + TaskArgs a1 = single_ref_args(0xC1, TensorArgType::OUTPUT); auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); TaskSlot slot = res.task_slot; @@ -782,8 +792,8 @@ TEST_F(GroupSchedulerFixture, GroupFailureWaitsForRunningMembersThenConsumes) { } TEST_F(GroupSchedulerFixture, InvalidGroupIndexFailsAndConsumesGroup) { - TaskArgs a0 = single_tensor_args(0xD0, TensorArgType::OUTPUT); - TaskArgs a1 = single_tensor_args(0xD1, TensorArgType::OUTPUT); + TaskArgs a0 = single_ref_args(0xD0, TensorArgType::OUTPUT); + TaskArgs a1 = single_ref_args(0xD1, TensorArgType::OUTPUT); auto res = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); TaskSlot slot = res.task_slot; @@ -805,7 +815,7 @@ TEST_F(GroupSchedulerFixture, InvalidGroupIndexFailsAndConsumesGroup) { } TEST_F(GroupSchedulerFixture, ExplicitTargetWithinEligibilityIsUsed) { - TaskArgs args = single_tensor_args(0xE0, TensorArgType::OUTPUT); + TaskArgs args = single_ref_args(0xE0, TensorArgType::OUTPUT); auto res = orch.submit_next_level(C(55), args, cfg, 1, {1}); TaskSlot slot = res.task_slot; @@ -820,16 +830,16 @@ TEST_F(GroupSchedulerFixture, ExplicitTargetWithinEligibilityIsUsed) { } TEST_F(GroupSchedulerFixture, BusyTargetDoesNotBlockAnotherWorkerQueue) { - auto running_args = single_tensor_args(0xE4, TensorArgType::OUTPUT); + auto running_args = single_ref_args(0xE4, TensorArgType::OUTPUT); auto running = orch.submit_next_level(C(62), running_args, cfg, 0); worker_a.wait_running(); ASSERT_TRUE(worker_a.is_running.load()); - auto blocked_args = single_tensor_args(0xE5, TensorArgType::OUTPUT); + auto blocked_args = single_ref_args(0xE5, TensorArgType::OUTPUT); auto blocked = orch.submit_next_level(C(63), blocked_args, cfg, 0); - auto blocked_second_args = single_tensor_args(0xE8, TensorArgType::OUTPUT); + auto blocked_second_args = single_ref_args(0xE8, TensorArgType::OUTPUT); auto blocked_second = orch.submit_next_level(C(67), blocked_second_args, cfg, 0); - auto independent_args = single_tensor_args(0xE6, TensorArgType::OUTPUT); + auto independent_args = single_ref_args(0xE6, TensorArgType::OUTPUT); auto independent = orch.submit_next_level(C(64), independent_args, cfg, 1); worker_b.wait_running(); @@ -862,9 +872,9 @@ TEST_F(GroupSchedulerFixture, BusyTargetDoesNotBlockAnotherWorkerQueue) { } TEST_F(GroupSchedulerFixture, DependencyReleaseUsesConsumerWorkerQueue) { - auto producer_args = single_tensor_args(0xE7, TensorArgType::OUTPUT); + auto producer_args = single_ref_args(0xE7, TensorArgType::OUTPUT); auto producer = orch.submit_next_level(C(65), producer_args, cfg, 0); - auto consumer_args = single_tensor_args(0xE7, TensorArgType::INPUT); + auto consumer_args = single_ref_args(0xE7, TensorArgType::INPUT); auto consumer = orch.submit_next_level(C(66), consumer_args, cfg, 1); EXPECT_EQ(S(consumer.task_slot).state.load(), TaskState::PENDING); @@ -882,12 +892,12 @@ TEST_F(GroupSchedulerFixture, DependencyReleaseUsesConsumerWorkerQueue) { } TEST_F(GroupSchedulerFixture, TargetMustBeInEligibleEndpointSet) { - TaskArgs args = single_tensor_args(0xE1, TensorArgType::OUTPUT); + TaskArgs args = single_ref_args(0xE1, TensorArgType::OUTPUT); EXPECT_THROW((void)orch.submit_next_level(C(56), args, cfg, 0, {1}), std::invalid_argument); } TEST_F(GroupSchedulerFixture, UnknownEligibleWorkerIdIsRejectedBeforeScheduling) { - TaskArgs args = single_tensor_args(0xE3, TensorArgType::OUTPUT); + TaskArgs args = single_ref_args(0xE3, TensorArgType::OUTPUT); EXPECT_THROW((void)orch.submit_next_level(C(59), args, cfg, 99, {99}), std::invalid_argument); } @@ -952,7 +962,7 @@ TEST(SchedulerWorkerTargetTest, NextLevelTargetUsesWorkerIdNotVectorIndex) { EXPECT_TRUE(consumed); }; - TaskArgs args = single_tensor_args(0xE2, TensorArgType::OUTPUT); + TaskArgs args = single_ref_args(0xE2, TensorArgType::OUTPUT); auto res = orch.submit_next_level(C(58), args, cfg, 9); worker_b.wait_running(); @@ -966,8 +976,8 @@ TEST(SchedulerWorkerTargetTest, NextLevelTargetUsesWorkerIdNotVectorIndex) { wait_consumed_slot(res.task_slot); - TaskArgs a0 = single_tensor_args(0xE6, TensorArgType::OUTPUT); - TaskArgs a1 = single_tensor_args(0xE7, TensorArgType::OUTPUT); + TaskArgs a0 = single_ref_args(0xE6, TensorArgType::OUTPUT); + TaskArgs a1 = single_ref_args(0xE7, TensorArgType::OUTPUT); auto group_res = orch.submit_next_level_group(C(61), {a0, a1}, cfg, {7, 9}, {{7}, {9}}); worker_a.wait_running(); @@ -988,12 +998,13 @@ TEST(SchedulerWorkerTargetTest, NextLevelTargetUsesWorkerIdNotVectorIndex) { TEST_F(GroupSchedulerFixture, RemoteSidecarRejectsLocalEndpointEligibility) { TaskArgs args; - Tensor tensor{}; - tensor.buffer.addr = 0; - tensor.ndims = 1; - tensor.shapes[0] = 1; - tensor.dtype = DataType::UINT8; - args.add_tensor(tensor, TensorArgType::OUTPUT); + BufferRef ref{}; + ref.handle.backend_kind = static_cast(BackendKind::REMOTE_SIDECAR); + ref.ndims = 1; + ref.shapes[0] = 1; + ref.strides[0] = 1; + ref.dtype = DataType::UINT8; + args.add_tensor(ref, TensorArgType::OUTPUT); RemoteTaskArgsSidecar sidecar; sidecar.tensors.resize(1); @@ -1093,7 +1104,7 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { TEST_F(MixedTypeSchedulerFixture, SubTaskDispatchesWhileNextLevelPoolSaturated) { // Submit a next-level task; the only chip worker begins running it and // stays blocked until we call complete() on it. - auto chip_args = single_tensor_args(0xAAA, TensorArgType::OUTPUT); + auto chip_args = single_ref_args(0xAAA, TensorArgType::OUTPUT); auto chip = orch.submit_next_level(C(20), chip_args, cfg, 0); next_level_worker.wait_running(); ASSERT_TRUE(next_level_worker.is_running.load()); @@ -1102,7 +1113,7 @@ TEST_F(MixedTypeSchedulerFixture, SubTaskDispatchesWhileNextLevelPoolSaturated) // shared ready queue this would block behind any next-level task sitting // in worker 0's directed FIFO. The independent shared SUB queue must // dispatch immediately to the idle SUB worker. - auto sub_args = single_tensor_args(0xBBB, TensorArgType::OUTPUT); + auto sub_args = single_ref_args(0xBBB, TensorArgType::OUTPUT); auto sub = orch.submit_sub(C(7), sub_args); sub_worker.wait_running(); @@ -1122,11 +1133,11 @@ TEST_F(MixedTypeSchedulerFixture, SubTaskDispatchesWhileNextLevelPoolSaturated) TEST_F(GroupSchedulerFixture, GroupDependencyChain) { // Group A (2 workers) produces an OUTPUT at key 0xCAFE. // Task B reads INPUT at the same key -- depends on group A. - TaskArgs a0 = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); - TaskArgs a1 = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); + TaskArgs a0 = single_ref_args(0xCAFE, TensorArgType::OUTPUT); + TaskArgs a1 = single_ref_args(0xCAFE, TensorArgType::OUTPUT); auto a = orch.submit_next_level_group(C(42), {a0, a1}, cfg, {0, 1}); - auto args_b = single_tensor_args(0xCAFE, TensorArgType::INPUT); + auto args_b = single_ref_args(0xCAFE, TensorArgType::INPUT); auto b = orch.submit_next_level(C(42), args_b, cfg, 0); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::PENDING); diff --git a/tests/ut/cpp/types/test_buffer_handle.cpp b/tests/ut/cpp/types/test_buffer_handle.cpp new file mode 100644 index 0000000000..2c6e0eff79 --- /dev/null +++ b/tests/ut/cpp/types/test_buffer_handle.cpp @@ -0,0 +1,235 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +// Wire-ABI tests for BufferHandle / BufferRef (buffer_handle.h), implementing the frozen logical +// schema in .docs/L3-new/worker-memory-model/bufferhandle-abi.md. Byte layout is pinned by +// static_assert in the header; these tests pin the sizes, enum values, and the blob codec. + +#include +#include +#include +#include +#include +#include + +#include + +#include "buffer_handle.h" +#include "task_args.h" + +namespace { + +CanonicalIdentity make_identity() { + CanonicalIdentity id{}; + for (uint32_t i = 0; i < OWNER_INSTANCE_ID_BYTES; ++i) + id.owner_instance_id[i] = static_cast(0xA0 + i); + id.buffer_id = 0x0102030405060708ULL; + id.generation = 7; + const char *path = "L4/L3[2]"; + id.path_len = static_cast(std::strlen(path)); + std::memcpy(id.owner_worker_path, path, id.path_len); + return id; +} + +BufferRef make_ref() { + BufferRef r{}; + r.handle.abi_version = BUFFER_ABI_VERSION; + r.handle.backend_kind = static_cast(BackendKind::POSIX_SHM); + r.handle.descriptor_version = BUFFER_DESCRIPTOR_VERSION; + r.handle.identity = make_identity(); + r.byte_offset = 4096; + r.ndims = 3; + r.shapes[0] = 2; + r.shapes[1] = 4; + r.shapes[2] = 8; + r.strides[0] = 32; + r.strides[1] = 8; + r.strides[2] = 1; + r.dtype = DataType::FLOAT16; + return r; +} + +// --- Layout / value contracts (frozen ABI) ----------------------------------------------------- + +TEST(BufferHandleAbi, StructSizesAreFrozen) { + EXPECT_EQ(sizeof(CanonicalIdentity), 96u); + EXPECT_EQ(sizeof(BufferRef), 272u); + EXPECT_EQ(sizeof(BufferHandleDescriptor), 216u); +} + +TEST(BufferHandleAbi, ConstantsAreFrozen) { + EXPECT_EQ(BUFFER_ABI_VERSION, 1); + EXPECT_EQ(BUFFER_DESCRIPTOR_VERSION, 1); + EXPECT_EQ(OWNER_INSTANCE_ID_BYTES, 16u); + EXPECT_EQ(PATH_MAX_BYTES, 64u); + EXPECT_EQ(DESC_MAX_BYTES, 96u); +} + +TEST(BufferHandleAbi, EnumValuesAreFrozen) { + EXPECT_EQ(static_cast(AddressSpace::HOST), 0); + EXPECT_EQ(static_cast(AddressSpace::DEVICE), 1); + EXPECT_EQ(static_cast(Visibility::PRIVATE), 0); + EXPECT_EQ(static_cast(Visibility::SHARED), 1); + EXPECT_EQ(static_cast(AccessMode::READ), 0); + EXPECT_EQ(static_cast(AccessMode::WRITE), 1); + EXPECT_EQ(static_cast(AccessMode::READWRITE), 2); + EXPECT_EQ(static_cast(BackendKind::FORK_SHM), 0); + EXPECT_EQ(static_cast(BackendKind::POSIX_SHM), 1); + EXPECT_EQ(static_cast(BackendKind::VMM_WINDOW), 2); + EXPECT_EQ(static_cast(BackendKind::REMOTE_SIDECAR), 3); + EXPECT_EQ(static_cast(BackendKind::DEVICE_MALLOC), 4); +} + +// --- memcpy round trip ------------------------------------------------------------------------- + +TEST(BufferHandleAbi, BufferRefSurvivesByteRoundTrip) { + BufferRef src = make_ref(); + uint8_t bytes[sizeof(BufferRef)]; + std::memcpy(bytes, &src, sizeof(BufferRef)); + BufferRef dst{}; + std::memcpy(&dst, bytes, sizeof(BufferRef)); + EXPECT_EQ(std::memcmp(&src, &dst, sizeof(BufferRef)), 0); + EXPECT_EQ(dst.byte_offset, 4096u); + EXPECT_EQ(dst.dtype, DataType::FLOAT16); + EXPECT_EQ(dst.strides[0], 32u); + EXPECT_EQ(dst.handle.identity, src.handle.identity); +} + +TEST(BufferHandleAbi, HandleDescriptorSurvivesByteRoundTrip) { + BufferHandleDescriptor src{}; + src.abi_version = BUFFER_ABI_VERSION; + src.address_space = static_cast(AddressSpace::DEVICE); + src.visibility = static_cast(Visibility::SHARED); + src.access = static_cast(AccessMode::READWRITE); + src.backend_kind = static_cast(BackendKind::POSIX_SHM); + src.descriptor_version = BUFFER_DESCRIPTOR_VERSION; + src.identity = make_identity(); + src.nbytes = 1 << 20; + const char *body = "psm_deadbeef"; + src.body_len = static_cast(std::strlen(body)); + std::memcpy(src.body, body, src.body_len); + + uint8_t bytes[sizeof(BufferHandleDescriptor)]; + std::memcpy(bytes, &src, sizeof(BufferHandleDescriptor)); + BufferHandleDescriptor dst{}; + std::memcpy(&dst, bytes, sizeof(BufferHandleDescriptor)); + EXPECT_EQ(std::memcmp(&src, &dst, sizeof(BufferHandleDescriptor)), 0); + EXPECT_EQ(dst.abi_version, BUFFER_ABI_VERSION); + EXPECT_EQ(dst.identity, src.identity); + EXPECT_EQ(std::string(dst.body, dst.body_len), "psm_deadbeef"); +} + +// --- canonical identity: the import-registry key ----------------------------------------------- + +TEST(BufferHandleAbi, IdentityDistinguishesGenerationOwnerPath) { + CanonicalIdentity a = make_identity(); + CanonicalIdentity b = make_identity(); + EXPECT_EQ(a, b); + + b.generation = a.generation + 1; // buffer_id reuse across generations (ABA) + EXPECT_NE(a, b); + + CanonicalIdentity c = make_identity(); + c.owner_instance_id[0] ^= 0xFF; // different owner incarnation nonce + EXPECT_NE(a, c); + + CanonicalIdentity e = make_identity(); + const char *p2 = "L4/L3[3]"; // different owner path (same length) + std::memcpy(e.owner_worker_path, p2, e.path_len); + EXPECT_NE(a, e); +} + +TEST(BufferHandleAbi, IdentityHashMatchesEquality) { + CanonicalIdentityHash h; + CanonicalIdentity a = make_identity(); + CanonicalIdentity b = make_identity(); + EXPECT_EQ(h(a), h(b)); + + b.generation = a.generation + 1; + EXPECT_NE(h(a), h(b)); // good hash separates the ABA case (not a strict requirement) +} + +// --- BufferRef wire blob: versioned length-prefixed round trip + rejection ---------------------- + +BufferRef make_ref_b() { + BufferRef r = make_ref(); + r.handle.identity.buffer_id = 99; + r.byte_offset = 0; + r.ndims = 1; + r.shapes[0] = 5; + r.strides[0] = 1; + r.dtype = DataType::INT32; + return r; +} + +TEST(BufferRefBlob, RoundTrip) { + BufferRef refs[2] = {make_ref(), make_ref_b()}; + uint64_t scalars[2] = {42, 0xC0FFEE}; + size_t sz = bufferref_blob_size(2, 2); + EXPECT_EQ(sz, BUFFERREF_BLOB_HEADER_SIZE + 2 * sizeof(BufferRef) + 2 * sizeof(uint64_t)); + + std::vector buf(sz); + write_bufferref_blob(buf.data(), refs, 2, scalars, 2); + + BufferRefBlobView v = read_bufferref_blob(buf.data(), sz); + ASSERT_EQ(v.ref_count, 2); + ASSERT_EQ(v.scalar_count, 2); + BufferRef r0 = v.ref(0); + BufferRef r1 = v.ref(1); + EXPECT_EQ(std::memcmp(&r0, &refs[0], sizeof(BufferRef)), 0); + EXPECT_EQ(std::memcmp(&r1, &refs[1], sizeof(BufferRef)), 0); + EXPECT_EQ(v.scalars[0], 42u); + EXPECT_EQ(v.scalars[1], 0xC0FFEEu); +} + +TEST(BufferRefBlob, EmptyBlob) { + size_t sz = bufferref_blob_size(0, 0); + EXPECT_EQ(sz, BUFFERREF_BLOB_HEADER_SIZE); + std::vector buf(sz); + write_bufferref_blob(buf.data(), nullptr, 0, nullptr, 0); + BufferRefBlobView v = read_bufferref_blob(buf.data(), sz); + EXPECT_EQ(v.ref_count, 0); + EXPECT_EQ(v.scalar_count, 0); +} + +TEST(BufferRefBlob, RejectsUnknownVersion) { + std::vector buf(bufferref_blob_size(0, 0)); + write_bufferref_blob(buf.data(), nullptr, 0, nullptr, 0); + uint32_t bad = BUFFER_ABI_VERSION + 1; + std::memcpy(buf.data(), &bad, sizeof(bad)); + EXPECT_THROW(read_bufferref_blob(buf.data(), buf.size()), std::runtime_error); +} + +TEST(BufferRefBlob, RejectsTruncatedCapacity) { + BufferRef refs[1] = {make_ref()}; + std::vector buf(bufferref_blob_size(1, 0)); + write_bufferref_blob(buf.data(), refs, 1, nullptr, 0); + EXPECT_THROW(read_bufferref_blob(buf.data(), BUFFERREF_BLOB_HEADER_SIZE), std::runtime_error); + EXPECT_THROW(read_bufferref_blob(buf.data(), 4), std::runtime_error); +} + +TEST(BufferRefBlob, RejectsNegativeCount) { + std::vector buf(64, 0); + uint32_t ver = BUFFER_ABI_VERSION; + std::memcpy(buf.data() + 0, &ver, sizeof(ver)); + int32_t neg = -1; + std::memcpy(buf.data() + 4, &neg, sizeof(neg)); // ref_count = -1 + EXPECT_THROW(read_bufferref_blob(buf.data(), buf.size()), std::runtime_error); +} + +TEST(BufferRefBlob, RejectsNonZeroReserved) { + std::vector buf(bufferref_blob_size(0, 0)); + write_bufferref_blob(buf.data(), nullptr, 0, nullptr, 0); + uint32_t dirty = 1; + std::memcpy(buf.data() + 12, &dirty, sizeof(dirty)); // reserved header word + EXPECT_THROW(read_bufferref_blob(buf.data(), buf.size()), std::runtime_error); +} + +} // namespace diff --git a/tests/ut/cpp/types/test_child_memory.cpp b/tests/ut/cpp/types/test_child_memory.cpp index e3b63bcc91..845ac40363 100644 --- a/tests/ut/cpp/types/test_child_memory.cpp +++ b/tests/ut/cpp/types/test_child_memory.cpp @@ -25,8 +25,8 @@ TEST(ChildMemory, TensorAbiSize) { EXPECT_EQ(sizeof(Tensor), 128u); } TEST(ChildMemory, DefaultIsZero) { Tensor t{}; - EXPECT_EQ(t.child_memory, 0); - EXPECT_FALSE(t.is_child_memory()); + EXPECT_EQ(t.address_space, AddressSpace::HOST); + EXPECT_FALSE(t.is_device_memory()); } TEST(ChildMemory, SetChildMemory) { @@ -35,9 +35,9 @@ TEST(ChildMemory, SetChildMemory) { t.shapes[0] = 16; t.ndims = 1; t.dtype = DataType::FLOAT32; - t.child_memory = 1; + t.address_space = AddressSpace::DEVICE; - EXPECT_TRUE(t.is_child_memory()); + EXPECT_TRUE(t.is_device_memory()); EXPECT_EQ(t.buffer.addr, 0xDEAD0000u); EXPECT_EQ(t.nbytes(), 16u * 4u); } @@ -47,23 +47,23 @@ TEST(ChildMemory, SetChildMemory) { // --------------------------------------------------------------------------- TEST(ChildMemory, BlobRoundtripPreservesChildMemory) { - TaskArgs args; + ChipStorageTaskArgs args; // the Tensor-blob round trip is the L2/materialized form, not the BufferRef wire Tensor host_t{}; host_t.buffer.addr = 0x1000; host_t.shapes[0] = 4; host_t.ndims = 1; host_t.dtype = DataType::FLOAT32; - host_t.child_memory = 0; - args.add_tensor(host_t, TensorArgType::INPUT); + host_t.address_space = AddressSpace::HOST; + args.add_tensor(host_t); Tensor dev_t{}; dev_t.buffer.addr = 0x2000; dev_t.shapes[0] = 8; dev_t.ndims = 1; dev_t.dtype = DataType::FLOAT16; - dev_t.child_memory = 1; - args.add_tensor(dev_t, TensorArgType::INPUT); + dev_t.address_space = AddressSpace::DEVICE; + args.add_tensor(dev_t); args.add_scalar(42); @@ -77,11 +77,11 @@ TEST(ChildMemory, BlobRoundtripPreservesChildMemory) { ASSERT_EQ(view.tensor_count, 2); ASSERT_EQ(view.scalar_count, 1); - EXPECT_EQ(view.tensors(0).child_memory, 0); - EXPECT_FALSE(view.tensors(0).is_child_memory()); + EXPECT_EQ(view.tensors(0).address_space, AddressSpace::HOST); + EXPECT_FALSE(view.tensors(0).is_device_memory()); - EXPECT_EQ(view.tensors(1).child_memory, 1); - EXPECT_TRUE(view.tensors(1).is_child_memory()); + EXPECT_EQ(view.tensors(1).address_space, AddressSpace::DEVICE); + EXPECT_TRUE(view.tensors(1).is_device_memory()); EXPECT_EQ(view.tensors(1).buffer.addr, 0x2000u); } @@ -95,13 +95,13 @@ TEST(ChildMemory, ViewToChipStoragePreservesChildMemory) { tensors[0].shapes[0] = 1; tensors[0].ndims = 1; tensors[0].dtype = DataType::INT32; - tensors[0].child_memory = 0; + tensors[0].address_space = AddressSpace::HOST; tensors[1].buffer.addr = 0xB000; tensors[1].shapes[0] = 2; tensors[1].ndims = 1; tensors[1].dtype = DataType::INT32; - tensors[1].child_memory = 1; + tensors[1].address_space = AddressSpace::DEVICE; uint64_t scalars[] = {99}; TaskArgsView view{2, 1, reinterpret_cast(tensors), scalars}; @@ -109,8 +109,8 @@ TEST(ChildMemory, ViewToChipStoragePreservesChildMemory) { ChipStorageTaskArgs chip = view_to_chip_storage(view); ASSERT_EQ(chip.tensor_count(), 2); - EXPECT_FALSE(chip.tensor(0).is_child_memory()); - EXPECT_TRUE(chip.tensor(1).is_child_memory()); + EXPECT_FALSE(chip.tensor(0).is_device_memory()); + EXPECT_TRUE(chip.tensor(1).is_device_memory()); EXPECT_EQ(chip.tensor(1).buffer.addr, 0xB000u); } @@ -129,7 +129,7 @@ TEST(ChildMemory, SkipLogicSimulation) { host_t.shapes[0] = 4; host_t.ndims = 1; host_t.dtype = DataType::FLOAT32; - host_t.child_memory = 0; + host_t.address_space = AddressSpace::HOST; args.add_tensor(host_t); Tensor dev_t{}; @@ -137,7 +137,7 @@ TEST(ChildMemory, SkipLogicSimulation) { dev_t.shapes[0] = 8; dev_t.ndims = 1; dev_t.dtype = DataType::FLOAT32; - dev_t.child_memory = 1; + dev_t.address_space = AddressSpace::DEVICE; args.add_tensor(dev_t); int malloc_count = 0; @@ -145,7 +145,7 @@ TEST(ChildMemory, SkipLogicSimulation) { for (int i = 0; i < args.tensor_count(); i++) { Tensor t = args.tensor(i); - if (t.is_child_memory()) { + if (t.is_device_memory()) { passthrough_count++; } else { malloc_count++; diff --git a/tests/ut/py/test_buffer_handle.py b/tests/ut/py/test_buffer_handle.py new file mode 100644 index 0000000000..8d09fd11fd --- /dev/null +++ b/tests/ut/py/test_buffer_handle.py @@ -0,0 +1,399 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Unit tests for simpler.buffer_handle: identity/descriptor pack-unpack + create/import round trip.""" + +import ctypes + +import pytest +from _task_interface import ( + BUFFER_HANDLE_DESCRIPTOR_BYTES, + CANONICAL_IDENTITY_BYTES, + OWNER_INSTANCE_ID_BYTES, + PATH_MAX_BYTES, + DataType, + materialize_bufferref_blob, + read_args_from_blob, +) +from simpler.buffer_handle import ( + AccessMode, + AddressSpace, + BackendKind, + BufferHandleDescriptor, + BufferRef, + CanonicalIdentity, + ImportRegistry, + Visibility, + create_host_shared_buffer, + mint_owner_instance_id, + pack_bufferref_blob, + re_export, + wrap_device_malloc, + wrap_fork_inherited, +) + +_OID = bytes(range(0xA0, 0xA0 + OWNER_INSTANCE_ID_BYTES)) + + +def _identity(oid=_OID, buffer_id=7, path="L4/L3[2]", generation=2): + return CanonicalIdentity(oid, buffer_id, path, generation) + + +def test_identity_roundtrip_various_paths(): + for path in ["", "L4", "L4/L3[2]", "L4/L3[2]/L2[5]"]: + ident = _identity(path=path) + raw = ident.pack() + assert len(raw) == CANONICAL_IDENTITY_BYTES + assert CanonicalIdentity.unpack(raw) == ident + + +def test_identity_rejects_bad_oid_width(): + with pytest.raises(ValueError): + CanonicalIdentity(b"\x00" * 8, 1, "L4", 0) # 8 != OWNER_INSTANCE_ID_BYTES + + +def test_identity_path_over_limit_rejected(): + with pytest.raises(ValueError): + CanonicalIdentity(_OID, 1, "x" * (PATH_MAX_BYTES + 1), 0) + + +def test_identity_distinguishes_generation_owner_path(): + 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(path="L4/L3[3]") # different path + + +def test_descriptor_roundtrip_host_and_device(): + host = BufferHandleDescriptor( + identity=_identity(), + address_space=AddressSpace.HOST, + visibility=Visibility.SHARED, + access=AccessMode.READWRITE, + backend_kind=BackendKind.POSIX_SHM, + nbytes=4096, + body=b"psm_deadbeef", + ) + raw = host.pack() + assert len(raw) == BUFFER_HANDLE_DESCRIPTOR_BYTES + assert BufferHandleDescriptor.unpack(raw) == host + + dev = BufferHandleDescriptor( + identity=_identity(buffer_id=99), + address_space=AddressSpace.DEVICE, + visibility=Visibility.PRIVATE, + access=AccessMode.READ, + backend_kind=BackendKind.VMM_WINDOW, + nbytes=1 << 20, + body=(0x7F00ABCD).to_bytes(8, "little"), + ) + assert BufferHandleDescriptor.unpack(dev.pack()) == dev + + +def test_descriptor_rejects_unknown_version(): + raw = bytearray( + BufferHandleDescriptor( + identity=_identity(), + address_space=AddressSpace.HOST, + visibility=Visibility.SHARED, + access=AccessMode.READWRITE, + backend_kind=BackendKind.POSIX_SHM, + nbytes=8, + ).pack() + ) + raw[0] = raw[0] + 1 # bump abi_version (u16 @ offset 0) + with pytest.raises(ValueError, match="abi_version"): + BufferHandleDescriptor.unpack(bytes(raw)) + + +def test_descriptor_rejects_oversized_body(): + with pytest.raises(ValueError, match="body"): + BufferHandleDescriptor( + identity=_identity(), + address_space=AddressSpace.HOST, + visibility=Visibility.SHARED, + 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_bufferref_view_algebra(): + oid = mint_owner_instance_id() + h = create_host_shared_buffer(nbytes=1024, owner_instance_id=oid, buffer_id=1) + try: + # handle.ref(shape, dtype) is a contiguous full view (row-major strides). + v = h.ref(shapes=(4, 8), dtype=DataType.FLOAT32.value) + assert v.strides == (8, 1) + assert v.is_contiguous + assert v.numel() == 32 + assert v.byte_offset == 0 + + # slice: inner-dim slice keeps the stride, shifts the byte_offset, breaks contiguity. + s = v.slice(1, 2, 6) + assert s.shapes == (4, 4) + assert s.strides == (8, 1) + assert s.byte_offset == 2 * 1 * 4 + assert not s.is_contiguous + # slice with a step multiplies the stride. + s2 = v.slice(0, 0, 4, step=2) + assert s2.shapes == (2, 8) + assert s2.strides == (16, 1) + + # transpose / permute: swap/reorder shapes+strides, unconstrained (strided ok). + t = v.transpose(0, 1) + assert t.shapes == (8, 4) + assert t.strides == (1, 8) + assert not t.is_contiguous + assert v.permute((1, 0)).strides == (1, 8) + + # view: sub-region by per-dim offset, strides unchanged. + vv = v.view((2, 3), (1, 2)) + assert vv.shapes == (2, 3) + assert vv.strides == (8, 1) + assert vv.byte_offset == (1 * 8 + 2 * 1) * 4 + + # reshape: contiguous only. + assert v.reshape((32,)).strides == (1,) + with pytest.raises(ValueError, match="contiguous"): + t.reshape((32,)) + finally: + h.close() + + +def test_bufferref_unpack_roundtrip(): + oid = mint_owner_instance_id() + h = create_host_shared_buffer(64, oid, buffer_id=1, owner_worker_path="L3") + try: + ref = h.ref(shapes=(2, 4), dtype=DataType.FLOAT16.value, byte_offset=8) + assert BufferRef.unpack(ref.pack()) == ref + finally: + h.close() + + +def test_re_export_preserves_identity_same_backing_no_map(): + # Frozen model §5/§8: canonical identity is invariant across every edge. Re-exporting an L4-owned + # backing for forwarding keeps the SOURCE identity (owner_instance_id / path / buffer_id / + # generation) and the same backing, only stripping the mapping. + l4 = mint_owner_instance_id() + src = create_host_shared_buffer(64, l4, buffer_id=7, owner_worker_path="L4") + try: + sdesc = src.to_descriptor() + hp = re_export(sdesc) + assert hp.identity.pack() == src.identity.pack() # identity invariant across the edge + assert hp.backend_kind == BackendKind.POSIX_SHM + assert hp.body == sdesc.body and hp.nbytes == 64 # same backing + assert hp.shm is None and hp.base == 0 # no map (lazy — a compute leaf maps) + # a ref built from H' carries the source identity + the same shm body, so L2 can materialize it + r = hp.ref(shapes=(16,), dtype=DataType.FLOAT32.value) + assert BufferRef.unpack(r.pack()).handle.identity.pack() == src.identity.pack() + finally: + src.close() + + +def test_device_malloc_wrap_materialize(): + # A device pointer (from orch.malloc) wrapped as DEVICE_MALLOC: materializes to the pointer with + # no map, address_space DEVICE (-> a child_memory Tensor). + oid = mint_owner_instance_id() + h = wrap_device_malloc(0xDEAD0000, 4096, oid, buffer_id=3, owner_worker_path="L3") + assert h.backend_kind == BackendKind.DEVICE_MALLOC + assert h.address_space == AddressSpace.DEVICE + assert h.shm is None and h.base == 0xDEAD0000 + reg = ImportRegistry() + imp = reg.materialize(h.to_descriptor()) + assert imp.base == 0xDEAD0000 + assert imp.address_space == AddressSpace.DEVICE + assert imp.shm is None + + +def test_fork_inherited_zero_copy_materialize(): + # A pre-fork COW-inherited allocation: the FORK_SHM body is the base VA, materialized in place + # (no shm, no copy). In-process the VA is trivially valid. + backing = ctypes.create_string_buffer(64) + addr = ctypes.addressof(backing) + oid = mint_owner_instance_id() + handle = wrap_fork_inherited(addr, 64, owner_instance_id=oid, buffer_id=5, owner_worker_path="L3") + assert handle.backend_kind == BackendKind.FORK_SHM + assert handle.access == AccessMode.READ + assert handle.shm is None + reg = ImportRegistry() + try: + ref = handle.ref(shapes=(16,), strides=(1,), dtype=DataType.INT32.value, byte_offset=4) + blob = pack_bufferref_blob([ref]) + src = ctypes.create_string_buffer(blob, len(blob)) + resolved = reg.materialize_blob(ctypes.addressof(src), len(blob)) + assert reg.resolve(handle.identity).base == addr # same VA, no mapping + tensor_blob = materialize_bufferref_blob(ctypes.addressof(src), len(blob), resolved) + dst = ctypes.create_string_buffer(tensor_blob, len(tensor_blob)) + args = read_args_from_blob(ctypes.addressof(dst)) + assert args.tensor(0).data == addr + 4 + finally: + reg.close() + handle.close() # no-op: FORK_SHM owns no shm + + +def test_materialize_remote_sidecar_rejected(): + desc = BufferHandleDescriptor( + identity=_identity(), + address_space=AddressSpace.HOST, + visibility=Visibility.SHARED, + 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) + + +def test_materialize_bufferref_blob_to_tensors(): + oid = mint_owner_instance_id() + h0 = create_host_shared_buffer(nbytes=64, owner_instance_id=oid, buffer_id=1) + h1 = create_host_shared_buffer(nbytes=128, owner_instance_id=oid, buffer_id=2, generation=3) + reg = ImportRegistry() + try: + # Self-describing: refs embed the full descriptor (built via BufferHandle.ref); the consumer + # materializes lazily from the blob (no prior register), map-once by identity. + ref0 = h0.ref(shapes=(4,), strides=(1,), dtype=DataType.FLOAT32.value) + ref1 = h1.ref(shapes=(2, 4), strides=(4, 1), dtype=DataType.FLOAT16.value, byte_offset=8) + assert ref0.handle == h0.to_descriptor() + blob = pack_bufferref_blob([ref0, ref1], scalars=(42,)) + src = ctypes.create_string_buffer(blob, len(blob)) + + resolved = reg.materialize_blob(ctypes.addressof(src), len(blob)) + tensor_blob = materialize_bufferref_blob(ctypes.addressof(src), len(blob), resolved) + dst = ctypes.create_string_buffer(tensor_blob, len(tensor_blob)) + args = read_args_from_blob(ctypes.addressof(dst)) + + assert args.tensor_count() == 2 + assert args.scalar_count() == 1 + assert args.tensor(0).data == reg.resolve(h0.identity).base + 0 + assert args.tensor(0).shapes == (4,) + assert args.tensor(0).child_memory is False # HOST + assert args.tensor(1).data == reg.resolve(h1.identity).base + 8 + assert args.tensor(1).shapes == (2, 4) + assert args.scalar(0) == 42 + finally: + reg.close() + h0.close() + h1.close() + + +def test_materialize_strided_views(): + # transpose / inner-slice produce non-contiguous refs that materialize to strided Tensors + # (matching the mainline strided Tensor wire), not rejected. + oid = mint_owner_instance_id() + h = create_host_shared_buffer(1024, oid, buffer_id=1) + reg = ImportRegistry() + try: + t = h.ref(shapes=(4, 8), dtype=DataType.FLOAT32.value).transpose(0, 1) # (8,4) strides (1,8) + s = h.ref(shapes=(4, 8), dtype=DataType.FLOAT32.value).slice(1, 2, 6) # (4,4) strides (8,1), off 8 + blob = pack_bufferref_blob([t, s]) + src = ctypes.create_string_buffer(blob, len(blob)) + resolved = reg.materialize_blob(ctypes.addressof(src), len(blob)) + tb = materialize_bufferref_blob(ctypes.addressof(src), len(blob), resolved) + dst = ctypes.create_string_buffer(tb, len(tb)) + args = read_args_from_blob(ctypes.addressof(dst)) + base = reg.resolve(h.identity).base + + tt = args.tensor(0) + assert tt.shapes == (8, 4) and tt.strides == (1, 8) and not tt.is_contiguous + assert tt.data == base + ss = args.tensor(1) + assert ss.shapes == (4, 4) and ss.strides == (8, 1) and not ss.is_contiguous + assert ss.data == base + 2 * 1 * 4 # slice(1,2,..) shifts byte_offset by start*stride*elem + finally: + reg.close() + h.close() + + +def test_materialize_rejects_unresolved_identity(): + oid = mint_owner_instance_id() + handle = create_host_shared_buffer(nbytes=32, owner_instance_id=oid, buffer_id=9) + try: + ref = BufferRef(handle.to_descriptor(), byte_offset=0, shapes=(8,), strides=(1,), dtype=DataType.INT32.value) + blob = pack_bufferref_blob([ref]) + src = ctypes.create_string_buffer(blob, len(blob)) + with pytest.raises(RuntimeError, match="identity"): + materialize_bufferref_blob(ctypes.addressof(src), len(blob), {}) # nothing resolved + finally: + handle.close() + + +@pytest.mark.parametrize( + "space,backend", + [ + (AddressSpace.HOST, BackendKind.VMM_WINDOW), + (AddressSpace.HOST, BackendKind.DEVICE_MALLOC), + (AddressSpace.DEVICE, BackendKind.FORK_SHM), + (AddressSpace.DEVICE, BackendKind.POSIX_SHM), + ], +) +def test_descriptor_rejects_bad_capability_combo(space, backend): + # §4.1 capability matrix: an unsupported address_space×backend_kind fails at construction (before + # dispatch, before it can ride the wire). + with pytest.raises(ValueError, match="capability"): + BufferHandleDescriptor( + identity=_identity(), + address_space=space, + visibility=Visibility.SHARED, + access=AccessMode.READWRITE, + backend_kind=backend, + nbytes=64, + body=b"", + ) + + +def test_descriptor_accepts_legal_combos(): + for space, backend in [ + (AddressSpace.HOST, BackendKind.FORK_SHM), + (AddressSpace.HOST, BackendKind.POSIX_SHM), + (AddressSpace.DEVICE, BackendKind.VMM_WINDOW), + (AddressSpace.DEVICE, BackendKind.DEVICE_MALLOC), + (AddressSpace.HOST, BackendKind.REMOTE_SIDECAR), + (AddressSpace.DEVICE, BackendKind.REMOTE_SIDECAR), + ]: + BufferHandleDescriptor(_identity(), space, Visibility.SHARED, AccessMode.READWRITE, backend, 64, b"") diff --git a/tests/ut/py/test_callable_identity.py b/tests/ut/py/test_callable_identity.py index e8ce434496..ba7df01db6 100644 --- a/tests/ut/py/test_callable_identity.py +++ b/tests/ut/py/test_callable_identity.py @@ -9,6 +9,7 @@ import ctypes import hashlib +import itertools import os import socket import subprocess @@ -21,6 +22,7 @@ import pytest import simpler.worker as worker_mod from simpler import callable_identity +from simpler.buffer_handle import mint_owner_instance_id, wrap_fork_inherited from simpler.callable_identity import ( CallableHandle, CallableKindName, @@ -57,8 +59,8 @@ RemoteBufferHandle, RemoteTensorRef, TaskArgs, - Tensor, TensorArgType, + get_element_size, ) from simpler.worker import ( RemoteCallable, @@ -68,6 +70,19 @@ _read_raw_payload_from_shm, ) +_LOCAL_REF_BID = itertools.count(1) + + +def _local_ref(addr, shapes, dtype): + """A local host (non-remote, non-child-memory) ``BufferRef`` at ``addr`` — a bare arg carrying no + sidecar. FORK_SHM (host) so it is not treated as a child device pointer by the dispatch guard.""" + nbytes = get_element_size(dtype) + for s in shapes: + nbytes *= int(s) + return wrap_fork_inherited(addr, nbytes, mint_owner_instance_id(), next(_LOCAL_REF_BID), "L3").ref( + tuple(shapes), int(dtype.value) + ) + def _py_target(args): return args @@ -778,7 +793,7 @@ def submit_next_level_group(self, *args): nbytes=16, ) args = TaskArgs() - args.add_tensor(RemoteTensorRef(buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) + args.add_ref(RemoteTensorRef(buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) orch.submit_next_level(handle, args, worker=worker_id) call = fake.submit_next_level_args @@ -791,7 +806,7 @@ def submit_next_level_group(self, *args): assert sidecar.tensors[0].desc.owner_worker_id == worker_id bare = TaskArgs() - bare.add_tensor(Tensor.make(0x1234, (1,), DataType.UINT8), TensorArgType.INPUT) + bare.add_ref(_local_ref(0x1234, (1,), DataType.UINT8), TensorArgType.INPUT) orch.submit_next_level(handle, bare, worker=worker_id) bare_sidecar = fake.submit_next_level_args[7] @@ -813,7 +828,7 @@ def submit_sub(self, *args): fake = FakeCOrchestrator() orch = Orchestrator(fake) # type: ignore[arg-type] args = TaskArgs() - args.add_tensor( + args.add_ref( RemoteTensorRef.host_inline(b"abcd", shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT, ) @@ -837,7 +852,7 @@ def submit_sub_group(self, *args): orch = Orchestrator(fake) # type: ignore[arg-type] local_args = TaskArgs() remote_args = TaskArgs() - remote_args.add_tensor( + remote_args.add_ref( RemoteTensorRef.host_inline(b"abcd", shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT, ) @@ -914,7 +929,7 @@ def submit_next_level(self, *args): access_flags=access_flags, ) args = TaskArgs() - args.add_tensor(RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), tag) + args.add_ref(RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), tag) with pytest.raises(ValueError, match="remote tensor .* access"): orch.submit_next_level(handle, args, worker=worker_id) @@ -955,7 +970,7 @@ def submit_next_level(self, *args): access_flags=3, ) args = TaskArgs() - args.add_tensor(RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), tag) + args.add_ref(RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), tag) orch.submit_next_level(handle, args, worker=worker_id) @@ -988,7 +1003,7 @@ def submit_next_level_group(self, *args): nbytes=16, ) args = TaskArgs() - args.add_tensor(RemoteTensorRef(buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) + args.add_ref(RemoteTensorRef(buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) orch.submit_next_level(handle, args, worker=worker_id1) assert fake.submit_next_level_args[6] == [worker_id1] @@ -1020,7 +1035,7 @@ def submit_next_level_group(self, *args): nbytes=16, ) args = TaskArgs() - args.add_tensor(RemoteTensorRef(buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) + args.add_ref(RemoteTensorRef(buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) with pytest.raises(ValueError, match="no eligible remote worker"): orch.submit_next_level(handle, args, worker=worker_id0) @@ -1059,9 +1074,9 @@ def submit_next_level_group(self, *args): nbytes=16, ) args0 = TaskArgs() - args0.add_tensor(RemoteTensorRef(buf0, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) + args0.add_ref(RemoteTensorRef(buf0, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) args1 = TaskArgs() - args1.add_tensor(RemoteTensorRef(buf1, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) + args1.add_ref(RemoteTensorRef(buf1, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) orch.submit_next_level_group(handle, [args0, args1], workers=[worker_id0, worker_id1]) assert fake.submit_next_level_group_args[6] == [[worker_id0], [worker_id1]] @@ -1327,7 +1342,7 @@ def test_remote_sim_buffer_copy_roundtrip(): def parent_orch(orch, _args, cfg): task_args = TaskArgs() - task_args.add_tensor( + task_args.add_ref( RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INOUT, ) @@ -1373,7 +1388,7 @@ def test_remote_sim_imported_buffer_runs_on_peer_worker(): def parent_orch(orch, _args, cfg): task_args = TaskArgs() - task_args.add_tensor( + task_args.add_ref( RemoteTensorRef(peer_buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INOUT, ) @@ -1781,12 +1796,12 @@ def test_remote_sim_failed_dependency_skips_consumer(): def parent_orch(orch, _args, cfg): producer_args = TaskArgs() - producer_args.add_tensor( + producer_args.add_ref( RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.OUTPUT, ) consumer_args = TaskArgs() - consumer_args.add_tensor( + consumer_args.add_ref( RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT, ) @@ -1852,11 +1867,11 @@ def test_remote_sim_input_free_is_deferred_until_slot_refs_drop(): def parent_orch(orch, _args, cfg): task_args = TaskArgs() - task_args.add_tensor( + task_args.add_ref( RemoteTensorRef(remote_in, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT, ) - task_args.add_tensor( + task_args.add_ref( RemoteTensorRef(remote_out, shape=(1,), dtype=DataType.UINT8), TensorArgType.OUTPUT, ) @@ -1894,11 +1909,11 @@ def test_remote_sim_host_inline_descriptor_roundtrip(): def parent_orch(orch, _args, cfg): task_args = TaskArgs() - task_args.add_tensor( + task_args.add_ref( RemoteTensorRef.host_inline(b"\x01\x02\x03\x04", shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT, ) - task_args.add_tensor( + task_args.add_ref( RemoteTensorRef(remote_out, shape=(1,), dtype=DataType.UINT8), TensorArgType.OUTPUT, ) diff --git a/tests/ut/py/test_scene_test_golden_hooks.py b/tests/ut/py/test_scene_test_golden_hooks.py index da01d1b7bf..32b9468744 100644 --- a/tests/ut/py/test_scene_test_golden_hooks.py +++ b/tests/ut/py/test_scene_test_golden_hooks.py @@ -52,7 +52,7 @@ def _drive(inst, case, *, cli_skip=False): """Run one case through the L2 path with the device-touching steps stubbed.""" inst._st_level = 2 with ( - patch.object(_SCENE_TEST_MOD, "_build_chip_task_args", return_value=(object(), list(_OUTPUT_NAMES))), + patch.object(_SCENE_TEST_MOD, "_build_l2_ref_args", return_value=(object(), list(_OUTPUT_NAMES))), patch.object(type(inst), "_build_config", return_value=object()), ): inst._run_and_validate_l2(MagicMock(), object(), case, skip_golden=cli_skip) diff --git a/tests/ut/py/test_scene_test_rehost.py b/tests/ut/py/test_scene_test_rehost.py index 4d4f5fb9a2..7482856901 100644 --- a/tests/ut/py/test_scene_test_rehost.py +++ b/tests/ut/py/test_scene_test_rehost.py @@ -26,31 +26,41 @@ from simpler_setup.scene_test import Scalar, TaskArgsBuilder, Tensor, _RehostedTaskArgs -class _FakeHostBuffer: - def __init__(self, nbytes: int): +class _FakeHandle: + """Stands in for a ``create_buffer`` BufferHandle: a POSIX shm the rehost view is built over.""" + + def __init__(self, nbytes: int, worker: _FakeWorker): self.shm = SharedMemory(create=True, size=nbytes) - self.buffer = self.shm.buf + self._worker = worker + self._closed = False + + @property + def buffer(self): + return self.shm.buf + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._worker.freed.append(self) + self.shm.close() + self.shm.unlink() class _FakeWorker: - """Stands in for a started L3 Worker's host-buffer allocator.""" + """Stands in for a started L3 Worker's ``create_buffer`` allocator.""" def __init__(self, fail_on_create: int | None = None): - self.created: list[_FakeHostBuffer] = [] - self.freed: list[_FakeHostBuffer] = [] + self.created: list[_FakeHandle] = [] + self.freed: list[_FakeHandle] = [] self._fail_on_create = fail_on_create - def create_host_buffer(self, nbytes: int) -> _FakeHostBuffer: + def create_buffer(self, nbytes: int) -> _FakeHandle: if self._fail_on_create is not None and len(self.created) >= self._fail_on_create: - raise RuntimeError("injected create_host_buffer failure") - buf = _FakeHostBuffer(nbytes) - self.created.append(buf) - return buf - - def free_host_buffer(self, buf: _FakeHostBuffer) -> None: - self.freed.append(buf) - buf.shm.close() - buf.shm.unlink() + raise RuntimeError("injected create_buffer failure") + handle = _FakeHandle(nbytes, self) + self.created.append(handle) + return handle def test_rehost_preserves_values_and_frees_lifo(): @@ -87,7 +97,7 @@ def test_rehost_partial_failure_rolls_back(): Tensor("c", torch.zeros(4, dtype=torch.float32)), ) w = _FakeWorker(fail_on_create=2) # third allocation fails - with pytest.raises(RuntimeError, match="injected create_host_buffer failure"): + with pytest.raises(RuntimeError, match="injected create_buffer failure"): _RehostedTaskArgs(w, ta) # The two successfully-created buffers are freed, and the builder is # restored to its original tensors (no half-rehosted state). diff --git a/tests/ut/py/test_task_interface.py b/tests/ut/py/test_task_interface.py index 9ed930c620..59f8160ba9 100644 --- a/tests/ut/py/test_task_interface.py +++ b/tests/ut/py/test_task_interface.py @@ -11,6 +11,7 @@ import ctypes import gc +import itertools import struct import weakref @@ -31,6 +32,15 @@ get_dtype_name, get_element_size, ) +from simpler.buffer_handle import ( + AccessMode, + BackendKind, + BufferRef, + create_host_shared_buffer, + mint_owner_instance_id, + wrap_device_malloc, + wrap_fork_inherited, +) from simpler.task_interface import ( RemoteAddressSpace, RemoteBufferExport, @@ -39,6 +49,28 @@ _remote_sidecar_for, ) +_REF_BID = itertools.count(1) + + +def _dev_ref(addr, shapes, dtype, tag=None): + """A DEVICE_MALLOC ``BufferRef`` at ``addr`` for wire-``TaskArgs`` container tests. + + Each call mints a fresh identity, so refs added to one TaskArgs stay distinct. The backend body + carries ``addr`` (u64 LE), the container-test counterpart of the old ``Tensor.make(addr, ...)``. + """ + nbytes = get_element_size(dtype) + for s in shapes: + nbytes *= int(s) + return wrap_device_malloc(addr, nbytes, mint_owner_instance_id(), next(_REF_BID), "L2").ref( + tuple(shapes), int(dtype.value) + ) + + +def _ref_addr(packed_ref) -> int: + """The device pointer carried in a DEVICE_MALLOC ref's backend body.""" + return int.from_bytes(BufferRef.unpack(packed_ref).handle.body[:8], "little") + + # ============================================================================ # DataType enum # ============================================================================ @@ -355,26 +387,24 @@ def test_empty(self): assert args.tensor_count() == 0 assert args.scalar_count() == 0 - def test_add_tensor_default_tag(self): + def test_add_ref_default_tag(self): args = TaskArgs() - t = Tensor.make(0xBEEF, (4, 8), DataType.FLOAT32) - args.add_tensor(t) + args.add_ref(_dev_ref(0xBEEF, (4, 8), DataType.FLOAT32)) assert args.tensor_count() == 1 assert args.tag(0) == TensorArgType.INPUT - def test_add_tensor_with_tag(self): + def test_add_ref_with_tag(self): args = TaskArgs() - t = Tensor.make(0xBEEF, (4, 8), DataType.FLOAT32) - args.add_tensor(t, TensorArgType.OUTPUT) + args.add_ref(_dev_ref(0xBEEF, (4, 8), DataType.FLOAT32), TensorArgType.OUTPUT) assert args.tag(0) == TensorArgType.OUTPUT - def test_multiple_tensors_with_tags(self): + def test_multiple_refs_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_ref(_dev_ref(0x1, (2,), DataType.INT32), TensorArgType.INPUT) + args.add_ref(_dev_ref(0x2, (3,), DataType.FLOAT16), TensorArgType.OUTPUT) + args.add_ref(_dev_ref(0x3, (4,), DataType.INT8), TensorArgType.INOUT) + args.add_ref(_dev_ref(0x4, (5,), DataType.FLOAT32), TensorArgType.OUTPUT_EXISTING) + args.add_ref(_dev_ref(0x5, (6,), DataType.INT32), TensorArgType.NO_DEP) assert args.tensor_count() == 5 assert args.tag(0) == TensorArgType.INPUT assert args.tag(1) == TensorArgType.OUTPUT @@ -384,7 +414,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_ref(_dev_ref(0x1, (2,), DataType.INT32)) assert args.tag(0) == TensorArgType.INPUT args.set_tag(0, TensorArgType.INOUT) assert args.tag(0) == TensorArgType.INOUT @@ -398,15 +428,15 @@ 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_ref(_dev_ref(0x1, (2,), DataType.INT32), TensorArgType.INPUT) + args.add_ref(_dev_ref(0x2, (3,), DataType.FLOAT16), TensorArgType.OUTPUT) args.add_scalar(99) args.add_scalar(100) assert args.tensor_count() == 2 assert args.scalar_count() == 2 assert len(args) == 4 - assert args.tensor(0).data == 0x1 - assert args.tensor(1).data == 0x2 + assert _ref_addr(args.ref(0)) == 0x1 + assert _ref_addr(args.ref(1)) == 0x2 assert args.scalar(0) == 99 assert args.scalar(1) == 100 @@ -414,16 +444,16 @@ 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_ref(_dev_ref(0x1, (2,), DataType.INT32)) - def test_tensor_access(self): + def test_ref_access(self): args = TaskArgs() - args.add_tensor(Tensor.make(0xA, (4,), DataType.FLOAT32)) - args.add_tensor(Tensor.make(0xB, (8,), DataType.INT32)) - assert args.tensor(0).data == 0xA - assert args.tensor(1).data == 0xB - assert args.tensor(0).shapes == (4,) - assert args.tensor(1).shapes == (8,) + args.add_ref(_dev_ref(0xA, (4,), DataType.FLOAT32)) + args.add_ref(_dev_ref(0xB, (8,), DataType.INT32)) + assert _ref_addr(args.ref(0)) == 0xA + assert _ref_addr(args.ref(1)) == 0xB + assert BufferRef.unpack(args.ref(0)).shapes == (4,) + assert BufferRef.unpack(args.ref(1)).shapes == (8,) def test_scalar_access(self): args = TaskArgs() @@ -435,7 +465,7 @@ def test_scalar_access(self): def test_tensor_out_of_range(self): args = TaskArgs() with pytest.raises((IndexError, RuntimeError)): - args.tensor(0) + args.ref(0) def test_scalar_out_of_range(self): args = TaskArgs() @@ -454,7 +484,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_ref(_dev_ref(0, (1,), DataType.INT8), TensorArgType.OUTPUT) args.add_scalar(42) args.clear() assert len(args) == 0 @@ -462,10 +492,10 @@ def test_clear(self): assert args.scalar_count() == 0 def test_no_capacity_limit_tensors(self): - """TaskArgs is vector-backed — no per-class capacity limit on tensors.""" + """TaskArgs is vector-backed — no per-class capacity limit on refs.""" args = TaskArgs() for i in range(20): - args.add_tensor(Tensor.make(i, (1,), DataType.INT8)) + args.add_ref(_dev_ref(i + 1, (1,), DataType.INT8)) assert args.tensor_count() == 20 def test_no_capacity_limit_scalars(self): @@ -493,11 +523,12 @@ def test_remote_buffer_ref_adds_zero_metadata_and_sidecar(self): ref = RemoteTensorRef(handle=handle, offset=8, shape=(4,), dtype=DataType.UINT8) args = TaskArgs() - args.add_tensor(ref, TensorArgType.OUTPUT) + args.add_ref(ref, TensorArgType.OUTPUT) args.add_scalar(9) assert args.tensor_count() == 1 - assert args.tensor(0).data == 0 + # An arg destined for a remote worker carries a REMOTE_SIDECAR placeholder ref (no local backing). + assert BufferRef.unpack(args.ref(0)).handle.backend_kind == BackendKind.REMOTE_SIDECAR assert args.tag(0) == TensorArgType.OUTPUT assert args.scalar(0) == 9 @@ -529,7 +560,7 @@ def make_remote_args_ref(): nbytes=64, ) args = TaskArgs() - args.add_tensor(RemoteTensorRef(handle=handle, shape=(4,), dtype=DataType.UINT8)) + args.add_ref(RemoteTensorRef(handle=handle, shape=(4,), dtype=DataType.UINT8)) assert args in task_interface_module._REMOTE_TASK_ARGS_STORAGE # noqa: SLF001 assert _remote_sidecar_for(args) is not None return weakref.ref(args) @@ -577,7 +608,7 @@ def test_remote_buffer_handle_is_opaque_to_public_constructor(self): def test_host_inline_ref_uses_inline_payload_arena(self): args = TaskArgs() - args.add_tensor( + args.add_ref( RemoteTensorRef.host_inline(b"abcd", shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT, ) @@ -958,3 +989,30 @@ def test_chip_storage_preserves_child_memory(self): out = args.tensor(0) assert out.child_memory is True assert out.data == 0x2000 + + +class TestAddRefAccessSubset: + # §3 access: add_ref rejects an arg whose TensorArgType needs access the handle does not grant. + def test_add_ref_rejects_write_arg_on_read_only_backing(self): + oid = mint_owner_instance_id() + h = wrap_fork_inherited(0x1000, 64, owner_instance_id=oid, buffer_id=1, access=AccessMode.READ) + ref = h.ref((16,), DataType.FLOAT32.value) # READ-only backing + args = TaskArgs() + with pytest.raises(ValueError, match="access"): + args.add_ref(ref.pack(), TensorArgType.OUTPUT_EXISTING) + with pytest.raises(ValueError, match="access"): + args.add_ref(ref.pack(), TensorArgType.INOUT) + args.add_ref(ref.pack(), TensorArgType.INPUT) # READ⊆READ ok + + def test_add_ref_accepts_matching_access(self): + oid = mint_owner_instance_id() + h = create_host_shared_buffer(64, oid, buffer_id=1) # READWRITE + ref = h.ref((16,), DataType.FLOAT32.value) + try: + args = TaskArgs() + args.add_ref(ref.pack(), TensorArgType.INPUT) # READ⊆RW + args.add_ref(ref.pack(), TensorArgType.OUTPUT_EXISTING) # WRITE⊆RW + args.add_ref(ref.pack(), TensorArgType.INOUT) # RW⊆RW + assert args.tensor_count() == 3 + finally: + h.close() 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 5196443dc8..dfac677d93 100644 --- a/tests/ut/py/test_worker/test_child_addr_guard.py +++ b/tests/ut/py/test_worker/test_child_addr_guard.py @@ -24,19 +24,44 @@ import pytest import simpler.orchestrator as orch_mod -from _task_interface import DataType, Tensor, TensorArgType +from _task_interface import DataType, TensorArgType +from simpler.buffer_handle import BufferHandle, mint_owner_instance_id, wrap_device_malloc, wrap_fork_inherited from simpler.orchestrator import Orchestrator from simpler.task_interface import TaskArgs from simpler.worker import Worker, _ChildProvEntry, _Lifecycle +_F32 = 0 # DataType.FLOAT32 value +_OID = mint_owner_instance_id() +_HOSTSRC = bytearray(64) # a writable host buffer for copy_to/copy_from + def _l3() -> Worker: return Worker(level=3, num_sub_workers=0, platform="a2a3sim", runtime="tensormap_and_ringbuffer") +def _l3_ready(malloc_ret: int = 0x1000, *, chips: int = 2) -> tuple[Worker, MagicMock]: + """An uninitialized L3 Worker forced READY with a mocked native worker, so the leased device-mem + ops run in isolation. Returns ``(worker, native_worker_mock)``; the mock's ``malloc`` returns + ``malloc_ret`` (the fabricated device ptr).""" + w = _l3() + w._lifecycle = _Lifecycle.READY + w._chip_shms = [None] * chips # type: ignore[list-item] + nw = MagicMock() + nw.malloc.return_value = malloc_ret + w._worker = nw + return w, nw + + +def _dev_handle(ptr: int, *, wid: int = 0, nbytes: int = 64) -> BufferHandle: + """A DEVICE_MALLOC handle keyed at ``(wid, ptr)`` — the free/copy provenance key.""" + return wrap_device_malloc(ptr, nbytes, _OID, buffer_id=ptr, owner_worker_id=wid) + + def _child_args(ptr: int, *, n: int = 16) -> TaskArgs: + # A DEVICE_MALLOC ref carries ``ptr`` in its backend body: the device-pointer provenance key. args = TaskArgs() - args.add_tensor(Tensor.make(ptr, (n,), DataType.FLOAT32, child_memory=True), TensorArgType.OUTPUT_EXISTING) + dev = wrap_device_malloc(ptr, n * 4, _OID, buffer_id=ptr) + args.add_ref(dev.ref(shapes=(n,), dtype=_F32), TensorArgType.OUTPUT_EXISTING) return args @@ -181,62 +206,58 @@ def test_unique_target_but_wrong_worker_rejected(self): # ---------------------------------------------------------------------------- -# Orchestrator malloc / free / copy — the L3 choke (also the orch.* bypass) +# Worker device memory — alloc_child_tensor / free / copy on a next-level worker +# (the Worker is the sole allocator; orch.* are thin delegates) # ---------------------------------------------------------------------------- -class TestOrchestratorMemoryOps: - def test_malloc_records_then_free_clears(self): - w = _l3() - fake = MagicMock() - fake.malloc.return_value = 0x1000 - o = Orchestrator(fake, w) - ptr = o.malloc(0, 64) - assert ptr == 0x1000 +class TestChildDeviceMemoryOps: + def test_alloc_child_records_then_free_clears(self): + w, nw = _l3_ready(0x1000) + h = w.alloc_child_tensor(0, (64,), DataType.UINT8) + assert h.base == 0x1000 assert (0, 0x1000) in w._child_alloc_prov - o.free(0, ptr) - fake.free.assert_called_once_with(0, 0x1000) + w.free(h) + nw.free.assert_called_once_with(0, 0x1000) assert (0, 0x1000) not in w._child_alloc_prov def test_free_wrong_worker_rejected_without_native_free(self): - w = _l3() - fake = MagicMock() - fake.malloc.return_value = 0x1000 - o = Orchestrator(fake, w) - o.malloc(1, 64) # allocated on worker 1 + w, nw = _l3_ready(0x1000) + w.alloc_child_tensor(1, (64,), DataType.UINT8) # allocated on worker 1 with pytest.raises(ValueError, match="not a live malloc base"): - o.free(0, 0x1000) # freed on worker 0 - fake.free.assert_not_called() + w.free(_dev_handle(0x1000, wid=0)) # freed on worker 0 + nw.free.assert_not_called() def test_copy_to_requires_live_device_dst(self): - w = _l3() - fake = MagicMock() - fake.malloc.return_value = 0x2000 - o = Orchestrator(fake, w) + w, nw = _l3_ready(0x2000) with pytest.raises(ValueError, match="not a live allocation"): - o.copy_to(0, 0x2000, 0xABCD, 64) - fake.copy_to.assert_not_called() - o.malloc(0, 64) - o.copy_to(0, 0x2000, 0xABCD, 64) - fake.copy_to.assert_called_once() + w.copy_to(_dev_handle(0x2000, wid=0), _HOSTSRC) + nw.copy_to.assert_not_called() + w.alloc_child_tensor(0, (64,), DataType.UINT8) + w.copy_to(_dev_handle(0x2000, wid=0), _HOSTSRC) + nw.copy_to.assert_called_once() def test_copy_from_requires_live_device_src(self): - w = _l3() - fake = MagicMock() - fake.malloc.return_value = 0x3000 - o = Orchestrator(fake, w) + w, nw = _l3_ready(0x3000) with pytest.raises(ValueError, match="not a live allocation"): - o.copy_from(0, 0xABCD, 0x3000, 64) - fake.copy_from.assert_not_called() + w.copy_from(_HOSTSRC, _dev_handle(0x3000, wid=0)) + nw.copy_from.assert_not_called() - def test_isolated_orchestrator_has_no_guard(self): - # An Orchestrator constructed without a Worker back-ref (test isolation) - # must not attempt provenance tracking. - fake = MagicMock() - fake.malloc.return_value = 0x1000 - o = Orchestrator(fake, None) - assert o.malloc(0, 64) == 0x1000 - o.free(0, 0x1000) # no raise, no table + def test_orch_delegates_to_worker(self): + # orch.alloc_child_tensor / free are thin wrappers over the bound Worker. + w, nw = _l3_ready(0x1000) + o = Orchestrator(MagicMock(), w) + h = o.alloc_child_tensor(0, (64,), DataType.UINT8) + assert (0, 0x1000) in w._child_alloc_prov + o.free(h) + nw.free.assert_called_once_with(0, 0x1000) + assert (0, 0x1000) not in w._child_alloc_prov + + def test_orch_without_worker_rejects_memory_ops(self): + # A worker-less Orchestrator can't allocate/free/copy — the impl lives on the Worker. + o = Orchestrator(MagicMock(), None) + with pytest.raises(RuntimeError, match="requires a Worker context"): + o.free(_dev_handle(0x1000)) # ---------------------------------------------------------------------------- @@ -277,12 +298,13 @@ def test_child_arg_to_wrong_worker_rejected(self, _fake_handle): fake.submit_next_level.assert_not_called() def test_host_only_args_are_not_guarded(self, _fake_handle): - # A submit with no child_memory tensor never touches provenance. + # A submit with no device (DEVICE_MALLOC) ref never touches provenance. w = _l3() fake = MagicMock() o = Orchestrator(fake, w) args = TaskArgs() - args.add_tensor(Tensor.make(0, (16,), DataType.FLOAT32, child_memory=False), TensorArgType.INPUT) + host = wrap_fork_inherited(0x9000, 64, _OID, buffer_id=0x9000) + args.add_ref(host.ref(shapes=(16,), dtype=_F32), TensorArgType.INPUT) o.submit_next_level(object(), args, None, worker=0) fake.submit_next_level.assert_called_once() @@ -380,37 +402,38 @@ def _l2(self) -> tuple[Worker, MagicMock]: def test_l2_malloc_records_and_free_clears(self): w, chip = self._l2() - ptr = w.malloc(64) - assert (0, ptr) in w._child_alloc_prov - w.free(ptr) + h = w.malloc(64) + assert h.base == 0x2000 + assert (0, 0x2000) in w._child_alloc_prov + w.free(h) chip.free.assert_called_once_with(0x2000) - assert (0, ptr) not in w._child_alloc_prov + assert (0, 0x2000) not in w._child_alloc_prov def test_l2_free_stale_rejected_without_native_free(self): w, chip = self._l2() - w.malloc(64) - w.free(0x2000) + h = w.malloc(64) + w.free(h) chip.free.reset_mock() with pytest.raises(ValueError, match="already-freed/stale"): - w.free(0x2000) + w.free(h) chip.free.assert_not_called() def test_l2_free_revokes_before_native_free(self): # L2 mirrors the L3 commit barrier: revoke before the native free. w, chip = self._l2() - w.malloc(64) + h = w.malloc(64) seen = {} chip.free.side_effect = lambda p: seen.__setitem__("live_at_native", (0, 0x2000) in w._child_alloc_prov) - w.free(0x2000) + w.free(h) assert seen["live_at_native"] is False def test_l2_copy_to_requires_live_dst(self): w, chip = self._l2() with pytest.raises(ValueError, match="not a live allocation"): - w.copy_to(0x2000, 0xABCD, 64) + w.copy_to(_dev_handle(0x2000, wid=0), _HOSTSRC) chip.copy_to.assert_not_called() w.malloc(64) - w.copy_to(0x2000, 0xABCD, 64) + w.copy_to(_dev_handle(0x2000, wid=0), _HOSTSRC) chip.copy_to.assert_called_once() @@ -421,56 +444,45 @@ def test_l2_copy_to_requires_live_dst(self): class TestProvenanceTransactions: - def test_orch_malloc_native_error_records_nothing(self): + def test_alloc_child_native_error_records_nothing(self): # Provenance is recorded only after the backend malloc succeeds. - w = _l3() - fake = MagicMock() - fake.malloc.side_effect = RuntimeError("device OOM") - o = Orchestrator(fake, w) + w, nw = _l3_ready() + nw.malloc.side_effect = RuntimeError("device OOM") with pytest.raises(RuntimeError, match="device OOM"): - o.malloc(0, 64) + w.alloc_child_tensor(0, (64,), DataType.UINT8) assert w._child_alloc_prov == {} - def test_orch_free_revokes_before_native_free(self): + def test_free_revokes_before_native_free(self): # Safety-first commit barrier: provenance is revoked BEFORE the native # free, so an async unwind after a successful free cannot leave a freed # address live. - w = _l3() - fake = MagicMock() - fake.malloc.return_value = 0x1000 - o = Orchestrator(fake, w) - o.malloc(0, 64) + w, nw = _l3_ready(0x1000) + h = w.alloc_child_tensor(0, (64,), DataType.UINT8) seen = {} - fake.free.side_effect = lambda wid, p: seen.__setitem__("live_at_native", (0, 0x1000) in w._child_alloc_prov) - o.free(0, 0x1000) + nw.free.side_effect = lambda wid, p: seen.__setitem__("live_at_native", (0, 0x1000) in w._child_alloc_prov) + w.free(h) assert seen["live_at_native"] is False # already revoked when native free runs - def test_orch_free_native_error_revokes_provenance_safe_first(self): + def test_free_native_error_revokes_provenance_safe_first(self): # A native free that fails becomes a terminal leak — provenance is # revoked, never re-authorized. No retry (the address is no longer a # live malloc base). - w = _l3() - fake = MagicMock() - fake.malloc.return_value = 0x1000 - o = Orchestrator(fake, w) - o.malloc(0, 64) - fake.free.side_effect = RuntimeError("free failed") + w, nw = _l3_ready(0x1000) + h = w.alloc_child_tensor(0, (64,), DataType.UINT8) + nw.free.side_effect = RuntimeError("free failed") with pytest.raises(RuntimeError, match="free failed"): - o.free(0, 0x1000) + w.free(h) assert (0, 0x1000) not in w._child_alloc_prov # revoked (terminal leak) - fake.free.side_effect = None + nw.free.side_effect = None with pytest.raises(ValueError, match="not a live malloc base"): - o.free(0, 0x1000) + w.free(h) def test_free_holds_lock_across_native_free(self): # Deterministic mutual-exclusion check: the native free runs while # _child_prov_lock is held, so a concurrent free/copy/dispatch cannot # interleave with a half-completed free. - w = _l3() - fake = MagicMock() - fake.malloc.return_value = 0x1000 - o = Orchestrator(fake, w) - o.malloc(0, 64) + w, nw = _l3_ready(0x1000) + h = w.alloc_child_tensor(0, (64,), DataType.UINT8) held = {} def _sf(wid, p): @@ -479,8 +491,8 @@ def _sf(wid, p): if acquired: w._child_prov_lock.release() - fake.free.side_effect = _sf - o.free(0, 0x1000) + nw.free.side_effect = _sf + w.free(h) assert held["locked_during_native"] is True def test_capture_refs_after_provenance_analysis(self, _fake_handle, monkeypatch): diff --git a/tests/ut/py/test_worker/test_dynamic_alloc_hw.py b/tests/ut/py/test_worker/test_dynamic_alloc_hw.py index d90ae33468..63728ab1b0 100644 --- a/tests/ut/py/test_worker/test_dynamic_alloc_hw.py +++ b/tests/ut/py/test_worker/test_dynamic_alloc_hw.py @@ -77,7 +77,7 @@ def orch_fn(orch, _args, _cfg): "domain_size": tp[chip_idx].domain_size, "device_ctx": int(tp[chip_idx].device_ctx), "local_window_base": int(tp[chip_idx].local_window_base), - "buffer_ptrs": dict(tp[chip_idx].buffer_ptrs), + "buffer_bases": {name: h.base for name, h in tp[chip_idx].buffers.items()}, } for chip_idx in tp.workers } @@ -115,7 +115,7 @@ def orch_fn(orch, _args, _cfg): assert ctx["device_ctx"] != 0, f"chip {chip_idx}: device_ctx is 0" assert ctx["local_window_base"] != 0, f"chip {chip_idx}: local_window_base is 0" # Buffers are carved sequentially from the local pool. - ptrs = ctx["buffer_ptrs"] + ptrs = ctx["buffer_bases"] assert isinstance(ptrs, dict) assert ptrs["scratch"] == ctx["local_window_base"] assert ptrs["signal"] == ctx["local_window_base"] + 64 diff --git a/tests/ut/py/test_worker/test_dynamic_alloc_sim.py b/tests/ut/py/test_worker/test_dynamic_alloc_sim.py index 7e0cbe4b6b..164b8b83f0 100644 --- a/tests/ut/py/test_worker/test_dynamic_alloc_sim.py +++ b/tests/ut/py/test_worker/test_dynamic_alloc_sim.py @@ -92,7 +92,7 @@ def orch_fn(orch, _args, _cfg): "domain_size": tp[chip_idx].domain_size, "device_ctx": int(tp[chip_idx].device_ctx), "local_window_base": int(tp[chip_idx].local_window_base), - "buffer_ptrs": dict(tp[chip_idx].buffer_ptrs), + "buffer_bases": {name: h.base for name, h in tp[chip_idx].buffers.items()}, } for chip_idx in tp.workers } @@ -134,8 +134,8 @@ def orch_fn(orch, _args, _cfg): # (i.e. non-zero). assert ctx0["local_window_base"] != 0 assert ctx1["local_window_base"] != 0 - scratch0 = ctx0["buffer_ptrs"] - scratch1 = ctx1["buffer_ptrs"] + scratch0 = ctx0["buffer_bases"] + scratch1 = ctx1["buffer_bases"] assert isinstance(scratch0, dict) and scratch0["scratch"] == ctx0["local_window_base"] assert isinstance(scratch1, dict) and scratch1["scratch"] == ctx1["local_window_base"] diff --git a/tests/ut/py/test_worker/test_group_task.py b/tests/ut/py/test_worker/test_group_task.py index 0671b8ccc9..92e46add53 100644 --- a/tests/ut/py/test_worker/test_group_task.py +++ b/tests/ut/py/test_worker/test_group_task.py @@ -27,14 +27,16 @@ import struct from multiprocessing.shared_memory import SharedMemory +from simpler.buffer_handle import mint_owner_instance_id, wrap_fork_inherited from simpler.task_interface import ( - DataType, TaskArgs, - Tensor, TensorArgType, ) from simpler.worker import Worker +_U8 = 5 # DataType.UINT8 value +_OID = mint_owner_instance_id() + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -54,12 +56,13 @@ def _read(shm: SharedMemory) -> int: def _sync_args(ptr: int, tag: TensorArgType) -> TaskArgs: """Build a TaskArgs whose only purpose is to register a synthetic - tensor-pointer key with the TensorMap so a downstream task can wire a - dep on it. SUB callables don't actually read tensors, so the pointer - value just needs to be a unique non-zero key. + dependency key with the TensorMap so a downstream task can wire a + dep on it. SUB callables don't actually read data, so ``ptr`` just + needs to yield a unique canonical identity (here the handle buffer_id). """ args = TaskArgs() - args.add_tensor(Tensor.make(ptr, (1,), DataType.UINT8), tag) + ref = wrap_fork_inherited(ptr, 1, _OID, buffer_id=ptr).ref(shapes=(1,), dtype=_U8) + args.add_ref(ref, 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 deleted file mode 100644 index fde974cc0b..0000000000 --- a/tests/ut/py/test_worker/test_host_addr_rewrite.py +++ /dev/null @@ -1,73 +0,0 @@ -# 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. -# ----------------------------------------------------------------------------------------------------------- -"""Host-pointer blob rewrite must not touch child_memory (device) tensors. - -``_rewrite_blob_host_addrs`` redirects registered host pointers in a task-args -blob into a forked child's own mapping. A tensor whose ``buffer.addr`` is a -child-owned device pointer (``child_memory == 1``) carries a device VA, which -may numerically fall inside a registered host range yet must never be rewritten -— doing so corrupts the device pointer. The rewrite therefore keys on the -per-tensor ``child_memory`` flag, not on the numeric address alone. -""" - -import struct - -from _task_interface import TENSOR_CHILD_MEMORY_OFFSET, TENSOR_STRIDE_BYTES -from simpler.worker import _BLOB_HEADER_BYTES, _rewrite_blob_host_addrs - -_PARENT_LO = 0x7F00_0000_0000 -_PARENT_HI = 0x7F00_0010_0000 -_CHILD_BASE = 0x7E00_0000_0000 - - -def _make_blob(tensors: list[tuple[int, int]]) -> bytearray: - """Build a task-args blob: [int32 T][int32 S][Tensor*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) - are populated; the rest of each 128-byte tensor stays zero. - """ - n = len(tensors) - buf = bytearray(_BLOB_HEADER_BYTES + n * TENSOR_STRIDE_BYTES) - struct.pack_into(" int: - off = _BLOB_HEADER_BYTES + i * TENSOR_STRIDE_BYTES - return struct.unpack_from("= 64 - finally: - buf.buffer.release() - w.free_host_buffer(buf) - finally: - w.close() - - def test_childless_l3_rejects_create_host_buffer(self): - # No callable registered: a childless L3 with a pre-registered callable - # is rejected earlier, at init() (see TestEligibleTargetPrecheck). - w = Worker(level=3, num_sub_workers=0) - w.init() - try: - with pytest.raises(RuntimeError, match="at least one forked chip or sub child"): - w.create_host_buffer(64) - finally: - w.close() diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index ad7a4cbc61..2689b9eee9 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -2196,7 +2196,7 @@ def orch(o, args, cfg): # Mixed with submits. with o.scope(): inner = o.alloc((32,), DataType.FLOAT32) - assert inner.data != 0 + assert inner.base != 0 hw.run(orch) hw.close() @@ -2244,12 +2244,13 @@ def test_alloc_returns_valid_tensor(self): def orch(o, args, cfg): inter = o.alloc((64,), DataType.FLOAT32) - captured.append((inter.data, inter.ndims, inter.shapes[0])) + ref = inter.ref((64,), DataType.FLOAT32) + captured.append((inter.base, ref.ndims, ref.shapes[0])) # Tag as OUTPUT in some submit so the synthetic alloc slot has a # downstream consumer (otherwise scope_end consumes alone — still fine). sub_args = TaskArgs() - sub_args.add_tensor(inter, TensorArgType.INPUT) + sub_args.add_ref(ref, TensorArgType.INPUT) o.submit_sub(handle, sub_args) hw.run(orch) @@ -2281,13 +2282,13 @@ def orch(o, args, cfg): # do TensorMap.lookup. Plain OUTPUT / OUTPUT_EXISTING are # pure inserts and would leave no dep on the alloc slot. p_args = TaskArgs() - p_args.add_tensor(inter, TensorArgType.INOUT) + p_args.add_ref(inter.ref((128,), DataType.FLOAT32), TensorArgType.INOUT) o.submit_sub(producer_handle, p_args) - # Consumer tags inter as INPUT — tensormap.lookup finds the - # producer slot, dep wired automatically. + # Consumer tags inter as INPUT — dep inference keys on the ref's + # canonical identity (shared with the producer), dep wired automatically. c_args = TaskArgs() - c_args.add_tensor(inter, TensorArgType.INPUT) + c_args.add_ref(inter.ref((128,), DataType.FLOAT32), TensorArgType.INPUT) o.submit_sub(consumer_handle, c_args) hw.run(orch) @@ -2328,7 +2329,7 @@ def test_alloc_across_runs_does_not_leak(self): def orch(o, args, cfg): inter = o.alloc((64,), DataType.FLOAT32) args = TaskArgs() - args.add_tensor(inter, TensorArgType.INPUT) + args.add_ref(inter.ref((64,), DataType.FLOAT32), TensorArgType.INPUT) o.submit_sub(handle, args) for _ in range(8): @@ -2348,30 +2349,30 @@ 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 + """Sub callable receives MappedArgs with correct tensor count and shape.""" + from simpler.buffer_handle import mint_owner_instance_id, wrap_fork_inherited # noqa: PLC0415 result_shm, result_buf = _make_shared_counter() try: hw = Worker(level=3, num_sub_workers=1) def check_args(args): - # Verify args decoded correctly: 1 tensor, shape (4,), FLOAT32 - if args.tensor_count() == 1 and args.scalar_count() == 0: - t = args.tensor(0) - if t.ndims == 1 and t.shapes[0] == 4: + # Verify args decoded correctly: 1 tensor, shape (4,), and no scalars. + if len(args) == 1 and args.scalar_count() == 0: + t = args[0] + if len(t.shapes) == 1 and t.shapes[0] == 4: _increment_counter(result_buf) handle = hw.register(check_args) hw.init() - # Use a synthetic non-zero pointer — sub callable only checks metadata, - # doesn't dereference the pointer. - ct = Tensor.make(0xCAFE0000, (4,), DataType.FLOAT32) + # Use a synthetic non-zero pointer — the sub callable only checks metadata (shapes), + # never dereferences the buffer, so a FORK_SHM ref over a fake VA is enough. + cref = wrap_fork_inherited(0xCAFE0000, 16, mint_owner_instance_id(), 1, "L3").ref((4,), DataType.FLOAT32) def orch(o, args, cfg): sub_args = TaskArgs() - sub_args.add_tensor(ct, TensorArgType.INPUT) + sub_args.add_ref(cref, TensorArgType.INPUT) o.submit_sub(handle, sub_args) hw.run(orch) 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 fbcd7a95c4..7f278b1986 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 @@ -8,6 +8,7 @@ # ----------------------------------------------------------------------------------------------------------- import ctypes +import itertools import math import struct from dataclasses import dataclass @@ -17,6 +18,7 @@ import pytest from simpler import l3_l2_orch_comm from simpler import worker as worker_module +from simpler.buffer_handle import AccessMode, CanonicalIdentity, mint_owner_instance_id, wrap_fork_inherited from simpler.l3_l2_message_queue import ( L3L2_QUEUE_COUNTER_BYTES, L3L2_QUEUE_DESC_SLOT_BYTES, @@ -36,9 +38,20 @@ WaitCmp, ) from simpler.orchestrator import Orchestrator -from simpler.task_interface import DataType, Tensor, get_element_size +from simpler.task_interface import DataType, get_element_size from simpler.worker import _IDLE, _OFF_STATE, Worker, _buffer_field_addr, _mailbox_store_i32 +_DESC_BID = itertools.count(1) + + +def _fake_alloc_handle(orch, nbytes): + """A FORK_SHM BufferHandle over a bare _FakeCOrch alloc — mirrors Orchestrator.alloc for the + low-level tests that drive L3L2Queue with a fake C orch directly.""" + oid, bid = mint_owner_instance_id(), next(_DESC_BID) + identity = CanonicalIdentity(oid, bid, "L3", 0) + va = int(orch.alloc([nbytes], DataType.UINT8, identity.pack())) + return wrap_fork_inherited(va, nbytes, oid, bid, "L3", access=AccessMode.READWRITE) + @dataclass(frozen=True) class _FakeRequest: @@ -103,15 +116,14 @@ def __init__(self): self._buffers = [] self.fail_next_alloc = False - def alloc(self, shape, dtype): + def alloc(self, shape, dtype, identity): if self.fail_next_alloc: self.fail_next_alloc = False raise RuntimeError("injected allocation failure") nbytes = math.prod(int(x) for x in shape) * int(get_element_size(dtype)) - storage_t = ctypes.c_uint8 * nbytes - storage = storage_t() + storage = (ctypes.c_uint8 * nbytes)() self._buffers.append(storage) - return Tensor.make(ctypes.addressof(storage), tuple(int(x) for x in shape), dtype) + return ctypes.addressof(storage) class _FakeClient: @@ -400,9 +412,9 @@ def test_create_l3_l2_queue_allocates_region_and_exposes_l2_task_scalars(): def test_create_l3_l2_queue_frees_region_on_post_region_alloc_failure(): orch, worker, shm, _fake_client = _make_orchestrator() - original_alloc = orch._o.alloc + original_alloc_ref = orch._o.alloc - def fail_alloc(_shape, _dtype): + def fail_alloc(_shape, _dtype, _identity): raise RuntimeError("injected alloc failure") orch._o.alloc = fail_alloc @@ -413,7 +425,7 @@ def fail_alloc(_shape, _dtype): assert len(worker._live_l3_l2_regions) == 1 assert worker._live_l3_l2_regions[0]._released is True finally: - orch._o.alloc = original_alloc + orch._o.alloc = original_alloc_ref _close(worker, shm) @@ -526,9 +538,9 @@ def test_l3_host_mapped_queue_ordinary_input_uses_direct_payload_write(monkeypat orch, region, layout, - orch.alloc([24], DataType.UINT8), - orch.alloc([8], DataType.UINT8), - orch.alloc([L3L2_QUEUE_DESC_SLOT_BYTES], DataType.UINT8), + _fake_alloc_handle(orch, 24), + _fake_alloc_handle(orch, 8), + _fake_alloc_handle(orch, L3L2_QUEUE_DESC_SLOT_BYTES), ) alloc_count = len(orch._buffers) payload_writes: list[tuple[int, bytes]] = [] @@ -586,7 +598,7 @@ def test_output_read_into_registered_tensor_uses_fast_path_and_release_notifies_ queue.output.read_into(handle, output) queue.output.release(handle) - assert ctypes.string_at(int(output.data), 16) == b"abcdefghijklmnop" + assert ctypes.string_at(int(output.base), 16) == b"abcdefghijklmnop" assert fake_client.counters[queue.layout.output_desc_head_offset] == 1 finally: _close(worker, shm) @@ -603,7 +615,7 @@ def test_dequeue_into_reads_and_releases_output(): assert message.seq == 1 assert message.opcode == L3L2QueueOpcode.DATA - assert ctypes.string_at(int(output.data), 16) == b"abcdefghijklmnop" + assert ctypes.string_at(int(output.base), 16) == b"abcdefghijklmnop" assert fake_client.counters[queue.layout.output_desc_head_offset] == 1 finally: _close(worker, shm) @@ -619,7 +631,7 @@ def test_output_error_opcode_is_delivered_without_poison(): message = queue.output.dequeue_into(output, timeout=0.001) assert message.opcode == L3L2QueueOpcode.ERROR - assert ctypes.string_at(int(output.data), 12) == b"error-detail" + assert ctypes.string_at(int(output.base), 12) == b"error-detail" assert fake_client.counters[queue.layout.output_desc_head_offset] == 1 assert fake_client.counters.get(L3L2_QUEUE_L3_ABORT_FLAG_OFFSET, 0) == 0 finally: 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 10531fb690..53630c824e 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 @@ -17,12 +17,13 @@ import pytest from simpler import l3_l2_orch_comm from simpler import worker as worker_module +from simpler.buffer_handle import mint_owner_instance_id, wrap_fork_inherited from simpler.l3_l2_orch_comm import ( NotifyOp, SignalTestResult, WaitCmp, ) -from simpler.task_interface import DataType, Tensor +from simpler.task_interface import DataType from simpler.worker import ( _IDLE, _OFF_STATE, @@ -219,7 +220,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 = wrap_fork_inherited(0x1234_0000, 16, mint_owner_instance_id(), 1, "L3") region.payload_write(0, payload, nbytes=8) region.payload_read(8, payload, nbytes=8) result = region.counter(64).test(7, WaitCmp.EQ) @@ -570,7 +571,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 = wrap_fork_inherited(0x1234_0000, 16, mint_owner_instance_id(), 1, "L3") with pytest.raises(RuntimeError, match="copy failed"): region.payload_write(0, payload, nbytes=8) with pytest.raises(RuntimeError, match="poisoned"): @@ -632,7 +633,7 @@ def test_sim_worker_region_payload_roundtrip(platform): def orch(orch_handle, _args, _cfg): host = orch_handle.alloc([16], DataType.UINT8) buf_t = ctypes.c_uint8 * 16 - buf = buf_t.from_address(int(host.data)) + buf = buf_t.from_address(int(host.base)) for i in range(16): buf[i] = (i + 41) & 0xFF region = orch_handle.create_l3_l2_region(worker_id=0, payload_bytes=16, counter_bytes=128) diff --git a/tests/ut/py/test_worker/test_l4_recursive.py b/tests/ut/py/test_worker/test_l4_recursive.py index 969919aa44..a01cbf7a9f 100644 --- a/tests/ut/py/test_worker/test_l4_recursive.py +++ b/tests/ut/py/test_worker/test_l4_recursive.py @@ -19,10 +19,20 @@ from multiprocessing.shared_memory import SharedMemory import pytest +from _task_interface import DataType +from simpler.buffer_handle import mint_owner_instance_id, wrap_device_malloc from simpler.callable_identity import CallableHandle from simpler.task_interface import CallConfig, TaskArgs from simpler.worker import Worker +_OID = mint_owner_instance_id() +_HOSTBUF = bytearray(64) + + +def _dev_handle(ptr: int, *, wid: int = 0): + return wrap_device_malloc(ptr, 64, _OID, buffer_id=ptr, owner_worker_id=wid) + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -128,25 +138,26 @@ def test_add_worker_with_device_ids_rejected(self): with pytest.raises(RuntimeError, match="cannot be combined with device_ids"): w4.add_worker(child) - def test_malloc_on_l4_raises_index_error(self): - # L4 has no chip mailboxes — `Worker.malloc` must surface IndexError - # rather than silently dispatch CTRL_MALLOC to a next_level (L3 worker) - # child whose `_child_worker_loop` doesn't recognise CTRL_MALLOC and - # would return a garbage pointer from an uninitialised mailbox result - # slot. + def test_device_mem_on_l4_rejected(self): + # L4 has no chip mailboxes — device memory ops must be rejected rather than silently dispatch + # CTRL_MALLOC/FREE/COPY to a next_level (L3 worker) child whose `_child_worker_loop` doesn't + # recognise them and would return a garbage pointer from an uninitialised mailbox result slot. + # `malloc` is L2-only (TypeError); the child-device ops guard the worker id (IndexError). l3_child = Worker(level=3, num_sub_workers=0) w4 = Worker(level=4, num_sub_workers=0) w4.add_worker(l3_child) w4.init() try: + with pytest.raises(TypeError, match="L2-only"): + w4.malloc(1024) with pytest.raises(IndexError, match="out of range"): - w4.malloc(1024, worker_id=0) + w4.alloc_child_tensor(0, (256,), DataType.FLOAT32) with pytest.raises(IndexError, match="out of range"): - w4.free(0xDEADBEEF, worker_id=0) + w4.free(_dev_handle(0xDEADBEEF, wid=0)) with pytest.raises(IndexError, match="out of range"): - w4.copy_to(0xDEAD, 0xBEEF, 64, worker_id=0) + w4.copy_to(_dev_handle(0xDEAD, wid=0), _HOSTBUF) with pytest.raises(IndexError, match="out of range"): - w4.copy_from(0xDEAD, 0xBEEF, 64, worker_id=0) + w4.copy_from(_HOSTBUF, _dev_handle(0xBEEF, wid=0)) finally: w4.close() diff --git a/tests/ut/py/test_worker/test_l4_reexport.py b/tests/ut/py/test_worker/test_l4_reexport.py new file mode 100644 index 0000000000..473e5c89de --- /dev/null +++ b/tests/ut/py/test_worker/test_l4_reexport.py @@ -0,0 +1,96 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""L4 -> L3 -> sub re-export end-to-end (P1-B B3). + +An L4-owned create_buffer backing reaches L3 as a BufferRef; L3's orch forwards it re-exported to a +handle H' that keeps L4's canonical identity (no BufferRef pass-through, no map on the forwarding hop), +to its sub; the sub (a compute leaf) maps H' and writes through it. The write lands in the shared +backing L4 owns. +No NPU device — L3 uses a SubWorker. +""" + +import torch +from simpler.buffer_handle import ( + AccessMode, + AddressSpace, + BackendKind, + BufferHandleDescriptor, + BufferRef, + CanonicalIdentity, + Visibility, +) +from simpler.task_interface import CallConfig, TaskArgs, TensorArgType +from simpler.worker import Worker + +_F32 = 0 # DataType.FLOAT32 value + + +def test_l4_l3_reexport_to_sub(): + def l3_sub(args): + a = torch.frombuffer(args[0].buffer, dtype=torch.float32, count=4) + a.add_(1.0) # compute leaf writes through the mapped, re-exported backing + + l3 = Worker(level=3, num_sub_workers=1) + l3_sub_handle = l3.register(l3_sub) + + def l3_orch(orch, args, config): + # args are re-exported BufferRefs owned by L3 (each level sees only its own handles). + sa = TaskArgs() + sa.add_ref(args[0], TensorArgType.INOUT) + orch.submit_sub(l3_sub_handle, sa) + + w4 = Worker(level=4, num_sub_workers=1) + w4.register(lambda args: None) # a sub child so create_buffer has a forked child to satisfy + l3_orch_handle = w4.register(l3_orch) + w4.add_worker(l3) + w4.init() + + t = None + try: + buf_h = w4.create_buffer(16) # L4-owned POSIX shm, allocated post-init + shm = buf_h.shm + assert shm is not None + t = torch.frombuffer(shm.buf, dtype=torch.float32, count=4) + t.fill_(5.0) + + def l4_orch(orch, args, config): + ta = TaskArgs() + ta.add_ref(buf_h.ref(shapes=(4,), dtype=_F32), TensorArgType.INOUT) + orch.submit_next_level(l3_orch_handle, ta, CallConfig(), worker=0) + + w4.run(l4_orch) + assert torch.allclose(t, torch.full((4,), 6.0)), t.tolist() + finally: + t = None + w4.close() + + +def test_reexport_preserves_canonical_identity(): + # Frozen model §5/§8: the canonical identity is invariant across every edge. Re-exporting a + # forwarded upper-level backing keeps the SOURCE (owner_instance_id, owner_worker_path, buffer_id, + # generation) — only the mapping is stripped (base=0, shm=None) on the forwarding hop. Dependency + # inference keys on the identity, so an alias / retain-release must not split across L4→L3→L2. + src_identity = CanonicalIdentity(b"\x9f" * 16, 42, "L4", 1) + source = BufferHandleDescriptor( + identity=src_identity, + address_space=AddressSpace.HOST, + visibility=Visibility.SHARED, + access=AccessMode.READWRITE, + backend_kind=BackendKind.FORK_SHM, + nbytes=1024, + body=(0xDEAD0000).to_bytes(8, "little"), + ) + w = Worker(level=3, num_sub_workers=0) + h_prime = w._reexport(source) + assert h_prime.identity.pack() == src_identity.pack() # identity invariant across the edge + assert h_prime.base == 0 and h_prime.shm is None # no map on the forwarding hop + assert h_prime.backend_kind == source.backend_kind and h_prime.body == source.body + # A ref built over H' carries the same identity a downstream consumer keys deps on. + r = BufferRef.unpack(h_prime.ref((4,), _F32).pack()) + assert r.handle.identity.pack() == src_identity.pack() diff --git a/tests/ut/py/test_worker/test_startup_readiness.py b/tests/ut/py/test_worker/test_startup_readiness.py index 82be5ed568..55b0ab9bef 100644 --- a/tests/ut/py/test_worker/test_startup_readiness.py +++ b/tests/ut/py/test_worker/test_startup_readiness.py @@ -689,6 +689,14 @@ def owner_body(): class _FakeChipOk: """Stand-in for a ChipWorker whose init succeeds — no NPU touched.""" + class _Impl: + # The L2 in-process path materializes BufferRef args to a Tensor blob and dispatches + # through `_impl.run_from_blob`; stub it so no NPU is touched. + def run_from_blob(self, *_a, **_k): + pass + + _impl = _Impl() + def init(self, *_a, **_k): pass @@ -938,7 +946,7 @@ def test_reap_deadline_starts_after_shutdown_broadcast(self, monkeypatch): w._worker = types.SimpleNamespace(close=lambda: None) # look "started" for the L3 branch # A slow pre-child cleanup step (runs before the SHUTDOWN broadcast). - monkeypatch.setattr(Worker, "_release_all_host_buffers", lambda self: time.sleep(0.6)) + monkeypatch.setattr(Worker, "_release_all_buffer_handles", lambda self: time.sleep(0.6)) captured: dict = {} def capture_reap(groups, deadline): diff --git a/tests/ut/py/test_worker/test_sub_bufferref.py b/tests/ut/py/test_worker/test_sub_bufferref.py new file mode 100644 index 0000000000..b726cbc75d --- /dev/null +++ b/tests/ut/py/test_worker/test_sub_bufferref.py @@ -0,0 +1,246 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Sub-worker task args over the BufferRef wire (P1-B B3). + +A Python sub callable is a compute leaf: it receives its args as MappedArgs (each backing mapped into +the sub child, map-once) and computes with torch.frombuffer(arg.buffer, ...). Writes land in the +shared backing the owner sees — no C++ Tensor involved. This is exactly the case a closure cannot +serve: a post-init create_buffer shm is not mapped in the pre-forked sub child, so the buffer must +arrive via args and be mapped from its Ref. +""" + +import ctypes + +import pytest +import torch +from simpler.buffer_handle import ( + BackendKind, + BufferRef, + ImportRegistry, + mint_owner_instance_id, + pack_bufferref_blob, + wrap_fork_inherited, +) +from simpler.task_interface import ( + CallConfig, + DataType, + TaskArgs, + TensorArgType, + read_args_from_blob, +) +from simpler.worker import Worker + +_F32 = 0 # DataType.FLOAT32 value + + +def test_alloc_shared_tensor_returns_managed_fork_handle(): + hw = Worker(level=3, num_sub_workers=1) + hw.init() + captured = {} + try: + # alloc_shared_tensor allocates from the orchestrator's HeapRing, so it must run inside + # an active building run — drive it through a no-submit orch fn. + def orch(o, _args, cfg): + captured["h"] = hw.alloc_shared_tensor((4, 8), DataType.FLOAT32) + + hw.run(orch, args=None, config=CallConfig()) + h = captured["h"] + assert h.nbytes == 4 * 8 * 4 # prod(shape) * element_size + # Backed by the orchestrator's HeapRing VA (FORK_SHM, no POSIX shm), runtime-managed. + assert h.backend_kind == BackendKind.FORK_SHM + assert h.shm is None and h.base != 0 + finally: + hw.close() + + +def test_create_buffer_at_l2_needs_no_child(): + # An L2 leaf has no forked children — it materializes the ref in-process itself — so create_buffer + # must not require a child. The handle is a usable POSIX-shm backing. + w = Worker(level=2) + h = w._create_buffer_locked(64) + try: + assert h.nbytes == 64 + assert h.shm is not None + t = torch.frombuffer(h.shm.buf, dtype=torch.float32, count=4) + t.fill_(3.0) + assert t.tolist() == [3.0, 3.0, 3.0, 3.0] + t = None + # The level guard now admits L2: the public create_buffer no longer TypeErrors on level, it + # reaches the READY check (this Worker is uninitialized). + with pytest.raises(RuntimeError, match="READY"): + w.create_buffer(64) + finally: + h.close() + + +def test_l2_run_materializes_bufferref_to_tensor_blob(): + # An L2 leaf consumes its own BufferRef args: _run_l2_materialized resolves each ref to a local + # base and hands the runtime a Tensor blob (write_blob format), exactly like a chip child, minus + # the mailbox. Capture that blob via a fake ChipWorker and decode it to prove the materialization. + captured = {} + + class _FakeImpl: + def run_from_blob(self, cid, ptr, cap, cfg): + captured["cid"] = cid + captured["blob"] = ctypes.string_at(ptr, cap) + + class _FakeChip: + _impl = _FakeImpl() + + w = Worker(level=2) + w._chip_worker = _FakeChip() # type: ignore[assignment] + h = w._create_buffer_locked(16) # 4 x f32 + try: + shm = h.shm + assert shm is not None + torch.frombuffer(shm.buf, dtype=torch.float32, count=4).fill_(7.0) + ta = TaskArgs() + ta.add_ref(h.ref(shapes=(4,), dtype=_F32), TensorArgType.INPUT) + w._run_l2_materialized(3, ta, CallConfig()) + + assert captured["cid"] == 3 + decode_buf = ctypes.create_string_buffer(captured["blob"], len(captured["blob"])) + args = read_args_from_blob(ctypes.addressof(decode_buf)) + assert args.tensor_count() == 1 + t = args.tensor(0) + assert tuple(t.shapes[: t.ndims]) == (4,) + assert t.data != 0 # resolved to a real local base + # The materialized base maps the same physical pages the owner wrote through. + mapped = torch.frombuffer((ctypes.c_float * 4).from_address(t.data), dtype=torch.float32, count=4) + assert mapped.tolist() == [7.0, 7.0, 7.0, 7.0] + finally: + h.close() + w._close_l2_import_registry() + + +def test_mapped_args_from_blob_delivers_tensors_and_scalars(): + # A sub callable's args expose both the mapped tensors (args[i].buffer) and the blob's scalars + # (args.scalar_count() / args.scalar(i)) — a sub task built with add_ref + add_scalar reaches both. + backing = (ctypes.c_float * 4)(1.0, 2.0, 3.0, 4.0) + va = ctypes.addressof(backing) + ref = wrap_fork_inherited(va, 16, mint_owner_instance_id(), 1, "L3").ref((4,), _F32) + blob = pack_bufferref_blob([ref], (17, 99)) + buf = ctypes.create_string_buffer(blob, len(blob)) + args = ImportRegistry().mapped_args_from_blob(ctypes.addressof(buf), len(blob)) + assert len(args) == 1 and args.tensor_count() == 1 + assert args.scalar_count() == 2 + assert args.scalar(0) == 17 and args.scalar(1) == 99 + assert torch.frombuffer(args[0].buffer, dtype=torch.float32, count=4).tolist() == [1.0, 2.0, 3.0, 4.0] + + +def test_make_ref_arg_memoizes_handle_per_storage(): + # Every ref over the same torch storage shares one FORK_SHM handle/identity (so deps key on it); + # a view's byte_offset places it within that shared backing. + w = Worker(level=3, num_sub_workers=0) # no init needed: make_ref_arg only reads owner-side state + t = torch.zeros(16, dtype=torch.float32) + r1 = BufferRef.unpack(w.make_ref_arg(t, shapes=(16,), dtype=_F32).pack()) + r2 = BufferRef.unpack(w.make_ref_arg(t[4:12], shapes=(8,), dtype=_F32).pack()) # slice of same storage + assert r1.handle.backend_kind == BackendKind.FORK_SHM + assert r1.handle.identity.pack() == r2.handle.identity.pack() # one memoized handle + assert r1.byte_offset == 0 + assert r2.byte_offset == 16 # t[4:12] starts 4 float32 = 16 B into the storage + assert len(w._fork_tensor_handles) == 1 + + +def test_make_ref_arg_share_memory_output_across_fork(): + # A pre-fork share_memory_() tensor is MAP_SHARED: a forked child's writes land in the same + # physical pages the parent reads. make_ref_arg names it (FORK_SHM), the sub child writes through + # the inherited VA, and the parent sees the result — the fork-inherited read-write path. + t = torch.zeros(4, dtype=torch.float32).share_memory_() # allocated + shared BEFORE init/fork + + def sub_fn(args): + a = torch.frombuffer(args[0].buffer, dtype=torch.float32, count=4) + a.add_(1.0) + + hw = Worker(level=3, num_sub_workers=1) + handle = hw.register(sub_fn) + hw.init() # forks the sub child, inheriting t's MAP_SHARED mapping at the same VA + try: + t.fill_(5.0) + + def orch(o, _args, cfg): + sa = TaskArgs() + sa.add_ref(hw.make_ref_arg(t, shapes=(4,), dtype=_F32), TensorArgType.INOUT) + o.submit_sub(handle, sa) + + hw.run(orch, args=None, config=CallConfig()) + assert torch.allclose(t, torch.full((4,), 6.0)), t.tolist() + finally: + hw.close() + + +def test_alloc_shared_tensor_managed_intermediate_deps_and_share(): + # A runtime-managed HeapRing intermediate (alloc_shared_tensor) named as a BufferRef: sub-A writes + # it (INOUT → depends on the alloc slot + becomes producer), sub-B reads it (INPUT → depends on the + # producer) and copies it into a create_buffer output the parent reads. The dep wires via the ref's + # canonical identity; the ring VA is MAP_SHARED so B sees A's write; the slot auto-reclaims at scope + # end (drain completes, no hang). Result crosses back via the shared `out` buffer, not a dict (the + # subs run in forked processes). + def producer(args): + torch.frombuffer(args[0].buffer, dtype=torch.float32, count=4).fill_(9.0) + + def consumer(args): + src = torch.frombuffer(args[0].buffer, dtype=torch.float32, count=4) + torch.frombuffer(args[1].buffer, dtype=torch.float32, count=4).copy_(src) + + hw = Worker(level=3, num_sub_workers=2) + ph = hw.register(producer) + ch = hw.register(consumer) + hw.init() + out_t = None + try: + out_h = hw.create_buffer(16) # 4 x f32 — the parent-observable output + shm = out_h.shm + assert shm is not None + out_t = torch.frombuffer(shm.buf, dtype=torch.float32, count=4) + out_t.zero_() + + def orch(o, _args, cfg): + inter = hw.alloc_shared_tensor((4,), DataType.FLOAT32) # worker captured in the closure + pa = TaskArgs() + pa.add_ref(inter.ref(shapes=(4,), dtype=_F32), TensorArgType.INOUT) + o.submit_sub(ph, pa) + ca = TaskArgs() + ca.add_ref(inter.ref(shapes=(4,), dtype=_F32), TensorArgType.INPUT) + ca.add_ref(out_h.ref(shapes=(4,), dtype=_F32), TensorArgType.OUTPUT_EXISTING) + o.submit_sub(ch, ca) + + hw.run(orch, args=None, config=CallConfig()) + assert torch.allclose(out_t, torch.full((4,), 9.0)), out_t.tolist() + finally: + out_t = None + hw.close() + + +def test_sub_worker_mapped_arg_readwrite(): + def sub_fn(args): + a = torch.frombuffer(args[0].buffer, dtype=torch.float32, count=4) + a.add_(1.0) # write through the mapped shared buffer + + hw = Worker(level=3, num_sub_workers=1) + handle = hw.register(sub_fn) + hw.init() + t = None + try: + buf_h = hw.create_buffer(16) # 4 x float32, POSIX shm allocated post-init + shm = buf_h.shm + assert shm is not None + t = torch.frombuffer(shm.buf, dtype=torch.float32, count=4) + t.fill_(5.0) + + def orch(o, args, cfg): + sa = TaskArgs() + sa.add_ref(buf_h.ref(shapes=(4,), dtype=_F32), TensorArgType.INOUT) + o.submit_sub(handle, sa) + + hw.run(orch, args=None, config=CallConfig()) + assert torch.allclose(t, torch.full((4,), 6.0)), t.tolist() + finally: + t = None + hw.close()