diff --git a/docs/dfx/dep-gen.md b/docs/dfx/dep-gen.md index 58f1cc4d68..25d8c5cd0f 100644 --- a/docs/dfx/dep-gen.md +++ b/docs/dfx/dep-gen.md @@ -56,8 +56,9 @@ inputs to each submit are captured and the graph is reconstructed afterwards. `dep_gen_replay_emit_deps_json` runs every record back through *two* parallel host-resident `PTO2TensorMap` instances that evolve in lockstep: - **Oracle pass** drives the canonical `compute_task_fanin` template - from `pto_dep_compute.h` and collects the producer-id set the - runtime would have emitted. + from `pto_dep_compute.h` and collects the producer-id → `DepFlags` + mapping the runtime would have emitted (flags OR-accumulated per + producer). - **Annotated pass** runs an inlined mirror of STEP A (creator retention) + STEP B (tensormap lookup) against the second map, with a wider callback so each edge gets recorded with its @@ -66,8 +67,9 @@ inputs to each submit are captured and the graph is reconstructed afterwards. (explicit deps), STEP 3 (creator retention + tensormap lookup), STEP 4 (register outputs). Per-successor dedup matches `PTO2FaninBuilder::append_fanin_or_fail`. After both passes finish per - record, the replay asserts the two producer-id sets are equal; if they - diverge, `deps.json` is not written and the function returns non-zero. + record, the replay asserts the two producer-id → `DepFlags` mappings are + equal (same producers and same per-producer flags); if they diverge, + `deps.json` is not written and the function returns non-zero. This is the guarantee against silent shotgun modifications — anyone who changes `compute_task_fanin` semantics will trip the gate immediately and know to update the annotated mirror. @@ -151,10 +153,12 @@ The standard SceneTest path ], "edges": [ {"pred": "0", "succ": "4294967296", "arg": 0, "source": "creator", + "flags": ["wait", "retain"], "tensor_id": "13451765318376212391", "consumer_dtype": "FLOAT32", "consumer_shape": [16384], "consumer_start_offset": "0", "consumer_strides": [1]}, {"pred": "4294967296", "succ": "4294967298", "arg": 0, "source": "tensormap", + "flags": ["wait"], "overlap": "covered", "tensor_id": "9514117477438350967", "consumer_dtype": "FLOAT32", "consumer_shape": [16384], @@ -210,6 +214,7 @@ Each edge is `{pred, succ}` plus annotation. Fields: | `pred`, `succ` | uint64 (string) | always | `PTO2TaskId::raw` of producer and consumer | | `arg` | int32 | always | Consumer's arg-slot index; `-1` for `explicit` source | | `source` | string | always | `explicit` (from `explicit_deps[]`), `creator` (`owner_task_id` retention), or `tensormap` (overlap lookup hit) | +| `flags` | string array | always | Subset of `["wait", "retain"]` — the edge's `DepFlags`. `wait` = ordering (readiness); `retain` = producer lifetime held until the consumer releases. `creator` edges are `["wait","retain"]`; `tensormap` edges `["wait"]`. `explicit` edges are **always recorded as `["wait","retain"]`**: the `DepGenRecord` does not carry per-dep kinds, so a replayed explicit dep cannot distinguish the ordering-only `CoreTaskArgsWithDeps::add_dep_wait()` API from the default. The differential gate is unaffected (both passes read the same constant). At runtime an `add_dep_wait()` edge is genuinely `["wait"]`; that distinction is a known replay limitation, not written to `deps.json`. | | `overlap` | string | `source=tensormap` | `covered` (producer slice fully contains consumer slice) or `other` | | `tensor_id` | uint64 (string) | not `explicit` | Identity of the underlying tensor; cross-references `tensors[]` | | `consumer_dtype` | string | not `explicit` | Element type the consumer reads as | 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 59720c97f1..d81ee13e2d 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -436,8 +436,8 @@ Key members: | 2 | Initialize task descriptor + slot state, copy parameters | | 3 | **Lookup**: for each INPUT/INOUT param, search TensorMap for producers; collect producer pointers in `PTO2FaninBuilder` | | 4 | **Insert**: register OUTPUT/INOUT args in TensorMap | -| 5 | **Record fanin metadata**: store producer pointers in `payload->fanin_inline_slot_states[]` (+ spill pool if >64); claim each live producer by incrementing `fanout_count` under that producer's `fanout_lock`. This step runs **before** `payload.init()`. | -| 6 | **Orch-side wiring / ready publish**: the orchestrator wires live fanout edges into the per-ring dep_pool; zero-fanin and already-completed fanin tasks publish directly to ready queues | +| 5 | **Record fanin metadata**: store producer edges (slot pointer + `DepFlags` packed in the low bits) in `payload->fanin_inline_edges[]` (+ spill pool if >64); claim each live producer by incrementing `fanout_count` under that producer's `fanout_lock`. Creator edges are `DEP_WAIT\|DEP_RETAIN`, tensormap-modifier edges `DEP_WAIT`. This step runs **before** `payload.init()`. | +| 6 | **Orch-side wiring / ready publish**: the orchestrator wires live fanout edges into the per-ring dep_pool; zero-fanin and already-completed fanin tasks publish directly to ready queues. Only `DEP_WAIT` edges gate readiness — they count toward `fanin_count` and are linked onto the producer's `fanout_head` for completion notification. A `DEP_WAIT`-only edge releases its submit→wire retention pin **at wiring** (and on the already-completed fast path), so its producer can be CONSUMED without waiting for this consumer; a `DEP_RETAIN` edge keeps the pin until this consumer's `on_task_release`. A hypothetical `RETAIN`-only edge (none exist yet) would neither gate readiness nor link a fanout node — it only holds the lifetime pin. | > **Note**: Fanout wiring is now completed before publish in the orchestrator submit path. > Scheduler threads consume ready queues directly. 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 2ab6ba4fba..0a9384d44f 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 @@ -20,9 +20,9 @@ * * ORACLE pass (read-only contract): * Drives `compute_task_fanin` (the same template the device orchestrator - * uses in pto_orchestrator.cpp:submit_task) against `tm_oracle`. Emits - * only PTO2TaskId values — the canonical set of producer IDs the runtime - * would have wired. We never widen this template's emit signature: this + * uses in pto_orchestrator.cpp:submit_task) against `tm_oracle`. Its emit + * fires with (PTO2TaskId, DepFlags) — the canonical (producer, WAIT/RETAIN) + * mapping the runtime would have wired, OR-accumulated per producer. This * pass IS the contract, and any future change to `compute_task_fanin` * automatically refreshes the oracle. * @@ -33,12 +33,13 @@ * replay can record per-edge tensor metadata (producer/consumer * shape/offset, dtype, version). * - * After both passes finish per record, we compare the producer-ID set the - * oracle emitted to the producer-ID set the annot pass emitted. They MUST - * match. If they diverge, deps.json is not written and the function returns - * non-zero — this is the "no shotgun modifications" guarantee: anyone who - * changes `compute_task_fanin` will trip this gate immediately and know to - * mirror the change in the annot pass. + * After both passes finish per record, we compare the (producer -> DepFlags) + * mapping the oracle emitted to the one the annot pass emitted. They MUST + * match on both the producer set and each producer's accumulated flags. If they + * diverge, deps.json is not written and the function returns non-zero — this is + * the "no shotgun modifications" guarantee: anyone who changes + * `compute_task_fanin`'s producers or edge flags trips this gate immediately and + * knows to mirror the change in the annot pass. * * STEP 1 (explicit_deps) is emitted at the call site (per pto_dep_compute.h's * "kept at call site" note); both passes run the same explicit-deps loop, so @@ -130,6 +131,21 @@ const char *edge_source_str(EdgeSource s) { return "unknown"; } +// JSON array of the DepFlags bits set on an edge, e.g. ["wait","retain"]. +void write_dep_flags(std::ostream &out, DepFlags flags) { + out << '['; + bool first = true; + if (dep_has_wait(flags)) { + out << "\"wait\""; + first = false; + } + if (dep_has_retain(flags)) { + if (!first) out << ','; + out << "\"retain\""; + } + out << ']'; +} + const char *overlap_status_str(OverlapStatus s) { switch (s) { case OverlapStatus::COVERED: @@ -154,6 +170,7 @@ struct EdgeAnnot { uint64_t succ; int32_t consumer_arg_idx; // -1 for EXPLICIT (not tied to a tensor arg) 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 // Consumer side (the ChipTensor the submitting task is reading). @@ -373,6 +390,8 @@ bool write_deps_json( out << "{\"pred\":\"" << e.pred << "\",\"succ\":\"" << e.succ << '"'; out << ",\"arg\":" << e.consumer_arg_idx; out << ",\"source\":\"" << edge_source_str(e.source) << '"'; + out << ",\"flags\":"; + write_dep_flags(out, e.flags); if (e.source == EdgeSource::TENSORMAP) { out << ",\"overlap\":\"" << overlap_status_str(e.overlap) << '"'; } @@ -516,13 +535,14 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c TensorRef tref_buf[CORE_MAX_TENSOR_ARGS]; TensorArgType atype_buf[CORE_MAX_TENSOR_ARGS]; - // Per-record dedup of producer IDs — must match runtime's + // Per-record producer ID -> accumulated DepFlags — must match runtime's // PTO2FaninBuilder::append_fanin_or_fail semantics, which collapses STEP 1 // (explicit_deps) + STEP A (creator retention) + STEP B (tensormap lookup) - // into a single per-task fanin list. Both oracle and annot use this same - // semantics so the divergence check is meaningful. - std::unordered_set oracle_preds; - std::unordered_set annot_preds; + // into a single per-task fanin edge and OR-accumulates its flags. Both oracle + // and annot use this same semantics so the divergence check compares the + // (producer, flags) mapping rather than the producer-ID set alone. + std::unordered_map oracle_preds; + std::unordered_map annot_preds; // Scratch buffer for assembling full dep lists across overflow chains. // Declared outside the loop so it can be reused (clear() keeps capacity). @@ -683,24 +703,29 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c // gathered base+chain buffer on overflow path). for (int32_t i = 0; i < dc; i++) { uint64_t pred_raw = deps_data[i]; - if (oracle_preds.insert(pred_raw).second) { - // First time this pred is seen at runtime call site. - } - if (annot_preds.insert(pred_raw).second) { + // Explicit deps are recorded conservatively as DEP_WAIT|DEP_RETAIN; + // the DepGenRecord does not carry per-dep kinds, matching Arg's + // set_dependencies default. + oracle_preds[pred_raw] |= (DEP_WAIT | DEP_RETAIN); + bool first = annot_preds.find(pred_raw) == annot_preds.end(); + annot_preds[pred_raw] |= (DEP_WAIT | DEP_RETAIN); + if (first) { EdgeAnnot e{}; e.pred = pred_raw; e.succ = rec.task_id; e.consumer_arg_idx = -1; e.source = EdgeSource::EXPLICIT; + e.flags = DEP_WAIT | DEP_RETAIN; annot_edges.push_back(e); } } // ============ ORACLE pass — drive compute_task_fanin ============ - bool ok = compute_task_fanin(inputs, tm_oracle, in_manual_scope, [&](PTO2TaskId producer) -> bool { - oracle_preds.insert(producer.raw); - return true; - }); + bool ok = + compute_task_fanin(inputs, tm_oracle, in_manual_scope, [&](PTO2TaskId producer, DepFlags kind) -> bool { + oracle_preds[producer.raw] |= kind; + return true; + }); if (!ok) { LOG_ERROR("dep_gen replay: compute_task_fanin returned fatal at task_id=%" PRIu64, rec.task_id); tm_oracle.destroy(); @@ -713,7 +738,9 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c inputs, tm_annot, in_manual_scope, // emit_creator(producer, arg_idx, consumer_tensor) [&](PTO2TaskId producer, int32_t arg_idx, const ChipTensor &consumer) { - if (!annot_preds.insert(producer.raw).second) { + bool first = annot_preds.find(producer.raw) == annot_preds.end(); + annot_preds[producer.raw] |= (DEP_WAIT | DEP_RETAIN); + if (!first) { return; // already covered by an earlier emit on this record } EdgeAnnot e{}; @@ -721,6 +748,7 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c e.succ = rec.task_id; e.consumer_arg_idx = arg_idx; e.source = EdgeSource::CREATOR; + e.flags = DEP_WAIT | DEP_RETAIN; e.tensor_id = make_tensor_id(consumer.buffer.addr, consumer.version); fill_consumer(e, consumer); annot_edges.push_back(e); @@ -735,12 +763,13 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c // producers, both yield their own edges. The producer-id-set // comparison below uses annot_preds, which dedups by pred // only, matching runtime PTO2FaninBuilder semantics. - annot_preds.insert(producer.raw); + annot_preds[producer.raw] |= DEP_WAIT; EdgeAnnot e{}; e.pred = producer.raw; e.succ = rec.task_id; e.consumer_arg_idx = arg_idx; e.source = EdgeSource::TENSORMAP; + e.flags = DEP_WAIT; e.overlap = status; e.tensor_id = make_tensor_id(entry.buffer_addr, entry.version); fill_consumer(e, consumer); @@ -755,15 +784,21 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c "dep_gen replay: DIVERGENCE at task_id=%" PRIu64 " (rec_idx=%zu): oracle has %zu preds, annot has %zu", rec.task_id, rec_i, oracle_preds.size(), annot_preds.size() ); - // Log the symmetric difference for debugging. - for (uint64_t p : oracle_preds) { - if (annot_preds.find(p) == annot_preds.end()) { - LOG_ERROR(" only-in-oracle pred: %" PRIu64, p); + // Log the symmetric difference (missing preds and flag mismatches). + for (const auto &[p, f] : oracle_preds) { + auto it = annot_preds.find(p); + if (it == annot_preds.end()) { + LOG_ERROR(" only-in-oracle pred: %" PRIu64 " flags=%u", p, static_cast(f)); + } else if (it->second != f) { + LOG_ERROR( + " flags mismatch pred: %" PRIu64 " oracle=%u annot=%u", p, static_cast(f), + static_cast(it->second) + ); } } - for (uint64_t p : annot_preds) { + for (const auto &[p, f] : annot_preds) { if (oracle_preds.find(p) == oracle_preds.end()) { - LOG_ERROR(" only-in-annot pred: %" PRIu64, p); + LOG_ERROR(" only-in-annot pred: %" PRIu64 " flags=%u", p, static_cast(f)); } } tm_oracle.destroy(); 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 f8fa586ca5..057cfb8cb0 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 @@ -71,26 +71,29 @@ class CoreTaskArgsWithDeps : private CoreTaskArgs { // the convenience layer reach dependencies only through add_dep() below. /** - * Append one or more dependencies to the bundled buffer. May be called - * multiple times; deps accumulate. Variadic accepts any non-zero number - * of PTO2TaskId arguments. + * Append one or more RETAIN dependencies to the bundled buffer: the producer + * is kept alive (its slot/output buffer retained) until this consumer + * completes. This is the conservative default — use it when the consumer reads + * a tensor whose buffer the producer allocated. May be called multiple times; + * deps accumulate. Variadic accepts any non-zero number of PTO2TaskId args. * * Overflow (more than MAX_DEP_COUNT total) records an error on the * underlying Arg; the error surfaces at submit time. */ template void add_dep(Ids... ids) { - static_assert(sizeof...(Ids) >= 1, "add_dep: at least one task id is required"); - static_assert( - (std::is_same_v, PTO2TaskId> && ...), "add_dep: all arguments must be PTO2TaskId" - ); - if (count_ + sizeof...(Ids) > MAX_DEP_COUNT) { - CoreTaskArgs::set_error( - "CoreTaskArgsWithDeps::add_dep: dep count exceeds MAX_DEP_COUNT (bump the template arg)" - ); - return; - } - ((deps_[count_++] = ids), ...); + add_dep_impl(DEP_WAIT | DEP_RETAIN, ids...); + } + + /** + * Append one or more ordering-only (DEP_WAIT) dependencies: the producer is + * NOT retained and may be reclaimed as soon as it completes and notifies this + * consumer. Use this only when the consumer merely orders after the producer + * and does not read a buffer the producer allocated. + */ + template + void add_dep_wait(Ids... ids) { + add_dep_impl(DEP_WAIT, ids...); } /** @@ -113,12 +116,29 @@ class CoreTaskArgsWithDeps : private CoreTaskArgs { */ CoreTaskArgs &finalize_for_submit() { CoreTaskArgs::set_dependencies(nullptr, 0); - CoreTaskArgs::set_dependencies(deps_, count_); + CoreTaskArgs::set_dependencies_with_kinds(deps_, kinds_, count_); return *this; } private: + template + void add_dep_impl(DepFlags kind, Ids... ids) { + static_assert(sizeof...(Ids) >= 1, "add_dep/add_dep_wait: at least one task id is required"); + static_assert( + (std::is_same_v, PTO2TaskId> && ...), + "add_dep/add_dep_wait: all arguments must be PTO2TaskId" + ); + if (count_ + sizeof...(Ids) > MAX_DEP_COUNT) { + CoreTaskArgs::set_error( + "CoreTaskArgsWithDeps::add_dep/add_dep_wait: dep count exceeds MAX_DEP_COUNT (bump the template arg)" + ); + return; + } + ((kinds_[count_] = kind, deps_[count_] = ids, ++count_), ...); + } + PTO2TaskId deps_[MAX_DEP_COUNT]; + DepFlags kinds_[MAX_DEP_COUNT]; uint32_t count_ = 0; }; 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 ee1ee7b753..16a92380a7 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 @@ -29,7 +29,11 @@ * the minor structural overlap. Replay handles STEP 1 with a one-line loop of its own. * * The Emit callback contract: - * bool emit(PTO2TaskId producer); + * bool emit(PTO2TaskId producer, DepFlags kind); + * - kind is DEP_WAIT|DEP_RETAIN for a Step-A creator edge (the consumer reads + * the producer's allocated buffer, so the producer is retained) and DEP_WAIT + * for a Step-B modifier edge (ordering only; the buffer was allocated + * elsewhere). Duplicate producers OR-accumulate their flags. * - return true to continue (whether or not the producer was actually recorded — * producer-not-alive / dedup-hit / etc. all return true silently) * - return false to signal fatal (e.g. fanin spill overflow); caller bails @@ -93,10 +97,11 @@ compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_m const ChipTensor *tensor = &inputs.tensors[i].ref(); - // Step A: creator retention — all existing tensors extend their creator lifetime. + // Step A: creator retention — reading a tensor retains its allocator, so + // the creator edge carries both ordering and lifetime. PTO2TaskId owner = tensor->owner_task_id; if (owner.is_valid()) { - if (!emit(owner)) { + if (!emit(owner, DEP_WAIT | DEP_RETAIN)) { return false; } } @@ -111,7 +116,17 @@ 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 { - if (!emit(entry.producer_task_id)) { + // 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)) { fatal = true; return false; // stop iteration } 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 35399df5da..b543298766 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp @@ -237,16 +237,18 @@ static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { struct PTO2FaninBuilder { PTO2FaninBuilder(PTO2OrchestratorState *orch, PTO2FaninPool &spill_pool, uint32_t seen_epoch) : count(0), + wait_count(0), spill_start(0), orch(orch), seen_epoch(seen_epoch), spill_pool(spill_pool) {} - int32_t count{0}; + int32_t count{0}; // total fanin edges (all flag combinations) + int32_t wait_count{0}; // edges carrying DEP_WAIT — sizes readiness accounting int32_t spill_start{0}; PTO2OrchestratorState *orch{nullptr}; uint32_t seen_epoch{0}; PTO2FaninPool &spill_pool; - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP]; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP]; template PTO2FaninForEachReturn for_each(Fn &&fn) const { @@ -265,11 +267,57 @@ struct PTO2FaninBuilder { seen[slot] = seen_epoch; return false; } + + // Append a new edge (caller has already claimed the producer's fanout pin). + void push_edge(PTO2FaninSpillEntry &entry, PTO2TaskSlotState *prod_state, DepFlags kind) { + entry.set(prod_state, kind); + count++; + if (dep_has_wait(kind)) { + wait_count++; + } + } + + // Dedup path: a producer already recorded this submission is reached again + // (e.g. as both a creator and a modifier). OR the new flags into the existing + // edge instead of adding a second one — no extra fanout pin is claimed. + void or_flags_into_existing(PTO2TaskSlotState *prod_state, DepFlags kind) { + for (int32_t i = 0; i < count; i++) { + PTO2FaninSpillEntry &entry = entry_at(i); + if (entry.slot_state() == prod_state) { + accumulate_flags(entry, kind); + return; + } + } + // mark_seen reported this producer already owns an edge this submission, + // so it must be present. A miss means the seen-set and the builder + // disagree — a logic error, not a runtime condition. + always_assert(false && "or_flags_into_existing: deduped producer missing from the fanin builder"); + } + +private: + // The i-th appended fanin edge: inline for i < PTO2_FANIN_INLINE_CAP, else in + // the spill ring at the same linear->physical position for_each_fanin_storage + // walks. Single source of the wrap arithmetic. + PTO2FaninSpillEntry &entry_at(int32_t i) { + if (i < PTO2_FANIN_INLINE_CAP) { + return inline_slots[i]; + } + int32_t spill_idx = i - PTO2_FANIN_INLINE_CAP; + return spill_pool.base[(spill_start % spill_pool.capacity + spill_idx) % spill_pool.capacity]; + } + + void accumulate_flags(PTO2FaninSpillEntry &entry, DepFlags kind) { + bool had_wait = dep_has_wait(entry.flags()); + entry.add_flags(kind); + if (!had_wait && dep_has_wait(kind)) { + wait_count++; + } + } }; static bool append_fanin_or_fail( PTO2OrchestratorState *orch, uint8_t prod_ring, int32_t prod_slot, PTO2TaskSlotState *prod_state, - PTO2TaskId producer_task_id, PTO2FaninBuilder *fanin_builder, uint8_t ring_id + PTO2TaskId producer_task_id, PTO2FaninBuilder *fanin_builder, uint8_t ring_id, DepFlags kind ) { // Decide-and-claim under the producer's fanout_lock. Two conditions make this // resolved slot a non-dependency, and both must be checked together with the @@ -281,12 +329,13 @@ static bool append_fanin_or_fail( // producer; ++'ing it would corrupt an unrelated task. // (2) Already CONSUMED in place — finished, output ready, no real edge. // In either case, adding it to the fanin and bumping fanout_count would leave - // a stale ++/release pair (Orch-side wiring drops the fanout edge but keeps - // the fanin slot, so on_task_release still release_producer()'s it) that - // desyncs the slot's refcount (rc != fc) and wedges in-order reclaim. Claiming a live - // producer under the lock pins it: fanout_count now counts us, so it cannot - // reach CONSUMED (rc == fc) until we release it in on_task_release, keeping the - // slot's generation stable until then. check_and_handle_consumed flips + // a stale ++/release pair that desyncs the slot's refcount (rc != fc) and + // wedges in-order reclaim. Every edge (regardless of DepFlags) claims one + // fanout_count++ here: it pins the producer's slot across the submit->wire + // window so it cannot be CONSUMED + reused before wire_fanin_task links the + // consumer. The pin's release differs by kind — an ordering-only (DEP_WAIT + // without DEP_RETAIN) edge releases it at wiring; a DEP_RETAIN edge holds it + // until the consumer's on_task_release. check_and_handle_consumed flips // COMPLETED->CONSUMED under the same lock, so the check and the ++ are atomic // against the consume. fanout_count is lock-protected per the // PTO2TaskSlotState contract. @@ -295,14 +344,16 @@ static bool append_fanin_or_fail( // gone check. mark_seen keys only on (ring, slot); a stale owner that resolves // to a reused slot must not record it as seen, or a later dependency on the // live generation in the same submission would hit mark_seen and be skipped - // without claiming it (dropped edge). Marking only when !gone keeps the dedup - // keyed to the live producer, and doing it before the ++ still suppresses a - // double-count for a producer named twice in one submission. + // without claiming it (dropped edge). A duplicate live producer claims no new + // pin and adds no new edge; its flags are OR-accumulated into the existing + // edge so the stronger of several discovery reasons (creator vs modifier vs + // explicit) always wins, independent of the order they are reached in. prod_state->lock_fanout(); PTO2TaskState pstate = prod_state->task_state.load(std::memory_order_acquire); bool gone = prod_state->task == nullptr || prod_state->task->task_id.local() != producer_task_id.local() || pstate == PTO2_TASK_CONSUMED; - bool claim = !gone && !fanin_builder->mark_seen(prod_ring, prod_slot); + bool already_seen = !gone && fanin_builder->mark_seen(prod_ring, prod_slot); + bool claim = !gone && !already_seen; int32_t fanout_now = -1; if (claim) { // Low bits hold the consumer count; bit31 is the scope ref. The consumer @@ -329,14 +380,18 @@ static bool append_fanin_or_fail( static_cast(producer_task_id.ring()), producer_task_id.local(), PTO2_DEP_DEGREE_DEBUG_THRESHOLD ); } - // gone (stale/consumed) or an already-seen duplicate live producer: no new - // fanin edge either way. - if (!claim) { + // Stale/consumed producer: no edge at all. + if (gone) { + return true; + } + // Duplicate live producer: fold the flags into the edge already recorded. + if (already_seen) { + fanin_builder->or_flags_into_existing(prod_state, kind); return true; } if (fanin_builder->count < PTO2_FANIN_INLINE_CAP) { - fanin_builder->inline_slots[fanin_builder->count++] = prod_state; + fanin_builder->push_edge(fanin_builder->inline_slots[fanin_builder->count], prod_state, kind); return true; } @@ -354,21 +409,24 @@ static bool append_fanin_or_fail( if (fanin_builder->count == PTO2_FANIN_INLINE_CAP) { fanin_builder->spill_start = spill_idx; } - entry->slot_state = prod_state; - fanin_builder->count++; + fanin_builder->push_edge(*entry, prod_state, kind); return true; } static bool all_claimed_fanin_completed(const PTO2FaninBuilder &fanin_builder) { if (fanin_builder.count == 0) return true; - return fanin_builder.for_each([](PTO2TaskSlotState *producer) -> bool { + // Only DEP_WAIT edges gate readiness; a retention-only edge never blocks + // dispatch, so it is treated as satisfied here. + return fanin_builder.for_each([](PTO2TaskSlotState *producer, DepFlags flags) -> bool { + if (!dep_has_wait(flags)) return true; return producer != nullptr && producer->task_state.load(std::memory_order_acquire) >= PTO2_TASK_COMPLETED; }); } static bool all_claimed_fanin_allow_early_resolve(const PTO2FaninBuilder &fanin_builder) { if (fanin_builder.count == 0) return true; - return fanin_builder.for_each([](PTO2TaskSlotState *producer) -> bool { + return fanin_builder.for_each([](PTO2TaskSlotState *producer, DepFlags flags) -> bool { + if (!dep_has_wait(flags)) return true; return producer != nullptr && producer->task_attrs.allow_early_resolve(); }); } @@ -393,23 +451,40 @@ void PTO2OrchestratorState::wire_fanin_task(PTO2TaskSlotState &slot_state, int32 int32_t completed_fanin = 0; int32_t early_propagated = 0; bool early_disqualified = false; - for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer) { - producer->lock_fanout(); - int32_t pstate = producer->task_state.load(std::memory_order_acquire); - if (!early_disqualified && !producer->task_attrs.allow_early_resolve()) { - early_disqualified = true; - } - if (pstate >= PTO2_TASK_COMPLETED) { - completed_fanin++; - } else { - producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, &slot_state); - // The marker shares fanout_lock with propagation's snapshot. A set - // marker means this edge is outside that snapshot and needs a seed. - if (!early_disqualified && producer->has_dispatch_propagated()) { - early_propagated++; + for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer, DepFlags flags) { + // Only DEP_WAIT edges contribute to readiness: they gate fanin and are + // linked onto the producer's fanout_head for completion notification. A + // retention-only edge carries no ordering, so it is skipped here. + if (dep_has_wait(flags)) { + producer->lock_fanout(); + int32_t pstate = producer->task_state.load(std::memory_order_acquire); + if (!early_disqualified && !producer->task_attrs.allow_early_resolve()) { + early_disqualified = true; } + if (pstate >= PTO2_TASK_COMPLETED) { + completed_fanin++; + } else { + producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, &slot_state); + // The marker shares fanout_lock with propagation's snapshot. A set + // marker means this edge is outside that snapshot and needs a seed. + if (!early_disqualified && producer->has_dispatch_propagated()) { + early_propagated++; + } + } + producer->unlock_fanout(); + } + // The submit->wire pin protects an ordering-only edge only until the + // consumer is linked. With the consumer now on fanout_head (or already + // seen as completed), release it so the producer can be CONSUMED without + // waiting for this consumer. A DEP_RETAIN edge keeps the pin until the + // consumer's on_task_release. + if (dep_has_wait(flags) && !dep_has_retain(flags)) { + // Wiring-phase atomics (this release, plus the lock_fanout / dep_pool.prepend / + // fanin_refcount ops around it) are not bucketed: g_orch_args_atomic_count + // covers the submit/dep-claim phase only, whose g_orch_args_cycle window has + // already closed by the time wiring runs. + sched->release_producer(*producer); } - producer->unlock_fanout(); }); // Completed producers and edges outside a producer's one-shot propagation @@ -897,7 +972,8 @@ static TaskOutputTensors submit_task_common( int32_t dep_slot = dep_ring.get_slot_by_task_id(dep_local_task_id); PTO2TaskSlotState *producer_slot_state = &dep_ring.get_slot_state_by_slot(dep_slot); if (!append_fanin_or_fail( - orch, dep_ring_id, dep_slot, producer_slot_state, dep_task_id, &fanin_builder, ring_id + orch, dep_ring_id, dep_slot, producer_slot_state, dep_task_id, &fanin_builder, ring_id, + args.explicit_dep_kind(i) )) { return result; } @@ -909,12 +985,14 @@ static TaskOutputTensors submit_task_common( args.explicit_deps_data(), }; - auto runtime_emit = [&](PTO2TaskId producer_task_id) -> bool { + auto runtime_emit = [&](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); - return append_fanin_or_fail(orch, prod_ring, prod_slot, prod_state, producer_task_id, &fanin_builder, ring_id); + 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)) { @@ -953,6 +1031,16 @@ static TaskOutputTensors submit_task_common( // here) is what prevents a producer from transitioning to CONSUMED between // the dependency decision and the claim. int32_t inline_count = std::min(fanin_builder.count, PTO2_FANIN_INLINE_CAP); + // Every fanin edge produced here carries DEP_WAIT (creator = WAIT|RETAIN, + // modifier = WAIT, explicit defaults to WAIT|RETAIN or opts into WAIT), so + // wait_count == count. fanin_actual_count therefore doubles as the WAIT-edge + // count that the early-dispatch threshold (dispatch_fanin, which counts only + // WAIT producers) is compared against. A future RETAIN-only edge would break + // that equality and must carry its own WAIT-edge count for that comparison. + always_assert( + fanin_builder.wait_count == fanin_builder.count && + "fanin_actual_count is the early-dispatch WAIT denominator; a non-WAIT edge needs a separate count" + ); // Store fanin metadata in payload for scheduler to iterate payload.fanin_actual_count = fanin_builder.count; // fanin_builder.count is finalized here and submit runs once per task, so @@ -967,7 +1055,7 @@ static TaskOutputTensors submit_task_common( payload.fanin_spill_start = fanin_builder.spill_start; payload.fanin_spill_pool = &fanin_builder.spill_pool; for (int i = 0; i < inline_count; i++) { - payload.fanin_inline_slot_states[i] = fanin_builder.inline_slots[i]; + payload.fanin_inline_edges[i] = fanin_builder.inline_slots[i]; } payload.init(args, result, prepared.alloc_result, layout); @@ -1018,16 +1106,24 @@ static TaskOutputTensors submit_task_common( orch->mark_dep_pool_position(cur_slot_state); sched->push_ready_routed(&cur_slot_state); } else if (all_claimed_fanin_completed(fanin_builder)) { - int32_t ready_seed = fanin_builder.count + 1; + int32_t ready_seed = fanin_builder.wait_count + 1; cur_slot_state.fanin_count = ready_seed; if (all_claimed_fanin_allow_early_resolve(fanin_builder)) { payload.dispatch_fanin.store(fanin_builder.count, std::memory_order_release); } cur_slot_state.fanin_refcount.store(ready_seed, std::memory_order_release); + // wire_fanin_task is skipped here, so its ordering-only pin release runs + // on this path too: an edge without retention drops its submit->wire pin + // so the (already completed) producer can be CONSUMED. + for_each_fanin_slot_state(payload, [&](PTO2TaskSlotState *producer, DepFlags flags) { + if (dep_has_wait(flags) && !dep_has_retain(flags)) { + sched->release_producer(*producer); // wiring-phase atomic, not bucketed (see wire_fanin_task) + } + }); orch->mark_dep_pool_position(cur_slot_state); sched->push_ready_routed(&cur_slot_state); } else { - if (!orch_wire_live_fanin_task(orch, cur_slot_state, fanin_builder.count)) { + if (!orch_wire_live_fanin_task(orch, cur_slot_state, fanin_builder.wait_count)) { return result; } } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h index eb35e7cfac..b376e99ab2 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h @@ -519,7 +519,7 @@ struct PTO2FaninPool { tail = 1; high_water = 0; reclaim_task_cursor = 0; - base[0].slot_state = nullptr; + base[0].clear(); error_code_ptr = in_error_code_ptr; } @@ -528,7 +528,7 @@ struct PTO2FaninPool { tail = 1; high_water = 0; reclaim_task_cursor = 0; - base[0].slot_state = nullptr; + base[0].clear(); error_code_ptr = in_error_code_ptr; } @@ -575,14 +575,16 @@ struct PTO2FaninPool { }; template -using PTO2FaninCallbackResult = std::invoke_result_t; +using PTO2FaninCallbackResult = std::invoke_result_t; template using PTO2FaninForEachReturn = std::conditional_t, void>, void, bool>; +// Visit each fanin edge as (producer slot, DepFlags). Inline and spill entries +// share the packed PTO2FaninSpillEntry layout, so both are unpacked the same way. template inline PTO2FaninForEachReturn for_each_fanin_storage( - InlineSlots &&inline_slot_states, int32_t fanin_count, int32_t spill_start, PTO2FaninPool &spill_pool, Fn &&fn + InlineSlots &&inline_edges, int32_t fanin_count, int32_t spill_start, PTO2FaninPool &spill_pool, Fn &&fn ) { using FaninCallbackResult = PTO2FaninCallbackResult; static_assert( @@ -593,7 +595,7 @@ inline PTO2FaninForEachReturn for_each_fanin_storage( if constexpr (std::is_void_v) { int32_t inline_count = std::min(fanin_count, PTO2_FANIN_INLINE_CAP); for (int32_t i = 0; i < inline_count; i++) { - fn(inline_slot_states[i]); + fn(inline_edges[i].slot_state(), inline_edges[i].flags()); } int32_t spill_count = fanin_count - inline_count; @@ -605,18 +607,18 @@ inline PTO2FaninForEachReturn for_each_fanin_storage( int32_t first_count = std::min(spill_count, spill_pool.capacity - start_idx); PTO2FaninSpillEntry *first = spill_pool.base + start_idx; for (int32_t i = 0; i < first_count; i++) { - fn(first[i].slot_state); + fn(first[i].slot_state(), first[i].flags()); } int32_t second_count = spill_count - first_count; for (int32_t i = 0; i < second_count; i++) { - fn(spill_pool.base[i].slot_state); + fn(spill_pool.base[i].slot_state(), spill_pool.base[i].flags()); } return; } else { int32_t inline_count = std::min(fanin_count, PTO2_FANIN_INLINE_CAP); for (int32_t i = 0; i < inline_count; i++) { - if (!fn(inline_slot_states[i])) { + if (!fn(inline_edges[i].slot_state(), inline_edges[i].flags())) { return false; } } @@ -630,14 +632,14 @@ inline PTO2FaninForEachReturn for_each_fanin_storage( int32_t first_count = std::min(spill_count, spill_pool.capacity - start_idx); PTO2FaninSpillEntry *first = spill_pool.base + start_idx; for (int32_t i = 0; i < first_count; i++) { - if (!fn(first[i].slot_state)) { + if (!fn(first[i].slot_state(), first[i].flags())) { return false; } } int32_t second_count = spill_count - first_count; for (int32_t i = 0; i < second_count; i++) { - if (!fn(spill_pool.base[i].slot_state)) { + if (!fn(spill_pool.base[i].slot_state(), spill_pool.base[i].flags())) { return false; } } @@ -648,8 +650,8 @@ inline PTO2FaninForEachReturn for_each_fanin_storage( template inline PTO2FaninForEachReturn for_each_fanin_slot_state(const PTO2TaskPayload &payload, Fn &&fn) { return for_each_fanin_storage( - payload.fanin_inline_slot_states, payload.fanin_actual_count, payload.fanin_spill_start, - *payload.fanin_spill_pool, static_cast(fn) + payload.fanin_inline_edges, payload.fanin_actual_count, payload.fanin_spill_start, *payload.fanin_spill_pool, + static_cast(fn) ); } 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 fc115b388e..af7d4742fa 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 @@ -162,8 +162,28 @@ struct PTO2OutputLayout { */ struct PTO2TaskSlotState; // Forward declaration struct PTO2FaninPool; // Forward declaration + +// One fanin edge: a producer slot pointer with the per-edge DepFlags packed into +// bits 0..1 of the pointer. PTO2TaskSlotState is alignas(64), so bits 0..5 are +// always zero in a real pointer. Used for both the inline fanin array and the +// spill pool, keeping sizeof == sizeof(uintptr_t) so neither footprint grows. struct PTO2FaninSpillEntry { - PTO2TaskSlotState *slot_state; + static constexpr uintptr_t FLAG_MASK = 0x3; + // No in-class initializer: the type stays trivially default-constructible so + // the payload's fanin_inline_edges[64] and the orchestrator's per-submit + // builder array are default-initialized at zero cost (only entries [0, count) + // are written via set()). Value-init (`PTO2FaninSpillEntry{}`) still zeroes it. + uintptr_t packed; + + PTO2TaskSlotState *slot_state() const { return reinterpret_cast(packed & ~FLAG_MASK); } + DepFlags flags() const { return static_cast(packed & FLAG_MASK); } + // Only bits within FLAG_MASK are stored; any bit outside it (a malformed + // DepFlags value) is masked off so it can never corrupt the slot pointer. + void set(PTO2TaskSlotState *s, DepFlags f) { + packed = reinterpret_cast(s) | (static_cast(f) & FLAG_MASK); + } + void add_flags(DepFlags f) { packed |= (static_cast(f) & FLAG_MASK); } + void clear() { packed = 0; } }; static_assert(sizeof(PTO2FaninSpillEntry) == sizeof(uintptr_t)); @@ -255,7 +275,10 @@ struct PTO2TaskPayload { int32_t fanin_actual_count{0}; // Actual fanin count (without the +1 redundance) int32_t fanin_spill_start{0}; // Linear start index in fanin spill pool (0 = no spill) PTO2FaninPool *fanin_spill_pool{nullptr}; - PTO2TaskSlotState *fanin_inline_slot_states[PTO2_FANIN_INLINE_CAP]; + // Inline fanin edges (producer slot + packed DepFlags). Spill beyond + // PTO2_FANIN_INLINE_CAP goes to fanin_spill_pool. Same packed layout as the + // spill entries, so the array footprint is unchanged. + PTO2FaninSpillEntry fanin_inline_edges[PTO2_FANIN_INLINE_CAP]; // Early-dispatch metadata (AICPU-side only). Ordered by descending // alignment so the block packs without internal padding. Cache line 8 // contains the rarely-touched fanin tail rather than the hot tensor/scalar @@ -378,9 +401,7 @@ struct PTO2TaskPayload { // PTO2TaskPayload layout verification (offsetof requires complete type). static_assert(offsetof(PTO2TaskPayload, fanin_spill_pool) == 16, "spill pool pointer layout drift"); -static_assert( - offsetof(PTO2TaskPayload, fanin_inline_slot_states) == 24, "inline fanin array must follow spill metadata" -); +static_assert(offsetof(PTO2TaskPayload, fanin_inline_edges) == 24, "inline fanin array must follow spill metadata"); static_assert( offsetof(PTO2TaskPayload, predicate) == 576, "dispatch predicate occupies cache line 9 at fixed byte 576 (before tensors, never moves)" @@ -631,3 +652,9 @@ struct alignas(64) PTO2TaskSlotState { }; static_assert(sizeof(PTO2TaskSlotState) == 64); +// PTO2FaninSpillEntry packs DepFlags into the low bits of a PTO2TaskSlotState*. +// That is only lossless while the type is aligned past those tag bits. +static_assert( + alignof(PTO2TaskSlotState) > PTO2FaninSpillEntry::FLAG_MASK, + "PTO2TaskSlotState alignment must exceed PTO2FaninSpillEntry::FLAG_MASK so the packed DepFlags bits are free" +); 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 51ed7eb2ac..f8d86d6bac 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h @@ -70,6 +70,42 @@ enum class PTO2ScopeMode : uint8_t { MANUAL = 1, }; +/** + * Orthogonal dependency-edge semantics, stored per fanin edge. + * + * DEP_WAIT — readiness/ordering: the consumer cannot become ready until the + * producer completes. Contributes to the consumer's fanin count + * and to the producer's fanout notification list. + * DEP_RETAIN — lifetime: the producer's slot and packed output buffer must + * stay alive until the consumer releases it. Holds a + * fanout_count reference from submit until on_task_release. + * + * The two are independent; an edge may carry either, both, or (transiently + * during construction) neither. A creator edge is DEP_WAIT|DEP_RETAIN (the + * consumer reads the producer's allocated buffer). A tensormap modifier edge is + * DEP_WAIT only (the buffer was allocated elsewhere, so only ordering is + * needed). When one (producer, consumer) pair is discovered for several + * reasons, the flags are OR-accumulated. + */ +enum DepFlags : uint8_t { + DEP_NONE = 0, + DEP_WAIT = 1u << 0, + DEP_RETAIN = 1u << 1, +}; + +constexpr DepFlags operator|(DepFlags a, DepFlags b) { + return static_cast(static_cast(a) | static_cast(b)); +} +constexpr DepFlags operator&(DepFlags a, DepFlags b) { + return static_cast(static_cast(a) & static_cast(b)); +} +constexpr DepFlags &operator|=(DepFlags &a, DepFlags b) { + a = a | b; + return a; +} +constexpr bool dep_has_wait(DepFlags f) { return (f & DEP_WAIT) != DEP_NONE; } +constexpr bool dep_has_retain(DepFlags f) { return (f & DEP_RETAIN) != DEP_NONE; } + /** * TaskOutputTensors — returned by submit, holds materialized output ChipTensors. * @@ -286,6 +322,7 @@ struct Arg : TaskArgsTpl { #endif explicit_deps_ = nullptr; explicit_dep_count_ = 0; + explicit_dep_kinds_ = nullptr; allow_early_resolve_ = false; predicate_ = CoreTaskPredicate{}; task_timing_slot_ = TASK_TIMING_SLOT_NONE; @@ -396,6 +433,7 @@ struct Arg : TaskArgsTpl { if (count == 0) { explicit_deps_ = nullptr; explicit_dep_count_ = 0; + explicit_dep_kinds_ = nullptr; return; } if (deps == nullptr) { @@ -408,6 +446,35 @@ struct Arg : TaskArgsTpl { } explicit_deps_ = deps; explicit_dep_count_ = count; + explicit_dep_kinds_ = nullptr; + } + + /** + * Same as set_dependencies() but attaches a per-dep DepFlags array so the + * caller can mark an explicit dep as ordering-only (DEP_WAIT) rather than the + * conservative DEP_WAIT|DEP_RETAIN default. The kinds array is caller-owned + * and must outlive submit (same lifetime rule as deps). A null kinds array + * with count > 0 is rejected; use set_dependencies() for the + * conservative DEP_WAIT|DEP_RETAIN default. + */ + void set_dependencies_with_kinds(const PTO2TaskId *deps, const DepFlags *kinds, uint32_t count) { + if (count == 0) { + explicit_deps_ = nullptr; + explicit_dep_count_ = 0; + explicit_dep_kinds_ = nullptr; + return; + } + if (deps == nullptr || kinds == nullptr) { + set_error("set_dependencies_with_kinds: deps and kinds must not be null when count > 0"); + return; + } + if (explicit_deps_ != nullptr) { + set_error("set_dependencies_with_kinds: may be called at most once per Arg"); + return; + } + explicit_deps_ = deps; + explicit_dep_kinds_ = kinds; + explicit_dep_count_ = count; } uint32_t explicit_dep_count() const { return explicit_dep_count_; } @@ -417,6 +484,16 @@ struct Arg : TaskArgsTpl { return explicit_deps_[index]; } + /** + * Flags of the i-th explicit dep. Returns DEP_WAIT|DEP_RETAIN when no kinds + * array is attached (plain set_dependencies), so an explicit dep protecting a + * runtime-created output is never silently weakened to ordering-only. + */ + DepFlags explicit_dep_kind(uint32_t index) const { + always_assert(index < explicit_dep_count_); + return explicit_dep_kinds_ ? explicit_dep_kinds_[index] : (DEP_WAIT | DEP_RETAIN); + } + const PTO2TaskId *explicit_deps_data() const { return explicit_deps_; } /** @@ -516,6 +593,7 @@ struct Arg : TaskArgsTpl { DumpArgSelection dump_arg_selection_; #endif const PTO2TaskId *explicit_deps_{nullptr}; + const DepFlags *explicit_dep_kinds_{nullptr}; uint32_t explicit_dep_count_{0}; #if SIMPLER_DFX template diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index 75aa3dc007..c3809a0880 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -1303,7 +1303,8 @@ struct PTO2SchedulerState { /** * Cold path: release producers (fanin traversal) + check self for CONSUMED. - * Returns fanin edge count for profiling. + * Returns the number of retained (DEP_RETAIN) producers actually released — + * ordering-only edges dropped their pin at wiring and are skipped here. */ #if SIMPLER_SCHED_PROFILING @@ -1318,7 +1319,16 @@ struct PTO2SchedulerState { int32_t on_task_release(PTO2TaskSlotState &slot_state) { #endif PTO2TaskPayload *payload = slot_state.payload; - for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer_slot_state) { + int32_t released = 0; + // Only DEP_RETAIN edges still hold a fanout pin at completion: an + // ordering-only edge released its submit->wire pin at wiring, so releasing + // it again here would over-count fanout_refcount against fanout_count and + // break the rc == fc consume invariant. + for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer_slot_state, DepFlags flags) { + if (!dep_has_retain(flags)) { + return; + } + released++; #if SIMPLER_SCHED_PROFILING release_producer(*producer_slot_state, fanin_atomics); #else @@ -1340,7 +1350,7 @@ struct PTO2SchedulerState { #else check_and_handle_consumed(slot_state); #endif - return payload->fanin_actual_count; + return released; } // === Cold-path API (defined in pto_scheduler.cpp) === 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 130644295a..fc52dd4b9c 100644 --- a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp @@ -304,7 +304,7 @@ namespace { // host_build_graph is host-orchestration-first: the HOST dlopens the // orchestration .so and runs it to completion. The shared memory + arena carry // host-DDR cross-task pointers (slot_state.task/payload, -// payload.fanin_inline_slot_states[], dep_pool/ready queues); the host relocates them to +// payload.fanin_local_ids[], dep_pool/ready queues); the host relocates them to // their final device addresses (relocate_host_orch_image, below) BEFORE the H2D // copy, so the device receives a fully device-addressed image and schedules // only — no on-device pointer fixup. @@ -375,7 +375,7 @@ struct HostOrchEntryPoints { // and boots scheduler-only with no on-device pointer fixup. // // Relocated pointers span TWO regions with DIFFERENT deltas: the SM block -// (slot_state.task/.payload, fanin_inline_slot_states[], dep-entry.slot_state, +// (slot_state.task/.payload, fanin_local_ids[], dep-entry.slot_state, // ready-queue slot.slot_state) and the arena block (slot_state.fanout_head, // dep-entry.next point into the SM but live in the arena). // Rather than track which delta each field needs, relocate() classifies every 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 150a0f6781..fca7d1369d 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -437,8 +437,8 @@ Key members: | 2 | Initialize task descriptor + slot state, copy parameters | | 3 | **Lookup**: for each INPUT/INOUT param, search TensorMap for producers; collect producer pointers in `PTO2FaninBuilder` | | 4 | **Insert**: register OUTPUT/INOUT args in TensorMap | -| 5 | **Record fanin metadata**: store producer pointers in `payload->fanin_inline_slot_states[]` (+ spill pool if >64); claim each live producer by incrementing `fanout_count` under that producer's `fanout_lock`. This step runs **before** `payload.init()`. | -| 6 | **Orch-side wiring / ready publish**: the orchestrator wires live fanout edges into the per-ring dep_pool; zero-fanin and already-completed fanin tasks publish directly to ready queues | +| 5 | **Record fanin metadata**: store producer edges (slot pointer + `DepFlags` packed in the low bits) in `payload->fanin_inline_edges[]` (+ spill pool if >64); claim each live producer by incrementing `fanout_count` under that producer's `fanout_lock`. Creator edges are `DEP_WAIT\|DEP_RETAIN`, tensormap-modifier edges `DEP_WAIT`. This step runs **before** `payload.init()`. | +| 6 | **Orch-side wiring / ready publish**: the orchestrator wires live fanout edges into the per-ring dep_pool; zero-fanin and already-completed fanin tasks publish directly to ready queues. Only `DEP_WAIT` edges gate readiness — they count toward `fanin_count` and are linked onto the producer's `fanout_head` for completion notification. A `DEP_WAIT`-only edge releases its submit→wire retention pin **at wiring** (and on the already-completed fast path), so its producer can be CONSUMED without waiting for this consumer; a `DEP_RETAIN` edge keeps the pin until this consumer's `on_task_release`. A hypothetical `RETAIN`-only edge (none exist yet) would neither gate readiness nor link a fanout node — it only holds the lifetime pin. | > **Note**: Fanout wiring is now completed before publish in the orchestrator submit path. > Scheduler threads consume ready queues directly. 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 2ab6ba4fba..0a9384d44f 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 @@ -20,9 +20,9 @@ * * ORACLE pass (read-only contract): * Drives `compute_task_fanin` (the same template the device orchestrator - * uses in pto_orchestrator.cpp:submit_task) against `tm_oracle`. Emits - * only PTO2TaskId values — the canonical set of producer IDs the runtime - * would have wired. We never widen this template's emit signature: this + * uses in pto_orchestrator.cpp:submit_task) against `tm_oracle`. Its emit + * fires with (PTO2TaskId, DepFlags) — the canonical (producer, WAIT/RETAIN) + * mapping the runtime would have wired, OR-accumulated per producer. This * pass IS the contract, and any future change to `compute_task_fanin` * automatically refreshes the oracle. * @@ -33,12 +33,13 @@ * replay can record per-edge tensor metadata (producer/consumer * shape/offset, dtype, version). * - * After both passes finish per record, we compare the producer-ID set the - * oracle emitted to the producer-ID set the annot pass emitted. They MUST - * match. If they diverge, deps.json is not written and the function returns - * non-zero — this is the "no shotgun modifications" guarantee: anyone who - * changes `compute_task_fanin` will trip this gate immediately and know to - * mirror the change in the annot pass. + * After both passes finish per record, we compare the (producer -> DepFlags) + * mapping the oracle emitted to the one the annot pass emitted. They MUST + * match on both the producer set and each producer's accumulated flags. If they + * diverge, deps.json is not written and the function returns non-zero — this is + * the "no shotgun modifications" guarantee: anyone who changes + * `compute_task_fanin`'s producers or edge flags trips this gate immediately and + * knows to mirror the change in the annot pass. * * STEP 1 (explicit_deps) is emitted at the call site (per pto_dep_compute.h's * "kept at call site" note); both passes run the same explicit-deps loop, so @@ -130,6 +131,21 @@ const char *edge_source_str(EdgeSource s) { return "unknown"; } +// JSON array of the DepFlags bits set on an edge, e.g. ["wait","retain"]. +void write_dep_flags(std::ostream &out, DepFlags flags) { + out << '['; + bool first = true; + if (dep_has_wait(flags)) { + out << "\"wait\""; + first = false; + } + if (dep_has_retain(flags)) { + if (!first) out << ','; + out << "\"retain\""; + } + out << ']'; +} + const char *overlap_status_str(OverlapStatus s) { switch (s) { case OverlapStatus::COVERED: @@ -154,6 +170,7 @@ struct EdgeAnnot { uint64_t succ; int32_t consumer_arg_idx; // -1 for EXPLICIT (not tied to a tensor arg) 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 // Consumer side (the ChipTensor the submitting task is reading). @@ -373,6 +390,8 @@ bool write_deps_json( out << "{\"pred\":\"" << e.pred << "\",\"succ\":\"" << e.succ << '"'; out << ",\"arg\":" << e.consumer_arg_idx; out << ",\"source\":\"" << edge_source_str(e.source) << '"'; + out << ",\"flags\":"; + write_dep_flags(out, e.flags); if (e.source == EdgeSource::TENSORMAP) { out << ",\"overlap\":\"" << overlap_status_str(e.overlap) << '"'; } @@ -516,13 +535,14 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c TensorRef tref_buf[CORE_MAX_TENSOR_ARGS]; TensorArgType atype_buf[CORE_MAX_TENSOR_ARGS]; - // Per-record dedup of producer IDs — must match runtime's + // Per-record producer ID -> accumulated DepFlags — must match runtime's // PTO2FaninBuilder::append_fanin_or_fail semantics, which collapses STEP 1 // (explicit_deps) + STEP A (creator retention) + STEP B (tensormap lookup) - // into a single per-task fanin list. Both oracle and annot use this same - // semantics so the divergence check is meaningful. - std::unordered_set oracle_preds; - std::unordered_set annot_preds; + // into a single per-task fanin edge and OR-accumulates its flags. Both oracle + // and annot use this same semantics so the divergence check compares the + // (producer, flags) mapping rather than the producer-ID set alone. + std::unordered_map oracle_preds; + std::unordered_map annot_preds; // Scratch buffer for assembling full dep lists across overflow chains. // Declared outside the loop so it can be reused (clear() keeps capacity). @@ -683,24 +703,29 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c // gathered base+chain buffer on overflow path). for (int32_t i = 0; i < dc; i++) { uint64_t pred_raw = deps_data[i]; - if (oracle_preds.insert(pred_raw).second) { - // First time this pred is seen at runtime call site. - } - if (annot_preds.insert(pred_raw).second) { + // Explicit deps are recorded conservatively as DEP_WAIT|DEP_RETAIN; + // the DepGenRecord does not carry per-dep kinds, matching Arg's + // set_dependencies default. + oracle_preds[pred_raw] |= (DEP_WAIT | DEP_RETAIN); + bool first = annot_preds.find(pred_raw) == annot_preds.end(); + annot_preds[pred_raw] |= (DEP_WAIT | DEP_RETAIN); + if (first) { EdgeAnnot e{}; e.pred = pred_raw; e.succ = rec.task_id; e.consumer_arg_idx = -1; e.source = EdgeSource::EXPLICIT; + e.flags = DEP_WAIT | DEP_RETAIN; annot_edges.push_back(e); } } // ============ ORACLE pass — drive compute_task_fanin ============ - bool ok = compute_task_fanin(inputs, tm_oracle, in_manual_scope, [&](PTO2TaskId producer) -> bool { - oracle_preds.insert(producer.raw); - return true; - }); + bool ok = + compute_task_fanin(inputs, tm_oracle, in_manual_scope, [&](PTO2TaskId producer, DepFlags kind) -> bool { + oracle_preds[producer.raw] |= kind; + return true; + }); if (!ok) { LOG_ERROR("dep_gen replay: compute_task_fanin returned fatal at task_id=%" PRIu64, rec.task_id); tm_oracle.destroy(); @@ -713,7 +738,9 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c inputs, tm_annot, in_manual_scope, // emit_creator(producer, arg_idx, consumer_tensor) [&](PTO2TaskId producer, int32_t arg_idx, const ChipTensor &consumer) { - if (!annot_preds.insert(producer.raw).second) { + bool first = annot_preds.find(producer.raw) == annot_preds.end(); + annot_preds[producer.raw] |= (DEP_WAIT | DEP_RETAIN); + if (!first) { return; // already covered by an earlier emit on this record } EdgeAnnot e{}; @@ -721,6 +748,7 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c e.succ = rec.task_id; e.consumer_arg_idx = arg_idx; e.source = EdgeSource::CREATOR; + e.flags = DEP_WAIT | DEP_RETAIN; e.tensor_id = make_tensor_id(consumer.buffer.addr, consumer.version); fill_consumer(e, consumer); annot_edges.push_back(e); @@ -735,12 +763,13 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c // producers, both yield their own edges. The producer-id-set // comparison below uses annot_preds, which dedups by pred // only, matching runtime PTO2FaninBuilder semantics. - annot_preds.insert(producer.raw); + annot_preds[producer.raw] |= DEP_WAIT; EdgeAnnot e{}; e.pred = producer.raw; e.succ = rec.task_id; e.consumer_arg_idx = arg_idx; e.source = EdgeSource::TENSORMAP; + e.flags = DEP_WAIT; e.overlap = status; e.tensor_id = make_tensor_id(entry.buffer_addr, entry.version); fill_consumer(e, consumer); @@ -755,15 +784,21 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c "dep_gen replay: DIVERGENCE at task_id=%" PRIu64 " (rec_idx=%zu): oracle has %zu preds, annot has %zu", rec.task_id, rec_i, oracle_preds.size(), annot_preds.size() ); - // Log the symmetric difference for debugging. - for (uint64_t p : oracle_preds) { - if (annot_preds.find(p) == annot_preds.end()) { - LOG_ERROR(" only-in-oracle pred: %" PRIu64, p); + // Log the symmetric difference (missing preds and flag mismatches). + for (const auto &[p, f] : oracle_preds) { + auto it = annot_preds.find(p); + if (it == annot_preds.end()) { + LOG_ERROR(" only-in-oracle pred: %" PRIu64 " flags=%u", p, static_cast(f)); + } else if (it->second != f) { + LOG_ERROR( + " flags mismatch pred: %" PRIu64 " oracle=%u annot=%u", p, static_cast(f), + static_cast(it->second) + ); } } - for (uint64_t p : annot_preds) { + for (const auto &[p, f] : annot_preds) { if (oracle_preds.find(p) == oracle_preds.end()) { - LOG_ERROR(" only-in-annot pred: %" PRIu64, p); + LOG_ERROR(" only-in-annot pred: %" PRIu64 " flags=%u", p, static_cast(f)); } } tm_oracle.destroy(); 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 f8fa586ca5..057cfb8cb0 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 @@ -71,26 +71,29 @@ class CoreTaskArgsWithDeps : private CoreTaskArgs { // the convenience layer reach dependencies only through add_dep() below. /** - * Append one or more dependencies to the bundled buffer. May be called - * multiple times; deps accumulate. Variadic accepts any non-zero number - * of PTO2TaskId arguments. + * Append one or more RETAIN dependencies to the bundled buffer: the producer + * is kept alive (its slot/output buffer retained) until this consumer + * completes. This is the conservative default — use it when the consumer reads + * a tensor whose buffer the producer allocated. May be called multiple times; + * deps accumulate. Variadic accepts any non-zero number of PTO2TaskId args. * * Overflow (more than MAX_DEP_COUNT total) records an error on the * underlying Arg; the error surfaces at submit time. */ template void add_dep(Ids... ids) { - static_assert(sizeof...(Ids) >= 1, "add_dep: at least one task id is required"); - static_assert( - (std::is_same_v, PTO2TaskId> && ...), "add_dep: all arguments must be PTO2TaskId" - ); - if (count_ + sizeof...(Ids) > MAX_DEP_COUNT) { - CoreTaskArgs::set_error( - "CoreTaskArgsWithDeps::add_dep: dep count exceeds MAX_DEP_COUNT (bump the template arg)" - ); - return; - } - ((deps_[count_++] = ids), ...); + add_dep_impl(DEP_WAIT | DEP_RETAIN, ids...); + } + + /** + * Append one or more ordering-only (DEP_WAIT) dependencies: the producer is + * NOT retained and may be reclaimed as soon as it completes and notifies this + * consumer. Use this only when the consumer merely orders after the producer + * and does not read a buffer the producer allocated. + */ + template + void add_dep_wait(Ids... ids) { + add_dep_impl(DEP_WAIT, ids...); } /** @@ -113,12 +116,29 @@ class CoreTaskArgsWithDeps : private CoreTaskArgs { */ CoreTaskArgs &finalize_for_submit() { CoreTaskArgs::set_dependencies(nullptr, 0); - CoreTaskArgs::set_dependencies(deps_, count_); + CoreTaskArgs::set_dependencies_with_kinds(deps_, kinds_, count_); return *this; } private: + template + void add_dep_impl(DepFlags kind, Ids... ids) { + static_assert(sizeof...(Ids) >= 1, "add_dep/add_dep_wait: at least one task id is required"); + static_assert( + (std::is_same_v, PTO2TaskId> && ...), + "add_dep/add_dep_wait: all arguments must be PTO2TaskId" + ); + if (count_ + sizeof...(Ids) > MAX_DEP_COUNT) { + CoreTaskArgs::set_error( + "CoreTaskArgsWithDeps::add_dep/add_dep_wait: dep count exceeds MAX_DEP_COUNT (bump the template arg)" + ); + return; + } + ((kinds_[count_] = kind, deps_[count_] = ids, ++count_), ...); + } + PTO2TaskId deps_[MAX_DEP_COUNT]; + DepFlags kinds_[MAX_DEP_COUNT]; uint32_t count_ = 0; }; 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 ee1ee7b753..16a92380a7 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 @@ -29,7 +29,11 @@ * the minor structural overlap. Replay handles STEP 1 with a one-line loop of its own. * * The Emit callback contract: - * bool emit(PTO2TaskId producer); + * bool emit(PTO2TaskId producer, DepFlags kind); + * - kind is DEP_WAIT|DEP_RETAIN for a Step-A creator edge (the consumer reads + * the producer's allocated buffer, so the producer is retained) and DEP_WAIT + * for a Step-B modifier edge (ordering only; the buffer was allocated + * elsewhere). Duplicate producers OR-accumulate their flags. * - return true to continue (whether or not the producer was actually recorded — * producer-not-alive / dedup-hit / etc. all return true silently) * - return false to signal fatal (e.g. fanin spill overflow); caller bails @@ -93,10 +97,11 @@ compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_m const ChipTensor *tensor = &inputs.tensors[i].ref(); - // Step A: creator retention — all existing tensors extend their creator lifetime. + // Step A: creator retention — reading a tensor retains its allocator, so + // the creator edge carries both ordering and lifetime. PTO2TaskId owner = tensor->owner_task_id; if (owner.is_valid()) { - if (!emit(owner)) { + if (!emit(owner, DEP_WAIT | DEP_RETAIN)) { return false; } } @@ -111,7 +116,17 @@ 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 { - if (!emit(entry.producer_task_id)) { + // 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)) { fatal = true; return false; // stop iteration } 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 f8f9040547..8983c895b9 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp @@ -231,16 +231,18 @@ static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { struct PTO2FaninBuilder { PTO2FaninBuilder(PTO2OrchestratorState *orch, PTO2FaninPool &spill_pool, uint32_t seen_epoch) : count(0), + wait_count(0), spill_start(0), orch(orch), seen_epoch(seen_epoch), spill_pool(spill_pool) {} - int32_t count{0}; + int32_t count{0}; // total fanin edges (all flag combinations) + int32_t wait_count{0}; // edges carrying DEP_WAIT — sizes readiness accounting int32_t spill_start{0}; PTO2OrchestratorState *orch{nullptr}; uint32_t seen_epoch{0}; PTO2FaninPool &spill_pool; - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP]; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP]; template PTO2FaninForEachReturn for_each(Fn &&fn) const { @@ -259,11 +261,57 @@ struct PTO2FaninBuilder { seen[slot] = seen_epoch; return false; } + + // Append a new edge (caller has already claimed the producer's fanout pin). + void push_edge(PTO2FaninSpillEntry &entry, PTO2TaskSlotState *prod_state, DepFlags kind) { + entry.set(prod_state, kind); + count++; + if (dep_has_wait(kind)) { + wait_count++; + } + } + + // Dedup path: a producer already recorded this submission is reached again + // (e.g. as both a creator and a modifier). OR the new flags into the existing + // edge instead of adding a second one — no extra fanout pin is claimed. + void or_flags_into_existing(PTO2TaskSlotState *prod_state, DepFlags kind) { + for (int32_t i = 0; i < count; i++) { + PTO2FaninSpillEntry &entry = entry_at(i); + if (entry.slot_state() == prod_state) { + accumulate_flags(entry, kind); + return; + } + } + // mark_seen reported this producer already owns an edge this submission, + // so it must be present. A miss means the seen-set and the builder + // disagree — a logic error, not a runtime condition. + always_assert(false && "or_flags_into_existing: deduped producer missing from the fanin builder"); + } + +private: + // The i-th appended fanin edge: inline for i < PTO2_FANIN_INLINE_CAP, else in + // the spill ring at the same linear->physical position for_each_fanin_storage + // walks. Single source of the wrap arithmetic. + PTO2FaninSpillEntry &entry_at(int32_t i) { + if (i < PTO2_FANIN_INLINE_CAP) { + return inline_slots[i]; + } + int32_t spill_idx = i - PTO2_FANIN_INLINE_CAP; + return spill_pool.base[(spill_start % spill_pool.capacity + spill_idx) % spill_pool.capacity]; + } + + void accumulate_flags(PTO2FaninSpillEntry &entry, DepFlags kind) { + bool had_wait = dep_has_wait(entry.flags()); + entry.add_flags(kind); + if (!had_wait && dep_has_wait(kind)) { + wait_count++; + } + } }; static bool append_fanin_or_fail( PTO2OrchestratorState *orch, uint8_t prod_ring, int32_t prod_slot, PTO2TaskSlotState *prod_state, - PTO2TaskId producer_task_id, PTO2FaninBuilder *fanin_builder, uint8_t ring_id + PTO2TaskId producer_task_id, PTO2FaninBuilder *fanin_builder, uint8_t ring_id, DepFlags kind ) { // Decide-and-claim under the producer's fanout_lock. Two conditions make this // resolved slot a non-dependency, and both must be checked together with the @@ -275,12 +323,13 @@ static bool append_fanin_or_fail( // producer; ++'ing it would corrupt an unrelated task. // (2) Already CONSUMED in place — finished, output ready, no real edge. // In either case, adding it to the fanin and bumping fanout_count would leave - // a stale ++/release pair (Orch-side wiring drops the fanout edge but keeps - // the fanin slot, so on_task_release still release_producer()'s it) that - // desyncs the slot's refcount (rc != fc) and wedges in-order reclaim. Claiming a live - // producer under the lock pins it: fanout_count now counts us, so it cannot - // reach CONSUMED (rc == fc) until we release it in on_task_release, keeping the - // slot's generation stable until then. check_and_handle_consumed flips + // a stale ++/release pair that desyncs the slot's refcount (rc != fc) and + // wedges in-order reclaim. Every edge (regardless of DepFlags) claims one + // fanout_count++ here: it pins the producer's slot across the submit->wire + // window so it cannot be CONSUMED + reused before wire_fanin_task links the + // consumer. The pin's release differs by kind — an ordering-only (DEP_WAIT + // without DEP_RETAIN) edge releases it at wiring; a DEP_RETAIN edge holds it + // until the consumer's on_task_release. check_and_handle_consumed flips // COMPLETED->CONSUMED under the same lock, so the check and the ++ are atomic // against the consume. fanout_count is lock-protected per the // PTO2TaskSlotState contract. @@ -289,13 +338,15 @@ static bool append_fanin_or_fail( // gone check. mark_seen keys only on (ring, slot); a stale owner that resolves // to a reused slot must not record it as seen, or a later dependency on the // live generation in the same submission would hit mark_seen and be skipped - // without claiming it (dropped edge). Marking only when !gone keeps the dedup - // keyed to the live producer, and doing it before the ++ still suppresses a - // double-count for a producer named twice in one submission. + // without claiming it (dropped edge). A duplicate live producer claims no new + // pin and adds no new edge; its flags are OR-accumulated into the existing + // edge so the stronger of several discovery reasons (creator vs modifier vs + // explicit) always wins, independent of the order they are reached in. prod_state->lock_fanout(); bool gone = prod_state->task == nullptr || prod_state->task->task_id.local() != producer_task_id.local() || prod_state->task_state.load(std::memory_order_acquire) == PTO2_TASK_CONSUMED; - bool claim = !gone && !fanin_builder->mark_seen(prod_ring, prod_slot); + bool already_seen = !gone && fanin_builder->mark_seen(prod_ring, prod_slot); + bool claim = !gone && !already_seen; int32_t fanout_now = -1; if (claim) { // Low bits hold the consumer count; bit31 is the scope ref. The consumer @@ -322,14 +373,18 @@ static bool append_fanin_or_fail( static_cast(producer_task_id.ring()), producer_task_id.local(), PTO2_DEP_DEGREE_DEBUG_THRESHOLD ); } - // gone (stale/consumed) or an already-seen duplicate live producer: no new - // fanin edge either way. - if (!claim) { + // Stale/consumed producer: no edge at all. + if (gone) { + return true; + } + // Duplicate live producer: fold the flags into the edge already recorded. + if (already_seen) { + fanin_builder->or_flags_into_existing(prod_state, kind); return true; } if (fanin_builder->count < PTO2_FANIN_INLINE_CAP) { - fanin_builder->inline_slots[fanin_builder->count++] = prod_state; + fanin_builder->push_edge(fanin_builder->inline_slots[fanin_builder->count], prod_state, kind); return true; } @@ -347,21 +402,24 @@ static bool append_fanin_or_fail( if (fanin_builder->count == PTO2_FANIN_INLINE_CAP) { fanin_builder->spill_start = spill_idx; } - entry->slot_state = prod_state; - fanin_builder->count++; + fanin_builder->push_edge(*entry, prod_state, kind); return true; } static bool all_claimed_fanin_completed(const PTO2FaninBuilder &fanin_builder) { if (fanin_builder.count == 0) return true; - return fanin_builder.for_each([](PTO2TaskSlotState *producer) -> bool { + // Only DEP_WAIT edges gate readiness; a retention-only edge never blocks + // dispatch, so it is treated as satisfied here. + return fanin_builder.for_each([](PTO2TaskSlotState *producer, DepFlags flags) -> bool { + if (!dep_has_wait(flags)) return true; return producer != nullptr && producer->task_state.load(std::memory_order_acquire) >= PTO2_TASK_COMPLETED; }); } static bool all_claimed_fanin_allow_early_resolve(const PTO2FaninBuilder &fanin_builder) { if (fanin_builder.count == 0) return true; - return fanin_builder.for_each([](PTO2TaskSlotState *producer) -> bool { + return fanin_builder.for_each([](PTO2TaskSlotState *producer, DepFlags flags) -> bool { + if (!dep_has_wait(flags)) return true; return producer != nullptr && producer->task_attrs.allow_early_resolve(); }); } @@ -387,23 +445,40 @@ void PTO2OrchestratorState::wire_fanin_task(PTO2TaskSlotState &slot_state, int32 int32_t early_propagated = 0; // Direct-only (#1285): one unflagged producer => consumer never early-dispatches. bool early_disqualified = false; - for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer) { - producer->lock_fanout(); - int32_t pstate = producer->task_state.load(std::memory_order_acquire); - if (!early_disqualified && !producer->task_attrs.allow_early_resolve()) { - early_disqualified = true; - } - if (pstate >= PTO2_TASK_COMPLETED) { - completed_fanin++; - } else { - producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, &slot_state); - // Marker shares fanout_lock with propagation's snapshot. A set marker - // means this edge is outside that snapshot and needs a seed. - if (!early_disqualified && producer->has_dispatch_propagated()) { - early_propagated++; + for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer, DepFlags flags) { + // Only DEP_WAIT edges contribute to readiness: they gate fanin and are + // linked onto the producer's fanout_head for completion notification. A + // retention-only edge carries no ordering, so it is skipped here. + if (dep_has_wait(flags)) { + producer->lock_fanout(); + int32_t pstate = producer->task_state.load(std::memory_order_acquire); + if (!early_disqualified && !producer->task_attrs.allow_early_resolve()) { + early_disqualified = true; } + if (pstate >= PTO2_TASK_COMPLETED) { + completed_fanin++; + } else { + producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, &slot_state); + // Marker shares fanout_lock with propagation's snapshot. A set marker + // means this edge is outside that snapshot and needs a seed. + if (!early_disqualified && producer->has_dispatch_propagated()) { + early_propagated++; + } + } + producer->unlock_fanout(); + } + // The submit->wire pin protects an ordering-only edge only until the + // consumer is linked. With the consumer now on fanout_head (or already + // seen as completed), release it so the producer can be CONSUMED without + // waiting for this consumer. A DEP_RETAIN edge keeps the pin until the + // consumer's on_task_release. + if (dep_has_wait(flags) && !dep_has_retain(flags)) { + // Wiring-phase atomics (this release, plus the lock_fanout / dep_pool.prepend / + // fanin_refcount ops around it) are not bucketed: g_orch_args_atomic_count + // covers the submit/dep-claim phase only, whose g_orch_args_cycle window has + // already closed by the time wiring runs. + sched->release_producer(*producer); } - producer->unlock_fanout(); }); // Seed only when every direct producer is codegen-flagged; one unflagged @@ -886,7 +961,8 @@ static TaskOutputTensors submit_task_common( int32_t dep_slot = dep_ring.get_slot_by_task_id(dep_local_task_id); PTO2TaskSlotState *producer_slot_state = &dep_ring.get_slot_state_by_slot(dep_slot); if (!append_fanin_or_fail( - orch, dep_ring_id, dep_slot, producer_slot_state, dep_task_id, &fanin_builder, ring_id + orch, dep_ring_id, dep_slot, producer_slot_state, dep_task_id, &fanin_builder, ring_id, + args.explicit_dep_kind(i) )) { return result; } @@ -898,12 +974,14 @@ static TaskOutputTensors submit_task_common( args.explicit_deps_data(), }; - auto runtime_emit = [&](PTO2TaskId producer_task_id) -> bool { + auto runtime_emit = [&](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); - return append_fanin_or_fail(orch, prod_ring, prod_slot, prod_state, producer_task_id, &fanin_builder, ring_id); + 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)) { @@ -942,6 +1020,16 @@ static TaskOutputTensors submit_task_common( // here) is what prevents a producer from transitioning to CONSUMED between // the dependency decision and the claim. int32_t inline_count = std::min(fanin_builder.count, PTO2_FANIN_INLINE_CAP); + // Every fanin edge produced here carries DEP_WAIT (creator = WAIT|RETAIN, + // modifier = WAIT, explicit defaults to WAIT|RETAIN or opts into WAIT), so + // wait_count == count. fanin_actual_count therefore doubles as the WAIT-edge + // count that the early-dispatch threshold (dispatch_fanin, which counts only + // WAIT producers) is compared against. A future RETAIN-only edge would break + // that equality and must carry its own WAIT-edge count for that comparison. + always_assert( + fanin_builder.wait_count == fanin_builder.count && + "fanin_actual_count is the early-dispatch WAIT denominator; a non-WAIT edge needs a separate count" + ); // Store fanin metadata in payload for scheduler to iterate payload.fanin_actual_count = fanin_builder.count; // fanin_builder.count is finalized here and submit runs once per task, so @@ -956,7 +1044,7 @@ static TaskOutputTensors submit_task_common( payload.fanin_spill_start = fanin_builder.spill_start; payload.fanin_spill_pool = &fanin_builder.spill_pool; for (int i = 0; i < inline_count; i++) { - payload.fanin_inline_slot_states[i] = fanin_builder.inline_slots[i]; + payload.fanin_inline_edges[i] = fanin_builder.inline_slots[i]; } payload.init(args, result, prepared.alloc_result, layout); @@ -1007,16 +1095,24 @@ static TaskOutputTensors submit_task_common( orch->mark_dep_pool_position(cur_slot_state); sched->push_ready_routed(&cur_slot_state); } else if (all_claimed_fanin_completed(fanin_builder)) { - int32_t ready_seed = fanin_builder.count + 1; + int32_t ready_seed = fanin_builder.wait_count + 1; cur_slot_state.fanin_count = ready_seed; if (all_claimed_fanin_allow_early_resolve(fanin_builder)) { payload.dispatch_fanin.store(fanin_builder.count, std::memory_order_release); } cur_slot_state.fanin_refcount.store(ready_seed, std::memory_order_release); + // wire_fanin_task is skipped here, so its ordering-only pin release runs + // on this path too: an edge without retention drops its submit->wire pin + // so the (already completed) producer can be CONSUMED. + for_each_fanin_slot_state(payload, [&](PTO2TaskSlotState *producer, DepFlags flags) { + if (dep_has_wait(flags) && !dep_has_retain(flags)) { + sched->release_producer(*producer); // wiring-phase atomic, not bucketed (see wire_fanin_task) + } + }); orch->mark_dep_pool_position(cur_slot_state); sched->push_ready_routed(&cur_slot_state); } else { - if (!orch_wire_live_fanin_task(orch, cur_slot_state, fanin_builder.count)) { + if (!orch_wire_live_fanin_task(orch, cur_slot_state, fanin_builder.wait_count)) { return result; } } diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h index eb35e7cfac..b376e99ab2 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h @@ -519,7 +519,7 @@ struct PTO2FaninPool { tail = 1; high_water = 0; reclaim_task_cursor = 0; - base[0].slot_state = nullptr; + base[0].clear(); error_code_ptr = in_error_code_ptr; } @@ -528,7 +528,7 @@ struct PTO2FaninPool { tail = 1; high_water = 0; reclaim_task_cursor = 0; - base[0].slot_state = nullptr; + base[0].clear(); error_code_ptr = in_error_code_ptr; } @@ -575,14 +575,16 @@ struct PTO2FaninPool { }; template -using PTO2FaninCallbackResult = std::invoke_result_t; +using PTO2FaninCallbackResult = std::invoke_result_t; template using PTO2FaninForEachReturn = std::conditional_t, void>, void, bool>; +// Visit each fanin edge as (producer slot, DepFlags). Inline and spill entries +// share the packed PTO2FaninSpillEntry layout, so both are unpacked the same way. template inline PTO2FaninForEachReturn for_each_fanin_storage( - InlineSlots &&inline_slot_states, int32_t fanin_count, int32_t spill_start, PTO2FaninPool &spill_pool, Fn &&fn + InlineSlots &&inline_edges, int32_t fanin_count, int32_t spill_start, PTO2FaninPool &spill_pool, Fn &&fn ) { using FaninCallbackResult = PTO2FaninCallbackResult; static_assert( @@ -593,7 +595,7 @@ inline PTO2FaninForEachReturn for_each_fanin_storage( if constexpr (std::is_void_v) { int32_t inline_count = std::min(fanin_count, PTO2_FANIN_INLINE_CAP); for (int32_t i = 0; i < inline_count; i++) { - fn(inline_slot_states[i]); + fn(inline_edges[i].slot_state(), inline_edges[i].flags()); } int32_t spill_count = fanin_count - inline_count; @@ -605,18 +607,18 @@ inline PTO2FaninForEachReturn for_each_fanin_storage( int32_t first_count = std::min(spill_count, spill_pool.capacity - start_idx); PTO2FaninSpillEntry *first = spill_pool.base + start_idx; for (int32_t i = 0; i < first_count; i++) { - fn(first[i].slot_state); + fn(first[i].slot_state(), first[i].flags()); } int32_t second_count = spill_count - first_count; for (int32_t i = 0; i < second_count; i++) { - fn(spill_pool.base[i].slot_state); + fn(spill_pool.base[i].slot_state(), spill_pool.base[i].flags()); } return; } else { int32_t inline_count = std::min(fanin_count, PTO2_FANIN_INLINE_CAP); for (int32_t i = 0; i < inline_count; i++) { - if (!fn(inline_slot_states[i])) { + if (!fn(inline_edges[i].slot_state(), inline_edges[i].flags())) { return false; } } @@ -630,14 +632,14 @@ inline PTO2FaninForEachReturn for_each_fanin_storage( int32_t first_count = std::min(spill_count, spill_pool.capacity - start_idx); PTO2FaninSpillEntry *first = spill_pool.base + start_idx; for (int32_t i = 0; i < first_count; i++) { - if (!fn(first[i].slot_state)) { + if (!fn(first[i].slot_state(), first[i].flags())) { return false; } } int32_t second_count = spill_count - first_count; for (int32_t i = 0; i < second_count; i++) { - if (!fn(spill_pool.base[i].slot_state)) { + if (!fn(spill_pool.base[i].slot_state(), spill_pool.base[i].flags())) { return false; } } @@ -648,8 +650,8 @@ inline PTO2FaninForEachReturn for_each_fanin_storage( template inline PTO2FaninForEachReturn for_each_fanin_slot_state(const PTO2TaskPayload &payload, Fn &&fn) { return for_each_fanin_storage( - payload.fanin_inline_slot_states, payload.fanin_actual_count, payload.fanin_spill_start, - *payload.fanin_spill_pool, static_cast(fn) + payload.fanin_inline_edges, payload.fanin_actual_count, payload.fanin_spill_start, *payload.fanin_spill_pool, + static_cast(fn) ); } 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 c334ba50de..5898b9c069 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 @@ -162,8 +162,28 @@ struct PTO2OutputLayout { */ struct PTO2TaskSlotState; // Forward declaration struct PTO2FaninPool; // Forward declaration + +// One fanin edge: a producer slot pointer with the per-edge DepFlags packed into +// bits 0..1 of the pointer. PTO2TaskSlotState is alignas(64), so bits 0..5 are +// always zero in a real pointer. Used for both the inline fanin array and the +// spill pool, keeping sizeof == sizeof(uintptr_t) so neither footprint grows. struct PTO2FaninSpillEntry { - PTO2TaskSlotState *slot_state; + static constexpr uintptr_t FLAG_MASK = 0x3; + // No in-class initializer: the type stays trivially default-constructible so + // the payload's fanin_inline_edges[64] and the orchestrator's per-submit + // builder array are default-initialized at zero cost (only entries [0, count) + // are written via set()). Value-init (`PTO2FaninSpillEntry{}`) still zeroes it. + uintptr_t packed; + + PTO2TaskSlotState *slot_state() const { return reinterpret_cast(packed & ~FLAG_MASK); } + DepFlags flags() const { return static_cast(packed & FLAG_MASK); } + // Only bits within FLAG_MASK are stored; any bit outside it (a malformed + // DepFlags value) is masked off so it can never corrupt the slot pointer. + void set(PTO2TaskSlotState *s, DepFlags f) { + packed = reinterpret_cast(s) | (static_cast(f) & FLAG_MASK); + } + void add_flags(DepFlags f) { packed |= (static_cast(f) & FLAG_MASK); } + void clear() { packed = 0; } }; static_assert(sizeof(PTO2FaninSpillEntry) == sizeof(uintptr_t)); @@ -255,7 +275,10 @@ struct PTO2TaskPayload { int32_t fanin_actual_count{0}; // Actual fanin count (without the +1 redundance) int32_t fanin_spill_start{0}; // Linear start index in fanin spill pool (0 = no spill) PTO2FaninPool *fanin_spill_pool{nullptr}; - PTO2TaskSlotState *fanin_inline_slot_states[PTO2_FANIN_INLINE_CAP]; + // Inline fanin edges (producer slot + packed DepFlags). Spill beyond + // PTO2_FANIN_INLINE_CAP goes to fanin_spill_pool. Same packed layout as the + // spill entries, so the array footprint is unchanged. + PTO2FaninSpillEntry fanin_inline_edges[PTO2_FANIN_INLINE_CAP]; // Early-dispatch metadata (AICPU-side only). Ordered by descending // alignment so the block packs without internal padding. Cache line 8 // contains the rarely-touched fanin tail rather than the hot tensor/scalar @@ -378,9 +401,7 @@ struct PTO2TaskPayload { // PTO2TaskPayload layout verification (offsetof requires complete type). static_assert(offsetof(PTO2TaskPayload, fanin_spill_pool) == 16, "spill pool pointer layout drift"); -static_assert( - offsetof(PTO2TaskPayload, fanin_inline_slot_states) == 24, "inline fanin array must follow spill metadata" -); +static_assert(offsetof(PTO2TaskPayload, fanin_inline_edges) == 24, "inline fanin array must follow spill metadata"); static_assert( offsetof(PTO2TaskPayload, predicate) == 576, "dispatch predicate occupies cache line 9 at fixed byte 576 (before tensors, never moves)" @@ -631,3 +652,9 @@ struct alignas(64) PTO2TaskSlotState { }; static_assert(sizeof(PTO2TaskSlotState) == 64); +// PTO2FaninSpillEntry packs DepFlags into the low bits of a PTO2TaskSlotState*. +// That is only lossless while the type is aligned past those tag bits. +static_assert( + alignof(PTO2TaskSlotState) > PTO2FaninSpillEntry::FLAG_MASK, + "PTO2TaskSlotState alignment must exceed PTO2FaninSpillEntry::FLAG_MASK so the packed DepFlags bits are free" +); 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 fcb727ea2e..57f239b242 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h @@ -71,6 +71,42 @@ enum class PTO2ScopeMode : uint8_t { MANUAL = 1, }; +/** + * Orthogonal dependency-edge semantics, stored per fanin edge. + * + * DEP_WAIT — readiness/ordering: the consumer cannot become ready until the + * producer completes. Contributes to the consumer's fanin count + * and to the producer's fanout notification list. + * DEP_RETAIN — lifetime: the producer's slot and packed output buffer must + * stay alive until the consumer releases it. Holds a + * fanout_count reference from submit until on_task_release. + * + * The two are independent; an edge may carry either, both, or (transiently + * during construction) neither. A creator edge is DEP_WAIT|DEP_RETAIN (the + * consumer reads the producer's allocated buffer). A tensormap modifier edge is + * DEP_WAIT only (the buffer was allocated elsewhere, so only ordering is + * needed). When one (producer, consumer) pair is discovered for several + * reasons, the flags are OR-accumulated. + */ +enum DepFlags : uint8_t { + DEP_NONE = 0, + DEP_WAIT = 1u << 0, + DEP_RETAIN = 1u << 1, +}; + +constexpr DepFlags operator|(DepFlags a, DepFlags b) { + return static_cast(static_cast(a) | static_cast(b)); +} +constexpr DepFlags operator&(DepFlags a, DepFlags b) { + return static_cast(static_cast(a) & static_cast(b)); +} +constexpr DepFlags &operator|=(DepFlags &a, DepFlags b) { + a = a | b; + return a; +} +constexpr bool dep_has_wait(DepFlags f) { return (f & DEP_WAIT) != DEP_NONE; } +constexpr bool dep_has_retain(DepFlags f) { return (f & DEP_RETAIN) != DEP_NONE; } + /** * TaskOutputTensors — returned by submit, holds materialized output ChipTensors. * @@ -287,6 +323,7 @@ struct Arg : TaskArgsTpl { #endif explicit_deps_ = nullptr; explicit_dep_count_ = 0; + explicit_dep_kinds_ = nullptr; allow_early_resolve_ = false; predicate_ = CoreTaskPredicate{}; task_timing_slot_ = TASK_TIMING_SLOT_NONE; @@ -397,6 +434,7 @@ struct Arg : TaskArgsTpl { if (count == 0) { explicit_deps_ = nullptr; explicit_dep_count_ = 0; + explicit_dep_kinds_ = nullptr; return; } if (deps == nullptr) { @@ -409,6 +447,35 @@ struct Arg : TaskArgsTpl { } explicit_deps_ = deps; explicit_dep_count_ = count; + explicit_dep_kinds_ = nullptr; + } + + /** + * Same as set_dependencies() but attaches a per-dep DepFlags array so the + * caller can mark an explicit dep as ordering-only (DEP_WAIT) rather than the + * conservative DEP_WAIT|DEP_RETAIN default. The kinds array is caller-owned + * and must outlive submit (same lifetime rule as deps). A null kinds array + * with count > 0 is rejected; use set_dependencies() for the + * conservative DEP_WAIT|DEP_RETAIN default. + */ + void set_dependencies_with_kinds(const PTO2TaskId *deps, const DepFlags *kinds, uint32_t count) { + if (count == 0) { + explicit_deps_ = nullptr; + explicit_dep_count_ = 0; + explicit_dep_kinds_ = nullptr; + return; + } + if (deps == nullptr || kinds == nullptr) { + set_error("set_dependencies_with_kinds: deps and kinds must not be null when count > 0"); + return; + } + if (explicit_deps_ != nullptr) { + set_error("set_dependencies_with_kinds: may be called at most once per Arg"); + return; + } + explicit_deps_ = deps; + explicit_dep_kinds_ = kinds; + explicit_dep_count_ = count; } uint32_t explicit_dep_count() const { return explicit_dep_count_; } @@ -418,6 +485,16 @@ struct Arg : TaskArgsTpl { return explicit_deps_[index]; } + /** + * Flags of the i-th explicit dep. Returns DEP_WAIT|DEP_RETAIN when no kinds + * array is attached (plain set_dependencies), so an explicit dep protecting a + * runtime-created output is never silently weakened to ordering-only. + */ + DepFlags explicit_dep_kind(uint32_t index) const { + always_assert(index < explicit_dep_count_); + return explicit_dep_kinds_ ? explicit_dep_kinds_[index] : (DEP_WAIT | DEP_RETAIN); + } + const PTO2TaskId *explicit_deps_data() const { return explicit_deps_; } /** @@ -517,6 +594,7 @@ struct Arg : TaskArgsTpl { DumpArgSelection dump_arg_selection_; #endif const PTO2TaskId *explicit_deps_{nullptr}; + const DepFlags *explicit_dep_kinds_{nullptr}; uint32_t explicit_dep_count_{0}; #if SIMPLER_DFX template diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index 73d9fe03e3..f6deaac7fe 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -1170,7 +1170,8 @@ struct PTO2SchedulerState { /** * Cold path: release producers (fanin traversal) + check self for CONSUMED. - * Returns fanin edge count for profiling. + * Returns the number of retained (DEP_RETAIN) producers actually released — + * ordering-only edges dropped their pin at wiring and are skipped here. */ #if SIMPLER_SCHED_PROFILING @@ -1185,7 +1186,16 @@ struct PTO2SchedulerState { int32_t on_task_release(PTO2TaskSlotState &slot_state) { #endif PTO2TaskPayload *payload = slot_state.payload; - for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer_slot_state) { + int32_t released = 0; + // Only DEP_RETAIN edges still hold a fanout pin at completion: an + // ordering-only edge released its submit->wire pin at wiring, so releasing + // it again here would over-count fanout_refcount against fanout_count and + // break the rc == fc consume invariant. + for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer_slot_state, DepFlags flags) { + if (!dep_has_retain(flags)) { + return; + } + released++; #if SIMPLER_SCHED_PROFILING release_producer(*producer_slot_state, fanin_atomics); #else @@ -1207,7 +1217,7 @@ struct PTO2SchedulerState { #else check_and_handle_consumed(slot_state); #endif - return payload->fanin_actual_count; + return released; } // === Cold-path API (defined in pto_scheduler.cpp) === diff --git a/tests/ut/cpp/a2a3/test_fanin_pool.cpp b/tests/ut/cpp/a2a3/test_fanin_pool.cpp index f901c0d7ad..8dcbdc8e05 100644 --- a/tests/ut/cpp/a2a3/test_fanin_pool.cpp +++ b/tests/ut/cpp/a2a3/test_fanin_pool.cpp @@ -41,7 +41,7 @@ class FaninPoolTest : public ::testing::Test { PTO2FaninPool pool{}; void SetUp() override { - entries.assign(POOL_CAP, PTO2FaninSpillEntry{nullptr}); + entries.assign(POOL_CAP, PTO2FaninSpillEntry{}); error_code.store(PTO2_ERROR_NONE); pool.init(entries.data(), POOL_CAP, &error_code); } @@ -165,7 +165,7 @@ class ForEachFaninTest : public ::testing::Test { alignas(64) PTO2TaskSlotState slots[64]; void SetUp() override { - spill_entries.assign(POOL_CAP, PTO2FaninSpillEntry{nullptr}); + spill_entries.assign(POOL_CAP, PTO2FaninSpillEntry{}); error_code.store(PTO2_ERROR_NONE); spill_pool.init(spill_entries.data(), POOL_CAP, &error_code); memset(slots, 0, sizeof(slots)); @@ -173,13 +173,13 @@ class ForEachFaninTest : public ::testing::Test { }; TEST_F(ForEachFaninTest, InlineOnlyVoid) { - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < 5; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } std::vector visited; - for_each_fanin_storage(inline_slots, 5, 0, spill_pool, [&](PTO2TaskSlotState *s) { + for_each_fanin_storage(inline_slots, 5, 0, spill_pool, [&](PTO2TaskSlotState *s, DepFlags) { visited.push_back(s); }); @@ -190,13 +190,13 @@ TEST_F(ForEachFaninTest, InlineOnlyVoid) { } TEST_F(ForEachFaninTest, InlineOnlyBoolEarlyReturn) { - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < 5; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } int count = 0; - bool result = for_each_fanin_storage(inline_slots, 5, 0, spill_pool, [&](PTO2TaskSlotState *) -> bool { + bool result = for_each_fanin_storage(inline_slots, 5, 0, spill_pool, [&](PTO2TaskSlotState *, DepFlags) -> bool { count++; return count < 3; // stop after 3rd }); @@ -206,12 +206,12 @@ TEST_F(ForEachFaninTest, InlineOnlyBoolEarlyReturn) { } TEST_F(ForEachFaninTest, InlineOnlyBoolAllTrue) { - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < 3; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } - bool result = for_each_fanin_storage(inline_slots, 3, 0, spill_pool, [](PTO2TaskSlotState *) -> bool { + bool result = for_each_fanin_storage(inline_slots, 3, 0, spill_pool, [](PTO2TaskSlotState *, DepFlags) -> bool { return true; }); @@ -219,9 +219,9 @@ TEST_F(ForEachFaninTest, InlineOnlyBoolAllTrue) { } TEST_F(ForEachFaninTest, ZeroFanin) { - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; int count = 0; - for_each_fanin_storage(inline_slots, 0, 0, spill_pool, [&](PTO2TaskSlotState *) { + for_each_fanin_storage(inline_slots, 0, 0, spill_pool, [&](PTO2TaskSlotState *, DepFlags) { count++; }); EXPECT_EQ(count, 0); @@ -233,20 +233,20 @@ TEST_F(ForEachFaninTest, ZeroFanin) { TEST_F(ForEachFaninTest, SpillNoWrap) { // 18 fanins = 16 inline + 2 spill - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < PTO2_FANIN_INLINE_CAP; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } // Allocate 2 spill entries auto *s0 = spill_pool.alloc(); int32_t spill_start = spill_pool.top - 1; - s0->slot_state = &slots[16]; + s0->set(&slots[16], DEP_WAIT | DEP_RETAIN); auto *s1 = spill_pool.alloc(); - s1->slot_state = &slots[17]; + s1->set(&slots[17], DEP_WAIT | DEP_RETAIN); std::vector visited; - for_each_fanin_storage(inline_slots, 18, spill_start, spill_pool, [&](PTO2TaskSlotState *s) { + for_each_fanin_storage(inline_slots, 18, spill_start, spill_pool, [&](PTO2TaskSlotState *s, DepFlags) { visited.push_back(s); }); @@ -268,9 +268,9 @@ TEST_F(ForEachFaninTest, SpillWithWrap) { spill_pool.top = POOL_CAP - 2; spill_pool.tail = POOL_CAP - 2; - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < PTO2_FANIN_INLINE_CAP; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } // 4 spill entries: indices 30, 31, 0, 1 (wraps around) @@ -278,11 +278,11 @@ TEST_F(ForEachFaninTest, SpillWithWrap) { for (int i = 0; i < 4; i++) { auto *e = spill_pool.alloc(); ASSERT_NE(e, nullptr); - e->slot_state = &slots[16 + i]; + e->set(&slots[16 + i], DEP_WAIT | DEP_RETAIN); } std::vector visited; - for_each_fanin_storage(inline_slots, 20, spill_start, spill_pool, [&](PTO2TaskSlotState *s) { + for_each_fanin_storage(inline_slots, 20, spill_start, spill_pool, [&](PTO2TaskSlotState *s, DepFlags) { visited.push_back(s); }); @@ -302,23 +302,61 @@ TEST_F(ForEachFaninTest, SpillWithWrap) { // ============================================================================= TEST_F(ForEachFaninTest, SpillBoolEarlyReturnInSpillRegion) { - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < PTO2_FANIN_INLINE_CAP; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } int32_t spill_start = spill_pool.top; for (int i = 0; i < 4; i++) { auto *e = spill_pool.alloc(); - e->slot_state = &slots[16 + i]; + e->set(&slots[16 + i], DEP_WAIT); } int count = 0; - bool result = for_each_fanin_storage(inline_slots, 20, spill_start, spill_pool, [&](PTO2TaskSlotState *) -> bool { - count++; - return count < 17; // stop on 17th (first spill entry) - }); + bool result = + for_each_fanin_storage(inline_slots, 20, spill_start, spill_pool, [&](PTO2TaskSlotState *, DepFlags) -> bool { + count++; + return count < 17; // stop on 17th (first spill entry) + }); EXPECT_FALSE(result); EXPECT_EQ(count, 17); } + +// ============================================================================= +// PTO2FaninSpillEntry: DepFlags packing round-trips across inline and spill +// ============================================================================= + +TEST_F(ForEachFaninTest, DepFlagsRoundTripInlineAndSpill) { + // Fill all 64 inline edges; slots 0..2 carry distinct flag combinations. + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + inline_slots[0].set(&slots[0], DEP_WAIT); + inline_slots[1].set(&slots[1], DEP_RETAIN); + inline_slots[2].set(&slots[2], DEP_WAIT | DEP_RETAIN); + for (int i = 3; i < PTO2_FANIN_INLINE_CAP; i++) { + inline_slots[i].set(&slots[i], DEP_WAIT); + } + + // Two edges beyond the inline cap spill; they carry RETAIN-only and WAIT|RETAIN. + auto *s0 = spill_pool.alloc(); + int32_t spill_start = spill_pool.top - 1; + s0->set(&slots[0], DEP_RETAIN); + auto *s1 = spill_pool.alloc(); + s1->set(&slots[1], DEP_WAIT | DEP_RETAIN); + + const int32_t total = PTO2_FANIN_INLINE_CAP + 2; // 64 inline + 2 spill + std::vector flags; + for_each_fanin_storage(inline_slots, total, spill_start, spill_pool, [&](PTO2TaskSlotState *, DepFlags f) { + flags.push_back(f); + }); + + ASSERT_EQ(flags.size(), static_cast(total)); + // Inline flags survive. + EXPECT_EQ(flags[0], DEP_WAIT); + EXPECT_EQ(flags[1], DEP_RETAIN); + EXPECT_EQ(flags[2], DEP_WAIT | DEP_RETAIN); + // Spill flags survive. + EXPECT_EQ(flags[PTO2_FANIN_INLINE_CAP + 0], DEP_RETAIN); + EXPECT_EQ(flags[PTO2_FANIN_INLINE_CAP + 1], DEP_WAIT | DEP_RETAIN); +} diff --git a/tests/ut/cpp/a2a3/test_orchestrator_fanin.cpp b/tests/ut/cpp/a2a3/test_orchestrator_fanin.cpp index 75c4dec5fe..6a5f4f34e2 100644 --- a/tests/ut/cpp/a2a3/test_orchestrator_fanin.cpp +++ b/tests/ut/cpp/a2a3/test_orchestrator_fanin.cpp @@ -87,13 +87,140 @@ TEST_F(OrchestratorFaninTest, DuplicateExplicitProducerAddsOneFanin) { ASSERT_NE(consumer_slot.payload, nullptr); EXPECT_EQ(consumer_slot.payload->fanin_actual_count, 1); - EXPECT_EQ(consumer_slot.payload->fanin_inline_slot_states[0], &producer_slot); + EXPECT_EQ(consumer_slot.payload->fanin_inline_edges[0].slot_state(), &producer_slot); + // A plain set_dependencies() dep is conservative RETAIN: DEP_WAIT|DEP_RETAIN. + EXPECT_EQ(consumer_slot.payload->fanin_inline_edges[0].flags(), DEP_WAIT | DEP_RETAIN); // fanout_count is bit-packed: bit31 (PTO2_FANOUT_SCOPE_BIT) is the owning-scope // ref, low bits the consumer count. The duplicate explicit dep is deduped to a // single consumer, so this is scope + 1. EXPECT_EQ(producer_slot.fanout_count, PTO2_FANOUT_SCOPE_BIT + 1); } +// An explicit ordering-only dep (the primitive add_dep_wait() lowers to) yields a +// DEP_WAIT edge, not the conservative DEP_WAIT|DEP_RETAIN default. +TEST_F(OrchestratorFaninTest, ExplicitWaitDepProducesWaitOnlyEdge) { + orch.begin_scope(); + + CoreTaskArgs producer_args; + TaskOutputTensors producer = orch.submit_dummy_task(producer_args); + ASSERT_TRUE(producer.task_id().is_valid()); + + PTO2TaskId deps[] = {producer.task_id()}; + DepFlags kinds[] = {DEP_WAIT}; + CoreTaskArgs consumer_args; + consumer_args.set_dependencies_with_kinds(deps, kinds, 1); + TaskOutputTensors consumer = orch.submit_dummy_task(consumer_args); + ASSERT_TRUE(consumer.task_id().is_valid()); + + auto &consumer_slot = + sm_handle->header->rings[consumer.task_id().ring()].get_slot_state_by_task_id(consumer.task_id().local()); + ASSERT_NE(consumer_slot.payload, nullptr); + ASSERT_EQ(consumer_slot.payload->fanin_actual_count, 1); + EXPECT_EQ(consumer_slot.payload->fanin_inline_edges[0].flags(), DEP_WAIT); +} + +// The same producer reached with different kinds OR-accumulates into one edge: +// WAIT-only first, then WAIT|RETAIN folds RETAIN in, claiming exactly one pin. +TEST_F(OrchestratorFaninTest, DuplicateProducerOrAccumulatesFlags) { + orch.begin_scope(); + + CoreTaskArgs producer_args; + TaskOutputTensors producer = orch.submit_dummy_task(producer_args); + ASSERT_TRUE(producer.task_id().is_valid()); + + PTO2TaskId deps[] = {producer.task_id(), producer.task_id()}; + DepFlags kinds[] = {DEP_WAIT, DEP_WAIT | DEP_RETAIN}; + CoreTaskArgs consumer_args; + consumer_args.set_dependencies_with_kinds(deps, kinds, 2); + TaskOutputTensors consumer = orch.submit_dummy_task(consumer_args); + ASSERT_TRUE(consumer.task_id().is_valid()); + + auto &producer_slot = + sm_handle->header->rings[producer.task_id().ring()].get_slot_state_by_task_id(producer.task_id().local()); + auto &consumer_slot = + sm_handle->header->rings[consumer.task_id().ring()].get_slot_state_by_task_id(consumer.task_id().local()); + ASSERT_NE(consumer_slot.payload, nullptr); + ASSERT_EQ(consumer_slot.payload->fanin_actual_count, 1); + EXPECT_EQ(consumer_slot.payload->fanin_inline_edges[0].flags(), DEP_WAIT | DEP_RETAIN); + EXPECT_EQ(producer_slot.fanout_count, PTO2_FANOUT_SCOPE_BIT + 1); +} + +// The duplicate lands in the spill region (>64 fanin), exercising +// or_flags_into_existing's spill lookup: the dup folds (65 edges, not 66), claims +// exactly one pin, and OR-accumulates its flags into the spilled edge. +TEST_F(OrchestratorFaninTest, DuplicateProducerInSpillRegionDedups) { + orch.begin_scope(); + + constexpr int kProducers = PTO2_FANIN_INLINE_CAP + 1; // 65: the last one spills + std::vector producers; + producers.reserve(kProducers); + for (int i = 0; i < kProducers; i++) { + CoreTaskArgs a; + producers.push_back(orch.submit_dummy_task(a)); + ASSERT_TRUE(producers.back().task_id().is_valid()); + } + + std::vector deps; + std::vector kinds; + deps.reserve(kProducers + 1); + kinds.reserve(kProducers + 1); + for (auto &p : producers) { + deps.push_back(p.task_id()); + kinds.push_back(DEP_WAIT); // the 65th (first spill edge) starts WAIT-only + } + deps.push_back(producers.back().task_id()); // duplicate the spilled 65th ... + kinds.push_back(DEP_WAIT | DEP_RETAIN); // ... contributing RETAIN via the fold + + CoreTaskArgs consumer_args; + consumer_args.set_dependencies_with_kinds(deps.data(), kinds.data(), static_cast(deps.size())); + TaskOutputTensors consumer = orch.submit_dummy_task(consumer_args); + ASSERT_TRUE(consumer.task_id().is_valid()); + + auto &consumer_slot = + sm_handle->header->rings[consumer.task_id().ring()].get_slot_state_by_task_id(consumer.task_id().local()); + ASSERT_NE(consumer_slot.payload, nullptr); + PTO2TaskPayload *payload = consumer_slot.payload; + EXPECT_EQ(payload->fanin_actual_count, kProducers); // duplicate folded, not 66 + + PTO2TaskId dup = producers.back().task_id(); + auto &dup_slot = sm_handle->header->rings[dup.ring()].get_slot_state_by_task_id(dup.local()); + EXPECT_EQ(dup_slot.fanout_count, PTO2_FANOUT_SCOPE_BIT + 1); // one pin, not two + + // The first spilled edge is the duplicated producer; its flags OR-folded to + // WAIT|RETAIN across the two discovery kinds. + ASSERT_NE(payload->fanin_spill_pool, nullptr); + PTO2FaninPool &spill_pool = *payload->fanin_spill_pool; + PTO2FaninSpillEntry &spill_edge = spill_pool.base[payload->fanin_spill_start % spill_pool.capacity]; + EXPECT_EQ(spill_edge.slot_state(), &dup_slot); + EXPECT_EQ(spill_edge.flags(), DEP_WAIT | DEP_RETAIN); +} + +// The all-completed fast path (wire_fanin_task skipped) still drops an +// ordering-only producer's submit->wire pin. +TEST_F(OrchestratorFaninTest, AllCompletedFastPathReleasesWaitOnlyPin) { + orch.begin_scope(); + + CoreTaskArgs producer_args; + TaskOutputTensors producer = orch.submit_dummy_task(producer_args); + ASSERT_TRUE(producer.task_id().is_valid()); + auto &producer_slot = + sm_handle->header->rings[producer.task_id().ring()].get_slot_state_by_task_id(producer.task_id().local()); + // COMPLETED but not consumed (the open scope still pins it): the consumer takes + // the all-completed fast path. + producer_slot.task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); + int32_t rc_before = producer_slot.fanout_refcount.load(); + + PTO2TaskId deps[] = {producer.task_id()}; + DepFlags kinds[] = {DEP_WAIT}; // ordering-only + CoreTaskArgs consumer_args; + consumer_args.set_dependencies_with_kinds(deps, kinds, 1); + TaskOutputTensors consumer = orch.submit_dummy_task(consumer_args); + ASSERT_TRUE(consumer.task_id().is_valid()); + + // The fast path released the ordering-only pin. + EXPECT_EQ(producer_slot.fanout_refcount.load(), rc_before + 1); +} + TEST_F(OrchestratorFaninTest, SubmitPathHeapDeadlockLogReportsRingAndRealHeapState) { std::vector create_infos; create_infos.reserve(8); diff --git a/tests/ut/cpp/a2a3/test_wiring.cpp b/tests/ut/cpp/a2a3/test_wiring.cpp index 8f781c914a..f316d89004 100644 --- a/tests/ut/cpp/a2a3/test_wiring.cpp +++ b/tests/ut/cpp/a2a3/test_wiring.cpp @@ -164,8 +164,8 @@ TEST_F(WiringTest, WireTaskAllProducersEarlyFinished) { // Consumer task with 2 fanins init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 2; - payload.fanin_inline_slot_states[0] = &producer_slots[0]; - payload.fanin_inline_slot_states[1] = &producer_slots[1]; + payload.fanin_inline_edges[0].set(&producer_slots[0], DEP_WAIT | DEP_RETAIN); + payload.fanin_inline_edges[1].set(&producer_slots[1], DEP_WAIT | DEP_RETAIN); task_slot.payload = &payload; task_slot.task = &desc; @@ -201,8 +201,8 @@ TEST_F(WiringTest, WireTaskProducersPendingTaskNotReady) { init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 2; - payload.fanin_inline_slot_states[0] = &producer_slots[0]; - payload.fanin_inline_slot_states[1] = &producer_slots[1]; + payload.fanin_inline_edges[0].set(&producer_slots[0], DEP_WAIT | DEP_RETAIN); + payload.fanin_inline_edges[1].set(&producer_slots[1], DEP_WAIT | DEP_RETAIN); task_slot.payload = &payload; task_slot.task = &desc; @@ -264,7 +264,7 @@ TEST_F(WiringTest, WireTaskMixedProducerStates) { init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 3; for (int i = 0; i < 3; i++) { - payload.fanin_inline_slot_states[i] = &producers[i]; + payload.fanin_inline_edges[i].set(&producers[i], DEP_WAIT | DEP_RETAIN); } task_slot.payload = &payload; task_slot.task = &desc; @@ -306,8 +306,8 @@ TEST_F(WiringTest, WireTaskAllFlaggedPrecompletedSeedsDispatchFanin) { init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 2; - payload.fanin_inline_slot_states[0] = &producer_slots[0]; - payload.fanin_inline_slot_states[1] = &producer_slots[1]; + payload.fanin_inline_edges[0].set(&producer_slots[0], DEP_WAIT | DEP_RETAIN); + payload.fanin_inline_edges[1].set(&producer_slots[1], DEP_WAIT | DEP_RETAIN); task_slot.payload = &payload; task_slot.task = &desc; @@ -334,7 +334,7 @@ TEST_F(WiringTest, WireTaskUnflaggedPrecompletedProducerDoesNotSeed) { init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 1; - payload.fanin_inline_slot_states[0] = &producer; + payload.fanin_inline_edges[0].set(&producer, DEP_WAIT | DEP_RETAIN); task_slot.payload = &payload; task_slot.task = &desc; @@ -359,8 +359,8 @@ TEST_F(WiringTest, WireTaskOneUnflaggedProducerDisqualifiesSeed) { init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 2; - payload.fanin_inline_slot_states[0] = &producers[0]; - payload.fanin_inline_slot_states[1] = &producers[1]; + payload.fanin_inline_edges[0].set(&producers[0], DEP_WAIT | DEP_RETAIN); + payload.fanin_inline_edges[1].set(&producers[1], DEP_WAIT | DEP_RETAIN); task_slot.payload = &payload; task_slot.task = &desc; @@ -385,7 +385,7 @@ TEST_F(WiringTest, EarlyDispatchWaitsForAllProducerBlocksPublished) { init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 1; - payload.fanin_inline_slot_states[0] = &producer; + payload.fanin_inline_edges[0].set(&producer, DEP_WAIT | DEP_RETAIN); task_slot.payload = &payload; task_slot.task = &desc; @@ -420,7 +420,7 @@ TEST_F(WiringTest, LateWiredFullyPublishedProducerStillSeedsEarlyDispatch) { init_slot(consumer, PTO2_TASK_PENDING, 0, 1); consumer_payload.fanin_actual_count = 1; - consumer_payload.fanin_inline_slot_states[0] = &producer; + consumer_payload.fanin_inline_edges[0].set(&producer, DEP_WAIT | DEP_RETAIN); consumer.payload = &consumer_payload; consumer.task = &consumer_desc; @@ -449,7 +449,7 @@ TEST_F(WiringTest, WiringSeedEnqueuesAfterConcurrentPropagation) { init_slot(consumer, PTO2_TASK_PENDING, 0, 1); consumer_payload.fanin_actual_count = 3; for (int i = 0; i < 3; i++) - consumer_payload.fanin_inline_slot_states[i] = &producers[i]; + consumer_payload.fanin_inline_edges[i].set(&producers[i], DEP_WAIT | DEP_RETAIN); consumer.payload = &consumer_payload; consumer.task = &consumer_desc; @@ -913,8 +913,8 @@ TEST_F(WiringTest, EarlyDispatchBlockedByUnflaggedProducer) { init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 2; - payload.fanin_inline_slot_states[0] = &p_flagged; - payload.fanin_inline_slot_states[1] = &q_unflagged; + payload.fanin_inline_edges[0].set(&p_flagged, DEP_WAIT | DEP_RETAIN); + payload.fanin_inline_edges[1].set(&q_unflagged, DEP_WAIT | DEP_RETAIN); task_slot.payload = &payload; task_slot.task = &desc; @@ -974,8 +974,8 @@ TEST_F(WiringTest, FlaggedPrecompletedCreatorTransparentToEarlyDispatch) { init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 2; - payload.fanin_inline_slot_states[0] = &creator; - payload.fanin_inline_slot_states[1] = &compute; + payload.fanin_inline_edges[0].set(&creator, DEP_WAIT | DEP_RETAIN); + payload.fanin_inline_edges[1].set(&compute, DEP_WAIT | DEP_RETAIN); task_slot.payload = &payload; task_slot.task = &desc; @@ -1055,8 +1055,8 @@ TEST_F(WiringTest, OnTaskReleaseReleasesProducers) { init_slot(task_slot, PTO2_TASK_COMPLETED, 3, 1); payload.fanin_actual_count = 2; - payload.fanin_inline_slot_states[0] = &producers[0]; - payload.fanin_inline_slot_states[1] = &producers[1]; + payload.fanin_inline_edges[0].set(&producers[0], DEP_WAIT | DEP_RETAIN); + payload.fanin_inline_edges[1].set(&producers[1], DEP_WAIT | DEP_RETAIN); // Need a valid fanin_spill_pool even though we don't spill PTO2FaninPool dummy_pool{}; PTO2FaninSpillEntry dummy_entries[4]; @@ -1078,6 +1078,93 @@ TEST_F(WiringTest, OnTaskReleaseReleasesProducers) { EXPECT_EQ(producers[1].task_state.load(), PTO2_TASK_CONSUMED); } +// ============================================================================= +// WAIT/RETAIN split (issue #1375): an ordering-only (DEP_WAIT) producer drops +// its submit->wire pin at wiring; a retention (DEP_WAIT|DEP_RETAIN) producer +// keeps it until on_task_release. Both are linked for completion notification. +// ============================================================================= + +TEST_F(WiringTest, OrderingOnlyReleasedAtWiringRetentionHeldUntilRelease) { + alignas(64) PTO2TaskSlotState task_slot; + alignas(64) PTO2TaskSlotState wait_producer; // DEP_WAIT only (modifier) + alignas(64) PTO2TaskSlotState retain_producer; // DEP_WAIT|DEP_RETAIN (creator) + alignas(64) PTO2TaskPayload payload; + memset(&payload, 0, sizeof(payload)); + PTO2TaskDescriptor desc{}; + + // Both live (PENDING) with a single submit pin (fanout_count = 1). + init_slot(wait_producer, PTO2_TASK_PENDING, 1, 1); + init_slot(retain_producer, PTO2_TASK_PENDING, 1, 1); + + init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); + payload.fanin_actual_count = 2; + payload.fanin_inline_edges[0].set(&wait_producer, DEP_WAIT); + payload.fanin_inline_edges[1].set(&retain_producer, DEP_WAIT | DEP_RETAIN); + PTO2FaninPool dummy_pool{}; + PTO2FaninSpillEntry dummy_entries[4]; + std::atomic dummy_error{PTO2_ERROR_NONE}; + dummy_pool.init(dummy_entries, 4, &dummy_error); + payload.fanin_spill_pool = &dummy_pool; + task_slot.payload = &payload; + task_slot.task = &desc; + + // Both WAIT edges gate readiness (wfanin = 2) and both link onto fanout_head. + wire_fanin(task_slot, 2); + EXPECT_NE(wait_producer.fanout_head, nullptr); + EXPECT_NE(retain_producer.fanout_head, nullptr); + + // Ordering-only pin released at wiring; retention pin still held. + EXPECT_EQ(wait_producer.fanout_refcount.load(), 1); + EXPECT_EQ(retain_producer.fanout_refcount.load(), 0); + + // Release: only the retention edge releases here; the ordering edge is not + // released a second time. + sched.on_task_release(task_slot); + EXPECT_EQ(wait_producer.fanout_refcount.load(), 1); + EXPECT_EQ(retain_producer.fanout_refcount.load(), 1); +} + +// on_task_release must honor per-edge flags in the spill region too: a spilled +// DEP_RETAIN edge is released; inline ordering-only edges are skipped. +TEST_F(WiringTest, ReleaseHonorsRetainFlagInSpillRegion) { + alignas(64) PTO2TaskSlotState filler; // 64 inline DEP_WAIT-only edges + alignas(64) PTO2TaskSlotState spill_retain; // 1 spilled DEP_RETAIN edge + alignas(64) PTO2TaskSlotState task_slot; + alignas(64) PTO2TaskPayload payload; + memset(&payload, 0, sizeof(payload)); + PTO2TaskDescriptor desc{}; + + // filler carries a large fanout_count so releasing it can never consume it. + init_slot(filler, PTO2_TASK_COMPLETED, 1, 100); + init_slot(spill_retain, PTO2_TASK_COMPLETED, 1, 1); + init_slot(task_slot, PTO2_TASK_COMPLETED, 0, 1); + + for (int i = 0; i < PTO2_FANIN_INLINE_CAP; i++) { + payload.fanin_inline_edges[i].set(&filler, DEP_WAIT); + } + PTO2FaninPool spill_pool{}; + PTO2FaninSpillEntry spill_entries[4]; + std::atomic err{PTO2_ERROR_NONE}; + spill_pool.init(spill_entries, 4, &err); + auto *e = spill_pool.alloc(); + int32_t spill_start = spill_pool.top - 1; + e->set(&spill_retain, DEP_WAIT | DEP_RETAIN); + + payload.fanin_actual_count = PTO2_FANIN_INLINE_CAP + 1; + payload.fanin_spill_start = spill_start; + payload.fanin_spill_pool = &spill_pool; + task_slot.payload = &payload; + task_slot.task = &desc; + + sched.on_task_release(task_slot); + + // Ordering-only inline edges are skipped; filler is untouched. + EXPECT_EQ(filler.fanout_refcount.load(), 0); + // The spilled retention edge is released (and consumed: rc == fc, COMPLETED). + EXPECT_EQ(spill_retain.fanout_refcount.load(), 1); + EXPECT_EQ(spill_retain.task_state.load(), PTO2_TASK_CONSUMED); +} + // ============================================================================= // advance_ring_pointers: scans CONSUMED slots, resets, advances last_alive // ============================================================================= diff --git a/tests/ut/cpp/a5/test_fanin_pool.cpp b/tests/ut/cpp/a5/test_fanin_pool.cpp index f901c0d7ad..8dcbdc8e05 100644 --- a/tests/ut/cpp/a5/test_fanin_pool.cpp +++ b/tests/ut/cpp/a5/test_fanin_pool.cpp @@ -41,7 +41,7 @@ class FaninPoolTest : public ::testing::Test { PTO2FaninPool pool{}; void SetUp() override { - entries.assign(POOL_CAP, PTO2FaninSpillEntry{nullptr}); + entries.assign(POOL_CAP, PTO2FaninSpillEntry{}); error_code.store(PTO2_ERROR_NONE); pool.init(entries.data(), POOL_CAP, &error_code); } @@ -165,7 +165,7 @@ class ForEachFaninTest : public ::testing::Test { alignas(64) PTO2TaskSlotState slots[64]; void SetUp() override { - spill_entries.assign(POOL_CAP, PTO2FaninSpillEntry{nullptr}); + spill_entries.assign(POOL_CAP, PTO2FaninSpillEntry{}); error_code.store(PTO2_ERROR_NONE); spill_pool.init(spill_entries.data(), POOL_CAP, &error_code); memset(slots, 0, sizeof(slots)); @@ -173,13 +173,13 @@ class ForEachFaninTest : public ::testing::Test { }; TEST_F(ForEachFaninTest, InlineOnlyVoid) { - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < 5; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } std::vector visited; - for_each_fanin_storage(inline_slots, 5, 0, spill_pool, [&](PTO2TaskSlotState *s) { + for_each_fanin_storage(inline_slots, 5, 0, spill_pool, [&](PTO2TaskSlotState *s, DepFlags) { visited.push_back(s); }); @@ -190,13 +190,13 @@ TEST_F(ForEachFaninTest, InlineOnlyVoid) { } TEST_F(ForEachFaninTest, InlineOnlyBoolEarlyReturn) { - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < 5; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } int count = 0; - bool result = for_each_fanin_storage(inline_slots, 5, 0, spill_pool, [&](PTO2TaskSlotState *) -> bool { + bool result = for_each_fanin_storage(inline_slots, 5, 0, spill_pool, [&](PTO2TaskSlotState *, DepFlags) -> bool { count++; return count < 3; // stop after 3rd }); @@ -206,12 +206,12 @@ TEST_F(ForEachFaninTest, InlineOnlyBoolEarlyReturn) { } TEST_F(ForEachFaninTest, InlineOnlyBoolAllTrue) { - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < 3; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } - bool result = for_each_fanin_storage(inline_slots, 3, 0, spill_pool, [](PTO2TaskSlotState *) -> bool { + bool result = for_each_fanin_storage(inline_slots, 3, 0, spill_pool, [](PTO2TaskSlotState *, DepFlags) -> bool { return true; }); @@ -219,9 +219,9 @@ TEST_F(ForEachFaninTest, InlineOnlyBoolAllTrue) { } TEST_F(ForEachFaninTest, ZeroFanin) { - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; int count = 0; - for_each_fanin_storage(inline_slots, 0, 0, spill_pool, [&](PTO2TaskSlotState *) { + for_each_fanin_storage(inline_slots, 0, 0, spill_pool, [&](PTO2TaskSlotState *, DepFlags) { count++; }); EXPECT_EQ(count, 0); @@ -233,20 +233,20 @@ TEST_F(ForEachFaninTest, ZeroFanin) { TEST_F(ForEachFaninTest, SpillNoWrap) { // 18 fanins = 16 inline + 2 spill - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < PTO2_FANIN_INLINE_CAP; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } // Allocate 2 spill entries auto *s0 = spill_pool.alloc(); int32_t spill_start = spill_pool.top - 1; - s0->slot_state = &slots[16]; + s0->set(&slots[16], DEP_WAIT | DEP_RETAIN); auto *s1 = spill_pool.alloc(); - s1->slot_state = &slots[17]; + s1->set(&slots[17], DEP_WAIT | DEP_RETAIN); std::vector visited; - for_each_fanin_storage(inline_slots, 18, spill_start, spill_pool, [&](PTO2TaskSlotState *s) { + for_each_fanin_storage(inline_slots, 18, spill_start, spill_pool, [&](PTO2TaskSlotState *s, DepFlags) { visited.push_back(s); }); @@ -268,9 +268,9 @@ TEST_F(ForEachFaninTest, SpillWithWrap) { spill_pool.top = POOL_CAP - 2; spill_pool.tail = POOL_CAP - 2; - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < PTO2_FANIN_INLINE_CAP; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } // 4 spill entries: indices 30, 31, 0, 1 (wraps around) @@ -278,11 +278,11 @@ TEST_F(ForEachFaninTest, SpillWithWrap) { for (int i = 0; i < 4; i++) { auto *e = spill_pool.alloc(); ASSERT_NE(e, nullptr); - e->slot_state = &slots[16 + i]; + e->set(&slots[16 + i], DEP_WAIT | DEP_RETAIN); } std::vector visited; - for_each_fanin_storage(inline_slots, 20, spill_start, spill_pool, [&](PTO2TaskSlotState *s) { + for_each_fanin_storage(inline_slots, 20, spill_start, spill_pool, [&](PTO2TaskSlotState *s, DepFlags) { visited.push_back(s); }); @@ -302,23 +302,61 @@ TEST_F(ForEachFaninTest, SpillWithWrap) { // ============================================================================= TEST_F(ForEachFaninTest, SpillBoolEarlyReturnInSpillRegion) { - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; for (int i = 0; i < PTO2_FANIN_INLINE_CAP; i++) { - inline_slots[i] = &slots[i]; + inline_slots[i].set(&slots[i], DEP_WAIT); } int32_t spill_start = spill_pool.top; for (int i = 0; i < 4; i++) { auto *e = spill_pool.alloc(); - e->slot_state = &slots[16 + i]; + e->set(&slots[16 + i], DEP_WAIT); } int count = 0; - bool result = for_each_fanin_storage(inline_slots, 20, spill_start, spill_pool, [&](PTO2TaskSlotState *) -> bool { - count++; - return count < 17; // stop on 17th (first spill entry) - }); + bool result = + for_each_fanin_storage(inline_slots, 20, spill_start, spill_pool, [&](PTO2TaskSlotState *, DepFlags) -> bool { + count++; + return count < 17; // stop on 17th (first spill entry) + }); EXPECT_FALSE(result); EXPECT_EQ(count, 17); } + +// ============================================================================= +// PTO2FaninSpillEntry: DepFlags packing round-trips across inline and spill +// ============================================================================= + +TEST_F(ForEachFaninTest, DepFlagsRoundTripInlineAndSpill) { + // Fill all 64 inline edges; slots 0..2 carry distinct flag combinations. + PTO2FaninSpillEntry inline_slots[PTO2_FANIN_INLINE_CAP] = {}; + inline_slots[0].set(&slots[0], DEP_WAIT); + inline_slots[1].set(&slots[1], DEP_RETAIN); + inline_slots[2].set(&slots[2], DEP_WAIT | DEP_RETAIN); + for (int i = 3; i < PTO2_FANIN_INLINE_CAP; i++) { + inline_slots[i].set(&slots[i], DEP_WAIT); + } + + // Two edges beyond the inline cap spill; they carry RETAIN-only and WAIT|RETAIN. + auto *s0 = spill_pool.alloc(); + int32_t spill_start = spill_pool.top - 1; + s0->set(&slots[0], DEP_RETAIN); + auto *s1 = spill_pool.alloc(); + s1->set(&slots[1], DEP_WAIT | DEP_RETAIN); + + const int32_t total = PTO2_FANIN_INLINE_CAP + 2; // 64 inline + 2 spill + std::vector flags; + for_each_fanin_storage(inline_slots, total, spill_start, spill_pool, [&](PTO2TaskSlotState *, DepFlags f) { + flags.push_back(f); + }); + + ASSERT_EQ(flags.size(), static_cast(total)); + // Inline flags survive. + EXPECT_EQ(flags[0], DEP_WAIT); + EXPECT_EQ(flags[1], DEP_RETAIN); + EXPECT_EQ(flags[2], DEP_WAIT | DEP_RETAIN); + // Spill flags survive. + EXPECT_EQ(flags[PTO2_FANIN_INLINE_CAP + 0], DEP_RETAIN); + EXPECT_EQ(flags[PTO2_FANIN_INLINE_CAP + 1], DEP_WAIT | DEP_RETAIN); +} diff --git a/tests/ut/cpp/a5/test_orchestrator_fanin.cpp b/tests/ut/cpp/a5/test_orchestrator_fanin.cpp index 75c4dec5fe..6a5f4f34e2 100644 --- a/tests/ut/cpp/a5/test_orchestrator_fanin.cpp +++ b/tests/ut/cpp/a5/test_orchestrator_fanin.cpp @@ -87,13 +87,140 @@ TEST_F(OrchestratorFaninTest, DuplicateExplicitProducerAddsOneFanin) { ASSERT_NE(consumer_slot.payload, nullptr); EXPECT_EQ(consumer_slot.payload->fanin_actual_count, 1); - EXPECT_EQ(consumer_slot.payload->fanin_inline_slot_states[0], &producer_slot); + EXPECT_EQ(consumer_slot.payload->fanin_inline_edges[0].slot_state(), &producer_slot); + // A plain set_dependencies() dep is conservative RETAIN: DEP_WAIT|DEP_RETAIN. + EXPECT_EQ(consumer_slot.payload->fanin_inline_edges[0].flags(), DEP_WAIT | DEP_RETAIN); // fanout_count is bit-packed: bit31 (PTO2_FANOUT_SCOPE_BIT) is the owning-scope // ref, low bits the consumer count. The duplicate explicit dep is deduped to a // single consumer, so this is scope + 1. EXPECT_EQ(producer_slot.fanout_count, PTO2_FANOUT_SCOPE_BIT + 1); } +// An explicit ordering-only dep (the primitive add_dep_wait() lowers to) yields a +// DEP_WAIT edge, not the conservative DEP_WAIT|DEP_RETAIN default. +TEST_F(OrchestratorFaninTest, ExplicitWaitDepProducesWaitOnlyEdge) { + orch.begin_scope(); + + CoreTaskArgs producer_args; + TaskOutputTensors producer = orch.submit_dummy_task(producer_args); + ASSERT_TRUE(producer.task_id().is_valid()); + + PTO2TaskId deps[] = {producer.task_id()}; + DepFlags kinds[] = {DEP_WAIT}; + CoreTaskArgs consumer_args; + consumer_args.set_dependencies_with_kinds(deps, kinds, 1); + TaskOutputTensors consumer = orch.submit_dummy_task(consumer_args); + ASSERT_TRUE(consumer.task_id().is_valid()); + + auto &consumer_slot = + sm_handle->header->rings[consumer.task_id().ring()].get_slot_state_by_task_id(consumer.task_id().local()); + ASSERT_NE(consumer_slot.payload, nullptr); + ASSERT_EQ(consumer_slot.payload->fanin_actual_count, 1); + EXPECT_EQ(consumer_slot.payload->fanin_inline_edges[0].flags(), DEP_WAIT); +} + +// The same producer reached with different kinds OR-accumulates into one edge: +// WAIT-only first, then WAIT|RETAIN folds RETAIN in, claiming exactly one pin. +TEST_F(OrchestratorFaninTest, DuplicateProducerOrAccumulatesFlags) { + orch.begin_scope(); + + CoreTaskArgs producer_args; + TaskOutputTensors producer = orch.submit_dummy_task(producer_args); + ASSERT_TRUE(producer.task_id().is_valid()); + + PTO2TaskId deps[] = {producer.task_id(), producer.task_id()}; + DepFlags kinds[] = {DEP_WAIT, DEP_WAIT | DEP_RETAIN}; + CoreTaskArgs consumer_args; + consumer_args.set_dependencies_with_kinds(deps, kinds, 2); + TaskOutputTensors consumer = orch.submit_dummy_task(consumer_args); + ASSERT_TRUE(consumer.task_id().is_valid()); + + auto &producer_slot = + sm_handle->header->rings[producer.task_id().ring()].get_slot_state_by_task_id(producer.task_id().local()); + auto &consumer_slot = + sm_handle->header->rings[consumer.task_id().ring()].get_slot_state_by_task_id(consumer.task_id().local()); + ASSERT_NE(consumer_slot.payload, nullptr); + ASSERT_EQ(consumer_slot.payload->fanin_actual_count, 1); + EXPECT_EQ(consumer_slot.payload->fanin_inline_edges[0].flags(), DEP_WAIT | DEP_RETAIN); + EXPECT_EQ(producer_slot.fanout_count, PTO2_FANOUT_SCOPE_BIT + 1); +} + +// The duplicate lands in the spill region (>64 fanin), exercising +// or_flags_into_existing's spill lookup: the dup folds (65 edges, not 66), claims +// exactly one pin, and OR-accumulates its flags into the spilled edge. +TEST_F(OrchestratorFaninTest, DuplicateProducerInSpillRegionDedups) { + orch.begin_scope(); + + constexpr int kProducers = PTO2_FANIN_INLINE_CAP + 1; // 65: the last one spills + std::vector producers; + producers.reserve(kProducers); + for (int i = 0; i < kProducers; i++) { + CoreTaskArgs a; + producers.push_back(orch.submit_dummy_task(a)); + ASSERT_TRUE(producers.back().task_id().is_valid()); + } + + std::vector deps; + std::vector kinds; + deps.reserve(kProducers + 1); + kinds.reserve(kProducers + 1); + for (auto &p : producers) { + deps.push_back(p.task_id()); + kinds.push_back(DEP_WAIT); // the 65th (first spill edge) starts WAIT-only + } + deps.push_back(producers.back().task_id()); // duplicate the spilled 65th ... + kinds.push_back(DEP_WAIT | DEP_RETAIN); // ... contributing RETAIN via the fold + + CoreTaskArgs consumer_args; + consumer_args.set_dependencies_with_kinds(deps.data(), kinds.data(), static_cast(deps.size())); + TaskOutputTensors consumer = orch.submit_dummy_task(consumer_args); + ASSERT_TRUE(consumer.task_id().is_valid()); + + auto &consumer_slot = + sm_handle->header->rings[consumer.task_id().ring()].get_slot_state_by_task_id(consumer.task_id().local()); + ASSERT_NE(consumer_slot.payload, nullptr); + PTO2TaskPayload *payload = consumer_slot.payload; + EXPECT_EQ(payload->fanin_actual_count, kProducers); // duplicate folded, not 66 + + PTO2TaskId dup = producers.back().task_id(); + auto &dup_slot = sm_handle->header->rings[dup.ring()].get_slot_state_by_task_id(dup.local()); + EXPECT_EQ(dup_slot.fanout_count, PTO2_FANOUT_SCOPE_BIT + 1); // one pin, not two + + // The first spilled edge is the duplicated producer; its flags OR-folded to + // WAIT|RETAIN across the two discovery kinds. + ASSERT_NE(payload->fanin_spill_pool, nullptr); + PTO2FaninPool &spill_pool = *payload->fanin_spill_pool; + PTO2FaninSpillEntry &spill_edge = spill_pool.base[payload->fanin_spill_start % spill_pool.capacity]; + EXPECT_EQ(spill_edge.slot_state(), &dup_slot); + EXPECT_EQ(spill_edge.flags(), DEP_WAIT | DEP_RETAIN); +} + +// The all-completed fast path (wire_fanin_task skipped) still drops an +// ordering-only producer's submit->wire pin. +TEST_F(OrchestratorFaninTest, AllCompletedFastPathReleasesWaitOnlyPin) { + orch.begin_scope(); + + CoreTaskArgs producer_args; + TaskOutputTensors producer = orch.submit_dummy_task(producer_args); + ASSERT_TRUE(producer.task_id().is_valid()); + auto &producer_slot = + sm_handle->header->rings[producer.task_id().ring()].get_slot_state_by_task_id(producer.task_id().local()); + // COMPLETED but not consumed (the open scope still pins it): the consumer takes + // the all-completed fast path. + producer_slot.task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); + int32_t rc_before = producer_slot.fanout_refcount.load(); + + PTO2TaskId deps[] = {producer.task_id()}; + DepFlags kinds[] = {DEP_WAIT}; // ordering-only + CoreTaskArgs consumer_args; + consumer_args.set_dependencies_with_kinds(deps, kinds, 1); + TaskOutputTensors consumer = orch.submit_dummy_task(consumer_args); + ASSERT_TRUE(consumer.task_id().is_valid()); + + // The fast path released the ordering-only pin. + EXPECT_EQ(producer_slot.fanout_refcount.load(), rc_before + 1); +} + TEST_F(OrchestratorFaninTest, SubmitPathHeapDeadlockLogReportsRingAndRealHeapState) { std::vector create_infos; create_infos.reserve(8); diff --git a/tests/ut/cpp/a5/test_wiring.cpp b/tests/ut/cpp/a5/test_wiring.cpp index 65db0ac110..fe21d91df2 100644 --- a/tests/ut/cpp/a5/test_wiring.cpp +++ b/tests/ut/cpp/a5/test_wiring.cpp @@ -160,8 +160,8 @@ TEST_F(WiringTest, WireTaskAllProducersEarlyFinished) { // Consumer task with 2 fanins init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 2; - payload.fanin_inline_slot_states[0] = &producer_slots[0]; - payload.fanin_inline_slot_states[1] = &producer_slots[1]; + payload.fanin_inline_edges[0].set(&producer_slots[0], DEP_WAIT | DEP_RETAIN); + payload.fanin_inline_edges[1].set(&producer_slots[1], DEP_WAIT | DEP_RETAIN); task_slot.payload = &payload; task_slot.task = &desc; @@ -197,8 +197,8 @@ TEST_F(WiringTest, WireTaskProducersPendingTaskNotReady) { init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 2; - payload.fanin_inline_slot_states[0] = &producer_slots[0]; - payload.fanin_inline_slot_states[1] = &producer_slots[1]; + payload.fanin_inline_edges[0].set(&producer_slots[0], DEP_WAIT | DEP_RETAIN); + payload.fanin_inline_edges[1].set(&producer_slots[1], DEP_WAIT | DEP_RETAIN); task_slot.payload = &payload; task_slot.task = &desc; @@ -240,7 +240,7 @@ TEST_F(WiringTest, WireTaskMixedProducerStates) { init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); payload.fanin_actual_count = 3; for (int i = 0; i < 3; i++) { - payload.fanin_inline_slot_states[i] = &producers[i]; + payload.fanin_inline_edges[i].set(&producers[i], DEP_WAIT | DEP_RETAIN); } task_slot.payload = &payload; task_slot.task = &desc; @@ -450,8 +450,8 @@ TEST_F(WiringTest, OnTaskReleaseReleasesProducers) { init_slot(task_slot, PTO2_TASK_COMPLETED, 3, 1); payload.fanin_actual_count = 2; - payload.fanin_inline_slot_states[0] = &producers[0]; - payload.fanin_inline_slot_states[1] = &producers[1]; + payload.fanin_inline_edges[0].set(&producers[0], DEP_WAIT | DEP_RETAIN); + payload.fanin_inline_edges[1].set(&producers[1], DEP_WAIT | DEP_RETAIN); // Need a valid fanin_spill_pool even though we don't spill PTO2FaninPool dummy_pool{}; PTO2FaninSpillEntry dummy_entries[4]; @@ -473,6 +473,93 @@ TEST_F(WiringTest, OnTaskReleaseReleasesProducers) { EXPECT_EQ(producers[1].task_state.load(), PTO2_TASK_CONSUMED); } +// ============================================================================= +// WAIT/RETAIN split (issue #1375): an ordering-only (DEP_WAIT) producer drops +// its submit->wire pin at wiring; a retention (DEP_WAIT|DEP_RETAIN) producer +// keeps it until on_task_release. Both are linked for completion notification. +// ============================================================================= + +TEST_F(WiringTest, OrderingOnlyReleasedAtWiringRetentionHeldUntilRelease) { + alignas(64) PTO2TaskSlotState task_slot; + alignas(64) PTO2TaskSlotState wait_producer; // DEP_WAIT only (modifier) + alignas(64) PTO2TaskSlotState retain_producer; // DEP_WAIT|DEP_RETAIN (creator) + alignas(64) PTO2TaskPayload payload; + memset(&payload, 0, sizeof(payload)); + PTO2TaskDescriptor desc{}; + + // Both live (PENDING) with a single submit pin (fanout_count = 1). + init_slot(wait_producer, PTO2_TASK_PENDING, 1, 1); + init_slot(retain_producer, PTO2_TASK_PENDING, 1, 1); + + init_slot(task_slot, PTO2_TASK_PENDING, 0, 1); + payload.fanin_actual_count = 2; + payload.fanin_inline_edges[0].set(&wait_producer, DEP_WAIT); + payload.fanin_inline_edges[1].set(&retain_producer, DEP_WAIT | DEP_RETAIN); + PTO2FaninPool dummy_pool{}; + PTO2FaninSpillEntry dummy_entries[4]; + std::atomic dummy_error{PTO2_ERROR_NONE}; + dummy_pool.init(dummy_entries, 4, &dummy_error); + payload.fanin_spill_pool = &dummy_pool; + task_slot.payload = &payload; + task_slot.task = &desc; + + // Both WAIT edges gate readiness (wfanin = 2) and both link onto fanout_head. + wire_fanin(task_slot, 2); + EXPECT_NE(wait_producer.fanout_head, nullptr); + EXPECT_NE(retain_producer.fanout_head, nullptr); + + // Ordering-only pin released at wiring; retention pin still held. + EXPECT_EQ(wait_producer.fanout_refcount.load(), 1); + EXPECT_EQ(retain_producer.fanout_refcount.load(), 0); + + // Release: only the retention edge releases here; the ordering edge is not + // released a second time. + sched.on_task_release(task_slot); + EXPECT_EQ(wait_producer.fanout_refcount.load(), 1); + EXPECT_EQ(retain_producer.fanout_refcount.load(), 1); +} + +// on_task_release must honor per-edge flags in the spill region too: a spilled +// DEP_RETAIN edge is released; inline ordering-only edges are skipped. +TEST_F(WiringTest, ReleaseHonorsRetainFlagInSpillRegion) { + alignas(64) PTO2TaskSlotState filler; // 64 inline DEP_WAIT-only edges + alignas(64) PTO2TaskSlotState spill_retain; // 1 spilled DEP_RETAIN edge + alignas(64) PTO2TaskSlotState task_slot; + alignas(64) PTO2TaskPayload payload; + memset(&payload, 0, sizeof(payload)); + PTO2TaskDescriptor desc{}; + + // filler carries a large fanout_count so releasing it can never consume it. + init_slot(filler, PTO2_TASK_COMPLETED, 1, 100); + init_slot(spill_retain, PTO2_TASK_COMPLETED, 1, 1); + init_slot(task_slot, PTO2_TASK_COMPLETED, 0, 1); + + for (int i = 0; i < PTO2_FANIN_INLINE_CAP; i++) { + payload.fanin_inline_edges[i].set(&filler, DEP_WAIT); + } + PTO2FaninPool spill_pool{}; + PTO2FaninSpillEntry spill_entries[4]; + std::atomic err{PTO2_ERROR_NONE}; + spill_pool.init(spill_entries, 4, &err); + auto *e = spill_pool.alloc(); + int32_t spill_start = spill_pool.top - 1; + e->set(&spill_retain, DEP_WAIT | DEP_RETAIN); + + payload.fanin_actual_count = PTO2_FANIN_INLINE_CAP + 1; + payload.fanin_spill_start = spill_start; + payload.fanin_spill_pool = &spill_pool; + task_slot.payload = &payload; + task_slot.task = &desc; + + sched.on_task_release(task_slot); + + // Ordering-only inline edges are skipped; filler is untouched. + EXPECT_EQ(filler.fanout_refcount.load(), 0); + // The spilled retention edge is released (and consumed: rc == fc, COMPLETED). + EXPECT_EQ(spill_retain.fanout_refcount.load(), 1); + EXPECT_EQ(spill_retain.task_state.load(), PTO2_TASK_CONSUMED); +} + // ============================================================================= // advance_ring_pointers: scans CONSUMED slots, resets, advances last_alive // =============================================================================