Skip to content

Fix: harden the Buffer/Tensor wire gate and the create_buffer child check - #1703

Merged
ChaoWao merged 3 commits into
hw-native-sys:mainfrom
ChaoWao:fix-buffer-abi-validation-gaps
Aug 5, 2026
Merged

Fix: harden the Buffer/Tensor wire gate and the create_buffer child check#1703
ChaoWao merged 3 commits into
hw-native-sys:mainfrom
ChaoWao:fix-buffer-abi-validation-gaps

Conversation

@ChaoWao

@ChaoWao ChaoWao commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Follow-up to #1599: the gate the ABI rests on

#1599 froze the byte layout and shipped validate_tensor as the receive-side gate — every later
boundary is meant to run it, and after the wire flip it runs on every argument that arrives. Three
things in it do not hold up, and a frozen ABI is the wrong place to leave them. None of this
changes a field, an offset, or an enum value: every static_assert in buffer.h is untouched.

validate_tensor accepts a view that spans exabytes over a 4-byte backing

tensor_extent_bytes summed (shapes[i]-1)*strides[i] into a uint64. Both fields are u32, so a
single product already reaches ~2^64, and the multiply by the element size overflows on top of that.
The extent wrapped to a small value and the bound check passed. One dimension is enough, and it
is reachable straight from Python:

h.tensor(shapes=(2147483649,), dtype=DataType.FLOAT32, strides=(2147483648,))  # on a 4-byte buffer
before:  tensor_extent_bytes = 4        validate_tensor(nbytes=4) -> ACCEPTED   # true span ~16 EiB
after:   tensor_extent_bytes = SAT      validate_tensor(nbytes=4) -> rejected

The arithmetic saturates, and an extent that lands on the sentinel is refused by name rather
than compared against nbytes — otherwise a descriptor claiming an absurd nbytes buys the view
back. This breaks design invariant 4 (view footprint ⊆ nbytes), which the design doc lists as
executable.

tensors_overlap inherited the same wrap, and there the consequence is worse than a rejected
argument: two fully overlapping views compare as disjoint, i.e. a missed dependency edge — the
old-address-key failure mode the identity key exists to remove. Its end offsets saturate too.

The 4096-iteration arbitrary-bytes pass did not catch this because its survivor assertions covered
magic / generation / body_len / ndims but not the footprint invariant — the one most at
risk. It asserts it now.

create_buffer refuses an L4 whose children are local L3 Workers

The L3+ child check counted _chip_shms and _sub_shms only. A next-level Worker child is a forked
process that maps a POSIX_SHM backing by name exactly as a chip or sub child does, so it can
consume the buffer. _next_level_shms counts now.

create_buffer had no test at all; tests/ut/py/test_worker/test_create_buffer.py covers the gate
in each child shape, the childless refusal, the L2 leaf, id uniqueness within one incarnation, and
the per-buffer best-effort release (a failing close is reported, its entry kept for the journal to
retry).

wrap_fork_inherited inferred the backend from access

FORK_SHM if access != READ else FORK_COW. The two are opposite kernel write semantics, not two
spellings of one grant — putting the physical semantics in backend_kind rather than leaving it to
be inferred from an orthogonal field is the reason they are separate tags. Inferring it makes a
read-only MAP_SHARED backing inexpressible: it is tagged FORK_COW, whose READ-only rule then
locks it there. The caller holds the mmap and is the only party that knows which it is, so it states
the tag; the default pair stays the safe one. The function had no caller, so nothing depends on the
old signature.

④ Text that describes machinery which does not exist

  • ImportRegistry.materialize pointed a caller holding raw bytes at BufferDescriptor.unpack, and
    tests/ut/py/test_buffer.py said it pins a pack/unpack round trip. Neither exists
    withholding the encoding is exactly what makes construction the only way in.
  • validate_tensor's comment claimed it stands behind materialization. It does not:
    ImportRegistry.materialize takes an already-decoded descriptor and adds no endpoint check. The
    comment now says so and points at the separate change.
  • docs/buffer-abi.md named the wire type simpler.task_interface.Tensor, contradicting the same
    page's status note, the module, and test_wire_tensor_stays_off_the_public_submit_surface.
  • The h.shm.buf sample now records that it is transitional and why: byte access belongs on the
    view, since a device backing has no shm and code written against one forks by backend.
  • BufferDescriptor::operator== bounded a memcmp by an unvalidated body_len; clamped. Not
    reachable today (construction validates), but it is the only length field in the header that
    bounds a read, and the cutover is what makes raw decoded descriptors exist.

Verification

  • cpput test_buffer15/15, including the two new overflow / overlap cases and the
    strengthened arbitrary-bytes pass
  • pyut full suite — 1142 passed, 13 skipped
  • tests/ut/py/test_buffer.py + the new test_create_buffer.py — 31 passed
  • Build green (nanobind extension + all four arch×runtime trees)
  • clang-format / ruff check / ruff format clean; no added line over 120 chars
  • Hardware (onboard) — not run; no runtime path is touched

Deliberately not here

Each is a separate concern, and the first two are the ones I would take next:

  • The endpoint × address_space matrix. materialize still returns a device VA to a host
    endpoint. That is a behavioural gate with its own test surface, not a comment fix.
  • The dead surface question. wrap_vmm_window, remote_sidecar_tensor, host_ptr_nbytes and
    tensors_overlap still have no caller. Whether they survive the freeze is the author's call, not
    something to settle in a fix PR.
  • BufferDescriptor / Tensor bind value equality without value hashing, and their __eq__ is not
    nb::is_operator(), so a comparison against a foreign type raises TypeError instead of
    returning NotImplemented.
  • remote_sidecar_tensor builds owner_instance_id from owner_worker_id, against the full-width
    random draw its own minting function requires.

🤖 Generated with Claude Code

@coderabbitai

coderabbitai Bot commented Aug 5, 2026

Copy link
Copy Markdown

Review Change Stack

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: e7569230-60cd-48f7-a43d-3af802b04dc4

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

The changes harden tensor extent and descriptor handling, add explicit fork backend selection, clarify buffer materialization rules, and allow valid next-level worker topologies to create shared buffers. Tests cover overflow, backend constraints, worker topology, identity, and cleanup.

Changes

Buffer ABI and worker integration

Layer / File(s) Summary
Tensor extent safety
src/common/task_interface/buffer.h, tests/ut/cpp/types/test_buffer.cpp, tests/ut/py/test_buffer.py
Tensor arithmetic now saturates on 64-bit overflow. Validation rejects unrepresentable extents. Equality bounds descriptor comparisons, and overlap checks prevent offset wraparound.
Fork backend contracts and materialization
python/simpler/buffer.py, docs/buffer-abi.md, tests/ut/py/test_buffer.py
wrap_fork_inherited accepts an explicit backend kind. Materialization boundaries and host-access rules are documented. Tests cover shared-memory, copy-on-write, and writable-descriptor constraints.
Hierarchical worker buffer lifecycle
python/simpler/worker.py, tests/ut/py/test_worker/test_create_buffer.py
Worker validation recognizes next-level forked children. Tests cover topology, buffer identity, size validation, and cleanup behavior.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

Poem

A rabbit checks each tensor’s span,
And bounds each buffer where it can.
Forked paths now choose their kind,
Shared or copied, well-defined.
Worker children hop in line,
While cleanup keeps the burrow fine.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 28.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly summarizes the two primary fixes: buffer/Tensor hardening and the create_buffer child check.
Description check ✅ Passed The description directly explains the buffer validation, worker topology, backend selection, documentation, tests, and verification changes.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/ut/cpp/types/test_buffer.cpp`:
- Around line 287-293: Update the test setup before the `b = a` assignment so
`a.byte_offset` is 1 instead of 0, ensuring the end-offset addition overflows
and the existing `tensor_extent_bytes(a)` and `tensors_overlap(a, b)` assertions
cover that branch.
🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 65703234-c971-493d-ad83-d5c835207ad5

📥 Commits

Reviewing files that changed from the base of the PR and between b535fa2 and 42d6ea7.

📒 Files selected for processing (7)
  • docs/buffer-abi.md
  • python/simpler/buffer.py
  • python/simpler/worker.py
  • src/common/task_interface/buffer.h
  • tests/ut/cpp/types/test_buffer.cpp
  • tests/ut/py/test_buffer.py
  • tests/ut/py/test_worker/test_create_buffer.py

Comment thread tests/ut/cpp/types/test_buffer.cpp
ChaoWao and others added 3 commits August 5, 2026 06:04
`tensor_extent_bytes` summed `(shapes[i]-1)*strides[i]` into a uint64. Both
fields are u32, so a single product already reaches ~2^64 and the multiply by
the element size overflows on top of it: the extent wrapped to a small value and
`validate_tensor` accepted the view as in-bounds. One dimension is enough —
shapes=(2147483649,), strides=(2147483648,), FLOAT32 validated against a 4-byte
backing while addressing ~16 EiB — and the path is reachable from Python through
`buffer.tensor(...)`.

The arithmetic now saturates, and an extent that lands on the saturation
sentinel is refused by name rather than compared against `nbytes`, so a
descriptor claiming an absurd `nbytes` cannot buy the view back. The frozen
field layout is untouched.

`tensors_overlap` inherited the same wrap, where the consequence is worse than a
rejected argument: two fully overlapping views compare as disjoint, which at the
dependency layer is a missed edge rather than an error. Its end offsets saturate
too.

`BufferDescriptor::operator==` bounded a memcmp by an unvalidated `body_len`. It
is clamped to DESC_MAX_BYTES, so the one length field in the header cannot bound
a read past the array it indexes.

The arbitrary-bytes pass asserted magic / generation / body_len / ndims on its
survivors but not the footprint invariant, which is why 4096 random blobs never
caught this; it asserts it now.

`validate_tensor`'s comment claimed the validator stands behind materialization
as well. It does not — `ImportRegistry.materialize` takes an already-decoded
descriptor and adds no endpoint check — and a header that freezes an ABI should
not overstate its own gate.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ccess

`wrap_fork_inherited` derived `backend_kind` from `access`, tagging FORK_SHM
whenever the grant was not READ. FORK_SHM and FORK_COW are opposite kernel write
semantics, not two spellings of one grant, and inferring either from the other
makes a read-only MAP_SHARED backing inexpressible: it is tagged FORK_COW, and
FORK_COW's READ-only rule then locks it there. The caller holds the mmap and is
the only party that knows which it is, so it states the tag; the default pair
stays the safe one, FORK_COW with READ. The function had no caller, so nothing
depends on the old signature.

Also corrects text describing machinery that does not exist.
`ImportRegistry.materialize` pointed a caller holding raw bytes at
`BufferDescriptor.unpack`, and the test module said it pins a `pack`/`unpack`
round trip; neither exists, and withholding the encoding is exactly what keeps
construction the only way in. `docs/buffer-abi.md` named the wire type
`simpler.task_interface.Tensor` while the same page's status note, the module
and a test all say it stays in `simpler.buffer` until the cutover. The
`h.shm.buf` sample now records that it is transitional: byte access belongs on
the view, since a device backing has no `shm` and code written against one
forks by backend.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…rkers

The L3+ child check counted chip and sub children only, so an L4 whose children
are local L3 Workers was refused. A next-level child is a forked process that
maps a POSIX_SHM backing by name exactly as a chip or sub child does, so it can
consume the buffer; `_next_level_shms` counts now, and the error message names
all three shapes.

`create_buffer` had no test. The new ones cover the gate in each child shape,
the childless L3+ refusal, the L2 leaf that needs no child at all, buffer id
uniqueness and the single nonce within one incarnation, and that a failing
close is reported while its registry entry stays for the cleanup journal to
retry.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@ChaoWao
ChaoWao force-pushed the fix-buffer-abi-validation-gaps branch from 42d6ea7 to 29b3527 Compare August 5, 2026 13:06
@ChaoWao
ChaoWao merged commit 02f90ef into hw-native-sys:main Aug 5, 2026
14 of 18 checks passed
@ChaoWao
ChaoWao deleted the fix-buffer-abi-validation-gaps branch August 5, 2026 13:50
YunjiQin added a commit to YunjiQin/simpler that referenced this pull request Aug 7, 2026
`TaskArgs` carried the GM-address-bearing `ChipTensor`, so a sender had to know
its receiver's address space: `_rewrite_blob_host_addrs` patched addresses by
numeric range and mis-rewrote device pointers that happened to fall inside a
registered host range, which a `child_memory` skip then papered over. This flips
the wire to the address-free `Tensor` hw-native-sys#1599 froze. A consumer resolves a backing
by canonical identity — map-once, exact — and no address crosses a process
boundary at all.

The wire is the mailbox blob that was already there. `TaskArgs`, `TaskArgsView`,
`write_blob` and `read_blob` change what they carry, not how many of them exist;
`ChipTensor` survives only inside `ChipStorageTaskArgs`, which the L2 leaf
materializes into and hands to runtime.so. `TaskArgsView::tensors` now runs
`validate_tensor` on every element it decodes: the element gained wire invariants
a decoder can check, and since the bound types expose their fields but not their
bytes, this is the only place a mailbox blob becomes a Tensor.

`ChipWorker` takes only `ChipStorageTaskArgs`. The blob round trip it sat behind
— C++ serialising bytes Python handed straight back to C++ — is gone, with the
`TaskArgsView` overloads that existed to unpack it.

A control-plane copy names both of its ends. `copy_to`/`copy_from` take a
`Buffer` on each side and write the two descriptors into the control frame; the
child resolves them through the same `ImportRegistry` its task arguments go
through. That replaces both the raw host address, meaningless across a fork once
the MAP_HOST pre-attach is gone, and the shm-name path that briefly stood in for
it — a second resolution rule for the one thing this ABI exists to resolve.

Every receive path keeps its container. A nested next-level child re-exports each
backing to a local handle and still hands its orchestration function a
`TaskArgs`, scalars included; a remote L3 runner resolves its sidecar descriptors
into one too, so an orchestration function can forward the args it was given
whichever way it was reached. Remote buffers are allocated through
`create_buffer`, so what a runner hands back carries an identity and a descriptor
rather than being a bare `SharedMemory` belonging to neither mechanism.

`simpler_setup.torch_interop.make_tensor_arg` becomes `make_chip_tensor_arg`: it
returns the chip POD where `Worker.make_tensor_arg` returns the wire `Tensor`,
and one name for both was the second public meaning rule 13 forbids.

`wrap_fork_inherited` callers name their backend explicitly now that hw-native-sys#1703
stopped inferring it from `access`. Each already knew the answer in a comment —
the HeapRing backings are MAP_SHARED, and `make_tensor_arg` follows the `shared`
it computes: at L2 the consumer is this process, so a write reaches the owner
trivially and FORK_COW's contract is the one that would be false there.
ChaoWao pushed a commit to YunjiQin/simpler that referenced this pull request Aug 7, 2026
`TaskArgs` carried the GM-address-bearing `ChipTensor`, so a sender had to know
its receiver's address space: `_rewrite_blob_host_addrs` patched addresses by
numeric range and mis-rewrote device pointers that happened to fall inside a
registered host range, which a `child_memory` skip then papered over. This flips
the wire to the address-free `Tensor` hw-native-sys#1599 froze. A consumer resolves a backing
by canonical identity — map-once, exact — and no address crosses a process
boundary at all.

The wire is the mailbox blob that was already there. `TaskArgs`, `TaskArgsView`,
`write_blob` and `read_blob` change what they carry, not how many of them exist;
`ChipTensor` survives only inside `ChipStorageTaskArgs`, which the L2 leaf
materializes into and hands to runtime.so. `TaskArgsView::tensors` now runs
`validate_tensor` on every element it decodes: the element gained wire invariants
a decoder can check, and since the bound types expose their fields but not their
bytes, this is the only place a mailbox blob becomes a Tensor.

`ChipWorker` takes only `ChipStorageTaskArgs`. The blob round trip it sat behind
— C++ serialising bytes Python handed straight back to C++ — is gone, with the
`TaskArgsView` overloads that existed to unpack it.

A control-plane copy names both of its ends. `copy_to`/`copy_from` take a
`Buffer` on each side and write the two descriptors into the control frame; the
child resolves them through the same `ImportRegistry` its task arguments go
through. That replaces both the raw host address, meaningless across a fork once
the MAP_HOST pre-attach is gone, and the shm-name path that briefly stood in for
it — a second resolution rule for the one thing this ABI exists to resolve.

Every receive path keeps its container. A nested next-level child re-exports each
backing to a local handle and still hands its orchestration function a
`TaskArgs`, scalars included; a remote L3 runner resolves its sidecar descriptors
into one too, so an orchestration function can forward the args it was given
whichever way it was reached. Remote buffers are allocated through
`create_buffer`, so what a runner hands back carries an identity and a descriptor
rather than being a bare `SharedMemory` belonging to neither mechanism.

`simpler_setup.torch_interop.make_tensor_arg` becomes `make_chip_tensor_arg`: it
returns the chip POD where `Worker.make_tensor_arg` returns the wire `Tensor`,
and one name for both was the second public meaning rule 13 forbids.

`wrap_fork_inherited` callers name their backend explicitly now that hw-native-sys#1703
stopped inferring it from `access`. Each already knew the answer in a comment —
the HeapRing backings are MAP_SHARED, and `make_tensor_arg` follows the `shared`
it computes: at L2 the consumer is this process, so a write reaches the owner
trivially and FORK_COW's contract is the one that would be false there.

`docs/buffer-abi.md` is the published page describing this wire, so it moves with
it. Its status note said `add_tensor` still takes a `ChipTensor` and that `Tensor`
is deliberately absent from `simpler.task_interface`; its scope section said the
dispatch wire is not connected and the submit-time checks are unreachable. All
three described the tree before this change. What remains absent is now stated
positively: the two other allocators, and the endpoint x `address_space` check
inside `materialize` — a device backing resolved there yields a pointer
meaningful only on its owner chip, and that is enforced today where a task is
submitted rather than where it is materialized.

Three names outlived what this change removes. A `reserve_slot` comment still
offered `create_host_buffer` as an example of an allocator the caller owns, after
the last call to it disappeared. `MappedArg.buffer` labelled `FORK_SHM` as
copy-on-write; it is the MAP_SHARED one, and that distinction is the whole reason
the two fork backends are separate — the corrected wording already sits in
`materialize` in the same file. The comments describing `ChipTensor::child_memory`
name a field now called `address_space`. The Python keyword stays `child_memory`
and now says why: it is the name of a u8 on the remote-L3 tensor wire, which
renaming the keyword would not change.
ChaoWao pushed a commit that referenced this pull request Aug 7, 2026
)

`TaskArgs` carried the GM-address-bearing `ChipTensor`, so a sender had to know
its receiver's address space: `_rewrite_blob_host_addrs` patched addresses by
numeric range and mis-rewrote device pointers that happened to fall inside a
registered host range, which a `child_memory` skip then papered over. This flips
the wire to the address-free `Tensor` #1599 froze. A consumer resolves a backing
by canonical identity — map-once, exact — and no address crosses a process
boundary at all.

The wire is the mailbox blob that was already there. `TaskArgs`, `TaskArgsView`,
`write_blob` and `read_blob` change what they carry, not how many of them exist;
`ChipTensor` survives only inside `ChipStorageTaskArgs`, which the L2 leaf
materializes into and hands to runtime.so. `TaskArgsView::tensors` now runs
`validate_tensor` on every element it decodes: the element gained wire invariants
a decoder can check, and since the bound types expose their fields but not their
bytes, this is the only place a mailbox blob becomes a Tensor.

`ChipWorker` takes only `ChipStorageTaskArgs`. The blob round trip it sat behind
— C++ serialising bytes Python handed straight back to C++ — is gone, with the
`TaskArgsView` overloads that existed to unpack it.

A control-plane copy names both of its ends. `copy_to`/`copy_from` take a
`Buffer` on each side and write the two descriptors into the control frame; the
child resolves them through the same `ImportRegistry` its task arguments go
through. That replaces both the raw host address, meaningless across a fork once
the MAP_HOST pre-attach is gone, and the shm-name path that briefly stood in for
it — a second resolution rule for the one thing this ABI exists to resolve.

Every receive path keeps its container. A nested next-level child re-exports each
backing to a local handle and still hands its orchestration function a
`TaskArgs`, scalars included; a remote L3 runner resolves its sidecar descriptors
into one too, so an orchestration function can forward the args it was given
whichever way it was reached. Remote buffers are allocated through
`create_buffer`, so what a runner hands back carries an identity and a descriptor
rather than being a bare `SharedMemory` belonging to neither mechanism.

`simpler_setup.torch_interop.make_tensor_arg` becomes `make_chip_tensor_arg`: it
returns the chip POD where `Worker.make_tensor_arg` returns the wire `Tensor`,
and one name for both was the second public meaning rule 13 forbids.

`wrap_fork_inherited` callers name their backend explicitly now that #1703
stopped inferring it from `access`. Each already knew the answer in a comment —
the HeapRing backings are MAP_SHARED, and `make_tensor_arg` follows the `shared`
it computes: at L2 the consumer is this process, so a write reaches the owner
trivially and FORK_COW's contract is the one that would be false there.

`docs/buffer-abi.md` is the published page describing this wire, so it moves with
it. Its status note said `add_tensor` still takes a `ChipTensor` and that `Tensor`
is deliberately absent from `simpler.task_interface`; its scope section said the
dispatch wire is not connected and the submit-time checks are unreachable. All
three described the tree before this change. What remains absent is now stated
positively: the two other allocators, and the endpoint x `address_space` check
inside `materialize` — a device backing resolved there yields a pointer
meaningful only on its owner chip, and that is enforced today where a task is
submitted rather than where it is materialized.

Three names outlived what this change removes. A `reserve_slot` comment still
offered `create_host_buffer` as an example of an allocator the caller owns, after
the last call to it disappeared. `MappedArg.buffer` labelled `FORK_SHM` as
copy-on-write; it is the MAP_SHARED one, and that distinction is the whole reason
the two fork backends are separate — the corrected wording already sits in
`materialize` in the same file. The comments describing `ChipTensor::child_memory`
name a field now called `address_space`. The Python keyword stays `child_memory`
and now says why: it is the name of a u8 on the remote-L3 tensor wire, which
renaming the keyword would not change.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant