From b609d69fda800d2f8fae286ac4c0317b6aa544cb Mon Sep 17 00:00:00 2001 From: yanghaoran29 Date: Fri, 14 Aug 2026 17:37:45 +0800 Subject: [PATCH] Fix TensorMap reader anti-dependencies Track opt-in reader accesses separately from writers across A2/A3 and A5 HBG/TMR, and make later writes acquire the corresponding WAR dependencies. Exercise tracked inputs in the existing scalar_data examples, refresh PTO-ISA, and enable the A5 simulator MX FP GEMM cases. --- docs/dfx/dep-gen.md | 15 +- docs/orchestrator.md | 15 +- docs/task-flow.md | 10 +- docs/war-anti-dependency.md | 188 +++++-------- .../a2a3/tensormap_and_ringbuffer/README.md | 10 +- .../scalar_data/README.md | 35 ++- .../orchestration/scalar_data_orch.cpp | 47 +--- .../scalar_data/test_scalar_data.py | 8 +- .../a5/tensormap_and_ringbuffer/README.md | 9 +- .../scalar_data/README.md | 35 ++- .../orchestration/scalar_data_orch.cpp | 47 +--- .../scalar_data/test_scalar_data.py | 8 +- .../host_build_graph/docs/RUNTIME_LOGIC.md | 21 +- .../docs/SCALAR_DATA_ACCESS.md | 45 +-- .../host/dep_gen_host_graph.cpp | 54 +++- .../host_build_graph/host/runtime_maker.cpp | 22 +- .../orchestration/pto_arg_with_deps.h | 1 + .../runtime/dep_gen_host_graph.h | 13 +- .../orchestrator_core/pto_orchestrator.cpp | 165 +++++++++-- .../orchestrator_core/pto_runtime2.cpp | 55 ++-- .../runtime/pto_dep_compute.h | 88 +++--- .../runtime/pto_orchestrator.h | 8 +- .../runtime/pto_runtime2_types.h | 24 +- .../host_build_graph/runtime/pto_tensormap.h | 162 ++++++++--- .../host_build_graph/runtime/pto_types.h | 16 +- .../runtime/scheduler/pto_scheduler.cpp | 3 +- .../runtime/scheduler/pto_scheduler.h | 20 +- .../runtime/scheduler/scheduler_cold_path.cpp | 4 +- .../runtime/scheduler/scheduler_dispatch.cpp | 19 +- .../runtime/shared/pto_runtime2_init.cpp | 9 + .../runtime/shared/pto_tensormap.cpp | 81 ++++-- .../aicpu/aicpu_executor.cpp | 8 + .../docs/MULTI_RING.md | 6 +- .../docs/RUNTIME_LOGIC.md | 16 +- .../docs/SCALAR_DATA_ACCESS.md | 23 +- .../host/dep_gen_replay.cpp | 78 ++++-- .../host/dep_gen_replay.h | 7 +- .../orchestration/pto_arg_with_deps.h | 1 + .../orchestration/pto_orchestration_api.h | 17 +- .../runtime/pto_dep_compute.h | 129 ++++++--- .../runtime/pto_orchestrator.cpp | 28 +- .../runtime/pto_runtime2.cpp | 42 ++- .../runtime/pto_runtime2_types.h | 1 + .../runtime/pto_tensormap.h | 232 +++++++++++++--- .../runtime/pto_types.h | 11 +- .../runtime/shared/pto_tensormap.cpp | 95 +++++-- .../platform/onboard/host/device_runner.cpp | 42 ++- src/a5/platform/onboard/host/device_runner.h | 2 +- src/a5/platform/sim/host/device_runner.cpp | 41 ++- src/a5/platform/sim/host/device_runner.h | 2 +- .../host_build_graph/docs/RUNTIME_LOGIC.md | 21 +- .../docs/SCALAR_DATA_ACCESS.md | 45 +-- .../host/dep_gen_host_graph.cpp | 54 +++- .../host_build_graph/host/runtime_maker.cpp | 22 +- .../orchestration/pto_arg_with_deps.h | 1 + .../runtime/dep_gen_host_graph.h | 8 +- .../orchestrator_core/pto_orchestrator.cpp | 165 +++++++++-- .../orchestrator_core/pto_runtime2.cpp | 55 ++-- .../runtime/pto_dep_compute.h | 83 +++--- .../runtime/pto_orchestrator.h | 8 +- .../runtime/pto_runtime2_types.h | 19 +- .../host_build_graph/runtime/pto_tensormap.h | 162 ++++++++--- .../host_build_graph/runtime/pto_types.h | 11 +- .../runtime/scheduler/pto_scheduler.cpp | 3 +- .../runtime/scheduler/pto_scheduler.h | 20 +- .../runtime/scheduler/scheduler_cold_path.cpp | 4 +- .../runtime/scheduler/scheduler_dispatch.cpp | 19 +- .../runtime/shared/pto_runtime2_init.cpp | 9 + .../runtime/shared/pto_tensormap.cpp | 99 ++++--- .../aicpu/aicpu_executor.cpp | 8 + .../docs/MULTI_RING.md | 6 +- .../docs/RUNTIME_LOGIC.md | 16 +- .../docs/SCALAR_DATA_ACCESS.md | 23 +- .../host/dep_gen_replay.cpp | 78 ++++-- .../host/dep_gen_replay.h | 7 +- .../orchestration/pto_arg_with_deps.h | 1 + .../orchestration/pto_orchestration_api.h | 17 +- .../runtime/pto_dep_compute.h | 129 ++++++--- .../runtime/pto_orchestrator.cpp | 28 +- .../runtime/pto_runtime2.cpp | 31 ++- .../runtime/pto_runtime2_types.h | 1 + .../runtime/pto_tensormap.h | 232 +++++++++++++--- .../runtime/pto_types.h | 16 +- .../runtime/shared/pto_tensormap.cpp | 113 +++++--- src/common/hierarchical/orchestrator.cpp | 3 +- .../host_build_graph/docs/GRAPH_EXECUTION.md | 15 +- src/common/task_interface/task_args.h | 3 +- src/common/task_interface/tensor.h | 3 +- .../available_aicore_counts_orch.cpp | 20 +- .../reader_fanout/kernels/aiv/reader.cpp | 12 + .../reader_fanout/kernels/aiv/writer.cpp | 12 + .../orchestration/reader_fanout_orch.cpp | 12 + .../reader_fanout/test_reader_fanout.py | 83 ++++++ .../war_regression/kernels/aiv/reader.cpp | 12 + .../war_regression/kernels/aiv/writer.cpp | 12 + .../orchestration/war_regression_orch.cpp | 12 + .../war_regression/test_war_regression.py | 84 ++++++ .../available_aicore_counts_orch.cpp | 20 +- .../reader_fanout/kernels/aiv/reader.cpp | 12 + .../reader_fanout/kernels/aiv/writer.cpp | 12 + .../orchestration/reader_fanout_orch.cpp | 12 + .../reader_fanout/test_reader_fanout.py | 83 ++++++ .../war_regression/kernels/aiv/reader.cpp | 12 + .../war_regression/kernels/aiv/writer.cpp | 12 + .../orchestration/war_regression_orch.cpp | 12 + .../war_regression/test_war_regression.py | 89 ++++++ .../reader_fanout/kernels/aiv/reader.cpp | 12 + .../reader_fanout/kernels/aiv/writer.cpp | 12 + .../orchestration/reader_fanout_orch.cpp | 12 + .../reader_fanout/test_reader_fanout.py | 88 ++++++ .../war_regression/kernels/aiv/reader.cpp | 12 + .../war_regression/kernels/aiv/writer.cpp | 12 + .../orchestration/war_regression_orch.cpp | 12 + .../war_regression/test_war_regression.py | 84 ++++++ .../available_aicore_counts_orch.cpp | 20 +- .../mx_fp_gemm/kernels/aic/kernel_fp_gemm.cpp | 6 +- .../mx_fp_gemm/test_mx_fp_gemm.py | 6 +- .../reader_fanout/kernels/aiv/reader.cpp | 12 + .../reader_fanout/kernels/aiv/writer.cpp | 12 + .../orchestration/reader_fanout_orch.cpp | 12 + .../reader_fanout/test_reader_fanout.py | 88 ++++++ .../spmd_paged_attention_orch.cpp | 16 +- .../war_regression/kernels/aiv/reader.cpp | 12 + .../war_regression/kernels/aiv/writer.cpp | 12 + .../orchestration/war_regression_orch.cpp | 12 + .../war_regression/test_war_regression.py | 89 ++++++ .../common/reader_fanout/reader_fanout_orch.h | 102 +++++++ tests/st/common/reader_fanout/reader_kernel.h | 56 ++++ tests/st/common/reader_fanout/writer_kernel.h | 51 ++++ .../common/war_regression/war_reader_kernel.h | 54 ++++ .../war_regression/war_regression_orch.h | 67 +++++ .../common/war_regression/war_writer_kernel.h | 47 ++++ .../orchestration/require_sync_start_orch.cpp | 20 +- tests/ut/cpp/CMakeLists.txt | 4 + tests/ut/cpp/a2a3/test_dep_gen_host_graph.cpp | 54 ++++ tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp | 94 +++++++ tests/ut/cpp/a2a3/test_hbg_tensormap.cpp | 148 +++++++++- tests/ut/cpp/a2a3/test_orchestrator_fanin.cpp | 81 ++++++ tests/ut/cpp/a2a3/test_tensormap.cpp | 261 +++++++++++++++++- tests/ut/cpp/a5/test_hbg_submit_poison.cpp | 94 +++++++ tests/ut/cpp/a5/test_orchestrator_fanin.cpp | 81 ++++++ tests/ut/cpp/a5/test_tensormap.cpp | 261 +++++++++++++++++- .../common/test_hbg_graph_submit_failure.cpp | 29 +- 143 files changed, 4935 insertions(+), 1189 deletions(-) create mode 100644 tests/st/a2a3/host_build_graph/reader_fanout/kernels/aiv/reader.cpp create mode 100644 tests/st/a2a3/host_build_graph/reader_fanout/kernels/aiv/writer.cpp create mode 100644 tests/st/a2a3/host_build_graph/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp create mode 100644 tests/st/a2a3/host_build_graph/reader_fanout/test_reader_fanout.py create mode 100644 tests/st/a2a3/host_build_graph/war_regression/kernels/aiv/reader.cpp create mode 100644 tests/st/a2a3/host_build_graph/war_regression/kernels/aiv/writer.cpp create mode 100644 tests/st/a2a3/host_build_graph/war_regression/kernels/orchestration/war_regression_orch.cpp create mode 100644 tests/st/a2a3/host_build_graph/war_regression/test_war_regression.py create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/reader.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/writer.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/test_reader_fanout.py create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/aiv/reader.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/aiv/writer.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/orchestration/war_regression_orch.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/war_regression/test_war_regression.py create mode 100644 tests/st/a5/host_build_graph/reader_fanout/kernels/aiv/reader.cpp create mode 100644 tests/st/a5/host_build_graph/reader_fanout/kernels/aiv/writer.cpp create mode 100644 tests/st/a5/host_build_graph/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp create mode 100644 tests/st/a5/host_build_graph/reader_fanout/test_reader_fanout.py create mode 100644 tests/st/a5/host_build_graph/war_regression/kernels/aiv/reader.cpp create mode 100644 tests/st/a5/host_build_graph/war_regression/kernels/aiv/writer.cpp create mode 100644 tests/st/a5/host_build_graph/war_regression/kernels/orchestration/war_regression_orch.cpp create mode 100644 tests/st/a5/host_build_graph/war_regression/test_war_regression.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/reader.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/writer.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/reader_fanout/test_reader_fanout.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/aiv/reader.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/aiv/writer.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/orchestration/war_regression_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/war_regression/test_war_regression.py create mode 100644 tests/st/common/reader_fanout/reader_fanout_orch.h create mode 100644 tests/st/common/reader_fanout/reader_kernel.h create mode 100644 tests/st/common/reader_fanout/writer_kernel.h create mode 100644 tests/st/common/war_regression/war_reader_kernel.h create mode 100644 tests/st/common/war_regression/war_regression_orch.h create mode 100644 tests/st/common/war_regression/war_writer_kernel.h diff --git a/docs/dfx/dep-gen.md b/docs/dfx/dep-gen.md index 25d8c5cd0f..bbaad07002 100644 --- a/docs/dfx/dep-gen.md +++ b/docs/dfx/dep-gen.md @@ -65,7 +65,7 @@ inputs to each submit are captured and the graph is reconstructed afterwards. tensor metadata (producer/consumer shape + offset, dtype, version). Per-record semantics mirror runtime `submit_task` exactly: STEP 1 (explicit deps), STEP 3 (creator retention + tensormap lookup), - STEP 4 (register outputs). Per-successor dedup matches + STEP 4 (register reader/writer accesses). Per-successor dedup matches `PTO2FaninBuilder::append_fanin_or_fail`. After both passes finish per record, the replay asserts the two producer-id → `DepFlags` mappings are equal (same producers and same per-producer flags); if they diverge, @@ -86,8 +86,10 @@ nothing to capture-then-reconstruct. - **Capture point.** `submit_task_common` opens the task's entry, and `compute_task_fanin`'s `Annotate` hooks fire on each producer the runtime actually wires: creator retention (Step A) and tensormap lookup (Step B), plus - the declared dependencies at STEP 1. The edges are the runtime's own, not a - replay's inference, so they cannot drift from `compute_task_fanin` semantics. + the declared dependencies at STEP 1. A host-write node additionally records + its runtime-derived writer-consumer drain with source + `host_write_consumer`. The edges are the runtime's own, not a replay's + inference, so they cannot drift from dependency construction semantics. - **No ring, no collector, no replay.** The device-side dep_gen writer, its shared-memory ring, and the drain thread are all skipped (`dep_gen_host_graph_active()` tells the runner). Nothing is dropped under @@ -159,7 +161,7 @@ The standard SceneTest path "consumer_start_offset": "0", "consumer_strides": [1]}, {"pred": "4294967296", "succ": "4294967298", "arg": 0, "source": "tensormap", "flags": ["wait"], - "overlap": "covered", + "overlap": "covered", "hazard": "WAR", "access_kind": "READER", "tensor_id": "9514117477438350967", "consumer_dtype": "FLOAT32", "consumer_shape": [16384], "consumer_start_offset": "0", "consumer_strides": [1], @@ -178,6 +180,11 @@ silently lose precision if encoded as numbers. Python consumers pass these through `int(v)` which accepts either form, so the schema is JS-safe without burdening Python. +TensorMap edges additionally contain `hazard` (`RAW`, `WAW`, or `WAR`) and +`access_kind` (`WRITER` or `READER`). These fields describe why the consumer +must wait and which access index produced the match; creator and explicit +edges do not carry them. + Task ids encode `(ring_id << 32) | local_id` — the same layout as `PTO2TaskId::raw`: diff --git a/docs/orchestrator.md b/docs/orchestrator.md index 240f61f2c5..545c4b911f 100644 --- a/docs/orchestrator.md +++ b/docs/orchestrator.md @@ -560,16 +560,11 @@ private: first — `erase_task_outputs` therefore drops a key only while it still maps to the consumed slot, leaving the later writer's entry for new consumers to find. -- **WAR (write-after-read)** is not tracked directly. Read tasks don't - register in TensorMap; write tasks only look up current producer. If a - consumer reads `X` (recording fanin on producer P1) and then a later task - writes `X` (new producer P2 in TensorMap), there's no P1 → P2 edge. This is - correct: the reader only needs P1 to have completed, the new writer only - needs its own prior producer. Simultaneous read and write races are a user - bug, not a scheduler concern. When a workload genuinely needs the reader - ordered ahead of the overwrite, express it explicitly — see - [WAR anti-dependencies](war-anti-dependency.md) (issue #1306) for the - `add_dep` vs `INOUT` trade-off. +- **WAR (write-after-read)** is opt-in. Plain `INPUT` tasks do not register as + readers, so a later writer cannot discover them. `TRACKED_INPUT` publishes a + read-only reader entry; a later overlapping `INOUT`, existing-tensor + `OUTPUT`, or `set_tensor_data` then acquires the WAR dependency. See + [WAR anti-dependencies](war-anti-dependency.md) for the full access matrix. ### Thread safety diff --git a/docs/task-flow.md b/docs/task-flow.md index ce6f445a00..9af1669575 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -104,7 +104,7 @@ hierarchy levels. ```cpp class TaskArgs { std::vector tensors_; - std::vector tags_; // per-tensor: INPUT/OUTPUT/INOUT/OUTPUT_EXISTING/NO_DEP + std::vector tags_; // per-tensor access/dependency policy std::vector scalars_; public: void add_tensor(const ChipTensor&, TensorArgType tag = TensorArgType::INPUT); @@ -116,8 +116,12 @@ public: }; ``` -`TensorArgType` has five values (matches existing `tensor.h:45-51`): -`INPUT`, `OUTPUT`, `INOUT`, `OUTPUT_EXISTING`, `NO_DEP`. +The C++ `TensorArgType` has six values: `INPUT`, `OUTPUT`, `INOUT`, +`OUTPUT_EXISTING`, `NO_DEP`, and `TRACKED_INPUT`. `TRACKED_INPUT` has the same +read permission as `INPUT`; on A2/A3 and A5, both `host_build_graph` and +`tensormap_and_ringbuffer` publish it as a reader for a later WAR lookup. The +opt-in tag is currently a C++ `CoreTaskArgs` API and is not exposed by the +Python `TensorArgType` binding. For remote L3 submits, public Python still uses the same `TaskArgs` builder. `TaskArgs.add_tensor(RemoteTensorRef(...), tag)` appends a normal diff --git a/docs/war-anti-dependency.md b/docs/war-anti-dependency.md index 03927a1543..7755cad2ec 100644 --- a/docs/war-anti-dependency.md +++ b/docs/war-anti-dependency.md @@ -1,127 +1,81 @@ # WAR (Write-After-Read) Anti-Dependencies -**Decision record for issue #1306.** - -When a pure `INPUT` reader of a buffer is followed by a later task that -overwrites that same buffer, the runtime does **not** guarantee a -write-after-read (WAR) ordering on its own. This is a deliberate performance -trade-off, not a bug. Express the ordering explicitly — the recommended way is -a manual `add_dep`, not promoting the reader to `INOUT`. - -## The scenario +The runtime can track unfinished tensor readers and writers in TensorMap. A +pure reader that may overlap a later write must use `add_tracked_input()` so +the writer can discover it and create the WAR edge. ```text -R : reads X (add_input(X)) ── task R -W : writes X (add_inout(X) / add_output_existing(X)) ── task W, submitted later +W0: INOUT X ──RAW──> R0: TRACKED_INPUT X ──WAR──> W1: INOUT X ``` -`W` overwrites `X` while `R` may still be reading it. Correctness requires -`W` to wait for `R` (WAR / anti-dependency). The automatic dependency -generator does not create this edge for a pure `INPUT` reader. - -## Why the runtime does not track this automatically - -The automatic dep-gen in -[`pto_dep_compute.h`](../src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h) -tracks the two hazards that a producer-keyed map answers in O(1): - -- **RAW** (read-after-write): an `INPUT`/`INOUT` looks up the current writer - of `X` and takes an edge on it. -- **WAW** (write-after-write): an `INOUT`/`OUTPUT_EXISTING` replaces the - writer entry for `X` and takes an edge on the prior writer. - -Tracking **WAR** for pure readers is fundamentally different: a writer would -have to find *every reader that is still in flight*, which means keeping a -reader set per buffer and walking it on every write. Recording each pure read -as its own tensormap entry and walking the whole same-buffer chain on every -write is an `O(chain)` cost paid on the orchestration hot path — for an edge -that most workloads never need, because their reads are already ordered ahead -of the overwrite by the RAW/WAW chain that produced the new data. - -The orchestrator contract states this directly: -[`docs/orchestrator.md` §7 "Semantics"](orchestrator.md#semantics) — *"WAR is -not tracked directly … Simultaneous read and write races are a user bug, not a -scheduler concern."* This document is the how-to for the cases where you own -that WAR ordering. - -## Expressing the WAR edge — two options - -### Option A — promote the reader to `INOUT` - -Change the reader's argument from `add_input(X)` to `add_inout(X)`. An `INOUT` -access is treated as a writer: it registers `X` in the tensormap, so the later -write takes a WAW edge on it, and the host-side `set_tensor_data` path becomes -aware of the reader through the producer's `fanout_refcount`. - -**Cost — unnecessary read serialization.** Because `INOUT` is a write, two or -more readers of the same buffer no longer run concurrently: each becomes the -tensormap writer in turn, so the second reader takes a WAW edge on the first -and they serialize. A workload that reads `X` from several tasks in parallel -loses that parallelism purely to satisfy the anti-dependency. Reach for this -only when the reader genuinely also writes `X`, or when you specifically need -`set_tensor_data` on the host side to observe the reader. - -### Option B — manual `add_dep` (recommended) - -Capture the reader's task id and make the later writer depend on it -explicitly. This creates exactly the WAR edge and nothing else — the reads -stay pure `INPUT` and run concurrently. - -```cpp -// Reader: keep it a pure INPUT. -CoreTaskArgs r_args; -r_args.add_input(ext_X); -r_args.add_inout(ext_Y); -TaskOutputTensors r = rt_submit_aic_task(FUNC_READ_X, r_args); - -// Later writer: depend on the reader so the overwrite waits for the read. -// add_dep() lives on the convenience wrapper CoreTaskArgsWithDeps. -CoreTaskArgsWithDeps<> w_args; -w_args.add_inout(ext_X); -w_args.add_dep(r.task_id()); // <-- the WAR edge R -> W -rt_submit_aic_task(FUNC_WRITE_X, w_args); -``` +On A2/A3 and A5, both `host_build_graph` and `tensormap_and_ringbuffer` use the +same opt-in rule: plain `add_input()` queries prior writers but does not publish +a reader. Use `add_tracked_input()` when a later overlapping write must wait. + +## Access semantics + +| Argument | Queries | Registers | Meaning | +| -------- | ------- | --------- | ------- | +| `INPUT` | overlapping writers (RAW) | nothing | ordinary read-only access | +| `TRACKED_INPUT` | overlapping writers (RAW) | reader | read-only access that must order a later write | +| `INOUT` | overlapping writers (RAW/WAW) and readers (WAR) | writer | read-modify-write | +| existing-tensor `add_output` | overlapping readers (WAR) | writer | pure overwrite (`OUTPUT_EXISTING`) | +| runtime-created `add_output` | nothing | nothing | fresh allocation (`OUTPUT`) | +| `NO_DEP` | creator only | nothing | retains the allocator but skips TensorMap lookup/publication | + +Independent input tasks do not depend on each other and can still execute in +parallel. Accesses are registered only after the task's complete fanin has been +computed, preventing aliases within one task from creating a self-dependency. + +The annotation belongs on the reader, not on the later writer. A writer cannot +retroactively discover an earlier plain input that left no reader entry. + +`OUTPUT_EXISTING` retains its existing unordered-writer contract: it waits for +readers but does not acquire WAW edges on older writers. Consequently, a fully +covered reader entry can be retired after its WAR edge is created, while an +unordered writer entry must remain discoverable. + +## `INOUT` versus pure overwrite + +Use `INOUT` when the final value can depend on the target's old contents, such +as accumulation, an in-place operator, a partial update that preserves other +elements, or conditional writeback. Use existing-tensor `add_output(X)` when a +task fully determines the bytes it writes without reading the previous value. + +## Host access + +- `get_tensor_data()` is a host read and waits only for overlapping writers. +- `set_tensor_data()` is a host write and waits for overlapping writers, the + existing writer-consumer drain, and every overlapping tracked reader task to + complete. TMR performs the wait synchronously; HBG emits an equivalent host + write graph node. Neither waits for a reader's downstream consumers. + +Use `add_tracked_input()` for any pure reader followed by an overlapping +`set_tensor_data()` call. + +## Explicit overrides + +Manual scopes skip automatic dependency computation. Tensors marked +`manual_dep` and `add_no_dep()` retain a valid creator but skip automatic +reader/writer lookup and registration. In these modes, use +`CoreTaskArgs::set_dependencies()` or `CoreTaskArgsWithDeps::add_dep()` to state +the required task ordering explicitly. An explicit task edge is not a +buffer-keyed access record, so it cannot make a task visible to a later host +`set_tensor_data()` call. + +## Capacity and diagnostics + +Tracked-reader and writer entries share the existing 65,536-entry pool and +retire at the existing `CONSUMED` watermark. The indexes have separate bucket +heads so reader fan-out does not make subsequent readers scan older readers. +HBG keeps 128 fanins inline and spills additional deduplicated fanins into its +scheduler pool; fanins are never silently truncated. -`add_dep` is the convenience layer over the primitive -`CoreTaskArgs::set_dependencies(ptr, count)`; both are documented in -[`pto_arg_with_deps.h`](../src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h). -Multiple readers each contribute one `add_dep(reader.task_id())` on the writer -and still run in parallel with each other — only the writer waits. - -**`add_dep` is a task-to-task edge only — it is invisible to the host-side -`set_tensor_data`.** `set_tensor_data(X)` is a host write, not a task in the -graph, so a "writer depends on reader" edge places no constraint on it. Its -only channel for discovering in-flight readers is buffer-keyed: it looks up -`X`'s producer in the TensorMap and waits for that producer *and its -consumers* (`wait_for_tensor_ready` with `wait_for_consumers=true`, see -[`pto_runtime2.cpp`](../src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp)). -`add_dep` writes no TensorMap entry for `X` and does not touch that producer's -fanout, so a reader wired only through `add_dep` is not in the set -`set_tensor_data` waits on. If a buffer that a reader touches may later be -written from the host via `set_tensor_data`, that reader must use `add_inout` -(Option A) to register itself in the TensorMap — `add_dep` cannot substitute -here. - -## Recommendation - -Prefer **Option B (`add_dep`)**. It is precise (one edge, no side effects on -the tensormap), and it preserves read parallelism. Use **Option A (`INOUT`)** -only when the tensor is semantically read-modify-write anyway, or when the -reader must be visible to the host-side `set_tensor_data` WAR guard (see the -`set_tensor_data` note in -[`pto_orchestration_api.h`](../src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_orchestration_api.h)). - -| Concern | Option A: `INOUT` | Option B: `add_dep` | -| ------- | ----------------- | ------------------- | -| Creates the WAR edge | ✓ | ✓ | -| Keeps readers concurrent | ✗ (readers serialize as WAW) | ✓ | -| Visible to host `set_tensor_data` | ✓ | ✗ | -| Extra tensormap entry per read | ✓ | ✗ | -| Effort | change one arg tag | capture id + one `add_dep` | +Dependency capture labels TensorMap edges with `hazard` (`RAW`, `WAW`, or +`WAR`) and `access_kind` (`READER` or `WRITER`). TensorMap also retains +reader/writer live counts and high-water marks for capacity diagnosis. ## See also -- [`docs/orchestrator.md` §7](orchestrator.md#semantics) — TensorMap RAW/WAW/WAR - semantics. -- [`docs/manual-scope.md`](manual-scope.md) — manual scopes, where automatic - dep tracking is off and `add_dep` is the primary ordering tool. +- [Dependency generation DFX](dfx/dep-gen.md) +- [Manual scopes](manual-scope.md) diff --git a/examples/a2a3/tensormap_and_ringbuffer/README.md b/examples/a2a3/tensormap_and_ringbuffer/README.md index be90dca38f..4b3dc9a4a6 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/README.md +++ b/examples/a2a3/tensormap_and_ringbuffer/README.md @@ -18,7 +18,7 @@ For the `Worker` API underneath the framework, see | Example | What it teaches | | ------- | --------------- | | [`vector_example/`](vector_example/) | The smallest complete kernel: `f = (a+b+1)*(a+b+2) + (a+b)`. Runs on sim. | -| [`scalar_data/`](scalar_data/) | Orchestration-level data manipulation — `get_tensor_data` / `set_tensor_data` round-trips, runtime-created outputs with initial values, and automatic WAW / WAR waits. Also the reference for the one case where they **don't** fire: `add_input` on an external tensor registers no TensorMap entry, so a later `set_tensor_data` races the reader. | +| [`scalar_data/`](scalar_data/) | Orchestration-level data manipulation — `get_tensor_data` / `set_tensor_data` round-trips, runtime-created outputs with initial values, and opt-in WAR tracking through `add_tracked_input`. | ## Compute @@ -70,14 +70,14 @@ Wrap hardware runs in `task-submit` on a shared box; see ## Relationship to `examples/a5/` -Five examples exist under both architectures with the same name: +Six examples exist under both architectures with the same name: `vector_example`, `paged_attention`, `paged_attention_manual_scope`, -`paged_attention_unroll_manual_scope`, and `sdma_async_completion_demo`. They +`paged_attention_unroll_manual_scope`, `scalar_data`, and +`sdma_async_completion_demo`. They are ports of each other and differ mainly in tile shapes and platform strings — `vector_example` differs by two lines. When you change one, check whether its sibling needs the same change. Only here: `benchmark_bgemm` (a5 has `bgemm` instead), `deepseek_v4_flash_decode`, `merge_pipeline_barrier`, -`paged_attention_ringbuffer`, `prefetch_async_demo`, `qwen3_14b_decode`, -`scalar_data`. +`paged_attention_ringbuffer`, `prefetch_async_demo`, `qwen3_14b_decode`. diff --git a/examples/a2a3/tensormap_and_ringbuffer/scalar_data/README.md b/examples/a2a3/tensormap_and_ringbuffer/scalar_data/README.md index a0ef3d0d06..4cbd75aef5 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/scalar_data/README.md +++ b/examples/a2a3/tensormap_and_ringbuffer/scalar_data/README.md @@ -10,22 +10,21 @@ using, who waits?** ## The answer, and the trap -`set_tensor_data` auto-waits — but only for tensors TensorMap knows about. +`set_tensor_data` waits for overlapping writers, their existing consumers, +and readers explicitly published with `add_tracked_input`. -| Tensor kind | Added as | Does `set_tensor_data` wait? | -| ----------- | -------- | ---------------------------- | -| runtime-created | anything | **Yes.** TensorMap holds a producer entry, so it waits for the producer to complete (WAW) and for `fanout_refcount` to drain (WAR). | -| external | `add_output` / `add_inout` | **Yes.** These create a TensorMap entry, so the same lookup works. | -| external | `add_input` | **No — and this is a data race.** `add_input` on an external tensor creates no TensorMap entry. `set_tensor_data` finds no producer, so it writes immediately, while the reader kernel may still be running. | +| Reader annotation | Published in TensorMap? | Does a later overlapping `set_tensor_data` wait? | +| ----------------- | ----------------------- | ------------------------------------------------ | +| `add_tracked_input` | Yes | **Yes.** The host write waits for the reader task to complete. | +| `add_input` | No | **No.** The host write cannot discover an unfinished pure reader. | -That last row is the whole reason this example exists. There is no error and no -warning; the value is simply overwritten under a running kernel. **On an -external tensor, use `add_inout` when a later `set_tensor_data` must not race -the reader**, even though `add_input` describes the kernel's access accurately. +The annotation belongs on the reader: a later writer cannot recover an earlier +plain input that was never published. Keep ordinary `add_input` for reads that +cannot overlap a later write, and use `add_tracked_input` when they can. -Step 12 in the orchestration demonstrates the safe form: submit `noop` with -`ext_b` as an output (`noop` reads nothing), then `set_tensor_data(ext_b, 55.0)` -— which now waits, because the output registered an entry. +Step 11 submits a real `kernel_add` reader before overwriting its input. Step 12 +uses `noop` to isolate external tracked-reader registration from kernel work; +plain `add_input` would leave the host write with no reader entry to wait on. ## The thirteen steps @@ -40,9 +39,9 @@ landing a value in `check[0..8]`: | 8 | Plain arithmetic in the orchestration on a value it read back | `check[4]` = 79.0 | | 9 | `set_tensor_data` → `get_tensor_data` round-trip | `check[5]` = 42.0 | | 10 | Orchestration → AICore RAW: write 10.0, then a kernel reads it | `check[6]` = 12.0 | -| 11 | WAW + WAR on a runtime-created tensor — `set_tensor_data` waits for both the producer and the consumer | `check[7]` = 88.0 | -| 12 | External WAR done correctly, via `add_inout` | `check[8]` = 55.0 | -| 13 | `result = a + b`, external output through `add_inout` | `result` | +| 11 | Internal tracked reader followed by a host write | `check[7]` = 88.0 | +| 12 | External tracked-reader registration followed by a host write | `check[8]` = 55.0 | +| 13 | `result = a + b`, external output through `add_output` | `result` | `kernel_noop` exists precisely because several steps need a task that *registers a dependency* without touching data. @@ -54,8 +53,8 @@ auto selection. Note the comment on the `check` tensor in `generate_args`: it is **exactly 9 slots**, matching `check[0..8]`. Output-tensor slots are not seeded from the -host, so a tenth slot would read undefined device memory and the golden -comparison would fail nondeterministically. +host, so an unwritten slot would contain undefined device memory and make the +golden comparison nondeterministic. ## Run diff --git a/examples/a2a3/tensormap_and_ringbuffer/scalar_data/kernels/orchestration/scalar_data_orch.cpp b/examples/a2a3/tensormap_and_ringbuffer/scalar_data/kernels/orchestration/scalar_data_orch.cpp index 7bf760a420..70bc6bcf49 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/scalar_data/kernels/orchestration/scalar_data_orch.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/scalar_data/kernels/orchestration/scalar_data_orch.cpp @@ -11,7 +11,7 @@ /** * Scalar Data Dependency Test Orchestration * - * End-to-end test for get_tensor_data, set_tensor_data, and add_inout + * End-to-end test for get_tensor_data, set_tensor_data, add_inout, and add_tracked_input * with runtime-created outputs and initial value support. * * Flow: @@ -25,9 +25,9 @@ * 8. check[4] = 2.0 + 77.0 = 79.0 (orchestration arithmetic) * 9. set_tensor_data(scalar_tensor, {0}, 42.0), get_tensor_data → check[5] = 42.0 * 10. Orch set_tensor_data(d, {0}, 10.0) → kernel_add(d, a) → check[6] = 12.0 - * 11. WAW+WAR: kernel_add reads c → set_tensor_data(c, 88.0) auto-waits → check[7] = 88.0 - * 12. External WAR with INOUT: noop(ext_b as INOUT) → set_tensor_data(ext_b) → check[8] = 55.0 - * 13. result = a + b (kernel_add, external output via INOUT) + * 11. Internal WAR: tracked kernel_add reads c → set_tensor_data(c, 88.0) waits → check[7] = 88.0 + * 12. External WAR: tracked noop on ext_b → set_tensor_data(ext_b, 55.0) waits → check[8] = 55.0 + * 13. result = a + b (kernel_add, external output via add_output) */ #include @@ -185,28 +185,17 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const Chip // ========================================================= // Step 11: WAW + WAR on internal tensor - // c was written by Step 1 (kernel_add, TensorMap has producer entry). - // Submit a new kernel that reads c as INPUT (creates consumer dep). - // Then set_tensor_data(c) — no manual get_tensor_data sync. - // set_tensor_data internally waits for: - // - WAW: producer (Step 1) COMPLETED - // - WAR: consumer (this kernel) done (fanout_refcount check) - // - // NOTE on external tensors: ext_a was read by Step 1 as INPUT, - // but TensorMap has no producer entry for ext_a (only consumers). - // set_tensor_data(ext_a) would NOT detect the reader — data race. - // To ensure WAR safety on external tensors, use add_inout() - // instead of add_input() so TensorMap tracks the access chain. + // A tracked reader consumes c while set_tensor_data prepares to overwrite it. + // The host write waits for both c's earlier writer and the tracked reader. // ========================================================= { CoreTaskArgs args; - args.add_input(c); + args.add_tracked_input(c); args.add_input(ext_b); args.add_output(inter_ci); (void)rt_submit_aiv_task(FUNC_ADD, args); // NOLINT(readability/casting) } - // set_tensor_data auto-waits for producer + consumer before writing idx[0] = 0; set_tensor_data(c, 1, idx, 88.0f); float waw_val = get_tensor_data(c, 1, idx); @@ -216,32 +205,20 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const Chip set_tensor_data(ext_check, 1, check_idx, waw_val); // ========================================================= - // Step 12: External tensor WAR — must use add_output or add_inout, not add_input - // - // For external tensors, using add_input() does NOT create a - // TensorMap entry. set_tensor_data would then write immediately - // without waiting for the reader kernel — a WAR data race. - // - // Using add_output() (or add_inout()) creates a TensorMap entry, - // enabling set_tensor_data to detect the producer via TensorMap lookup - // and wait for fanout_refcount (all consumers done). - // - // Here we submit noop with ext_b as write-only output (noop doesn't - // read data), then set_tensor_data overwrites ext_b[0] = 55.0. - // set_tensor_data auto-waits for the noop to complete. + // Step 12: WAR on an external tensor + // Plain add_input would not publish the reader. add_tracked_input keeps + // the argument read-only while making the reader visible to the host write. // ========================================================= { CoreTaskArgs args; - args.add_output(ext_b); // write-only: creates TensorMap entry (not add_input!) + args.add_tracked_input(ext_b); rt_submit_aiv_task(FUNC_NOOP, args); } idx[0] = 0; set_tensor_data(ext_b, 1, idx, 55.0f); float ext_war_val = get_tensor_data(ext_b, 1, idx); - LOG_INFO( - "External WAR (INOUT): set_tensor_data(ext_b, 55.0) = %f (expected 55.0)", static_cast(ext_war_val) - ); + LOG_INFO("External WAR: set_tensor_data(ext_b, 55.0) = %f (expected 55.0)", static_cast(ext_war_val)); check_idx[0] = 8; set_tensor_data(ext_check, 1, check_idx, ext_war_val); diff --git a/examples/a2a3/tensormap_and_ringbuffer/scalar_data/test_scalar_data.py b/examples/a2a3/tensormap_and_ringbuffer/scalar_data/test_scalar_data.py index 20dc2d2a62..f404d97609 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/scalar_data/test_scalar_data.py +++ b/examples/a2a3/tensormap_and_ringbuffer/scalar_data/test_scalar_data.py @@ -7,7 +7,7 @@ # 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. # ----------------------------------------------------------------------------------------------------------- -"""Scalar data dependency test: GetTensorData, SetTensorData, add_inout. +"""Scalar data dependency test: GetTensorData, SetTensorData, and tracked readers. Tests orchestration-level data manipulation: scalar initialization, Get/Set round-trips, WAW+WAR dependency auto-wait, and external tensor WAR. @@ -21,7 +21,7 @@ @scene_test(level=2, runtime="tensormap_and_ringbuffer") class TestScalarData(SceneTestCase): - """Scalar data dependency: Get/SetTensorData, add_inout with initial value.""" + """Scalar data dependency: host access, initial values, and tracked readers.""" CALLABLE = { "orchestration": { @@ -77,8 +77,8 @@ def compute_golden(self, args, params): args.check[4] = 79.0 # orchestration arithmetic: 2.0 + 77.0 args.check[5] = 42.0 # Orch set->get round-trip: SetTensorData then GetTensorData args.check[6] = 12.0 # Orch->AICore RAW: SetTensorData(d,10.0) + kernel_add(d,a) -> 10.0+2.0 - args.check[7] = 88.0 # WAW+WAR: kernel reads c, SetTensorData(c,88.0) auto-waits - args.check[8] = 55.0 # External WAR: noop(ext_b INOUT) -> SetTensorData(ext_b,55.0) auto-waits + args.check[7] = 88.0 # host write to c completes after the tracked reader + args.check[8] = 55.0 # host write to external b completes after the tracked reader if __name__ == "__main__": diff --git a/examples/a5/tensormap_and_ringbuffer/README.md b/examples/a5/tensormap_and_ringbuffer/README.md index 0de941d811..3c786aaa9f 100644 --- a/examples/a5/tensormap_and_ringbuffer/README.md +++ b/examples/a5/tensormap_and_ringbuffer/README.md @@ -63,9 +63,10 @@ refuses a wrong-arch invocation before any device lock is taken. ## Relationship to `examples/a2a3/` -Five examples exist under both architectures with the same name: +Six examples exist under both architectures with the same name: `vector_example`, `paged_attention`, `paged_attention_manual_scope`, -`paged_attention_unroll_manual_scope`, and `sdma_async_completion_demo`. They +`paged_attention_unroll_manual_scope`, `scalar_data`, and +`sdma_async_completion_demo`. They are ports of each other and differ mainly in tile shapes and platform strings — `vector_example` differs by two lines. When you change one, check whether its sibling needs the same change. @@ -73,8 +74,8 @@ sibling needs the same change. Only here: `bgemm` (a2a3 has `benchmark_bgemm` instead) and `urma_deferred_completion_demo`. a2a3 additionally carries `merge_pipeline_barrier`, `paged_attention_ringbuffer`, `prefetch_async_demo`, -and `scalar_data`, none of which have an a5 port — `tests/st` tracks that gap -separately (PR #1450 ports the `tests/st` side). +none of which have an a5 port — `tests/st` tracks that gap separately (PR #1450 +ports the `tests/st` side). `qwen3_14b_decode/` is an **a5 port in progress** (a2a3 source → a5). It compiles and runs on a5 after the a2a3→a5 kernel diffs diff --git a/examples/a5/tensormap_and_ringbuffer/scalar_data/README.md b/examples/a5/tensormap_and_ringbuffer/scalar_data/README.md index 6213bb03db..94228e1eff 100644 --- a/examples/a5/tensormap_and_ringbuffer/scalar_data/README.md +++ b/examples/a5/tensormap_and_ringbuffer/scalar_data/README.md @@ -10,22 +10,21 @@ using, who waits?** ## The answer, and the trap -`set_tensor_data` auto-waits — but only for tensors TensorMap knows about. +`set_tensor_data` waits for overlapping writers, their existing consumers, +and readers explicitly published with `add_tracked_input`. -| Tensor kind | Added as | Does `set_tensor_data` wait? | -| ----------- | -------- | ---------------------------- | -| runtime-created | anything | **Yes.** TensorMap holds a producer entry, so it waits for the producer to complete (WAW) and for `fanout_refcount` to drain (WAR). | -| external | `add_output` / `add_inout` | **Yes.** These create a TensorMap entry, so the same lookup works. | -| external | `add_input` | **No — and this is a data race.** `add_input` on an external tensor creates no TensorMap entry. `set_tensor_data` finds no producer, so it writes immediately, while the reader kernel may still be running. | +| Reader annotation | Published in TensorMap? | Does a later overlapping `set_tensor_data` wait? | +| ----------------- | ----------------------- | ------------------------------------------------ | +| `add_tracked_input` | Yes | **Yes.** The host write waits for the reader task to complete. | +| `add_input` | No | **No.** The host write cannot discover an unfinished pure reader. | -That last row is the whole reason this example exists. There is no error and no -warning; the value is simply overwritten under a running kernel. **On an -external tensor, use `add_inout` when a later `set_tensor_data` must not race -the reader**, even though `add_input` describes the kernel's access accurately. +The annotation belongs on the reader: a later writer cannot recover an earlier +plain input that was never published. Keep ordinary `add_input` for reads that +cannot overlap a later write, and use `add_tracked_input` when they can. -Step 12 in the orchestration demonstrates the safe form: submit `noop` with -`ext_b` as an output (`noop` reads nothing), then `set_tensor_data(ext_b, 55.0)` -— which now waits, because the output registered an entry. +Step 11 submits a real `kernel_add` reader before overwriting its input. Step 12 +uses `noop` to isolate external tracked-reader registration from kernel work; +plain `add_input` would leave the host write with no reader entry to wait on. ## The thirteen steps @@ -40,9 +39,9 @@ landing a value in `check[0..8]`: | 8 | Plain arithmetic in the orchestration on a value it read back | `check[4]` = 79.0 | | 9 | `set_tensor_data` → `get_tensor_data` round-trip | `check[5]` = 42.0 | | 10 | Orchestration → AICore RAW: write 10.0, then a kernel reads it | `check[6]` = 12.0 | -| 11 | WAW + WAR on a runtime-created tensor — `set_tensor_data` waits for both the producer and the consumer | `check[7]` = 88.0 | -| 12 | External WAR done correctly, via `add_inout` | `check[8]` = 55.0 | -| 13 | `result = a + b`, external output through `add_inout` | `result` | +| 11 | Internal tracked reader followed by a host write | `check[7]` = 88.0 | +| 12 | External tracked-reader registration followed by a host write | `check[8]` = 55.0 | +| 13 | `result = a + b`, external output through `add_output` | `result` | `kernel_noop` exists precisely because several steps need a task that *registers a dependency* without touching data. @@ -54,8 +53,8 @@ auto selection. Note the comment on the `check` tensor in `generate_args`: it is **exactly 9 slots**, matching `check[0..8]`. Output-tensor slots are not seeded from the -host, so a tenth slot would read undefined device memory and the golden -comparison would fail nondeterministically. +host, so an unwritten slot would contain undefined device memory and make the +golden comparison nondeterministic. ## Run diff --git a/examples/a5/tensormap_and_ringbuffer/scalar_data/kernels/orchestration/scalar_data_orch.cpp b/examples/a5/tensormap_and_ringbuffer/scalar_data/kernels/orchestration/scalar_data_orch.cpp index 7bf760a420..70bc6bcf49 100644 --- a/examples/a5/tensormap_and_ringbuffer/scalar_data/kernels/orchestration/scalar_data_orch.cpp +++ b/examples/a5/tensormap_and_ringbuffer/scalar_data/kernels/orchestration/scalar_data_orch.cpp @@ -11,7 +11,7 @@ /** * Scalar Data Dependency Test Orchestration * - * End-to-end test for get_tensor_data, set_tensor_data, and add_inout + * End-to-end test for get_tensor_data, set_tensor_data, add_inout, and add_tracked_input * with runtime-created outputs and initial value support. * * Flow: @@ -25,9 +25,9 @@ * 8. check[4] = 2.0 + 77.0 = 79.0 (orchestration arithmetic) * 9. set_tensor_data(scalar_tensor, {0}, 42.0), get_tensor_data → check[5] = 42.0 * 10. Orch set_tensor_data(d, {0}, 10.0) → kernel_add(d, a) → check[6] = 12.0 - * 11. WAW+WAR: kernel_add reads c → set_tensor_data(c, 88.0) auto-waits → check[7] = 88.0 - * 12. External WAR with INOUT: noop(ext_b as INOUT) → set_tensor_data(ext_b) → check[8] = 55.0 - * 13. result = a + b (kernel_add, external output via INOUT) + * 11. Internal WAR: tracked kernel_add reads c → set_tensor_data(c, 88.0) waits → check[7] = 88.0 + * 12. External WAR: tracked noop on ext_b → set_tensor_data(ext_b, 55.0) waits → check[8] = 55.0 + * 13. result = a + b (kernel_add, external output via add_output) */ #include @@ -185,28 +185,17 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const Chip // ========================================================= // Step 11: WAW + WAR on internal tensor - // c was written by Step 1 (kernel_add, TensorMap has producer entry). - // Submit a new kernel that reads c as INPUT (creates consumer dep). - // Then set_tensor_data(c) — no manual get_tensor_data sync. - // set_tensor_data internally waits for: - // - WAW: producer (Step 1) COMPLETED - // - WAR: consumer (this kernel) done (fanout_refcount check) - // - // NOTE on external tensors: ext_a was read by Step 1 as INPUT, - // but TensorMap has no producer entry for ext_a (only consumers). - // set_tensor_data(ext_a) would NOT detect the reader — data race. - // To ensure WAR safety on external tensors, use add_inout() - // instead of add_input() so TensorMap tracks the access chain. + // A tracked reader consumes c while set_tensor_data prepares to overwrite it. + // The host write waits for both c's earlier writer and the tracked reader. // ========================================================= { CoreTaskArgs args; - args.add_input(c); + args.add_tracked_input(c); args.add_input(ext_b); args.add_output(inter_ci); (void)rt_submit_aiv_task(FUNC_ADD, args); // NOLINT(readability/casting) } - // set_tensor_data auto-waits for producer + consumer before writing idx[0] = 0; set_tensor_data(c, 1, idx, 88.0f); float waw_val = get_tensor_data(c, 1, idx); @@ -216,32 +205,20 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const Chip set_tensor_data(ext_check, 1, check_idx, waw_val); // ========================================================= - // Step 12: External tensor WAR — must use add_output or add_inout, not add_input - // - // For external tensors, using add_input() does NOT create a - // TensorMap entry. set_tensor_data would then write immediately - // without waiting for the reader kernel — a WAR data race. - // - // Using add_output() (or add_inout()) creates a TensorMap entry, - // enabling set_tensor_data to detect the producer via TensorMap lookup - // and wait for fanout_refcount (all consumers done). - // - // Here we submit noop with ext_b as write-only output (noop doesn't - // read data), then set_tensor_data overwrites ext_b[0] = 55.0. - // set_tensor_data auto-waits for the noop to complete. + // Step 12: WAR on an external tensor + // Plain add_input would not publish the reader. add_tracked_input keeps + // the argument read-only while making the reader visible to the host write. // ========================================================= { CoreTaskArgs args; - args.add_output(ext_b); // write-only: creates TensorMap entry (not add_input!) + args.add_tracked_input(ext_b); rt_submit_aiv_task(FUNC_NOOP, args); } idx[0] = 0; set_tensor_data(ext_b, 1, idx, 55.0f); float ext_war_val = get_tensor_data(ext_b, 1, idx); - LOG_INFO( - "External WAR (INOUT): set_tensor_data(ext_b, 55.0) = %f (expected 55.0)", static_cast(ext_war_val) - ); + LOG_INFO("External WAR: set_tensor_data(ext_b, 55.0) = %f (expected 55.0)", static_cast(ext_war_val)); check_idx[0] = 8; set_tensor_data(ext_check, 1, check_idx, ext_war_val); diff --git a/examples/a5/tensormap_and_ringbuffer/scalar_data/test_scalar_data.py b/examples/a5/tensormap_and_ringbuffer/scalar_data/test_scalar_data.py index 5a2b159e5a..260986c723 100644 --- a/examples/a5/tensormap_and_ringbuffer/scalar_data/test_scalar_data.py +++ b/examples/a5/tensormap_and_ringbuffer/scalar_data/test_scalar_data.py @@ -7,7 +7,7 @@ # 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. # ----------------------------------------------------------------------------------------------------------- -"""Scalar data dependency test: GetTensorData, SetTensorData, add_inout. +"""Scalar data dependency test: GetTensorData, SetTensorData, and tracked readers. Tests orchestration-level data manipulation: scalar initialization, Get/Set round-trips, WAW+WAR dependency auto-wait, and external tensor WAR. @@ -21,7 +21,7 @@ @scene_test(level=2, runtime="tensormap_and_ringbuffer") class TestScalarData(SceneTestCase): - """Scalar data dependency: Get/SetTensorData, add_inout with initial value.""" + """Scalar data dependency: host access, initial values, and tracked readers.""" CALLABLE = { "orchestration": { @@ -77,8 +77,8 @@ def compute_golden(self, args, params): args.check[4] = 79.0 # orchestration arithmetic: 2.0 + 77.0 args.check[5] = 42.0 # Orch set->get round-trip: SetTensorData then GetTensorData args.check[6] = 12.0 # Orch->AICore RAW: SetTensorData(d,10.0) + kernel_add(d,a) -> 10.0+2.0 - args.check[7] = 88.0 # WAW+WAR: kernel reads c, SetTensorData(c,88.0) auto-waits - args.check[8] = 55.0 # External WAR: noop(ext_b INOUT) -> SetTensorData(ext_b,55.0) auto-waits + args.check[7] = 88.0 # host write to c completes after the tracked reader + args.check[8] = 55.0 # host write to external b completes after the tracked reader if __name__ == "__main__": diff --git a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md index d50bb8b511..dc2268415e 100644 --- a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md @@ -165,13 +165,14 @@ AIC/MIX tasks use the available cluster count. ### 5.2 TensorMap and Fanins -TensorMap maps tensor regions to producer task IDs. For every task: +TensorMap indexes unfinished reader and writer accesses. For every task: -1. INPUT/INOUT regions look up overlapping producers. -2. Explicit and discovered producers are deduplicated into +1. INPUT/TRACKED_INPUT/INOUT regions look up overlapping writers. +2. INOUT/OUTPUT_EXISTING regions look up overlapping tracked readers. +3. Explicit and discovered producers are deduplicated into `fanin_local_ids[]`. -3. OUTPUT/INOUT regions register the new task as producer. -4. Each producer tracks its highest consumer local ID for completion metadata. +4. TRACKED_INPUT registers a reader; INOUT/OUTPUT_EXISTING register writers. +5. Each producer tracks its highest consumer local ID for completion metadata. There is no fanout adjacency or dependency pool. A per-slot completion flag is the readiness truth on device. @@ -216,10 +217,12 @@ could strand the drain protocol, so the active path relies on that invariant. ## 8. Scalar Access During Construction -`get_tensor_data` and `set_tensor_data` operate on registered host views of -external tensors. They cannot wait for a submitted device producer because the -device scheduler starts only after orchestration returns. Runtime-created graph- -heap outputs also have no host view. +`get_tensor_data` reads registered host views of external tensors and therefore +cannot wait for a submitted device producer: the scheduler starts only after +orchestration returns. Before any task is submitted, `set_tensor_data` updates +that host view directly. After submission begins, it emits a scheduler-local +host-write node that waits for conflicting writers, their direct consumers, and +tracked readers; this also supports runtime-created graph-heap outputs. Producer references are checked against the complete bound descriptor ID before a slot is used, preventing masked-slot aliasing. See diff --git a/src/a2a3/runtime/host_build_graph/docs/SCALAR_DATA_ACCESS.md b/src/a2a3/runtime/host_build_graph/docs/SCALAR_DATA_ACCESS.md index c71dfb5036..aaca485238 100644 --- a/src/a2a3/runtime/host_build_graph/docs/SCALAR_DATA_ACCESS.md +++ b/src/a2a3/runtime/host_build_graph/docs/SCALAR_DATA_ACCESS.md @@ -9,15 +9,16 @@ external tensors; they do not interleave host code with device execution. | Tensor state | `get_tensor_data` | `set_tensor_data` | | ------------ | ----------------- | ----------------- | -| External tensor with no submitted producer | Reads the staged host value | Updates the staged host value | +| External tensor with no submitted producer | Reads the staged host value | Updates the host view before submission; otherwise emits a host-write node | | External control/output tensor not referenced by a task | Reads immediately | Writes immediately | -| Runtime-created output | Unsupported: no registered host view | Unsupported: no registered host view | -| Tensor owned by a submitted device task | Unsupported during graph construction | Unsupported during graph construction | +| Runtime-created output | Unsupported: no registered host view | Emits a host-write node after its producer | +| Tensor owned by a submitted device task | Unsupported during graph construction | Emits a host-write node after conflicting accesses | | Tensor with an invalid or stale owner task ID | Fails with `INVALID_ARGS` | Fails with `INVALID_ARGS` | -The supported write changes the data that will be copied to the device. Every -task in the graph observes that final staged value; submit order does not turn -the write into a barrier between kernels. +A write before the first submission changes the data staged for the device. A +later write becomes a dependency-aware scheduler node: it waits for overlapping +writers, their direct consumers, and tracked readers, then publishes itself as +the next writer for subsequent tasks. ## API @@ -28,9 +29,9 @@ int32_t value = get_tensor_data(control, 1, index); set_tensor_data(layout, 1, index, value + 1); ``` -Both tensors in this example must be external tensors staged by the host. A -common use is to read an input control value or publish runtime geometry into an -external layout tensor that no submitted task owns. +`control` must be an external tensor staged by the host. `layout` may also be +external, or it may be a runtime-created tensor when the write occurs after its +allocation task has been submitted. ## Why Device-Produced Values Cannot Be Read Here @@ -41,14 +42,14 @@ The execution order is: 3. The host relocates and copies the graph image to device memory. 4. AICPU schedulers boot and dispatch the graph. -A producer submitted in step 1 cannot become `COMPLETED` until step 4. Waiting -for that producer from the orchestration call cannot make progress. The runtime -keeps a timeout as a defensive failure backstop, but it is not a supported -synchronization mechanism. +A producer submitted in step 1 cannot become `COMPLETED` until step 4. A host +read that waits from orchestration therefore cannot make progress. A host write +does not wait synchronously; it is represented in the graph and executes in +step 4. -Runtime-created output buffers also live in the graph heap and have no host-view -registration. Even if their address is nonzero, host orchestration must not -dereference or modify them. +Runtime-created output buffers live in the graph heap and have no host-view +registration. Host orchestration must not dereference them, but a host-write +node may modify them on the device after their producer completes. ## Ownership Validation @@ -65,9 +66,11 @@ recorded. ## Practical Rules -- Use scalar access only on external, host-staged tensors with no device - producer or outstanding device consumer. -- Use tensor dependencies to order device tasks; do not use host scalar access - as a device synchronization barrier. +- Use `get_tensor_data` only on external, host-staged tensors with no device + producer. +- Mark pure readers with `add_tracked_input()` when a later host write overlaps. +- Use tensor dependencies to order device tasks; a host write is a graph node, + not a synchronous device barrier. - Pass values needed for graph construction as orchestration inputs or scalars. -- Keep device-produced values on the device or return them after the run. +- Keep device-produced values on the device or return them after the run; they + cannot drive host control flow during graph construction. diff --git a/src/a2a3/runtime/host_build_graph/host/dep_gen_host_graph.cpp b/src/a2a3/runtime/host_build_graph/host/dep_gen_host_graph.cpp index a11f05f3d5..5ee0f7cc14 100644 --- a/src/a2a3/runtime/host_build_graph/host/dep_gen_host_graph.cpp +++ b/src/a2a3/runtime/host_build_graph/host/dep_gen_host_graph.cpp @@ -18,9 +18,10 @@ * creator — creator retention on an existing tensor (STEP 3 Step A). * tensormap — a producer whose written slice overlaps what this task reads * (STEP 3 Step B); carries both slices' geometry. + * host_write_consumer — a direct consumer that a host write must wait for. * * Per-task producer dedup mirrors PTO2FaninBuilder::append_fanin_or_fail, which - * collapses all three sources into one fanin list: the first edge to name a + * collapses all four sources into one fanin list: the first edge to name a * producer is kept. tensormap edges are exempt — a second producer slice for the * same task is a distinct fact about the data flow, and viewers rely on seeing * every overlap. @@ -46,8 +47,8 @@ namespace { // Graph tables (serialized as tasks[] / tensors[] / edges[]) // --------------------------------------------------------------------------- -// Edge categories — matches the three places a runtime fanin edge is born. -enum class EdgeSource { EXPLICIT, CREATOR, TENSORMAP }; +// Edge categories — matches the four places a runtime fanin edge is born. +enum class EdgeSource { EXPLICIT, CREATOR, TENSORMAP, HOST_WRITE_CONSUMER }; const char *edge_source_str(EdgeSource s) { switch (s) { @@ -57,6 +58,8 @@ const char *edge_source_str(EdgeSource s) { return "creator"; case EdgeSource::TENSORMAP: return "tensormap"; + case EdgeSource::HOST_WRITE_CONSUMER: + return "host_write_consumer"; } return "unknown"; } @@ -73,6 +76,20 @@ const char *overlap_status_str(OverlapStatus s) { return "unknown"; } +const char *hazard_kind_str(TensorHazardKind kind) { + switch (kind) { + case TensorHazardKind::RAW: + return "RAW"; + case TensorHazardKind::WAW: + return "WAW"; + case TensorHazardKind::WAR: + return "WAR"; + } + return "unknown"; +} + +const char *access_kind_str(TensorAccessKind kind) { return kind == TensorAccessKind::READER ? "READER" : "WRITER"; } + const char *arg_type_str(TensorArgType t) { switch (t) { case TensorArgType::INPUT: @@ -83,6 +100,10 @@ const char *arg_type_str(TensorArgType t) { return "INOUT"; case TensorArgType::OUTPUT_EXISTING: return "OUTPUT_EXISTING"; + case TensorArgType::NO_DEP: + return "NO_DEP"; + case TensorArgType::TRACKED_INPUT: + return "TRACKED_INPUT"; } return "UNKNOWN"; } @@ -100,7 +121,9 @@ struct EdgeAnnot { int32_t consumer_arg_idx; // -1 for EXPLICIT (not tied to a tensor arg) EdgeSource source; OverlapStatus overlap; // only meaningful for TENSORMAP - uint64_t tensor_id; // 0 for EXPLICIT + TensorHazardKind hazard; + TensorAccessKind access_kind; + uint64_t tensor_id; // 0 for EXPLICIT // Consumer side (the ChipTensor the submitting task is reading). uint8_t consumer_dtype; uint32_t consumer_ndims; @@ -340,6 +363,8 @@ bool write_deps_json( out << ",\"source\":\"" << edge_source_str(e.source) << '"'; if (e.source == EdgeSource::TENSORMAP) { out << ",\"overlap\":\"" << overlap_status_str(e.overlap) << '"'; + out << ",\"hazard\":\"" << hazard_kind_str(e.hazard) << '"'; + out << ",\"access_kind\":\"" << access_kind_str(e.access_kind) << '"'; } if (e.source != EdgeSource::EXPLICIT) { out << ",\"tensor_id\":\"" << e.tensor_id << '"'; @@ -466,7 +491,7 @@ void dep_gen_host_graph_add_creator_edge(uint64_t producer_raw, int32_t arg_idx, void dep_gen_host_graph_add_tensormap_edge( uint64_t producer_raw, int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, - OverlapStatus overlap + OverlapStatus overlap, TensorHazardKind hazard ) { HostGraphState &s = state(); if (!s.enabled || !s.in_task) { @@ -481,12 +506,31 @@ void dep_gen_host_graph_add_tensormap_edge( e.consumer_arg_idx = arg_idx; e.source = EdgeSource::TENSORMAP; e.overlap = overlap; + e.hazard = hazard; + e.access_kind = entry.access_kind; e.tensor_id = make_tensor_id(entry.buffer_addr, entry.version); fill_consumer(e, consumer); fill_producer(e, entry); s.edges.push_back(e); } +void dep_gen_host_graph_add_host_write_consumer_edge( + uint64_t producer_raw, int32_t arg_idx, const ChipTensor &consumer +) { + HostGraphState &s = state(); + if (!s.enabled || !s.in_task || !s.task_preds.insert(producer_raw).second) { + return; + } + EdgeAnnot e{}; + e.pred = producer_raw; + e.succ = s.current_task_id; + e.consumer_arg_idx = arg_idx; + e.source = EdgeSource::HOST_WRITE_CONSUMER; + e.tensor_id = register_tensor(s, consumer); + fill_consumer(e, consumer); + s.edges.push_back(e); +} + // --------------------------------------------------------------------------- // Control surface // --------------------------------------------------------------------------- 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 572d3828a1..b0daa3b757 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -540,6 +540,26 @@ int32_t run_host_orchestration( rt_scope_end(rt); rt_orchestration_done(rt); +#if SIMPLER_ORCH_PROFILING + const PTO2OrchProfilingData profile = orchestrator_get_profiling(); + const uint64_t profile_total = profile.sync_cycle + profile.alloc_cycle + profile.args_cycle + + profile.lookup_cycle + profile.insert_cycle + profile.fanin_cycle + + profile.scope_end_cycle; + const double profile_us = static_cast(profile_total) * 1000000.0 / PLATFORM_PROF_SYS_CNT_FREQ; + LOG_INFO( + "Host Orchestrator Profiling: %" PRId64 " tasks, total=%.3fus, avg/task=%.3fus", + static_cast(profile.submit_count), profile_us, + profile.submit_count > 0 ? profile_us / profile.submit_count : 0.0 + ); +#if SIMPLER_TENSORMAP_PROFILING + rt->orchestrator.tensor_map.print_stats(); + LOG_INFO( + "Host fanin spill: top=%d used=%d capacity=%d", rt->scheduler.fanin_spill_top, rt->scheduler.fanin_spill_top, + rt->scheduler.fanin_spill_capacity + ); +#endif +#endif + const int32_t total_tasks = pto2_sm_layout::ring_current_task_index_addr(host_sm)->load(std::memory_order_acquire); if (!upload_graph_submissions(runtime, api, *graph_state)) return -1; @@ -929,7 +949,7 @@ extern "C" int bind_callable_to_runtime_impl( // on-device nodes that push past the host task count. const auto &sq = layout.sched; const size_t orch_start = layout.orch.off_fanin_seen_epoch; - const size_t orch_end = sq.off_ready_queue_slots[0]; + const size_t orch_end = sq.off_fanin_spill_ids; always_assert(orch_start <= orch_end); char *arena_host = static_cast(host_arena.base()); char *arena_dev = static_cast(runtime_arena_dev); diff --git a/src/a2a3/runtime/host_build_graph/orchestration/pto_arg_with_deps.h b/src/a2a3/runtime/host_build_graph/orchestration/pto_arg_with_deps.h index 536a6f8084..98ba192733 100644 --- a/src/a2a3/runtime/host_build_graph/orchestration/pto_arg_with_deps.h +++ b/src/a2a3/runtime/host_build_graph/orchestration/pto_arg_with_deps.h @@ -54,6 +54,7 @@ class CoreTaskArgsWithDeps : private CoreTaskArgs { using CoreTaskArgs::add_scalar; using CoreTaskArgs::add_scalars; using CoreTaskArgs::add_scalars_i32; + using CoreTaskArgs::add_tracked_input; using CoreTaskArgs::allow_early_resolve; // speculative early-dispatch hint (getter) using CoreTaskArgs::copy_scalars_from; using CoreTaskArgs::set_allow_early_resolve; // speculative early-dispatch hint (setter) diff --git a/src/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.h b/src/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.h index 8212f217d4..2720980355 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.h +++ b/src/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.h @@ -28,6 +28,7 @@ * add_explicit_edge() — STEP 1, per declared dependency * add_creator_edge() — STEP 3 Step A, per creator-retention producer * add_tensormap_edge() — STEP 3 Step B, per tensormap producer + * add_host_write_consumer_edge() — runtime-derived host-write ordering * end_task() — closes the task, after its last dependency step * * Control surface, called from the device runner (same host_runtime.so): @@ -54,8 +55,7 @@ * consumer (deps viewer, swimlane join) reads both runtimes' output the same way. */ -#ifndef SRC_A2A3_RUNTIME_HOST_BUILD_GRAPH_RUNTIME_DEP_GEN_HOST_GRAPH_H_ -#define SRC_A2A3_RUNTIME_HOST_BUILD_GRAPH_RUNTIME_DEP_GEN_HOST_GRAPH_H_ +#pragma once #include @@ -105,7 +105,12 @@ void dep_gen_host_graph_add_creator_edge(uint64_t producer_raw, int32_t arg_idx, /** STEP 3 Step B: a tensormap producer whose written slice this task reads. */ void dep_gen_host_graph_add_tensormap_edge( uint64_t producer_raw, int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, - OverlapStatus overlap + OverlapStatus overlap, TensorHazardKind hazard +); + +/** A direct consumer of an overlapping writer that must finish before a host write. */ +void dep_gen_host_graph_add_host_write_consumer_edge( + uint64_t producer_raw, int32_t arg_idx, const ChipTensor &consumer ); // --------------------------------------------------------------------------- @@ -130,5 +135,3 @@ bool dep_gen_host_graph_active(); */ int dep_gen_host_graph_emit(const char *deps_json_path); } - -#endif // SRC_A2A3_RUNTIME_HOST_BUILD_GRAPH_RUNTIME_DEP_GEN_HOST_GRAPH_H_ diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp index 54fc335dab..6e6f0f5834 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp @@ -73,8 +73,10 @@ __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_add_explicit __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_add_creator_edge(uint64_t, int32_t, const ChipTensor &) {} __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_add_tensormap_edge( - uint64_t, int32_t, const ChipTensor &, const PTO2TensorMapEntry &, OverlapStatus + uint64_t, int32_t, const ChipTensor &, const PTO2TensorMapEntry &, OverlapStatus, TensorHazardKind ) {} +__attribute__((weak, visibility("hidden"))) void +dep_gen_host_graph_add_host_write_consumer_edge(uint64_t, int32_t, const ChipTensor &) {} // Scope_stats enable gate, queried via the same predicate idiom as // dep_gen_host_graph_enabled above. The AICPU collector links the strong definition; @@ -687,10 +689,9 @@ static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { return next; } -// Polling: fanin is a flat array of position-independent producer local ids on -// the payload (no dep-pool spill, no producer pointers). The builder writes them -// directly into payload->fanin_local_ids as producers are appended, deduping by -// slot and hard-capping at PTO2_MAX_FANIN. self_local is this task's own local id +// Polling: fanin is a flat sequence of position-independent producer local ids. +// The builder writes the inline prefix to the payload and the remainder to the +// scheduler spill pool, deduping by slot. self_local is this task's own local id // (the consumer), used to bump each producer's last_consumer_local_id (the // reclaim gate the host wait_for_consumers polls via completed_watermark). struct PTO2FaninBuilder { @@ -699,7 +700,10 @@ struct PTO2FaninBuilder { orch(orch), seen_epoch(seen_epoch), self_local(self_local), - payload(payload) {} + payload(payload) { + payload->fanin_spill_start = 0; + payload->fanin_spill_count = 0; + } int32_t count{0}; PTO2OrchestratorState *orch{nullptr}; uint32_t seen_epoch{0}; @@ -735,11 +739,22 @@ static bool append_fanin_or_fail( if (fanin_builder->mark_seen(prod_ring, prod_slot)) { return true; } - if (fanin_builder->count >= PTO2_MAX_FANIN) { - orch_mark_fatal(orch, PTO2_ERROR_DEP_POOL_OVERFLOW); - return false; + const int32_t producer_local = static_cast(producer_task_id.local()); + if (fanin_builder->count < PTO2_MAX_FANIN) { + fanin_builder->payload->fanin_local_ids[fanin_builder->count] = producer_local; + } else { + PTO2SchedulerState *sched = orch->scheduler; + if (sched == nullptr || sched->fanin_spill_top >= sched->fanin_spill_capacity) { + orch_mark_fatal(orch, PTO2_ERROR_DEP_POOL_OVERFLOW); + return false; + } + if (fanin_builder->count == PTO2_MAX_FANIN) { + fanin_builder->payload->fanin_spill_start = sched->fanin_spill_top; + } + sched->fanin_spill_ids[sched->fanin_spill_top++] = producer_local; + fanin_builder->payload->fanin_spill_count++; } - fanin_builder->payload->fanin_local_ids[fanin_builder->count++] = static_cast(producer_task_id.local()); + fanin_builder->count++; // Reclaim gate: record this task as a consumer of the producer. The producer // slot retires once the per-ring completed_watermark reaches this consumer id. @@ -1058,7 +1073,7 @@ static bool ensure_tensormap_capacity(PTO2OrchestratorState *orch, int32_t neede return true; } // Progress is entries actually freed, NOT watermark movement: a ring can - // retire zero-output tasks (count_registrable_outputs == 0), advancing + // retire zero-access tasks (count_registrable_accesses == 0), advancing // last_task_alive without freeing any entry. Gating the backstop on // free_entries() keeps a wedged pool from dodging the timeout while some // unrelated ring keeps draining. @@ -1112,7 +1127,7 @@ static bool ensure_tensormap_capacity(PTO2OrchestratorState *orch, int32_t neede // Orch-side wiring/ready publication. static TaskOutputTensors submit_task_common( PTO2OrchestratorState *orch, const CoreTaskArgs &args, ActiveMask active_mask, TaskAttrs task_attrs, - int32_t aic_kernel_id, int32_t aiv0_kernel_id, int32_t aiv1_kernel_id + int32_t aic_kernel_id, int32_t aiv0_kernel_id, int32_t aiv1_kernel_id, TaskKind task_kind = TaskKind::KERNEL ) { CYCLE_COUNT_START(); TaskOutputTensors result; @@ -1127,6 +1142,7 @@ static TaskOutputTensors submit_task_common( PTO2TaskDescriptor &task = *prepared.task; PTO2TaskPayload &payload = *prepared.payload; result.set_task_id(task_id); + prepared.slot_state->task_kind = task_kind; // dep_gen capture point: open this task's graph entry before its dependency // steps run, so the edges STEP 1 / STEP 3 discover attach to it. The graph @@ -1192,6 +1208,74 @@ static TaskOutputTensors submit_task_common( } } + // A graph-local host write preserves the synchronous TMR contract: wait for + // each overlapping writer and every direct consumer of that writer. The + // normal OUTPUT_EXISTING lookup below adds tracked readers. HBG has the + // complete task graph resident, so direct consumers are recoverable by + // scanning earlier payload fanins without serializing unrelated tasks. + if (task_kind == TaskKind::HOST_WRITE) { + const int32_t self_local = static_cast(task_id.local()); + PTO2SharedMemoryRingHeader &producer_ring = orch->sm_header->ring; + const ChipTensor &write_tensor = args.tensor(0).ref(); + auto append_predecessor = [&](PTO2TaskId predecessor) -> bool { + const int32_t predecessor_local = static_cast(predecessor.local()); + const int32_t predecessor_slot = producer_ring.get_slot_by_task_id(predecessor_local); + PTO2TaskSlotState *predecessor_state = &producer_ring.get_slot_state_by_slot(predecessor_slot); + return append_fanin_or_fail( + orch, predecessor.ring(), predecessor_slot, predecessor_state, predecessor, &fanin_builder + ); + }; + auto append_consumers = [&](PTO2TaskId writer) -> bool { + const int32_t writer_local = static_cast(writer.local()); + for (int32_t consumer_local = 0; consumer_local < self_local; ++consumer_local) { + const int32_t consumer_slot = producer_ring.get_slot_by_task_id(consumer_local); + PTO2TaskSlotState *consumer_state = &producer_ring.get_slot_state_by_slot(consumer_slot); + if (consumer_state->payload == nullptr || consumer_state->task == nullptr || + consumer_state->task->task_id.local() != static_cast(consumer_local)) { + continue; + } + const PTO2TaskPayload &consumer_payload = *consumer_state->payload; + bool reads_writer = false; + for (int32_t i = 0; i < consumer_payload.fanin_count; ++i) { + if (orch->scheduler->fanin_local_id(consumer_payload, i) == writer_local) { + reads_writer = true; + break; + } + } + if (!reads_writer) continue; + const PTO2TaskId consumer = PTO2TaskId::make(task_id.ring(), static_cast(consumer_local)); + if (!append_predecessor(consumer)) return false; + if (capture_dep_graph) { + dep_gen_host_graph_add_host_write_consumer_edge(consumer.raw, 0, write_tensor); + } + } + return true; + }; + + if (write_tensor.owner_task_id.is_valid() && !append_consumers(write_tensor.owner_task_id)) { + return result; + } + bool host_write_fatal = false; + orch->tensor_map.lookup( + write_tensor, TensorAccessKind::WRITER, + [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { + if (!append_predecessor(entry.access_task_id) || !append_consumers(entry.access_task_id)) { + host_write_fatal = true; + return false; + } + if (capture_dep_graph) { + dep_gen_host_graph_add_tensormap_edge( + entry.access_task_id.raw, 0, write_tensor, entry, overlap_status, TensorHazardKind::WAW + ); + } + return true; + } + ); + if (host_write_fatal) { + return result; + } + } + // === STEP 3: Lookup inputs (creator retention + tensormap modifier lookup) === DepInputs dep_inputs{ args.tensor_count(), args.tensor_data(), args.tag_data(), static_cast(args.explicit_dep_count()), @@ -1214,9 +1298,12 @@ static TaskOutputTensors submit_task_common( dep_gen_host_graph_add_creator_edge(producer.raw, arg_idx, consumer); } void tensormap( - int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, OverlapStatus overlap + int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, OverlapStatus overlap, + TensorHazardKind hazard ) const { - dep_gen_host_graph_add_tensormap_edge(entry.producer_task_id.raw, arg_idx, consumer, entry, overlap); + dep_gen_host_graph_add_tensormap_edge( + entry.access_task_id.raw, arg_idx, consumer, entry, overlap, hazard + ); } }; const bool ok = @@ -1240,11 +1327,11 @@ static TaskOutputTensors submit_task_common( // is reclaimed as last_task_alive advances; an // exhausted pool back-pressures here (and detects a wedged watermark) rather // than tripping new_entry()'s hard assert mid-registration. - int32_t tensormap_needed = count_registrable_outputs(dep_inputs, orch->in_manual_scope()); + int32_t tensormap_needed = count_registrable_accesses(dep_inputs, orch->in_manual_scope()); if (tensormap_needed > 0 && !ensure_tensormap_capacity(orch, tensormap_needed)) { return result; } - register_task_outputs(dep_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); + register_task_accesses(dep_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); CYCLE_COUNT_LAP(g_orch_insert_cycle); @@ -1450,7 +1537,7 @@ bool graph_submit_definition( DepInputs boundary_inputs{ args.tensor_count(), args.tensor_data(), args.tag_data(), 0, nullptr, }; - const int32_t tensormap_needed = count_registrable_outputs(boundary_inputs, orch->in_manual_scope()); + const int32_t tensormap_needed = count_registrable_accesses(boundary_inputs, orch->in_manual_scope()); if (tensormap_needed > 0 && !ensure_tensormap_capacity(orch, tensormap_needed)) return false; if (!check_scope_can_accept_task(orch, allocator, 0)) return false; @@ -1491,7 +1578,7 @@ bool graph_submit_definition( return append_fanin_or_fail(orch, producer_id.ring(), producer_slot, producer, producer_id, &fanin_builder); }; if (!compute_task_fanin(boundary_inputs, orch->tensor_map, orch->in_manual_scope(), emit)) return false; - register_task_outputs(boundary_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); + register_task_accesses(boundary_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); payload.fanin_count = fanin_builder.count; pending.outer_slot = &slot; @@ -1530,6 +1617,14 @@ TaskOutputTensors graph_record_submit_node( if (node_index >= GRAPH_MAX_NODES || args.has_error || args.predicate().op != PredicateOp::NONE) { recording.unsupported = true; } + for (int32_t i = 0; i < args.tensor_count(); ++i) { + if (args.tag(i) == TensorArgType::TRACKED_INPUT) { + // The compact Graph format derives only producer-source fanins; it + // cannot publish an internal reader for a later WAR lookup. + recording.unsupported = true; + break; + } + } const PTO2OutputLayout layout = calculate_output_layout(args); void *packed_base = orch->ring.task_allocator.reserve_heap_scratch(layout.total_output_size); @@ -1733,7 +1828,6 @@ bool PTO2OrchestratorState::graph_end() { std::vector definition; if (!graph_build_definition(*recording, &definition)) { - debug_assert(false && "The recorded Graph contains a construct that Graph Execution does not support"); LOG_WARN("%s", "[GraphExecution] unsupported construct observed; falling back to the ordinary path"); return false; } @@ -1890,7 +1984,38 @@ TaskOutputTensors PTO2OrchestratorState::submit_dummy_task(const CoreTaskArgs &a } return submit_task_common( - orch, args, ActiveMask{}, task_attrs, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID + orch, args, ActiveMask{}, task_attrs, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID, TaskKind::DUMMY + ); +} + +void PTO2OrchestratorState::submit_host_write( + const ChipTensor &tensor, uint64_t address, uint64_t value, uint64_t size +) { + if (fatal) return; + if (size == 0 || size > sizeof(value)) { + report_fatal( + PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "host write size=%llu is outside [1,8]", (unsigned long long)size + ); + return; + } + + CoreTaskArgs args; + args.add_output(tensor); // OUTPUT_EXISTING: WAR lookup, then writer registration. + args.add_scalar(address); + args.add_scalar(value); + args.add_scalar(size); + + if (GraphHostState *state = graph_state_from(this); state != nullptr && state->recording != nullptr) { + // The cached Graph format has no host-write node. Preserve recording + // task-id arithmetic, then force graph_end to replay the ordinary path. + (void)submit_dummy_task(args); + graph_record_mark_unsupported(this); + return; + } + + (void)submit_task_common( + this, args, ActiveMask{}, TaskAttrs{}, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID, + TaskKind::HOST_WRITE ); } diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp index c122e8562a..d4ef32271b 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp @@ -154,6 +154,7 @@ wait_for_tensor_ready(PTO2Runtime *rt, const ChipTensor &tensor, bool wait_for_c // the second encounter. constexpr int kSegmentCap = 64; const PTO2TaskSlotState *seg[kSegmentCap]; + bool seg_wait_consumers[kSegmentCap]{}; int seg_count = 0; bool failed = false; @@ -250,22 +251,26 @@ wait_for_tensor_ready(PTO2Runtime *rt, const ChipTensor &tensor, bool wait_for_c for (int i = 0; i < seg_count; i++) { wait_one_producer(*seg[i]); if (failed) return; - if (!wait_for_consumers) continue; + if (!seg_wait_consumers[i]) continue; wait_one_consumers(*seg[i]); if (failed) return; } seg_count = 0; }; - auto try_push = [&](const PTO2TaskSlotState &s) { + auto try_push = [&](const PTO2TaskSlotState &s, bool include_consumers) { for (int j = 0; j < seg_count; j++) { - if (seg[j] == &s) return; // per-segment dedup + if (seg[j] == &s) { + seg_wait_consumers[j] = seg_wait_consumers[j] || include_consumers; + return; + } } if (seg_count == kSegmentCap) { flush_segment(); if (failed) return; } - seg[seg_count++] = &s; + seg[seg_count] = &s; + seg_wait_consumers[seg_count++] = include_consumers; }; auto do_wait = [&]() { @@ -273,17 +278,27 @@ wait_for_tensor_ready(PTO2Runtime *rt, const ChipTensor &tensor, bool wait_for_c if (owner.is_valid()) { const auto *slot = resolve_producer(owner); if (slot == nullptr) return; - try_push(*slot); + try_push(*slot, wait_for_consumers); } // Step B: modifier writer lookup (OverlapMap), direct callback orch.tensor_map.lookup(tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus) -> bool { - PTO2TaskId pid = entry.producer_task_id; + PTO2TaskId pid = entry.access_task_id; const auto *slot = resolve_producer(pid); if (slot == nullptr) return false; - try_push(*slot); + try_push(*slot, wait_for_consumers); return !failed; }); + if (wait_for_consumers && !failed) { + orch.tensor_map.lookup( + tensor, TensorAccessKind::READER, [&](PTO2TensorMapEntry &entry, OverlapStatus) -> bool { + const auto *slot = resolve_producer(entry.access_task_id); + if (slot == nullptr) return false; + try_push(*slot, false); + return !failed; + } + ); + } if (failed) return; flush_segment(); }; @@ -333,22 +348,24 @@ void set_tensor_data( return; } - // Wait for producer + all consumers before writing (WAW + WAR safety) - if (!wait_for_tensor_ready(rt, tensor, true, __FUNCTION__)) { - return; - } - uint64_t flat_offset = tensor.compute_flat_offset(indices, ndims); uint64_t elem_size = get_element_size(tensor.dtype); uint64_t elem_addr = tensor.buffer.addr + flat_offset * elem_size; - if (!host_tensor_write(rt->tensor_access, elem_addr, &value, elem_size)) { - rt->orchestrator.report_fatal( - PTO2_ERROR_INVALID_ARGS, __FUNCTION__, - "no writable host view for device address %#llx (%llu bytes): during host orchestration only tensors " - "the runtime staged are writable, not runtime-created or child-memory buffers", - (unsigned long long)elem_addr, (unsigned long long)elem_size - ); + if (rt->orchestrator.sm_header->ring.fc.current_task_index.load(std::memory_order_acquire) == 0) { + if (!host_tensor_write(rt->tensor_access, elem_addr, &value, elem_size)) { + rt->orchestrator.report_fatal( + PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "no writable host view for device address %#llx (%llu bytes)", + (unsigned long long)elem_addr, (unsigned long long)elem_size + ); + } + return; } + // HBG constructs the whole graph before device scheduling starts, so a + // synchronous host wait here can never observe a submitted reader finish. + // Represent the write as a scheduler-local graph node instead: normal + // OUTPUT_EXISTING dependency discovery gives it WAR predecessors and + // registers it as the next writer for tasks submitted afterwards. + rt->orchestrator.submit_host_write(tensor, elem_addr, value, elem_size); } // Ops-table entry that hands the call-site captured by PTO2ScopeGuard to the diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_dep_compute.h b/src/a2a3/runtime/host_build_graph/runtime/pto_dep_compute.h index befb683dc4..09cb860562 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_dep_compute.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_dep_compute.h @@ -16,11 +16,11 @@ * Two header-only template entry points: * * compute_task_fanin — STEP 3 in submit_task: per-tensor creator retention (Step A) - * + tensormap.lookup for INPUT/INOUT (Step B). Calls back into - * user-supplied `emit` for each producer it identifies. + * + TensorMap lookup for INPUT/TRACKED_INPUT/INOUT (Step B). + * Calls back into user-supplied `emit` for each producer. * - * register_task_outputs — STEP 4 in submit_task: tensormap.insert for INOUT and - * OUTPUT_EXISTING tensors. No callbacks. + * register_task_accesses — STEP 4 in submit_task: register TRACKED_INPUT readers and + * INOUT/OUTPUT_EXISTING writers. No callbacks. * * STEP 1 (explicit_deps) is intentionally left at the runtime call site because its * `last_task_alive` shortcut + unchecked slot lookup is subtly different from the @@ -49,8 +49,7 @@ * inlining and add ~5 ns/call to the orch hot path. */ -#ifndef SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_DEP_COMPUTE_H_ -#define SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_DEP_COMPUTE_H_ +#pragma once #include @@ -60,7 +59,7 @@ #include "tensor.h" /** - * View struct for inputs to compute_task_fanin / register_task_outputs. + * View struct for inputs to compute_task_fanin / register_task_accesses. * * Both runtime and replay assemble one of these from their own data sources * (runtime: from Arg accessors; replay: from SubmitTraceEntry fields). All @@ -80,7 +79,7 @@ struct DepInputs { */ struct NoDepAnnotate { void creator(int32_t, const ChipTensor &, PTO2TaskId) const {} - void tensormap(int32_t, const ChipTensor &, const PTO2TensorMapEntry &, OverlapStatus) const {} + void tensormap(int32_t, const ChipTensor &, const PTO2TensorMapEntry &, OverlapStatus, TensorHazardKind) const {} }; /** @@ -89,8 +88,8 @@ struct NoDepAnnotate { * * For each non-OUTPUT tensor: * - If owner_task_id is valid, emit(owner) - * - For INPUT/INOUT (and not manual_dep), tensor_map.lookup(*tensor) and emit - * each matching producer. INOUT+COVERED triggers tensor_map.remove_entry(entry). + * - For INPUT/TRACKED_INPUT/INOUT (and not manual_dep), query overlapping writers. + * - For INOUT/OUTPUT_EXISTING, query overlapping tracked readers. * * @return true on success (or producer-skipped-silently); false if emit signaled * fatal — caller should propagate (after any fatal bookkeeping done by emit). @@ -122,26 +121,37 @@ template annotate.creator(i, *tensor, owner); } - // Step B: only INPUT/INOUT need modifier dependency lookup. - if (ptype != TensorArgType::INPUT && ptype != TensorArgType::INOUT) { - continue; - } if (tensor->manual_dep) { continue; } bool fatal = false; - tensor_map.lookup(*tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { - if (!emit(entry.producer_task_id)) { - fatal = true; - return false; // stop iteration - } - annotate.tensormap(i, *tensor, entry, overlap_status); - if (ptype == TensorArgType::INOUT && overlap_status == OverlapStatus::COVERED) { - tensor_map.remove_entry(entry); - } - return true; - }); + auto lookup = [&](TensorAccessKind access_kind, TensorHazardKind hazard) { + tensor_map.lookup( + *tensor, access_kind, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { + if (!emit(entry.access_task_id)) { + fatal = true; + return false; // stop iteration + } + annotate.tensormap(i, *tensor, entry, overlap_status, hazard); + if (ptype == TensorArgType::INOUT && access_kind == TensorAccessKind::WRITER) { + annotate.tensormap(i, *tensor, entry, overlap_status, TensorHazardKind::WAW); + } + if ((ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) && + overlap_status == OverlapStatus::COVERED) { + tensor_map.remove_entry(entry); + } + return true; + } + ); + }; + if (ptype == TensorArgType::INPUT || ptype == TensorArgType::TRACKED_INPUT || ptype == TensorArgType::INOUT) { + lookup(TensorAccessKind::WRITER, TensorHazardKind::RAW); + } + if (!fatal && tensor_map.current_readers() != 0 && + (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING)) { + lookup(TensorAccessKind::READER, TensorHazardKind::WAR); + } if (fatal) { return false; } @@ -150,45 +160,47 @@ template } /** - * Register a task's outputs in the tensormap (STEP 4 in submit_task). - * - * For INOUT and OUTPUT_EXISTING tensors (excluding manual_dep), inserts the - * tensor into tensor_map keyed by its buffer.addr with `task_id` as producer. + * Register a task's accesses in the tensormap (STEP 4 in submit_task). * * No-op when in_manual_scope. */ inline void -register_task_outputs(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap &tensor_map, bool in_manual_scope) { +register_task_accesses(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap &tensor_map, bool in_manual_scope) { if (in_manual_scope) { return; } for (int32_t i = 0; i < inputs.tensor_count; i++) { TensorArgType ptype = inputs.arg_types[i]; - if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) { + if (ptype == TensorArgType::TRACKED_INPUT || ptype == TensorArgType::INOUT || + ptype == TensorArgType::OUTPUT_EXISTING) { const ChipTensor *tensor = &inputs.tensors[i].ref(); if (!tensor->manual_dep) { - tensor_map.insert(*tensor, task_id); + if (ptype == TensorArgType::TRACKED_INPUT) { + tensor_map.insert(*tensor, task_id, TensorAccessKind::READER); + } else { + tensor_map.insert(*tensor, task_id); + } } } } } /** - * Count the tensormap entries register_task_outputs() will insert for this task. + * Count the tensormap entries register_task_accesses() will insert for this task. * - * Mirrors register_task_outputs()'s selection exactly (INOUT / OUTPUT_EXISTING, - * excluding manual_dep), so the returned value is the precise number of + * Mirrors register_task_accesses() exactly, so the returned value is the precise number of * new_entry() calls that step makes. The orchestrator uses it to reserve pool * capacity before inserting. Returns 0 in a manual scope (no registration). */ -inline int32_t count_registrable_outputs(const DepInputs &inputs, bool in_manual_scope) { +inline int32_t count_registrable_accesses(const DepInputs &inputs, bool in_manual_scope) { if (in_manual_scope) { return 0; } int32_t needed = 0; for (int32_t i = 0; i < inputs.tensor_count; i++) { TensorArgType ptype = inputs.arg_types[i]; - if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) { + if (ptype == TensorArgType::TRACKED_INPUT || ptype == TensorArgType::INOUT || + ptype == TensorArgType::OUTPUT_EXISTING) { if (!inputs.tensors[i].ref().manual_dep) { needed++; } @@ -196,5 +208,3 @@ inline int32_t count_registrable_outputs(const DepInputs &inputs, bool in_manual } return needed; } - -#endif // SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_DEP_COMPUTE_H_ diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h b/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h index b669d237dc..e9cbddd547 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h @@ -9,7 +9,7 @@ * ----------------------------------------------------------------------------------------------------------- */ /** - * PTO Runtime2 - Orchestrator Interface + * host_build_graph - Orchestrator Interface * * The Orchestrator is responsible for: * 1. Executing the orchestration function (Turing-complete control flow) @@ -25,8 +25,7 @@ * Based on: docs/RUNTIME_LOGIC.md */ -#ifndef PTO_ORCHESTRATOR_H -#define PTO_ORCHESTRATOR_H +#pragma once #include "common/chip_swimlane_profiling.h" #include "utils/device_arena.h" @@ -161,6 +160,7 @@ struct PTO2OrchestratorState { void end_scope(); TaskOutputTensors submit_task(const MixedKernels &mixed_kernels, const CoreTaskArgs &args); TaskOutputTensors submit_dummy_task(const CoreTaskArgs &args); + void submit_host_write(const ChipTensor &tensor, uint64_t address, uint64_t value, uint64_t size); TaskOutputTensors alloc_tensors(const CoreTaskArgs &args); GraphScopeResult graph_begin(uint64_t graph_key, const CoreTaskArgs &args, uint64_t callable_hash); bool graph_end(); @@ -193,5 +193,3 @@ struct PTO2OrchProfilingData { PTO2OrchProfilingData orchestrator_get_profiling(); #endif - -#endif // PTO_ORCHESTRATOR_H diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h index 17770d2141..f22cdf04be 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h @@ -10,9 +10,9 @@ */ /** - * PTO Runtime2 - Core Type Definitions + * host_build_graph - Core Type Definitions * - * This header defines all fundamental types used by the PTO Runtime2 system: + * This header defines the fundamental host_build_graph runtime types: * - Configuration constants * - Worker types and task states * - ChipTensor regions and task parameters @@ -22,8 +22,7 @@ * Based on: docs/RUNTIME_LOGIC.md */ -#ifndef SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_RUNTIME2_TYPES_H_ -#define SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_RUNTIME2_TYPES_H_ +#pragma once #include #include @@ -85,6 +84,7 @@ #define PTO2_HEAP_SIZE (256 * 1024 * 1024) // 256MB #define PTO2_TENSORMAP_POOL_SIZE (65536) // TensorMap entry pool #define PTO2_TENSORMAP_NUM_BUCKETS 4096 // Power of 2 for fast hash (4096×8B=32KB fits L1) +#define PTO2_TENSORMAP_READER_NUM_BUCKETS 512 // Sparse tracked-reader sidecar (512×8B=4KB) // Scope management #define PTO2_MAX_SCOPE_DEPTH 64 // Maximum nesting depth @@ -109,11 +109,10 @@ // Fanin storage #define PTO2_FANIN_INLINE_CAP 64 -// Polling-scheduler inline fanin cap. The polling model stores producer -// dependencies as flat position-independent local-id integers on the payload -// (no dep-pool spill), so a task's fanin degree is hard-capped here. Must cover -// the worst-case fanin of any workload (paged_attention is the densest). +// Polling-scheduler inline fanin cap. Additional position-independent local ids +// continue in the scheduler-owned spill pool. #define PTO2_MAX_FANIN 128 +#define HBG_FANIN_SPILL_POOL_SIZE 262144 // Dependency-degree diagnostic: warn once when a task's fanin or a producer's // fanout first exceeds this degree, so dense dependency graphs surface without @@ -176,6 +175,7 @@ enum class TaskKind : uint8_t { DUMMY = 1, GRAPH = 2, GRAPH_NODE = 3, + HOST_WRITE = 4, }; struct PTO2OutputLayout { @@ -270,12 +270,14 @@ struct PTO2TaskPayload { // Producer dependencies as position-independent local task ids. Single-ring // hbg: every producer is ring 0, so no per-edge ring id is stored. Scanned // by fanin_satisfied / classify_fanin_state against the ring completion_flags. - // Hard-capped at PTO2_MAX_FANIN (no dep-pool spill). + // The first PTO2_MAX_FANIN ids are inline; additional ids use fanin_spill_*. int32_t fanin_local_ids[PTO2_MAX_FANIN]; // Reserved: preserves the early-dispatch block and tensors[] offsets. tensors // must stay at byte 576 (AICore arg-materialization contract), so this fanin // region keeps its original 528-byte footprint. - int32_t _fanin_reserved[3]; + int32_t fanin_spill_start{0}; + int32_t fanin_spill_count{0}; + int32_t _fanin_reserved{0}; // Early-dispatch metadata (AICPU-side only). Ordered by descending // alignment (8B mask, 4B fanin, then 2B/1B counters and flags) so the block packs with no // internal padding. Kept here after the fanin array (not moved up front): on @@ -589,5 +591,3 @@ static_assert(sizeof(PTO2TaskSlotState) == 64); // Sentinel marking a wake list as "owner already completed; no more // registrations accepted". Distinct from any real slot_state pointer. inline PTO2TaskSlotState *const WAKE_LIST_SENTINEL = reinterpret_cast(static_cast(0x1)); - -#endif // SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_RUNTIME2_TYPES_H_ diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_tensormap.h b/src/a2a3/runtime/host_build_graph/runtime/pto_tensormap.h index e925337307..9db9388940 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_tensormap.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_tensormap.h @@ -10,10 +10,10 @@ */ /** - * PTO Runtime2 - TensorMap Interface + * host_build_graph - TensorMap Interface * - * TensorMap provides producer lookup for dependency discovery: - * - Maps ChipTensor -> producer task ID + * TensorMap provides access lookup for dependency discovery: + * - Maps ChipTensor -> reader/writer task IDs * - Used by pto_submit_task() to find dependencies * * Key design features: @@ -57,6 +57,17 @@ enum class OverlapStatus { OTHER, }; +enum class TensorAccessKind : uint8_t { + READER = 0, + WRITER = 1, +}; + +enum class TensorHazardKind : uint8_t { + RAW = 0, + WAW = 1, + WAR = 2, +}; + struct Segment { uint64_t begin; uint64_t end; @@ -74,10 +85,12 @@ struct Segment { */ struct PTO2TensorMapLayout { size_t off_buckets; + size_t off_reader_buckets; size_t off_entry_pool; size_t off_free_entry_list; size_t off_task_entry_heads; int32_t num_buckets; + int32_t num_reader_buckets; int32_t pool_size; int32_t task_window_size; }; @@ -106,7 +119,7 @@ extern uint64_t g_insert_count; * the hash key, size in [8, 16) is unused by the entry — we repurpose it for * `next_in_bucket`). * - * buffer_addr / next_in_bucket / producer_task_id — chain traversal + match + * buffer_addr / next_in_bucket / access_task_id — chain traversal + match * start_offset — overlap byte range begin * version, ndims, dtype, manual_dep, is_contiguous — overlap fast path * shapes[5] — overlap comparison (line 1) @@ -126,15 +139,15 @@ struct alignas(64) PTO2TensorMapEntry { // === Cache line 1 (64B) — lookup hot path; mirrors ChipTensor line 1 from byte 16 === uint64_t buffer_addr; // 8B [0, 8): tensor base address (hash key, mirrors ChipTensor::buffer.addr) PTO2TensorMapEntry - *next_in_bucket; // 8B [8, 16): next entry in hash bucket chain (overlays ChipTensor::buffer.size) - PTO2TaskId producer_task_id; // 8B [16,24): mirrors ChipTensor::owner_task_id slot - uint64_t start_offset; // 8B [24,32): mirrors ChipTensor::start_offset (element offset) - int32_t version; // 4B [32,36): mirrors ChipTensor::version - uint32_t ndims; // 4B [36,40): mirrors ChipTensor::ndims - DataType dtype; // 1B [40,41): mirrors ChipTensor::dtype - bool manual_dep; // 1B [41,42): mirrors ChipTensor::manual_dep - bool is_contiguous; // 1B [42,43): mirrors ChipTensor::is_contiguous - uint8_t __padding1__; // 1B [43,44): mirrors ChipTensor padding + *next_in_bucket; // 8B [8, 16): next entry in hash bucket chain (overlays ChipTensor::buffer.size) + PTO2TaskId access_task_id; // 8B [16,24): task owning this access + uint64_t start_offset; // 8B [24,32): mirrors ChipTensor::start_offset (element offset) + int32_t version; // 4B [32,36): mirrors ChipTensor::version + uint32_t ndims; // 4B [36,40): mirrors ChipTensor::ndims + DataType dtype; // 1B [40,41): mirrors ChipTensor::dtype + bool manual_dep; // 1B [41,42): mirrors ChipTensor::manual_dep + bool is_contiguous; // 1B [42,43): mirrors ChipTensor::is_contiguous + TensorAccessKind access_kind; // 1B [43,44): reader or writer index uint32_t shapes[MAX_TENSOR_DIMS]; // 20B [44,64): mirrors ChipTensor::shapes // === Cache line 2 (64B) — chain manipulation + non-contiguous overlap data === @@ -151,7 +164,7 @@ struct alignas(64) PTO2TensorMapEntry { * Copy overlap-relevant fields from a ChipTensor into this entry. * * 64B memcpy of ChipTensor cache line 1 populates buffer_addr (byte [0,8)), - * producer_task_id, start_offset, version, ndims, dtype, manual_dep, + * access_task_id, start_offset, version, ndims, dtype, manual_dep, * is_contiguous and shapes[]. Byte [8,16) holds ChipTensor::buffer.size in * the source and gets written into next_in_bucket; that's harmless * because link_entry() overwrites next_in_bucket immediately after. @@ -336,7 +349,7 @@ struct alignas(64) PTO2TensorMapEntry { static_assert(sizeof(PTO2TensorMapEntry) == 128, "TensorMapEntry must be exactly 2 cache lines (128 bytes)"); static_assert(offsetof(PTO2TensorMapEntry, buffer_addr) == offsetof(ChipTensor, buffer.addr)); -static_assert(offsetof(PTO2TensorMapEntry, producer_task_id) == offsetof(ChipTensor, owner_task_id)); +static_assert(offsetof(PTO2TensorMapEntry, access_task_id) == offsetof(ChipTensor, owner_task_id)); static_assert(offsetof(PTO2TensorMapEntry, start_offset) == offsetof(ChipTensor, start_offset)); static_assert(offsetof(PTO2TensorMapEntry, version) == offsetof(ChipTensor, version)); static_assert(offsetof(PTO2TensorMapEntry, ndims) == offsetof(ChipTensor, ndims)); @@ -360,7 +373,7 @@ static_assert( struct PTO2TensorMap { // Hash table buckets (fixed size, power of 2) PTO2TensorMapEntry **buckets; // Array of offsets into entry_pool (-1 = empty) - int32_t num_buckets; // Must be power of 2 for fast modulo + int32_t num_buckets; // Must be power of 2 for fast modulo // Entry pool as ring buffer PTO2TensorMapEntry *entry_pool; // Ring buffer of entries @@ -380,6 +393,13 @@ struct PTO2TensorMap { // Cleanup progress (for periodic cleanup_retired) int32_t last_cleanup{}; + // Tracked-reader side index. Appended to preserve the original writer + // fields' offsets and cache locality. + PTO2TensorMapEntry **reader_buckets; + int32_t num_reader_buckets; + int32_t reader_used{0}; + int32_t reader_high_water{0}; + uint32_t get_task_local_id_slot(uint32_t task_local_id) const { return task_local_id & (task_window_size - 1); } // Accessors read by scope_stats_collector. Declared unconditionally so the @@ -387,6 +407,8 @@ struct PTO2TensorMap { // setter symbols must export for host dlsym; the probe call sites that use // these accessors stay gated by SIMPLER_DFX). int32_t current_used() const { return next_entry_idx - free_num; } + int32_t current_readers() const { return reader_used; } + int32_t current_writers() const { return current_used() - reader_used; } int32_t pool_capacity() const { return pool_size; } int32_t free_entries() const { return pool_size - current_used(); } @@ -407,7 +429,7 @@ struct PTO2TensorMap { // new_entry allocates a slot and initializes only its linkage (bucket_index // and the four link pointers) to the clean unlinked state; insert() assigns - // the tensor attributes and producer_task_id. + // the tensor attributes and access_task_id. PTO2TensorMapEntry *new_entry() { if (free_num > 0) { PTO2TensorMapEntry *res = free_entry_list[--free_num]; @@ -436,7 +458,8 @@ struct PTO2TensorMap { if (entry.prev_in_bucket == nullptr) { // Entry is the head of its bucket chain, update bucket head // Must compute hash BEFORE clearing tensor - buckets[entry.bucket_index] = entry.next_in_bucket; + PTO2TensorMapEntry **index = entry.access_kind == TensorAccessKind::READER ? reader_buckets : buckets; + index[entry.bucket_index] = entry.next_in_bucket; } else { entry.prev_in_bucket->next_in_bucket = entry.next_in_bucket; } @@ -446,6 +469,10 @@ struct PTO2TensorMap { entry.next_in_bucket->prev_in_bucket = entry.prev_in_bucket; } + if (entry.access_kind == TensorAccessKind::READER) { + always_assert(reader_used > 0); + reader_used--; + } free_entry_list[free_num++] = &entry; entry.bucket_index = -1; entry.next_in_bucket = nullptr; @@ -517,10 +544,13 @@ struct PTO2TensorMap { * @param tensor ChipTensor to look up * @param on_match Callback invoked for each overlapping entry */ - template - void lookup(const ChipTensor &tensor, Fn &&on_match) { - uint32_t bucket_index = hash(tensor.buffer.addr); - PTO2TensorMapEntry *cur_entry = buckets[bucket_index]; + template + void lookup_impl(const ChipTensor &tensor, Fn &&on_match) { + constexpr bool is_reader = access_kind == TensorAccessKind::READER; + int32_t bucket_count = is_reader ? num_reader_buckets : num_buckets; + uint32_t bucket_index = hash(tensor.buffer.addr, bucket_count); + PTO2TensorMapEntry **index = is_reader ? reader_buckets : buckets; + PTO2TensorMapEntry *cur_entry = index[bucket_index]; #if SIMPLER_TENSORMAP_PROFILING g_lookup_count++; @@ -572,19 +602,39 @@ struct PTO2TensorMap { #endif } + template + void lookup(const ChipTensor &tensor, Fn &&on_match) { + lookup_impl(tensor, static_cast(on_match)); + } + + template + void lookup(const ChipTensor &tensor, TensorAccessKind access_kind, Fn &&on_match) { + if (access_kind == TensorAccessKind::READER) { + lookup_impl(tensor, static_cast(on_match)); + } else { + lookup_impl(tensor, static_cast(on_match)); + } + } + /** - * Insert a new entry (called when task produces output) + * Insert a new reader or writer access entry. * * Allocates from ring buffer pool, may overwrite stale entries. * Inserts at head of hash bucket chain (maintains task_id ordering). * - * @param tensor ChipTensor produced - * @param producer_task_id Task ID of producer + * @param tensor ChipTensor accessed + * @param access_task_id Task ID owning the access */ - void insert(const ChipTensor &tensor, PTO2TaskId producer_task_id) { - PTO2TensorMapEntry *entry = new_entry(); - entry->copy_from_tensor(tensor); - link_entry(entry, tensor.buffer.addr, producer_task_id); + void insert(const ChipTensor &tensor, PTO2TaskId access_task_id, TensorAccessKind access_kind) { + if (access_kind == TensorAccessKind::READER) { + insert_impl(tensor, access_task_id); + } else { + insert_impl(tensor, access_task_id); + } + } + + void insert(const ChipTensor &tensor, PTO2TaskId access_task_id) { + insert_impl(tensor, access_task_id); } /** @@ -608,7 +658,7 @@ struct PTO2TensorMap { PTO2TensorMapEntry *cur_entry = task_entry_heads[task_slot]; while (cur_entry != nullptr) { PTO2TensorMapEntry *next_entry = cur_entry->next_in_task; // free_entry clears it - if (cur_entry->producer_task_id == retired_task) { + if (cur_entry->access_task_id == retired_task) { if (cur_entry->prev_in_task != nullptr) { cur_entry->prev_in_task->next_in_task = next_entry; } else { @@ -636,31 +686,42 @@ struct PTO2TensorMap { * addresses (low bits all-zero) still distribute evenly. We extract * the top log2(num_buckets) bits which carry the most entropy. */ - uint32_t hash(uint64_t key) { + uint32_t hash(uint64_t key, int32_t bucket_count) { key *= 0x9E3779B97F4A7C15ULL; - return static_cast(key >> (64 - __builtin_ctz(num_buckets))); + return static_cast(key >> (64 - __builtin_ctz(bucket_count))); } + uint32_t hash(uint64_t key) { return hash(key, num_buckets); } + /** * Link an initialized entry into bucket and task chains. */ - void link_entry(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId producer_task_id) { + template + void link_entry_impl(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id) { #if SIMPLER_TENSORMAP_PROFILING g_insert_count++; #endif - uint32_t bucket_index = hash(addr); - auto local_id = producer_task_id.local(); + constexpr bool is_reader = access_kind == TensorAccessKind::READER; + int32_t bucket_count = is_reader ? num_reader_buckets : num_buckets; + uint32_t bucket_index = hash(addr, bucket_count); + auto local_id = access_task_id.local(); int32_t task_slot = local_id & (task_window_size - 1); - entry->producer_task_id = producer_task_id; + entry->access_task_id = access_task_id; + entry->access_kind = access_kind; + if constexpr (is_reader) { + reader_used++; + if (reader_used > reader_high_water) reader_high_water = reader_used; + } // Insert at head of hash bucket + PTO2TensorMapEntry **index = is_reader ? reader_buckets : buckets; entry->bucket_index = bucket_index; - entry->next_in_bucket = buckets[bucket_index]; + entry->next_in_bucket = index[bucket_index]; if (entry->next_in_bucket != nullptr) { entry->next_in_bucket->prev_in_bucket = entry; } - buckets[bucket_index] = entry; + index[bucket_index] = entry; entry->prev_in_bucket = nullptr; // Link to task's entry list @@ -672,11 +733,32 @@ struct PTO2TensorMap { task_entry_heads[task_slot] = entry; } + template + void insert_impl(const ChipTensor &tensor, PTO2TaskId access_task_id) { + PTO2TensorMapEntry *entry = new_entry(); + entry->copy_from_tensor(tensor); + link_entry_impl(entry, tensor.buffer.addr, access_task_id); + } + + void link_entry(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id) { + link_entry_impl(entry, addr, access_task_id); + } + + void link_entry( + PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id, TensorAccessKind access_kind + ) { + if (access_kind == TensorAccessKind::READER) { + link_entry_impl(entry, addr, access_task_id); + } else { + link_entry_impl(entry, addr, access_task_id); + } + } + /** * Check if entry is valid (producer has not retired) */ bool entry_valid(const PTO2TensorMapEntry &entry) const { - return static_cast(entry.producer_task_id.local()) >= last_task_alive_cached; + return static_cast(entry.access_task_id.local()) >= last_task_alive_cached; } void remove_entry(PTO2TensorMapEntry &entry) { @@ -693,7 +775,7 @@ struct PTO2TensorMap { // Update predecessor's next pointer (O(1) via prev_in_task) if (entry.prev_in_task == nullptr) { // Entry is the head of its task chain, update task_entry_heads - int32_t local_id = static_cast(entry.producer_task_id.local()); + int32_t local_id = static_cast(entry.access_task_id.local()); int32_t task_slot = local_id & (task_window_size - 1); task_entry_heads[task_slot] = entry.next_in_task; } else { diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_types.h index 492c311a77..e010d57328 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_types.h @@ -22,8 +22,7 @@ * without type conflicts (Handshake, TensorPair, HostApi). */ -#ifndef SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_ -#define SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_ +#pragma once #include #include @@ -349,6 +348,15 @@ struct Arg : TaskArgsTpl { ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::INPUT, tensor_count_++), ...); } + template + void add_tracked_input(Args &&...args) { + assert_add_tensor_args(); + if (!check_add_tensor_capacity(static_cast(sizeof...(Args)))) { + return; + } + ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::TRACKED_INPUT, tensor_count_++), ...); + } + /// Batch add outputs — all ChipTensor or all TensorCreateInfo: /// add_output(ci1, ci2) — runtime allocates buffers (OUTPUT) /// add_output(t1, t2) — write-only existing tensors (OUTPUT_EXISTING) @@ -373,7 +381,7 @@ struct Arg : TaskArgsTpl { ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::INOUT, tensor_count_++), ...); } - /// No-dependency existing tensor: skips OverlapMap lookup, depends on creator only. + /// Existing tensor with creator retention but no automatic TensorMap lookup or publication. template void add_no_dep(Args &&...args) { assert_add_tensor_args(); @@ -708,5 +716,3 @@ struct ChipTaskArgs : Arg { } } }; - -#endif // SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_ diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp index 8c0a43fcd8..b72f4ed5ce 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp @@ -9,7 +9,7 @@ * ----------------------------------------------------------------------------------------------------------- */ /** - * PTO Runtime2 - Scheduler Implementation + * host_build_graph - Scheduler Implementation * * Implements scheduler state management, ready queues, and task lifecycle. * @@ -74,6 +74,7 @@ PTO2SchedProfilingData scheduler_get_profiling(int thread_idx) { void PTO2SchedulerState::print_stats() { PTO2SchedulerState *sched = this; LOG_DEBUG("=== Scheduler Statistics ==="); + LOG_DEBUG("fanin spill: %d / %d", sched->fanin_spill_top, sched->fanin_spill_capacity); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { if (sched->ring_sched_state.last_task_alive > 0) { LOG_DEBUG("Ring %d:", r); diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h index cd294f811a..ecaac7e366 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h @@ -10,7 +10,7 @@ */ /** - * PTO Runtime2 - Scheduler Interface + * host_build_graph - Scheduler Interface * * The Scheduler is responsible for: * 1. Maintaining per-resource-shape ready queues @@ -416,6 +416,7 @@ struct CompletionStats { * capacities used at layout time (init_from_layout reuses them). */ struct PTO2SchedulerLayout { + size_t off_fanin_spill_ids; size_t off_ready_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_ready_sync_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_dummy_ready_queue_slots; @@ -424,6 +425,7 @@ struct PTO2SchedulerLayout { size_t off_early_dispatch_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_early_sync_start_queue_slots; uint64_t ready_queue_capacity; + int32_t fanin_spill_capacity; }; /** @@ -436,6 +438,9 @@ struct PTO2SchedulerLayout { struct PTO2SchedulerState { // Shared memory access PTO2SharedMemoryHeader *sm_header; + int32_t *fanin_spill_ids; + int32_t fanin_spill_capacity; + int32_t fanin_spill_top; // Per-ring state struct alignas(64) RingSchedState { @@ -539,11 +544,16 @@ struct PTO2SchedulerState { // set its completion_flags byte. Single-ring: all producers are ring 0, so // there is no per-edge ring indirection. + int32_t fanin_local_id(const PTO2TaskPayload &p, int32_t index) const { + if (index < PTO2_MAX_FANIN) return p.fanin_local_ids[index]; + return fanin_spill_ids[p.fanin_spill_start + index - PTO2_MAX_FANIN]; + } + bool fanin_satisfied(const PTO2TaskSlotState *s) const { const PTO2TaskPayload &p = *s->payload; const PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; for (int32_t i = 0; i < p.fanin_count; i++) { - if (!ring.is_completion_flag_set(p.fanin_local_ids[i])) return false; + if (!ring.is_completion_flag_set(fanin_local_id(p, i))) return false; } return true; } @@ -556,7 +566,7 @@ struct PTO2SchedulerState { const PTO2TaskPayload &p = *s->payload; const PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; for (int32_t i = 0; i < p.fanin_count; i++) { - if (!ring.is_completion_flag_set(p.fanin_local_ids[i])) return i; + if (!ring.is_completion_flag_set(fanin_local_id(p, i))) return i; } return -1; } @@ -582,7 +592,7 @@ struct PTO2SchedulerState { push_ready_routed(consumer); return; } - producer = &ring.get_slot_state_by_task_id(consumer->payload->fanin_local_ids[state]); + producer = &ring.get_slot_state_by_task_id(fanin_local_id(*consumer->payload, state)); } } @@ -611,7 +621,7 @@ struct PTO2SchedulerState { if (state < 0) { push_ready_routed(waiter); } else { - register_wake(&ring.get_slot_state_by_task_id(waiter->payload->fanin_local_ids[state]), waiter); + register_wake(&ring.get_slot_state_by_task_id(fanin_local_id(*waiter->payload, state)), waiter); } waiter = next; } diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp index 436ab4838c..0b802089f8 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp @@ -246,7 +246,7 @@ void SchedulerContext::log_stall_diagnostics( int32_t rc = 0; if (slot_state.payload != nullptr) { for (int32_t k = 0; k < fi; k++) { - int32_t pid = slot_state.payload->fanin_local_ids[k]; + int32_t pid = sched_->fanin_local_id(*slot_state.payload, k); if (ring.is_completion_flag_set(pid, std::memory_order_relaxed)) rc++; } } @@ -1156,7 +1156,7 @@ void SchedulerContext::classify_partition(int32_t thread_idx, int32_t nthreads) if (state < 0) { sched_->push_ready_routed(&slot); } else { - int32_t prod_local = slot.payload->fanin_local_ids[state]; + int32_t prod_local = sched_->fanin_local_id(*slot.payload, state); sched_->register_wake(&ring.get_slot_state_by_task_id(prod_local), &slot); } } diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp index de87c5b9e0..5b192a118c 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp @@ -12,12 +12,14 @@ #include #include +#include #include #include "common.h" // debug_assert #include "common/unified_log.h" #include "aicpu/aicpu_device_config.h" +#include "aicpu/cache_maintenance.h" #include "aicpu/device_time.h" #include "aicpu/platform_regs.h" #include "callable.h" @@ -952,11 +954,22 @@ int32_t SchedulerContext::run_resolution_thread(Runtime *runtime, int32_t thread int dummy_got; while ((dummy_got = sched_->dummy_ready_queue.pop_batch(dummy_batch, DUMMY_DRAIN_BATCH)) > 0) { for (int di = 0; di < dummy_got; di++) { + PTO2TaskSlotState *ready = dummy_batch[di]; + if (ready->task_kind == TaskKind::HOST_WRITE) { + PTO2TaskPayload &payload = *ready->payload; + const uint64_t bytes = payload.scalars[2]; + if (payload.scalar_count != 3 || bytes == 0 || bytes > sizeof(uint64_t)) { + fail_scheduler(runtime, thread_idx, PTO2_ERROR_INVALID_ARGS); + break; + } + void *address = reinterpret_cast(static_cast(payload.scalars[0])); + std::memcpy(address, &payload.scalars[1], static_cast(bytes)); + cache_flush_range(address, static_cast(bytes)); + } #if SIMPLER_SCHED_PROFILING - PTO2SchedulerState::TaskCompletionOutcome outcome = - sched_->complete_task(*dummy_batch[di], thread_idx); + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(*ready, thread_idx); #else - PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(*dummy_batch[di]); + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(*ready); #endif if (outcome.error_code != PTO2_ERROR_NONE) { fail_scheduler(runtime, thread_idx, outcome.error_code); diff --git a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp index 0cde969e90..69cacd2854 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp @@ -104,6 +104,9 @@ void PTO2SchedulerState::RingSchedState::destroy() { ring = nullptr; } PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena) { PTO2SchedulerLayout layout{}; layout.ready_queue_capacity = PTO2_READY_QUEUE_SIZE; + layout.fanin_spill_capacity = HBG_FANIN_SPILL_POOL_SIZE; + layout.off_fanin_spill_ids = + arena.reserve(static_cast(layout.fanin_spill_capacity) * sizeof(int32_t), alignof(int32_t)); for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { layout.off_ready_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); @@ -128,6 +131,8 @@ bool PTO2SchedulerState::init_data_from_layout( ) { PTO2SchedulerState *sched = this; sched->sm_header = reinterpret_cast(sm_dev_base); + sched->fanin_spill_capacity = layout.fanin_spill_capacity; + sched->fanin_spill_top = 0; #if SIMPLER_SCHED_PROFILING sched->tasks_completed.store(0, std::memory_order_relaxed); sched->tasks_consumed.store(0, std::memory_order_relaxed); @@ -185,6 +190,7 @@ bool PTO2SchedulerState::init_data_from_layout( void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, DeviceArena &arena) { PTO2SchedulerState *sched = this; + sched->fanin_spill_ids = static_cast(arena.region_ptr(layout.off_fanin_spill_ids)); for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { ready_queue_wire_arena_pointers(&sched->ready_queues[i], arena, layout.off_ready_queue_slots[i]); } @@ -204,6 +210,9 @@ void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, void PTO2SchedulerState::destroy() { PTO2SchedulerState *sched = this; + sched->fanin_spill_ids = nullptr; + sched->fanin_spill_capacity = 0; + sched->fanin_spill_top = 0; sched->ring_sched_state.destroy(); for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { ready_queue_destroy(&sched->ready_queues[i]); diff --git a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp index b9cd461780..46d1655975 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp @@ -9,7 +9,7 @@ * ----------------------------------------------------------------------------------------------------------- */ /** - * PTO Runtime2 - TensorMap Implementation + * host_build_graph - TensorMap Implementation * * Implements TensorMap with ring buffer pool, lazy invalidation, * and chain truncation optimization. @@ -55,12 +55,17 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout( PTO2TensorMapLayout layout{}; layout.num_buckets = new_num_buckets; + layout.num_reader_buckets = + new_num_buckets < PTO2_TENSORMAP_READER_NUM_BUCKETS ? new_num_buckets : PTO2_TENSORMAP_READER_NUM_BUCKETS; layout.pool_size = new_pool_size; layout.task_window_size = new_task_window_size; layout.off_buckets = arena.reserve( static_cast(new_num_buckets) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) ); + layout.off_reader_buckets = arena.reserve( + static_cast(layout.num_reader_buckets) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) + ); layout.off_entry_pool = arena.reserve(static_cast(new_pool_size) * sizeof(PTO2TensorMapEntry), alignof(PTO2TensorMapEntry)); layout.off_free_entry_list = @@ -86,16 +91,21 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout_default(DeviceArena &arena, in */ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, DeviceArena &arena) { num_buckets = layout.num_buckets; + num_reader_buckets = layout.num_reader_buckets; pool_size = layout.pool_size; // Address arena regions for data writes; do not store these in struct // fields (wire_arena_pointers does that). auto *buckets_arena = static_cast(arena.region_ptr(layout.off_buckets)); + auto *reader_buckets_arena = static_cast(arena.region_ptr(layout.off_reader_buckets)); // buckets[]: empty == nullptr. for (int32_t i = 0; i < num_buckets; i++) { buckets_arena[i] = nullptr; } + for (int32_t i = 0; i < num_reader_buckets; i++) { + reader_buckets_arena[i] = nullptr; + } // Init-on-write: the entry pool is not pre-zeroed. new_entry() puts each // bump-allocated slot into the clean "unlinked" state (bucket_index == -1, @@ -108,6 +118,8 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev next_entry_idx = 0; free_num = 0; + reader_used = 0; + reader_high_water = 0; auto *heads_arena = static_cast(arena.region_ptr(layout.off_task_entry_heads)); for (int32_t i = 0; i < layout.task_window_size; i++) { @@ -122,6 +134,7 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev void PTO2TensorMap::wire_arena_pointers(const PTO2TensorMapLayout &layout, DeviceArena &arena) { buckets = static_cast(arena.region_ptr(layout.off_buckets)); + reader_buckets = static_cast(arena.region_ptr(layout.off_reader_buckets)); entry_pool = static_cast(arena.region_ptr(layout.off_entry_pool)); free_entry_list = static_cast(arena.region_ptr(layout.off_free_entry_list)); task_entry_heads = static_cast(arena.region_ptr(layout.off_task_entry_heads)); @@ -132,6 +145,7 @@ void PTO2TensorMap::destroy() { // stray post-destroy access trips a nullptr dereference instead of reading // a recycled allocation. buckets = nullptr; + reader_buckets = nullptr; entry_pool = nullptr; free_entry_list = nullptr; task_entry_heads = nullptr; @@ -144,10 +158,30 @@ void PTO2TensorMap::destroy() { void PTO2TensorMap::print_stats() { int32_t valid = 0; int32_t stale = 0; - int32_t empty_buckets = 0; - int32_t max_chain = 0; - int64_t total_chain = 0; - int32_t non_empty_buckets = 0; + + struct BucketStats { + int32_t empty{0}; + int32_t non_empty{0}; + int32_t max_chain{0}; + int64_t total_chain{0}; + }; + auto collect_bucket_stats = [&](PTO2TensorMapEntry **index, int32_t bucket_count) { + BucketStats stats; + for (int32_t b = 0; b < bucket_count; b++) { + int32_t chain_len = 0; + for (auto *entry = index[b]; entry != nullptr; entry = entry->next_in_bucket) { + chain_len++; + } + if (chain_len == 0) { + stats.empty++; + } else { + stats.non_empty++; + stats.total_chain += chain_len; + stats.max_chain = std::max(stats.max_chain, chain_len); + } + } + return stats; + }; // Count entries // Init-on-write: only [0, next_entry_idx) slots have ever been allocated and @@ -162,37 +196,26 @@ void PTO2TensorMap::print_stats() { } } - // Count bucket stats - for (int32_t b = 0; b < num_buckets; b++) { - int32_t chain_len = 0; - auto cur_entry = buckets[b]; - - while (cur_entry != nullptr) { - chain_len++; - cur_entry = cur_entry->next_in_bucket; - } - - if (chain_len == 0) { - empty_buckets++; - } else { - non_empty_buckets++; - total_chain += chain_len; - if (chain_len > max_chain) { - max_chain = chain_len; - } - } - } + const BucketStats writer_stats = collect_bucket_stats(buckets, num_buckets); + const BucketStats reader_stats = collect_bucket_stats(reader_buckets, num_reader_buckets); LOG_DEBUG("=== TensorMap Statistics ==="); LOG_DEBUG("Pool size: %d", pool_size); LOG_DEBUG("Pool next entry idx: %d", next_entry_idx); LOG_DEBUG("Pool free_num: %d", free_num); - LOG_DEBUG("Num buckets: %d", num_buckets); + LOG_DEBUG("Reader entries: live=%d high_water=%d", reader_used, reader_high_water); + LOG_DEBUG("Writer entries: live=%d", current_writers()); + LOG_DEBUG("Buckets: writer=%d reader=%d", num_buckets, num_reader_buckets); LOG_DEBUG("Valid entries: %d", valid); LOG_DEBUG("Stale entries: %d", stale); - LOG_DEBUG("Empty buckets: %d", empty_buckets); - LOG_DEBUG("Max chain len: %d", max_chain); - LOG_DEBUG("Avg chain len: %.2f", non_empty_buckets > 0 ? (float)total_chain / non_empty_buckets : 0); + LOG_DEBUG( + "Writer buckets: empty=%d max_chain=%d avg_chain=%.2f", writer_stats.empty, writer_stats.max_chain, + writer_stats.non_empty > 0 ? (float)writer_stats.total_chain / writer_stats.non_empty : 0 + ); + LOG_DEBUG( + "Reader buckets: empty=%d max_chain=%d avg_chain=%.2f", reader_stats.empty, reader_stats.max_chain, + reader_stats.non_empty > 0 ? (float)reader_stats.total_chain / reader_stats.non_empty : 0 + ); LOG_DEBUG("Last task alive: %d", last_task_alive_cached); LOG_DEBUG("============================"); } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp index c983053cd6..71d23a49f4 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp @@ -783,6 +783,14 @@ int32_t AicpuExecutor::run(Runtime *runtime) { static_cast(tp.overlap_checks), static_cast(tp.overlap_hits), tp.overlap_checks > 0 ? tp.overlap_hits * 100.0 / tp.overlap_checks : 0.0 ); + rt->orchestrator.tensor_map.print_stats(); + for (int32_t ring_id = 0; ring_id < PTO2_MAX_RING_DEPTH; ++ring_id) { + const PTO2FaninPool &pool = rt->orchestrator.rings[ring_id].fanin_pool; + LOG_INFO( + "Thread %d: Fanin spill ring=%d used=%d high_water=%d capacity=%d", thread_idx, ring_id, + pool.used(), pool.high_water, pool.capacity + ); + } #endif #endif // SIMPLER_ORCH_PROFILING diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md index 3594ccf259..b3bab2730e 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md @@ -41,7 +41,7 @@ Type changes: | Field | Before | After | | ----- | ------ | ----- | | `PTO2TaskDescriptor.task_id` | `int32_t` | `PTO2TaskId` | -| `PTO2TensorMapEntry.producer_task_id` | `int32_t` | `PTO2TaskId` | +| `PTO2TensorMapEntry.access_task_id` | `int32_t` | `PTO2TaskId` | | `PTO2TaskSlotState.ring_id` | N/A | `uint8_t` (new, denormalized for fast access) | ## 4. Data Structures @@ -155,8 +155,8 @@ Entry validity checks and `cleanup_retired` operate per-ring: ```cpp bool entry_valid(const PTO2TensorMapEntry& e) { - int32_t ring = e.producer_task_id.ring(); - int32_t local = e.producer_task_id.local(); + int32_t ring = e.access_task_id.ring(); + int32_t local = e.access_task_id.local(); return local >= last_task_alives[ring]; } ``` diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index d81ee13e2d..a3892977a4 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -320,24 +320,28 @@ This design avoids the complexity of ring-based wrapping while still being bound ### 5.4 Stale Entry Cleanup: Three-Layer Defense -TensorMap must ensure entries for retired tasks (`producer_task_id < last_task_alive`) are removed, so that: +TensorMap must ensure entries for retired tasks (`access_task_id < last_task_alive`) are removed, so that: - The pool does not grow unboundedly (capacity is finite) - Lookup performance does not degrade as stale entries accumulate in bucket chains Three complementary mechanisms achieve this: -**Layer 1 — Chain Truncation during Lookup** (lazy, per-bucket): +**Layer 1 — Stale Filtering during Lookup** (lazy, per-bucket): -Since `insert` always prepends to the bucket head, entries in each bucket chain are in **descending task_id order**. When `PTO2TensorMap::lookup` encounters the first stale entry (`producer_task_id < last_task_alive`), all subsequent entries in the chain are guaranteed stale too. The entire tail is truncated in one operation using `prev_in_bucket` pointers for O(1) unlinking. +`PTO2TensorMap::lookup` checks each entry against its source ring's +`last_task_alive` and skips stale entries. Entries from different rings may be +interleaved in one bucket, so a stale entry does not imply that the remaining +tail is stale and lookup must not truncate the chain. -This guarantees lookup only traverses valid entries — O(valid_entries_in_bucket), not O(total_entries). +This guarantees stale entries never create dependencies; periodic cleanup below +keeps traversal and pool use bounded. **Layer 2 — Periodic Batch Cleanup** (`cleanup_retired`, per-task): Every time the orchestrator submits a task (Step 0 of `PTO2OrchestratorState::submit_task`), it calls `PTO2TensorMap::sync_tensormap`. When `last_task_alive` has advanced by more than `PTO2_TENSORMAP_CLEANUP_INTERVAL` (default 64) tasks since the last cleanup, `PTO2TensorMap::cleanup_retired` runs: -This uses the **per-task entry chain** (`task_entry_head[task_slot]`) — each task's entries are doubly-linked together at insert time via `next_in_task`/`prev_in_task`. A slot's chain can hold more than one task's entries: a task at `local_id + N * window` reuses the slot and prepends to the chain already there, and cleanup can lag that reuse. Cleanup therefore walks the chain and frees only the entries whose `producer_task_id` matches the retiring task, unlinking each and leaving the rest linked — O(entries_in_slot), with no scan of the entire pool or all buckets. Freed entries are returned to `free_entry_list` for immediate reuse. +This uses the **per-task entry chain** (`task_entry_head[task_slot]`) — each task's entries are doubly-linked together at insert time via `next_in_task`/`prev_in_task`. A slot's chain can hold more than one task's entries: a task at `local_id + N * window` reuses the slot and prepends to the chain already there, and cleanup can lag that reuse. Cleanup therefore walks the chain and frees only the entries whose `access_task_id` matches the retiring task, unlinking each and leaving the rest linked — O(entries_in_slot), with no scan of the entire pool or all buckets. Freed entries are returned to `free_entry_list` for immediate reuse. **Layer 3 — Back-Pressure on Pool Exhaustion** (blocking): @@ -349,7 +353,7 @@ This forms a back-pressure mechanism analogous to the Task Ring's flow control. | Layer | Trigger | Method | Guarantees | | ----- | ------- | ------ | ---------- | -| Chain Truncation | Every lookup | Truncate stale tail of bucket chain | Lookup only visits valid entries | +| Stale Filtering | Every lookup | Skip entries retired in their source ring | Stale entries never create dependencies | | Periodic Cleanup | Every 64 retired tasks | Walk per-task chains, free entries | Pool capacity reclaimed in bounded time | | Pool Back-Pressure | Pool exhausted | Block until scheduler advances watermark | Hard capacity bound, no OOM | diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md index de13df3ec0..d32d01e1d1 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md @@ -38,10 +38,11 @@ addr null-check → TensorMap lookup → spin-wait producer COMPLETED → comput ### 3.2 set_tensor_data Flow ```text -addr null-check → TensorMap lookup → spin-wait producer COMPLETED → spin-wait consumers done → memcpy write +addr null-check → writer/reader lookup → wait conflicting accesses → memcpy write ``` -One extra step versus get_tensor_data: wait for all consumers to finish (`fanout_refcount >= fanout_count - 1`, excluding the scope reference). +Compared with `get_tensor_data`, a write also drains each writer's consumers +and waits for every overlapping `TRACKED_INPUT` task to complete. ### 3.3 Timeout @@ -104,7 +105,7 @@ Three actors: | - | ---------- | -------- | ------ | --------- | ----- | | 1 | Kernel write (OUTPUT) | Orch Read | RAW | spin-wait producer COMPLETED | Yes | | 2 | Kernel write (OUTPUT) | Orch Write | WAW | spin-wait producer COMPLETED | Yes | -| 3 | Kernel read (INPUT) | Orch Write | WAR | spin-wait fanout_refcount | **Needs INOUT** | +| 3 | Kernel read (TRACKED_INPUT) | Orch Write | WAR | spin-wait reader task | Yes | | 4 | Kernel read-write (INOUT) | Orch Read | RAW | spin-wait producer COMPLETED | Yes | | 5 | Kernel read-write (INOUT) | Orch Write | WAW+WAR | spin-wait producer + consumers | Yes | | 6 | Orch Write | Kernel read (INPUT) | RAW | blocking completes before next submit | Yes | @@ -114,11 +115,12 @@ Three actors: ### Key Design Points -**Scenario #3 is the only case requiring special attention**: +**Scenario #3 requires an explicit reader annotation**: -TensorMap tracks only producers (OUTPUT/INOUT), not pure INPUT consumers. If a tensor is only registered via `add_input()`, TensorMap has no producer entry for it. `set_tensor_data`'s `wait_for_tensor_ready()` finds no matching producer (the lookup callback never fires) and returns immediately — but the kernel may still be reading → **WAR data race**. - -**Solution**: For tensors that may later be written via `set_tensor_data`, use `add_inout()` instead of `add_input()`. INOUT registers a producer entry in TensorMap, enabling `set_tensor_data` to track all consumers through `fanout_refcount`. +Plain `add_input()` performs the RAW lookup but does not publish a reader in +A2/A3 `tensormap_and_ringbuffer`. If a later `set_tensor_data()` may overlap +that read, use `add_tracked_input()`. The tracked reader remains read-only but +is visible to the host writer's WAR lookup. **Scenarios #6–8 serial guarantee**: @@ -132,6 +134,9 @@ get/set_tensor_data are blocking calls, and orchestration is single-threaded ser | -------- | -------- | | External tensor never submitted as OUTPUT/INOUT | No TensorMap entry — get/set execute immediately | | External tensor previously submitted as OUTPUT/INOUT | TensorMap has producer entry — get/set spin-wait | -| External tensor submitted as INPUT, then set_tensor_data | **WAR risk** — must use INOUT instead (same as scenario #3) | +| External tensor submitted as INPUT, then set_tensor_data | **WAR risk** — the reader was not published | +| External tensor submitted as TRACKED_INPUT, then set_tensor_data | Host write waits for the reader task | -**Key rule**: If an external tensor will later be written via `set_tensor_data`, all prior kernel accesses must use `add_inout()`, not `add_input()`. +**Key rule**: If an external tensor will later be written via +`set_tensor_data`, every overlapping pure reader must use +`add_tracked_input()`. diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp index 0a9384d44f..2e0a89e455 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp @@ -45,13 +45,13 @@ * "kept at call site" note); both passes run the same explicit-deps loop, so * the comparison covers it too. * - * STEP 4 (`register_task_outputs`) runs on BOTH tensor maps after both passes + * STEP 4 (`register_task_accesses`) runs on BOTH tensor maps after both passes * complete, keeping `tm_oracle` and `tm_annot` bit-equivalent for the next * record's INOUT+COVERED `remove_entry` mutations. * * Pool sizing: replay never advances last_task_alive, so each tensor map's - * entry pool must accommodate every output write across the whole trace. We - * scan the record buffer once to count INOUT + OUTPUT_EXISTING slots and size + * entry pool must accommodate every tracked access across the whole trace. We + * scan the record buffer once to count INPUT + INOUT + OUTPUT_EXISTING slots and size * the pool accordingly. Both maps get the same size. */ @@ -89,12 +89,12 @@ int32_t ceil_pow2(int32_t v) { return v + 1; } -// Count INOUT + OUTPUT_EXISTING slots across the record buffer — -// register_task_outputs only inserts those, and skips entries with manual_dep +// Count TRACKED_INPUT + INOUT + OUTPUT_EXISTING slots across the record buffer — +// register_task_accesses inserts those, and skips entries with manual_dep // set. Counting both without inspecting manual_dep is a conservative upper // bound (manual_dep is rare; the small over-allocation pays for itself in // avoided pool exhaustion). -int32_t count_outputs(const DepGenRecord *records, size_t n) { +int32_t count_accesses(const DepGenRecord *records, size_t n) { int32_t total = 0; for (size_t i = 0; i < n; i++) { const DepGenRecord &r = records[i]; @@ -104,7 +104,7 @@ int32_t count_outputs(const DepGenRecord *records, size_t n) { if (r.flags & DEP_GEN_FLAG_OVERFLOW) continue; for (uint16_t j = 0; j < r.tensor_count; j++) { auto t = static_cast(r.arg_types[j]); - if (t == TensorArgType::INOUT || t == TensorArgType::OUTPUT_EXISTING) { + if (t == TensorArgType::TRACKED_INPUT || t == TensorArgType::INOUT || t == TensorArgType::OUTPUT_EXISTING) { total++; } } @@ -158,6 +158,20 @@ const char *overlap_status_str(OverlapStatus s) { return "unknown"; } +const char *hazard_kind_str(TensorHazardKind kind) { + switch (kind) { + case TensorHazardKind::RAW: + return "RAW"; + case TensorHazardKind::WAW: + return "WAW"; + case TensorHazardKind::WAR: + return "WAR"; + } + return "unknown"; +} + +const char *access_kind_str(TensorAccessKind kind) { return kind == TensorAccessKind::READER ? "READER" : "WRITER"; } + // One annotated edge. consumer_* always populated. producer_* populated for // TENSORMAP source only — the explicit/creator emit paths don't have a // matched tensormap entry to copy from. @@ -172,7 +186,9 @@ struct EdgeAnnot { EdgeSource source; DepFlags flags; // per-edge WAIT/RETAIN semantics carried into deps.json OverlapStatus overlap; // only meaningful for TENSORMAP - uint64_t tensor_id; // 0 for EXPLICIT + TensorHazardKind hazard; + TensorAccessKind access_kind; + uint64_t tensor_id; // 0 for EXPLICIT // Consumer side (the ChipTensor the submitting task is reading). uint8_t consumer_dtype; uint32_t consumer_ndims; @@ -228,6 +244,8 @@ const char *arg_type_str(TensorArgType t) { switch (t) { case TensorArgType::INPUT: return "INPUT"; + case TensorArgType::TRACKED_INPUT: + return "TRACKED_INPUT"; case TensorArgType::OUTPUT: return "OUTPUT"; case TensorArgType::INOUT: @@ -394,6 +412,8 @@ bool write_deps_json( write_dep_flags(out, e.flags); if (e.source == EdgeSource::TENSORMAP) { out << ",\"overlap\":\"" << overlap_status_str(e.overlap) << '"'; + out << ",\"hazard\":\"" << hazard_kind_str(e.hazard) << '"'; + out << ",\"access_kind\":\"" << access_kind_str(e.access_kind) << '"'; } if (e.source != EdgeSource::EXPLICIT) { out << ",\"tensor_id\":\"" << e.tensor_id << '"'; @@ -444,21 +464,29 @@ void annot_pass( emit_creator(owner, i, *tensor); } - // STEP B: tensormap lookup (only INPUT/INOUT, skip manual_dep). - if (ptype != TensorArgType::INPUT && ptype != TensorArgType::INOUT) { - continue; - } if (tensor->manual_dep) { continue; } - tensor_map.lookup(*tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { - emit_tensormap(entry.producer_task_id, i, *tensor, entry, overlap_status); - if (ptype == TensorArgType::INOUT && overlap_status == OverlapStatus::COVERED) { - tensor_map.remove_entry(entry); - } - return true; - }); + auto lookup = [&](TensorAccessKind access_kind, TensorHazardKind hazard) { + tensor_map.lookup(*tensor, access_kind, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) { + emit_tensormap(entry.access_task_id, i, *tensor, entry, overlap_status, hazard); + if (ptype == TensorArgType::INOUT && access_kind == TensorAccessKind::WRITER) { + emit_tensormap(entry.access_task_id, i, *tensor, entry, overlap_status, TensorHazardKind::WAW); + } + if ((ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) && + overlap_status == OverlapStatus::COVERED) { + tensor_map.remove_entry(entry); + } + return true; + }); + }; + if (ptype == TensorArgType::INPUT || ptype == TensorArgType::TRACKED_INPUT || ptype == TensorArgType::INOUT) { + lookup(TensorAccessKind::WRITER, TensorHazardKind::RAW); + } + if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) { + lookup(TensorAccessKind::READER, TensorHazardKind::WAR); + } } } @@ -496,8 +524,8 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c task_window_sizes[r] = ceil_pow2(need < 16 ? 16 : need); } - int32_t output_count = count_outputs(records, num_records); - int32_t pool_size = output_count + (output_count / 10) + 64; + int32_t access_count = count_accesses(records, num_records); + int32_t pool_size = access_count + (access_count / 10) + 64; if (pool_size < PTO2_TENSORMAP_POOL_SIZE) { pool_size = PTO2_TENSORMAP_POOL_SIZE; } @@ -755,7 +783,7 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c }, // emit_tensormap(producer, arg_idx, consumer_tensor, entry, status) [&](PTO2TaskId producer, int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, - OverlapStatus status) { + OverlapStatus status, TensorHazardKind hazard) { // Per-(succ, arg_idx, producer_buffer_addr, producer_version) // dedup gives us "the same producer slice fired twice for the // same consumer arg" collapse — but two distinct slices from @@ -771,6 +799,8 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c e.source = EdgeSource::TENSORMAP; e.flags = DEP_WAIT; e.overlap = status; + e.hazard = hazard; + e.access_kind = entry.access_kind; e.tensor_id = make_tensor_id(entry.buffer_addr, entry.version); fill_consumer(e, consumer); fill_producer(e, entry); @@ -807,8 +837,8 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c } // ============ STEP 4 — publish outputs on BOTH maps ============ - register_task_outputs(inputs, task_id, tm_oracle, in_manual_scope); - register_task_outputs(inputs, task_id, tm_annot, in_manual_scope); + register_task_accesses(inputs, task_id, tm_oracle, in_manual_scope); + register_task_accesses(inputs, task_id, tm_annot, in_manual_scope); } tm_oracle.destroy(); diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.h b/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.h index dd81a85ba1..7bac37de07 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.h @@ -15,7 +15,7 @@ * * Takes the records the host collector drained from the device ring buffer * (``DepGenCollector::records()``) and runs them back through a host-resident - * PTO2TensorMap using the same ``compute_task_fanin`` / ``register_task_outputs`` + * PTO2TensorMap using the same ``compute_task_fanin`` / ``register_task_accesses`` * primitives the device orchestrator uses, emitting the full * predecessor → successor edge list to deps.json. * @@ -70,8 +70,7 @@ * The replay is single-threaded and pure CPU: no device handle is required. */ -#ifndef SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_HOST_DEP_GEN_REPLAY_H_ -#define SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_HOST_DEP_GEN_REPLAY_H_ +#pragma once #include #include @@ -102,5 +101,3 @@ int dep_gen_replay_emit_deps_json(const struct DepGenRecord *records, size_t num #ifdef __cplusplus } // extern "C" #endif - -#endif // SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_HOST_DEP_GEN_REPLAY_H_ diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h b/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h index 057cfb8cb0..bbf57c8f35 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h @@ -54,6 +54,7 @@ class CoreTaskArgsWithDeps : private CoreTaskArgs { using CoreTaskArgs::add_scalar; using CoreTaskArgs::add_scalars; using CoreTaskArgs::add_scalars_i32; + using CoreTaskArgs::add_tracked_input; using CoreTaskArgs::allow_early_resolve; // early-dispatch hint (getter) using CoreTaskArgs::copy_scalars_from; using CoreTaskArgs::set_allow_early_resolve; // early-dispatch hint (setter) diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_orchestration_api.h b/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_orchestration_api.h index 16386b1a1e..6d446fdda7 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_orchestration_api.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_orchestration_api.h @@ -297,20 +297,9 @@ static inline T get_tensor_data(const ChipTensor &tensor, uint32_t ndims, const * set_tensor_data(tensor, 1, idx, raw_u64); // old usage unchanged * set_tensor_data(tensor, 1, idx, 42.0f); // typed write (T = float) * - * If the tensor has a producer in TensorMap, spin-waits until the producer - * and all its consumers complete before writing (WAW + WAR safety). - * External tensors (make_tensor_external) with no TensorMap entry are - * written immediately without waiting. - * - * Limitation: TensorMap only tracks producers (OUTPUT/INOUT), not consumers - * that used the tensor as INPUT. If a kernel reads this tensor as INPUT - * (not INOUT) and the tensor has no TensorMap producer entry, set_tensor_data - * cannot detect the reader and may cause a data race. - * - * To ensure WAR safety for all access patterns, use add_inout() instead of - * add_input() for kernel parameters that may later be written via - * set_tensor_data. INOUT creates a TensorMap entry that enables automatic - * consumer tracking via fanout_refcount. + * Spin-waits for overlapping writers and explicitly tracked readers before + * writing (WAW + WAR safety). A pure reader that may precede this call must + * use add_tracked_input(); plain add_input() does not publish a reader entry. * * The tensor must already have an allocated buffer (addr != 0). * For runtime-created outputs, call this only on the ChipTensor returned by diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h index 16a92380a7..8936242f5a 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h @@ -16,11 +16,11 @@ * Two header-only template entry points: * * compute_task_fanin — STEP 3 in submit_task: per-tensor creator retention (Step A) - * + tensormap.lookup for INPUT/INOUT (Step B). Calls back into + * + tensormap.lookup for INPUT/TRACKED_INPUT/INOUT (Step B). Calls back into * user-supplied `emit` for each producer it identifies. * - * register_task_outputs — STEP 4 in submit_task: tensormap.insert for INOUT and - * OUTPUT_EXISTING tensors. No callbacks. + * register_task_accesses — STEP 4 in submit_task: register TRACKED_INPUT readers and + * INOUT/OUTPUT_EXISTING writers. No callbacks. * * STEP 1 (explicit_deps) is intentionally left at the runtime call site because its * `last_task_alive` shortcut + unchecked slot lookup is subtly different from the @@ -47,6 +47,7 @@ #pragma once #include +#include #include "pto_task_id.h" #include "pto_tensormap.h" @@ -54,7 +55,7 @@ #include "tensor.h" /** - * View struct for inputs to compute_task_fanin / register_task_outputs. + * View struct for inputs to compute_task_fanin / register_task_accesses. * * Both runtime and replay assemble one of these from their own data sources * (runtime: from Arg accessors; replay: from SubmitTraceEntry fields). All @@ -68,21 +69,83 @@ struct DepInputs { const PTO2TaskId *explicit_deps; // length = explicit_dep_count (validity checked by caller) }; +template +[[nodiscard]] inline bool +emit_writer_dependency(Emit &emit, PTO2TaskId task_id, DepFlags flags) { + if constexpr (std::is_invocable_r_v) { + return emit(task_id, flags); + } else if constexpr (std::is_invocable_r_v) { + return emit(task_id, flags, TensorAccessKind::WRITER); + } else if constexpr (std::is_invocable_r_v) { + return emit(task_id, TensorAccessKind::WRITER); + } else { + return emit(task_id); + } +} + +template +[[nodiscard]] inline bool emit_reader_dependency(Emit &emit, PTO2TaskId task_id, DepFlags flags) { + if constexpr (std::is_invocable_r_v) { + return emit(task_id, flags); + } else if constexpr (std::is_invocable_r_v) { + return emit(task_id, flags, TensorAccessKind::READER); + } else if constexpr (std::is_invocable_r_v) { + return emit(task_id, TensorAccessKind::READER); + } else { + return emit(task_id); + } +} + +template +[[nodiscard]] __attribute__((noinline, cold)) bool +compute_reader_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, ReaderEmit &reader_emit) { + for (int32_t i = 0; i < inputs.tensor_count; i++) { + TensorArgType ptype = inputs.arg_types[i]; + if (ptype != TensorArgType::INOUT && ptype != TensorArgType::OUTPUT_EXISTING) { + continue; + } + const ChipTensor *tensor = &inputs.tensors[i].ref(); + if (tensor->manual_dep) { + continue; + } + bool fatal = false; + tensor_map.lookup( + *tensor, TensorAccessKind::READER, + [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { + if (!emit_reader_dependency(reader_emit, entry.access_task_id, DEP_WAIT)) { + fatal = true; + return false; + } + if (overlap_status == OverlapStatus::COVERED) { + tensor_map.remove_reader_entry(entry); + } + return true; + } + ); + if (fatal) { + return false; + } + } + return true; +} + /** * Compute fanin for a task being submitted (STEP 3: Step A creator retention + * Step B tensormap modifier lookup). * * For each non-OUTPUT tensor: * - If owner_task_id is valid, emit(owner) - * - For INPUT/INOUT (and not manual_dep), tensor_map.lookup(*tensor) and emit + * - For INPUT/TRACKED_INPUT/INOUT (and not manual_dep), tensor_map.lookup(*tensor) and emit * each matching producer. INOUT+COVERED triggers tensor_map.remove_entry(entry). * * @return true on success (or producer-skipped-silently); false if emit signaled * fatal — caller should propagate (after any fatal bookkeeping done by emit). */ -template +template [[nodiscard]] inline bool -compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_manual_scope, Emit emit) { +compute_task_fanin( + const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_manual_scope, Emit emit, ReaderEmit reader_emit +) { if (in_manual_scope) { return true; } @@ -101,13 +164,13 @@ compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_m // the creator edge carries both ordering and lifetime. PTO2TaskId owner = tensor->owner_task_id; if (owner.is_valid()) { - if (!emit(owner, DEP_WAIT | DEP_RETAIN)) { + if (!emit_writer_dependency(emit, owner, DEP_WAIT | DEP_RETAIN)) { return false; } } - // Step B: only INPUT/INOUT need modifier dependency lookup. - if (ptype != TensorArgType::INPUT && ptype != TensorArgType::INOUT) { + if (ptype != TensorArgType::INPUT && ptype != TensorArgType::INOUT && + ptype != TensorArgType::TRACKED_INPUT) { continue; } if (tensor->manual_dep) { @@ -116,22 +179,12 @@ compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_m bool fatal = false; tensor_map.lookup(*tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { - // Ordering-only (DEP_WAIT): a modifier only rewrote a buffer someone - // else allocated, so its lifetime rides that allocator's creator edge, - // not this modifier edge. Retention-safety invariant that makes this - // sound: only TensorArgType::OUTPUT tensors are allocated into the - // packed output heap, and a runtime-created OUTPUT always carries a - // valid owner_task_id — so its consumer takes a Step-A DEP_RETAIN edge - // to the allocator above. INOUT / OUTPUT_EXISTING buffers are never - // owned by a modifier. If a future layout put a modifier-owned buffer - // into the packed heap, this edge would have to become RETAIN or the - // producer could be reclaimed under a live reader (use-after-free). - if (!emit(entry.producer_task_id, DEP_WAIT)) { + if (!emit_writer_dependency(emit, entry.access_task_id, DEP_WAIT)) { fatal = true; - return false; // stop iteration + return false; } if (ptype == TensorArgType::INOUT && overlap_status == OverlapStatus::COVERED) { - tensor_map.remove_entry(entry); + tensor_map.remove_writer_entry(entry); } return true; }); @@ -139,19 +192,22 @@ compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_m return false; } } - return true; + return tensor_map.current_readers() == 0 || compute_reader_fanin(inputs, tensor_map, reader_emit); +} + +template +[[nodiscard]] inline bool +compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_manual_scope, Emit emit) { + return compute_task_fanin(inputs, tensor_map, in_manual_scope, emit, emit); } /** - * Register a task's outputs in the tensormap (STEP 4 in submit_task). - * - * For INOUT and OUTPUT_EXISTING tensors (excluding manual_dep), inserts the - * tensor into tensor_map keyed by its buffer.addr with `task_id` as producer. + * Register a task's accesses in the tensormap (STEP 4 in submit_task). * * No-op when in_manual_scope. */ inline void -register_task_outputs(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap &tensor_map, bool in_manual_scope) { +register_task_accesses(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap &tensor_map, bool in_manual_scope) { if (in_manual_scope) { return; } @@ -162,26 +218,31 @@ register_task_outputs(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap if (!tensor->manual_dep) { tensor_map.insert(*tensor, task_id); } + } else if (ptype == TensorArgType::TRACKED_INPUT) { + const ChipTensor *tensor = &inputs.tensors[i].ref(); + if (!tensor->manual_dep) { + tensor_map.insert(*tensor, task_id, TensorAccessKind::READER); + } } } } /** - * Count the tensormap entries register_task_outputs() will insert for this task. + * Count the tensormap entries register_task_accesses() will insert for this task. * - * Mirrors register_task_outputs()'s selection exactly (INOUT / OUTPUT_EXISTING, - * excluding manual_dep), so the returned value is the precise number of + * Mirrors register_task_accesses() exactly, so the returned value is the precise number of * new_entry() calls that step makes. The orchestrator uses it to reserve pool * capacity before inserting. Returns 0 in a manual scope (no registration). */ -inline int32_t count_registrable_outputs(const DepInputs &inputs, bool in_manual_scope) { +inline int32_t count_registrable_accesses(const DepInputs &inputs, bool in_manual_scope) { if (in_manual_scope) { return 0; } int32_t needed = 0; for (int32_t i = 0; i < inputs.tensor_count; i++) { TensorArgType ptype = inputs.arg_types[i]; - if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) { + if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING || + ptype == TensorArgType::TRACKED_INPUT) { if (!inputs.tensors[i].ref().manual_dep) { needed++; } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp index b543298766..7f7fd949a7 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp @@ -10,7 +10,7 @@ */ /** - * PTO Runtime2 - Orchestrator Implementation + * tensormap_and_ringbuffer - Orchestrator Implementation * * Implements orchestrator state management, scope handling, and task submission. * @@ -826,7 +826,7 @@ static bool ensure_tensormap_capacity(PTO2OrchestratorState *orch, int32_t neede return true; } // Progress is entries actually freed, NOT watermark movement: a ring can - // retire zero-output tasks (count_registrable_outputs == 0), advancing + // retire zero-access tasks (count_registrable_accesses == 0), advancing // last_task_alive without freeing any entry. Gating the backstop on // free_entries() keeps a wedged pool from dodging the timeout while some // unrelated ring keeps draining. @@ -908,7 +908,7 @@ static TaskOutputTensors submit_task_common( if (is_dep_gen_enabled()) { const void *tensor_ptrs[MAX_TENSOR_ARGS]; // TensorArgType is `enum class : int32_t` (4 bytes); the on-disk record - // packs arg_types as uint8_t[16] (5-value enum fits in a byte). Narrow + // packs arg_types as uint8_t[16] (the enum fits in a byte). Narrow // each tag here rather than letting the AICPU writer reinterpret a // 4×-wider array as bytes — that path silently lost two of every three // tags on little-endian and synthesized phantom self-edges in replay. @@ -995,7 +995,23 @@ static TaskOutputTensors submit_task_common( ); }; - if (!compute_task_fanin(dep_inputs, orch->tensor_map, orch->in_manual_scope(), runtime_emit)) { + auto runtime_emit_reader = [&](PTO2TaskId producer_task_id, DepFlags kind) -> bool { + uint8_t prod_ring = producer_task_id.ring(); + PTO2SharedMemoryRingHeader &producer_ring = orch->sm_header->rings[prod_ring]; + int32_t prod_slot = producer_ring.get_slot_by_task_id(static_cast(producer_task_id.local())); + PTO2TaskSlotState *prod_state = &producer_ring.get_slot_state_by_slot(prod_slot); + if (prod_state->task != nullptr && prod_state->task->task_id == producer_task_id && + prod_state->task_state.load(std::memory_order_acquire) >= PTO2_TASK_COMPLETED) { + return true; + } + return append_fanin_or_fail( + orch, prod_ring, prod_slot, prod_state, producer_task_id, &fanin_builder, ring_id, kind + ); + }; + + if (!compute_task_fanin( + dep_inputs, orch->tensor_map, orch->in_manual_scope(), runtime_emit, runtime_emit_reader + )) { return result; } @@ -1006,11 +1022,11 @@ static TaskOutputTensors submit_task_common( // is shared across rings and reclaimed as last_task_alive advances; an // exhausted pool back-pressures here (and detects a wedged watermark) rather // than tripping new_entry()'s hard assert mid-registration. - int32_t tensormap_needed = count_registrable_outputs(dep_inputs, orch->in_manual_scope()); + int32_t tensormap_needed = count_registrable_accesses(dep_inputs, orch->in_manual_scope()); if (tensormap_needed > 0 && !ensure_tensormap_capacity(orch, tensormap_needed)) { return result; } - register_task_outputs(dep_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); + register_task_accesses(dep_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); CYCLE_COUNT_LAP(g_orch_insert_cycle); diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp index f26c5f9993..62cd1e4b2e 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp @@ -10,7 +10,7 @@ */ /** - * PTO Runtime2 - Main Implementation + * tensormap_and_ringbuffer - Main Implementation * * Implements the unified runtime API that combines orchestrator and scheduler. * @@ -87,12 +87,9 @@ void rt_report_fatal(PTO2Runtime *rt, int32_t error_code, const char *func, cons va_end(args); } -// Wait for all producers of this tensor to be safe for data access. -// Checks owner metadata (lifecycle anchor) and OverlapMap (modifier writers). -// For reads: wait until each producer COMPLETED (done writing). -// For writes: also wait until all consumers done reading -// (consumer low bits of fanout_refcount >= consumer count, excluding the -// bit31 scope reference). +// Wait for all accesses that conflict with orchestration-side data access. +// Reads wait for owner and writer entries to complete. Writes additionally +// drain each writer's consumers and wait for tracked reader entries. // Uses cycle-based timeout (checked every 1024 spins). // Returns false on timeout (sets orch.fatal). MAYBE_UNINITIALIZED_BEGIN @@ -109,6 +106,7 @@ wait_for_tensor_ready(PTO2Runtime *rt, const ChipTensor &tensor, bool wait_for_c // the second encounter. constexpr int kSegmentCap = 64; const PTO2TaskSlotState *seg[kSegmentCap]; + bool seg_wait_consumers[kSegmentCap]{}; int seg_count = 0; bool failed = false; @@ -169,39 +167,53 @@ wait_for_tensor_ready(PTO2Runtime *rt, const ChipTensor &tensor, bool wait_for_c for (int i = 0; i < seg_count; i++) { wait_one_producer(*seg[i]); if (failed) return; - if (!wait_for_consumers) continue; + if (!seg_wait_consumers[i]) continue; wait_one_consumers(*seg[i]); if (failed) return; } seg_count = 0; }; - auto try_push = [&](const PTO2TaskSlotState &s) { + auto try_push = [&](const PTO2TaskSlotState &s, bool include_consumers) { for (int j = 0; j < seg_count; j++) { - if (seg[j] == &s) return; // per-segment dedup + if (seg[j] == &s) { + seg_wait_consumers[j] = seg_wait_consumers[j] || include_consumers; + return; + } } if (seg_count == kSegmentCap) { flush_segment(); if (failed) return; } - seg[seg_count++] = &s; + seg[seg_count] = &s; + seg_wait_consumers[seg_count++] = include_consumers; }; auto do_wait = [&]() { // Step A: creator retention — read owner directly from tensor metadata if (owner.is_valid()) { auto &s = orch.sm_header->rings[owner.ring()].get_slot_state_by_task_id(owner.local()); - try_push(s); + try_push(s, wait_for_consumers); if (failed) return; } // Step B: modifier writer lookup (OverlapMap), direct callback orch.tensor_map.lookup(tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus) -> bool { - PTO2TaskId pid = entry.producer_task_id; + PTO2TaskId pid = entry.access_task_id; auto &s = orch.sm_header->rings[pid.ring()].get_slot_state_by_task_id(pid.local()); - try_push(s); + try_push(s, wait_for_consumers); return !failed; }); + if (wait_for_consumers && !failed) { + orch.tensor_map.lookup( + tensor, TensorAccessKind::READER, [&](PTO2TensorMapEntry &entry, OverlapStatus) -> bool { + PTO2TaskId pid = entry.access_task_id; + auto &s = orch.sm_header->rings[pid.ring()].get_slot_state_by_task_id(pid.local()); + try_push(s, false); + return !failed; + } + ); + } if (failed) return; flush_segment(); }; @@ -243,7 +255,7 @@ void set_tensor_data( return; } - // Wait for producer + all consumers before writing (WAW + WAR safety) + // Wait for conflicting writers, their consumers, and tracked readers. if (!wait_for_tensor_ready(rt, tensor, true, __FUNCTION__)) { return; } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h index af7d4742fa..b13cc13c9c 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h @@ -73,6 +73,7 @@ #define PTO2_DEP_LIST_POOL_SIZE 16384 // Per-ring dependency list pool entries #define PTO2_TENSORMAP_POOL_SIZE (65536) // TensorMap entry pool #define PTO2_TENSORMAP_NUM_BUCKETS 4096 // Power of 2 for fast hash (4096×8B=32KB fits L1) +#define PTO2_TENSORMAP_READER_NUM_BUCKETS 512 // Sparse tracked-reader sidecar (512×12B=6KB with epochs) // Scope management #define PTO2_MAX_SCOPE_DEPTH 64 // Maximum nesting depth diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h index 2b708d318d..d3a7d880f3 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h @@ -10,10 +10,10 @@ */ /** - * PTO Runtime2 - TensorMap Interface + * tensormap_and_ringbuffer - TensorMap Interface * - * TensorMap provides producer lookup for dependency discovery: - * - Maps ChipTensor -> producer task ID + * TensorMap provides access lookup for dependency discovery: + * - Maps ChipTensor -> reader/writer task IDs * - Used by pto_submit_task() to find dependencies * * Key design features: @@ -57,6 +57,17 @@ enum class OverlapStatus { OTHER, }; +enum class TensorAccessKind : uint8_t { + READER = 0, + WRITER = 1, +}; + +enum class TensorHazardKind : uint8_t { + RAW = 0, + WAW = 1, + WAR = 2, +}; + struct Segment { uint64_t begin; uint64_t end; @@ -75,11 +86,14 @@ struct Segment { struct PTO2TensorMapLayout { size_t off_buckets; size_t off_bucket_epochs; + size_t off_reader_buckets; + size_t off_reader_bucket_epochs; size_t off_entry_pool; size_t off_free_entry_list; size_t off_task_entry_heads[PTO2_MAX_RING_DEPTH]; size_t off_task_entry_head_epochs[PTO2_MAX_RING_DEPTH]; int32_t num_buckets; + int32_t num_reader_buckets; int32_t pool_size; int32_t task_window_sizes[PTO2_MAX_RING_DEPTH]; }; @@ -108,7 +122,7 @@ extern uint64_t g_insert_count; * the hash key, size in [8, 16) is unused by the entry — we repurpose it for * `next_in_bucket`). * - * buffer_addr / next_in_bucket / producer_task_id — chain traversal + match + * buffer_addr / next_in_bucket / access_task_id — chain traversal + match * start_offset — overlap byte range begin * version, ndims, dtype, manual_dep, is_contiguous — overlap fast path * shapes[5] — overlap comparison (line 1) @@ -128,15 +142,15 @@ struct alignas(64) PTO2TensorMapEntry { // === Cache line 1 (64B) — lookup hot path; mirrors ChipTensor line 1 from byte 16 === uint64_t buffer_addr; // 8B [0, 8): tensor base address (hash key, mirrors ChipTensor::buffer.addr) PTO2TensorMapEntry - *next_in_bucket; // 8B [8, 16): next entry in hash bucket chain (overlays ChipTensor::buffer.size) - PTO2TaskId producer_task_id; // 8B [16,24): mirrors ChipTensor::owner_task_id slot - uint64_t start_offset; // 8B [24,32): mirrors ChipTensor::start_offset (element offset) - int32_t version; // 4B [32,36): mirrors ChipTensor::version - uint32_t ndims; // 4B [36,40): mirrors ChipTensor::ndims - DataType dtype; // 1B [40,41): mirrors ChipTensor::dtype - bool manual_dep; // 1B [41,42): mirrors ChipTensor::manual_dep - bool is_contiguous; // 1B [42,43): mirrors ChipTensor::is_contiguous - uint8_t __padding1__; // 1B [43,44): mirrors ChipTensor padding + *next_in_bucket; // 8B [8, 16): next entry in hash bucket chain (overlays ChipTensor::buffer.size) + PTO2TaskId access_task_id; // 8B [16,24): task owning this access + uint64_t start_offset; // 8B [24,32): mirrors ChipTensor::start_offset (element offset) + int32_t version; // 4B [32,36): mirrors ChipTensor::version + uint32_t ndims; // 4B [36,40): mirrors ChipTensor::ndims + DataType dtype; // 1B [40,41): mirrors ChipTensor::dtype + bool manual_dep; // 1B [41,42): mirrors ChipTensor::manual_dep + bool is_contiguous; // 1B [42,43): mirrors ChipTensor::is_contiguous + TensorAccessKind access_kind; // 1B [43,44): reader or writer index uint32_t shapes[MAX_TENSOR_DIMS]; // 20B [44,64): mirrors ChipTensor::shapes // === Cache line 2 (64B) — chain manipulation + non-contiguous overlap data === @@ -153,7 +167,7 @@ struct alignas(64) PTO2TensorMapEntry { * Copy overlap-relevant fields from a ChipTensor into this entry. * * 64B memcpy of ChipTensor cache line 1 populates buffer_addr (byte [0,8)), - * producer_task_id, start_offset, version, ndims, dtype, manual_dep, + * access_task_id, start_offset, version, ndims, dtype, manual_dep, * is_contiguous and shapes[]. Byte [8,16) holds ChipTensor::buffer.size in * the source and gets written into next_in_bucket; that's harmless * because link_entry() overwrites next_in_bucket immediately after. @@ -338,7 +352,7 @@ struct alignas(64) PTO2TensorMapEntry { static_assert(sizeof(PTO2TensorMapEntry) == 128, "TensorMapEntry must be exactly 2 cache lines (128 bytes)"); static_assert(offsetof(PTO2TensorMapEntry, buffer_addr) == offsetof(ChipTensor, buffer.addr)); -static_assert(offsetof(PTO2TensorMapEntry, producer_task_id) == offsetof(ChipTensor, owner_task_id)); +static_assert(offsetof(PTO2TensorMapEntry, access_task_id) == offsetof(ChipTensor, owner_task_id)); static_assert(offsetof(PTO2TensorMapEntry, start_offset) == offsetof(ChipTensor, start_offset)); static_assert(offsetof(PTO2TensorMapEntry, version) == offsetof(ChipTensor, version)); static_assert(offsetof(PTO2TensorMapEntry, ndims) == offsetof(ChipTensor, ndims)); @@ -385,6 +399,14 @@ struct PTO2TensorMap { // Per-ring cleanup progress (for periodic cleanup_retired) int32_t last_cleanup[PTO2_MAX_RING_DEPTH]{}; + // Tracked-reader side index. Keep these fields after the original writer + // state so the writer-only hot path retains main's field offsets/cache layout. + PTO2TensorMapEntry **reader_buckets; + uint32_t *reader_bucket_epochs; + int32_t num_reader_buckets; + int32_t reader_used{0}; + int32_t reader_high_water{0}; + uint32_t get_task_local_id_slot(uint8_t ring_id, uint32_t task_local_id) const { return task_local_id & (task_window_sizes[ring_id] - 1); } @@ -394,6 +416,8 @@ struct PTO2TensorMap { // setter symbols must export for host dlsym; the probe call sites that use // these accessors stay gated by SIMPLER_DFX). int32_t current_used() const { return next_entry_idx - free_num; } + int32_t current_readers() const { return reader_used; } + int32_t current_writers() const { return current_used() - reader_used; } int32_t pool_capacity() const { return pool_size; } int32_t free_entries() const { return pool_size - current_used(); } @@ -430,23 +454,37 @@ struct PTO2TensorMap { return res; } - void free_entry(PTO2TensorMapEntry &entry) { + void free_writer_entry(PTO2TensorMapEntry &entry) { always_assert(entry.bucket_index != -1); // must still be in a bucket - // Update predecessor's next pointer (O(1) via prev_in_bucket) if (entry.prev_in_bucket == nullptr) { - // Entry is the head of its bucket chain, update bucket head - // Must compute hash BEFORE clearing tensor buckets[entry.bucket_index] = entry.next_in_bucket; } else { entry.prev_in_bucket->next_in_bucket = entry.next_in_bucket; } - - // Update successor's prev pointer if (entry.next_in_bucket != nullptr) { entry.next_in_bucket->prev_in_bucket = entry.prev_in_bucket; } + free_entry_list[free_num++] = &entry; + entry.bucket_index = -1; + entry.next_in_bucket = nullptr; + entry.prev_in_bucket = nullptr; + entry.next_in_task = nullptr; + entry.prev_in_task = nullptr; + } + void free_reader_entry(PTO2TensorMapEntry &entry) { + always_assert(entry.bucket_index != -1); + if (entry.prev_in_bucket == nullptr) { + reader_buckets[entry.bucket_index] = entry.next_in_bucket; + } else { + entry.prev_in_bucket->next_in_bucket = entry.next_in_bucket; + } + if (entry.next_in_bucket != nullptr) { + entry.next_in_bucket->prev_in_bucket = entry.prev_in_bucket; + } + always_assert(reader_used > 0); + reader_used--; free_entry_list[free_num++] = &entry; entry.bucket_index = -1; entry.next_in_bucket = nullptr; @@ -455,6 +493,14 @@ struct PTO2TensorMap { entry.prev_in_task = nullptr; } + void free_entry(PTO2TensorMapEntry &entry) { + if (entry.access_kind == TensorAccessKind::READER) { + free_reader_entry(entry); + } else { + free_writer_entry(entry); + } + } + // ============================================================================= // TensorMap API // ============================================================================= @@ -521,13 +567,17 @@ struct PTO2TensorMap { * @param tensor ChipTensor to look up * @param on_match Callback invoked for each overlapping entry */ - template - void lookup(const ChipTensor &tensor, Fn &&on_match) { - uint32_t bucket_index = hash(tensor.buffer.addr); - if (bucket_epochs[bucket_index] != current_epoch) { + template + void lookup_impl(const ChipTensor &tensor, Fn &&on_match) { + constexpr bool is_reader = access_kind == TensorAccessKind::READER; + int32_t bucket_count = is_reader ? num_reader_buckets : num_buckets; + uint32_t bucket_index = hash(tensor.buffer.addr, bucket_count); + uint32_t *epochs = is_reader ? reader_bucket_epochs : bucket_epochs; + PTO2TensorMapEntry **index = is_reader ? reader_buckets : buckets; + if (epochs[bucket_index] != current_epoch) { return; } - PTO2TensorMapEntry *cur_entry = buckets[bucket_index]; + PTO2TensorMapEntry *cur_entry = index[bucket_index]; #if SIMPLER_TENSORMAP_PROFILING g_lookup_count++; @@ -579,19 +629,41 @@ struct PTO2TensorMap { #endif } + template + void lookup(const ChipTensor &tensor, Fn &&on_match) { + lookup_impl(tensor, static_cast(on_match)); + } + + template + void lookup(const ChipTensor &tensor, TensorAccessKind access_kind, Fn &&on_match) { + if (access_kind == TensorAccessKind::READER) { + lookup_impl(tensor, static_cast(on_match)); + } else { + lookup_impl(tensor, static_cast(on_match)); + } + } + /** - * Insert a new entry (called when task produces output) + * Insert a new reader or writer access entry. * * Allocates from ring buffer pool, may overwrite stale entries. * Inserts at head of hash bucket chain (maintains task_id ordering). * - * @param tensor ChipTensor produced - * @param producer_task_id Task ID of producer + * @param tensor ChipTensor accessed + * @param access_task_id Task ID owning the access */ - void insert(const ChipTensor &tensor, PTO2TaskId producer_task_id) { + void insert(const ChipTensor &tensor, PTO2TaskId access_task_id, TensorAccessKind access_kind) { + if (access_kind == TensorAccessKind::READER) { + insert_impl(tensor, access_task_id); + } else { + insert_impl(tensor, access_task_id); + } + } + + void insert(const ChipTensor &tensor, PTO2TaskId access_task_id) { PTO2TensorMapEntry *entry = new_entry(); entry->copy_from_tensor(tensor); - link_entry(entry, tensor.buffer.addr, producer_task_id); + link_entry(entry, tensor.buffer.addr, access_task_id); } /** @@ -618,7 +690,7 @@ struct PTO2TensorMap { PTO2TensorMapEntry *cur_entry = task_entry_heads[ring_id][task_slot]; while (cur_entry != nullptr) { PTO2TensorMapEntry *next_entry = cur_entry->next_in_task; // free_entry clears it - if (cur_entry->producer_task_id == retired_task) { + if (cur_entry->access_task_id == retired_task) { if (cur_entry->prev_in_task != nullptr) { cur_entry->prev_in_task->next_in_task = next_entry; } else { @@ -646,26 +718,81 @@ struct PTO2TensorMap { * addresses (low bits all-zero) still distribute evenly. We extract * the top log2(num_buckets) bits which carry the most entropy. */ - uint32_t hash(uint64_t key) { + uint32_t hash(uint64_t key, int32_t bucket_count) { key *= 0x9E3779B97F4A7C15ULL; - return static_cast(key >> (64 - __builtin_ctz(num_buckets))); + return static_cast(key >> (64 - __builtin_ctz(bucket_count))); } + uint32_t hash(uint64_t key) { return hash(key, num_buckets); } + /** * Link an initialized entry into bucket and task chains. */ - void link_entry(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId producer_task_id) { + template + void link_entry_impl(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id) { #if SIMPLER_TENSORMAP_PROFILING g_insert_count++; #endif - uint32_t bucket_index = hash(addr); - auto ring_id = producer_task_id.ring(); - auto local_id = producer_task_id.local(); + constexpr bool is_reader = access_kind == TensorAccessKind::READER; + int32_t bucket_count = is_reader ? num_reader_buckets : num_buckets; + uint32_t bucket_index = hash(addr, bucket_count); + auto ring_id = access_task_id.ring(); + auto local_id = access_task_id.local(); int32_t task_slot = local_id & (task_window_sizes[ring_id] - 1); - entry->producer_task_id = producer_task_id; + entry->access_task_id = access_task_id; + entry->access_kind = access_kind; + if constexpr (is_reader) { + reader_used++; + if (reader_used > reader_high_water) reader_high_water = reader_used; + } // Insert at head of hash bucket + PTO2TensorMapEntry **index = is_reader ? reader_buckets : buckets; + uint32_t *epochs = is_reader ? reader_bucket_epochs : bucket_epochs; + if (epochs[bucket_index] != current_epoch) { + index[bucket_index] = nullptr; + epochs[bucket_index] = current_epoch; + } + entry->bucket_index = bucket_index; + entry->next_in_bucket = index[bucket_index]; + if (entry->next_in_bucket != nullptr) { + entry->next_in_bucket->prev_in_bucket = entry; + } + index[bucket_index] = entry; + entry->prev_in_bucket = nullptr; + + // Link to task's entry list + if (task_entry_head_epochs[ring_id][task_slot] != current_epoch) { + task_entry_heads[ring_id][task_slot] = nullptr; + task_entry_head_epochs[ring_id][task_slot] = current_epoch; + } + entry->next_in_task = task_entry_heads[ring_id][task_slot]; + entry->prev_in_task = nullptr; + if (entry->next_in_task != nullptr) { + entry->next_in_task->prev_in_task = entry; + } + task_entry_heads[ring_id][task_slot] = entry; + } + + template + void insert_impl(const ChipTensor &tensor, PTO2TaskId access_task_id) { + PTO2TensorMapEntry *entry = new_entry(); + entry->copy_from_tensor(tensor); + link_entry_impl(entry, tensor.buffer.addr, access_task_id); + } + + void link_entry(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id) { +#if SIMPLER_TENSORMAP_PROFILING + g_insert_count++; +#endif + uint32_t bucket_index = hash(addr); + auto ring_id = access_task_id.ring(); + auto local_id = access_task_id.local(); + int32_t task_slot = local_id & (task_window_sizes[ring_id] - 1); + + entry->access_task_id = access_task_id; + entry->access_kind = TensorAccessKind::WRITER; if (bucket_epochs[bucket_index] != current_epoch) { buckets[bucket_index] = nullptr; bucket_epochs[bucket_index] = current_epoch; @@ -678,7 +805,6 @@ struct PTO2TensorMap { buckets[bucket_index] = entry; entry->prev_in_bucket = nullptr; - // Link to task's entry list if (task_entry_head_epochs[ring_id][task_slot] != current_epoch) { task_entry_heads[ring_id][task_slot] = nullptr; task_entry_head_epochs[ring_id][task_slot] = current_epoch; @@ -691,11 +817,21 @@ struct PTO2TensorMap { task_entry_heads[ring_id][task_slot] = entry; } + void link_entry( + PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id, TensorAccessKind access_kind + ) { + if (access_kind == TensorAccessKind::READER) { + link_entry_impl(entry, addr, access_task_id); + } else { + link_entry_impl(entry, addr, access_task_id); + } + } + /** * Check if entry is valid (producer has not retired) */ bool entry_valid(const PTO2TensorMapEntry &entry) const { - return static_cast(entry.producer_task_id.local()) >= last_task_alives[entry.producer_task_id.ring()]; + return static_cast(entry.access_task_id.local()) >= last_task_alives[entry.access_task_id.ring()]; } void remove_entry(PTO2TensorMapEntry &entry) { @@ -703,6 +839,16 @@ struct PTO2TensorMap { free_entry(entry); } + void remove_writer_entry(PTO2TensorMapEntry &entry) { + remove_from_task(entry); + free_writer_entry(entry); + } + + void remove_reader_entry(PTO2TensorMapEntry &entry) { + remove_from_task(entry); + free_reader_entry(entry); + } + /** * Remove entry from its task chain (O(1) with prev pointer) * Called during pool wrap-around to unlink reused entries. @@ -712,8 +858,8 @@ struct PTO2TensorMap { // Update predecessor's next pointer (O(1) via prev_in_task) if (entry.prev_in_task == nullptr) { // Entry is the head of its task chain, update task_entry_heads - int32_t ring_id = entry.producer_task_id.ring(); - int32_t local_id = static_cast(entry.producer_task_id.local()); + int32_t ring_id = entry.access_task_id.ring(); + int32_t local_id = static_cast(entry.access_task_id.local()); int32_t task_slot = local_id & (task_window_sizes[ring_id] - 1); task_entry_heads[ring_id][task_slot] = entry.next_in_task; } else { diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h index f8d86d6bac..c70655d96f 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h @@ -379,6 +379,15 @@ struct Arg : TaskArgsTpl { ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::INPUT, tensor_count_++), ...); } + template + void add_tracked_input(Args &&...args) { + assert_add_tensor_args(); + if (!check_add_tensor_capacity(static_cast(sizeof...(Args)))) { + return; + } + ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::TRACKED_INPUT, tensor_count_++), ...); + } + /// Batch add outputs — all ChipTensor or all TensorCreateInfo: /// add_output(ci1, ci2) — runtime allocates buffers (OUTPUT) /// add_output(t1, t2) — write-only existing tensors (OUTPUT_EXISTING) @@ -403,7 +412,7 @@ struct Arg : TaskArgsTpl { ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::INOUT, tensor_count_++), ...); } - /// No-dependency existing tensor: skips OverlapMap lookup, depends on creator only. + /// Existing tensor with creator retention but no automatic TensorMap lookup or publication. template void add_no_dep(Args &&...args) { assert_add_tensor_args(); diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp index 3cfa357a49..dfd3cd31dd 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp @@ -9,7 +9,7 @@ * ----------------------------------------------------------------------------------------------------------- */ /** - * PTO Runtime2 - TensorMap Implementation + * tensormap_and_ringbuffer - TensorMap Implementation * * Implements TensorMap with ring buffer pool, lazy invalidation, * and chain truncation optimization. @@ -56,6 +56,8 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout( PTO2TensorMapLayout layout{}; layout.num_buckets = new_num_buckets; + layout.num_reader_buckets = + new_num_buckets < PTO2_TENSORMAP_READER_NUM_BUCKETS ? new_num_buckets : PTO2_TENSORMAP_READER_NUM_BUCKETS; layout.pool_size = new_pool_size; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { layout.task_window_sizes[r] = new_task_window_sizes[r]; @@ -66,6 +68,11 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout( ); layout.off_bucket_epochs = arena.reserve(static_cast(new_num_buckets) * sizeof(uint32_t), alignof(uint32_t)); + layout.off_reader_buckets = arena.reserve( + static_cast(layout.num_reader_buckets) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) + ); + layout.off_reader_bucket_epochs = + arena.reserve(static_cast(layout.num_reader_buckets) * sizeof(uint32_t), alignof(uint32_t)); layout.off_entry_pool = arena.reserve(static_cast(new_pool_size) * sizeof(PTO2TensorMapEntry), alignof(PTO2TensorMapEntry)); layout.off_free_entry_list = @@ -87,12 +94,15 @@ PTO2TensorMap::reserve_layout_default(DeviceArena &arena, const int32_t new_task bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, DeviceArena &arena) { num_buckets = layout.num_buckets; + num_reader_buckets = layout.num_reader_buckets; pool_size = layout.pool_size; // Address arena regions for data writes; do not store these in struct // fields (wire_arena_pointers does that). auto *buckets_arena = static_cast(arena.region_ptr(layout.off_buckets)); auto *bucket_epochs_arena = static_cast(arena.region_ptr(layout.off_bucket_epochs)); + auto *reader_buckets_arena = static_cast(arena.region_ptr(layout.off_reader_buckets)); + auto *reader_bucket_epochs_arena = static_cast(arena.region_ptr(layout.off_reader_bucket_epochs)); auto *entry_pool_arena = static_cast(arena.region_ptr(layout.off_entry_pool)); auto *free_list_arena = static_cast(arena.region_ptr(layout.off_free_entry_list)); @@ -101,6 +111,10 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev buckets_arena[i] = nullptr; bucket_epochs_arena[i] = 0; } + for (int32_t i = 0; i < num_reader_buckets; i++) { + reader_buckets_arena[i] = nullptr; + reader_bucket_epochs_arena[i] = 0; + } // entry_pool: zero-init equivalent to the previous calloc(entry_pool, ...). // The pool's persistent invariant after init is "bucket_index == -1 means @@ -112,7 +126,7 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev entry_pool_arena[i].prev_in_bucket = nullptr; entry_pool_arena[i].next_in_task = nullptr; entry_pool_arena[i].prev_in_task = nullptr; - entry_pool_arena[i].producer_task_id = PTO2TaskId{}; + entry_pool_arena[i].access_task_id = PTO2TaskId{}; } // free_entry_list: zeroed (was calloc'd before); contents become meaningful @@ -121,6 +135,8 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev next_entry_idx = 0; free_num = 0; + reader_used = 0; + reader_high_water = 0; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { auto *heads_arena = static_cast(arena.region_ptr(layout.off_task_entry_heads[r])); @@ -139,13 +155,17 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev void PTO2TensorMap::reset_for_reuse(const PTO2TensorMapLayout &layout) { num_buckets = layout.num_buckets; + num_reader_buckets = layout.num_reader_buckets; pool_size = layout.pool_size; next_entry_idx = 0; free_num = 0; + reader_used = 0; + reader_high_water = 0; current_epoch++; if (current_epoch == 0) { current_epoch = 1; memset(bucket_epochs, 0, static_cast(layout.num_buckets) * sizeof(uint32_t)); + memset(reader_bucket_epochs, 0, static_cast(layout.num_reader_buckets) * sizeof(uint32_t)); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { memset(task_entry_head_epochs[r], 0, static_cast(layout.task_window_sizes[r]) * sizeof(uint32_t)); } @@ -161,6 +181,8 @@ void PTO2TensorMap::reset_for_reuse(const PTO2TensorMapLayout &layout) { void PTO2TensorMap::wire_arena_pointers(const PTO2TensorMapLayout &layout, DeviceArena &arena) { buckets = static_cast(arena.region_ptr(layout.off_buckets)); bucket_epochs = static_cast(arena.region_ptr(layout.off_bucket_epochs)); + reader_buckets = static_cast(arena.region_ptr(layout.off_reader_buckets)); + reader_bucket_epochs = static_cast(arena.region_ptr(layout.off_reader_bucket_epochs)); entry_pool = static_cast(arena.region_ptr(layout.off_entry_pool)); free_entry_list = static_cast(arena.region_ptr(layout.off_free_entry_list)); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { @@ -175,6 +197,8 @@ void PTO2TensorMap::destroy() { // a recycled allocation. buckets = nullptr; bucket_epochs = nullptr; + reader_buckets = nullptr; + reader_bucket_epochs = nullptr; entry_pool = nullptr; free_entry_list = nullptr; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { @@ -190,10 +214,32 @@ void PTO2TensorMap::destroy() { void PTO2TensorMap::print_stats() { int32_t valid = 0; int32_t stale = 0; - int32_t empty_buckets = 0; - int32_t max_chain = 0; - int64_t total_chain = 0; - int32_t non_empty_buckets = 0; + + struct BucketStats { + int32_t empty{0}; + int32_t non_empty{0}; + int32_t max_chain{0}; + int64_t total_chain{0}; + }; + auto collect_bucket_stats = [&](PTO2TensorMapEntry **index, uint32_t *epochs, int32_t bucket_count) { + BucketStats stats; + for (int32_t b = 0; b < bucket_count; b++) { + int32_t chain_len = 0; + if (epochs[b] == current_epoch) { + for (auto *entry = index[b]; entry != nullptr; entry = entry->next_in_bucket) { + chain_len++; + } + } + if (chain_len == 0) { + stats.empty++; + } else { + stats.non_empty++; + stats.total_chain += chain_len; + stats.max_chain = std::max(stats.max_chain, chain_len); + } + } + return stats; + }; // Count entries for (int32_t i = 0; i < pool_size; i++) { @@ -206,37 +252,26 @@ void PTO2TensorMap::print_stats() { } } - // Count bucket stats - for (int32_t b = 0; b < num_buckets; b++) { - int32_t chain_len = 0; - auto cur_entry = buckets[b]; - - while (cur_entry != nullptr) { - chain_len++; - cur_entry = cur_entry->next_in_bucket; - } - - if (chain_len == 0) { - empty_buckets++; - } else { - non_empty_buckets++; - total_chain += chain_len; - if (chain_len > max_chain) { - max_chain = chain_len; - } - } - } + const BucketStats writer_stats = collect_bucket_stats(buckets, bucket_epochs, num_buckets); + const BucketStats reader_stats = collect_bucket_stats(reader_buckets, reader_bucket_epochs, num_reader_buckets); LOG_DEBUG("=== TensorMap Statistics ==="); LOG_DEBUG("Pool size: %d", pool_size); LOG_DEBUG("Pool next entry idx: %d", next_entry_idx); LOG_DEBUG("Pool free_num: %d", free_num); - LOG_DEBUG("Num buckets: %d", num_buckets); + LOG_DEBUG("Reader entries: live=%d high_water=%d", reader_used, reader_high_water); + LOG_DEBUG("Writer entries: live=%d", current_writers()); + LOG_DEBUG("Buckets: writer=%d reader=%d", num_buckets, num_reader_buckets); LOG_DEBUG("Valid entries: %d", valid); LOG_DEBUG("Stale entries: %d", stale); - LOG_DEBUG("Empty buckets: %d", empty_buckets); - LOG_DEBUG("Max chain len: %d", max_chain); - LOG_DEBUG("Avg chain len: %.2f", non_empty_buckets > 0 ? (float)total_chain / non_empty_buckets : 0); + LOG_DEBUG( + "Writer buckets: empty=%d max_chain=%d avg_chain=%.2f", writer_stats.empty, writer_stats.max_chain, + writer_stats.non_empty > 0 ? (float)writer_stats.total_chain / writer_stats.non_empty : 0 + ); + LOG_DEBUG( + "Reader buckets: empty=%d max_chain=%d avg_chain=%.2f", reader_stats.empty, reader_stats.max_chain, + reader_stats.non_empty > 0 ? (float)reader_stats.total_chain / reader_stats.non_empty : 0 + ); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { LOG_DEBUG("Last task alive[%d]: %d", r, last_task_alives[r]); } diff --git a/src/a5/platform/onboard/host/device_runner.cpp b/src/a5/platform/onboard/host/device_runner.cpp index 5feb045751..2928948886 100644 --- a/src/a5/platform/onboard/host/device_runner.cpp +++ b/src/a5/platform/onboard/host/device_runner.cpp @@ -64,12 +64,24 @@ extern "C" __attribute__((weak, visibility("hidden"))) int dep_gen_replay_emit_d return -1; } +extern "C" __attribute__((weak, visibility("hidden"))) bool dep_gen_host_graph_active() { return false; } +extern "C" __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_set_enabled(bool /*enable*/) {} +extern "C" __attribute__((weak, visibility("hidden"))) int dep_gen_host_graph_emit(const char * /*deps_json_path*/) { + LOG_DEBUG("dep_gen host graph not implemented for this runtime — deps.json skipped"); + return -1; +} + // ============================================================================= // DeviceRunner Implementation // ============================================================================= DeviceRunner::~DeviceRunner() { finalize(); } +void DeviceRunner::set_dep_gen_enabled(bool enable) { + enable_dep_gen_ = enable; + dep_gen_host_graph_set_enabled(enable); +} + // `setup_static_arena`, `create_thread`, `attach_current_thread`, // `configure_aicore_op_timeout`, `ensure_device_initialized`, // `ensure_binaries_loaded`, `query_max_block_dim`, and `validate_block_dim` @@ -287,7 +299,9 @@ int DeviceRunner::prepare_execution( if (enable_dump_args_) SIMPLER_SET_DFX_FLAG(enable_profiling_flag, SIMPLER_DFX_FLAG_DUMP_ARGS); if (enable_chip_swimlane_) SIMPLER_SET_DFX_FLAG(enable_profiling_flag, SIMPLER_DFX_FLAG_CHIP_SWIMLANE); if (enable_pmu_) SIMPLER_SET_DFX_FLAG(enable_profiling_flag, SIMPLER_DFX_FLAG_PMU); - if (enable_dep_gen_) SIMPLER_SET_DFX_FLAG(enable_profiling_flag, SIMPLER_DFX_FLAG_DEP_GEN); + if (enable_dep_gen_ && !dep_gen_host_graph_active()) { + SIMPLER_SET_DFX_FLAG(enable_profiling_flag, SIMPLER_DFX_FLAG_DEP_GEN); + } if (enable_scope_stats_) SIMPLER_SET_DFX_FLAG(enable_profiling_flag, SIMPLER_DFX_FLAG_SCOPE_STATS); execution->kernel_args.args.enable_profiling_flag = enable_profiling_flag; @@ -399,7 +413,7 @@ int DeviceRunner::prepare_execution( } } - if (enable_dep_gen_) { + if (enable_dep_gen_ && !dep_gen_host_graph_active()) { rc = init_dep_gen(active_aicpu_num, device_id_, execution->kernel_args); if (rc != 0) { LOG_ERROR("init_dep_gen failed: %d", rc); @@ -454,7 +468,7 @@ DeviceRunner::launch_execution(std::unique_ptr prepared, Laun activate_launch_shape(runtime); (void)arm_device_wall_buffer(prepared->kernel_args); start_shared_collectors_for_run(); - if (enable_dep_gen_) { + if (enable_dep_gen_ && !dep_gen_host_graph_active()) { auto thread_factory = [this](std::function fn) { return create_thread(std::move(fn)); }; @@ -592,15 +606,21 @@ int DeviceRunner::drain_execution(ActiveExecution &active) { read_device_wall_ns(); teardown_shared_collectors_after_run(); - // a5-specific dep_gen teardown: stop + reconcile + replay emit. if (enable_dep_gen_) { - dep_gen_collector_.stop(); - if (dep_gen_collector_.reconcile_counters()) { - const auto &records = dep_gen_collector_.records(); - const std::string deps = make_deps_json_path(output_prefix_); - int replay_rc = dep_gen_replay_emit_deps_json(records.data(), records.size(), deps.c_str()); - if (replay_rc != 0) { - LOG_ERROR("dep_gen replay failed (%d) — deps.json not produced", replay_rc); + const std::string deps = make_deps_json_path(output_prefix_); + if (dep_gen_host_graph_active()) { + const int emit_rc = dep_gen_host_graph_emit(deps.c_str()); + if (emit_rc != 0) { + LOG_ERROR("dep_gen host graph emit failed (%d) — deps.json not produced", emit_rc); + } + } else { + dep_gen_collector_.stop(); + if (dep_gen_collector_.reconcile_counters()) { + const auto &records = dep_gen_collector_.records(); + const int replay_rc = dep_gen_replay_emit_deps_json(records.data(), records.size(), deps.c_str()); + if (replay_rc != 0) { + LOG_ERROR("dep_gen replay failed (%d) — deps.json not produced", replay_rc); + } } } } diff --git a/src/a5/platform/onboard/host/device_runner.h b/src/a5/platform/onboard/host/device_runner.h index fc2e3111b3..e43f6eb56a 100644 --- a/src/a5/platform/onboard/host/device_runner.h +++ b/src/a5/platform/onboard/host/device_runner.h @@ -106,7 +106,7 @@ class DeviceRunner : public DeviceRunnerBase { * a5 `dep_gen` enablement setter, overriding the base no-op. Captures * orchestrator submit_task inputs for offline replay into deps.json. */ - void set_dep_gen_enabled(bool enable) override { enable_dep_gen_ = enable; } + void set_dep_gen_enabled(bool enable) override; /** * Cleanup all resources diff --git a/src/a5/platform/sim/host/device_runner.cpp b/src/a5/platform/sim/host/device_runner.cpp index 56a9520cb1..7d880e40b1 100644 --- a/src/a5/platform/sim/host/device_runner.cpp +++ b/src/a5/platform/sim/host/device_runner.cpp @@ -58,6 +58,13 @@ extern "C" __attribute__((weak, visibility("hidden"))) int dep_gen_replay_emit_d return -1; } +extern "C" __attribute__((weak, visibility("hidden"))) bool dep_gen_host_graph_active() { return false; } +extern "C" __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_set_enabled(bool /*enable*/) {} +extern "C" __attribute__((weak, visibility("hidden"))) int dep_gen_host_graph_emit(const char * /*deps_json_path*/) { + LOG_DEBUG("dep_gen host graph not implemented for this runtime — deps.json skipped"); + return -1; +} + // a5 sim: malloc / free wrappers shared by the four profiling subsystems' // init_* methods. Plain function pointers convert implicitly into the // framework's std::function alloc / free shapes. Kept on the subclass (not @@ -92,6 +99,11 @@ struct DeviceRunner::ActiveRun { DeviceRunner::DeviceRunner() = default; DeviceRunner::~DeviceRunner() { finalize(); } +void DeviceRunner::set_dep_gen_enabled(bool enable) { + enable_dep_gen_ = enable; + dep_gen_host_graph_set_enabled(enable); +} + void DeviceRunner::cleanup_active_run() noexcept { if (active_run_ == nullptr) return; active_run_->join(); @@ -308,7 +320,7 @@ int DeviceRunner::prepare_execution( if (enable_pmu_) { SIMPLER_SET_DFX_FLAG(enable_profiling_flag, SIMPLER_DFX_FLAG_PMU); } - if (enable_dep_gen_) { + if (enable_dep_gen_ && !dep_gen_host_graph_active()) { SIMPLER_SET_DFX_FLAG(enable_profiling_flag, SIMPLER_DFX_FLAG_DEP_GEN); } if (enable_scope_stats_) { @@ -367,7 +379,7 @@ int DeviceRunner::prepare_execution( } } - if (enable_dep_gen_) { + if (enable_dep_gen_ && !dep_gen_host_graph_active()) { rc = init_dep_gen(launch_aicpu_num, device_id_); if (rc != 0) { LOG_ERROR("init_dep_gen failed: %d", rc); @@ -463,7 +475,7 @@ DeviceRunner::launch_execution(std::unique_ptr prepared, Laun set_platform_pmu_base_func_(kernel_args_.pmu_data_base); set_pmu_enabled_func_(enable_pmu_); set_platform_dep_gen_base_func_(kernel_args_.dep_gen_data_base); - set_dep_gen_enabled_func_(enable_dep_gen_); + set_dep_gen_enabled_func_(enable_dep_gen_ && !dep_gen_host_graph_active()); set_scope_stats_enabled_func_(enable_scope_stats_); set_platform_scope_stats_base_func_(kernel_args_.scope_stats_data_base); @@ -473,7 +485,7 @@ DeviceRunner::launch_execution(std::unique_ptr prepared, Laun if (enable_chip_swimlane_) chip_swimlane_collector_.start(thread_factory); if (enable_dump_args_) dump_collector_.start(thread_factory); if (enable_pmu_) pmu_collector_.start(thread_factory); - if (enable_dep_gen_) dep_gen_collector_.start(thread_factory); + if (enable_dep_gen_ && !dep_gen_host_graph_active()) dep_gen_collector_.start(thread_factory); if (enable_scope_stats_) scope_stats_collector_.start(thread_factory); if (kernel_args_.device_wall_data_base != 0) { @@ -618,13 +630,20 @@ int DeviceRunner::drain_execution(ActiveExecution &) { } if (enable_dep_gen_) { - dep_gen_collector_.stop(); - if (dep_gen_collector_.reconcile_counters()) { - const auto &records = dep_gen_collector_.records(); - const std::string deps = make_deps_json_path(output_prefix_); - int replay_rc = dep_gen_replay_emit_deps_json(records.data(), records.size(), deps.c_str()); - if (replay_rc != 0) { - LOG_ERROR("dep_gen replay failed (%d) — deps.json not produced", replay_rc); + const std::string deps = make_deps_json_path(output_prefix_); + if (dep_gen_host_graph_active()) { + const int emit_rc = dep_gen_host_graph_emit(deps.c_str()); + if (emit_rc != 0) { + LOG_ERROR("dep_gen host graph emit failed (%d) — deps.json not produced", emit_rc); + } + } else { + dep_gen_collector_.stop(); + if (dep_gen_collector_.reconcile_counters()) { + const auto &records = dep_gen_collector_.records(); + const int replay_rc = dep_gen_replay_emit_deps_json(records.data(), records.size(), deps.c_str()); + if (replay_rc != 0) { + LOG_ERROR("dep_gen replay failed (%d) — deps.json not produced", replay_rc); + } } } } diff --git a/src/a5/platform/sim/host/device_runner.h b/src/a5/platform/sim/host/device_runner.h index 17f2ee8e52..9e2c8c4f93 100644 --- a/src/a5/platform/sim/host/device_runner.h +++ b/src/a5/platform/sim/host/device_runner.h @@ -44,7 +44,7 @@ class DeviceRunner : public SimDeviceRunnerBase { int finalize() override; // a5 dep_gen enablement setter, overriding the base no-op (the c_api // unconditionally calls it). - void set_dep_gen_enabled(bool enable) override { enable_dep_gen_ = enable; } + void set_dep_gen_enabled(bool enable) override; private: struct ActiveRun; diff --git a/src/a5/runtime/host_build_graph/docs/RUNTIME_LOGIC.md b/src/a5/runtime/host_build_graph/docs/RUNTIME_LOGIC.md index d50bb8b511..dc2268415e 100644 --- a/src/a5/runtime/host_build_graph/docs/RUNTIME_LOGIC.md +++ b/src/a5/runtime/host_build_graph/docs/RUNTIME_LOGIC.md @@ -165,13 +165,14 @@ AIC/MIX tasks use the available cluster count. ### 5.2 TensorMap and Fanins -TensorMap maps tensor regions to producer task IDs. For every task: +TensorMap indexes unfinished reader and writer accesses. For every task: -1. INPUT/INOUT regions look up overlapping producers. -2. Explicit and discovered producers are deduplicated into +1. INPUT/TRACKED_INPUT/INOUT regions look up overlapping writers. +2. INOUT/OUTPUT_EXISTING regions look up overlapping tracked readers. +3. Explicit and discovered producers are deduplicated into `fanin_local_ids[]`. -3. OUTPUT/INOUT regions register the new task as producer. -4. Each producer tracks its highest consumer local ID for completion metadata. +4. TRACKED_INPUT registers a reader; INOUT/OUTPUT_EXISTING register writers. +5. Each producer tracks its highest consumer local ID for completion metadata. There is no fanout adjacency or dependency pool. A per-slot completion flag is the readiness truth on device. @@ -216,10 +217,12 @@ could strand the drain protocol, so the active path relies on that invariant. ## 8. Scalar Access During Construction -`get_tensor_data` and `set_tensor_data` operate on registered host views of -external tensors. They cannot wait for a submitted device producer because the -device scheduler starts only after orchestration returns. Runtime-created graph- -heap outputs also have no host view. +`get_tensor_data` reads registered host views of external tensors and therefore +cannot wait for a submitted device producer: the scheduler starts only after +orchestration returns. Before any task is submitted, `set_tensor_data` updates +that host view directly. After submission begins, it emits a scheduler-local +host-write node that waits for conflicting writers, their direct consumers, and +tracked readers; this also supports runtime-created graph-heap outputs. Producer references are checked against the complete bound descriptor ID before a slot is used, preventing masked-slot aliasing. See diff --git a/src/a5/runtime/host_build_graph/docs/SCALAR_DATA_ACCESS.md b/src/a5/runtime/host_build_graph/docs/SCALAR_DATA_ACCESS.md index c71dfb5036..aaca485238 100644 --- a/src/a5/runtime/host_build_graph/docs/SCALAR_DATA_ACCESS.md +++ b/src/a5/runtime/host_build_graph/docs/SCALAR_DATA_ACCESS.md @@ -9,15 +9,16 @@ external tensors; they do not interleave host code with device execution. | Tensor state | `get_tensor_data` | `set_tensor_data` | | ------------ | ----------------- | ----------------- | -| External tensor with no submitted producer | Reads the staged host value | Updates the staged host value | +| External tensor with no submitted producer | Reads the staged host value | Updates the host view before submission; otherwise emits a host-write node | | External control/output tensor not referenced by a task | Reads immediately | Writes immediately | -| Runtime-created output | Unsupported: no registered host view | Unsupported: no registered host view | -| Tensor owned by a submitted device task | Unsupported during graph construction | Unsupported during graph construction | +| Runtime-created output | Unsupported: no registered host view | Emits a host-write node after its producer | +| Tensor owned by a submitted device task | Unsupported during graph construction | Emits a host-write node after conflicting accesses | | Tensor with an invalid or stale owner task ID | Fails with `INVALID_ARGS` | Fails with `INVALID_ARGS` | -The supported write changes the data that will be copied to the device. Every -task in the graph observes that final staged value; submit order does not turn -the write into a barrier between kernels. +A write before the first submission changes the data staged for the device. A +later write becomes a dependency-aware scheduler node: it waits for overlapping +writers, their direct consumers, and tracked readers, then publishes itself as +the next writer for subsequent tasks. ## API @@ -28,9 +29,9 @@ int32_t value = get_tensor_data(control, 1, index); set_tensor_data(layout, 1, index, value + 1); ``` -Both tensors in this example must be external tensors staged by the host. A -common use is to read an input control value or publish runtime geometry into an -external layout tensor that no submitted task owns. +`control` must be an external tensor staged by the host. `layout` may also be +external, or it may be a runtime-created tensor when the write occurs after its +allocation task has been submitted. ## Why Device-Produced Values Cannot Be Read Here @@ -41,14 +42,14 @@ The execution order is: 3. The host relocates and copies the graph image to device memory. 4. AICPU schedulers boot and dispatch the graph. -A producer submitted in step 1 cannot become `COMPLETED` until step 4. Waiting -for that producer from the orchestration call cannot make progress. The runtime -keeps a timeout as a defensive failure backstop, but it is not a supported -synchronization mechanism. +A producer submitted in step 1 cannot become `COMPLETED` until step 4. A host +read that waits from orchestration therefore cannot make progress. A host write +does not wait synchronously; it is represented in the graph and executes in +step 4. -Runtime-created output buffers also live in the graph heap and have no host-view -registration. Even if their address is nonzero, host orchestration must not -dereference or modify them. +Runtime-created output buffers live in the graph heap and have no host-view +registration. Host orchestration must not dereference them, but a host-write +node may modify them on the device after their producer completes. ## Ownership Validation @@ -65,9 +66,11 @@ recorded. ## Practical Rules -- Use scalar access only on external, host-staged tensors with no device - producer or outstanding device consumer. -- Use tensor dependencies to order device tasks; do not use host scalar access - as a device synchronization barrier. +- Use `get_tensor_data` only on external, host-staged tensors with no device + producer. +- Mark pure readers with `add_tracked_input()` when a later host write overlaps. +- Use tensor dependencies to order device tasks; a host write is a graph node, + not a synchronous device barrier. - Pass values needed for graph construction as orchestration inputs or scalars. -- Keep device-produced values on the device or return them after the run. +- Keep device-produced values on the device or return them after the run; they + cannot drive host control flow during graph construction. diff --git a/src/a5/runtime/host_build_graph/host/dep_gen_host_graph.cpp b/src/a5/runtime/host_build_graph/host/dep_gen_host_graph.cpp index a11f05f3d5..5ee0f7cc14 100644 --- a/src/a5/runtime/host_build_graph/host/dep_gen_host_graph.cpp +++ b/src/a5/runtime/host_build_graph/host/dep_gen_host_graph.cpp @@ -18,9 +18,10 @@ * creator — creator retention on an existing tensor (STEP 3 Step A). * tensormap — a producer whose written slice overlaps what this task reads * (STEP 3 Step B); carries both slices' geometry. + * host_write_consumer — a direct consumer that a host write must wait for. * * Per-task producer dedup mirrors PTO2FaninBuilder::append_fanin_or_fail, which - * collapses all three sources into one fanin list: the first edge to name a + * collapses all four sources into one fanin list: the first edge to name a * producer is kept. tensormap edges are exempt — a second producer slice for the * same task is a distinct fact about the data flow, and viewers rely on seeing * every overlap. @@ -46,8 +47,8 @@ namespace { // Graph tables (serialized as tasks[] / tensors[] / edges[]) // --------------------------------------------------------------------------- -// Edge categories — matches the three places a runtime fanin edge is born. -enum class EdgeSource { EXPLICIT, CREATOR, TENSORMAP }; +// Edge categories — matches the four places a runtime fanin edge is born. +enum class EdgeSource { EXPLICIT, CREATOR, TENSORMAP, HOST_WRITE_CONSUMER }; const char *edge_source_str(EdgeSource s) { switch (s) { @@ -57,6 +58,8 @@ const char *edge_source_str(EdgeSource s) { return "creator"; case EdgeSource::TENSORMAP: return "tensormap"; + case EdgeSource::HOST_WRITE_CONSUMER: + return "host_write_consumer"; } return "unknown"; } @@ -73,6 +76,20 @@ const char *overlap_status_str(OverlapStatus s) { return "unknown"; } +const char *hazard_kind_str(TensorHazardKind kind) { + switch (kind) { + case TensorHazardKind::RAW: + return "RAW"; + case TensorHazardKind::WAW: + return "WAW"; + case TensorHazardKind::WAR: + return "WAR"; + } + return "unknown"; +} + +const char *access_kind_str(TensorAccessKind kind) { return kind == TensorAccessKind::READER ? "READER" : "WRITER"; } + const char *arg_type_str(TensorArgType t) { switch (t) { case TensorArgType::INPUT: @@ -83,6 +100,10 @@ const char *arg_type_str(TensorArgType t) { return "INOUT"; case TensorArgType::OUTPUT_EXISTING: return "OUTPUT_EXISTING"; + case TensorArgType::NO_DEP: + return "NO_DEP"; + case TensorArgType::TRACKED_INPUT: + return "TRACKED_INPUT"; } return "UNKNOWN"; } @@ -100,7 +121,9 @@ struct EdgeAnnot { int32_t consumer_arg_idx; // -1 for EXPLICIT (not tied to a tensor arg) EdgeSource source; OverlapStatus overlap; // only meaningful for TENSORMAP - uint64_t tensor_id; // 0 for EXPLICIT + TensorHazardKind hazard; + TensorAccessKind access_kind; + uint64_t tensor_id; // 0 for EXPLICIT // Consumer side (the ChipTensor the submitting task is reading). uint8_t consumer_dtype; uint32_t consumer_ndims; @@ -340,6 +363,8 @@ bool write_deps_json( out << ",\"source\":\"" << edge_source_str(e.source) << '"'; if (e.source == EdgeSource::TENSORMAP) { out << ",\"overlap\":\"" << overlap_status_str(e.overlap) << '"'; + out << ",\"hazard\":\"" << hazard_kind_str(e.hazard) << '"'; + out << ",\"access_kind\":\"" << access_kind_str(e.access_kind) << '"'; } if (e.source != EdgeSource::EXPLICIT) { out << ",\"tensor_id\":\"" << e.tensor_id << '"'; @@ -466,7 +491,7 @@ void dep_gen_host_graph_add_creator_edge(uint64_t producer_raw, int32_t arg_idx, void dep_gen_host_graph_add_tensormap_edge( uint64_t producer_raw, int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, - OverlapStatus overlap + OverlapStatus overlap, TensorHazardKind hazard ) { HostGraphState &s = state(); if (!s.enabled || !s.in_task) { @@ -481,12 +506,31 @@ void dep_gen_host_graph_add_tensormap_edge( e.consumer_arg_idx = arg_idx; e.source = EdgeSource::TENSORMAP; e.overlap = overlap; + e.hazard = hazard; + e.access_kind = entry.access_kind; e.tensor_id = make_tensor_id(entry.buffer_addr, entry.version); fill_consumer(e, consumer); fill_producer(e, entry); s.edges.push_back(e); } +void dep_gen_host_graph_add_host_write_consumer_edge( + uint64_t producer_raw, int32_t arg_idx, const ChipTensor &consumer +) { + HostGraphState &s = state(); + if (!s.enabled || !s.in_task || !s.task_preds.insert(producer_raw).second) { + return; + } + EdgeAnnot e{}; + e.pred = producer_raw; + e.succ = s.current_task_id; + e.consumer_arg_idx = arg_idx; + e.source = EdgeSource::HOST_WRITE_CONSUMER; + e.tensor_id = register_tensor(s, consumer); + fill_consumer(e, consumer); + s.edges.push_back(e); +} + // --------------------------------------------------------------------------- // Control surface // --------------------------------------------------------------------------- diff --git a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp index fc52dd4b9c..82c8c40a7a 100644 --- a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp @@ -596,6 +596,26 @@ int32_t run_host_orchestration( rt_scope_end(rt); rt_orchestration_done(rt); +#if SIMPLER_ORCH_PROFILING + const PTO2OrchProfilingData profile = orchestrator_get_profiling(); + const uint64_t profile_total = profile.sync_cycle + profile.alloc_cycle + profile.args_cycle + + profile.lookup_cycle + profile.insert_cycle + profile.fanin_cycle + + profile.scope_end_cycle; + const double profile_us = static_cast(profile_total) * 1000000.0 / PLATFORM_PROF_SYS_CNT_FREQ; + LOG_INFO( + "Host Orchestrator Profiling: %" PRId64 " tasks, total=%.3fus, avg/task=%.3fus", + static_cast(profile.submit_count), profile_us, + profile.submit_count > 0 ? profile_us / profile.submit_count : 0.0 + ); +#if SIMPLER_TENSORMAP_PROFILING + rt->orchestrator.tensor_map.print_stats(); + LOG_INFO( + "Host fanin spill: top=%d used=%d capacity=%d", rt->scheduler.fanin_spill_top, rt->scheduler.fanin_spill_top, + rt->scheduler.fanin_spill_capacity + ); +#endif +#endif + const int32_t total_tasks = pto2_sm_layout::ring_current_task_index_addr(host_sm)->load(std::memory_order_acquire); if (!upload_graph_submissions(runtime, api, *graph_state)) return -1; @@ -991,7 +1011,7 @@ extern "C" int bind_callable_to_runtime_impl( // on-device nodes that push past the host task count. const auto &sq = layout.sched; const size_t orch_start = layout.orch.off_fanin_seen_epoch; - const size_t orch_end = sq.off_ready_queue_slots[0]; + const size_t orch_end = sq.off_fanin_spill_ids; always_assert(orch_start <= orch_end); char *arena_host = static_cast(host_arena.base()); char *arena_dev = static_cast(runtime_arena_dev); diff --git a/src/a5/runtime/host_build_graph/orchestration/pto_arg_with_deps.h b/src/a5/runtime/host_build_graph/orchestration/pto_arg_with_deps.h index 536a6f8084..98ba192733 100644 --- a/src/a5/runtime/host_build_graph/orchestration/pto_arg_with_deps.h +++ b/src/a5/runtime/host_build_graph/orchestration/pto_arg_with_deps.h @@ -54,6 +54,7 @@ class CoreTaskArgsWithDeps : private CoreTaskArgs { using CoreTaskArgs::add_scalar; using CoreTaskArgs::add_scalars; using CoreTaskArgs::add_scalars_i32; + using CoreTaskArgs::add_tracked_input; using CoreTaskArgs::allow_early_resolve; // speculative early-dispatch hint (getter) using CoreTaskArgs::copy_scalars_from; using CoreTaskArgs::set_allow_early_resolve; // speculative early-dispatch hint (setter) diff --git a/src/a5/runtime/host_build_graph/runtime/dep_gen_host_graph.h b/src/a5/runtime/host_build_graph/runtime/dep_gen_host_graph.h index 3f86de4bd8..2720980355 100644 --- a/src/a5/runtime/host_build_graph/runtime/dep_gen_host_graph.h +++ b/src/a5/runtime/host_build_graph/runtime/dep_gen_host_graph.h @@ -28,6 +28,7 @@ * add_explicit_edge() — STEP 1, per declared dependency * add_creator_edge() — STEP 3 Step A, per creator-retention producer * add_tensormap_edge() — STEP 3 Step B, per tensormap producer + * add_host_write_consumer_edge() — runtime-derived host-write ordering * end_task() — closes the task, after its last dependency step * * Control surface, called from the device runner (same host_runtime.so): @@ -104,7 +105,12 @@ void dep_gen_host_graph_add_creator_edge(uint64_t producer_raw, int32_t arg_idx, /** STEP 3 Step B: a tensormap producer whose written slice this task reads. */ void dep_gen_host_graph_add_tensormap_edge( uint64_t producer_raw, int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, - OverlapStatus overlap + OverlapStatus overlap, TensorHazardKind hazard +); + +/** A direct consumer of an overlapping writer that must finish before a host write. */ +void dep_gen_host_graph_add_host_write_consumer_edge( + uint64_t producer_raw, int32_t arg_idx, const ChipTensor &consumer ); // --------------------------------------------------------------------------- diff --git a/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp b/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp index 54fc335dab..6e6f0f5834 100644 --- a/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp +++ b/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp @@ -73,8 +73,10 @@ __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_add_explicit __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_add_creator_edge(uint64_t, int32_t, const ChipTensor &) {} __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_add_tensormap_edge( - uint64_t, int32_t, const ChipTensor &, const PTO2TensorMapEntry &, OverlapStatus + uint64_t, int32_t, const ChipTensor &, const PTO2TensorMapEntry &, OverlapStatus, TensorHazardKind ) {} +__attribute__((weak, visibility("hidden"))) void +dep_gen_host_graph_add_host_write_consumer_edge(uint64_t, int32_t, const ChipTensor &) {} // Scope_stats enable gate, queried via the same predicate idiom as // dep_gen_host_graph_enabled above. The AICPU collector links the strong definition; @@ -687,10 +689,9 @@ static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { return next; } -// Polling: fanin is a flat array of position-independent producer local ids on -// the payload (no dep-pool spill, no producer pointers). The builder writes them -// directly into payload->fanin_local_ids as producers are appended, deduping by -// slot and hard-capping at PTO2_MAX_FANIN. self_local is this task's own local id +// Polling: fanin is a flat sequence of position-independent producer local ids. +// The builder writes the inline prefix to the payload and the remainder to the +// scheduler spill pool, deduping by slot. self_local is this task's own local id // (the consumer), used to bump each producer's last_consumer_local_id (the // reclaim gate the host wait_for_consumers polls via completed_watermark). struct PTO2FaninBuilder { @@ -699,7 +700,10 @@ struct PTO2FaninBuilder { orch(orch), seen_epoch(seen_epoch), self_local(self_local), - payload(payload) {} + payload(payload) { + payload->fanin_spill_start = 0; + payload->fanin_spill_count = 0; + } int32_t count{0}; PTO2OrchestratorState *orch{nullptr}; uint32_t seen_epoch{0}; @@ -735,11 +739,22 @@ static bool append_fanin_or_fail( if (fanin_builder->mark_seen(prod_ring, prod_slot)) { return true; } - if (fanin_builder->count >= PTO2_MAX_FANIN) { - orch_mark_fatal(orch, PTO2_ERROR_DEP_POOL_OVERFLOW); - return false; + const int32_t producer_local = static_cast(producer_task_id.local()); + if (fanin_builder->count < PTO2_MAX_FANIN) { + fanin_builder->payload->fanin_local_ids[fanin_builder->count] = producer_local; + } else { + PTO2SchedulerState *sched = orch->scheduler; + if (sched == nullptr || sched->fanin_spill_top >= sched->fanin_spill_capacity) { + orch_mark_fatal(orch, PTO2_ERROR_DEP_POOL_OVERFLOW); + return false; + } + if (fanin_builder->count == PTO2_MAX_FANIN) { + fanin_builder->payload->fanin_spill_start = sched->fanin_spill_top; + } + sched->fanin_spill_ids[sched->fanin_spill_top++] = producer_local; + fanin_builder->payload->fanin_spill_count++; } - fanin_builder->payload->fanin_local_ids[fanin_builder->count++] = static_cast(producer_task_id.local()); + fanin_builder->count++; // Reclaim gate: record this task as a consumer of the producer. The producer // slot retires once the per-ring completed_watermark reaches this consumer id. @@ -1058,7 +1073,7 @@ static bool ensure_tensormap_capacity(PTO2OrchestratorState *orch, int32_t neede return true; } // Progress is entries actually freed, NOT watermark movement: a ring can - // retire zero-output tasks (count_registrable_outputs == 0), advancing + // retire zero-access tasks (count_registrable_accesses == 0), advancing // last_task_alive without freeing any entry. Gating the backstop on // free_entries() keeps a wedged pool from dodging the timeout while some // unrelated ring keeps draining. @@ -1112,7 +1127,7 @@ static bool ensure_tensormap_capacity(PTO2OrchestratorState *orch, int32_t neede // Orch-side wiring/ready publication. static TaskOutputTensors submit_task_common( PTO2OrchestratorState *orch, const CoreTaskArgs &args, ActiveMask active_mask, TaskAttrs task_attrs, - int32_t aic_kernel_id, int32_t aiv0_kernel_id, int32_t aiv1_kernel_id + int32_t aic_kernel_id, int32_t aiv0_kernel_id, int32_t aiv1_kernel_id, TaskKind task_kind = TaskKind::KERNEL ) { CYCLE_COUNT_START(); TaskOutputTensors result; @@ -1127,6 +1142,7 @@ static TaskOutputTensors submit_task_common( PTO2TaskDescriptor &task = *prepared.task; PTO2TaskPayload &payload = *prepared.payload; result.set_task_id(task_id); + prepared.slot_state->task_kind = task_kind; // dep_gen capture point: open this task's graph entry before its dependency // steps run, so the edges STEP 1 / STEP 3 discover attach to it. The graph @@ -1192,6 +1208,74 @@ static TaskOutputTensors submit_task_common( } } + // A graph-local host write preserves the synchronous TMR contract: wait for + // each overlapping writer and every direct consumer of that writer. The + // normal OUTPUT_EXISTING lookup below adds tracked readers. HBG has the + // complete task graph resident, so direct consumers are recoverable by + // scanning earlier payload fanins without serializing unrelated tasks. + if (task_kind == TaskKind::HOST_WRITE) { + const int32_t self_local = static_cast(task_id.local()); + PTO2SharedMemoryRingHeader &producer_ring = orch->sm_header->ring; + const ChipTensor &write_tensor = args.tensor(0).ref(); + auto append_predecessor = [&](PTO2TaskId predecessor) -> bool { + const int32_t predecessor_local = static_cast(predecessor.local()); + const int32_t predecessor_slot = producer_ring.get_slot_by_task_id(predecessor_local); + PTO2TaskSlotState *predecessor_state = &producer_ring.get_slot_state_by_slot(predecessor_slot); + return append_fanin_or_fail( + orch, predecessor.ring(), predecessor_slot, predecessor_state, predecessor, &fanin_builder + ); + }; + auto append_consumers = [&](PTO2TaskId writer) -> bool { + const int32_t writer_local = static_cast(writer.local()); + for (int32_t consumer_local = 0; consumer_local < self_local; ++consumer_local) { + const int32_t consumer_slot = producer_ring.get_slot_by_task_id(consumer_local); + PTO2TaskSlotState *consumer_state = &producer_ring.get_slot_state_by_slot(consumer_slot); + if (consumer_state->payload == nullptr || consumer_state->task == nullptr || + consumer_state->task->task_id.local() != static_cast(consumer_local)) { + continue; + } + const PTO2TaskPayload &consumer_payload = *consumer_state->payload; + bool reads_writer = false; + for (int32_t i = 0; i < consumer_payload.fanin_count; ++i) { + if (orch->scheduler->fanin_local_id(consumer_payload, i) == writer_local) { + reads_writer = true; + break; + } + } + if (!reads_writer) continue; + const PTO2TaskId consumer = PTO2TaskId::make(task_id.ring(), static_cast(consumer_local)); + if (!append_predecessor(consumer)) return false; + if (capture_dep_graph) { + dep_gen_host_graph_add_host_write_consumer_edge(consumer.raw, 0, write_tensor); + } + } + return true; + }; + + if (write_tensor.owner_task_id.is_valid() && !append_consumers(write_tensor.owner_task_id)) { + return result; + } + bool host_write_fatal = false; + orch->tensor_map.lookup( + write_tensor, TensorAccessKind::WRITER, + [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { + if (!append_predecessor(entry.access_task_id) || !append_consumers(entry.access_task_id)) { + host_write_fatal = true; + return false; + } + if (capture_dep_graph) { + dep_gen_host_graph_add_tensormap_edge( + entry.access_task_id.raw, 0, write_tensor, entry, overlap_status, TensorHazardKind::WAW + ); + } + return true; + } + ); + if (host_write_fatal) { + return result; + } + } + // === STEP 3: Lookup inputs (creator retention + tensormap modifier lookup) === DepInputs dep_inputs{ args.tensor_count(), args.tensor_data(), args.tag_data(), static_cast(args.explicit_dep_count()), @@ -1214,9 +1298,12 @@ static TaskOutputTensors submit_task_common( dep_gen_host_graph_add_creator_edge(producer.raw, arg_idx, consumer); } void tensormap( - int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, OverlapStatus overlap + int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, OverlapStatus overlap, + TensorHazardKind hazard ) const { - dep_gen_host_graph_add_tensormap_edge(entry.producer_task_id.raw, arg_idx, consumer, entry, overlap); + dep_gen_host_graph_add_tensormap_edge( + entry.access_task_id.raw, arg_idx, consumer, entry, overlap, hazard + ); } }; const bool ok = @@ -1240,11 +1327,11 @@ static TaskOutputTensors submit_task_common( // is reclaimed as last_task_alive advances; an // exhausted pool back-pressures here (and detects a wedged watermark) rather // than tripping new_entry()'s hard assert mid-registration. - int32_t tensormap_needed = count_registrable_outputs(dep_inputs, orch->in_manual_scope()); + int32_t tensormap_needed = count_registrable_accesses(dep_inputs, orch->in_manual_scope()); if (tensormap_needed > 0 && !ensure_tensormap_capacity(orch, tensormap_needed)) { return result; } - register_task_outputs(dep_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); + register_task_accesses(dep_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); CYCLE_COUNT_LAP(g_orch_insert_cycle); @@ -1450,7 +1537,7 @@ bool graph_submit_definition( DepInputs boundary_inputs{ args.tensor_count(), args.tensor_data(), args.tag_data(), 0, nullptr, }; - const int32_t tensormap_needed = count_registrable_outputs(boundary_inputs, orch->in_manual_scope()); + const int32_t tensormap_needed = count_registrable_accesses(boundary_inputs, orch->in_manual_scope()); if (tensormap_needed > 0 && !ensure_tensormap_capacity(orch, tensormap_needed)) return false; if (!check_scope_can_accept_task(orch, allocator, 0)) return false; @@ -1491,7 +1578,7 @@ bool graph_submit_definition( return append_fanin_or_fail(orch, producer_id.ring(), producer_slot, producer, producer_id, &fanin_builder); }; if (!compute_task_fanin(boundary_inputs, orch->tensor_map, orch->in_manual_scope(), emit)) return false; - register_task_outputs(boundary_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); + register_task_accesses(boundary_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); payload.fanin_count = fanin_builder.count; pending.outer_slot = &slot; @@ -1530,6 +1617,14 @@ TaskOutputTensors graph_record_submit_node( if (node_index >= GRAPH_MAX_NODES || args.has_error || args.predicate().op != PredicateOp::NONE) { recording.unsupported = true; } + for (int32_t i = 0; i < args.tensor_count(); ++i) { + if (args.tag(i) == TensorArgType::TRACKED_INPUT) { + // The compact Graph format derives only producer-source fanins; it + // cannot publish an internal reader for a later WAR lookup. + recording.unsupported = true; + break; + } + } const PTO2OutputLayout layout = calculate_output_layout(args); void *packed_base = orch->ring.task_allocator.reserve_heap_scratch(layout.total_output_size); @@ -1733,7 +1828,6 @@ bool PTO2OrchestratorState::graph_end() { std::vector definition; if (!graph_build_definition(*recording, &definition)) { - debug_assert(false && "The recorded Graph contains a construct that Graph Execution does not support"); LOG_WARN("%s", "[GraphExecution] unsupported construct observed; falling back to the ordinary path"); return false; } @@ -1890,7 +1984,38 @@ TaskOutputTensors PTO2OrchestratorState::submit_dummy_task(const CoreTaskArgs &a } return submit_task_common( - orch, args, ActiveMask{}, task_attrs, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID + orch, args, ActiveMask{}, task_attrs, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID, TaskKind::DUMMY + ); +} + +void PTO2OrchestratorState::submit_host_write( + const ChipTensor &tensor, uint64_t address, uint64_t value, uint64_t size +) { + if (fatal) return; + if (size == 0 || size > sizeof(value)) { + report_fatal( + PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "host write size=%llu is outside [1,8]", (unsigned long long)size + ); + return; + } + + CoreTaskArgs args; + args.add_output(tensor); // OUTPUT_EXISTING: WAR lookup, then writer registration. + args.add_scalar(address); + args.add_scalar(value); + args.add_scalar(size); + + if (GraphHostState *state = graph_state_from(this); state != nullptr && state->recording != nullptr) { + // The cached Graph format has no host-write node. Preserve recording + // task-id arithmetic, then force graph_end to replay the ordinary path. + (void)submit_dummy_task(args); + graph_record_mark_unsupported(this); + return; + } + + (void)submit_task_common( + this, args, ActiveMask{}, TaskAttrs{}, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID, + TaskKind::HOST_WRITE ); } diff --git a/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp b/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp index c122e8562a..d4ef32271b 100644 --- a/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp +++ b/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp @@ -154,6 +154,7 @@ wait_for_tensor_ready(PTO2Runtime *rt, const ChipTensor &tensor, bool wait_for_c // the second encounter. constexpr int kSegmentCap = 64; const PTO2TaskSlotState *seg[kSegmentCap]; + bool seg_wait_consumers[kSegmentCap]{}; int seg_count = 0; bool failed = false; @@ -250,22 +251,26 @@ wait_for_tensor_ready(PTO2Runtime *rt, const ChipTensor &tensor, bool wait_for_c for (int i = 0; i < seg_count; i++) { wait_one_producer(*seg[i]); if (failed) return; - if (!wait_for_consumers) continue; + if (!seg_wait_consumers[i]) continue; wait_one_consumers(*seg[i]); if (failed) return; } seg_count = 0; }; - auto try_push = [&](const PTO2TaskSlotState &s) { + auto try_push = [&](const PTO2TaskSlotState &s, bool include_consumers) { for (int j = 0; j < seg_count; j++) { - if (seg[j] == &s) return; // per-segment dedup + if (seg[j] == &s) { + seg_wait_consumers[j] = seg_wait_consumers[j] || include_consumers; + return; + } } if (seg_count == kSegmentCap) { flush_segment(); if (failed) return; } - seg[seg_count++] = &s; + seg[seg_count] = &s; + seg_wait_consumers[seg_count++] = include_consumers; }; auto do_wait = [&]() { @@ -273,17 +278,27 @@ wait_for_tensor_ready(PTO2Runtime *rt, const ChipTensor &tensor, bool wait_for_c if (owner.is_valid()) { const auto *slot = resolve_producer(owner); if (slot == nullptr) return; - try_push(*slot); + try_push(*slot, wait_for_consumers); } // Step B: modifier writer lookup (OverlapMap), direct callback orch.tensor_map.lookup(tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus) -> bool { - PTO2TaskId pid = entry.producer_task_id; + PTO2TaskId pid = entry.access_task_id; const auto *slot = resolve_producer(pid); if (slot == nullptr) return false; - try_push(*slot); + try_push(*slot, wait_for_consumers); return !failed; }); + if (wait_for_consumers && !failed) { + orch.tensor_map.lookup( + tensor, TensorAccessKind::READER, [&](PTO2TensorMapEntry &entry, OverlapStatus) -> bool { + const auto *slot = resolve_producer(entry.access_task_id); + if (slot == nullptr) return false; + try_push(*slot, false); + return !failed; + } + ); + } if (failed) return; flush_segment(); }; @@ -333,22 +348,24 @@ void set_tensor_data( return; } - // Wait for producer + all consumers before writing (WAW + WAR safety) - if (!wait_for_tensor_ready(rt, tensor, true, __FUNCTION__)) { - return; - } - uint64_t flat_offset = tensor.compute_flat_offset(indices, ndims); uint64_t elem_size = get_element_size(tensor.dtype); uint64_t elem_addr = tensor.buffer.addr + flat_offset * elem_size; - if (!host_tensor_write(rt->tensor_access, elem_addr, &value, elem_size)) { - rt->orchestrator.report_fatal( - PTO2_ERROR_INVALID_ARGS, __FUNCTION__, - "no writable host view for device address %#llx (%llu bytes): during host orchestration only tensors " - "the runtime staged are writable, not runtime-created or child-memory buffers", - (unsigned long long)elem_addr, (unsigned long long)elem_size - ); + if (rt->orchestrator.sm_header->ring.fc.current_task_index.load(std::memory_order_acquire) == 0) { + if (!host_tensor_write(rt->tensor_access, elem_addr, &value, elem_size)) { + rt->orchestrator.report_fatal( + PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "no writable host view for device address %#llx (%llu bytes)", + (unsigned long long)elem_addr, (unsigned long long)elem_size + ); + } + return; } + // HBG constructs the whole graph before device scheduling starts, so a + // synchronous host wait here can never observe a submitted reader finish. + // Represent the write as a scheduler-local graph node instead: normal + // OUTPUT_EXISTING dependency discovery gives it WAR predecessors and + // registers it as the next writer for tasks submitted afterwards. + rt->orchestrator.submit_host_write(tensor, elem_addr, value, elem_size); } // Ops-table entry that hands the call-site captured by PTO2ScopeGuard to the diff --git a/src/a5/runtime/host_build_graph/runtime/pto_dep_compute.h b/src/a5/runtime/host_build_graph/runtime/pto_dep_compute.h index 1ee793d14c..09cb860562 100644 --- a/src/a5/runtime/host_build_graph/runtime/pto_dep_compute.h +++ b/src/a5/runtime/host_build_graph/runtime/pto_dep_compute.h @@ -16,11 +16,11 @@ * Two header-only template entry points: * * compute_task_fanin — STEP 3 in submit_task: per-tensor creator retention (Step A) - * + tensormap.lookup for INPUT/INOUT (Step B). Calls back into - * user-supplied `emit` for each producer it identifies. + * + TensorMap lookup for INPUT/TRACKED_INPUT/INOUT (Step B). + * Calls back into user-supplied `emit` for each producer. * - * register_task_outputs — STEP 4 in submit_task: tensormap.insert for INOUT and - * OUTPUT_EXISTING tensors. No callbacks. + * register_task_accesses — STEP 4 in submit_task: register TRACKED_INPUT readers and + * INOUT/OUTPUT_EXISTING writers. No callbacks. * * STEP 1 (explicit_deps) is intentionally left at the runtime call site because its * `last_task_alive` shortcut + unchecked slot lookup is subtly different from the @@ -59,7 +59,7 @@ #include "tensor.h" /** - * View struct for inputs to compute_task_fanin / register_task_outputs. + * View struct for inputs to compute_task_fanin / register_task_accesses. * * Both runtime and replay assemble one of these from their own data sources * (runtime: from Arg accessors; replay: from SubmitTraceEntry fields). All @@ -79,7 +79,7 @@ struct DepInputs { */ struct NoDepAnnotate { void creator(int32_t, const ChipTensor &, PTO2TaskId) const {} - void tensormap(int32_t, const ChipTensor &, const PTO2TensorMapEntry &, OverlapStatus) const {} + void tensormap(int32_t, const ChipTensor &, const PTO2TensorMapEntry &, OverlapStatus, TensorHazardKind) const {} }; /** @@ -88,8 +88,8 @@ struct NoDepAnnotate { * * For each non-OUTPUT tensor: * - If owner_task_id is valid, emit(owner) - * - For INPUT/INOUT (and not manual_dep), tensor_map.lookup(*tensor) and emit - * each matching producer. INOUT+COVERED triggers tensor_map.remove_entry(entry). + * - For INPUT/TRACKED_INPUT/INOUT (and not manual_dep), query overlapping writers. + * - For INOUT/OUTPUT_EXISTING, query overlapping tracked readers. * * @return true on success (or producer-skipped-silently); false if emit signaled * fatal — caller should propagate (after any fatal bookkeeping done by emit). @@ -121,26 +121,37 @@ template annotate.creator(i, *tensor, owner); } - // Step B: only INPUT/INOUT need modifier dependency lookup. - if (ptype != TensorArgType::INPUT && ptype != TensorArgType::INOUT) { - continue; - } if (tensor->manual_dep) { continue; } bool fatal = false; - tensor_map.lookup(*tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { - if (!emit(entry.producer_task_id)) { - fatal = true; - return false; // stop iteration - } - annotate.tensormap(i, *tensor, entry, overlap_status); - if (ptype == TensorArgType::INOUT && overlap_status == OverlapStatus::COVERED) { - tensor_map.remove_entry(entry); - } - return true; - }); + auto lookup = [&](TensorAccessKind access_kind, TensorHazardKind hazard) { + tensor_map.lookup( + *tensor, access_kind, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { + if (!emit(entry.access_task_id)) { + fatal = true; + return false; // stop iteration + } + annotate.tensormap(i, *tensor, entry, overlap_status, hazard); + if (ptype == TensorArgType::INOUT && access_kind == TensorAccessKind::WRITER) { + annotate.tensormap(i, *tensor, entry, overlap_status, TensorHazardKind::WAW); + } + if ((ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) && + overlap_status == OverlapStatus::COVERED) { + tensor_map.remove_entry(entry); + } + return true; + } + ); + }; + if (ptype == TensorArgType::INPUT || ptype == TensorArgType::TRACKED_INPUT || ptype == TensorArgType::INOUT) { + lookup(TensorAccessKind::WRITER, TensorHazardKind::RAW); + } + if (!fatal && tensor_map.current_readers() != 0 && + (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING)) { + lookup(TensorAccessKind::READER, TensorHazardKind::WAR); + } if (fatal) { return false; } @@ -149,45 +160,47 @@ template } /** - * Register a task's outputs in the tensormap (STEP 4 in submit_task). - * - * For INOUT and OUTPUT_EXISTING tensors (excluding manual_dep), inserts the - * tensor into tensor_map keyed by its buffer.addr with `task_id` as producer. + * Register a task's accesses in the tensormap (STEP 4 in submit_task). * * No-op when in_manual_scope. */ inline void -register_task_outputs(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap &tensor_map, bool in_manual_scope) { +register_task_accesses(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap &tensor_map, bool in_manual_scope) { if (in_manual_scope) { return; } for (int32_t i = 0; i < inputs.tensor_count; i++) { TensorArgType ptype = inputs.arg_types[i]; - if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) { + if (ptype == TensorArgType::TRACKED_INPUT || ptype == TensorArgType::INOUT || + ptype == TensorArgType::OUTPUT_EXISTING) { const ChipTensor *tensor = &inputs.tensors[i].ref(); if (!tensor->manual_dep) { - tensor_map.insert(*tensor, task_id); + if (ptype == TensorArgType::TRACKED_INPUT) { + tensor_map.insert(*tensor, task_id, TensorAccessKind::READER); + } else { + tensor_map.insert(*tensor, task_id); + } } } } } /** - * Count the tensormap entries register_task_outputs() will insert for this task. + * Count the tensormap entries register_task_accesses() will insert for this task. * - * Mirrors register_task_outputs()'s selection exactly (INOUT / OUTPUT_EXISTING, - * excluding manual_dep), so the returned value is the precise number of + * Mirrors register_task_accesses() exactly, so the returned value is the precise number of * new_entry() calls that step makes. The orchestrator uses it to reserve pool * capacity before inserting. Returns 0 in a manual scope (no registration). */ -inline int32_t count_registrable_outputs(const DepInputs &inputs, bool in_manual_scope) { +inline int32_t count_registrable_accesses(const DepInputs &inputs, bool in_manual_scope) { if (in_manual_scope) { return 0; } int32_t needed = 0; for (int32_t i = 0; i < inputs.tensor_count; i++) { TensorArgType ptype = inputs.arg_types[i]; - if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) { + if (ptype == TensorArgType::TRACKED_INPUT || ptype == TensorArgType::INOUT || + ptype == TensorArgType::OUTPUT_EXISTING) { if (!inputs.tensors[i].ref().manual_dep) { needed++; } diff --git a/src/a5/runtime/host_build_graph/runtime/pto_orchestrator.h b/src/a5/runtime/host_build_graph/runtime/pto_orchestrator.h index aec93828b5..ae55880329 100644 --- a/src/a5/runtime/host_build_graph/runtime/pto_orchestrator.h +++ b/src/a5/runtime/host_build_graph/runtime/pto_orchestrator.h @@ -9,7 +9,7 @@ * ----------------------------------------------------------------------------------------------------------- */ /** - * PTO Runtime2 - Orchestrator Interface + * host_build_graph - Orchestrator Interface * * The Orchestrator is responsible for: * 1. Executing the orchestration function (Turing-complete control flow) @@ -25,8 +25,7 @@ * Based on: docs/RUNTIME_LOGIC.md */ -#ifndef PTO_ORCHESTRATOR_H -#define PTO_ORCHESTRATOR_H +#pragma once #include "common/chip_swimlane_profiling.h" #include "utils/device_arena.h" @@ -160,6 +159,7 @@ struct PTO2OrchestratorState { void end_scope(); TaskOutputTensors submit_task(const MixedKernels &mixed_kernels, const CoreTaskArgs &args); TaskOutputTensors submit_dummy_task(const CoreTaskArgs &args); + void submit_host_write(const ChipTensor &tensor, uint64_t address, uint64_t value, uint64_t size); TaskOutputTensors alloc_tensors(const CoreTaskArgs &args); GraphScopeResult graph_begin(uint64_t graph_key, const CoreTaskArgs &args, uint64_t callable_hash); bool graph_end(); @@ -192,5 +192,3 @@ struct PTO2OrchProfilingData { PTO2OrchProfilingData orchestrator_get_profiling(); #endif - -#endif // PTO_ORCHESTRATOR_H diff --git a/src/a5/runtime/host_build_graph/runtime/pto_runtime2_types.h b/src/a5/runtime/host_build_graph/runtime/pto_runtime2_types.h index 32c0b23a65..4a26be452f 100644 --- a/src/a5/runtime/host_build_graph/runtime/pto_runtime2_types.h +++ b/src/a5/runtime/host_build_graph/runtime/pto_runtime2_types.h @@ -10,9 +10,9 @@ */ /** - * PTO Runtime2 - Core Type Definitions + * host_build_graph - Core Type Definitions * - * This header defines all fundamental types used by the PTO Runtime2 system: + * This header defines the fundamental host_build_graph runtime types: * - Configuration constants * - Worker types and task states * - ChipTensor regions and task parameters @@ -84,6 +84,7 @@ #define PTO2_HEAP_SIZE (256 * 1024 * 1024) // 256MB #define PTO2_TENSORMAP_POOL_SIZE (65536) // TensorMap entry pool #define PTO2_TENSORMAP_NUM_BUCKETS 4096 // Power of 2 for fast hash (4096×8B=32KB fits L1) +#define PTO2_TENSORMAP_READER_NUM_BUCKETS 512 // Sparse tracked-reader sidecar (512×8B=4KB) // Scope management #define PTO2_MAX_SCOPE_DEPTH 64 // Maximum nesting depth @@ -108,11 +109,10 @@ // Fanin storage #define PTO2_FANIN_INLINE_CAP 64 -// Polling-scheduler inline fanin cap. The polling model stores producer -// dependencies as flat position-independent local-id integers on the payload -// (no dep-pool spill), so a task's fanin degree is hard-capped here. Must cover -// the worst-case fanin of any workload (paged_attention is the densest). +// Polling-scheduler inline fanin cap. Additional position-independent local ids +// continue in the scheduler-owned spill pool. #define PTO2_MAX_FANIN 128 +#define HBG_FANIN_SPILL_POOL_SIZE 262144 // Dependency-degree diagnostic: warn once when a task's fanin or a producer's // fanout first exceeds this degree, so dense dependency graphs surface without @@ -175,6 +175,7 @@ enum class TaskKind : uint8_t { DUMMY = 1, GRAPH = 2, GRAPH_NODE = 3, + HOST_WRITE = 4, }; struct PTO2OutputLayout { @@ -269,12 +270,14 @@ struct PTO2TaskPayload { // Producer dependencies as position-independent local task ids. Single-ring // hbg: every producer is ring 0, so no per-edge ring id is stored. Scanned // by fanin_satisfied / classify_fanin_state against the ring completion_flags. - // Hard-capped at PTO2_MAX_FANIN (no dep-pool spill). + // The first PTO2_MAX_FANIN ids are inline; additional ids use fanin_spill_*. int32_t fanin_local_ids[PTO2_MAX_FANIN]; // Reserved: preserves the early-dispatch block and tensors[] offsets. tensors // must stay at byte 576 (AICore arg-materialization contract), so this fanin // region keeps its original 528-byte footprint. - int32_t _fanin_reserved[3]; + int32_t fanin_spill_start{0}; + int32_t fanin_spill_count{0}; + int32_t _fanin_reserved{0}; // Early-dispatch metadata (AICPU-side only). Ordered by descending // alignment (8B mask, 4B fanin, then 2B/1B counters and flags) so the block packs with no // internal padding. Kept here after the fanin array (not moved up front): on diff --git a/src/a5/runtime/host_build_graph/runtime/pto_tensormap.h b/src/a5/runtime/host_build_graph/runtime/pto_tensormap.h index e925337307..9db9388940 100644 --- a/src/a5/runtime/host_build_graph/runtime/pto_tensormap.h +++ b/src/a5/runtime/host_build_graph/runtime/pto_tensormap.h @@ -10,10 +10,10 @@ */ /** - * PTO Runtime2 - TensorMap Interface + * host_build_graph - TensorMap Interface * - * TensorMap provides producer lookup for dependency discovery: - * - Maps ChipTensor -> producer task ID + * TensorMap provides access lookup for dependency discovery: + * - Maps ChipTensor -> reader/writer task IDs * - Used by pto_submit_task() to find dependencies * * Key design features: @@ -57,6 +57,17 @@ enum class OverlapStatus { OTHER, }; +enum class TensorAccessKind : uint8_t { + READER = 0, + WRITER = 1, +}; + +enum class TensorHazardKind : uint8_t { + RAW = 0, + WAW = 1, + WAR = 2, +}; + struct Segment { uint64_t begin; uint64_t end; @@ -74,10 +85,12 @@ struct Segment { */ struct PTO2TensorMapLayout { size_t off_buckets; + size_t off_reader_buckets; size_t off_entry_pool; size_t off_free_entry_list; size_t off_task_entry_heads; int32_t num_buckets; + int32_t num_reader_buckets; int32_t pool_size; int32_t task_window_size; }; @@ -106,7 +119,7 @@ extern uint64_t g_insert_count; * the hash key, size in [8, 16) is unused by the entry — we repurpose it for * `next_in_bucket`). * - * buffer_addr / next_in_bucket / producer_task_id — chain traversal + match + * buffer_addr / next_in_bucket / access_task_id — chain traversal + match * start_offset — overlap byte range begin * version, ndims, dtype, manual_dep, is_contiguous — overlap fast path * shapes[5] — overlap comparison (line 1) @@ -126,15 +139,15 @@ struct alignas(64) PTO2TensorMapEntry { // === Cache line 1 (64B) — lookup hot path; mirrors ChipTensor line 1 from byte 16 === uint64_t buffer_addr; // 8B [0, 8): tensor base address (hash key, mirrors ChipTensor::buffer.addr) PTO2TensorMapEntry - *next_in_bucket; // 8B [8, 16): next entry in hash bucket chain (overlays ChipTensor::buffer.size) - PTO2TaskId producer_task_id; // 8B [16,24): mirrors ChipTensor::owner_task_id slot - uint64_t start_offset; // 8B [24,32): mirrors ChipTensor::start_offset (element offset) - int32_t version; // 4B [32,36): mirrors ChipTensor::version - uint32_t ndims; // 4B [36,40): mirrors ChipTensor::ndims - DataType dtype; // 1B [40,41): mirrors ChipTensor::dtype - bool manual_dep; // 1B [41,42): mirrors ChipTensor::manual_dep - bool is_contiguous; // 1B [42,43): mirrors ChipTensor::is_contiguous - uint8_t __padding1__; // 1B [43,44): mirrors ChipTensor padding + *next_in_bucket; // 8B [8, 16): next entry in hash bucket chain (overlays ChipTensor::buffer.size) + PTO2TaskId access_task_id; // 8B [16,24): task owning this access + uint64_t start_offset; // 8B [24,32): mirrors ChipTensor::start_offset (element offset) + int32_t version; // 4B [32,36): mirrors ChipTensor::version + uint32_t ndims; // 4B [36,40): mirrors ChipTensor::ndims + DataType dtype; // 1B [40,41): mirrors ChipTensor::dtype + bool manual_dep; // 1B [41,42): mirrors ChipTensor::manual_dep + bool is_contiguous; // 1B [42,43): mirrors ChipTensor::is_contiguous + TensorAccessKind access_kind; // 1B [43,44): reader or writer index uint32_t shapes[MAX_TENSOR_DIMS]; // 20B [44,64): mirrors ChipTensor::shapes // === Cache line 2 (64B) — chain manipulation + non-contiguous overlap data === @@ -151,7 +164,7 @@ struct alignas(64) PTO2TensorMapEntry { * Copy overlap-relevant fields from a ChipTensor into this entry. * * 64B memcpy of ChipTensor cache line 1 populates buffer_addr (byte [0,8)), - * producer_task_id, start_offset, version, ndims, dtype, manual_dep, + * access_task_id, start_offset, version, ndims, dtype, manual_dep, * is_contiguous and shapes[]. Byte [8,16) holds ChipTensor::buffer.size in * the source and gets written into next_in_bucket; that's harmless * because link_entry() overwrites next_in_bucket immediately after. @@ -336,7 +349,7 @@ struct alignas(64) PTO2TensorMapEntry { static_assert(sizeof(PTO2TensorMapEntry) == 128, "TensorMapEntry must be exactly 2 cache lines (128 bytes)"); static_assert(offsetof(PTO2TensorMapEntry, buffer_addr) == offsetof(ChipTensor, buffer.addr)); -static_assert(offsetof(PTO2TensorMapEntry, producer_task_id) == offsetof(ChipTensor, owner_task_id)); +static_assert(offsetof(PTO2TensorMapEntry, access_task_id) == offsetof(ChipTensor, owner_task_id)); static_assert(offsetof(PTO2TensorMapEntry, start_offset) == offsetof(ChipTensor, start_offset)); static_assert(offsetof(PTO2TensorMapEntry, version) == offsetof(ChipTensor, version)); static_assert(offsetof(PTO2TensorMapEntry, ndims) == offsetof(ChipTensor, ndims)); @@ -360,7 +373,7 @@ static_assert( struct PTO2TensorMap { // Hash table buckets (fixed size, power of 2) PTO2TensorMapEntry **buckets; // Array of offsets into entry_pool (-1 = empty) - int32_t num_buckets; // Must be power of 2 for fast modulo + int32_t num_buckets; // Must be power of 2 for fast modulo // Entry pool as ring buffer PTO2TensorMapEntry *entry_pool; // Ring buffer of entries @@ -380,6 +393,13 @@ struct PTO2TensorMap { // Cleanup progress (for periodic cleanup_retired) int32_t last_cleanup{}; + // Tracked-reader side index. Appended to preserve the original writer + // fields' offsets and cache locality. + PTO2TensorMapEntry **reader_buckets; + int32_t num_reader_buckets; + int32_t reader_used{0}; + int32_t reader_high_water{0}; + uint32_t get_task_local_id_slot(uint32_t task_local_id) const { return task_local_id & (task_window_size - 1); } // Accessors read by scope_stats_collector. Declared unconditionally so the @@ -387,6 +407,8 @@ struct PTO2TensorMap { // setter symbols must export for host dlsym; the probe call sites that use // these accessors stay gated by SIMPLER_DFX). int32_t current_used() const { return next_entry_idx - free_num; } + int32_t current_readers() const { return reader_used; } + int32_t current_writers() const { return current_used() - reader_used; } int32_t pool_capacity() const { return pool_size; } int32_t free_entries() const { return pool_size - current_used(); } @@ -407,7 +429,7 @@ struct PTO2TensorMap { // new_entry allocates a slot and initializes only its linkage (bucket_index // and the four link pointers) to the clean unlinked state; insert() assigns - // the tensor attributes and producer_task_id. + // the tensor attributes and access_task_id. PTO2TensorMapEntry *new_entry() { if (free_num > 0) { PTO2TensorMapEntry *res = free_entry_list[--free_num]; @@ -436,7 +458,8 @@ struct PTO2TensorMap { if (entry.prev_in_bucket == nullptr) { // Entry is the head of its bucket chain, update bucket head // Must compute hash BEFORE clearing tensor - buckets[entry.bucket_index] = entry.next_in_bucket; + PTO2TensorMapEntry **index = entry.access_kind == TensorAccessKind::READER ? reader_buckets : buckets; + index[entry.bucket_index] = entry.next_in_bucket; } else { entry.prev_in_bucket->next_in_bucket = entry.next_in_bucket; } @@ -446,6 +469,10 @@ struct PTO2TensorMap { entry.next_in_bucket->prev_in_bucket = entry.prev_in_bucket; } + if (entry.access_kind == TensorAccessKind::READER) { + always_assert(reader_used > 0); + reader_used--; + } free_entry_list[free_num++] = &entry; entry.bucket_index = -1; entry.next_in_bucket = nullptr; @@ -517,10 +544,13 @@ struct PTO2TensorMap { * @param tensor ChipTensor to look up * @param on_match Callback invoked for each overlapping entry */ - template - void lookup(const ChipTensor &tensor, Fn &&on_match) { - uint32_t bucket_index = hash(tensor.buffer.addr); - PTO2TensorMapEntry *cur_entry = buckets[bucket_index]; + template + void lookup_impl(const ChipTensor &tensor, Fn &&on_match) { + constexpr bool is_reader = access_kind == TensorAccessKind::READER; + int32_t bucket_count = is_reader ? num_reader_buckets : num_buckets; + uint32_t bucket_index = hash(tensor.buffer.addr, bucket_count); + PTO2TensorMapEntry **index = is_reader ? reader_buckets : buckets; + PTO2TensorMapEntry *cur_entry = index[bucket_index]; #if SIMPLER_TENSORMAP_PROFILING g_lookup_count++; @@ -572,19 +602,39 @@ struct PTO2TensorMap { #endif } + template + void lookup(const ChipTensor &tensor, Fn &&on_match) { + lookup_impl(tensor, static_cast(on_match)); + } + + template + void lookup(const ChipTensor &tensor, TensorAccessKind access_kind, Fn &&on_match) { + if (access_kind == TensorAccessKind::READER) { + lookup_impl(tensor, static_cast(on_match)); + } else { + lookup_impl(tensor, static_cast(on_match)); + } + } + /** - * Insert a new entry (called when task produces output) + * Insert a new reader or writer access entry. * * Allocates from ring buffer pool, may overwrite stale entries. * Inserts at head of hash bucket chain (maintains task_id ordering). * - * @param tensor ChipTensor produced - * @param producer_task_id Task ID of producer + * @param tensor ChipTensor accessed + * @param access_task_id Task ID owning the access */ - void insert(const ChipTensor &tensor, PTO2TaskId producer_task_id) { - PTO2TensorMapEntry *entry = new_entry(); - entry->copy_from_tensor(tensor); - link_entry(entry, tensor.buffer.addr, producer_task_id); + void insert(const ChipTensor &tensor, PTO2TaskId access_task_id, TensorAccessKind access_kind) { + if (access_kind == TensorAccessKind::READER) { + insert_impl(tensor, access_task_id); + } else { + insert_impl(tensor, access_task_id); + } + } + + void insert(const ChipTensor &tensor, PTO2TaskId access_task_id) { + insert_impl(tensor, access_task_id); } /** @@ -608,7 +658,7 @@ struct PTO2TensorMap { PTO2TensorMapEntry *cur_entry = task_entry_heads[task_slot]; while (cur_entry != nullptr) { PTO2TensorMapEntry *next_entry = cur_entry->next_in_task; // free_entry clears it - if (cur_entry->producer_task_id == retired_task) { + if (cur_entry->access_task_id == retired_task) { if (cur_entry->prev_in_task != nullptr) { cur_entry->prev_in_task->next_in_task = next_entry; } else { @@ -636,31 +686,42 @@ struct PTO2TensorMap { * addresses (low bits all-zero) still distribute evenly. We extract * the top log2(num_buckets) bits which carry the most entropy. */ - uint32_t hash(uint64_t key) { + uint32_t hash(uint64_t key, int32_t bucket_count) { key *= 0x9E3779B97F4A7C15ULL; - return static_cast(key >> (64 - __builtin_ctz(num_buckets))); + return static_cast(key >> (64 - __builtin_ctz(bucket_count))); } + uint32_t hash(uint64_t key) { return hash(key, num_buckets); } + /** * Link an initialized entry into bucket and task chains. */ - void link_entry(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId producer_task_id) { + template + void link_entry_impl(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id) { #if SIMPLER_TENSORMAP_PROFILING g_insert_count++; #endif - uint32_t bucket_index = hash(addr); - auto local_id = producer_task_id.local(); + constexpr bool is_reader = access_kind == TensorAccessKind::READER; + int32_t bucket_count = is_reader ? num_reader_buckets : num_buckets; + uint32_t bucket_index = hash(addr, bucket_count); + auto local_id = access_task_id.local(); int32_t task_slot = local_id & (task_window_size - 1); - entry->producer_task_id = producer_task_id; + entry->access_task_id = access_task_id; + entry->access_kind = access_kind; + if constexpr (is_reader) { + reader_used++; + if (reader_used > reader_high_water) reader_high_water = reader_used; + } // Insert at head of hash bucket + PTO2TensorMapEntry **index = is_reader ? reader_buckets : buckets; entry->bucket_index = bucket_index; - entry->next_in_bucket = buckets[bucket_index]; + entry->next_in_bucket = index[bucket_index]; if (entry->next_in_bucket != nullptr) { entry->next_in_bucket->prev_in_bucket = entry; } - buckets[bucket_index] = entry; + index[bucket_index] = entry; entry->prev_in_bucket = nullptr; // Link to task's entry list @@ -672,11 +733,32 @@ struct PTO2TensorMap { task_entry_heads[task_slot] = entry; } + template + void insert_impl(const ChipTensor &tensor, PTO2TaskId access_task_id) { + PTO2TensorMapEntry *entry = new_entry(); + entry->copy_from_tensor(tensor); + link_entry_impl(entry, tensor.buffer.addr, access_task_id); + } + + void link_entry(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id) { + link_entry_impl(entry, addr, access_task_id); + } + + void link_entry( + PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id, TensorAccessKind access_kind + ) { + if (access_kind == TensorAccessKind::READER) { + link_entry_impl(entry, addr, access_task_id); + } else { + link_entry_impl(entry, addr, access_task_id); + } + } + /** * Check if entry is valid (producer has not retired) */ bool entry_valid(const PTO2TensorMapEntry &entry) const { - return static_cast(entry.producer_task_id.local()) >= last_task_alive_cached; + return static_cast(entry.access_task_id.local()) >= last_task_alive_cached; } void remove_entry(PTO2TensorMapEntry &entry) { @@ -693,7 +775,7 @@ struct PTO2TensorMap { // Update predecessor's next pointer (O(1) via prev_in_task) if (entry.prev_in_task == nullptr) { // Entry is the head of its task chain, update task_entry_heads - int32_t local_id = static_cast(entry.producer_task_id.local()); + int32_t local_id = static_cast(entry.access_task_id.local()); int32_t task_slot = local_id & (task_window_size - 1); task_entry_heads[task_slot] = entry.next_in_task; } else { diff --git a/src/a5/runtime/host_build_graph/runtime/pto_types.h b/src/a5/runtime/host_build_graph/runtime/pto_types.h index 99a90edd19..e010d57328 100644 --- a/src/a5/runtime/host_build_graph/runtime/pto_types.h +++ b/src/a5/runtime/host_build_graph/runtime/pto_types.h @@ -348,6 +348,15 @@ struct Arg : TaskArgsTpl { ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::INPUT, tensor_count_++), ...); } + template + void add_tracked_input(Args &&...args) { + assert_add_tensor_args(); + if (!check_add_tensor_capacity(static_cast(sizeof...(Args)))) { + return; + } + ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::TRACKED_INPUT, tensor_count_++), ...); + } + /// Batch add outputs — all ChipTensor or all TensorCreateInfo: /// add_output(ci1, ci2) — runtime allocates buffers (OUTPUT) /// add_output(t1, t2) — write-only existing tensors (OUTPUT_EXISTING) @@ -372,7 +381,7 @@ struct Arg : TaskArgsTpl { ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::INOUT, tensor_count_++), ...); } - /// No-dependency existing tensor: skips OverlapMap lookup, depends on creator only. + /// Existing tensor with creator retention but no automatic TensorMap lookup or publication. template void add_no_dep(Args &&...args) { assert_add_tensor_args(); diff --git a/src/a5/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp b/src/a5/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp index 8c0a43fcd8..b72f4ed5ce 100644 --- a/src/a5/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp @@ -9,7 +9,7 @@ * ----------------------------------------------------------------------------------------------------------- */ /** - * PTO Runtime2 - Scheduler Implementation + * host_build_graph - Scheduler Implementation * * Implements scheduler state management, ready queues, and task lifecycle. * @@ -74,6 +74,7 @@ PTO2SchedProfilingData scheduler_get_profiling(int thread_idx) { void PTO2SchedulerState::print_stats() { PTO2SchedulerState *sched = this; LOG_DEBUG("=== Scheduler Statistics ==="); + LOG_DEBUG("fanin spill: %d / %d", sched->fanin_spill_top, sched->fanin_spill_capacity); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { if (sched->ring_sched_state.last_task_alive > 0) { LOG_DEBUG("Ring %d:", r); diff --git a/src/a5/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h b/src/a5/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h index cd294f811a..ecaac7e366 100644 --- a/src/a5/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h @@ -10,7 +10,7 @@ */ /** - * PTO Runtime2 - Scheduler Interface + * host_build_graph - Scheduler Interface * * The Scheduler is responsible for: * 1. Maintaining per-resource-shape ready queues @@ -416,6 +416,7 @@ struct CompletionStats { * capacities used at layout time (init_from_layout reuses them). */ struct PTO2SchedulerLayout { + size_t off_fanin_spill_ids; size_t off_ready_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_ready_sync_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_dummy_ready_queue_slots; @@ -424,6 +425,7 @@ struct PTO2SchedulerLayout { size_t off_early_dispatch_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_early_sync_start_queue_slots; uint64_t ready_queue_capacity; + int32_t fanin_spill_capacity; }; /** @@ -436,6 +438,9 @@ struct PTO2SchedulerLayout { struct PTO2SchedulerState { // Shared memory access PTO2SharedMemoryHeader *sm_header; + int32_t *fanin_spill_ids; + int32_t fanin_spill_capacity; + int32_t fanin_spill_top; // Per-ring state struct alignas(64) RingSchedState { @@ -539,11 +544,16 @@ struct PTO2SchedulerState { // set its completion_flags byte. Single-ring: all producers are ring 0, so // there is no per-edge ring indirection. + int32_t fanin_local_id(const PTO2TaskPayload &p, int32_t index) const { + if (index < PTO2_MAX_FANIN) return p.fanin_local_ids[index]; + return fanin_spill_ids[p.fanin_spill_start + index - PTO2_MAX_FANIN]; + } + bool fanin_satisfied(const PTO2TaskSlotState *s) const { const PTO2TaskPayload &p = *s->payload; const PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; for (int32_t i = 0; i < p.fanin_count; i++) { - if (!ring.is_completion_flag_set(p.fanin_local_ids[i])) return false; + if (!ring.is_completion_flag_set(fanin_local_id(p, i))) return false; } return true; } @@ -556,7 +566,7 @@ struct PTO2SchedulerState { const PTO2TaskPayload &p = *s->payload; const PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; for (int32_t i = 0; i < p.fanin_count; i++) { - if (!ring.is_completion_flag_set(p.fanin_local_ids[i])) return i; + if (!ring.is_completion_flag_set(fanin_local_id(p, i))) return i; } return -1; } @@ -582,7 +592,7 @@ struct PTO2SchedulerState { push_ready_routed(consumer); return; } - producer = &ring.get_slot_state_by_task_id(consumer->payload->fanin_local_ids[state]); + producer = &ring.get_slot_state_by_task_id(fanin_local_id(*consumer->payload, state)); } } @@ -611,7 +621,7 @@ struct PTO2SchedulerState { if (state < 0) { push_ready_routed(waiter); } else { - register_wake(&ring.get_slot_state_by_task_id(waiter->payload->fanin_local_ids[state]), waiter); + register_wake(&ring.get_slot_state_by_task_id(fanin_local_id(*waiter->payload, state)), waiter); } waiter = next; } diff --git a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp index 436ab4838c..0b802089f8 100644 --- a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp @@ -246,7 +246,7 @@ void SchedulerContext::log_stall_diagnostics( int32_t rc = 0; if (slot_state.payload != nullptr) { for (int32_t k = 0; k < fi; k++) { - int32_t pid = slot_state.payload->fanin_local_ids[k]; + int32_t pid = sched_->fanin_local_id(*slot_state.payload, k); if (ring.is_completion_flag_set(pid, std::memory_order_relaxed)) rc++; } } @@ -1156,7 +1156,7 @@ void SchedulerContext::classify_partition(int32_t thread_idx, int32_t nthreads) if (state < 0) { sched_->push_ready_routed(&slot); } else { - int32_t prod_local = slot.payload->fanin_local_ids[state]; + int32_t prod_local = sched_->fanin_local_id(*slot.payload, state); sched_->register_wake(&ring.get_slot_state_by_task_id(prod_local), &slot); } } diff --git a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp index de87c5b9e0..5b192a118c 100644 --- a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp @@ -12,12 +12,14 @@ #include #include +#include #include #include "common.h" // debug_assert #include "common/unified_log.h" #include "aicpu/aicpu_device_config.h" +#include "aicpu/cache_maintenance.h" #include "aicpu/device_time.h" #include "aicpu/platform_regs.h" #include "callable.h" @@ -952,11 +954,22 @@ int32_t SchedulerContext::run_resolution_thread(Runtime *runtime, int32_t thread int dummy_got; while ((dummy_got = sched_->dummy_ready_queue.pop_batch(dummy_batch, DUMMY_DRAIN_BATCH)) > 0) { for (int di = 0; di < dummy_got; di++) { + PTO2TaskSlotState *ready = dummy_batch[di]; + if (ready->task_kind == TaskKind::HOST_WRITE) { + PTO2TaskPayload &payload = *ready->payload; + const uint64_t bytes = payload.scalars[2]; + if (payload.scalar_count != 3 || bytes == 0 || bytes > sizeof(uint64_t)) { + fail_scheduler(runtime, thread_idx, PTO2_ERROR_INVALID_ARGS); + break; + } + void *address = reinterpret_cast(static_cast(payload.scalars[0])); + std::memcpy(address, &payload.scalars[1], static_cast(bytes)); + cache_flush_range(address, static_cast(bytes)); + } #if SIMPLER_SCHED_PROFILING - PTO2SchedulerState::TaskCompletionOutcome outcome = - sched_->complete_task(*dummy_batch[di], thread_idx); + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(*ready, thread_idx); #else - PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(*dummy_batch[di]); + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(*ready); #endif if (outcome.error_code != PTO2_ERROR_NONE) { fail_scheduler(runtime, thread_idx, outcome.error_code); diff --git a/src/a5/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp b/src/a5/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp index 0cde969e90..69cacd2854 100644 --- a/src/a5/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp +++ b/src/a5/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp @@ -104,6 +104,9 @@ void PTO2SchedulerState::RingSchedState::destroy() { ring = nullptr; } PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena) { PTO2SchedulerLayout layout{}; layout.ready_queue_capacity = PTO2_READY_QUEUE_SIZE; + layout.fanin_spill_capacity = HBG_FANIN_SPILL_POOL_SIZE; + layout.off_fanin_spill_ids = + arena.reserve(static_cast(layout.fanin_spill_capacity) * sizeof(int32_t), alignof(int32_t)); for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { layout.off_ready_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); @@ -128,6 +131,8 @@ bool PTO2SchedulerState::init_data_from_layout( ) { PTO2SchedulerState *sched = this; sched->sm_header = reinterpret_cast(sm_dev_base); + sched->fanin_spill_capacity = layout.fanin_spill_capacity; + sched->fanin_spill_top = 0; #if SIMPLER_SCHED_PROFILING sched->tasks_completed.store(0, std::memory_order_relaxed); sched->tasks_consumed.store(0, std::memory_order_relaxed); @@ -185,6 +190,7 @@ bool PTO2SchedulerState::init_data_from_layout( void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, DeviceArena &arena) { PTO2SchedulerState *sched = this; + sched->fanin_spill_ids = static_cast(arena.region_ptr(layout.off_fanin_spill_ids)); for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { ready_queue_wire_arena_pointers(&sched->ready_queues[i], arena, layout.off_ready_queue_slots[i]); } @@ -204,6 +210,9 @@ void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, void PTO2SchedulerState::destroy() { PTO2SchedulerState *sched = this; + sched->fanin_spill_ids = nullptr; + sched->fanin_spill_capacity = 0; + sched->fanin_spill_top = 0; sched->ring_sched_state.destroy(); for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { ready_queue_destroy(&sched->ready_queues[i]); diff --git a/src/a5/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp b/src/a5/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp index b9cd461780..2edda1c280 100644 --- a/src/a5/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp +++ b/src/a5/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp @@ -9,7 +9,7 @@ * ----------------------------------------------------------------------------------------------------------- */ /** - * PTO Runtime2 - TensorMap Implementation + * host_build_graph - TensorMap Implementation * * Implements TensorMap with ring buffer pool, lazy invalidation, * and chain truncation optimization. @@ -55,12 +55,17 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout( PTO2TensorMapLayout layout{}; layout.num_buckets = new_num_buckets; + layout.num_reader_buckets = + new_num_buckets < PTO2_TENSORMAP_READER_NUM_BUCKETS ? new_num_buckets : PTO2_TENSORMAP_READER_NUM_BUCKETS; layout.pool_size = new_pool_size; layout.task_window_size = new_task_window_size; layout.off_buckets = arena.reserve( static_cast(new_num_buckets) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) ); + layout.off_reader_buckets = arena.reserve( + static_cast(layout.num_reader_buckets) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) + ); layout.off_entry_pool = arena.reserve(static_cast(new_pool_size) * sizeof(PTO2TensorMapEntry), alignof(PTO2TensorMapEntry)); layout.off_free_entry_list = @@ -86,16 +91,21 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout_default(DeviceArena &arena, in */ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, DeviceArena &arena) { num_buckets = layout.num_buckets; + num_reader_buckets = layout.num_reader_buckets; pool_size = layout.pool_size; // Address arena regions for data writes; do not store these in struct // fields (wire_arena_pointers does that). auto *buckets_arena = static_cast(arena.region_ptr(layout.off_buckets)); + auto *reader_buckets_arena = static_cast(arena.region_ptr(layout.off_reader_buckets)); // buckets[]: empty == nullptr. for (int32_t i = 0; i < num_buckets; i++) { buckets_arena[i] = nullptr; } + for (int32_t i = 0; i < num_reader_buckets; i++) { + reader_buckets_arena[i] = nullptr; + } // Init-on-write: the entry pool is not pre-zeroed. new_entry() puts each // bump-allocated slot into the clean "unlinked" state (bucket_index == -1, @@ -108,6 +118,8 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev next_entry_idx = 0; free_num = 0; + reader_used = 0; + reader_high_water = 0; auto *heads_arena = static_cast(arena.region_ptr(layout.off_task_entry_heads)); for (int32_t i = 0; i < layout.task_window_size; i++) { @@ -122,6 +134,7 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev void PTO2TensorMap::wire_arena_pointers(const PTO2TensorMapLayout &layout, DeviceArena &arena) { buckets = static_cast(arena.region_ptr(layout.off_buckets)); + reader_buckets = static_cast(arena.region_ptr(layout.off_reader_buckets)); entry_pool = static_cast(arena.region_ptr(layout.off_entry_pool)); free_entry_list = static_cast(arena.region_ptr(layout.off_free_entry_list)); task_entry_heads = static_cast(arena.region_ptr(layout.off_task_entry_heads)); @@ -132,6 +145,7 @@ void PTO2TensorMap::destroy() { // stray post-destroy access trips a nullptr dereference instead of reading // a recycled allocation. buckets = nullptr; + reader_buckets = nullptr; entry_pool = nullptr; free_entry_list = nullptr; task_entry_heads = nullptr; @@ -144,10 +158,30 @@ void PTO2TensorMap::destroy() { void PTO2TensorMap::print_stats() { int32_t valid = 0; int32_t stale = 0; - int32_t empty_buckets = 0; - int32_t max_chain = 0; - int64_t total_chain = 0; - int32_t non_empty_buckets = 0; + + struct BucketStats { + int32_t empty{0}; + int32_t non_empty{0}; + int32_t max_chain{0}; + int64_t total_chain{0}; + }; + auto collect_bucket_stats = [&](PTO2TensorMapEntry **index, int32_t bucket_count) { + BucketStats stats; + for (int32_t b = 0; b < bucket_count; b++) { + int32_t chain_len = 0; + for (auto *entry = index[b]; entry != nullptr; entry = entry->next_in_bucket) { + chain_len++; + } + if (chain_len == 0) { + stats.empty++; + } else { + stats.non_empty++; + stats.total_chain += chain_len; + stats.max_chain = std::max(stats.max_chain, chain_len); + } + } + return stats; + }; // Count entries // Init-on-write: only [0, next_entry_idx) slots have ever been allocated and @@ -162,39 +196,28 @@ void PTO2TensorMap::print_stats() { } } - // Count bucket stats - for (int32_t b = 0; b < num_buckets; b++) { - int32_t chain_len = 0; - auto cur_entry = buckets[b]; - - while (cur_entry != nullptr) { - chain_len++; - cur_entry = cur_entry->next_in_bucket; - } - - if (chain_len == 0) { - empty_buckets++; - } else { - non_empty_buckets++; - total_chain += chain_len; - if (chain_len > max_chain) { - max_chain = chain_len; - } - } - } - - LOG_DEBUG("=== TensorMap Statistics ==="); - LOG_DEBUG("Pool size: %d", pool_size); - LOG_DEBUG("Pool next entry idx: %d", next_entry_idx); - LOG_DEBUG("Pool free_num: %d", free_num); - LOG_DEBUG("Num buckets: %d", num_buckets); - LOG_DEBUG("Valid entries: %d", valid); - LOG_DEBUG("Stale entries: %d", stale); - LOG_DEBUG("Empty buckets: %d", empty_buckets); - LOG_DEBUG("Max chain len: %d", max_chain); - LOG_DEBUG("Avg chain len: %.2f", non_empty_buckets > 0 ? (float)total_chain / non_empty_buckets : 0); - LOG_DEBUG("Last task alive: %d", last_task_alive_cached); - LOG_DEBUG("============================"); + const BucketStats writer_stats = collect_bucket_stats(buckets, num_buckets); + const BucketStats reader_stats = collect_bucket_stats(reader_buckets, num_reader_buckets); + + LOG_INFO("=== TensorMap Statistics ==="); + LOG_INFO("Pool size: %d", pool_size); + LOG_INFO("Pool next entry idx: %d", next_entry_idx); + LOG_INFO("Pool free_num: %d", free_num); + LOG_INFO("Reader entries: live=%d high_water=%d", reader_used, reader_high_water); + LOG_INFO("Writer entries: live=%d", current_writers()); + LOG_INFO("Buckets: writer=%d reader=%d", num_buckets, num_reader_buckets); + LOG_INFO("Valid entries: %d", valid); + LOG_INFO("Stale entries: %d", stale); + LOG_INFO( + "Writer buckets: empty=%d max_chain=%d avg_chain=%.2f", writer_stats.empty, writer_stats.max_chain, + writer_stats.non_empty > 0 ? (float)writer_stats.total_chain / writer_stats.non_empty : 0 + ); + LOG_INFO( + "Reader buckets: empty=%d max_chain=%d avg_chain=%.2f", reader_stats.empty, reader_stats.max_chain, + reader_stats.non_empty > 0 ? (float)reader_stats.total_chain / reader_stats.non_empty : 0 + ); + LOG_INFO("Last task alive: %d", last_task_alive_cached); + LOG_INFO("============================"); } int32_t PTO2TensorMap::valid_count() { diff --git a/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp b/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp index f767e8d0ec..f6d259a63b 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp @@ -788,6 +788,14 @@ int32_t AicpuExecutor::run(Runtime *runtime) { static_cast(tp.overlap_checks), static_cast(tp.overlap_hits), tp.overlap_checks > 0 ? tp.overlap_hits * 100.0 / tp.overlap_checks : 0.0 ); + rt->orchestrator.tensor_map.print_stats(); + for (int32_t ring_id = 0; ring_id < PTO2_MAX_RING_DEPTH; ++ring_id) { + const PTO2FaninPool &pool = rt->orchestrator.rings[ring_id].fanin_pool; + LOG_INFO( + "Thread %d: Fanin spill ring=%d used=%d high_water=%d capacity=%d", thread_idx, ring_id, + pool.used(), pool.high_water, pool.capacity + ); + } #endif #endif // SIMPLER_ORCH_PROFILING diff --git a/src/a5/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md b/src/a5/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md index 3594ccf259..b3bab2730e 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md +++ b/src/a5/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md @@ -41,7 +41,7 @@ Type changes: | Field | Before | After | | ----- | ------ | ----- | | `PTO2TaskDescriptor.task_id` | `int32_t` | `PTO2TaskId` | -| `PTO2TensorMapEntry.producer_task_id` | `int32_t` | `PTO2TaskId` | +| `PTO2TensorMapEntry.access_task_id` | `int32_t` | `PTO2TaskId` | | `PTO2TaskSlotState.ring_id` | N/A | `uint8_t` (new, denormalized for fast access) | ## 4. Data Structures @@ -155,8 +155,8 @@ Entry validity checks and `cleanup_retired` operate per-ring: ```cpp bool entry_valid(const PTO2TensorMapEntry& e) { - int32_t ring = e.producer_task_id.ring(); - int32_t local = e.producer_task_id.local(); + int32_t ring = e.access_task_id.ring(); + int32_t local = e.access_task_id.local(); return local >= last_task_alives[ring]; } ``` diff --git a/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index fca7d1369d..1a4c6cb5e4 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -321,24 +321,28 @@ This design avoids the complexity of ring-based wrapping while still being bound ### 5.4 Stale Entry Cleanup: Three-Layer Defense -TensorMap must ensure entries for retired tasks (`producer_task_id < last_task_alive`) are removed, so that: +TensorMap must ensure entries for retired tasks (`access_task_id < last_task_alive`) are removed, so that: - The pool does not grow unboundedly (capacity is finite) - Lookup performance does not degrade as stale entries accumulate in bucket chains Three complementary mechanisms achieve this: -**Layer 1 — Chain Truncation during Lookup** (lazy, per-bucket): +**Layer 1 — Stale Filtering during Lookup** (lazy, per-bucket): -Since `insert` always prepends to the bucket head, entries in each bucket chain are in **descending task_id order**. When `PTO2TensorMap::lookup` encounters the first stale entry (`producer_task_id < last_task_alive`), all subsequent entries in the chain are guaranteed stale too. The entire tail is truncated in one operation using `prev_in_bucket` pointers for O(1) unlinking. +`PTO2TensorMap::lookup` checks each entry against its source ring's +`last_task_alive` and skips stale entries. Entries from different rings may be +interleaved in one bucket, so a stale entry does not imply that the remaining +tail is stale and lookup must not truncate the chain. -This guarantees lookup only traverses valid entries — O(valid_entries_in_bucket), not O(total_entries). +This guarantees stale entries never create dependencies; periodic cleanup below +keeps traversal and pool use bounded. **Layer 2 — Periodic Batch Cleanup** (`cleanup_retired`, per-task): Every time the orchestrator submits a task (Step 0 of `PTO2OrchestratorState::submit_task`), it calls `PTO2TensorMap::sync_tensormap`. When `last_task_alive` has advanced by more than `PTO2_TENSORMAP_CLEANUP_INTERVAL` (default 64) tasks since the last cleanup, `PTO2TensorMap::cleanup_retired` runs: -This uses the **per-task entry chain** (`task_entry_head[task_slot]`) — each task's entries are doubly-linked together at insert time via `next_in_task`/`prev_in_task`. A slot's chain can hold more than one task's entries: a task at `local_id + N * window` reuses the slot and prepends to the chain already there, and cleanup can lag that reuse. Cleanup therefore walks the chain and frees only the entries whose `producer_task_id` matches the retiring task, unlinking each and leaving the rest linked — O(entries_in_slot), with no scan of the entire pool or all buckets. Freed entries are returned to `free_entry_list` for immediate reuse. +This uses the **per-task entry chain** (`task_entry_head[task_slot]`) — each task's entries are doubly-linked together at insert time via `next_in_task`/`prev_in_task`. A slot's chain can hold more than one task's entries: a task at `local_id + N * window` reuses the slot and prepends to the chain already there, and cleanup can lag that reuse. Cleanup therefore walks the chain and frees only the entries whose `access_task_id` matches the retiring task, unlinking each and leaving the rest linked — O(entries_in_slot), with no scan of the entire pool or all buckets. Freed entries are returned to `free_entry_list` for immediate reuse. **Layer 3 — Back-Pressure on Pool Exhaustion** (blocking): @@ -350,7 +354,7 @@ This forms a back-pressure mechanism analogous to the Task Ring's flow control. | Layer | Trigger | Method | Guarantees | | ----- | ------- | ------ | ---------- | -| Chain Truncation | Every lookup | Truncate stale tail of bucket chain | Lookup only visits valid entries | +| Stale Filtering | Every lookup | Skip entries retired in their source ring | Stale entries never create dependencies | | Periodic Cleanup | Every 64 retired tasks | Walk per-task chains, free entries | Pool capacity reclaimed in bounded time | | Pool Back-Pressure | Pool exhausted | Block until scheduler advances watermark | Hard capacity bound, no OOM | diff --git a/src/a5/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md b/src/a5/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md index de13df3ec0..f49eaff5a3 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md +++ b/src/a5/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md @@ -38,10 +38,11 @@ addr null-check → TensorMap lookup → spin-wait producer COMPLETED → comput ### 3.2 set_tensor_data Flow ```text -addr null-check → TensorMap lookup → spin-wait producer COMPLETED → spin-wait consumers done → memcpy write +addr null-check → writer/reader lookup → wait conflicting accesses → memcpy write ``` -One extra step versus get_tensor_data: wait for all consumers to finish (`fanout_refcount >= fanout_count - 1`, excluding the scope reference). +Compared with `get_tensor_data`, a write also drains each writer's consumers +and waits for every overlapping `TRACKED_INPUT` task to complete. ### 3.3 Timeout @@ -104,7 +105,7 @@ Three actors: | - | ---------- | -------- | ------ | --------- | ----- | | 1 | Kernel write (OUTPUT) | Orch Read | RAW | spin-wait producer COMPLETED | Yes | | 2 | Kernel write (OUTPUT) | Orch Write | WAW | spin-wait producer COMPLETED | Yes | -| 3 | Kernel read (INPUT) | Orch Write | WAR | spin-wait fanout_refcount | **Needs INOUT** | +| 3 | Kernel read (TRACKED_INPUT) | Orch Write | WAR | spin-wait reader task | Yes | | 4 | Kernel read-write (INOUT) | Orch Read | RAW | spin-wait producer COMPLETED | Yes | | 5 | Kernel read-write (INOUT) | Orch Write | WAW+WAR | spin-wait producer + consumers | Yes | | 6 | Orch Write | Kernel read (INPUT) | RAW | blocking completes before next submit | Yes | @@ -114,11 +115,12 @@ Three actors: ### Key Design Points -**Scenario #3 is the only case requiring special attention**: +**Scenario #3 requires an explicit reader annotation**: -TensorMap tracks only producers (OUTPUT/INOUT), not pure INPUT consumers. If a tensor is only registered via `add_input()`, TensorMap has no producer entry for it. `set_tensor_data`'s `wait_for_tensor_ready()` finds no matching producer (the lookup callback never fires) and returns immediately — but the kernel may still be reading → **WAR data race**. - -**Solution**: For tensors that may later be written via `set_tensor_data`, use `add_inout()` instead of `add_input()`. INOUT registers a producer entry in TensorMap, enabling `set_tensor_data` to track all consumers through `fanout_refcount`. +Plain `add_input()` performs the RAW lookup but does not publish a reader in +A5 `tensormap_and_ringbuffer`. If a later `set_tensor_data()` may overlap that +read, use `add_tracked_input()`. The tracked reader remains read-only but is +visible to the host writer's WAR lookup. **Scenarios #6–8 serial guarantee**: @@ -132,6 +134,9 @@ get/set_tensor_data are blocking calls, and orchestration is single-threaded ser | -------- | -------- | | External tensor never submitted as OUTPUT/INOUT | No TensorMap entry — get/set execute immediately | | External tensor previously submitted as OUTPUT/INOUT | TensorMap has producer entry — get/set spin-wait | -| External tensor submitted as INPUT, then set_tensor_data | **WAR risk** — must use INOUT instead (same as scenario #3) | +| External tensor submitted as INPUT, then set_tensor_data | **WAR risk** — the reader was not published | +| External tensor submitted as TRACKED_INPUT, then set_tensor_data | Host write waits for the reader task | -**Key rule**: If an external tensor will later be written via `set_tensor_data`, all prior kernel accesses must use `add_inout()`, not `add_input()`. +**Key rule**: If an external tensor will later be written via +`set_tensor_data`, every overlapping pure reader must use +`add_tracked_input()`. diff --git a/src/a5/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp b/src/a5/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp index 0a9384d44f..2e0a89e455 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp @@ -45,13 +45,13 @@ * "kept at call site" note); both passes run the same explicit-deps loop, so * the comparison covers it too. * - * STEP 4 (`register_task_outputs`) runs on BOTH tensor maps after both passes + * STEP 4 (`register_task_accesses`) runs on BOTH tensor maps after both passes * complete, keeping `tm_oracle` and `tm_annot` bit-equivalent for the next * record's INOUT+COVERED `remove_entry` mutations. * * Pool sizing: replay never advances last_task_alive, so each tensor map's - * entry pool must accommodate every output write across the whole trace. We - * scan the record buffer once to count INOUT + OUTPUT_EXISTING slots and size + * entry pool must accommodate every tracked access across the whole trace. We + * scan the record buffer once to count INPUT + INOUT + OUTPUT_EXISTING slots and size * the pool accordingly. Both maps get the same size. */ @@ -89,12 +89,12 @@ int32_t ceil_pow2(int32_t v) { return v + 1; } -// Count INOUT + OUTPUT_EXISTING slots across the record buffer — -// register_task_outputs only inserts those, and skips entries with manual_dep +// Count TRACKED_INPUT + INOUT + OUTPUT_EXISTING slots across the record buffer — +// register_task_accesses inserts those, and skips entries with manual_dep // set. Counting both without inspecting manual_dep is a conservative upper // bound (manual_dep is rare; the small over-allocation pays for itself in // avoided pool exhaustion). -int32_t count_outputs(const DepGenRecord *records, size_t n) { +int32_t count_accesses(const DepGenRecord *records, size_t n) { int32_t total = 0; for (size_t i = 0; i < n; i++) { const DepGenRecord &r = records[i]; @@ -104,7 +104,7 @@ int32_t count_outputs(const DepGenRecord *records, size_t n) { if (r.flags & DEP_GEN_FLAG_OVERFLOW) continue; for (uint16_t j = 0; j < r.tensor_count; j++) { auto t = static_cast(r.arg_types[j]); - if (t == TensorArgType::INOUT || t == TensorArgType::OUTPUT_EXISTING) { + if (t == TensorArgType::TRACKED_INPUT || t == TensorArgType::INOUT || t == TensorArgType::OUTPUT_EXISTING) { total++; } } @@ -158,6 +158,20 @@ const char *overlap_status_str(OverlapStatus s) { return "unknown"; } +const char *hazard_kind_str(TensorHazardKind kind) { + switch (kind) { + case TensorHazardKind::RAW: + return "RAW"; + case TensorHazardKind::WAW: + return "WAW"; + case TensorHazardKind::WAR: + return "WAR"; + } + return "unknown"; +} + +const char *access_kind_str(TensorAccessKind kind) { return kind == TensorAccessKind::READER ? "READER" : "WRITER"; } + // One annotated edge. consumer_* always populated. producer_* populated for // TENSORMAP source only — the explicit/creator emit paths don't have a // matched tensormap entry to copy from. @@ -172,7 +186,9 @@ struct EdgeAnnot { EdgeSource source; DepFlags flags; // per-edge WAIT/RETAIN semantics carried into deps.json OverlapStatus overlap; // only meaningful for TENSORMAP - uint64_t tensor_id; // 0 for EXPLICIT + TensorHazardKind hazard; + TensorAccessKind access_kind; + uint64_t tensor_id; // 0 for EXPLICIT // Consumer side (the ChipTensor the submitting task is reading). uint8_t consumer_dtype; uint32_t consumer_ndims; @@ -228,6 +244,8 @@ const char *arg_type_str(TensorArgType t) { switch (t) { case TensorArgType::INPUT: return "INPUT"; + case TensorArgType::TRACKED_INPUT: + return "TRACKED_INPUT"; case TensorArgType::OUTPUT: return "OUTPUT"; case TensorArgType::INOUT: @@ -394,6 +412,8 @@ bool write_deps_json( write_dep_flags(out, e.flags); if (e.source == EdgeSource::TENSORMAP) { out << ",\"overlap\":\"" << overlap_status_str(e.overlap) << '"'; + out << ",\"hazard\":\"" << hazard_kind_str(e.hazard) << '"'; + out << ",\"access_kind\":\"" << access_kind_str(e.access_kind) << '"'; } if (e.source != EdgeSource::EXPLICIT) { out << ",\"tensor_id\":\"" << e.tensor_id << '"'; @@ -444,21 +464,29 @@ void annot_pass( emit_creator(owner, i, *tensor); } - // STEP B: tensormap lookup (only INPUT/INOUT, skip manual_dep). - if (ptype != TensorArgType::INPUT && ptype != TensorArgType::INOUT) { - continue; - } if (tensor->manual_dep) { continue; } - tensor_map.lookup(*tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { - emit_tensormap(entry.producer_task_id, i, *tensor, entry, overlap_status); - if (ptype == TensorArgType::INOUT && overlap_status == OverlapStatus::COVERED) { - tensor_map.remove_entry(entry); - } - return true; - }); + auto lookup = [&](TensorAccessKind access_kind, TensorHazardKind hazard) { + tensor_map.lookup(*tensor, access_kind, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) { + emit_tensormap(entry.access_task_id, i, *tensor, entry, overlap_status, hazard); + if (ptype == TensorArgType::INOUT && access_kind == TensorAccessKind::WRITER) { + emit_tensormap(entry.access_task_id, i, *tensor, entry, overlap_status, TensorHazardKind::WAW); + } + if ((ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) && + overlap_status == OverlapStatus::COVERED) { + tensor_map.remove_entry(entry); + } + return true; + }); + }; + if (ptype == TensorArgType::INPUT || ptype == TensorArgType::TRACKED_INPUT || ptype == TensorArgType::INOUT) { + lookup(TensorAccessKind::WRITER, TensorHazardKind::RAW); + } + if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) { + lookup(TensorAccessKind::READER, TensorHazardKind::WAR); + } } } @@ -496,8 +524,8 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c task_window_sizes[r] = ceil_pow2(need < 16 ? 16 : need); } - int32_t output_count = count_outputs(records, num_records); - int32_t pool_size = output_count + (output_count / 10) + 64; + int32_t access_count = count_accesses(records, num_records); + int32_t pool_size = access_count + (access_count / 10) + 64; if (pool_size < PTO2_TENSORMAP_POOL_SIZE) { pool_size = PTO2_TENSORMAP_POOL_SIZE; } @@ -755,7 +783,7 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c }, // emit_tensormap(producer, arg_idx, consumer_tensor, entry, status) [&](PTO2TaskId producer, int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, - OverlapStatus status) { + OverlapStatus status, TensorHazardKind hazard) { // Per-(succ, arg_idx, producer_buffer_addr, producer_version) // dedup gives us "the same producer slice fired twice for the // same consumer arg" collapse — but two distinct slices from @@ -771,6 +799,8 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c e.source = EdgeSource::TENSORMAP; e.flags = DEP_WAIT; e.overlap = status; + e.hazard = hazard; + e.access_kind = entry.access_kind; e.tensor_id = make_tensor_id(entry.buffer_addr, entry.version); fill_consumer(e, consumer); fill_producer(e, entry); @@ -807,8 +837,8 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c } // ============ STEP 4 — publish outputs on BOTH maps ============ - register_task_outputs(inputs, task_id, tm_oracle, in_manual_scope); - register_task_outputs(inputs, task_id, tm_annot, in_manual_scope); + register_task_accesses(inputs, task_id, tm_oracle, in_manual_scope); + register_task_accesses(inputs, task_id, tm_annot, in_manual_scope); } tm_oracle.destroy(); diff --git a/src/a5/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.h b/src/a5/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.h index 06acb95a3d..7bac37de07 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.h @@ -15,7 +15,7 @@ * * Takes the records the host collector drained from the device ring buffer * (``DepGenCollector::records()``) and runs them back through a host-resident - * PTO2TensorMap using the same ``compute_task_fanin`` / ``register_task_outputs`` + * PTO2TensorMap using the same ``compute_task_fanin`` / ``register_task_accesses`` * primitives the device orchestrator uses, emitting the full * predecessor → successor edge list to deps.json. * @@ -70,8 +70,7 @@ * The replay is single-threaded and pure CPU: no device handle is required. */ -#ifndef SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_HOST_DEP_GEN_REPLAY_H_ -#define SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_HOST_DEP_GEN_REPLAY_H_ +#pragma once #include #include @@ -102,5 +101,3 @@ int dep_gen_replay_emit_deps_json(const struct DepGenRecord *records, size_t num #ifdef __cplusplus } // extern "C" #endif - -#endif // SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_HOST_DEP_GEN_REPLAY_H_ diff --git a/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h b/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h index 057cfb8cb0..bbf57c8f35 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h @@ -54,6 +54,7 @@ class CoreTaskArgsWithDeps : private CoreTaskArgs { using CoreTaskArgs::add_scalar; using CoreTaskArgs::add_scalars; using CoreTaskArgs::add_scalars_i32; + using CoreTaskArgs::add_tracked_input; using CoreTaskArgs::allow_early_resolve; // early-dispatch hint (getter) using CoreTaskArgs::copy_scalars_from; using CoreTaskArgs::set_allow_early_resolve; // early-dispatch hint (setter) diff --git a/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_orchestration_api.h b/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_orchestration_api.h index 16386b1a1e..6d446fdda7 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_orchestration_api.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_orchestration_api.h @@ -297,20 +297,9 @@ static inline T get_tensor_data(const ChipTensor &tensor, uint32_t ndims, const * set_tensor_data(tensor, 1, idx, raw_u64); // old usage unchanged * set_tensor_data(tensor, 1, idx, 42.0f); // typed write (T = float) * - * If the tensor has a producer in TensorMap, spin-waits until the producer - * and all its consumers complete before writing (WAW + WAR safety). - * External tensors (make_tensor_external) with no TensorMap entry are - * written immediately without waiting. - * - * Limitation: TensorMap only tracks producers (OUTPUT/INOUT), not consumers - * that used the tensor as INPUT. If a kernel reads this tensor as INPUT - * (not INOUT) and the tensor has no TensorMap producer entry, set_tensor_data - * cannot detect the reader and may cause a data race. - * - * To ensure WAR safety for all access patterns, use add_inout() instead of - * add_input() for kernel parameters that may later be written via - * set_tensor_data. INOUT creates a TensorMap entry that enables automatic - * consumer tracking via fanout_refcount. + * Spin-waits for overlapping writers and explicitly tracked readers before + * writing (WAW + WAR safety). A pure reader that may precede this call must + * use add_tracked_input(); plain add_input() does not publish a reader entry. * * The tensor must already have an allocated buffer (addr != 0). * For runtime-created outputs, call this only on the ChipTensor returned by diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h index 16a92380a7..8936242f5a 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h @@ -16,11 +16,11 @@ * Two header-only template entry points: * * compute_task_fanin — STEP 3 in submit_task: per-tensor creator retention (Step A) - * + tensormap.lookup for INPUT/INOUT (Step B). Calls back into + * + tensormap.lookup for INPUT/TRACKED_INPUT/INOUT (Step B). Calls back into * user-supplied `emit` for each producer it identifies. * - * register_task_outputs — STEP 4 in submit_task: tensormap.insert for INOUT and - * OUTPUT_EXISTING tensors. No callbacks. + * register_task_accesses — STEP 4 in submit_task: register TRACKED_INPUT readers and + * INOUT/OUTPUT_EXISTING writers. No callbacks. * * STEP 1 (explicit_deps) is intentionally left at the runtime call site because its * `last_task_alive` shortcut + unchecked slot lookup is subtly different from the @@ -47,6 +47,7 @@ #pragma once #include +#include #include "pto_task_id.h" #include "pto_tensormap.h" @@ -54,7 +55,7 @@ #include "tensor.h" /** - * View struct for inputs to compute_task_fanin / register_task_outputs. + * View struct for inputs to compute_task_fanin / register_task_accesses. * * Both runtime and replay assemble one of these from their own data sources * (runtime: from Arg accessors; replay: from SubmitTraceEntry fields). All @@ -68,21 +69,83 @@ struct DepInputs { const PTO2TaskId *explicit_deps; // length = explicit_dep_count (validity checked by caller) }; +template +[[nodiscard]] inline bool +emit_writer_dependency(Emit &emit, PTO2TaskId task_id, DepFlags flags) { + if constexpr (std::is_invocable_r_v) { + return emit(task_id, flags); + } else if constexpr (std::is_invocable_r_v) { + return emit(task_id, flags, TensorAccessKind::WRITER); + } else if constexpr (std::is_invocable_r_v) { + return emit(task_id, TensorAccessKind::WRITER); + } else { + return emit(task_id); + } +} + +template +[[nodiscard]] inline bool emit_reader_dependency(Emit &emit, PTO2TaskId task_id, DepFlags flags) { + if constexpr (std::is_invocable_r_v) { + return emit(task_id, flags); + } else if constexpr (std::is_invocable_r_v) { + return emit(task_id, flags, TensorAccessKind::READER); + } else if constexpr (std::is_invocable_r_v) { + return emit(task_id, TensorAccessKind::READER); + } else { + return emit(task_id); + } +} + +template +[[nodiscard]] __attribute__((noinline, cold)) bool +compute_reader_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, ReaderEmit &reader_emit) { + for (int32_t i = 0; i < inputs.tensor_count; i++) { + TensorArgType ptype = inputs.arg_types[i]; + if (ptype != TensorArgType::INOUT && ptype != TensorArgType::OUTPUT_EXISTING) { + continue; + } + const ChipTensor *tensor = &inputs.tensors[i].ref(); + if (tensor->manual_dep) { + continue; + } + bool fatal = false; + tensor_map.lookup( + *tensor, TensorAccessKind::READER, + [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { + if (!emit_reader_dependency(reader_emit, entry.access_task_id, DEP_WAIT)) { + fatal = true; + return false; + } + if (overlap_status == OverlapStatus::COVERED) { + tensor_map.remove_reader_entry(entry); + } + return true; + } + ); + if (fatal) { + return false; + } + } + return true; +} + /** * Compute fanin for a task being submitted (STEP 3: Step A creator retention + * Step B tensormap modifier lookup). * * For each non-OUTPUT tensor: * - If owner_task_id is valid, emit(owner) - * - For INPUT/INOUT (and not manual_dep), tensor_map.lookup(*tensor) and emit + * - For INPUT/TRACKED_INPUT/INOUT (and not manual_dep), tensor_map.lookup(*tensor) and emit * each matching producer. INOUT+COVERED triggers tensor_map.remove_entry(entry). * * @return true on success (or producer-skipped-silently); false if emit signaled * fatal — caller should propagate (after any fatal bookkeeping done by emit). */ -template +template [[nodiscard]] inline bool -compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_manual_scope, Emit emit) { +compute_task_fanin( + const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_manual_scope, Emit emit, ReaderEmit reader_emit +) { if (in_manual_scope) { return true; } @@ -101,13 +164,13 @@ compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_m // the creator edge carries both ordering and lifetime. PTO2TaskId owner = tensor->owner_task_id; if (owner.is_valid()) { - if (!emit(owner, DEP_WAIT | DEP_RETAIN)) { + if (!emit_writer_dependency(emit, owner, DEP_WAIT | DEP_RETAIN)) { return false; } } - // Step B: only INPUT/INOUT need modifier dependency lookup. - if (ptype != TensorArgType::INPUT && ptype != TensorArgType::INOUT) { + if (ptype != TensorArgType::INPUT && ptype != TensorArgType::INOUT && + ptype != TensorArgType::TRACKED_INPUT) { continue; } if (tensor->manual_dep) { @@ -116,22 +179,12 @@ compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_m bool fatal = false; tensor_map.lookup(*tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { - // Ordering-only (DEP_WAIT): a modifier only rewrote a buffer someone - // else allocated, so its lifetime rides that allocator's creator edge, - // not this modifier edge. Retention-safety invariant that makes this - // sound: only TensorArgType::OUTPUT tensors are allocated into the - // packed output heap, and a runtime-created OUTPUT always carries a - // valid owner_task_id — so its consumer takes a Step-A DEP_RETAIN edge - // to the allocator above. INOUT / OUTPUT_EXISTING buffers are never - // owned by a modifier. If a future layout put a modifier-owned buffer - // into the packed heap, this edge would have to become RETAIN or the - // producer could be reclaimed under a live reader (use-after-free). - if (!emit(entry.producer_task_id, DEP_WAIT)) { + if (!emit_writer_dependency(emit, entry.access_task_id, DEP_WAIT)) { fatal = true; - return false; // stop iteration + return false; } if (ptype == TensorArgType::INOUT && overlap_status == OverlapStatus::COVERED) { - tensor_map.remove_entry(entry); + tensor_map.remove_writer_entry(entry); } return true; }); @@ -139,19 +192,22 @@ compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_m return false; } } - return true; + return tensor_map.current_readers() == 0 || compute_reader_fanin(inputs, tensor_map, reader_emit); +} + +template +[[nodiscard]] inline bool +compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_manual_scope, Emit emit) { + return compute_task_fanin(inputs, tensor_map, in_manual_scope, emit, emit); } /** - * Register a task's outputs in the tensormap (STEP 4 in submit_task). - * - * For INOUT and OUTPUT_EXISTING tensors (excluding manual_dep), inserts the - * tensor into tensor_map keyed by its buffer.addr with `task_id` as producer. + * Register a task's accesses in the tensormap (STEP 4 in submit_task). * * No-op when in_manual_scope. */ inline void -register_task_outputs(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap &tensor_map, bool in_manual_scope) { +register_task_accesses(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap &tensor_map, bool in_manual_scope) { if (in_manual_scope) { return; } @@ -162,26 +218,31 @@ register_task_outputs(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap if (!tensor->manual_dep) { tensor_map.insert(*tensor, task_id); } + } else if (ptype == TensorArgType::TRACKED_INPUT) { + const ChipTensor *tensor = &inputs.tensors[i].ref(); + if (!tensor->manual_dep) { + tensor_map.insert(*tensor, task_id, TensorAccessKind::READER); + } } } } /** - * Count the tensormap entries register_task_outputs() will insert for this task. + * Count the tensormap entries register_task_accesses() will insert for this task. * - * Mirrors register_task_outputs()'s selection exactly (INOUT / OUTPUT_EXISTING, - * excluding manual_dep), so the returned value is the precise number of + * Mirrors register_task_accesses() exactly, so the returned value is the precise number of * new_entry() calls that step makes. The orchestrator uses it to reserve pool * capacity before inserting. Returns 0 in a manual scope (no registration). */ -inline int32_t count_registrable_outputs(const DepInputs &inputs, bool in_manual_scope) { +inline int32_t count_registrable_accesses(const DepInputs &inputs, bool in_manual_scope) { if (in_manual_scope) { return 0; } int32_t needed = 0; for (int32_t i = 0; i < inputs.tensor_count; i++) { TensorArgType ptype = inputs.arg_types[i]; - if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) { + if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING || + ptype == TensorArgType::TRACKED_INPUT) { if (!inputs.tensors[i].ref().manual_dep) { needed++; } diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp index 8983c895b9..063e804de2 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp @@ -10,7 +10,7 @@ */ /** - * PTO Runtime2 - Orchestrator Implementation + * tensormap_and_ringbuffer - Orchestrator Implementation * * Implements orchestrator state management, scope handling, and task submission. * @@ -815,7 +815,7 @@ static bool ensure_tensormap_capacity(PTO2OrchestratorState *orch, int32_t neede return true; } // Progress is entries actually freed, NOT watermark movement: a ring can - // retire zero-output tasks (count_registrable_outputs == 0), advancing + // retire zero-access tasks (count_registrable_accesses == 0), advancing // last_task_alive without freeing any entry. Gating the backstop on // free_entries() keeps a wedged pool from dodging the timeout while some // unrelated ring keeps draining. @@ -897,7 +897,7 @@ static TaskOutputTensors submit_task_common( if (is_dep_gen_enabled()) { const void *tensor_ptrs[MAX_TENSOR_ARGS]; // TensorArgType is `enum class : int32_t` (4 bytes); the on-disk record - // packs arg_types as uint8_t[16] (5-value enum fits in a byte). Narrow + // packs arg_types as uint8_t[16] (the enum fits in a byte). Narrow // each tag here rather than letting the AICPU writer reinterpret a // 4×-wider array as bytes — that path silently lost two of every three // tags on little-endian and synthesized phantom self-edges in replay. @@ -984,7 +984,23 @@ static TaskOutputTensors submit_task_common( ); }; - if (!compute_task_fanin(dep_inputs, orch->tensor_map, orch->in_manual_scope(), runtime_emit)) { + auto runtime_emit_reader = [&](PTO2TaskId producer_task_id, DepFlags kind) -> bool { + uint8_t prod_ring = producer_task_id.ring(); + PTO2SharedMemoryRingHeader &producer_ring = orch->sm_header->rings[prod_ring]; + int32_t prod_slot = producer_ring.get_slot_by_task_id(static_cast(producer_task_id.local())); + PTO2TaskSlotState *prod_state = &producer_ring.get_slot_state_by_slot(prod_slot); + if (prod_state->task != nullptr && prod_state->task->task_id == producer_task_id && + prod_state->task_state.load(std::memory_order_acquire) >= PTO2_TASK_COMPLETED) { + return true; + } + return append_fanin_or_fail( + orch, prod_ring, prod_slot, prod_state, producer_task_id, &fanin_builder, ring_id, kind + ); + }; + + if (!compute_task_fanin( + dep_inputs, orch->tensor_map, orch->in_manual_scope(), runtime_emit, runtime_emit_reader + )) { return result; } @@ -995,11 +1011,11 @@ static TaskOutputTensors submit_task_common( // is shared across rings and reclaimed as last_task_alive advances; an // exhausted pool back-pressures here (and detects a wedged watermark) rather // than tripping new_entry()'s hard assert mid-registration. - int32_t tensormap_needed = count_registrable_outputs(dep_inputs, orch->in_manual_scope()); + int32_t tensormap_needed = count_registrable_accesses(dep_inputs, orch->in_manual_scope()); if (tensormap_needed > 0 && !ensure_tensormap_capacity(orch, tensormap_needed)) { return result; } - register_task_outputs(dep_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); + register_task_accesses(dep_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); CYCLE_COUNT_LAP(g_orch_insert_cycle); diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp index f26c5f9993..6d36570e76 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp @@ -10,7 +10,7 @@ */ /** - * PTO Runtime2 - Main Implementation + * tensormap_and_ringbuffer - Main Implementation * * Implements the unified runtime API that combines orchestrator and scheduler. * @@ -109,6 +109,7 @@ wait_for_tensor_ready(PTO2Runtime *rt, const ChipTensor &tensor, bool wait_for_c // the second encounter. constexpr int kSegmentCap = 64; const PTO2TaskSlotState *seg[kSegmentCap]; + bool seg_wait_consumers[kSegmentCap]{}; int seg_count = 0; bool failed = false; @@ -169,39 +170,53 @@ wait_for_tensor_ready(PTO2Runtime *rt, const ChipTensor &tensor, bool wait_for_c for (int i = 0; i < seg_count; i++) { wait_one_producer(*seg[i]); if (failed) return; - if (!wait_for_consumers) continue; + if (!seg_wait_consumers[i]) continue; wait_one_consumers(*seg[i]); if (failed) return; } seg_count = 0; }; - auto try_push = [&](const PTO2TaskSlotState &s) { + auto try_push = [&](const PTO2TaskSlotState &s, bool include_consumers) { for (int j = 0; j < seg_count; j++) { - if (seg[j] == &s) return; // per-segment dedup + if (seg[j] == &s) { + seg_wait_consumers[j] = seg_wait_consumers[j] || include_consumers; + return; + } } if (seg_count == kSegmentCap) { flush_segment(); if (failed) return; } - seg[seg_count++] = &s; + seg[seg_count] = &s; + seg_wait_consumers[seg_count++] = include_consumers; }; auto do_wait = [&]() { // Step A: creator retention — read owner directly from tensor metadata if (owner.is_valid()) { auto &s = orch.sm_header->rings[owner.ring()].get_slot_state_by_task_id(owner.local()); - try_push(s); + try_push(s, wait_for_consumers); if (failed) return; } // Step B: modifier writer lookup (OverlapMap), direct callback orch.tensor_map.lookup(tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus) -> bool { - PTO2TaskId pid = entry.producer_task_id; + PTO2TaskId pid = entry.access_task_id; auto &s = orch.sm_header->rings[pid.ring()].get_slot_state_by_task_id(pid.local()); - try_push(s); + try_push(s, wait_for_consumers); return !failed; }); + if (wait_for_consumers && !failed) { + orch.tensor_map.lookup( + tensor, TensorAccessKind::READER, [&](PTO2TensorMapEntry &entry, OverlapStatus) -> bool { + PTO2TaskId pid = entry.access_task_id; + auto &s = orch.sm_header->rings[pid.ring()].get_slot_state_by_task_id(pid.local()); + try_push(s, false); + return !failed; + } + ); + } if (failed) return; flush_segment(); }; diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h index 5898b9c069..94e368ba2f 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h @@ -73,6 +73,7 @@ #define PTO2_DEP_LIST_POOL_SIZE 16384 // Per-ring dependency list pool entries #define PTO2_TENSORMAP_POOL_SIZE (65536) // TensorMap entry pool #define PTO2_TENSORMAP_NUM_BUCKETS 4096 // Power of 2 for fast hash (4096×8B=32KB fits L1) +#define PTO2_TENSORMAP_READER_NUM_BUCKETS 512 // Sparse tracked-reader sidecar (512×12B=6KB with epochs) // Scope management #define PTO2_MAX_SCOPE_DEPTH 64 // Maximum nesting depth diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h index 2b708d318d..d3a7d880f3 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h @@ -10,10 +10,10 @@ */ /** - * PTO Runtime2 - TensorMap Interface + * tensormap_and_ringbuffer - TensorMap Interface * - * TensorMap provides producer lookup for dependency discovery: - * - Maps ChipTensor -> producer task ID + * TensorMap provides access lookup for dependency discovery: + * - Maps ChipTensor -> reader/writer task IDs * - Used by pto_submit_task() to find dependencies * * Key design features: @@ -57,6 +57,17 @@ enum class OverlapStatus { OTHER, }; +enum class TensorAccessKind : uint8_t { + READER = 0, + WRITER = 1, +}; + +enum class TensorHazardKind : uint8_t { + RAW = 0, + WAW = 1, + WAR = 2, +}; + struct Segment { uint64_t begin; uint64_t end; @@ -75,11 +86,14 @@ struct Segment { struct PTO2TensorMapLayout { size_t off_buckets; size_t off_bucket_epochs; + size_t off_reader_buckets; + size_t off_reader_bucket_epochs; size_t off_entry_pool; size_t off_free_entry_list; size_t off_task_entry_heads[PTO2_MAX_RING_DEPTH]; size_t off_task_entry_head_epochs[PTO2_MAX_RING_DEPTH]; int32_t num_buckets; + int32_t num_reader_buckets; int32_t pool_size; int32_t task_window_sizes[PTO2_MAX_RING_DEPTH]; }; @@ -108,7 +122,7 @@ extern uint64_t g_insert_count; * the hash key, size in [8, 16) is unused by the entry — we repurpose it for * `next_in_bucket`). * - * buffer_addr / next_in_bucket / producer_task_id — chain traversal + match + * buffer_addr / next_in_bucket / access_task_id — chain traversal + match * start_offset — overlap byte range begin * version, ndims, dtype, manual_dep, is_contiguous — overlap fast path * shapes[5] — overlap comparison (line 1) @@ -128,15 +142,15 @@ struct alignas(64) PTO2TensorMapEntry { // === Cache line 1 (64B) — lookup hot path; mirrors ChipTensor line 1 from byte 16 === uint64_t buffer_addr; // 8B [0, 8): tensor base address (hash key, mirrors ChipTensor::buffer.addr) PTO2TensorMapEntry - *next_in_bucket; // 8B [8, 16): next entry in hash bucket chain (overlays ChipTensor::buffer.size) - PTO2TaskId producer_task_id; // 8B [16,24): mirrors ChipTensor::owner_task_id slot - uint64_t start_offset; // 8B [24,32): mirrors ChipTensor::start_offset (element offset) - int32_t version; // 4B [32,36): mirrors ChipTensor::version - uint32_t ndims; // 4B [36,40): mirrors ChipTensor::ndims - DataType dtype; // 1B [40,41): mirrors ChipTensor::dtype - bool manual_dep; // 1B [41,42): mirrors ChipTensor::manual_dep - bool is_contiguous; // 1B [42,43): mirrors ChipTensor::is_contiguous - uint8_t __padding1__; // 1B [43,44): mirrors ChipTensor padding + *next_in_bucket; // 8B [8, 16): next entry in hash bucket chain (overlays ChipTensor::buffer.size) + PTO2TaskId access_task_id; // 8B [16,24): task owning this access + uint64_t start_offset; // 8B [24,32): mirrors ChipTensor::start_offset (element offset) + int32_t version; // 4B [32,36): mirrors ChipTensor::version + uint32_t ndims; // 4B [36,40): mirrors ChipTensor::ndims + DataType dtype; // 1B [40,41): mirrors ChipTensor::dtype + bool manual_dep; // 1B [41,42): mirrors ChipTensor::manual_dep + bool is_contiguous; // 1B [42,43): mirrors ChipTensor::is_contiguous + TensorAccessKind access_kind; // 1B [43,44): reader or writer index uint32_t shapes[MAX_TENSOR_DIMS]; // 20B [44,64): mirrors ChipTensor::shapes // === Cache line 2 (64B) — chain manipulation + non-contiguous overlap data === @@ -153,7 +167,7 @@ struct alignas(64) PTO2TensorMapEntry { * Copy overlap-relevant fields from a ChipTensor into this entry. * * 64B memcpy of ChipTensor cache line 1 populates buffer_addr (byte [0,8)), - * producer_task_id, start_offset, version, ndims, dtype, manual_dep, + * access_task_id, start_offset, version, ndims, dtype, manual_dep, * is_contiguous and shapes[]. Byte [8,16) holds ChipTensor::buffer.size in * the source and gets written into next_in_bucket; that's harmless * because link_entry() overwrites next_in_bucket immediately after. @@ -338,7 +352,7 @@ struct alignas(64) PTO2TensorMapEntry { static_assert(sizeof(PTO2TensorMapEntry) == 128, "TensorMapEntry must be exactly 2 cache lines (128 bytes)"); static_assert(offsetof(PTO2TensorMapEntry, buffer_addr) == offsetof(ChipTensor, buffer.addr)); -static_assert(offsetof(PTO2TensorMapEntry, producer_task_id) == offsetof(ChipTensor, owner_task_id)); +static_assert(offsetof(PTO2TensorMapEntry, access_task_id) == offsetof(ChipTensor, owner_task_id)); static_assert(offsetof(PTO2TensorMapEntry, start_offset) == offsetof(ChipTensor, start_offset)); static_assert(offsetof(PTO2TensorMapEntry, version) == offsetof(ChipTensor, version)); static_assert(offsetof(PTO2TensorMapEntry, ndims) == offsetof(ChipTensor, ndims)); @@ -385,6 +399,14 @@ struct PTO2TensorMap { // Per-ring cleanup progress (for periodic cleanup_retired) int32_t last_cleanup[PTO2_MAX_RING_DEPTH]{}; + // Tracked-reader side index. Keep these fields after the original writer + // state so the writer-only hot path retains main's field offsets/cache layout. + PTO2TensorMapEntry **reader_buckets; + uint32_t *reader_bucket_epochs; + int32_t num_reader_buckets; + int32_t reader_used{0}; + int32_t reader_high_water{0}; + uint32_t get_task_local_id_slot(uint8_t ring_id, uint32_t task_local_id) const { return task_local_id & (task_window_sizes[ring_id] - 1); } @@ -394,6 +416,8 @@ struct PTO2TensorMap { // setter symbols must export for host dlsym; the probe call sites that use // these accessors stay gated by SIMPLER_DFX). int32_t current_used() const { return next_entry_idx - free_num; } + int32_t current_readers() const { return reader_used; } + int32_t current_writers() const { return current_used() - reader_used; } int32_t pool_capacity() const { return pool_size; } int32_t free_entries() const { return pool_size - current_used(); } @@ -430,23 +454,37 @@ struct PTO2TensorMap { return res; } - void free_entry(PTO2TensorMapEntry &entry) { + void free_writer_entry(PTO2TensorMapEntry &entry) { always_assert(entry.bucket_index != -1); // must still be in a bucket - // Update predecessor's next pointer (O(1) via prev_in_bucket) if (entry.prev_in_bucket == nullptr) { - // Entry is the head of its bucket chain, update bucket head - // Must compute hash BEFORE clearing tensor buckets[entry.bucket_index] = entry.next_in_bucket; } else { entry.prev_in_bucket->next_in_bucket = entry.next_in_bucket; } - - // Update successor's prev pointer if (entry.next_in_bucket != nullptr) { entry.next_in_bucket->prev_in_bucket = entry.prev_in_bucket; } + free_entry_list[free_num++] = &entry; + entry.bucket_index = -1; + entry.next_in_bucket = nullptr; + entry.prev_in_bucket = nullptr; + entry.next_in_task = nullptr; + entry.prev_in_task = nullptr; + } + void free_reader_entry(PTO2TensorMapEntry &entry) { + always_assert(entry.bucket_index != -1); + if (entry.prev_in_bucket == nullptr) { + reader_buckets[entry.bucket_index] = entry.next_in_bucket; + } else { + entry.prev_in_bucket->next_in_bucket = entry.next_in_bucket; + } + if (entry.next_in_bucket != nullptr) { + entry.next_in_bucket->prev_in_bucket = entry.prev_in_bucket; + } + always_assert(reader_used > 0); + reader_used--; free_entry_list[free_num++] = &entry; entry.bucket_index = -1; entry.next_in_bucket = nullptr; @@ -455,6 +493,14 @@ struct PTO2TensorMap { entry.prev_in_task = nullptr; } + void free_entry(PTO2TensorMapEntry &entry) { + if (entry.access_kind == TensorAccessKind::READER) { + free_reader_entry(entry); + } else { + free_writer_entry(entry); + } + } + // ============================================================================= // TensorMap API // ============================================================================= @@ -521,13 +567,17 @@ struct PTO2TensorMap { * @param tensor ChipTensor to look up * @param on_match Callback invoked for each overlapping entry */ - template - void lookup(const ChipTensor &tensor, Fn &&on_match) { - uint32_t bucket_index = hash(tensor.buffer.addr); - if (bucket_epochs[bucket_index] != current_epoch) { + template + void lookup_impl(const ChipTensor &tensor, Fn &&on_match) { + constexpr bool is_reader = access_kind == TensorAccessKind::READER; + int32_t bucket_count = is_reader ? num_reader_buckets : num_buckets; + uint32_t bucket_index = hash(tensor.buffer.addr, bucket_count); + uint32_t *epochs = is_reader ? reader_bucket_epochs : bucket_epochs; + PTO2TensorMapEntry **index = is_reader ? reader_buckets : buckets; + if (epochs[bucket_index] != current_epoch) { return; } - PTO2TensorMapEntry *cur_entry = buckets[bucket_index]; + PTO2TensorMapEntry *cur_entry = index[bucket_index]; #if SIMPLER_TENSORMAP_PROFILING g_lookup_count++; @@ -579,19 +629,41 @@ struct PTO2TensorMap { #endif } + template + void lookup(const ChipTensor &tensor, Fn &&on_match) { + lookup_impl(tensor, static_cast(on_match)); + } + + template + void lookup(const ChipTensor &tensor, TensorAccessKind access_kind, Fn &&on_match) { + if (access_kind == TensorAccessKind::READER) { + lookup_impl(tensor, static_cast(on_match)); + } else { + lookup_impl(tensor, static_cast(on_match)); + } + } + /** - * Insert a new entry (called when task produces output) + * Insert a new reader or writer access entry. * * Allocates from ring buffer pool, may overwrite stale entries. * Inserts at head of hash bucket chain (maintains task_id ordering). * - * @param tensor ChipTensor produced - * @param producer_task_id Task ID of producer + * @param tensor ChipTensor accessed + * @param access_task_id Task ID owning the access */ - void insert(const ChipTensor &tensor, PTO2TaskId producer_task_id) { + void insert(const ChipTensor &tensor, PTO2TaskId access_task_id, TensorAccessKind access_kind) { + if (access_kind == TensorAccessKind::READER) { + insert_impl(tensor, access_task_id); + } else { + insert_impl(tensor, access_task_id); + } + } + + void insert(const ChipTensor &tensor, PTO2TaskId access_task_id) { PTO2TensorMapEntry *entry = new_entry(); entry->copy_from_tensor(tensor); - link_entry(entry, tensor.buffer.addr, producer_task_id); + link_entry(entry, tensor.buffer.addr, access_task_id); } /** @@ -618,7 +690,7 @@ struct PTO2TensorMap { PTO2TensorMapEntry *cur_entry = task_entry_heads[ring_id][task_slot]; while (cur_entry != nullptr) { PTO2TensorMapEntry *next_entry = cur_entry->next_in_task; // free_entry clears it - if (cur_entry->producer_task_id == retired_task) { + if (cur_entry->access_task_id == retired_task) { if (cur_entry->prev_in_task != nullptr) { cur_entry->prev_in_task->next_in_task = next_entry; } else { @@ -646,26 +718,81 @@ struct PTO2TensorMap { * addresses (low bits all-zero) still distribute evenly. We extract * the top log2(num_buckets) bits which carry the most entropy. */ - uint32_t hash(uint64_t key) { + uint32_t hash(uint64_t key, int32_t bucket_count) { key *= 0x9E3779B97F4A7C15ULL; - return static_cast(key >> (64 - __builtin_ctz(num_buckets))); + return static_cast(key >> (64 - __builtin_ctz(bucket_count))); } + uint32_t hash(uint64_t key) { return hash(key, num_buckets); } + /** * Link an initialized entry into bucket and task chains. */ - void link_entry(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId producer_task_id) { + template + void link_entry_impl(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id) { #if SIMPLER_TENSORMAP_PROFILING g_insert_count++; #endif - uint32_t bucket_index = hash(addr); - auto ring_id = producer_task_id.ring(); - auto local_id = producer_task_id.local(); + constexpr bool is_reader = access_kind == TensorAccessKind::READER; + int32_t bucket_count = is_reader ? num_reader_buckets : num_buckets; + uint32_t bucket_index = hash(addr, bucket_count); + auto ring_id = access_task_id.ring(); + auto local_id = access_task_id.local(); int32_t task_slot = local_id & (task_window_sizes[ring_id] - 1); - entry->producer_task_id = producer_task_id; + entry->access_task_id = access_task_id; + entry->access_kind = access_kind; + if constexpr (is_reader) { + reader_used++; + if (reader_used > reader_high_water) reader_high_water = reader_used; + } // Insert at head of hash bucket + PTO2TensorMapEntry **index = is_reader ? reader_buckets : buckets; + uint32_t *epochs = is_reader ? reader_bucket_epochs : bucket_epochs; + if (epochs[bucket_index] != current_epoch) { + index[bucket_index] = nullptr; + epochs[bucket_index] = current_epoch; + } + entry->bucket_index = bucket_index; + entry->next_in_bucket = index[bucket_index]; + if (entry->next_in_bucket != nullptr) { + entry->next_in_bucket->prev_in_bucket = entry; + } + index[bucket_index] = entry; + entry->prev_in_bucket = nullptr; + + // Link to task's entry list + if (task_entry_head_epochs[ring_id][task_slot] != current_epoch) { + task_entry_heads[ring_id][task_slot] = nullptr; + task_entry_head_epochs[ring_id][task_slot] = current_epoch; + } + entry->next_in_task = task_entry_heads[ring_id][task_slot]; + entry->prev_in_task = nullptr; + if (entry->next_in_task != nullptr) { + entry->next_in_task->prev_in_task = entry; + } + task_entry_heads[ring_id][task_slot] = entry; + } + + template + void insert_impl(const ChipTensor &tensor, PTO2TaskId access_task_id) { + PTO2TensorMapEntry *entry = new_entry(); + entry->copy_from_tensor(tensor); + link_entry_impl(entry, tensor.buffer.addr, access_task_id); + } + + void link_entry(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id) { +#if SIMPLER_TENSORMAP_PROFILING + g_insert_count++; +#endif + uint32_t bucket_index = hash(addr); + auto ring_id = access_task_id.ring(); + auto local_id = access_task_id.local(); + int32_t task_slot = local_id & (task_window_sizes[ring_id] - 1); + + entry->access_task_id = access_task_id; + entry->access_kind = TensorAccessKind::WRITER; if (bucket_epochs[bucket_index] != current_epoch) { buckets[bucket_index] = nullptr; bucket_epochs[bucket_index] = current_epoch; @@ -678,7 +805,6 @@ struct PTO2TensorMap { buckets[bucket_index] = entry; entry->prev_in_bucket = nullptr; - // Link to task's entry list if (task_entry_head_epochs[ring_id][task_slot] != current_epoch) { task_entry_heads[ring_id][task_slot] = nullptr; task_entry_head_epochs[ring_id][task_slot] = current_epoch; @@ -691,11 +817,21 @@ struct PTO2TensorMap { task_entry_heads[ring_id][task_slot] = entry; } + void link_entry( + PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId access_task_id, TensorAccessKind access_kind + ) { + if (access_kind == TensorAccessKind::READER) { + link_entry_impl(entry, addr, access_task_id); + } else { + link_entry_impl(entry, addr, access_task_id); + } + } + /** * Check if entry is valid (producer has not retired) */ bool entry_valid(const PTO2TensorMapEntry &entry) const { - return static_cast(entry.producer_task_id.local()) >= last_task_alives[entry.producer_task_id.ring()]; + return static_cast(entry.access_task_id.local()) >= last_task_alives[entry.access_task_id.ring()]; } void remove_entry(PTO2TensorMapEntry &entry) { @@ -703,6 +839,16 @@ struct PTO2TensorMap { free_entry(entry); } + void remove_writer_entry(PTO2TensorMapEntry &entry) { + remove_from_task(entry); + free_writer_entry(entry); + } + + void remove_reader_entry(PTO2TensorMapEntry &entry) { + remove_from_task(entry); + free_reader_entry(entry); + } + /** * Remove entry from its task chain (O(1) with prev pointer) * Called during pool wrap-around to unlink reused entries. @@ -712,8 +858,8 @@ struct PTO2TensorMap { // Update predecessor's next pointer (O(1) via prev_in_task) if (entry.prev_in_task == nullptr) { // Entry is the head of its task chain, update task_entry_heads - int32_t ring_id = entry.producer_task_id.ring(); - int32_t local_id = static_cast(entry.producer_task_id.local()); + int32_t ring_id = entry.access_task_id.ring(); + int32_t local_id = static_cast(entry.access_task_id.local()); int32_t task_slot = local_id & (task_window_sizes[ring_id] - 1); task_entry_heads[ring_id][task_slot] = entry.next_in_task; } else { diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h index 57f239b242..c70655d96f 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h @@ -22,8 +22,7 @@ * without type conflicts (Handshake, TensorLease, HostApi). */ -#ifndef SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_ -#define SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_ +#pragma once #include #include @@ -380,6 +379,15 @@ struct Arg : TaskArgsTpl { ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::INPUT, tensor_count_++), ...); } + template + void add_tracked_input(Args &&...args) { + assert_add_tensor_args(); + if (!check_add_tensor_capacity(static_cast(sizeof...(Args)))) { + return; + } + ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::TRACKED_INPUT, tensor_count_++), ...); + } + /// Batch add outputs — all ChipTensor or all TensorCreateInfo: /// add_output(ci1, ci2) — runtime allocates buffers (OUTPUT) /// add_output(t1, t2) — write-only existing tensors (OUTPUT_EXISTING) @@ -404,7 +412,7 @@ struct Arg : TaskArgsTpl { ((tensors_[tensor_count_] = &args, tags_[tensor_count_] = TensorArgType::INOUT, tensor_count_++), ...); } - /// No-dependency existing tensor: skips OverlapMap lookup, depends on creator only. + /// Existing tensor with creator retention but no automatic TensorMap lookup or publication. template void add_no_dep(Args &&...args) { assert_add_tensor_args(); @@ -739,5 +747,3 @@ struct ChipTaskArgs : Arg { } } }; - -#endif // SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_ diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp index 3cfa357a49..e27c340ea5 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp @@ -9,7 +9,7 @@ * ----------------------------------------------------------------------------------------------------------- */ /** - * PTO Runtime2 - TensorMap Implementation + * tensormap_and_ringbuffer - TensorMap Implementation * * Implements TensorMap with ring buffer pool, lazy invalidation, * and chain truncation optimization. @@ -56,6 +56,8 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout( PTO2TensorMapLayout layout{}; layout.num_buckets = new_num_buckets; + layout.num_reader_buckets = + new_num_buckets < PTO2_TENSORMAP_READER_NUM_BUCKETS ? new_num_buckets : PTO2_TENSORMAP_READER_NUM_BUCKETS; layout.pool_size = new_pool_size; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { layout.task_window_sizes[r] = new_task_window_sizes[r]; @@ -66,6 +68,11 @@ PTO2TensorMapLayout PTO2TensorMap::reserve_layout( ); layout.off_bucket_epochs = arena.reserve(static_cast(new_num_buckets) * sizeof(uint32_t), alignof(uint32_t)); + layout.off_reader_buckets = arena.reserve( + static_cast(layout.num_reader_buckets) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) + ); + layout.off_reader_bucket_epochs = + arena.reserve(static_cast(layout.num_reader_buckets) * sizeof(uint32_t), alignof(uint32_t)); layout.off_entry_pool = arena.reserve(static_cast(new_pool_size) * sizeof(PTO2TensorMapEntry), alignof(PTO2TensorMapEntry)); layout.off_free_entry_list = @@ -87,12 +94,15 @@ PTO2TensorMap::reserve_layout_default(DeviceArena &arena, const int32_t new_task bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, DeviceArena &arena) { num_buckets = layout.num_buckets; + num_reader_buckets = layout.num_reader_buckets; pool_size = layout.pool_size; // Address arena regions for data writes; do not store these in struct // fields (wire_arena_pointers does that). auto *buckets_arena = static_cast(arena.region_ptr(layout.off_buckets)); auto *bucket_epochs_arena = static_cast(arena.region_ptr(layout.off_bucket_epochs)); + auto *reader_buckets_arena = static_cast(arena.region_ptr(layout.off_reader_buckets)); + auto *reader_bucket_epochs_arena = static_cast(arena.region_ptr(layout.off_reader_bucket_epochs)); auto *entry_pool_arena = static_cast(arena.region_ptr(layout.off_entry_pool)); auto *free_list_arena = static_cast(arena.region_ptr(layout.off_free_entry_list)); @@ -101,6 +111,10 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev buckets_arena[i] = nullptr; bucket_epochs_arena[i] = 0; } + for (int32_t i = 0; i < num_reader_buckets; i++) { + reader_buckets_arena[i] = nullptr; + reader_bucket_epochs_arena[i] = 0; + } // entry_pool: zero-init equivalent to the previous calloc(entry_pool, ...). // The pool's persistent invariant after init is "bucket_index == -1 means @@ -112,7 +126,7 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev entry_pool_arena[i].prev_in_bucket = nullptr; entry_pool_arena[i].next_in_task = nullptr; entry_pool_arena[i].prev_in_task = nullptr; - entry_pool_arena[i].producer_task_id = PTO2TaskId{}; + entry_pool_arena[i].access_task_id = PTO2TaskId{}; } // free_entry_list: zeroed (was calloc'd before); contents become meaningful @@ -121,6 +135,8 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev next_entry_idx = 0; free_num = 0; + reader_used = 0; + reader_high_water = 0; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { auto *heads_arena = static_cast(arena.region_ptr(layout.off_task_entry_heads[r])); @@ -139,13 +155,17 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev void PTO2TensorMap::reset_for_reuse(const PTO2TensorMapLayout &layout) { num_buckets = layout.num_buckets; + num_reader_buckets = layout.num_reader_buckets; pool_size = layout.pool_size; next_entry_idx = 0; free_num = 0; + reader_used = 0; + reader_high_water = 0; current_epoch++; if (current_epoch == 0) { current_epoch = 1; memset(bucket_epochs, 0, static_cast(layout.num_buckets) * sizeof(uint32_t)); + memset(reader_bucket_epochs, 0, static_cast(layout.num_reader_buckets) * sizeof(uint32_t)); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { memset(task_entry_head_epochs[r], 0, static_cast(layout.task_window_sizes[r]) * sizeof(uint32_t)); } @@ -161,6 +181,8 @@ void PTO2TensorMap::reset_for_reuse(const PTO2TensorMapLayout &layout) { void PTO2TensorMap::wire_arena_pointers(const PTO2TensorMapLayout &layout, DeviceArena &arena) { buckets = static_cast(arena.region_ptr(layout.off_buckets)); bucket_epochs = static_cast(arena.region_ptr(layout.off_bucket_epochs)); + reader_buckets = static_cast(arena.region_ptr(layout.off_reader_buckets)); + reader_bucket_epochs = static_cast(arena.region_ptr(layout.off_reader_bucket_epochs)); entry_pool = static_cast(arena.region_ptr(layout.off_entry_pool)); free_entry_list = static_cast(arena.region_ptr(layout.off_free_entry_list)); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { @@ -175,6 +197,8 @@ void PTO2TensorMap::destroy() { // a recycled allocation. buckets = nullptr; bucket_epochs = nullptr; + reader_buckets = nullptr; + reader_bucket_epochs = nullptr; entry_pool = nullptr; free_entry_list = nullptr; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { @@ -190,10 +214,32 @@ void PTO2TensorMap::destroy() { void PTO2TensorMap::print_stats() { int32_t valid = 0; int32_t stale = 0; - int32_t empty_buckets = 0; - int32_t max_chain = 0; - int64_t total_chain = 0; - int32_t non_empty_buckets = 0; + + struct BucketStats { + int32_t empty{0}; + int32_t non_empty{0}; + int32_t max_chain{0}; + int64_t total_chain{0}; + }; + auto collect_bucket_stats = [&](PTO2TensorMapEntry **index, uint32_t *epochs, int32_t bucket_count) { + BucketStats stats; + for (int32_t b = 0; b < bucket_count; b++) { + int32_t chain_len = 0; + if (epochs[b] == current_epoch) { + for (auto *entry = index[b]; entry != nullptr; entry = entry->next_in_bucket) { + chain_len++; + } + } + if (chain_len == 0) { + stats.empty++; + } else { + stats.non_empty++; + stats.total_chain += chain_len; + stats.max_chain = std::max(stats.max_chain, chain_len); + } + } + return stats; + }; // Count entries for (int32_t i = 0; i < pool_size; i++) { @@ -206,41 +252,30 @@ void PTO2TensorMap::print_stats() { } } - // Count bucket stats - for (int32_t b = 0; b < num_buckets; b++) { - int32_t chain_len = 0; - auto cur_entry = buckets[b]; - - while (cur_entry != nullptr) { - chain_len++; - cur_entry = cur_entry->next_in_bucket; - } - - if (chain_len == 0) { - empty_buckets++; - } else { - non_empty_buckets++; - total_chain += chain_len; - if (chain_len > max_chain) { - max_chain = chain_len; - } - } - } - - LOG_DEBUG("=== TensorMap Statistics ==="); - LOG_DEBUG("Pool size: %d", pool_size); - LOG_DEBUG("Pool next entry idx: %d", next_entry_idx); - LOG_DEBUG("Pool free_num: %d", free_num); - LOG_DEBUG("Num buckets: %d", num_buckets); - LOG_DEBUG("Valid entries: %d", valid); - LOG_DEBUG("Stale entries: %d", stale); - LOG_DEBUG("Empty buckets: %d", empty_buckets); - LOG_DEBUG("Max chain len: %d", max_chain); - LOG_DEBUG("Avg chain len: %.2f", non_empty_buckets > 0 ? (float)total_chain / non_empty_buckets : 0); + const BucketStats writer_stats = collect_bucket_stats(buckets, bucket_epochs, num_buckets); + const BucketStats reader_stats = collect_bucket_stats(reader_buckets, reader_bucket_epochs, num_reader_buckets); + + LOG_INFO("=== TensorMap Statistics ==="); + LOG_INFO("Pool size: %d", pool_size); + LOG_INFO("Pool next entry idx: %d", next_entry_idx); + LOG_INFO("Pool free_num: %d", free_num); + LOG_INFO("Reader entries: live=%d high_water=%d", reader_used, reader_high_water); + LOG_INFO("Writer entries: live=%d", current_writers()); + LOG_INFO("Buckets: writer=%d reader=%d", num_buckets, num_reader_buckets); + LOG_INFO("Valid entries: %d", valid); + LOG_INFO("Stale entries: %d", stale); + LOG_INFO( + "Writer buckets: empty=%d max_chain=%d avg_chain=%.2f", writer_stats.empty, writer_stats.max_chain, + writer_stats.non_empty > 0 ? (float)writer_stats.total_chain / writer_stats.non_empty : 0 + ); + LOG_INFO( + "Reader buckets: empty=%d max_chain=%d avg_chain=%.2f", reader_stats.empty, reader_stats.max_chain, + reader_stats.non_empty > 0 ? (float)reader_stats.total_chain / reader_stats.non_empty : 0 + ); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - LOG_DEBUG("Last task alive[%d]: %d", r, last_task_alives[r]); + LOG_INFO("Last task alive[%d]: %d", r, last_task_alives[r]); } - LOG_DEBUG("============================"); + LOG_INFO("============================"); } int32_t PTO2TensorMap::valid_count() { diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index 3a3045c9f1..f84d69a1d8 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -1224,7 +1224,8 @@ void Orchestrator::infer_deps( } TensorArgType tag = a.tag(i); switch (tag) { - case TensorArgType::INPUT: { + case TensorArgType::INPUT: + case TensorArgType::TRACKED_INPUT: { TaskSlot prod = tensormap_->lookup(run_id, key); if (prod != INVALID_SLOT) add_unique_producer(prod); break; diff --git a/src/common/host_build_graph/docs/GRAPH_EXECUTION.md b/src/common/host_build_graph/docs/GRAPH_EXECUTION.md index 771548716d..4dd1106fd2 100644 --- a/src/common/host_build_graph/docs/GRAPH_EXECUTION.md +++ b/src/common/host_build_graph/docs/GRAPH_EXECUTION.md @@ -110,19 +110,20 @@ with the unmodified boundary value. - Boundary ChipTensor shape, stride, dtype, size, direction, contiguity, and alias partition must match the first invocation. - Internal task scalars with no boundary source are fixed Definition data. -- Boundary storage is caller-owned. `INPUT`, `INOUT`, `OUTPUT_EXISTING`, and - `NO_DEP` are supported. A boundary `TensorCreateInfo` tagged `OUTPUT` is not. +- Boundary storage is caller-owned. `INPUT`, `TRACKED_INPUT`, `INOUT`, + `OUTPUT_EXISTING`, and `NO_DEP` are supported. A boundary `TensorCreateInfo` + tagged `OUTPUT` is not. +- An internal node tagged `TRACKED_INPUT` forces the ordinary path because the + compact Definition cannot publish an internal reader for a later WAR lookup. - Early-resolve hints apply while recording the first invocation. Replayed internal nodes use the saved completion topology without the hint. - A recorded task may depend on a Graph-external producer when that producer is the creator of a boundary ChipTensor. The outer Graph owns that dependency on replay; arbitrary cross-boundary explicit dependencies remain unsupported. -Structural or alias mismatch logs a warning and executes the Graph function -normally for that invocation. It never reuses heap offsets recorded for a -different shape. Debug builds also assert at these unsupported boundaries so -development catches a violated fixed-shape contract immediately; the ordinary -path remains the defensive release-build behavior. +Structural, alias, or unsupported-access mismatch logs a warning and executes +the Graph function normally for that invocation. It never reuses heap offsets +recorded for a different shape. ## Qwen decoder-layer example diff --git a/src/common/task_interface/task_args.h b/src/common/task_interface/task_args.h index 2276c7fbb7..c0f57f0f3b 100644 --- a/src/common/task_interface/task_args.h +++ b/src/common/task_interface/task_args.h @@ -302,7 +302,7 @@ inline TaskArgsView read_blob(const uint8_t *src, size_t capacity) { // ============================================================================ // access ⊆ granted: an arg's TensorArgType may only request what the backing grants. -// INPUT -> READ, OUTPUT_EXISTING -> WRITE, INOUT -> READWRITE; READWRITE grants everything. +// INPUT/TRACKED_INPUT -> READ, OUTPUT_EXISTING -> WRITE, INOUT -> READWRITE; READWRITE grants everything. // NO_DEP / OUTPUT are unconstrained. // Catches e.g. a READ-only copy-on-write backing tagged OUTPUT_EXISTING, whose writes in a forked // child would silently never reach the parent. @@ -312,6 +312,7 @@ inline bool access_permits(uint8_t granted, TensorArgType tag) { }; switch (tag) { case TensorArgType::INPUT: + case TensorArgType::TRACKED_INPUT: return granted_has(AccessMode::READ); case TensorArgType::OUTPUT_EXISTING: return granted_has(AccessMode::WRITE); diff --git a/src/common/task_interface/tensor.h b/src/common/task_interface/tensor.h index 5b3d4e14c7..601098df15 100644 --- a/src/common/task_interface/tensor.h +++ b/src/common/task_interface/tensor.h @@ -44,8 +44,9 @@ enum class TensorArgType : int32_t { INPUT = 0, // Read-only input buffer OUTPUT = 1, // Write-only output buffer (runtime allocates) INOUT = 2, // Read-then-write: modifier for downstream - OUTPUT_EXISTING = 3, // Write-only existing tensor: skips OverlapMap lookup, depends on creator + OUTPUT_EXISTING = 3, // Write-only existing tensor: queries readers and publishes a writer NO_DEP = 4, // No-dependency existing tensor: skips OverlapMap lookup, no publish + TRACKED_INPUT = 5, // Read-only input published for later WAR dependency lookup }; // `OverlapStatus` / `Segment` (overlap geometry) live in the runtime diff --git a/tests/st/a2a3/host_build_graph/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp b/tests/st/a2a3/host_build_graph/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp index 79ea60b041..43515babbb 100644 --- a/tests/st/a2a3/host_build_graph/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp +++ b/tests/st/a2a3/host_build_graph/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp @@ -31,18 +31,18 @@ #define FUNC_SPMD_MIX_AIV0 1 #define FUNC_SPMD_MIX_AIV1 2 -// PTO2LaunchSpec spells the SPMD block-count setter differently per arch -// (a5: set_core_num, a2a3: set_block_num) for the same field. Bridge it so this -// one fixture compiles on both; keyed off the arch's pto_types.h include guard, -// which pto_orchestration_api.h pulls in transitively. -static inline void set_block_count(CoreTaskArgs &args, int16_t n) { -#if defined(SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_) - args.launch_spec.set_core_num(n); -#else - args.launch_spec.set_block_num(n); -#endif +template +static inline auto set_spmd_count(Spec &spec, int16_t n) -> decltype(spec.set_block_num(n), void()) { + spec.set_block_num(n); } +template +static inline auto set_spmd_count(Spec &spec, int16_t n) -> decltype(spec.set_core_num(n), void()) { + spec.set_core_num(n); +} + +static inline void set_block_count(CoreTaskArgs &args, int16_t n) { set_spmd_count(args.launch_spec, n); } + extern "C" { __attribute__((visibility("default"))) PTO2OrchestrationConfig diff --git a/tests/st/a2a3/host_build_graph/reader_fanout/kernels/aiv/reader.cpp b/tests/st/a2a3/host_build_graph/reader_fanout/kernels/aiv/reader.cpp new file mode 100644 index 0000000000..50aaf92531 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/reader_fanout/kernels/aiv/reader.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/reader_kernel.h" diff --git a/tests/st/a2a3/host_build_graph/reader_fanout/kernels/aiv/writer.cpp b/tests/st/a2a3/host_build_graph/reader_fanout/kernels/aiv/writer.cpp new file mode 100644 index 0000000000..c71b7674e8 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/reader_fanout/kernels/aiv/writer.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/writer_kernel.h" diff --git a/tests/st/a2a3/host_build_graph/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp b/tests/st/a2a3/host_build_graph/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp new file mode 100644 index 0000000000..2dfdf36578 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/reader_fanout_orch.h" diff --git a/tests/st/a2a3/host_build_graph/reader_fanout/test_reader_fanout.py b/tests/st/a2a3/host_build_graph/reader_fanout/test_reader_fanout.py new file mode 100644 index 0000000000..dffbf5ed31 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/reader_fanout/test_reader_fanout.py @@ -0,0 +1,83 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Reader fan-out followed by a pure overwrite writer (WAR regression).""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, TensorArg, scene_test + + +@scene_test(level=2, runtime="host_build_graph") +class TestReaderFanoutHbgA2A3(SceneTestCase): + RTOL = 0 + ATOL = 0 + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/reader_fanout_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.OUT, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aiv/reader.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/writer.cpp", + "core_type": "aiv", + "signature": [D.OUT, D.OUT], + }, + ], + } + CASES = [ + { + "name": "Correctness256", + "platforms": ["a2a3sim", "a2a3"], + "params": {"reader_count": 256, "spin_iters": 100000, "host_write": False}, + }, + { + "name": "HostWrite256", + "platforms": ["a2a3sim", "a2a3"], + "params": {"reader_count": 256, "spin_iters": 100000, "host_write": True}, + }, + *[ + { + "name": f"Fanout{count}", + "platforms": ["a2a3sim", "a2a3"], + "manual": True, + "params": {"reader_count": count, "spin_iters": 0, "host_write": False}, + } + for count in (1, 8, 32, 64, 128, 256) + ], + ] + + def generate_args(self, params): + return TaskArgsBuilder( + TensorArg("x", torch.ones(16, dtype=torch.float32)), + TensorArg("reader_outputs", torch.zeros((257, 16), dtype=torch.float32)), + TensorArg("padding", torch.zeros(1, dtype=torch.float32)), + Scalar("reader_count", int(params["reader_count"])), + Scalar("spin_iters", int(params["spin_iters"])), + Scalar("host_write", int(params["host_write"])), + ) + + def compute_golden(self, args, params): + args.reader_outputs[: params["reader_count"]] = 1.0 + args.reader_outputs[params["reader_count"]] = 2.0 + args.x[0] = 2.0 + args.padding[0] = 3.0 + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/host_build_graph/war_regression/kernels/aiv/reader.cpp b/tests/st/a2a3/host_build_graph/war_regression/kernels/aiv/reader.cpp new file mode 100644 index 0000000000..092843983f --- /dev/null +++ b/tests/st/a2a3/host_build_graph/war_regression/kernels/aiv/reader.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_reader_kernel.h" diff --git a/tests/st/a2a3/host_build_graph/war_regression/kernels/aiv/writer.cpp b/tests/st/a2a3/host_build_graph/war_regression/kernels/aiv/writer.cpp new file mode 100644 index 0000000000..750e1ff95b --- /dev/null +++ b/tests/st/a2a3/host_build_graph/war_regression/kernels/aiv/writer.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_writer_kernel.h" diff --git a/tests/st/a2a3/host_build_graph/war_regression/kernels/orchestration/war_regression_orch.cpp b/tests/st/a2a3/host_build_graph/war_regression/kernels/orchestration/war_regression_orch.cpp new file mode 100644 index 0000000000..0d60f099d4 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/war_regression/kernels/orchestration/war_regression_orch.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_regression_orch.h" diff --git a/tests/st/a2a3/host_build_graph/war_regression/test_war_regression.py b/tests/st/a2a3/host_build_graph/war_regression/test_war_regression.py new file mode 100644 index 0000000000..e3be893c5b --- /dev/null +++ b/tests/st/a2a3/host_build_graph/war_regression/test_war_regression.py @@ -0,0 +1,84 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""A2/A3 HBG loop-carried write-after-read regression.""" + +import json +import time + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, TensorArg, scene_test +from simpler_setup.scene_test import _outputs_dir, _sanitize_for_filename + + +@scene_test(level=2, runtime="host_build_graph") +class TestWarRegressionHbgA2A3(SceneTestCase): + RTOL = 0 + ATOL = 0 + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/war_regression_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.OUT], + }, + "incores": [ + {"func_id": 0, "source": "kernels/aiv/writer.cpp", "core_type": "aiv", "signature": [D.INOUT]}, + {"func_id": 1, "source": "kernels/aiv/reader.cpp", "core_type": "aiv", "signature": [D.IN, D.OUT]}, + ], + } + CASES = [ + { + "name": "LoopCarried", + "platforms": ["a2a3sim", "a2a3"], + "params": {"cross_ring": False, "war_pred": 1, "war_succ": 2}, + }, + ] + + def generate_args(self, params): + return TaskArgsBuilder( + TensorArg("buffer", torch.zeros(16, dtype=torch.float32)), + TensorArg("outputs", torch.zeros((2, 16), dtype=torch.float32)), + Scalar("spin_iters", 100000), + Scalar("cross_ring", int(params["cross_ring"])), + ) + + def compute_golden(self, args, params): + args.buffer[:] = 2.0 + args.outputs[0] = 1.0 + args.outputs[1] = 2.0 + + def test_run(self, st_platform, st_worker, request): + run_marker = int(time.time()) + super().test_run(st_platform, st_worker, request) + if not self._effective_enable_dep_gen(request): + return + for case in self.CASES: + if st_platform in case.get("platforms", []): + self._validate_war_edge(case, run_marker) + + def _validate_war_edge(self, case, run_marker): + label = _sanitize_for_filename(f"TestWarRegressionHbgA2A3_{case['name']}") + matches = [p for p in _outputs_dir().glob(f"{label}_*") if p.stat().st_mtime >= run_marker] + assert matches, f"no dep_gen output for {case['name']}" + deps_path = max(matches, key=lambda p: p.stat().st_mtime) / "deps.json" + assert deps_path.exists(), f"missing {deps_path}" + edges = json.loads(deps_path.read_text()).get("edges", []) + expected = (case["params"]["war_pred"], case["params"]["war_succ"]) + war_edges = { + (int(edge["pred"]), int(edge["succ"])) + for edge in edges + if edge.get("source") == "tensormap" and edge.get("hazard") == "WAR" and edge.get("access_kind") == "READER" + } + assert expected in war_edges, f"missing READER WAR edge {expected}; got {war_edges}" + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp index 79ea60b041..43515babbb 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp +++ b/tests/st/a2a3/tensormap_and_ringbuffer/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp @@ -31,18 +31,18 @@ #define FUNC_SPMD_MIX_AIV0 1 #define FUNC_SPMD_MIX_AIV1 2 -// PTO2LaunchSpec spells the SPMD block-count setter differently per arch -// (a5: set_core_num, a2a3: set_block_num) for the same field. Bridge it so this -// one fixture compiles on both; keyed off the arch's pto_types.h include guard, -// which pto_orchestration_api.h pulls in transitively. -static inline void set_block_count(CoreTaskArgs &args, int16_t n) { -#if defined(SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_) - args.launch_spec.set_core_num(n); -#else - args.launch_spec.set_block_num(n); -#endif +template +static inline auto set_spmd_count(Spec &spec, int16_t n) -> decltype(spec.set_block_num(n), void()) { + spec.set_block_num(n); } +template +static inline auto set_spmd_count(Spec &spec, int16_t n) -> decltype(spec.set_core_num(n), void()) { + spec.set_core_num(n); +} + +static inline void set_block_count(CoreTaskArgs &args, int16_t n) { set_spmd_count(args.launch_spec, n); } + extern "C" { __attribute__((visibility("default"))) PTO2OrchestrationConfig diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/reader.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/reader.cpp new file mode 100644 index 0000000000..50aaf92531 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/reader.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/reader_kernel.h" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/writer.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/writer.cpp new file mode 100644 index 0000000000..c71b7674e8 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/writer.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/writer_kernel.h" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp new file mode 100644 index 0000000000..2dfdf36578 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/reader_fanout_orch.h" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/test_reader_fanout.py b/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/test_reader_fanout.py new file mode 100644 index 0000000000..696b16a765 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/reader_fanout/test_reader_fanout.py @@ -0,0 +1,83 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Reader fan-out followed by a pure overwrite writer (WAR regression).""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, TensorArg, scene_test + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestReaderFanoutTmrA2A3(SceneTestCase): + RTOL = 0 + ATOL = 0 + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/reader_fanout_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.OUT, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aiv/reader.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/writer.cpp", + "core_type": "aiv", + "signature": [D.OUT, D.OUT], + }, + ], + } + CASES = [ + { + "name": "Correctness256", + "platforms": ["a2a3sim", "a2a3"], + "params": {"reader_count": 256, "spin_iters": 100000, "host_write": False}, + }, + { + "name": "HostWrite256", + "platforms": ["a2a3sim", "a2a3"], + "params": {"reader_count": 256, "spin_iters": 100000, "host_write": True}, + }, + *[ + { + "name": f"Fanout{count}", + "platforms": ["a2a3sim", "a2a3"], + "manual": True, + "params": {"reader_count": count, "spin_iters": 0, "host_write": False}, + } + for count in (1, 8, 32, 64, 128, 256) + ], + ] + + def generate_args(self, params): + return TaskArgsBuilder( + TensorArg("x", torch.ones(16, dtype=torch.float32)), + TensorArg("reader_outputs", torch.zeros((257, 16), dtype=torch.float32)), + TensorArg("padding", torch.zeros(1, dtype=torch.float32)), + Scalar("reader_count", int(params["reader_count"])), + Scalar("spin_iters", int(params["spin_iters"])), + Scalar("host_write", int(params["host_write"])), + ) + + def compute_golden(self, args, params): + args.reader_outputs[: params["reader_count"]] = 1.0 + args.reader_outputs[params["reader_count"]] = 2.0 + args.x[0] = 2.0 + args.padding[0] = 3.0 + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/aiv/reader.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/aiv/reader.cpp new file mode 100644 index 0000000000..092843983f --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/aiv/reader.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_reader_kernel.h" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/aiv/writer.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/aiv/writer.cpp new file mode 100644 index 0000000000..750e1ff95b --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/aiv/writer.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_writer_kernel.h" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/orchestration/war_regression_orch.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/orchestration/war_regression_orch.cpp new file mode 100644 index 0000000000..0d60f099d4 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/kernels/orchestration/war_regression_orch.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_regression_orch.h" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/test_war_regression.py b/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/test_war_regression.py new file mode 100644 index 0000000000..d71355db05 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/war_regression/test_war_regression.py @@ -0,0 +1,89 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""A2/A3 loop-carried and cross-ring write-after-read regressions.""" + +import json +import time + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, TensorArg, scene_test +from simpler_setup.scene_test import _outputs_dir, _sanitize_for_filename + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestWarRegressionTmrA2A3(SceneTestCase): + RTOL = 0 + ATOL = 0 + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/war_regression_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.OUT], + }, + "incores": [ + {"func_id": 0, "source": "kernels/aiv/writer.cpp", "core_type": "aiv", "signature": [D.INOUT]}, + {"func_id": 1, "source": "kernels/aiv/reader.cpp", "core_type": "aiv", "signature": [D.IN, D.OUT]}, + ], + } + CASES = [ + { + "name": "LoopCarried", + "platforms": ["a2a3sim", "a2a3"], + "params": {"cross_ring": False, "war_pred": 1, "war_succ": 2}, + }, + { + "name": "CrossRing", + "platforms": ["a2a3sim", "a2a3"], + "params": {"cross_ring": True, "war_pred": 1 << 32, "war_succ": 1}, + }, + ] + + def generate_args(self, params): + return TaskArgsBuilder( + TensorArg("buffer", torch.zeros(16, dtype=torch.float32)), + TensorArg("outputs", torch.zeros((2, 16), dtype=torch.float32)), + Scalar("spin_iters", 100000), + Scalar("cross_ring", int(params["cross_ring"])), + ) + + def compute_golden(self, args, params): + args.buffer[:] = 2.0 + args.outputs[0] = 1.0 + args.outputs[1] = 2.0 + + def test_run(self, st_platform, st_worker, request): + run_marker = int(time.time()) + super().test_run(st_platform, st_worker, request) + if not self._effective_enable_dep_gen(request): + return + for case in self.CASES: + if st_platform in case.get("platforms", []): + self._validate_war_edge(case, run_marker) + + def _validate_war_edge(self, case, run_marker): + label = _sanitize_for_filename(f"TestWarRegressionTmrA2A3_{case['name']}") + matches = [p for p in _outputs_dir().glob(f"{label}_*") if p.stat().st_mtime >= run_marker] + assert matches, f"no dep_gen output for {case['name']}" + deps_path = max(matches, key=lambda p: p.stat().st_mtime) / "deps.json" + assert deps_path.exists(), f"missing {deps_path}" + edges = json.loads(deps_path.read_text()).get("edges", []) + expected = (case["params"]["war_pred"], case["params"]["war_succ"]) + war_edges = { + (int(edge["pred"]), int(edge["succ"])) + for edge in edges + if edge.get("source") == "tensormap" and edge.get("hazard") == "WAR" and edge.get("access_kind") == "READER" + } + assert expected in war_edges, f"missing READER WAR edge {expected}; got {war_edges}" + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/host_build_graph/reader_fanout/kernels/aiv/reader.cpp b/tests/st/a5/host_build_graph/reader_fanout/kernels/aiv/reader.cpp new file mode 100644 index 0000000000..50aaf92531 --- /dev/null +++ b/tests/st/a5/host_build_graph/reader_fanout/kernels/aiv/reader.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/reader_kernel.h" diff --git a/tests/st/a5/host_build_graph/reader_fanout/kernels/aiv/writer.cpp b/tests/st/a5/host_build_graph/reader_fanout/kernels/aiv/writer.cpp new file mode 100644 index 0000000000..c71b7674e8 --- /dev/null +++ b/tests/st/a5/host_build_graph/reader_fanout/kernels/aiv/writer.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/writer_kernel.h" diff --git a/tests/st/a5/host_build_graph/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp b/tests/st/a5/host_build_graph/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp new file mode 100644 index 0000000000..2dfdf36578 --- /dev/null +++ b/tests/st/a5/host_build_graph/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/reader_fanout_orch.h" diff --git a/tests/st/a5/host_build_graph/reader_fanout/test_reader_fanout.py b/tests/st/a5/host_build_graph/reader_fanout/test_reader_fanout.py new file mode 100644 index 0000000000..edc5d7f224 --- /dev/null +++ b/tests/st/a5/host_build_graph/reader_fanout/test_reader_fanout.py @@ -0,0 +1,88 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Reader fan-out followed by a pure overwrite writer (WAR regression).""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, TensorArg, scene_test + + +@scene_test(level=2, runtime="host_build_graph") +class TestReaderFanoutHbgA5(SceneTestCase): + RTOL = 0 + ATOL = 0 + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/reader_fanout_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.OUT, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aiv/reader.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/writer.cpp", + "core_type": "aiv", + "signature": [D.OUT, D.OUT], + }, + ], + } + CASES = [ + { + "name": "Correctness256", + "platforms": ["a5sim", "a5"], + "params": {"reader_count": 256, "spin_iters": 100000, "host_write": False}, + }, + { + "name": "HostWrite256", + "platforms": ["a5sim", "a5"], + "params": {"reader_count": 256, "spin_iters": 100000, "host_write": True}, + }, + { + "name": "HostRead256", + "platforms": ["a5sim", "a5"], + "params": {"reader_count": 256, "spin_iters": 100000, "host_write": 2}, + }, + *[ + { + "name": f"Fanout{count}", + "platforms": ["a5sim", "a5"], + "manual": True, + "params": {"reader_count": count, "spin_iters": 0, "host_write": False}, + } + for count in (1, 8, 32, 64, 128, 256) + ], + ] + + def generate_args(self, params): + return TaskArgsBuilder( + TensorArg("x", torch.ones(16, dtype=torch.float32)), + TensorArg("reader_outputs", torch.zeros((257, 16), dtype=torch.float32)), + TensorArg("padding", torch.zeros(1, dtype=torch.float32)), + Scalar("reader_count", int(params["reader_count"])), + Scalar("spin_iters", int(params["spin_iters"])), + Scalar("host_write", int(params["host_write"])), + ) + + def compute_golden(self, args, params): + args.reader_outputs[: params["reader_count"]] = 1.0 + args.reader_outputs[params["reader_count"]] = 2.0 + args.x[0] = 2.0 + args.padding[0] = 3.0 + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/host_build_graph/war_regression/kernels/aiv/reader.cpp b/tests/st/a5/host_build_graph/war_regression/kernels/aiv/reader.cpp new file mode 100644 index 0000000000..092843983f --- /dev/null +++ b/tests/st/a5/host_build_graph/war_regression/kernels/aiv/reader.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_reader_kernel.h" diff --git a/tests/st/a5/host_build_graph/war_regression/kernels/aiv/writer.cpp b/tests/st/a5/host_build_graph/war_regression/kernels/aiv/writer.cpp new file mode 100644 index 0000000000..750e1ff95b --- /dev/null +++ b/tests/st/a5/host_build_graph/war_regression/kernels/aiv/writer.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_writer_kernel.h" diff --git a/tests/st/a5/host_build_graph/war_regression/kernels/orchestration/war_regression_orch.cpp b/tests/st/a5/host_build_graph/war_regression/kernels/orchestration/war_regression_orch.cpp new file mode 100644 index 0000000000..0d60f099d4 --- /dev/null +++ b/tests/st/a5/host_build_graph/war_regression/kernels/orchestration/war_regression_orch.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_regression_orch.h" diff --git a/tests/st/a5/host_build_graph/war_regression/test_war_regression.py b/tests/st/a5/host_build_graph/war_regression/test_war_regression.py new file mode 100644 index 0000000000..b5c1d20e3b --- /dev/null +++ b/tests/st/a5/host_build_graph/war_regression/test_war_regression.py @@ -0,0 +1,84 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""A5 HBG loop-carried write-after-read regression.""" + +import json +import time + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, TensorArg, scene_test +from simpler_setup.scene_test import _outputs_dir, _sanitize_for_filename + + +@scene_test(level=2, runtime="host_build_graph") +class TestWarRegressionHbgA5(SceneTestCase): + RTOL = 0 + ATOL = 0 + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/war_regression_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.OUT], + }, + "incores": [ + {"func_id": 0, "source": "kernels/aiv/writer.cpp", "core_type": "aiv", "signature": [D.INOUT]}, + {"func_id": 1, "source": "kernels/aiv/reader.cpp", "core_type": "aiv", "signature": [D.IN, D.OUT]}, + ], + } + CASES = [ + { + "name": "LoopCarried", + "platforms": ["a5sim", "a5"], + "params": {"cross_ring": False, "war_pred": 1, "war_succ": 2}, + }, + ] + + def generate_args(self, params): + return TaskArgsBuilder( + TensorArg("buffer", torch.zeros(16, dtype=torch.float32)), + TensorArg("outputs", torch.zeros((2, 16), dtype=torch.float32)), + Scalar("spin_iters", 100000), + Scalar("cross_ring", int(params["cross_ring"])), + ) + + def compute_golden(self, args, params): + args.buffer[:] = 2.0 + args.outputs[0] = 1.0 + args.outputs[1] = 2.0 + + def test_run(self, st_platform, st_worker, request): + run_marker = int(time.time()) + super().test_run(st_platform, st_worker, request) + if not self._effective_enable_dep_gen(request): + return + for case in self.CASES: + if st_platform in case.get("platforms", []): + self._validate_war_edge(case, run_marker) + + def _validate_war_edge(self, case, run_marker): + label = _sanitize_for_filename(f"TestWarRegressionHbgA5_{case['name']}") + matches = [p for p in _outputs_dir().glob(f"{label}_*") if p.stat().st_mtime >= run_marker] + assert matches, f"no dep_gen output for {case['name']}" + deps_path = max(matches, key=lambda p: p.stat().st_mtime) / "deps.json" + assert deps_path.exists(), f"missing {deps_path}" + edges = json.loads(deps_path.read_text()).get("edges", []) + expected = (case["params"]["war_pred"], case["params"]["war_succ"]) + war_edges = { + (int(edge["pred"]), int(edge["succ"])) + for edge in edges + if edge.get("source") == "tensormap" and edge.get("hazard") == "WAR" and edge.get("access_kind") == "READER" + } + assert expected in war_edges, f"missing READER WAR edge {expected}; got {war_edges}" + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp index 79ea60b041..43515babbb 100644 --- a/tests/st/a5/tensormap_and_ringbuffer/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp +++ b/tests/st/a5/tensormap_and_ringbuffer/available_aicore_counts/kernels/orchestration/available_aicore_counts_orch.cpp @@ -31,18 +31,18 @@ #define FUNC_SPMD_MIX_AIV0 1 #define FUNC_SPMD_MIX_AIV1 2 -// PTO2LaunchSpec spells the SPMD block-count setter differently per arch -// (a5: set_core_num, a2a3: set_block_num) for the same field. Bridge it so this -// one fixture compiles on both; keyed off the arch's pto_types.h include guard, -// which pto_orchestration_api.h pulls in transitively. -static inline void set_block_count(CoreTaskArgs &args, int16_t n) { -#if defined(SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_) - args.launch_spec.set_core_num(n); -#else - args.launch_spec.set_block_num(n); -#endif +template +static inline auto set_spmd_count(Spec &spec, int16_t n) -> decltype(spec.set_block_num(n), void()) { + spec.set_block_num(n); } +template +static inline auto set_spmd_count(Spec &spec, int16_t n) -> decltype(spec.set_core_num(n), void()) { + spec.set_core_num(n); +} + +static inline void set_block_count(CoreTaskArgs &args, int16_t n) { set_spmd_count(args.launch_spec, n); } + extern "C" { __attribute__((visibility("default"))) PTO2OrchestrationConfig diff --git a/tests/st/a5/tensormap_and_ringbuffer/mx_fp_gemm/kernels/aic/kernel_fp_gemm.cpp b/tests/st/a5/tensormap_and_ringbuffer/mx_fp_gemm/kernels/aic/kernel_fp_gemm.cpp index 61e599ecf1..2512f7c10e 100644 --- a/tests/st/a5/tensormap_and_ringbuffer/mx_fp_gemm/kernels/aic/kernel_fp_gemm.cpp +++ b/tests/st/a5/tensormap_and_ringbuffer/mx_fp_gemm/kernels/aic/kernel_fp_gemm.cpp @@ -10,7 +10,7 @@ */ /** - * MXFP8 / MXFP4 matmul via pto-isa ``TMATMUL_MX`` (A5 Cube, onboard only). + * MXFP8 / MXFP4 matmul via pto-isa ``TMATMUL_MX``. * * Fixed tile: M=128, K=64, N=64 (K aligned to 64; scale K = K/32 = 2). * mode (args[5]): @@ -18,8 +18,6 @@ * 1 — MXFP4: A/B float4_e2m1x2 packed GM (logical [M,K]/[K,N]); same scales * * Pattern mirrors pto-isa ``tmatmul_mx_kernel.cpp`` RunTMATMULMX (no bias). - * a5sim is not supported: CPU stub TLOAD lacks MX_A_ZZ / MX_B_NN. - * * Args: [A, As, B, Bs, C, mode] */ @@ -27,8 +25,6 @@ #include #include #include -#include - #include "tensor.h" using namespace pto; diff --git a/tests/st/a5/tensormap_and_ringbuffer/mx_fp_gemm/test_mx_fp_gemm.py b/tests/st/a5/tensormap_and_ringbuffer/mx_fp_gemm/test_mx_fp_gemm.py index 61642e339e..2a7eac703d 100644 --- a/tests/st/a5/tensormap_and_ringbuffer/mx_fp_gemm/test_mx_fp_gemm.py +++ b/tests/st/a5/tensormap_and_ringbuffer/mx_fp_gemm/test_mx_fp_gemm.py @@ -7,7 +7,7 @@ # 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. # ----------------------------------------------------------------------------------------------------------- -"""mx_fp_gemm: A5 MXFP8/MXFP4 ``TMATMUL_MX`` ST via task interface (onboard only). +"""mx_fp_gemm: A5 MXFP8/MXFP4 ``TMATMUL_MX`` ST via task interface. Device AIC kernel calls pto-isa ``TMATMUL_MX`` (not scalar MAC). Host inputs + golden follow pto-isa ``tmatmul_mx`` host-prequant sample: @@ -36,7 +36,7 @@ def _cases(): cases.append( { "name": "MXFP8_TMATMUL_MX", - "platforms": ["a5"], + "platforms": ["a5sim", "a5"], "config": {"aicpu_thread_num": 2, "block_dim": 1}, "params": {"mode": 0}, } @@ -45,7 +45,7 @@ def _cases(): cases.append( { "name": "MXFP4_TMATMUL_MX", - "platforms": ["a5"], + "platforms": ["a5sim", "a5"], "config": {"aicpu_thread_num": 2, "block_dim": 1}, "params": {"mode": 1}, } diff --git a/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/reader.cpp b/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/reader.cpp new file mode 100644 index 0000000000..50aaf92531 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/reader.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/reader_kernel.h" diff --git a/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/writer.cpp b/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/writer.cpp new file mode 100644 index 0000000000..c71b7674e8 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/aiv/writer.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/writer_kernel.h" diff --git a/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp new file mode 100644 index 0000000000..2dfdf36578 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/kernels/orchestration/reader_fanout_orch.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/reader_fanout/reader_fanout_orch.h" diff --git a/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/test_reader_fanout.py b/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/test_reader_fanout.py new file mode 100644 index 0000000000..72e2efa7f1 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/reader_fanout/test_reader_fanout.py @@ -0,0 +1,88 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Reader fan-out followed by a pure overwrite writer (WAR regression).""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, TensorArg, scene_test + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestReaderFanoutTmrA5(SceneTestCase): + RTOL = 0 + ATOL = 0 + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/reader_fanout_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.OUT, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "kernels/aiv/reader.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "kernels/aiv/writer.cpp", + "core_type": "aiv", + "signature": [D.OUT, D.OUT], + }, + ], + } + CASES = [ + { + "name": "Correctness256", + "platforms": ["a5sim", "a5"], + "params": {"reader_count": 256, "spin_iters": 100000, "host_write": False}, + }, + { + "name": "HostWrite256", + "platforms": ["a5sim", "a5"], + "params": {"reader_count": 256, "spin_iters": 100000, "host_write": True}, + }, + { + "name": "HostRead256", + "platforms": ["a5sim", "a5"], + "params": {"reader_count": 256, "spin_iters": 100000, "host_write": 2}, + }, + *[ + { + "name": f"Fanout{count}", + "platforms": ["a5sim", "a5"], + "manual": True, + "params": {"reader_count": count, "spin_iters": 0, "host_write": False}, + } + for count in (1, 8, 32, 64, 128, 256) + ], + ] + + def generate_args(self, params): + return TaskArgsBuilder( + TensorArg("x", torch.ones(16, dtype=torch.float32)), + TensorArg("reader_outputs", torch.zeros((257, 16), dtype=torch.float32)), + TensorArg("padding", torch.zeros(1, dtype=torch.float32)), + Scalar("reader_count", int(params["reader_count"])), + Scalar("spin_iters", int(params["spin_iters"])), + Scalar("host_write", int(params["host_write"])), + ) + + def compute_golden(self, args, params): + args.reader_outputs[: params["reader_count"]] = 1.0 + args.reader_outputs[params["reader_count"]] = 2.0 + args.x[0] = 2.0 + args.padding[0] = 3.0 + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp index 8ae2296dc8..6e91974491 100644 --- a/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp @@ -52,14 +52,18 @@ static constexpr uint32_t PIJ_SLOT_SIZE = MAX_Q_TILE * MAX_BLOCK_SIZE * sizeof(u static constexpr uint32_t OI_SLOT_SIZE = MAX_Q_TILE * HEAD_DIM * sizeof(float); static constexpr uint32_t FIFO_DEPTH = 2; -static inline void set_block_count(CoreTaskArgs &args, int16_t n) { -#if defined(SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_) - args.launch_spec.set_core_num(n); -#else - args.launch_spec.set_block_num(n); -#endif +template +static inline auto set_spmd_count(Spec &spec, int16_t n) -> decltype(spec.set_block_num(n), void()) { + spec.set_block_num(n); } +template +static inline auto set_spmd_count(Spec &spec, int16_t n) -> decltype(spec.set_core_num(n), void()) { + spec.set_core_num(n); +} + +static inline void set_block_count(CoreTaskArgs &args, int16_t n) { set_spmd_count(args.launch_spec, n); } + extern "C" { __attribute__((visibility("default"))) PTO2OrchestrationConfig diff --git a/tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/aiv/reader.cpp b/tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/aiv/reader.cpp new file mode 100644 index 0000000000..092843983f --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/aiv/reader.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_reader_kernel.h" diff --git a/tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/aiv/writer.cpp b/tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/aiv/writer.cpp new file mode 100644 index 0000000000..750e1ff95b --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/aiv/writer.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_writer_kernel.h" diff --git a/tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/orchestration/war_regression_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/orchestration/war_regression_orch.cpp new file mode 100644 index 0000000000..0d60f099d4 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/war_regression/kernels/orchestration/war_regression_orch.cpp @@ -0,0 +1,12 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "../../../../../common/war_regression/war_regression_orch.h" diff --git a/tests/st/a5/tensormap_and_ringbuffer/war_regression/test_war_regression.py b/tests/st/a5/tensormap_and_ringbuffer/war_regression/test_war_regression.py new file mode 100644 index 0000000000..ab0c242e16 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/war_regression/test_war_regression.py @@ -0,0 +1,89 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""A5 loop-carried and cross-ring write-after-read regressions.""" + +import json +import time + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, TensorArg, scene_test +from simpler_setup.scene_test import _outputs_dir, _sanitize_for_filename + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestWarRegressionTmrA5(SceneTestCase): + RTOL = 0 + ATOL = 0 + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/war_regression_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.OUT], + }, + "incores": [ + {"func_id": 0, "source": "kernels/aiv/writer.cpp", "core_type": "aiv", "signature": [D.INOUT]}, + {"func_id": 1, "source": "kernels/aiv/reader.cpp", "core_type": "aiv", "signature": [D.IN, D.OUT]}, + ], + } + CASES = [ + { + "name": "LoopCarried", + "platforms": ["a5sim", "a5"], + "params": {"cross_ring": False, "war_pred": 1, "war_succ": 2}, + }, + { + "name": "CrossRing", + "platforms": ["a5sim", "a5"], + "params": {"cross_ring": True, "war_pred": 1 << 32, "war_succ": 1}, + }, + ] + + def generate_args(self, params): + return TaskArgsBuilder( + TensorArg("buffer", torch.zeros(16, dtype=torch.float32)), + TensorArg("outputs", torch.zeros((2, 16), dtype=torch.float32)), + Scalar("spin_iters", 100000), + Scalar("cross_ring", int(params["cross_ring"])), + ) + + def compute_golden(self, args, params): + args.buffer[:] = 2.0 + args.outputs[0] = 1.0 + args.outputs[1] = 2.0 + + def test_run(self, st_platform, st_worker, request): + run_marker = int(time.time()) + super().test_run(st_platform, st_worker, request) + if not self._effective_enable_dep_gen(request): + return + for case in self.CASES: + if st_platform in case.get("platforms", []): + self._validate_war_edge(case, run_marker) + + def _validate_war_edge(self, case, run_marker): + label = _sanitize_for_filename(f"TestWarRegressionTmrA5_{case['name']}") + matches = [p for p in _outputs_dir().glob(f"{label}_*") if p.stat().st_mtime >= run_marker] + assert matches, f"no dep_gen output for {case['name']}" + deps_path = max(matches, key=lambda p: p.stat().st_mtime) / "deps.json" + assert deps_path.exists(), f"missing {deps_path}" + edges = json.loads(deps_path.read_text()).get("edges", []) + expected = (case["params"]["war_pred"], case["params"]["war_succ"]) + war_edges = { + (int(edge["pred"]), int(edge["succ"])) + for edge in edges + if edge.get("source") == "tensormap" and edge.get("hazard") == "WAR" and edge.get("access_kind") == "READER" + } + assert expected in war_edges, f"missing READER WAR edge {expected}; got {war_edges}" + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/common/reader_fanout/reader_fanout_orch.h b/tests/st/common/reader_fanout/reader_fanout_orch.h new file mode 100644 index 0000000000..e184013e9b --- /dev/null +++ b/tests/st/common/reader_fanout/reader_fanout_orch.h @@ -0,0 +1,102 @@ +/* + * 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 + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +namespace reader_fanout_test { + +constexpr int32_t FUNC_READER = 0; +constexpr int32_t FUNC_WRITER = 1; + +} // namespace reader_fanout_test + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig +aicpu_orchestration_config(const ChipTaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 6, // X, reader outputs, padding, reader_count, spin_iters, host access mode + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const ChipTaskArgs &orch_args) { + const ChipTensor &x = orch_args.tensor(0).ref(); + const ChipTensor &reader_outputs = orch_args.tensor(1).ref(); + const ChipTensor &padding = orch_args.tensor(2).ref(); + const int32_t reader_count = static_cast(orch_args.scalar(0)); + const int32_t spin_iters = static_cast(orch_args.scalar(1)); + const int32_t host_access = static_cast(orch_args.scalar(2)); + + if (reader_count <= 0 || reader_count > static_cast(reader_outputs.shapes[0])) { + rt_report_fatal( + PTO2_ERROR_INVALID_ARGS, "reader_count=%d output_size=%u", reader_count, reader_outputs.shapes[0] + ); + return; + } + + // Each reader writes a disjoint result, but all of them read the same X. + // They therefore depend on a prior X writer (if any), never on each other. + for (int32_t i = 0; i < reader_count; ++i) { + uint32_t shape[] = {1, reader_outputs.shapes[1]}; + uint32_t offset[] = {static_cast(i), 0}; + ChipTensor output = reader_outputs.view(shape, offset); + CoreTaskArgs args; + args.add_tracked_input(x); + args.add_output(output); + args.add_scalar(spin_iters); + rt_submit_aiv_task(reader_fanout_test::FUNC_READER, args); + } + + if (host_access == 1) { + // Host-side writes have the same WAR obligation. They wait for each + // reader task itself, but do not wait for those readers' consumers. + const uint32_t index[] = {0}; + set_tensor_data(x, 1, index, 2.0f); + set_tensor_data(padding, 1, index, 3.0f); + + } else { + if (host_access == 2) { + // A host read waits only for the last writer. The submitted readers + // above are peers, so waiting for them would deadlock HBG before + // device execution starts. + const uint32_t index[] = {0}; + const float observed = get_tensor_data(x, 1, index); + if (observed != 1.0f) { + rt_report_fatal(PTO2_ERROR_INVALID_ARGS, "host read observed %.1f, expected 1.0", observed); + return; + } + } + // A pure device overwrite does not read X, but it must wait for every + // live reader before changing the storage. + CoreTaskArgs writer_args; + writer_args.add_output(x); + writer_args.add_output(padding); + rt_submit_aiv_task(reader_fanout_test::FUNC_WRITER, writer_args); + } + + // A later reader must observe the new writer entry, including when the + // writer is the HBG scheduler-local host-write node. + uint32_t shape[] = {1, reader_outputs.shapes[1]}; + uint32_t offset[] = {static_cast(reader_count), 0}; + ChipTensor verify_output = reader_outputs.view(shape, offset); + CoreTaskArgs verify_args; + verify_args.add_input(x); + verify_args.add_output(verify_output); + verify_args.add_scalar(0); + rt_submit_aiv_task(reader_fanout_test::FUNC_READER, verify_args); +} + +} // extern "C" diff --git a/tests/st/common/reader_fanout/reader_kernel.h b/tests/st/common/reader_fanout/reader_kernel.h new file mode 100644 index 0000000000..4f6e352692 --- /dev/null +++ b/tests/st/common/reader_fanout/reader_kernel.h @@ -0,0 +1,56 @@ +/* + * 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 + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ ChipTensor *input_tensor = reinterpret_cast<__gm__ ChipTensor *>(args[0]); + __gm__ ChipTensor *output_tensor = reinterpret_cast<__gm__ ChipTensor *>(args[1]); + __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; + __gm__ float *output = reinterpret_cast<__gm__ float *>(output_tensor->buffer.addr) + output_tensor->start_offset; + const int32_t spin_iters = static_cast(args[2]); + + volatile int32_t delay = 0; + for (int32_t i = 0; i < spin_iters; ++i) { + delay++; + } + (void)delay; + for (int32_t i = 0; i < 16; ++i) { + output[i] = input[0]; + } + dcci(&output[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/common/reader_fanout/writer_kernel.h b/tests/st/common/reader_fanout/writer_kernel.h new file mode 100644 index 0000000000..53843aab33 --- /dev/null +++ b/tests/st/common/reader_fanout/writer_kernel.h @@ -0,0 +1,51 @@ +/* + * 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 + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ ChipTensor *x_tensor = reinterpret_cast<__gm__ ChipTensor *>(args[0]); + __gm__ ChipTensor *padding_tensor = reinterpret_cast<__gm__ ChipTensor *>(args[1]); + __gm__ float *x = reinterpret_cast<__gm__ float *>(x_tensor->buffer.addr) + x_tensor->start_offset; + __gm__ float *padding = + reinterpret_cast<__gm__ float *>(padding_tensor->buffer.addr) + padding_tensor->start_offset; + + x[0] = 2.0f; + padding[0] = 3.0f; + dcci(&x[0], SINGLE_CACHE_LINE, CACHELINE_OUT); + dcci(&padding[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/common/war_regression/war_reader_kernel.h b/tests/st/common/war_regression/war_reader_kernel.h new file mode 100644 index 0000000000..2e0e84e86b --- /dev/null +++ b/tests/st/common/war_regression/war_reader_kernel.h @@ -0,0 +1,54 @@ +/* + * 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 + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ ChipTensor *input_tensor = reinterpret_cast<__gm__ ChipTensor *>(args[0]); + __gm__ ChipTensor *output_tensor = reinterpret_cast<__gm__ ChipTensor *>(args[1]); + __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; + __gm__ float *output = reinterpret_cast<__gm__ float *>(output_tensor->buffer.addr) + output_tensor->start_offset; + const int32_t spin_iters = static_cast(args[2]); + + volatile int32_t delay = 0; + for (int32_t i = 0; i < spin_iters; ++i) + delay++; + (void)delay; + for (int32_t i = 0; i < 16; ++i) + output[i] = input[i]; + dcci(&output[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/common/war_regression/war_regression_orch.h b/tests/st/common/war_regression/war_regression_orch.h new file mode 100644 index 0000000000..87c1eeaced --- /dev/null +++ b/tests/st/common/war_regression/war_regression_orch.h @@ -0,0 +1,67 @@ +/* + * 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 + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +namespace war_regression_test { +constexpr int32_t FUNC_WRITER = 0; +constexpr int32_t FUNC_READER = 1; +} // namespace war_regression_test + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig +aicpu_orchestration_config(const ChipTaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{.expected_arg_count = 4}; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const ChipTaskArgs &orch_args) { + const ChipTensor &buffer = orch_args.tensor(0).ref(); + const ChipTensor &outputs = orch_args.tensor(1).ref(); + const int32_t spin_iters = static_cast(orch_args.scalar(0)); + const bool cross_ring = static_cast(orch_args.scalar(1)); + uint32_t row_shape[] = {1, 16}; + uint32_t row0_offset[] = {0, 0}; + uint32_t row1_offset[] = {1, 0}; + ChipTensor first_output = outputs.view(row_shape, row0_offset); + ChipTensor second_output = outputs.view(row_shape, row1_offset); + + CoreTaskArgs first_writer_args; + first_writer_args.add_inout(buffer); + first_writer_args.add_scalar(1); + rt_submit_aiv_task(war_regression_test::FUNC_WRITER, first_writer_args); + + if (cross_ring) rt_scope_begin(); + CoreTaskArgs first_reader_args; + first_reader_args.add_tracked_input(buffer); + first_reader_args.add_output(first_output); + first_reader_args.add_scalar(spin_iters); + rt_submit_aiv_task(war_regression_test::FUNC_READER, first_reader_args); + if (cross_ring) rt_scope_end(); + + CoreTaskArgs second_writer_args; + second_writer_args.add_inout(buffer); + second_writer_args.add_scalar(2); + rt_submit_aiv_task(war_regression_test::FUNC_WRITER, second_writer_args); + + CoreTaskArgs second_reader_args; + second_reader_args.add_input(buffer); + second_reader_args.add_output(second_output); + second_reader_args.add_scalar(0); + rt_submit_aiv_task(war_regression_test::FUNC_READER, second_reader_args); +} + +} // extern "C" diff --git a/tests/st/common/war_regression/war_writer_kernel.h b/tests/st/common/war_regression/war_writer_kernel.h new file mode 100644 index 0000000000..9e38976963 --- /dev/null +++ b/tests/st/common/war_regression/war_writer_kernel.h @@ -0,0 +1,47 @@ +/* + * 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 + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ ChipTensor *tensor = reinterpret_cast<__gm__ ChipTensor *>(args[0]); + __gm__ float *data = reinterpret_cast<__gm__ float *>(tensor->buffer.addr) + tensor->start_offset; + const float value = static_cast(args[1]); + for (int32_t i = 0; i < 16; ++i) + data[i] = value; + dcci(&data[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/runtime_fatal_codes/kernels/orchestration/require_sync_start_orch.cpp b/tests/st/runtime_fatal_codes/kernels/orchestration/require_sync_start_orch.cpp index baa7ea75a4..760bbfb75c 100644 --- a/tests/st/runtime_fatal_codes/kernels/orchestration/require_sync_start_orch.cpp +++ b/tests/st/runtime_fatal_codes/kernels/orchestration/require_sync_start_orch.cpp @@ -26,18 +26,18 @@ #define FUNC_NOOP_KERNEL 0 -// PTO2LaunchSpec spells the SPMD block-count setter differently per arch -// (a5: set_core_num, a2a3: set_block_num) for the same field. Bridge it so this -// one fixture compiles on both; keyed off the arch's pto_types.h include guard, -// which pto_orchestration_api.h pulls in transitively. -static inline void set_block_count(CoreTaskArgs &args, int16_t n) { -#if defined(SRC_A5_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_TYPES_H_) - args.launch_spec.set_core_num(n); -#else - args.launch_spec.set_block_num(n); -#endif +template +static inline auto set_spmd_count(Spec &spec, int16_t n) -> decltype(spec.set_block_num(n), void()) { + spec.set_block_num(n); } +template +static inline auto set_spmd_count(Spec &spec, int16_t n) -> decltype(spec.set_core_num(n), void()) { + spec.set_core_num(n); +} + +static inline void set_block_count(CoreTaskArgs &args, int16_t n) { set_spmd_count(args.launch_spec, n); } + extern "C" { __attribute__((visibility("default"))) PTO2OrchestrationConfig diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index 27c9a908ef..bc38ae5761 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -798,6 +798,7 @@ set(HBG_RUNTIME_DIR ${CMAKE_SOURCE_DIR}/../../../src/a2a3/runtime/host_build_gra set(A5_HBG_RUNTIME_DIR ${CMAKE_SOURCE_DIR}/../../../src/a5/runtime/host_build_graph/runtime) add_a2a3_hbg_runtime_test(test_hbg_task_allocator a2a3/test_task_allocator.cpp) add_a2a3_hbg_runtime_test(test_hbg_tensormap a2a3/test_hbg_tensormap.cpp) +add_a5_hbg_runtime_test(test_a5_hbg_tensormap a2a3/test_hbg_tensormap.cpp) add_a2a3_hbg_runtime_test(test_hbg_dep_gen_host_graph a2a3/test_dep_gen_host_graph.cpp) add_a2a3_hbg_runtime_test(test_hbg_tensor_access a2a3/test_hbg_tensor_access.cpp) add_a2a3_hbg_runtime_test(test_hbg_core_tracker common/test_hbg_core_tracker.cpp) @@ -877,6 +878,9 @@ target_sources(test_a5_graph_activation PRIVATE target_sources(test_hbg_tensormap PRIVATE ${CMAKE_SOURCE_DIR}/../../../src/a2a3/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp ) +target_sources(test_a5_hbg_tensormap PRIVATE + ${CMAKE_SOURCE_DIR}/../../../src/a5/runtime/host_build_graph/runtime/shared/pto_tensormap.cpp +) target_sources(test_hbg_dep_gen_host_graph PRIVATE ${CMAKE_SOURCE_DIR}/../../../src/a2a3/runtime/host_build_graph/host/dep_gen_host_graph.cpp ) diff --git a/tests/ut/cpp/a2a3/test_dep_gen_host_graph.cpp b/tests/ut/cpp/a2a3/test_dep_gen_host_graph.cpp index 7d0f53ad0b..3144465254 100644 --- a/tests/ut/cpp/a2a3/test_dep_gen_host_graph.cpp +++ b/tests/ut/cpp/a2a3/test_dep_gen_host_graph.cpp @@ -94,3 +94,57 @@ TEST(DepGenHostGraphTest, BeginCaptureClearsThePreviousRunsGraph) { EXPECT_NE(json.find("\"task_id\":\"202\""), std::string::npos); std::filesystem::remove(path); } + +TEST(DepGenHostGraphTest, TensorMapEdgesIdentifyWarAndReaderAccess) { + const std::filesystem::path path = output_path("war_annotation"); + std::filesystem::remove(path); + uint32_t shape[] = {16}; + ChipTensor tensor = + make_tensor_external(reinterpret_cast(0x1000), shape, 1, DataType::FLOAT32, false, /*version=*/0); + TensorRef tensor_ref; + tensor_ref = &tensor; + TensorArgType arg_type = TensorArgType::OUTPUT_EXISTING; + const int32_t kernel_ids[3] = {1, -1, -1}; + PTO2TensorMapEntry reader{}; + reader.copy_from_tensor(tensor); + reader.access_task_id = PTO2TaskId::make(0, 7); + reader.access_kind = TensorAccessKind::READER; + + dep_gen_host_graph_set_enabled(true); + dep_gen_host_graph_begin_capture(); + dep_gen_host_graph_begin_task(8, false, false, kernel_ids, 1, 1, &tensor_ref, &arg_type); + dep_gen_host_graph_add_tensormap_edge( + reader.access_task_id.raw, 0, tensor, reader, OverlapStatus::COVERED, TensorHazardKind::WAR + ); + dep_gen_host_graph_end_task(); + ASSERT_EQ(dep_gen_host_graph_emit(path.c_str()), 0); + + const std::string json = read_file(path); + EXPECT_NE(json.find("\"hazard\":\"WAR\""), std::string::npos); + EXPECT_NE(json.find("\"access_kind\":\"READER\""), std::string::npos); + std::filesystem::remove(path); +} + +TEST(DepGenHostGraphTest, HostWriteConsumerEdgeIsNotLabeledExplicit) { + const std::filesystem::path path = output_path("host_write_consumer"); + std::filesystem::remove(path); + uint32_t shape[] = {16}; + ChipTensor tensor = + make_tensor_external(reinterpret_cast(0x2000), shape, 1, DataType::FLOAT32, false, /*version=*/0); + TensorRef tensor_ref; + tensor_ref = &tensor; + TensorArgType arg_type = TensorArgType::OUTPUT_EXISTING; + const int32_t kernel_ids[3] = {-1, -1, -1}; + + dep_gen_host_graph_set_enabled(true); + dep_gen_host_graph_begin_capture(); + dep_gen_host_graph_begin_task(9, false, false, kernel_ids, 1, 1, &tensor_ref, &arg_type); + dep_gen_host_graph_add_host_write_consumer_edge(8, 0, tensor); + dep_gen_host_graph_end_task(); + ASSERT_EQ(dep_gen_host_graph_emit(path.c_str()), 0); + + const std::string json = read_file(path); + EXPECT_NE(json.find("\"source\":\"host_write_consumer\""), std::string::npos); + EXPECT_EQ(json.find("\"source\":\"explicit\""), std::string::npos); + std::filesystem::remove(path); +} diff --git a/tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp b/tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp index 7b019204ba..2b2d218fc0 100644 --- a/tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp +++ b/tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp @@ -26,6 +26,7 @@ #include +#include #include #include #include @@ -88,6 +89,40 @@ class HbgSubmitPoisonTest : public ::testing::Test { std::memset(ring.slot_states, POISON, n * sizeof(PTO2TaskSlotState)); std::memset(ring.completion_flags, POISON, n * sizeof(std::atomic)); } + + void verify_reader_fanout(int32_t reader_count) { + uint32_t shape[] = {16}; + std::vector storage(shape[0]); + ChipTensor tensor = make_tensor_external(storage.data(), shape, 1, DataType::FLOAT32, false, /*version=*/0); + std::vector readers; + readers.reserve(reader_count); + + orch.begin_scope(); + for (int32_t i = 0; i < reader_count; ++i) { + CoreTaskArgs args; + args.add_tracked_input(tensor); + TaskOutputTensors reader = orch.submit_dummy_task(args); + ASSERT_TRUE(reader.task_id().is_valid()); + readers.push_back(reader.task_id()); + } + CoreTaskArgs writer_args; + writer_args.add_inout(tensor); + TaskOutputTensors writer = orch.submit_dummy_task(writer_args); + ASSERT_TRUE(writer.task_id().is_valid()); + orch.end_scope(); + + auto &ring = sm_handle->header->ring; + const PTO2TaskPayload &payload = ring.task_payloads[ring.get_slot_by_task_id(writer.task_id().local())]; + ASSERT_EQ(payload.fanin_count, reader_count); + const int32_t expected_spill = std::max(0, reader_count - PTO2_MAX_FANIN); + ASSERT_EQ(payload.fanin_spill_count, expected_spill); + EXPECT_EQ(sched.fanin_spill_top, expected_spill); + for (int32_t i = 0; i < reader_count; ++i) { + SCOPED_TRACE(testing::Message() << "fanin index=" << i); + EXPECT_EQ(sched.fanin_local_id(payload, i), static_cast(readers[reader_count - 1 - i].local())); + } + EXPECT_EQ(sm_handle->header->orch_error_code.load(std::memory_order_acquire), PTO2_ERROR_NONE); + } }; TEST_F(HbgSubmitPoisonTest, EveryDeviceReadFieldIsWrittenOverPoison) { @@ -183,3 +218,62 @@ TEST_F(HbgSubmitPoisonTest, EveryDeviceReadFieldIsWrittenOverPoison) { EXPECT_EQ(cons_pl.fanin_count, 1); EXPECT_EQ(cons_pl.fanin_local_ids[0], static_cast(root.task_id().local())); } + +TEST_F(HbgSubmitPoisonTest, ReaderFanoutBelowInlineBoundaryHasNoSpill) { verify_reader_fanout(64); } + +TEST_F(HbgSubmitPoisonTest, ReaderFanoutAtInlineBoundaryHasNoSpill) { verify_reader_fanout(PTO2_MAX_FANIN); } + +TEST_F(HbgSubmitPoisonTest, ReaderFanoutAboveInlineBoundarySpillsOneEdge) { verify_reader_fanout(PTO2_MAX_FANIN + 1); } + +TEST_F(HbgSubmitPoisonTest, ReaderFanoutBeyondInlineFaninSpillsWithoutDroppingWarEdges) { verify_reader_fanout(256); } + +TEST_F(HbgSubmitPoisonTest, HostWriteWaitsOnlyForConflictingWriterConsumersAndTrackedReaders) { + uint32_t shape[] = {16}; + std::vector storage(shape[0]); + ChipTensor tensor = make_tensor_external(storage.data(), shape, 1, DataType::FLOAT32, false); + + orch.begin_scope(); + CoreTaskArgs unrelated_before_args; + const PTO2TaskId unrelated_before = orch.submit_dummy_task(unrelated_before_args).task_id(); + + CoreTaskArgs writer_args; + writer_args.add_output(tensor); + const PTO2TaskId writer = orch.submit_dummy_task(writer_args).task_id(); + + CoreTaskArgs consumer_args; + consumer_args.add_input(tensor); + const PTO2TaskId consumer = orch.submit_dummy_task(consumer_args).task_id(); + + CoreTaskArgs tracked_reader_args; + tracked_reader_args.add_tracked_input(tensor); + const PTO2TaskId tracked_reader = orch.submit_dummy_task(tracked_reader_args).task_id(); + + CoreTaskArgs unrelated_after_args; + const PTO2TaskId unrelated_after = orch.submit_dummy_task(unrelated_after_args).task_id(); + + orch.submit_host_write(tensor, tensor.buffer.addr, 7, sizeof(uint32_t)); + const int32_t host_write_local = sm_handle->header->ring.fc.current_task_index.load(std::memory_order_acquire) - 1; + orch.submit_host_write(tensor, tensor.buffer.addr, 8, sizeof(uint32_t)); + const int32_t second_host_write_local = + sm_handle->header->ring.fc.current_task_index.load(std::memory_order_acquire) - 1; + orch.end_scope(); + + const auto &payload = sm_handle->header->ring.get_slot_state_by_task_id(host_write_local).payload[0]; + ASSERT_EQ(payload.fanin_count, 3); + std::vector fanins; + for (int32_t i = 0; i < payload.fanin_count; ++i) + fanins.push_back(sched.fanin_local_id(payload, i)); + EXPECT_NE(std::find(fanins.begin(), fanins.end(), static_cast(writer.local())), fanins.end()); + EXPECT_NE(std::find(fanins.begin(), fanins.end(), static_cast(consumer.local())), fanins.end()); + EXPECT_NE(std::find(fanins.begin(), fanins.end(), static_cast(tracked_reader.local())), fanins.end()); + EXPECT_EQ(std::find(fanins.begin(), fanins.end(), static_cast(unrelated_before.local())), fanins.end()); + EXPECT_EQ(std::find(fanins.begin(), fanins.end(), static_cast(unrelated_after.local())), fanins.end()); + + const auto &second_payload = sm_handle->header->ring.get_slot_state_by_task_id(second_host_write_local).payload[0]; + ASSERT_EQ(second_payload.fanin_count, 4); + std::vector second_fanins; + for (int32_t i = 0; i < second_payload.fanin_count; ++i) { + second_fanins.push_back(sched.fanin_local_id(second_payload, i)); + } + EXPECT_NE(std::find(second_fanins.begin(), second_fanins.end(), host_write_local), second_fanins.end()); +} diff --git a/tests/ut/cpp/a2a3/test_hbg_tensormap.cpp b/tests/ut/cpp/a2a3/test_hbg_tensormap.cpp index 62c2ecfc72..550141ad4a 100644 --- a/tests/ut/cpp/a2a3/test_hbg_tensormap.cpp +++ b/tests/ut/cpp/a2a3/test_hbg_tensormap.cpp @@ -22,6 +22,7 @@ #include #include "utils/device_arena.h" +#include "pto_dep_compute.h" #include "pto_tensormap.h" namespace { @@ -43,6 +44,21 @@ void run_lookup(PTO2TensorMap &tmap, const ChipTensor &tensor, TestLookupResult }); } +void run_lookup(PTO2TensorMap &tmap, const ChipTensor &tensor, TensorAccessKind access_kind, TestLookupResult &out) { + tmap.lookup(tensor, access_kind, [&](PTO2TensorMapEntry &e, OverlapStatus s) -> bool { + out.entries.push_back({&e, s}); + out.count++; + return true; + }); +} + +DepInputs dep_inputs(const CoreTaskArgs &args) { + return DepInputs{ + args.tensor_count(), args.tensor_data(), args.tag_data(), static_cast(args.explicit_dep_count()), + args.explicit_deps_data(), + }; +} + ChipTensor make_test_tensor(uint64_t addr, uint32_t shape0) { uint32_t shapes[MAX_TENSOR_DIMS] = {shape0}; return make_tensor_external(reinterpret_cast(addr), shapes, 1, DataType::FLOAT32, false, 0); @@ -70,6 +86,13 @@ class HbgTensorMapTest : public ::testing::Test { } }; +TEST(HbgTensorMapLayoutTest, DefaultLayoutCapsSparseReaderBuckets) { + DeviceArena arena; + auto layout = PTO2TensorMap::reserve_layout_default(arena, 32); + EXPECT_EQ(layout.num_buckets, PTO2_TENSORMAP_NUM_BUCKETS); + EXPECT_EQ(layout.num_reader_buckets, PTO2_TENSORMAP_READER_NUM_BUCKETS); +} + TEST_F(HbgTensorMapTest, CleanupRetiredRemovesEntriesForRetiredTasks) { ChipTensor t = make_test_tensor(0x1000, 256); tmap.insert(t, PTO2TaskId::make(0, 0)); @@ -83,7 +106,7 @@ TEST_F(HbgTensorMapTest, CleanupRetiredRemovesEntriesForRetiredTasks) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, 2)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(0, 2)); } TEST_F(HbgTensorMapTest, CleanupRetiredFreesEveryOutputOfOneTask) { @@ -119,7 +142,128 @@ TEST_F(HbgTensorMapTest, CleanupRetiredSparesLaterTaskReusingSlot) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, WINDOW_SIZE)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(0, WINDOW_SIZE)); +} + +TEST_F(HbgTensorMapTest, ReaderAndWriterIndexesImplementRawWarAndOutputExistingOrdering) { + ChipTensor tensor = make_test_tensor(0x4000, 256); + const PTO2TaskId writer0 = PTO2TaskId::make(0, 1); + const PTO2TaskId reader0 = PTO2TaskId::make(0, 2); + const PTO2TaskId writer1 = PTO2TaskId::make(0, 3); + const PTO2TaskId reader1 = PTO2TaskId::make(0, 4); + const PTO2TaskId output_existing = PTO2TaskId::make(0, 5); + tmap.insert(tensor, writer0, TensorAccessKind::WRITER); + + CoreTaskArgs untracked_input_args; + untracked_input_args.add_input(tensor); + std::vector fanin; + auto emit = [&](PTO2TaskId id) { + fanin.push_back(id); + return true; + }; + ASSERT_TRUE(compute_task_fanin(dep_inputs(untracked_input_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer0})); + EXPECT_EQ(count_registrable_accesses(dep_inputs(untracked_input_args), false), 0); + register_task_accesses(dep_inputs(untracked_input_args), reader0, tmap, false); + + TestLookupResult readers; + run_lookup(tmap, tensor, TensorAccessKind::READER, readers); + EXPECT_EQ(readers.count, 0); + + CoreTaskArgs input_args; + input_args.add_tracked_input(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(input_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer0})); + EXPECT_EQ(count_registrable_accesses(dep_inputs(input_args), false), 1); + register_task_accesses(dep_inputs(input_args), reader0, tmap, false); + + readers = {}; + run_lookup(tmap, tensor, TensorAccessKind::READER, readers); + ASSERT_EQ(readers.count, 1); + EXPECT_EQ(readers.entries[0].entry->access_task_id, reader0); + TestLookupResult writers; + run_lookup(tmap, tensor, TensorAccessKind::WRITER, writers); + ASSERT_EQ(writers.count, 1); + EXPECT_EQ(writers.entries[0].entry->access_task_id, writer0); + + CoreTaskArgs inout_args; + inout_args.add_inout(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(inout_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer0, reader0})); + register_task_accesses(dep_inputs(inout_args), writer1, tmap, false); + + CoreTaskArgs second_input_args; + second_input_args.add_tracked_input(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(second_input_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer1})); + register_task_accesses(dep_inputs(second_input_args), reader1, tmap, false); + + CoreTaskArgs output_args; + output_args.add_output(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(output_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({reader1})); + register_task_accesses(dep_inputs(output_args), output_existing, tmap, false); + + readers = {}; + run_lookup(tmap, tensor, TensorAccessKind::READER, readers); + EXPECT_EQ(readers.count, 0); + writers = {}; + run_lookup(tmap, tensor, TensorAccessKind::WRITER, writers); + ASSERT_EQ(writers.count, 2); + EXPECT_EQ(writers.entries[0].entry->access_task_id, output_existing); + EXPECT_EQ(writers.entries[1].entry->access_task_id, writer1); +} + +TEST_F(HbgTensorMapTest, TrackedAccessesSkipManualScopeAndManualTensor) { + ChipTensor tensor = make_test_tensor(0x5000, 64); + tmap.insert(tensor, PTO2TaskId::make(0, 1)); + int emitted = 0; + auto emit = [&](PTO2TaskId) { + emitted++; + return true; + }; + + CoreTaskArgs manual_scope_args; + manual_scope_args.add_tracked_input(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(manual_scope_args), tmap, true, emit)); + register_task_accesses(dep_inputs(manual_scope_args), PTO2TaskId::make(0, 2), tmap, true); + + tensor.manual_dep = true; + CoreTaskArgs manual_tensor_args; + manual_tensor_args.add_tracked_input(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(manual_tensor_args), tmap, false, emit)); + register_task_accesses(dep_inputs(manual_tensor_args), PTO2TaskId::make(0, 3), tmap, false); + + EXPECT_EQ(emitted, 0); + EXPECT_EQ(tmap.valid_count(), 1); + EXPECT_EQ(tmap.current_readers(), 0); + EXPECT_EQ(tmap.current_writers(), 1); +} + +TEST_F(HbgTensorMapTest, NoDepRetainsCreatorWithoutTensorMapAccess) { + ChipTensor tensor = make_test_tensor(0x5100, 64); + const PTO2TaskId owner = PTO2TaskId::make(0, 7); + tensor.owner_task_id = owner; + tmap.insert(tensor, PTO2TaskId::make(0, 1)); + std::vector fanin; + auto emit = [&](PTO2TaskId id) { + fanin.push_back(id); + return true; + }; + + CoreTaskArgs no_dep_args; + no_dep_args.add_no_dep(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(no_dep_args), tmap, false, emit)); + register_task_accesses(dep_inputs(no_dep_args), PTO2TaskId::make(0, 4), tmap, false); + + EXPECT_EQ(fanin, std::vector({owner})); + EXPECT_EQ(tmap.valid_count(), 1); + EXPECT_EQ(tmap.current_readers(), 0); + EXPECT_EQ(tmap.current_writers(), 1); } } // namespace diff --git a/tests/ut/cpp/a2a3/test_orchestrator_fanin.cpp b/tests/ut/cpp/a2a3/test_orchestrator_fanin.cpp index 6a5f4f34e2..448a049c0d 100644 --- a/tests/ut/cpp/a2a3/test_orchestrator_fanin.cpp +++ b/tests/ut/cpp/a2a3/test_orchestrator_fanin.cpp @@ -221,6 +221,87 @@ TEST_F(OrchestratorFaninTest, AllCompletedFastPathReleasesWaitOnlyPin) { EXPECT_EQ(producer_slot.fanout_refcount.load(), rc_before + 1); } +TEST_F(OrchestratorFaninTest, DuplicateReaderRegistrationAddsOneWarFanin) { + orch.begin_scope(); + uint32_t shape[] = {16}; + ChipTensor tensor = make_tensor_external(gm_heap.data(), shape, 1, DataType::FLOAT32, false); + + CoreTaskArgs reader_args; + reader_args.add_tracked_input(tensor); + reader_args.add_tracked_input(tensor); + TaskOutputTensors reader = orch.submit_dummy_task(reader_args); + ASSERT_TRUE(reader.task_id().is_valid()); + + CoreTaskArgs writer_args; + writer_args.add_output(tensor); + TaskOutputTensors writer = orch.submit_dummy_task(writer_args); + ASSERT_TRUE(writer.task_id().is_valid()); + + auto &reader_slot = + sm_handle->header->rings[reader.task_id().ring()].get_slot_state_by_task_id(reader.task_id().local()); + auto &writer_slot = + sm_handle->header->rings[writer.task_id().ring()].get_slot_state_by_task_id(writer.task_id().local()); + ASSERT_NE(writer_slot.payload, nullptr); + EXPECT_EQ(writer_slot.payload->fanin_actual_count, 1); + EXPECT_EQ(writer_slot.payload->fanin_inline_edges[0].slot_state(), &reader_slot); +} + +TEST_F(OrchestratorFaninTest, CompletedReaderDoesNotPinWarFanin) { + orch.begin_scope(); + uint32_t shape[] = {16}; + ChipTensor tensor = make_tensor_external(gm_heap.data(), shape, 1, DataType::FLOAT32, false); + + CoreTaskArgs reader_args; + reader_args.add_tracked_input(tensor); + TaskOutputTensors reader = orch.submit_dummy_task(reader_args); + ASSERT_TRUE(reader.task_id().is_valid()); + + auto &reader_slot = + sm_handle->header->rings[reader.task_id().ring()].get_slot_state_by_task_id(reader.task_id().local()); + reader_slot.task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); + const uint32_t fanout_before = reader_slot.fanout_count; + + CoreTaskArgs writer_args; + writer_args.add_output(tensor); + TaskOutputTensors writer = orch.submit_dummy_task(writer_args); + ASSERT_TRUE(writer.task_id().is_valid()); + + auto &writer_slot = + sm_handle->header->rings[writer.task_id().ring()].get_slot_state_by_task_id(writer.task_id().local()); + ASSERT_NE(writer_slot.payload, nullptr); + EXPECT_EQ(writer_slot.payload->fanin_actual_count, 0); + EXPECT_EQ(reader_slot.fanout_count, fanout_before); +} + +TEST_F(OrchestratorFaninTest, NestedReaderCreatesCrossRingWarToOuterWriter) { + orch.begin_scope(); + uint32_t shape[] = {16}; + ChipTensor tensor = make_tensor_external(gm_heap.data(), shape, 1, DataType::FLOAT32, false); + + CoreTaskArgs first_writer_args; + first_writer_args.add_output(tensor); + TaskOutputTensors first_writer = orch.submit_dummy_task(first_writer_args); + ASSERT_EQ(first_writer.task_id(), PTO2TaskId::make(0, 0)); + + orch.begin_scope(); + CoreTaskArgs reader_args; + reader_args.add_tracked_input(tensor); + TaskOutputTensors reader = orch.submit_dummy_task(reader_args); + ASSERT_EQ(reader.task_id(), PTO2TaskId::make(1, 0)); + orch.end_scope(); + + CoreTaskArgs next_writer_args; + next_writer_args.add_output(tensor); + TaskOutputTensors next_writer = orch.submit_dummy_task(next_writer_args); + ASSERT_EQ(next_writer.task_id(), PTO2TaskId::make(0, 1)); + + auto &reader_slot = sm_handle->header->rings[1].get_slot_state_by_task_id(0); + auto &writer_slot = sm_handle->header->rings[0].get_slot_state_by_task_id(1); + ASSERT_NE(writer_slot.payload, nullptr); + ASSERT_EQ(writer_slot.payload->fanin_actual_count, 1); + EXPECT_EQ(writer_slot.payload->fanin_inline_edges[0].slot_state(), &reader_slot); +} + TEST_F(OrchestratorFaninTest, SubmitPathHeapDeadlockLogReportsRingAndRealHeapState) { std::vector create_infos; create_infos.reserve(8); diff --git a/tests/ut/cpp/a2a3/test_tensormap.cpp b/tests/ut/cpp/a2a3/test_tensormap.cpp index 9cea13db58..6ebc9c9126 100644 --- a/tests/ut/cpp/a2a3/test_tensormap.cpp +++ b/tests/ut/cpp/a2a3/test_tensormap.cpp @@ -29,6 +29,7 @@ #include #include "utils/device_arena.h" +#include "pto_dep_compute.h" #include "pto_orchestration_api.h" #include "pto_tensormap.h" @@ -56,6 +57,22 @@ static void run_lookup(PTO2TensorMap &tmap, const ChipTensor &tensor, TestLookup }); } +static void +run_lookup(PTO2TensorMap &tmap, const ChipTensor &tensor, TensorAccessKind access_kind, TestLookupResult &out) { + tmap.lookup(tensor, access_kind, [&](PTO2TensorMapEntry &e, OverlapStatus s) -> bool { + out.entries.push_back({&e, s}); + out.count++; + return true; + }); +} + +static DepInputs dep_inputs(const CoreTaskArgs &args) { + return DepInputs{ + args.tensor_count(), args.tensor_data(), args.tag_data(), static_cast(args.explicit_dep_count()), + args.explicit_deps_data(), + }; +} + static ChipTensor make_test_tensor(uint64_t addr, uint32_t shape0, uint32_t ndims = 1, int32_t version = 0) { uint32_t shapes[MAX_TENSOR_DIMS] = {shape0}; return make_tensor_external(reinterpret_cast(addr), shapes, ndims, DataType::FLOAT32, false, version); @@ -99,12 +116,21 @@ class TensorMapTest : public ::testing::Test { TEST_F(TensorMapTest, InitValidState) { EXPECT_EQ(tmap.num_buckets, NUM_BUCKETS); + EXPECT_EQ(tmap.num_reader_buckets, NUM_BUCKETS); EXPECT_EQ(tmap.pool_size, POOL_SIZE); EXPECT_EQ(tmap.next_entry_idx, 0); EXPECT_EQ(tmap.free_num, 0); EXPECT_EQ(tmap.valid_count(), 0); } +TEST(TensorMapLayoutTest, DefaultLayoutCapsSparseReaderBuckets) { + DeviceArena arena; + int32_t window_sizes[PTO2_MAX_RING_DEPTH] = {32, 32, 32, 32}; + auto layout = PTO2TensorMap::reserve_layout_default(arena, window_sizes); + EXPECT_EQ(layout.num_buckets, PTO2_TENSORMAP_NUM_BUCKETS); + EXPECT_EQ(layout.num_reader_buckets, PTO2_TENSORMAP_READER_NUM_BUCKETS); +} + TEST_F(TensorMapTest, InitRequiresPowerOfTwoBuckets) { // Non-power-of-2 bucket counts trip an always_assert inside reserve_layout // (asserting EXPECT_DEATH is impossible in release builds where @@ -158,7 +184,7 @@ TEST_F(TensorMapTest, InsertThenLookupFindsProducer) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, tid); + EXPECT_EQ(result.entries[0].entry->access_task_id, tid); } TEST_F(TensorMapTest, LookupEmptyReturnsZero) { @@ -192,12 +218,12 @@ TEST_F(TensorMapTest, InsertDifferentBuffersNoCollision) { TestLookupResult r1; run_lookup(tmap, t1, r1); EXPECT_EQ(r1.count, 1); - EXPECT_EQ(r1.entries[0].entry->producer_task_id, PTO2TaskId::make(0, 0)); + EXPECT_EQ(r1.entries[0].entry->access_task_id, PTO2TaskId::make(0, 0)); TestLookupResult r2; run_lookup(tmap, t2, r2); EXPECT_EQ(r2.count, 1); - EXPECT_EQ(r2.entries[0].entry->producer_task_id, PTO2TaskId::make(0, 1)); + EXPECT_EQ(r2.entries[0].entry->access_task_id, PTO2TaskId::make(0, 1)); } // ============================================================================= @@ -332,7 +358,7 @@ TEST_F(TensorMapTest, StaleEntriesSkippedDuringLookup) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, 1)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(0, 1)); } TEST_F(TensorMapTest, StaleEntriesNotTruncatedAcrossRings) { @@ -348,7 +374,7 @@ TEST_F(TensorMapTest, StaleEntriesNotTruncatedAcrossRings) { run_lookup(tmap, t, result); // Ring 1 task 0 still valid, ring 0 task 0 invalidated ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(1, 0)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(1, 0)); } // ============================================================================= @@ -370,7 +396,7 @@ TEST_F(TensorMapTest, CleanupRetiredRemovesEntriesForRetiredTasks) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, 2)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(0, 2)); } TEST_F(TensorMapTest, CleanupRetiredPreservesOtherRings) { @@ -385,7 +411,7 @@ TEST_F(TensorMapTest, CleanupRetiredPreservesOtherRings) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(1, 0)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(1, 0)); } TEST_F(TensorMapTest, CleanupRetiredFreesEntriesToPool) { @@ -423,7 +449,7 @@ TEST_F(TensorMapTest, CleanupRetiredSparesLaterTaskReusingSlot) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, WINDOW_SIZE)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(0, WINDOW_SIZE)); } // ============================================================================= @@ -447,7 +473,7 @@ TEST_F(TensorMapTest, MultiRingIndependentLookup) { TestLookupResult result2; run_lookup(tmap, t, result2); EXPECT_EQ(result2.count, 1); - EXPECT_EQ(result2.entries[0].entry->producer_task_id, PTO2TaskId::make(1, 3)); + EXPECT_EQ(result2.entries[0].entry->access_task_id, PTO2TaskId::make(1, 3)); } // ============================================================================= @@ -543,12 +569,227 @@ TEST_F(TensorMapTest, RemoveMiddleEntryPreservesChain) { std::set found_locals; for (int i = 0; i < result.count; i++) { - found_locals.insert(result.entries[i].entry->producer_task_id.local()); + found_locals.insert(result.entries[i].entry->access_task_id.local()); } EXPECT_TRUE(found_locals.count(0)); EXPECT_TRUE(found_locals.count(2)); } +TEST_F(TensorMapTest, ReaderAndWriterIndexesImplementRawWarAndOutputExistingOrdering) { + ChipTensor tensor = make_test_tensor(0x4000, 256); + const PTO2TaskId writer0 = PTO2TaskId::make(0, 1); + const PTO2TaskId reader0 = PTO2TaskId::make(1, 2); + const PTO2TaskId writer1 = PTO2TaskId::make(0, 3); + const PTO2TaskId reader1 = PTO2TaskId::make(2, 4); + const PTO2TaskId output_existing = PTO2TaskId::make(0, 5); + tmap.insert(tensor, writer0, TensorAccessKind::WRITER); + + std::vector fanin; + auto emit = [&](PTO2TaskId id) { + fanin.push_back(id); + return true; + }; + + CoreTaskArgs untracked_input_args; + untracked_input_args.add_input(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(untracked_input_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer0})); + EXPECT_EQ(count_registrable_accesses(dep_inputs(untracked_input_args), false), 0); + register_task_accesses(dep_inputs(untracked_input_args), reader0, tmap, false); + + TestLookupResult readers; + run_lookup(tmap, tensor, TensorAccessKind::READER, readers); + EXPECT_EQ(readers.count, 0); + + CoreTaskArgs input_args; + input_args.add_tracked_input(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(input_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer0})); + EXPECT_EQ(count_registrable_accesses(dep_inputs(input_args), false), 1); + register_task_accesses(dep_inputs(input_args), reader0, tmap, false); + + readers = {}; + run_lookup(tmap, tensor, TensorAccessKind::READER, readers); + ASSERT_EQ(readers.count, 1); + EXPECT_EQ(readers.entries[0].entry->access_task_id, reader0); + TestLookupResult writers; + run_lookup(tmap, tensor, TensorAccessKind::WRITER, writers); + ASSERT_EQ(writers.count, 1); + EXPECT_EQ(writers.entries[0].entry->access_task_id, writer0); + + CoreTaskArgs inout_args; + inout_args.add_inout(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(inout_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer0, reader0})); + register_task_accesses(dep_inputs(inout_args), writer1, tmap, false); + + CoreTaskArgs second_input_args; + second_input_args.add_tracked_input(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(second_input_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer1})); + register_task_accesses(dep_inputs(second_input_args), reader1, tmap, false); + + CoreTaskArgs output_args; + output_args.add_output(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(output_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({reader1})); + register_task_accesses(dep_inputs(output_args), output_existing, tmap, false); + + readers = {}; + run_lookup(tmap, tensor, TensorAccessKind::READER, readers); + EXPECT_EQ(readers.count, 0); + writers = {}; + run_lookup(tmap, tensor, TensorAccessKind::WRITER, writers); + ASSERT_EQ(writers.count, 2); + EXPECT_EQ(writers.entries[0].entry->access_task_id, output_existing); + EXPECT_EQ(writers.entries[1].entry->access_task_id, writer1); +} + +TEST_F(TensorMapTest, ReadersAreIndependentAndWriterWaitsForEveryReader) { + ChipTensor tensor = make_test_tensor(0x4100, 256); + const PTO2TaskId reader0 = PTO2TaskId::make(0, 1); + const PTO2TaskId reader1 = PTO2TaskId::make(1, 2); + + CoreTaskArgs input_args; + input_args.add_tracked_input(tensor); + std::vector fanin; + auto emit = [&](PTO2TaskId id) { + fanin.push_back(id); + return true; + }; + ASSERT_TRUE(compute_task_fanin(dep_inputs(input_args), tmap, false, emit)); + EXPECT_TRUE(fanin.empty()); + register_task_accesses(dep_inputs(input_args), reader0, tmap, false); + + ASSERT_TRUE(compute_task_fanin(dep_inputs(input_args), tmap, false, emit)); + EXPECT_TRUE(fanin.empty()); + register_task_accesses(dep_inputs(input_args), reader1, tmap, false); + + CoreTaskArgs writer_args; + writer_args.add_output(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(writer_args), tmap, false, emit)); + std::set fanin_raw; + for (PTO2TaskId id : fanin) + fanin_raw.insert(id.raw); + EXPECT_EQ(fanin_raw, std::set({reader0.raw, reader1.raw})); + EXPECT_EQ(tmap.reader_high_water, 2); +} + +TEST_F(TensorMapTest, DisjointWriteKeepsReaderAndCoveringWriteRetiresIt) { + ChipTensor base = make_test_tensor_2d(0x4200, 256, 1); + uint32_t reader_shape[] = {128, 1}; + uint32_t reader_offset[] = {0, 0}; + ChipTensor reader_view = base.view(reader_shape, reader_offset); + + CoreTaskArgs reader_args; + reader_args.add_tracked_input(reader_view); + register_task_accesses(dep_inputs(reader_args), PTO2TaskId::make(0, 1), tmap, false); + + uint32_t disjoint_offset[] = {128, 0}; + ChipTensor disjoint = base.view(reader_shape, disjoint_offset); + CoreTaskArgs disjoint_args; + disjoint_args.add_output(disjoint); + std::vector fanin; + auto emit = [&](PTO2TaskId id) { + fanin.push_back(id); + return true; + }; + ASSERT_TRUE(compute_task_fanin(dep_inputs(disjoint_args), tmap, false, emit)); + EXPECT_TRUE(fanin.empty()); + register_task_accesses(dep_inputs(disjoint_args), PTO2TaskId::make(0, 2), tmap, false); + EXPECT_EQ(tmap.current_readers(), 1); + + uint32_t partial_shape[] = {128, 1}; + uint32_t partial_offset[] = {64, 0}; + ChipTensor partial = base.view(partial_shape, partial_offset); + CoreTaskArgs partial_args; + partial_args.add_output(partial); + ASSERT_TRUE(compute_task_fanin(dep_inputs(partial_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({PTO2TaskId::make(0, 1)})); + register_task_accesses(dep_inputs(partial_args), PTO2TaskId::make(0, 3), tmap, false); + EXPECT_EQ(tmap.current_readers(), 1); + + CoreTaskArgs covering_args; + covering_args.add_output(base); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(covering_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({PTO2TaskId::make(0, 1)})); + register_task_accesses(dep_inputs(covering_args), PTO2TaskId::make(0, 4), tmap, false); + EXPECT_EQ(tmap.current_readers(), 0); +} + +TEST_F(TensorMapTest, RetiredReaderDoesNotCreateStaleWar) { + ChipTensor tensor = make_test_tensor(0x4300, 256); + CoreTaskArgs reader_args; + reader_args.add_tracked_input(tensor); + register_task_accesses(dep_inputs(reader_args), PTO2TaskId::make(0, 0), tmap, false); + ASSERT_EQ(tmap.current_readers(), 1); + + tmap.cleanup_retired(0, 0, 1); + ASSERT_EQ(tmap.current_readers(), 0); + + CoreTaskArgs writer_args; + writer_args.add_output(tensor); + int emitted = 0; + ASSERT_TRUE(compute_task_fanin(dep_inputs(writer_args), tmap, false, [&](PTO2TaskId) { + emitted++; + return true; + })); + EXPECT_EQ(emitted, 0); +} + +TEST_F(TensorMapTest, TrackedAccessesSkipManualScopeAndManualTensor) { + ChipTensor tensor = make_test_tensor(0x5000, 64); + tmap.insert(tensor, PTO2TaskId::make(0, 1)); + int emitted = 0; + auto emit = [&](PTO2TaskId) { + emitted++; + return true; + }; + + CoreTaskArgs manual_scope_args; + manual_scope_args.add_tracked_input(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(manual_scope_args), tmap, true, emit)); + register_task_accesses(dep_inputs(manual_scope_args), PTO2TaskId::make(0, 2), tmap, true); + + tensor.manual_dep = true; + CoreTaskArgs manual_tensor_args; + manual_tensor_args.add_tracked_input(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(manual_tensor_args), tmap, false, emit)); + register_task_accesses(dep_inputs(manual_tensor_args), PTO2TaskId::make(0, 3), tmap, false); + + EXPECT_EQ(emitted, 0); + EXPECT_EQ(tmap.valid_count(), 1); + EXPECT_EQ(tmap.current_readers(), 0); + EXPECT_EQ(tmap.current_writers(), 1); +} + +TEST_F(TensorMapTest, NoDepRetainsCreatorWithoutTensorMapAccess) { + ChipTensor tensor = make_test_tensor(0x5100, 64); + const PTO2TaskId owner = PTO2TaskId::make(2, 7); + tensor.owner_task_id = owner; + tmap.insert(tensor, PTO2TaskId::make(0, 1)); + std::vector fanin; + auto emit = [&](PTO2TaskId id) { + fanin.push_back(id); + return true; + }; + + CoreTaskArgs no_dep_args; + no_dep_args.add_no_dep(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(no_dep_args), tmap, false, emit)); + register_task_accesses(dep_inputs(no_dep_args), PTO2TaskId::make(0, 4), tmap, false); + + EXPECT_EQ(fanin, std::vector({owner})); + EXPECT_EQ(tmap.valid_count(), 1); + EXPECT_EQ(tmap.current_readers(), 0); + EXPECT_EQ(tmap.current_writers(), 1); +} + // ============================================================================= // PTO2TaskId encoding/decoding // ============================================================================= diff --git a/tests/ut/cpp/a5/test_hbg_submit_poison.cpp b/tests/ut/cpp/a5/test_hbg_submit_poison.cpp index 7b019204ba..2b2d218fc0 100644 --- a/tests/ut/cpp/a5/test_hbg_submit_poison.cpp +++ b/tests/ut/cpp/a5/test_hbg_submit_poison.cpp @@ -26,6 +26,7 @@ #include +#include #include #include #include @@ -88,6 +89,40 @@ class HbgSubmitPoisonTest : public ::testing::Test { std::memset(ring.slot_states, POISON, n * sizeof(PTO2TaskSlotState)); std::memset(ring.completion_flags, POISON, n * sizeof(std::atomic)); } + + void verify_reader_fanout(int32_t reader_count) { + uint32_t shape[] = {16}; + std::vector storage(shape[0]); + ChipTensor tensor = make_tensor_external(storage.data(), shape, 1, DataType::FLOAT32, false, /*version=*/0); + std::vector readers; + readers.reserve(reader_count); + + orch.begin_scope(); + for (int32_t i = 0; i < reader_count; ++i) { + CoreTaskArgs args; + args.add_tracked_input(tensor); + TaskOutputTensors reader = orch.submit_dummy_task(args); + ASSERT_TRUE(reader.task_id().is_valid()); + readers.push_back(reader.task_id()); + } + CoreTaskArgs writer_args; + writer_args.add_inout(tensor); + TaskOutputTensors writer = orch.submit_dummy_task(writer_args); + ASSERT_TRUE(writer.task_id().is_valid()); + orch.end_scope(); + + auto &ring = sm_handle->header->ring; + const PTO2TaskPayload &payload = ring.task_payloads[ring.get_slot_by_task_id(writer.task_id().local())]; + ASSERT_EQ(payload.fanin_count, reader_count); + const int32_t expected_spill = std::max(0, reader_count - PTO2_MAX_FANIN); + ASSERT_EQ(payload.fanin_spill_count, expected_spill); + EXPECT_EQ(sched.fanin_spill_top, expected_spill); + for (int32_t i = 0; i < reader_count; ++i) { + SCOPED_TRACE(testing::Message() << "fanin index=" << i); + EXPECT_EQ(sched.fanin_local_id(payload, i), static_cast(readers[reader_count - 1 - i].local())); + } + EXPECT_EQ(sm_handle->header->orch_error_code.load(std::memory_order_acquire), PTO2_ERROR_NONE); + } }; TEST_F(HbgSubmitPoisonTest, EveryDeviceReadFieldIsWrittenOverPoison) { @@ -183,3 +218,62 @@ TEST_F(HbgSubmitPoisonTest, EveryDeviceReadFieldIsWrittenOverPoison) { EXPECT_EQ(cons_pl.fanin_count, 1); EXPECT_EQ(cons_pl.fanin_local_ids[0], static_cast(root.task_id().local())); } + +TEST_F(HbgSubmitPoisonTest, ReaderFanoutBelowInlineBoundaryHasNoSpill) { verify_reader_fanout(64); } + +TEST_F(HbgSubmitPoisonTest, ReaderFanoutAtInlineBoundaryHasNoSpill) { verify_reader_fanout(PTO2_MAX_FANIN); } + +TEST_F(HbgSubmitPoisonTest, ReaderFanoutAboveInlineBoundarySpillsOneEdge) { verify_reader_fanout(PTO2_MAX_FANIN + 1); } + +TEST_F(HbgSubmitPoisonTest, ReaderFanoutBeyondInlineFaninSpillsWithoutDroppingWarEdges) { verify_reader_fanout(256); } + +TEST_F(HbgSubmitPoisonTest, HostWriteWaitsOnlyForConflictingWriterConsumersAndTrackedReaders) { + uint32_t shape[] = {16}; + std::vector storage(shape[0]); + ChipTensor tensor = make_tensor_external(storage.data(), shape, 1, DataType::FLOAT32, false); + + orch.begin_scope(); + CoreTaskArgs unrelated_before_args; + const PTO2TaskId unrelated_before = orch.submit_dummy_task(unrelated_before_args).task_id(); + + CoreTaskArgs writer_args; + writer_args.add_output(tensor); + const PTO2TaskId writer = orch.submit_dummy_task(writer_args).task_id(); + + CoreTaskArgs consumer_args; + consumer_args.add_input(tensor); + const PTO2TaskId consumer = orch.submit_dummy_task(consumer_args).task_id(); + + CoreTaskArgs tracked_reader_args; + tracked_reader_args.add_tracked_input(tensor); + const PTO2TaskId tracked_reader = orch.submit_dummy_task(tracked_reader_args).task_id(); + + CoreTaskArgs unrelated_after_args; + const PTO2TaskId unrelated_after = orch.submit_dummy_task(unrelated_after_args).task_id(); + + orch.submit_host_write(tensor, tensor.buffer.addr, 7, sizeof(uint32_t)); + const int32_t host_write_local = sm_handle->header->ring.fc.current_task_index.load(std::memory_order_acquire) - 1; + orch.submit_host_write(tensor, tensor.buffer.addr, 8, sizeof(uint32_t)); + const int32_t second_host_write_local = + sm_handle->header->ring.fc.current_task_index.load(std::memory_order_acquire) - 1; + orch.end_scope(); + + const auto &payload = sm_handle->header->ring.get_slot_state_by_task_id(host_write_local).payload[0]; + ASSERT_EQ(payload.fanin_count, 3); + std::vector fanins; + for (int32_t i = 0; i < payload.fanin_count; ++i) + fanins.push_back(sched.fanin_local_id(payload, i)); + EXPECT_NE(std::find(fanins.begin(), fanins.end(), static_cast(writer.local())), fanins.end()); + EXPECT_NE(std::find(fanins.begin(), fanins.end(), static_cast(consumer.local())), fanins.end()); + EXPECT_NE(std::find(fanins.begin(), fanins.end(), static_cast(tracked_reader.local())), fanins.end()); + EXPECT_EQ(std::find(fanins.begin(), fanins.end(), static_cast(unrelated_before.local())), fanins.end()); + EXPECT_EQ(std::find(fanins.begin(), fanins.end(), static_cast(unrelated_after.local())), fanins.end()); + + const auto &second_payload = sm_handle->header->ring.get_slot_state_by_task_id(second_host_write_local).payload[0]; + ASSERT_EQ(second_payload.fanin_count, 4); + std::vector second_fanins; + for (int32_t i = 0; i < second_payload.fanin_count; ++i) { + second_fanins.push_back(sched.fanin_local_id(second_payload, i)); + } + EXPECT_NE(std::find(second_fanins.begin(), second_fanins.end(), host_write_local), second_fanins.end()); +} diff --git a/tests/ut/cpp/a5/test_orchestrator_fanin.cpp b/tests/ut/cpp/a5/test_orchestrator_fanin.cpp index 6a5f4f34e2..448a049c0d 100644 --- a/tests/ut/cpp/a5/test_orchestrator_fanin.cpp +++ b/tests/ut/cpp/a5/test_orchestrator_fanin.cpp @@ -221,6 +221,87 @@ TEST_F(OrchestratorFaninTest, AllCompletedFastPathReleasesWaitOnlyPin) { EXPECT_EQ(producer_slot.fanout_refcount.load(), rc_before + 1); } +TEST_F(OrchestratorFaninTest, DuplicateReaderRegistrationAddsOneWarFanin) { + orch.begin_scope(); + uint32_t shape[] = {16}; + ChipTensor tensor = make_tensor_external(gm_heap.data(), shape, 1, DataType::FLOAT32, false); + + CoreTaskArgs reader_args; + reader_args.add_tracked_input(tensor); + reader_args.add_tracked_input(tensor); + TaskOutputTensors reader = orch.submit_dummy_task(reader_args); + ASSERT_TRUE(reader.task_id().is_valid()); + + CoreTaskArgs writer_args; + writer_args.add_output(tensor); + TaskOutputTensors writer = orch.submit_dummy_task(writer_args); + ASSERT_TRUE(writer.task_id().is_valid()); + + auto &reader_slot = + sm_handle->header->rings[reader.task_id().ring()].get_slot_state_by_task_id(reader.task_id().local()); + auto &writer_slot = + sm_handle->header->rings[writer.task_id().ring()].get_slot_state_by_task_id(writer.task_id().local()); + ASSERT_NE(writer_slot.payload, nullptr); + EXPECT_EQ(writer_slot.payload->fanin_actual_count, 1); + EXPECT_EQ(writer_slot.payload->fanin_inline_edges[0].slot_state(), &reader_slot); +} + +TEST_F(OrchestratorFaninTest, CompletedReaderDoesNotPinWarFanin) { + orch.begin_scope(); + uint32_t shape[] = {16}; + ChipTensor tensor = make_tensor_external(gm_heap.data(), shape, 1, DataType::FLOAT32, false); + + CoreTaskArgs reader_args; + reader_args.add_tracked_input(tensor); + TaskOutputTensors reader = orch.submit_dummy_task(reader_args); + ASSERT_TRUE(reader.task_id().is_valid()); + + auto &reader_slot = + sm_handle->header->rings[reader.task_id().ring()].get_slot_state_by_task_id(reader.task_id().local()); + reader_slot.task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); + const uint32_t fanout_before = reader_slot.fanout_count; + + CoreTaskArgs writer_args; + writer_args.add_output(tensor); + TaskOutputTensors writer = orch.submit_dummy_task(writer_args); + ASSERT_TRUE(writer.task_id().is_valid()); + + auto &writer_slot = + sm_handle->header->rings[writer.task_id().ring()].get_slot_state_by_task_id(writer.task_id().local()); + ASSERT_NE(writer_slot.payload, nullptr); + EXPECT_EQ(writer_slot.payload->fanin_actual_count, 0); + EXPECT_EQ(reader_slot.fanout_count, fanout_before); +} + +TEST_F(OrchestratorFaninTest, NestedReaderCreatesCrossRingWarToOuterWriter) { + orch.begin_scope(); + uint32_t shape[] = {16}; + ChipTensor tensor = make_tensor_external(gm_heap.data(), shape, 1, DataType::FLOAT32, false); + + CoreTaskArgs first_writer_args; + first_writer_args.add_output(tensor); + TaskOutputTensors first_writer = orch.submit_dummy_task(first_writer_args); + ASSERT_EQ(first_writer.task_id(), PTO2TaskId::make(0, 0)); + + orch.begin_scope(); + CoreTaskArgs reader_args; + reader_args.add_tracked_input(tensor); + TaskOutputTensors reader = orch.submit_dummy_task(reader_args); + ASSERT_EQ(reader.task_id(), PTO2TaskId::make(1, 0)); + orch.end_scope(); + + CoreTaskArgs next_writer_args; + next_writer_args.add_output(tensor); + TaskOutputTensors next_writer = orch.submit_dummy_task(next_writer_args); + ASSERT_EQ(next_writer.task_id(), PTO2TaskId::make(0, 1)); + + auto &reader_slot = sm_handle->header->rings[1].get_slot_state_by_task_id(0); + auto &writer_slot = sm_handle->header->rings[0].get_slot_state_by_task_id(1); + ASSERT_NE(writer_slot.payload, nullptr); + ASSERT_EQ(writer_slot.payload->fanin_actual_count, 1); + EXPECT_EQ(writer_slot.payload->fanin_inline_edges[0].slot_state(), &reader_slot); +} + TEST_F(OrchestratorFaninTest, SubmitPathHeapDeadlockLogReportsRingAndRealHeapState) { std::vector create_infos; create_infos.reserve(8); diff --git a/tests/ut/cpp/a5/test_tensormap.cpp b/tests/ut/cpp/a5/test_tensormap.cpp index 54f920ee01..0670071709 100644 --- a/tests/ut/cpp/a5/test_tensormap.cpp +++ b/tests/ut/cpp/a5/test_tensormap.cpp @@ -29,6 +29,7 @@ #include #include "utils/device_arena.h" +#include "pto_dep_compute.h" #include "pto_orchestration_api.h" #include "pto_tensormap.h" @@ -56,6 +57,22 @@ static void run_lookup(PTO2TensorMap &tmap, const ChipTensor &tensor, TestLookup }); } +static void +run_lookup(PTO2TensorMap &tmap, const ChipTensor &tensor, TensorAccessKind access_kind, TestLookupResult &out) { + tmap.lookup(tensor, access_kind, [&](PTO2TensorMapEntry &e, OverlapStatus s) -> bool { + out.entries.push_back({&e, s}); + out.count++; + return true; + }); +} + +static DepInputs dep_inputs(const CoreTaskArgs &args) { + return DepInputs{ + args.tensor_count(), args.tensor_data(), args.tag_data(), static_cast(args.explicit_dep_count()), + args.explicit_deps_data(), + }; +} + static ChipTensor make_test_tensor(uint64_t addr, uint32_t shape0, uint32_t ndims = 1, int32_t version = 0) { uint32_t shapes[MAX_TENSOR_DIMS] = {shape0}; return make_tensor_external(reinterpret_cast(addr), shapes, ndims, DataType::FLOAT32, false, version); @@ -99,12 +116,21 @@ class TensorMapTest : public ::testing::Test { TEST_F(TensorMapTest, InitValidState) { EXPECT_EQ(tmap.num_buckets, NUM_BUCKETS); + EXPECT_EQ(tmap.num_reader_buckets, NUM_BUCKETS); EXPECT_EQ(tmap.pool_size, POOL_SIZE); EXPECT_EQ(tmap.next_entry_idx, 0); EXPECT_EQ(tmap.free_num, 0); EXPECT_EQ(tmap.valid_count(), 0); } +TEST(TensorMapLayoutTest, DefaultLayoutCapsSparseReaderBuckets) { + DeviceArena arena; + int32_t window_sizes[PTO2_MAX_RING_DEPTH] = {32, 32, 32, 32}; + auto layout = PTO2TensorMap::reserve_layout_default(arena, window_sizes); + EXPECT_EQ(layout.num_buckets, PTO2_TENSORMAP_NUM_BUCKETS); + EXPECT_EQ(layout.num_reader_buckets, PTO2_TENSORMAP_READER_NUM_BUCKETS); +} + TEST_F(TensorMapTest, InitWithPowerOfTwoBucketsSucceeds) { // The reject path for non-power-of-2 bucket counts is enforced via an // always_assert inside reserve_layout. It is not asserted here because @@ -159,7 +185,7 @@ TEST_F(TensorMapTest, InsertThenLookupFindsProducer) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, tid); + EXPECT_EQ(result.entries[0].entry->access_task_id, tid); } TEST_F(TensorMapTest, LookupEmptyReturnsZero) { @@ -193,12 +219,12 @@ TEST_F(TensorMapTest, InsertDifferentBuffersNoCollision) { TestLookupResult r1; run_lookup(tmap, t1, r1); EXPECT_EQ(r1.count, 1); - EXPECT_EQ(r1.entries[0].entry->producer_task_id, PTO2TaskId::make(0, 0)); + EXPECT_EQ(r1.entries[0].entry->access_task_id, PTO2TaskId::make(0, 0)); TestLookupResult r2; run_lookup(tmap, t2, r2); EXPECT_EQ(r2.count, 1); - EXPECT_EQ(r2.entries[0].entry->producer_task_id, PTO2TaskId::make(0, 1)); + EXPECT_EQ(r2.entries[0].entry->access_task_id, PTO2TaskId::make(0, 1)); } // ============================================================================= @@ -333,7 +359,7 @@ TEST_F(TensorMapTest, StaleEntriesSkippedDuringLookup) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, 1)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(0, 1)); } TEST_F(TensorMapTest, StaleEntriesNotTruncatedAcrossRings) { @@ -349,7 +375,7 @@ TEST_F(TensorMapTest, StaleEntriesNotTruncatedAcrossRings) { run_lookup(tmap, t, result); // Ring 1 task 0 still valid, ring 0 task 0 invalidated ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(1, 0)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(1, 0)); } // ============================================================================= @@ -371,7 +397,7 @@ TEST_F(TensorMapTest, CleanupRetiredRemovesEntriesForRetiredTasks) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, 2)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(0, 2)); } TEST_F(TensorMapTest, CleanupRetiredPreservesOtherRings) { @@ -386,7 +412,7 @@ TEST_F(TensorMapTest, CleanupRetiredPreservesOtherRings) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(1, 0)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(1, 0)); } TEST_F(TensorMapTest, CleanupRetiredFreesEntriesToPool) { @@ -424,7 +450,7 @@ TEST_F(TensorMapTest, CleanupRetiredSparesLaterTaskReusingSlot) { TestLookupResult result; run_lookup(tmap, t, result); ASSERT_EQ(result.count, 1); - EXPECT_EQ(result.entries[0].entry->producer_task_id, PTO2TaskId::make(0, WINDOW_SIZE)); + EXPECT_EQ(result.entries[0].entry->access_task_id, PTO2TaskId::make(0, WINDOW_SIZE)); } // ============================================================================= @@ -448,7 +474,7 @@ TEST_F(TensorMapTest, MultiRingIndependentLookup) { TestLookupResult result2; run_lookup(tmap, t, result2); EXPECT_EQ(result2.count, 1); - EXPECT_EQ(result2.entries[0].entry->producer_task_id, PTO2TaskId::make(1, 3)); + EXPECT_EQ(result2.entries[0].entry->access_task_id, PTO2TaskId::make(1, 3)); } // ============================================================================= @@ -544,12 +570,227 @@ TEST_F(TensorMapTest, RemoveMiddleEntryPreservesChain) { std::set found_locals; for (int i = 0; i < result.count; i++) { - found_locals.insert(result.entries[i].entry->producer_task_id.local()); + found_locals.insert(result.entries[i].entry->access_task_id.local()); } EXPECT_TRUE(found_locals.count(0)); EXPECT_TRUE(found_locals.count(2)); } +TEST_F(TensorMapTest, ReaderAndWriterIndexesImplementRawWarAndOutputExistingOrdering) { + ChipTensor tensor = make_test_tensor(0x4000, 256); + const PTO2TaskId writer0 = PTO2TaskId::make(0, 1); + const PTO2TaskId reader0 = PTO2TaskId::make(1, 2); + const PTO2TaskId writer1 = PTO2TaskId::make(0, 3); + const PTO2TaskId reader1 = PTO2TaskId::make(2, 4); + const PTO2TaskId output_existing = PTO2TaskId::make(0, 5); + tmap.insert(tensor, writer0, TensorAccessKind::WRITER); + + std::vector fanin; + auto emit = [&](PTO2TaskId id) { + fanin.push_back(id); + return true; + }; + + CoreTaskArgs untracked_input_args; + untracked_input_args.add_input(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(untracked_input_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer0})); + EXPECT_EQ(count_registrable_accesses(dep_inputs(untracked_input_args), false), 0); + register_task_accesses(dep_inputs(untracked_input_args), reader0, tmap, false); + + TestLookupResult readers; + run_lookup(tmap, tensor, TensorAccessKind::READER, readers); + EXPECT_EQ(readers.count, 0); + + CoreTaskArgs input_args; + input_args.add_tracked_input(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(input_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer0})); + EXPECT_EQ(count_registrable_accesses(dep_inputs(input_args), false), 1); + register_task_accesses(dep_inputs(input_args), reader0, tmap, false); + + readers = {}; + run_lookup(tmap, tensor, TensorAccessKind::READER, readers); + ASSERT_EQ(readers.count, 1); + EXPECT_EQ(readers.entries[0].entry->access_task_id, reader0); + TestLookupResult writers; + run_lookup(tmap, tensor, TensorAccessKind::WRITER, writers); + ASSERT_EQ(writers.count, 1); + EXPECT_EQ(writers.entries[0].entry->access_task_id, writer0); + + CoreTaskArgs inout_args; + inout_args.add_inout(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(inout_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer0, reader0})); + register_task_accesses(dep_inputs(inout_args), writer1, tmap, false); + + CoreTaskArgs second_input_args; + second_input_args.add_tracked_input(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(second_input_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({writer1})); + register_task_accesses(dep_inputs(second_input_args), reader1, tmap, false); + + CoreTaskArgs output_args; + output_args.add_output(tensor); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(output_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({reader1})); + register_task_accesses(dep_inputs(output_args), output_existing, tmap, false); + + readers = {}; + run_lookup(tmap, tensor, TensorAccessKind::READER, readers); + EXPECT_EQ(readers.count, 0); + writers = {}; + run_lookup(tmap, tensor, TensorAccessKind::WRITER, writers); + ASSERT_EQ(writers.count, 2); + EXPECT_EQ(writers.entries[0].entry->access_task_id, output_existing); + EXPECT_EQ(writers.entries[1].entry->access_task_id, writer1); +} + +TEST_F(TensorMapTest, ReadersAreIndependentAndWriterWaitsForEveryReader) { + ChipTensor tensor = make_test_tensor(0x4100, 256); + const PTO2TaskId reader0 = PTO2TaskId::make(0, 1); + const PTO2TaskId reader1 = PTO2TaskId::make(1, 2); + + CoreTaskArgs input_args; + input_args.add_tracked_input(tensor); + std::vector fanin; + auto emit = [&](PTO2TaskId id) { + fanin.push_back(id); + return true; + }; + ASSERT_TRUE(compute_task_fanin(dep_inputs(input_args), tmap, false, emit)); + EXPECT_TRUE(fanin.empty()); + register_task_accesses(dep_inputs(input_args), reader0, tmap, false); + + ASSERT_TRUE(compute_task_fanin(dep_inputs(input_args), tmap, false, emit)); + EXPECT_TRUE(fanin.empty()); + register_task_accesses(dep_inputs(input_args), reader1, tmap, false); + + CoreTaskArgs writer_args; + writer_args.add_output(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(writer_args), tmap, false, emit)); + std::set fanin_raw; + for (PTO2TaskId id : fanin) + fanin_raw.insert(id.raw); + EXPECT_EQ(fanin_raw, std::set({reader0.raw, reader1.raw})); + EXPECT_EQ(tmap.reader_high_water, 2); +} + +TEST_F(TensorMapTest, DisjointWriteKeepsReaderAndCoveringWriteRetiresIt) { + ChipTensor base = make_test_tensor_2d(0x4200, 256, 1); + uint32_t reader_shape[] = {128, 1}; + uint32_t reader_offset[] = {0, 0}; + ChipTensor reader_view = base.view(reader_shape, reader_offset); + + CoreTaskArgs reader_args; + reader_args.add_tracked_input(reader_view); + register_task_accesses(dep_inputs(reader_args), PTO2TaskId::make(0, 1), tmap, false); + + uint32_t disjoint_offset[] = {128, 0}; + ChipTensor disjoint = base.view(reader_shape, disjoint_offset); + CoreTaskArgs disjoint_args; + disjoint_args.add_output(disjoint); + std::vector fanin; + auto emit = [&](PTO2TaskId id) { + fanin.push_back(id); + return true; + }; + ASSERT_TRUE(compute_task_fanin(dep_inputs(disjoint_args), tmap, false, emit)); + EXPECT_TRUE(fanin.empty()); + register_task_accesses(dep_inputs(disjoint_args), PTO2TaskId::make(0, 2), tmap, false); + EXPECT_EQ(tmap.current_readers(), 1); + + uint32_t partial_shape[] = {128, 1}; + uint32_t partial_offset[] = {64, 0}; + ChipTensor partial = base.view(partial_shape, partial_offset); + CoreTaskArgs partial_args; + partial_args.add_output(partial); + ASSERT_TRUE(compute_task_fanin(dep_inputs(partial_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({PTO2TaskId::make(0, 1)})); + register_task_accesses(dep_inputs(partial_args), PTO2TaskId::make(0, 3), tmap, false); + EXPECT_EQ(tmap.current_readers(), 1); + + CoreTaskArgs covering_args; + covering_args.add_output(base); + fanin.clear(); + ASSERT_TRUE(compute_task_fanin(dep_inputs(covering_args), tmap, false, emit)); + ASSERT_EQ(fanin, std::vector({PTO2TaskId::make(0, 1)})); + register_task_accesses(dep_inputs(covering_args), PTO2TaskId::make(0, 4), tmap, false); + EXPECT_EQ(tmap.current_readers(), 0); +} + +TEST_F(TensorMapTest, RetiredReaderDoesNotCreateStaleWar) { + ChipTensor tensor = make_test_tensor(0x4300, 256); + CoreTaskArgs reader_args; + reader_args.add_tracked_input(tensor); + register_task_accesses(dep_inputs(reader_args), PTO2TaskId::make(0, 0), tmap, false); + ASSERT_EQ(tmap.current_readers(), 1); + + tmap.cleanup_retired(0, 0, 1); + ASSERT_EQ(tmap.current_readers(), 0); + + CoreTaskArgs writer_args; + writer_args.add_output(tensor); + int emitted = 0; + ASSERT_TRUE(compute_task_fanin(dep_inputs(writer_args), tmap, false, [&](PTO2TaskId) { + emitted++; + return true; + })); + EXPECT_EQ(emitted, 0); +} + +TEST_F(TensorMapTest, TrackedAccessesSkipManualScopeAndManualTensor) { + ChipTensor tensor = make_test_tensor(0x5000, 64); + tmap.insert(tensor, PTO2TaskId::make(0, 1)); + int emitted = 0; + auto emit = [&](PTO2TaskId) { + emitted++; + return true; + }; + + CoreTaskArgs manual_scope_args; + manual_scope_args.add_tracked_input(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(manual_scope_args), tmap, true, emit)); + register_task_accesses(dep_inputs(manual_scope_args), PTO2TaskId::make(0, 2), tmap, true); + + tensor.manual_dep = true; + CoreTaskArgs manual_tensor_args; + manual_tensor_args.add_tracked_input(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(manual_tensor_args), tmap, false, emit)); + register_task_accesses(dep_inputs(manual_tensor_args), PTO2TaskId::make(0, 3), tmap, false); + + EXPECT_EQ(emitted, 0); + EXPECT_EQ(tmap.valid_count(), 1); + EXPECT_EQ(tmap.current_readers(), 0); + EXPECT_EQ(tmap.current_writers(), 1); +} + +TEST_F(TensorMapTest, NoDepRetainsCreatorWithoutTensorMapAccess) { + ChipTensor tensor = make_test_tensor(0x5100, 64); + const PTO2TaskId owner = PTO2TaskId::make(2, 7); + tensor.owner_task_id = owner; + tmap.insert(tensor, PTO2TaskId::make(0, 1)); + std::vector fanin; + auto emit = [&](PTO2TaskId id) { + fanin.push_back(id); + return true; + }; + + CoreTaskArgs no_dep_args; + no_dep_args.add_no_dep(tensor); + ASSERT_TRUE(compute_task_fanin(dep_inputs(no_dep_args), tmap, false, emit)); + register_task_accesses(dep_inputs(no_dep_args), PTO2TaskId::make(0, 4), tmap, false); + + EXPECT_EQ(fanin, std::vector({owner})); + EXPECT_EQ(tmap.valid_count(), 1); + EXPECT_EQ(tmap.current_readers(), 0); + EXPECT_EQ(tmap.current_writers(), 1); +} + // ============================================================================= // PTO2TaskId encoding/decoding // ============================================================================= diff --git a/tests/ut/cpp/common/test_hbg_graph_submit_failure.cpp b/tests/ut/cpp/common/test_hbg_graph_submit_failure.cpp index 3651b68f15..25358a381b 100644 --- a/tests/ut/cpp/common/test_hbg_graph_submit_failure.cpp +++ b/tests/ut/cpp/common/test_hbg_graph_submit_failure.cpp @@ -63,7 +63,7 @@ class HbgGraphSubmitFailureTest : public ::testing::Test { } }; -TEST_F(HbgGraphSubmitFailureTest, FaninFailureLatchesFatalWithoutPartialUpload) { +TEST_F(HbgGraphSubmitFailureTest, FaninSpillFailureLatchesFatalWithoutPartialUpload) { std::array storage{}; uint32_t shape[] = {static_cast(storage.size())}; ChipTensor boundary = make_tensor_external(storage.data(), shape, 1); @@ -87,6 +87,8 @@ TEST_F(HbgGraphSubmitFailureTest, FaninFailureLatchesFatalWithoutPartialUpload) ASSERT_TRUE(orch.submit_dummy_task(producer_args).task_id().is_valid()); } + // A full spill pool rejects the first dependency beyond the inline payload. + sched.fanin_spill_capacity = sched.fanin_spill_top; const GraphScopeResult replay = orch.graph_begin(0x1715, boundary_args, 0x1736); EXPECT_TRUE(replay.execute_block); @@ -96,3 +98,28 @@ TEST_F(HbgGraphSubmitFailureTest, FaninFailureLatchesFatalWithoutPartialUpload) EXPECT_EQ(sm_handle->header->orch_error_code.load(std::memory_order_acquire), PTO2_ERROR_DEP_POOL_OVERFLOW); EXPECT_EQ(graph_host_upload_count(*graph_state), uploads_before_failure); } + +TEST_F(HbgGraphSubmitFailureTest, TrackedInputForcesOrdinaryPathFallback) { + std::array storage{}; + uint32_t shape[] = {static_cast(storage.size())}; + ChipTensor boundary = make_tensor_external(storage.data(), shape, 1); + + orch.begin_scope(); + CoreTaskArgs boundary_args; + boundary_args.add_input(boundary); + const GraphScopeResult graph = orch.graph_begin(0x1810, boundary_args, 0x1306); + ASSERT_TRUE(graph.recording); + + CoreTaskArgs node_args; + node_args.add_tracked_input(boundary); + ASSERT_TRUE(orch.submit_dummy_task(node_args).task_id().is_valid()); + + EXPECT_FALSE(orch.graph_end()); + EXPECT_EQ(graph_host_upload_count(*graph_state), 0U); + EXPECT_EQ(sm_handle->header->ring.fc.current_task_index.load(std::memory_order_acquire), 0); + + const TaskOutputTensors fallback = orch.submit_dummy_task(node_args); + EXPECT_TRUE(fallback.task_id().is_valid()); + EXPECT_EQ(sm_handle->header->ring.fc.current_task_index.load(std::memory_order_acquire), 1); + orch.end_scope(); +}