From a635be9955f14be85badc3e4a656754267cf50bb Mon Sep 17 00:00:00 2001 From: ChaoZheng109 Date: Mon, 10 Aug 2026 23:59:30 -0700 Subject: [PATCH] hbg: record Graph off the ring on A5 so the first invocation is one GRAPH task MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports a2a3 host_build_graph PR #1731 to A5. Mirrors #1731 line-for-line. With #1777 in place, the scalar-source machinery #1731's recording path depends on now exists on A5, so the port is direct: - PTO2TaskAllocator gains reserve_heap_scratch / restore_heap_top. - The recording pass routes submit_task / submit_dummy_task / alloc_tensors to graph_record_submit_node, which records each node's layout off the ring — output buffers come from heap scratch (released in graph_end) instead of a task-window slot. No ring slot, tensormap, fanin-pool entry, or upload is produced for an internal node. - A scope inside a Graph body is a no-op during recording. - graph_end rolls the heap back to the recording watermark, compacts and caches the Definition, and emits the single outer GRAPH task. It now returns bool; when the outer task cannot be placed, the body re-runs on the ordinary path so its work is still submitted. - The graph_end ops-table entry changes void -> bool in both mirrored ops structs. The first invocation now occupies one ring slot and one heap block instead of one per internal node. Verified: full no-hardware C++ unit suite (92/92) and the A5 graph_execution a5sim scenes (3/3). --- .../host_build_graph/docs/GRAPH_EXECUTION.md | 25 +- .../orchestration/pto_orchestration_api.h | 24 +- .../orchestrator_core/pto_orchestrator.cpp | 359 +++++++++++------- .../orchestrator_core/pto_runtime2.cpp | 4 +- .../runtime/pto_orchestrator.h | 2 +- .../runtime/pto_ring_buffer.h | 26 ++ .../host_build_graph/runtime/pto_runtime2.h | 2 +- 7 files changed, 293 insertions(+), 149 deletions(-) diff --git a/src/a5/runtime/host_build_graph/docs/GRAPH_EXECUTION.md b/src/a5/runtime/host_build_graph/docs/GRAPH_EXECUTION.md index ae8df28c67..9d854d1aea 100644 --- a/src/a5/runtime/host_build_graph/docs/GRAPH_EXECUTION.md +++ b/src/a5/runtime/host_build_graph/docs/GRAPH_EXECUTION.md @@ -4,10 +4,18 @@ Graph Execution is available only in the `host_build_graph` runtime. A Graph is a composite incore task: it is submitted and completed once like an AIC, AIV, MIX, or SPMD task, but contains a recorded task DAG. -The first invocation executes normally and records the DAG. A later invocation -places one `GRAPH` task in the host task window. The device Scheduler expands -the saved topology and dispatches its internal nodes; the Host Orchestrator does -not submit those nodes again. +Every invocation places exactly one `GRAPH` task in the host task window. The +first invocation records the DAG off the ring — its internal submissions build +host-only node metadata and reserve scratch output buffers instead of consuming +task-window slots — then emits the outer `GRAPH` task from the freshly built +Definition. Later invocations reuse the cached Definition and emit the same one +`GRAPH` task directly. In both cases the device Scheduler expands the saved +topology and dispatches the internal nodes; the Host Orchestrator never submits +those nodes as ring tasks. + +A recording that hits an unsupported construct is discarded and the body re-runs +on the ordinary task-submit path so its work is still submitted; the internal +nodes then occupy the ring only for that one fallback invocation. ## API @@ -175,10 +183,11 @@ void decode_three_layers( } ``` -The first layer records ordinary task submissions. Layers two and three submit -one Graph task each when their ChipTensor metadata and boundary scalar count -match. Each replay patches the current layer's `token_position`; its value is -not part of the Graph key. +All three layers submit one Graph task each: the first records the sub-DAG off +the ring and emits its Graph task, layers two and three replay the cached +Definition when their ChipTensor metadata and boundary scalar count match. Each +invocation patches the current layer's `token_position`; it is a dynamic +boundary scalar refreshed on every submission and is not part of the Graph key. ## Definition diff --git a/src/a5/runtime/host_build_graph/orchestration/pto_orchestration_api.h b/src/a5/runtime/host_build_graph/orchestration/pto_orchestration_api.h index 435d5d7427..77598e3188 100644 --- a/src/a5/runtime/host_build_graph/orchestration/pto_orchestration_api.h +++ b/src/a5/runtime/host_build_graph/orchestration/pto_orchestration_api.h @@ -94,7 +94,7 @@ typedef struct PTO2RuntimeOps { int32_t (*available_cluster_count)(PTO2Runtime *rt); int32_t (*available_aiv_count)(PTO2Runtime *rt); GraphScopeResult (*graph_begin)(PTO2Runtime *rt, uint64_t graph_key, const CoreTaskArgs &args); - void (*graph_end)(PTO2Runtime *rt); + bool (*graph_end)(PTO2Runtime *rt); void (*graph_commit)(PTO2Runtime *rt); // Stash the call-site of the next PTO2ScopeGuard so the [ScopeStats] @@ -221,12 +221,17 @@ static inline GraphScopeResult rt_graph_begin(uint64_t graph_key, const CoreTask return rt->ops->graph_begin(rt, graph_key, args); } -static inline void rt_graph_end() { +// Finish the recording pass. Returns true when the recorded sub-DAG was emitted +// as a single outer GRAPH task (the body must not run again); false when the +// recording was unsupported and the caller must re-run the body on the ordinary +// path. A fatal runtime is terminal, so it reports true to suppress a pointless +// re-run. +static inline bool rt_graph_end() { PTO2Runtime *rt = current_runtime(); if (rt->ops->is_fatal(rt) || rt->ops->graph_end == nullptr) { - return; + return true; } - rt->ops->graph_end(rt); + return rt->ops->graph_end(rt); } static inline void rt_graph_commit() { @@ -408,10 +413,17 @@ static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const C return GraphSubmitResult{}; } GraphScopeResult result = rt_graph_begin(graph_key, args); - if (result.execute_block) invoke(); if (result.recording) { - rt_graph_end(); + // First invocation: record the sub-DAG off the ring, then emit one outer + // GRAPH task. If the recording hit an unsupported construct, fall back and + // run the body on the ordinary path so its work is still submitted. + invoke(); + if (!rt_graph_end()) invoke(); + } else if (result.execute_block) { + // Un-cacheable at begin, or the Definition cache is full: ordinary path. + invoke(); } + // Cache hit: execute_block and recording are both false; the body is skipped. rt_graph_commit(); return result; } diff --git a/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp b/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp index a78e426f21..3f0f8dc9e6 100644 --- a/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp +++ b/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp @@ -307,12 +307,18 @@ struct GraphRecordedNode { struct GraphRecording { uint64_t full_key{0}; int32_t start_local_task_id{0}; - std::optional current_task_index; + // Heap allocation pointer captured at graph_begin. Internal nodes reserve + // scratch output buffers past it during the recording pass; graph_end rolls + // the heap back here so the outer GRAPH task reclaims the same region. + uint64_t heap_watermark{0}; + // The Graph boundary CoreTaskArgs, owned by the caller for the whole + // rt_submit_graph_impl call (recording pass + graph_end). graph_end replays + // it through graph_submit_definition to emit the outer GRAPH task with the + // current invocation's boundary tensor addresses. + const CoreTaskArgs *boundary_args{nullptr}; bool unsupported{false}; - std::vector current_fanins; std::vector boundary_tensors; std::vector boundary_types; - const CoreTaskArgs *boundary_args{nullptr}; std::vector nodes; }; @@ -423,127 +429,11 @@ bool graph_classify_tensor( return false; } -void graph_record_begin_task(PTO2OrchestratorState *orch, PTO2TaskId task_id) { - GraphHostState *state = graph_state_from(orch); - if (state == nullptr || state->recording == nullptr || state->recording->unsupported) return; - GraphRecording &recording = *state->recording; - const int32_t index = static_cast(task_id.local()) - recording.start_local_task_id; - if (index < 0 || index >= static_cast(GRAPH_MAX_NODES) || - index != static_cast(recording.nodes.size())) { - recording.unsupported = true; - return; - } - recording.current_task_index = static_cast(index); - recording.current_fanins.clear(); -} - -void graph_record_note_fanin(PTO2OrchestratorState *orch, PTO2TaskSlotState *producer) { - GraphHostState *state = graph_state_from(orch); - if (state == nullptr || state->recording == nullptr || state->recording->unsupported) return; - GraphRecording &recording = *state->recording; - if (producer == nullptr || producer->task == nullptr || !recording.current_task_index.has_value()) { - recording.unsupported = true; - return; - } - const int32_t producer_index = - static_cast(producer->task->task_id.local()) - recording.start_local_task_id; - if (producer_index >= 0 && static_cast(producer_index) >= *recording.current_task_index) { - recording.unsupported = true; - return; - } - if (producer_index >= 0) recording.current_fanins.push_back(static_cast(producer_index)); -} - void graph_record_mark_unsupported(PTO2OrchestratorState *orch) { GraphHostState *state = graph_state_from(orch); if (state != nullptr && state->recording != nullptr) state->recording->unsupported = true; } -void graph_record_task( - PTO2OrchestratorState *orch, PTO2TaskId task_id, const PTO2TaskDescriptor &task, const PTO2TaskPayload &payload, - const PTO2TaskSlotState &slot, const CoreTaskArgs &args -) { - GraphHostState *state = graph_state_from(orch); - if (state == nullptr || state->recording == nullptr || state->recording->unsupported) return; - GraphRecording &recording = *state->recording; - const int32_t task_index = static_cast(task_id.local()) - recording.start_local_task_id; - if (task_index < 0 || !recording.current_task_index.has_value() || - static_cast(task_index) != *recording.current_task_index || - static_cast(task_index) != recording.nodes.size() || args.predicate().op != PredicateOp::NONE) { - recording.unsupported = true; - return; - } - for (uint32_t i = 0; i < args.explicit_dep_count(); ++i) { - const PTO2TaskId dep = args.explicit_dep(i); - const int32_t dep_index = static_cast(dep.local()) - recording.start_local_task_id; - if (!dep.is_valid() || dep.ring() != 0 || dep_index >= task_index) { - recording.unsupported = true; - return; - } - if (dep_index < 0) { - const bool represented_by_boundary = std::any_of( - recording.boundary_tensors.begin(), recording.boundary_tensors.end(), [dep](const ChipTensor &tensor) { - return tensor.owner_task_id == dep; - } - ); - if (!represented_by_boundary) { - recording.unsupported = true; - return; - } - } - } - - GraphRecordedNode node; - std::copy_n(std::begin(task.kernel_id), PTO2_SUBTASK_SLOT_COUNT, node.kernel_ids.begin()); - node.active_mask = slot.active_mask; - node.task_attrs = slot.task_attrs; - node.task_attrs.set_early_resolve(false); - node.logical_block_num = slot.logical_block_num; - node.total_required_subtasks = slot.total_required_subtasks; - const uintptr_t packed_base = reinterpret_cast(task.packed_buffer_base); - const uintptr_t packed_end = reinterpret_cast(task.packed_buffer_end); - if (packed_end < packed_base) { - recording.unsupported = true; - return; - } - node.total_output_size = packed_end - packed_base; - node.record_packed_base = packed_base; - node.tensors.assign(payload.tensors, payload.tensors + payload.tensor_count); - node.tensor_sources.resize(static_cast(payload.tensor_count)); - node.scalars.assign(payload.scalars, payload.scalars + payload.scalar_count); - node.scalar_sources.resize(static_cast(payload.scalar_count)); - if (args.scalar_count() != payload.scalar_count) { - recording.unsupported = true; - return; - } - for (int32_t i = 0; i < payload.tensor_count; ++i) { - if (!graph_classify_tensor( - recording, node, task_index, payload.tensors[i], &node.tensor_sources[static_cast(i)] - )) { - recording.unsupported = true; - return; - } - } - for (int32_t i = 0; i < payload.scalar_count; ++i) { - GraphRecordedScalarSourceRef source = graph_classify_scalar(recording, args, i); - if (source.source == GraphRecordedScalarSource::INVALIDATED_BOUNDARY) { - recording.unsupported = true; - return; - } - node.scalar_sources[static_cast(i)] = source; - } - for (size_t producer : recording.current_fanins) { - if (producer >= static_cast(task_index)) { - recording.unsupported = true; - return; - } - node.internal_fanins.push_back(producer); - } - recording.nodes.push_back(std::move(node)); - recording.current_task_index.reset(); - recording.current_fanins.clear(); -} - GraphBoundarySignature graph_boundary_signature(const ChipTensor &tensor, TensorArgType type, uint16_t alias_rep) { GraphBoundarySignature signature{}; signature.buffer_size = tensor.buffer.size; @@ -843,7 +733,6 @@ static bool append_fanin_or_fail( if (fanin_builder->mark_seen(prod_ring, prod_slot)) { return true; } - graph_record_note_fanin(orch, prod_state); if (fanin_builder->count >= PTO2_MAX_FANIN) { orch_mark_fatal(orch, PTO2_ERROR_DEP_POOL_OVERFLOW); return false; @@ -949,8 +838,6 @@ static bool prepare_task( out->task = &orch->sm_header->ring.task_descriptors[out->alloc_result.slot]; out->payload = &orch->sm_header->ring.task_payloads[out->alloc_result.slot]; - graph_record_begin_task(orch, out->task_id); - out->payload->prefetch(args.tensor_count(), args.scalar_count()); // Re-bind payload/task pointers each submit. Value is per-slot constant @@ -1017,6 +904,14 @@ void PTO2OrchestratorState::begin_scope(PTO2ScopeMode mode) { if (orch->fatal) { return; } + // A Graph replays as a flat DAG with no scope structure: scope boundaries only + // shape scheduling on the ring, and the shadow-record path submits no ring + // tasks. Recorded ordering is preserved by the nodes' explicit dependencies + // and tensor-source classification, so a scope inside a Graph body is a no-op + // during recording and must not touch the real scope stack. + if (GraphHostState *state = graph_state_from(orch); state != nullptr && state->recording != nullptr) { + return; + } assert(orch->scope_stack_top < static_cast(orch->scope_stack_capacity - 1) && "Scope stack overflow"); if (mode == PTO2ScopeMode::AUTO && orch->in_manual_scope()) { report_fatal(PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "auto scope nested inside manual scope is not supported"); @@ -1052,6 +947,11 @@ void PTO2OrchestratorState::end_scope() { if (orch->fatal) { return; } + // Matches begin_scope: a scope inside a Graph body is a no-op during the + // shadow-record pass, so it must not touch the scope stack. + if (GraphHostState *state = graph_state_from(orch); state != nullptr && state->recording != nullptr) { + return; + } assert(orch->scope_stack_top >= 0 && "Scope stack underflow"); // Snapshot the ring start/end BEFORE the orchestrator drains pending tasks @@ -1396,7 +1296,6 @@ static TaskOutputTensors submit_task_common( // of position-independent integers, none of this needs host->device pointer // relocation. payload.fanin_count = fanin_builder.count; - graph_record_task(orch, task_id, task, payload, *prepared.slot_state, args); (void)sched; CYCLE_COUNT_LAP(g_orch_fanin_cycle); @@ -1606,6 +1505,150 @@ bool graph_submit_definition( return true; } +// Record one internal Graph node during the recording pass without consuming a +// ring task-window slot. Builds the node's metadata and materialized outputs +// exactly as submit_task_common would, but reserves output buffers from heap +// scratch (released in graph_end) and derives internal fanins from tensor-source +// classification — so no ring slot, tensormap entry, fanin-pool entry, or upload +// is produced for the node. The whole sub-DAG is later replayed by the single +// outer GRAPH task graph_end emits. The returned TaskOutputTensors borrow the +// node's own tensor storage; moving the node into recording.nodes keeps those +// addresses valid because the inner buffer is transferred, not copied. +TaskOutputTensors graph_record_submit_node( + PTO2OrchestratorState *orch, const CoreTaskArgs &args, ActiveMask active_mask, TaskAttrs task_attrs, + int32_t aic_kernel_id, int32_t aiv0_kernel_id, int32_t aiv1_kernel_id +) { + TaskOutputTensors result; + GraphRecording &recording = *graph_state_from(orch)->recording; + + const size_t node_index = recording.nodes.size(); + // A recorded node's index equals its local task id minus the recording + // baseline, so its synthetic id keeps classification and explicit-dep + // arithmetic identical to the ordinary path. + const PTO2TaskId task_id = + PTO2TaskId::make(0, static_cast(recording.start_local_task_id) + static_cast(node_index)); + result.set_task_id(task_id); + + if (node_index >= GRAPH_MAX_NODES || args.has_error || args.predicate().op != PredicateOp::NONE) { + recording.unsupported = true; + } + + const PTO2OutputLayout layout = calculate_output_layout(args); + void *packed_base = orch->ring.task_allocator.reserve_heap_scratch(layout.total_output_size); + if (layout.total_output_size > 0 && packed_base == nullptr) { + // No scratch storage for this node's outputs: the sub-DAG cannot be + // recorded, so graph_end falls back and the body runs on the ordinary + // path where the same heap pressure is reported through alloc(). + recording.unsupported = true; + return result; + } + const uint64_t aligned_output = + layout.total_output_size > 0 ? PTO2_ALIGN_UP(static_cast(layout.total_output_size), PTO2_ALIGN_SIZE) : + 0; + + GraphRecordedNode node; + node.kernel_ids[static_cast(PTO2SubtaskSlot::AIC)] = aic_kernel_id; + node.kernel_ids[static_cast(PTO2SubtaskSlot::AIV0)] = aiv0_kernel_id; + node.kernel_ids[static_cast(PTO2SubtaskSlot::AIV1)] = aiv1_kernel_id; + node.active_mask = active_mask; + node.task_attrs = task_attrs; + node.task_attrs.set_early_resolve(false); + node.logical_block_num = args.launch_spec.block_num(); + // Mirror prepare_task's contract: block_num must be positive and the subtask + // count must fit int16_t. An out-of-contract value marks the recording + // unsupported so graph_end falls back to the ordinary path, which reports + // PTO2_ERROR_INVALID_ARGS, rather than baking a truncated or negative count + // into the cached Definition (which the device would expand into a node that + // never completes). + const int32_t required_subtasks = + static_cast(node.logical_block_num) * __builtin_popcount(active_mask.core_mask()); + if (node.logical_block_num <= 0 || required_subtasks > std::numeric_limits::max()) { + recording.unsupported = true; + node.total_required_subtasks = 0; + } else { + node.total_required_subtasks = static_cast(required_subtasks); + } + node.record_packed_base = reinterpret_cast(packed_base); + node.total_output_size = aligned_output; + + // Build the tensor list exactly as PTO2TaskPayload::init: inputs/inouts copy + // the caller's ChipTensor; outputs materialize from the create-info onto the + // scratch buffer and carry this node's owner id. + const int32_t tensor_count = args.tensor_count(); + node.tensors.resize(static_cast(tensor_count)); + for (int32_t i = 0; i < tensor_count; ++i) { + ChipTensor &slot_tensor = node.tensors[static_cast(i)]; + if (args.tag(i) != TensorArgType::OUTPUT) { + slot_tensor.copy(args.tensor(i).ref()); + } else { + init_tensor_from_create_info( + slot_tensor, args.tensor(i).create_info(), + reinterpret_cast(reinterpret_cast(packed_base) + layout.offsets[i]), + layout.buffer_sizes[i] + ); + slot_tensor.owner_task_id = task_id; + } + } + // Materialize output refs only after node.tensors is fully sized so the + // borrowed addresses stay stable across the move into recording.nodes. + for (int32_t i = 0; i < tensor_count; ++i) { + if (args.tag(i) == TensorArgType::OUTPUT) result.materialize_output(node.tensors[static_cast(i)]); + } + node.scalars.assign(args.scalars(), args.scalars() + args.scalar_count()); + + // Classify each scalar's source: a plain literal is static Definition data, + // while a value copied from a boundary scalar is refreshed on replay. A + // mutable tracked boundary scalar is not supported and falls back. + node.scalar_sources.resize(static_cast(args.scalar_count())); + for (int32_t i = 0; i < args.scalar_count(); ++i) { + GraphRecordedScalarSourceRef source = graph_classify_scalar(recording, args, i); + if (source.source == GraphRecordedScalarSource::INVALIDATED_BOUNDARY) recording.unsupported = true; + node.scalar_sources[static_cast(i)] = source; + } + + // Classify each tensor's source, then derive internal fanins from the + // INTERNAL classifications plus any explicit internal dependency. + node.tensor_sources.resize(static_cast(tensor_count)); + for (int32_t i = 0; i < tensor_count; ++i) { + if (!graph_classify_tensor( + recording, node, static_cast(node_index), node.tensors[static_cast(i)], + &node.tensor_sources[static_cast(i)] + )) { + recording.unsupported = true; + } + } + auto add_fanin = [&node](size_t producer) { + if (std::find(node.internal_fanins.begin(), node.internal_fanins.end(), producer) == + node.internal_fanins.end()) { + node.internal_fanins.push_back(producer); + } + }; + for (const GraphRecordedTensorSourceRef &source : node.tensor_sources) { + if (source.source == GraphRecordedTensorSource::INTERNAL) add_fanin(source.source_index); + } + for (uint32_t i = 0; i < args.explicit_dep_count(); ++i) { + const PTO2TaskId dep = args.explicit_dep(i); + const int32_t dep_index = static_cast(dep.local()) - recording.start_local_task_id; + if (!dep.is_valid() || dep.ring() != 0 || dep_index >= static_cast(node_index)) { + recording.unsupported = true; + continue; + } + if (dep_index < 0) { + const bool represented_by_boundary = std::any_of( + recording.boundary_tensors.begin(), recording.boundary_tensors.end(), [dep](const ChipTensor &tensor) { + return tensor.owner_task_id == dep; + } + ); + if (!represented_by_boundary) recording.unsupported = true; + } else { + add_fanin(static_cast(dep_index)); + } + } + + recording.nodes.push_back(std::move(node)); + return result; +} + } // namespace GraphScopeResult @@ -1654,6 +1697,7 @@ PTO2OrchestratorState::graph_begin(uint64_t graph_key, const CoreTaskArgs &args, auto recording = std::make_unique(); recording->full_key = full_key; recording->start_local_task_id = orch->ring.task_allocator.active_count(); + recording->heap_watermark = orch->ring.task_allocator.heap_top(); args.anchor_scalar_sources(); recording->boundary_args = &args; recording->boundary_tensors.reserve(static_cast(args.tensor_count())); @@ -1663,27 +1707,58 @@ PTO2OrchestratorState::graph_begin(uint64_t graph_key, const CoreTaskArgs &args, recording->boundary_types.push_back(args.tag(i)); } state->recording = std::move(recording); + // The body runs through the shadow-record path, not the ring, so it does not + // execute as ordinary tasks. + result.execute_block = false; result.recording = true; return result; } -void PTO2OrchestratorState::graph_end() { +// Finish the recording pass. Returns true when the sub-DAG was compacted into a +// Definition and emitted as a single outer GRAPH task; false when the recording +// hit an unsupported construct or the outer task could not be placed, in which +// case the caller re-runs the body on the ordinary path so its work is still +// submitted. +bool PTO2OrchestratorState::graph_end() { GraphHostState *state = graph_state_from(this); - if (state == nullptr || state->recording == nullptr) return; + if (state == nullptr || state->recording == nullptr) return false; std::unique_ptr recording = std::move(state->recording); + // Release the scratch output buffers the internal nodes reserved; the outer + // GRAPH task reclaims the same heap region as one block below (or the + // ordinary fallback re-uses it per task). + this->ring.task_allocator.restore_heap_top(recording->heap_watermark); + std::vector definition; if (!graph_build_definition(*recording, &definition)) { debug_assert(false && "The recorded Graph contains a construct that Graph Execution does not support"); - LOG_WARN("%s", "[GraphExecution] unsupported construct observed; definition was not cached"); - return; + LOG_WARN("%s", "[GraphExecution] unsupported construct observed; falling back to the ordinary path"); + return false; } const GraphDefinition *header = graph_definition(definition); - if (header == nullptr) return; + if (header == nullptr) return false; LOG_DEBUG( "[GraphExecution] define key=0x%llx nodes=%u bytes=%u", static_cast(header->full_key), header->task_count, header->total_bytes ); - state->definitions.emplace(header->full_key, std::move(definition)); + auto inserted = state->definitions.emplace(header->full_key, std::move(definition)); + const std::vector &cached = inserted.first->second; + + // Emit the single outer GRAPH task for this first invocation exactly as a + // cache hit would, so the recording run occupies one ring slot and one heap + // block instead of one per internal node. The Definition stays cached for + // subsequent invocations even if this placement fails. + PTO2TaskId submitted = PTO2TaskId::invalid(); + if (recording->boundary_args == nullptr || + !graph_submit_definition(this, state, cached, *recording->boundary_args, &submitted)) { + return false; + } +#if SIMPLER_DFX + g_orch_submit_idx++; +#if SIMPLER_ORCH_PROFILING + g_orch_submit_count++; +#endif +#endif + return true; } void PTO2OrchestratorState::graph_commit() {} @@ -1762,6 +1837,13 @@ TaskOutputTensors PTO2OrchestratorState::submit_task(const MixedKernels &mixed_k task_attrs.set_predicate(); } + if (GraphHostState *state = graph_state_from(orch); state != nullptr && state->recording != nullptr) { + return graph_record_submit_node( + orch, args, active_mask, task_attrs, normalized.aic_kernel_id, normalized.aiv0_kernel_id, + normalized.aiv1_kernel_id + ); + } + return submit_task_common( orch, args, active_mask, task_attrs, normalized.aic_kernel_id, normalized.aiv0_kernel_id, normalized.aiv1_kernel_id @@ -1798,6 +1880,12 @@ TaskOutputTensors PTO2OrchestratorState::submit_dummy_task(const CoreTaskArgs &a task_attrs.set_early_resolve(args.allow_early_resolve()); task_attrs.set_timing_slot(args.task_timing_slot()); + if (GraphHostState *state = graph_state_from(orch); state != nullptr && state->recording != nullptr) { + return graph_record_submit_node( + orch, args, ActiveMask{}, task_attrs, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID + ); + } + return submit_task_common( orch, args, ActiveMask{}, task_attrs, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID ); @@ -1805,7 +1893,6 @@ TaskOutputTensors PTO2OrchestratorState::submit_dummy_task(const CoreTaskArgs &a TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const CoreTaskArgs &args) { auto *orch = this; - graph_record_mark_unsupported(orch); // Orchestration API should short-circuit after fatal, but keep this entry // robust as a no-op in case a caller reaches it directly. if (orch->fatal) { @@ -1839,6 +1926,18 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const CoreTaskArgs &args) return TaskOutputTensors{}; } + // Runtime-allocated outputs cannot be replayed by a Graph, so a Graph body + // that calls alloc_tensors is unsupported. Still materialize the outputs so + // the recording body chains correctly, then poison the recording; graph_end + // falls back and the body re-runs on the ordinary path. + if (GraphHostState *state = graph_state_from(orch); state != nullptr && state->recording != nullptr) { + TaskOutputTensors result = graph_record_submit_node( + orch, args, ActiveMask{}, TaskAttrs{}, INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID + ); + graph_record_mark_unsupported(orch); + return result; + } + PTO2OutputLayout layout = calculate_output_layout(args); PTO2PreparedTask prepared; // Kernel-less alloc task: no active subtasks, no dispatch-time attributes. The diff --git a/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp b/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp index d978a3835d..c122e8562a 100644 --- a/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp +++ b/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp @@ -103,9 +103,7 @@ static GraphScopeResult graph_begin_impl(PTO2Runtime *rt, uint64_t graph_key, co return rt->orchestrator.graph_begin(graph_key, args, rt->active_callable_hash); } -static void graph_end_impl(PTO2Runtime *rt) { - if (rt != nullptr) rt->orchestrator.graph_end(); -} +static bool graph_end_impl(PTO2Runtime *rt) { return rt != nullptr && rt->orchestrator.graph_end(); } static void graph_commit_impl(PTO2Runtime *rt) { if (rt != nullptr) rt->orchestrator.graph_commit(); diff --git a/src/a5/runtime/host_build_graph/runtime/pto_orchestrator.h b/src/a5/runtime/host_build_graph/runtime/pto_orchestrator.h index 30fa01e1b9..aec93828b5 100644 --- a/src/a5/runtime/host_build_graph/runtime/pto_orchestrator.h +++ b/src/a5/runtime/host_build_graph/runtime/pto_orchestrator.h @@ -162,7 +162,7 @@ struct PTO2OrchestratorState { TaskOutputTensors submit_dummy_task(const CoreTaskArgs &args); TaskOutputTensors alloc_tensors(const CoreTaskArgs &args); GraphScopeResult graph_begin(uint64_t graph_key, const CoreTaskArgs &args, uint64_t callable_hash); - void graph_end(); + bool graph_end(); void graph_commit(); void mark_done(); }; diff --git a/src/a5/runtime/host_build_graph/runtime/pto_ring_buffer.h b/src/a5/runtime/host_build_graph/runtime/pto_ring_buffer.h index 5da724b4c2..f3c7679f4f 100644 --- a/src/a5/runtime/host_build_graph/runtime/pto_ring_buffer.h +++ b/src/a5/runtime/host_build_graph/runtime/pto_ring_buffer.h @@ -250,6 +250,32 @@ class PTO2TaskAllocator { return (heap_top_ + heap_size_ - heap_tail_) % heap_size_; } + // Reserve output-buffer bytes without claiming a task-window slot. Graph + // recording gives each internal node a disjoint, valid heap address so + // tensor-source classification works, then releases the whole range with + // restore_heap_top() once the consolidated outer GRAPH task reclaims it as a + // single block. The reservation is strictly forward: it never takes + // try_bump_heap's wrap branches, because those also mutate heap_tail_ and + // heap_rebase_anchor_task_id_, which restore_heap_top() cannot undo. On + // exhaustion it returns nullptr and leaves allocator state unchanged; + // graph_record_submit_node treats that as an unsupported recording and falls + // back to the ordinary path. A zero size returns the current position. + void *reserve_heap_scratch(int32_t output_size) { + uint64_t aligned_size = + output_size > 0 ? PTO2_ALIGN_UP(static_cast(output_size), PTO2_ALIGN_SIZE) : 0; + uint64_t top = heap_top_; + if (aligned_size == 0) return static_cast(heap_base_) + top; + if (top < heap_tail_ || heap_size_ - top < aligned_size) return nullptr; + heap_top_ = top + aligned_size; + return static_cast(heap_base_) + top; + } + + // Roll the heap allocation pointer back to a value previously returned by + // heap_top(). Valid only while no allocation has been consumed since the + // snapshot (heap_tail_ unchanged), which holds throughout host graph + // construction. + void restore_heap_top(uint64_t top) { heap_top_ = top; } + private: // --- Task Ring --- PTO2TaskDescriptor *descriptors_ = nullptr; diff --git a/src/a5/runtime/host_build_graph/runtime/pto_runtime2.h b/src/a5/runtime/host_build_graph/runtime/pto_runtime2.h index 024900fcc3..cb4a4987bb 100644 --- a/src/a5/runtime/host_build_graph/runtime/pto_runtime2.h +++ b/src/a5/runtime/host_build_graph/runtime/pto_runtime2.h @@ -97,7 +97,7 @@ struct PTO2RuntimeOps { int32_t (*available_cluster_count)(PTO2Runtime *rt); int32_t (*available_aiv_count)(PTO2Runtime *rt); GraphScopeResult (*graph_begin)(PTO2Runtime *rt, uint64_t graph_key, const CoreTaskArgs &args); - void (*graph_end)(PTO2Runtime *rt); + bool (*graph_end)(PTO2Runtime *rt); void (*graph_commit)(PTO2Runtime *rt); // Stash the call-site captured by PTO2ScopeGuard into the [ScopeStats] // collector. Always present in the struct to keep ops-table layout stable