Refactor: split the L2 tensor argument from each runtime's working tensor - #1974
Merged
ChaoWao merged 1 commit intoAug 25, 2026
Merged
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (9)
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthrough
ChangesTensor metadata and geometry
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: ⚪ Minimal · up to This refactor preserves the documented tensor behavior and passes the supplied build and test coverage; no actionable merge-blocking risk remains beyond normal checks. Sequence Diagram(s)sequenceDiagram
participant TensorCreateInfo
participant TensorMapEntry
participant ChipTensor
TensorCreateInfo->>ChipTensor: materialize metadata and shapes
ChipTensor->>ChipTensor: refresh derived row-major geometry
ChipTensor->>TensorMapEntry: copy address, metadata, and shapes
TensorMapEntry->>TensorMapEntry: retain linkage and populate geometry
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
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. Comment |
ChaoWao
force-pushed
the
refactor/tensor-layout-decouple
branch
from
August 24, 2026 07:51
9b747a3 to
53df0d9
Compare
ChaoWao
changed the title
ChaoWao
force-pushed
the
refactor/tensor-layout-decouple
branch
7 times, most recently
from
August 25, 2026 07:14
cef6dd0 to
8f95650
Compare
…nsor An argument arrives at L2 as a `ChipTensor`, and the runtime then decides things about it: which task produced it, the version its OverlapMap keys on, whether dependency tracking is creator-only. Those decisions lived on `ChipTensor` itself, in `src/common/task_interface/`, so one type served the boundary and both runtimes' working state. `create_from_chip_args` recorded the mismatch as an assertion — `debug_assert(!t.manual_dep && t.version == 0)`, two fields that are meaningless on an argument. `docs/buffer-abi.md` recorded it too, calling `ChipTensor` "L2 leaf, internal" and saying "You never build a `ChipTensor`" while `ChipWorker.run` took a container of them and `nb::class_<ChipTensor>` exposed a `make()` factory. The fusion dates to hw-native-sys#1093, which needed the strided view hw-native-sys#808 had just given the runtime's tensor to reach the argument boundary, and got there by promoting the runtime-private header into `task_interface/` rather than adding strides to the 40 B `ContinuousTensor` it replaced. It cost 40 B → 128 B per wire tensor and a mailbox doubling for a capability that PR noted was not yet used: two months on, `make_tensor_arg` and `make_chip_tensor_arg` still reject non-contiguous tensors. with a cost argument that applies equally here; nobody undid it at L2. `ChipTensor` (72 B) is a task argument: a resolved buffer and a strided view. struct ChipTensor { PTOBufferHandle buffer; uint64_t start_offset; uint32_t shapes[MAX_TENSOR_DIMS]; uint32_t strides[MAX_TENSOR_DIMS]; uint32_t ndims; DataType dtype; AddressSpace address_space; }; `simpler::hbg::Tensor` and `simpler::tmr::Tensor` (128 B) are the same geometry plus `owner_task_id`, `version`, `manual_dep` and the two caches derived from the geometry. One copy each, shared across architectures, placed the way `src/common/host_build_graph/graph_cache.h` already established. `simpler::hbg` was an existing namespace. `Runtime::set_orch_args` is the single place one becomes the other. Both runtimes call it only from `host/runtime_maker.cpp`, so adoption happens on the host before any orchestration runs — including for `tensormap_and_ringbuffer`, whose AICPU executor now reads an already-adopted `simpler::tmr::EntryArgsStorage` instead of converting per run. From there inward nothing holds the argument form: orchestration, the payload, the TensorMap and the kernels all name their runtime's type. A `tensor.h` in each runtime's `runtime/` directory sits first on that runtime's include path, so existing `#include "tensor.h"` lines resolve to both types without an include sweep. That shim names the two headers by their path under `src/common`, which every platform target but `aicore` already had on its include path; the four `aicore` CMakeLists gain the line their `aicpu` and `host` siblings carry. `is_contiguous` and `extent_elem` become methods that compute from the geometry; a runtime that reads them per task caches them on its own `Tensor`. The view ops and the copy helpers that maintained those caches move with them — an argument is not a thing you take views of, and nothing outside a runtime called them. `init_external` and both factories lose their `manual_dep` / `version` parameters, which every caller passed as `false` / `0`. `PTO2TaskPayload` is no longer a friend, and the boundary header no longer includes `task_id.h`: it names no TaskId at all. `sizeof(ChipStorageTaskArgs)` falls from ~33.8 KB to 19464 B at the 256-tensor cap. `MAILBOX_SIZE` is unaffected — since hw-native-sys#1729 the frame is sized by the 144 B wire `Tensor`. `ChipTensor`, `PTO2TensorMapEntry` and `TensorCreateInfo` shared a byte-level layout so three copy paths could each be a single 64-byte `memcpy`, held by 21 hand-written `offsetof` assertions across four trees. Two of the three were shaped backwards to satisfy it: `TensorCreateInfo` carried `__pad0__` / `__pad2__` / `__pad_flags__` purely to occupy `ChipTensor::buffer` / `::owner_task_id` / `::address_space`, and the entry's `memcpy` wrote `ChipTensor::buffer.size` into a `PTO2TensorMapEntry *` field, its comment noting this was "harmless because `link_entry()` overwrites `next_in_bucket` immediately after". Each struct now assigns the fields it wants, by name. The optimization the `memcpy` existed for survives as intent: a canonically contiguous source has the derived pair recomputed from `shapes` rather than read across. `PTO2TensorMapEntry::copy_tensor_create_info` is deleted — it had no callers in any tree, and was the only reason `TensorCreateInfo` had to be punnable onto an entry. The entry keeps its own two-cache-line split and the assertions describing it; what goes is the claim that its bytes agree with a different struct's. A source that names one runtime's `Tensor` cannot be compiled under the other, and 57 scene-test sources were: a `host_build_graph` case naming the `tensormap_and_ringbuffer` case's kernels, and two `tensormap_and_ringbuffer` classes taking an hbg class's `CALLABLE` verbatim. That worked only because `ChipTensor` happened to be identical in both trees — the coupling this change removes. Each borrowing case now carries its own sources, 73 files, and names local paths. The two class-reuse cases rebase the inherited `CALLABLE` onto their own `kernels/` copy through a `_rebase_callable` helper, so the Python that drives them stays in one file while the C++ each runtime compiles is its own. `task_timing_slots` carries the runtime in a function default rather than a `@scene_test` decorator, so its sources move under `kernels/<runtime>/` and the helper builds the path from the runtime it is driving. The cost is real: `spmd_multiblock_mix`, `alternating_matmul_add`, `batch_paged_attention`, `paged_attention_unroll`, `benchmark_bgemm` and the task-timing kernels now exist once per runtime and the copies can drift. A case that asserts something about `host_build_graph` should not depend on a file the other runtime owns. `examples/` orchestration and kernel sources lose their `Generated by PyPTO IR Compiler` marker: they are this repository's sources, and this change edits them. 364 of those kernels declare a ptoas helper as `template <typename ChipTensor>`, shadowing the type with a parameter name. The parameter is `TensorT` now: a generated helper that means "any tensor with .data()" should not name one, and the shadowing is what made the rename ambiguous in the first place. ## Kernels name no runtime A kernel reads a payload element and its code is identical whichever orchestrator filled it, so naming a runtime there carries no information — and several kernels are compiled under both. `examples/a2a3/host_build_graph/deepseek_v4_flash_decode` re-points all 368 of the `tensormap_and_ringbuffer` case's incores at that case's directory; `test_task_timing_e2e` drives one AIV kernel under both. Each runtime's `runtime/tensor.h` therefore exports using TaskTensor = simpler::{hbg,tmr}::Tensor; and every kernel names `TaskTensor`. It resolves per translation unit, so it is one type per build, not a third type. The alternative was duplicating 13 MB of generated MoE kernels into the hbg tree, where the two copies would drift. The name is not `Tensor`: `buffer.h`'s L3+ wire `Tensor` is visible in every orchestration translation unit, and one spelling meaning two types by context is what this change exists to remove. Cross-runtime *orchestration* sources need no name — `const auto &a = orch_args.tensor(0).ref()` is enough, and only kernels must spell the type inside a `reinterpret_cast`. `tensor_create_info.h`'s doc paragraph said `host_build_graph` has no initial-value fill "unlike the tensormap_and_ringbuffer copy of this header", which hw-native-sys#1981 made false by removing that fill; both copies now describe what is there. The copies were taken from sources whose spellings `main` has since retired — hw-native-sys#1980 finished that retirement repo-wide — so they carry the current ones. `test_hbg_sm_compaction` compared `GraphTensor` against `ChipTensor` to justify packing Graph boundaries into the payload's tensor slots, and read that pool as `ChipTensor`. The pool is the payload's, so both now name the runtime's `Tensor` — the assertion was true only while one type served both roles. Verification: full rebuild of both runtimes on both architectures, cpput 119/119, pyut 1908 passed / 14 skipped, a2a3sim 21 cases and a5sim 17 cases green. The sim sweeps build only the cases they run, so every orchestration source was also syntax-checked against its runtime's include path — 129 (source, runtime, arch) compiles, which is what caught a copy under host_build_graph/ still calling simpler::tmr::make_tensor_external. A second check reads which runtime *strings* a Python driver mentions rather than its @scene_test decorator; that is what caught task_timing_slots. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
An argument arrives at L2 as a
ChipTensor, and the runtime then decides thingsabout it: which task produced it, the version its OverlapMap keys on, whether
dependency tracking is creator-only. Those decisions lived on
ChipTensoritself,in
src/common/task_interface/, so one type served the boundary and bothruntimes' working state.
The code already recorded the mismatch:
create_from_chip_argsasserted!t.manual_dep && t.version == 0— two fieldsthat are meaningless on an argument.
docs/buffer-abi.md:33calledChipTensor"L2 leaf, internal" and:35said"You never build a
ChipTensor", whileChipWorker.runtook a container of themand
nb::class_<ChipTensor>exposed amake()factory.Where it came from
#1093 needed the strided view #808 had just given the runtime's tensor to reach the
argument boundary, and got there by promoting the runtime-private header into
task_interface/rather than adding strides to the 40 BContinuousTensoritreplaced. That cost 40 B → 128 B per wire tensor plus a
MAILBOX_SIZEdoubling,for a capability the PR itself noted was unused — and two months on
make_tensor_arg/make_chip_tensor_argstill reject non-contiguous tensors.#1729 undid the fusion on the L3 wire with a cost argument that applies equally
here; nobody undid it at L2.
The two types
ChipTensor(72 B) is a task argument — a resolved buffer and a strided view:simpler::hbg::Tensor/simpler::tmr::Tensor(128 B) are the same geometryplus
owner_task_id,version,manual_depand the two caches derived from it.One copy each, shared across architectures, placed the way
src/common/host_build_graph/graph_cache.halready established;simpler::hbgwasan existing namespace.
Runtime::set_orch_argsis the single place one becomes the other. Bothruntimes call it only from
host/runtime_maker.cpp, so adoption happens on thehost before any orchestration runs — including for
tensormap_and_ringbuffer,whose AICPU executor now reads an already-adopted
simpler::tmr::EntryArgsStorageinstead of converting per run. From there inwardnothing holds the argument form.
A
tensor.hin each runtime'sruntime/directory sits first on that runtime'sinclude path, so existing
#include "tensor.h"lines resolve to both types withoutan include sweep.
What the boundary type dropped
owner_task_id,version,manual_depis_contiguous,extent_elem_cachemanual_dep/versionfactory parametersfalse/0#include "task_id.h",friend PTO2TaskPayloadsizeof(ChipStorageTaskArgs): ~33.8 KB → 19464 B at the 256-tensor cap.MAILBOX_SIZEis unaffected — since #1729 the frame is sized by the 144 B wireTensor.Prerequisite: three structs sharing one byte layout
ChipTensor,PTO2TensorMapEntryandTensorCreateInfoshared a byte-levellayout so three copy paths could each be a single 64-byte
memcpy, held by 21hand-written
offsetofassertions across four trees. Two of the three wereshaped backwards to satisfy it:
TensorCreateInfocarried__pad0__/__pad2__/__pad_flags__purely tooccupy
ChipTensor::buffer/::owner_task_id/::address_space.memcpywroteChipTensor::buffer.sizeinto aPTO2TensorMapEntry *field, its comment noting this was "harmless becauselink_entry()overwritesnext_in_bucketimmediately after".Each struct now assigns the fields it wants, by name. The optimization the
memcpyexisted for survives as intent: a canonically contiguous source has the derived
pair recomputed from
shapesrather than read across.PTO2TensorMapEntry::copy_tensor_create_infois deleted — no callers in anytree, and the only reason
TensorCreateInfohad to be punnable onto an entry. Theentry keeps its own two-cache-line split and the assertions describing it.
Consequence: each case owns the sources it compiles
A source that names one runtime's
Tensorcannot be compiled under the other, and57 scene-test sources were: a
host_build_graphcase naming thetensormap_and_ringbuffercase's kernels, and twotensormap_and_ringbufferclasses taking an hbg class's
CALLABLEverbatim. That worked only becauseChipTensorhappened to be identical in both trees — the coupling this PR removes.Each borrowing case now carries its own sources (67 files) and names local paths.
The two class-reuse cases rebase the inherited
CALLABLEonto their ownkernels/copy via a
_rebase_callablehelper, so the Python that drives them stays in onefile while the C++ each runtime compiles is its own.
The cost is real and worth stating:
spmd_multiblock_mix,alternating_matmul_add,batch_paged_attention,paged_attention_unrollandbenchmark_bgemmnow exist once per runtime, and the copies can drift. Thejudgement is that a case asserting something about
host_build_graphshould notdepend on a file the other runtime owns.
Incidental
examples/orchestration and kernel sources lose theirGenerated by PyPTO IR Compilermarker: they are this repository's sources, and this PR edits them.tensor_create_info.hsaidhost_build_graphhas no initial-value fill "unlikethe tensormap_and_ringbuffer copy of this header", which The recorder thread owns its recording storage #1981 made false by
removing that fill. Both copies now describe what is there.
docs/buffer-abi.mdmoves with the wire it describes: type table, field tableand size figures.
Testing
static_assertholds-LE requires_hardware)dies held by other users with four jobs queued ahead
An earlier revision of this branch (layout decoupling only) passed
st-onboard-a2a3,st-onboard-a5andst-network1-onboard-a2a3.🤖 Generated with Claude Code