Skip to content

tensormap_and_ringbuffer: split fanin edges into orthogonal WAIT/RETAIN flags (#1375) #1806

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChaoZheng109 merged 1 commit into from
Aug 13, 2026
Merged

tensormap_and_ringbuffer: split fanin edges into orthogonal WAIT/RETAIN flags (#1375) #1806

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions docs/dfx/dep-gen.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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 |
Expand Down
4 changes: 2 additions & 2 deletions src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
95 changes: 65 additions & 30 deletions src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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).
Expand Down Expand Up @@ -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) << '"';
}
Expand Down Expand Up @@ -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<uint64_t> oracle_preds;
std::unordered_set<uint64_t> 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<uint64_t, DepFlags> oracle_preds;
std::unordered_map<uint64_t, DepFlags> annot_preds;

// Scratch buffer for assembling full dep lists across overflow chains.
// Declared outside the loop so it can be reused (clear() keeps capacity).
Expand Down Expand Up @@ -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();
Expand All @@ -713,14 +738,17 @@ 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{};
e.pred = producer.raw;
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);
Expand All @@ -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);
Expand All @@ -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<unsigned>(f));
} else if (it->second != f) {
LOG_ERROR(
" flags mismatch pred: %" PRIu64 " oracle=%u annot=%u", p, static_cast<unsigned>(f),
static_cast<unsigned>(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<unsigned>(f));
}
}
tm_oracle.destroy();
Expand Down
50 changes: 35 additions & 15 deletions src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <typename... Ids>
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<std::decay_t<Ids>, 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 <typename... Ids>
void add_dep_wait(Ids... ids) {
add_dep_impl(DEP_WAIT, ids...);
}

/**
Expand All @@ -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 <typename... Ids>
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<std::decay_t<Ids>, 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;
};

Expand Down
Loading
Loading