From 67df8ecc7729c0328d8d783c8e377a9647db4d05 Mon Sep 17 00:00:00 2001 From: yanghaoran29 Date: Tue, 18 Aug 2026 19:00:50 +0800 Subject: [PATCH] Perf: pin Graph PODs and H2D each layer synchronously Rebased onto main after the shared Graph Definition upload (#1874) landed: submissions now reference a runner-retained Definition device object, so the eager uploader creates that object lazily on first reference instead of a batched pre-pass. The eager-upload hook and pinned bump arena now live on the run-owned GraphHostState instead of process globals, with an RAII guard clearing the hook on every exit path, so overlapping orchestrations cannot clobber each other's uploader or bump cursor. The pinned arena is allocated through dlopen-resolved aclrt entry points, falling back to plain host memory where no Ascend toolkit is present (sim builds, CI runners). Retained device submission storage moves from a process-lifetime static map with a lossy packed key to a new runner-owned HostApi op (acquire_graph_submission_buffer) keyed by (graph_key, occurrence), released at Worker finalization like the execution and Definition buffers. An eager-upload failure after the outer GRAPH task is published now latches EXPLICIT_ORCH_FATAL instead of returning an ordinary-path fallback, which would have re-submitted the graph body on top of a half-uploaded task. Both a2a3 and a5 host_build_graph carry the change. --- .../host_build_graph/host/runtime_maker.cpp | 196 +++++++++++++++--- .../orchestration/pto_orchestration_api.h | 7 +- .../orchestrator_core/pto_orchestrator.cpp | 160 ++++++++++++-- .../host_build_graph/host/runtime_maker.cpp | 195 ++++++++++++++--- .../orchestration/pto_orchestration_api.h | 7 +- .../orchestrator_core/pto_orchestrator.cpp | 160 ++++++++++++-- .../host_build_graph/graph_host_state.h | 19 ++ src/common/platform/include/common/host_api.h | 16 ++ .../platform/onboard/host/c_api_shared.cpp | 13 ++ .../onboard/host/device_runner_base.cpp | 46 +++- .../onboard/host/device_runner_base.h | 6 + src/common/platform/sim/host/c_api_shared.cpp | 13 ++ .../platform/sim/host/device_runner_base.cpp | 46 +++- .../platform/sim/host/device_runner_base.h | 6 + 14 files changed, 774 insertions(+), 116 deletions(-) diff --git a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp index b1c3089cbb..8fea765e7a 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -44,6 +44,7 @@ #include #include #include +#include #include #include #include @@ -416,27 +417,95 @@ static bool relocate_host_orch_image( return ok; } -bool upload_graph_submissions( - Runtime *runtime, const HostApi *api, GraphHostState &graph_state, uint64_t &uploaded_bytes -) { +// Retained pinned bump arena for Graph POD images. Orch writes each layer +// in place; upload_one synchronously H2Ds that POD (source is already pinned). +// The CANN entry points are resolved with dlopen so the same source builds on +// hosts without an Ascend toolkit (sim runners, CI); there the arena falls +// back to plain memory, which costs nothing because sim copies are memcpy. +struct GraphPinnedPack { + void *host = nullptr; + size_t cap = 0; + bool pinned = false; +}; + +static GraphPinnedPack g_graph_pack; +static std::mutex g_graph_pack_mu; +constexpr size_t kGraphPinnedArenaBytes = 16ull * 1024ull * 1024ull; + +namespace { +using AclrtMallocHostFn = int (*)(void **, size_t); +using AclrtFreeHostFn = int (*)(void *); + +AclrtMallocHostFn graph_aclrt_malloc_host() { + void *handle = dlopen("libascendcl.so", RTLD_NOW | RTLD_LOCAL); + if (handle == nullptr) return nullptr; + return reinterpret_cast(dlsym(handle, "aclrtMallocHost")); +} + +AclrtFreeHostFn graph_aclrt_free_host() { + void *handle = dlopen("libascendcl.so", RTLD_NOW | RTLD_LOCAL); + if (handle == nullptr) return nullptr; + return reinterpret_cast(dlsym(handle, "aclrtFreeHost")); +} +} // namespace + +static void free_graph_pack(GraphPinnedPack &pack) { + if (pack.host == nullptr) return; + if (pack.pinned) { + AclrtFreeHostFn free_host = graph_aclrt_free_host(); + if (free_host != nullptr) (void)free_host(pack.host); + } else { + delete[] static_cast(pack.host); + } + pack.host = nullptr; + pack.cap = 0; + pack.pinned = false; +} + +static bool ensure_graph_pinned_pack(size_t cap) { + std::lock_guard lock(g_graph_pack_mu); + if (g_graph_pack.cap >= cap && g_graph_pack.host != nullptr) { + return true; + } + free_graph_pack(g_graph_pack); + AclrtMallocHostFn malloc_host = graph_aclrt_malloc_host(); + if (malloc_host != nullptr) { + void *host = nullptr; + if (malloc_host(&host, cap) != 0 || host == nullptr) { + LOG_ERROR("host-orch: aclrtMallocHost(%zu) failed", cap); + return false; + } + std::memset(host, 0, cap); + g_graph_pack.host = host; + g_graph_pack.pinned = true; + } else { + g_graph_pack.host = new std::byte[cap]{}; + g_graph_pack.pinned = false; + } + g_graph_pack.cap = cap; + return true; +} + +struct GraphPodH2d { + const HostApi *api = nullptr; std::unordered_map occurrences; - uploaded_bytes = 0; - const size_t count = graph_host_upload_count(graph_state); - // Pass 1: upload each distinct Definition once as a shared device object - // ([GraphDefinitionHeader][Definition image]) keyed by content identity. - // Submissions reference the object's GM address, so this pass completing - // before any submission is uploaded is what makes the reference safe — - // the device boots only after both passes. - GraphHostDefinitionList definitions = graph_host_definitions(graph_state); + struct UploadedDefinition { void *device_object; // GM address; host must not dereference const GraphDefinition *host_view; // the host-side image the object was built from }; std::unordered_map definition_objects; - for (const GraphHostDefinition &entry : definitions.entries) { - if (entry.data == nullptr || entry.bytes < sizeof(GraphDefinition)) continue; + + // Upload one distinct Definition as a shared device object + // ([GraphDefinitionHeader][Definition image]) keyed by content identity. + // Submissions reference the object's GM address, so the object existing + // before any submission referencing it is uploaded is what makes the + // reference safe — the device boots only after both are done. + bool ensure_definition_object(const GraphHostDefinition &entry) { + if (entry.data == nullptr || entry.bytes < sizeof(GraphDefinition)) return false; const auto *definition = reinterpret_cast(entry.data); - if (definition->total_bytes != entry.bytes || definition->full_key != entry.full_key) continue; + if (definition->total_bytes != entry.bytes || definition->full_key != entry.full_key) return false; + if (definition_objects.count(definition->content_hash) != 0) return true; const size_t object_bytes = sizeof(GraphDefinitionHeader) + entry.bytes; void *object = api->acquire_graph_definition_buffer(entry.full_key, object_bytes, alignof(GraphDefinitionHeader)); @@ -462,11 +531,27 @@ bool upload_graph_submissions( return false; } definition_objects.emplace(definition->content_hash, UploadedDefinition{object, definition}); - uploaded_bytes += object_bytes; + return true; + } + + // Returns the Definition entry a submission references, or nullopt when the + // host state holds no valid Definition for it. + std::optional find_definition(const GraphHostState &graph_state, uint64_t content_hash) { + GraphHostDefinitionList definitions = graph_host_definitions(const_cast(graph_state)); + for (const GraphHostDefinition &entry : definitions.entries) { + if (entry.bytes < sizeof(GraphDefinition) || entry.data == nullptr) continue; + const auto *definition = reinterpret_cast(entry.data); + if (definition->content_hash == content_hash && definition->total_bytes == entry.bytes) { + return entry; + } + } + return std::nullopt; } - // Pass 2: per-submission execution storage + the small reference image. - for (size_t index = 0; index < count; ++index) { + bool upload_one(GraphHostState &graph_state, size_t index) { + if (graph_host_upload_h2d_done(graph_state, index)) { + return true; + } std::optional upload = graph_host_upload(graph_state, index); if (!upload.has_value() || upload->outer_slot == nullptr || upload->data == nullptr || upload->bytes < sizeof(GraphSubmission) || upload->outer_slot->task_kind != TaskKind::GRAPH || @@ -480,9 +565,19 @@ bool upload_graph_submissions( return false; } auto object_it = definition_objects.find(submission->definition_hash); - if (object_it == definition_objects.end() || object_it->second.device_object == nullptr) { - LOG_ERROR("host-orch: Graph submission has no uploaded Definition object"); - return false; + if (object_it == definition_objects.end()) { + // Eager uploads run during orch entry, before any batched pass, so + // the Definition object for this submission may not exist yet. + std::optional entry = find_definition(graph_state, submission->definition_hash); + if (!entry.has_value() || !ensure_definition_object(*entry)) { + LOG_ERROR("host-orch: Graph submission has no uploadable Definition"); + return false; + } + object_it = definition_objects.find(submission->definition_hash); + if (object_it == definition_objects.end()) { + LOG_ERROR("host-orch: Graph submission has no uploaded Definition object"); + return false; + } } // Capacities come from the host-side Definition image the device // object was built from; the GM object itself is never dereferenced @@ -503,10 +598,7 @@ bool upload_graph_submissions( submission->graph_key, occurrence, execution_bytes, alignof(GraphNodeStorage) ); if (execution_storage == nullptr) { - LOG_ERROR( - "host-orch: failed to retain %zu bytes for Graph execution key=%#llx occurrence=%u", execution_bytes, - static_cast(submission->graph_key), occurrence - ); + LOG_ERROR("host-orch: failed to retain Graph execution storage"); return false; } submission->definition_addr = reinterpret_cast(object_it->second.device_object); @@ -515,19 +607,41 @@ bool upload_graph_submissions( submission->local_execution = 0; submission->activation_gate = 0; - void *device_submission = api->device_malloc(upload->bytes); + // Retained runner-owned POD storage keyed by (graph_key, occurrence): + // reused across runs while capacity fits, released at Worker + // finalization — so it must not enter tensor_pairs_, which validate + // frees every round. + void *device_submission = api->acquire_graph_submission_buffer( + submission->graph_key, occurrence, upload->bytes, alignof(GraphSubmission) + ); if (device_submission == nullptr) { - LOG_ERROR("host-orch: failed to allocate %zu bytes for Graph submission", upload->bytes); + LOG_ERROR("host-orch: failed to retain Graph submission"); return false; } if (api->copy_to_device(device_submission, upload->data, upload->bytes) != 0) { LOG_ERROR("host-orch: failed to upload Graph submission POD image"); - api->device_free(device_submission); return false; } + upload->outer_slot->graph_context = device_submission; - runtime->tensor_pairs_.push_back({nullptr, device_submission, upload->bytes, false}); - uploaded_bytes += static_cast(upload->bytes); + graph_host_mark_upload_h2d_done(graph_state, index); + return true; + } + + static bool eager_cb(void *ctx, GraphHostState &state, size_t index) { + return static_cast(ctx)->upload_one(state, index); + } +}; + +static bool upload_leftover_graph_submissions(GraphPodH2d &h2d, GraphHostState &graph_state) { + const size_t count = graph_host_upload_count(graph_state); + for (size_t index = 0; index < count; ++index) { + if (graph_host_upload_h2d_done(graph_state, index)) { + continue; + } + if (!h2d.upload_one(graph_state, index)) { + return false; + } } return true; } @@ -583,6 +697,22 @@ int32_t run_host_orchestration( } GraphHostStateBinding graph_binding(rt->orchestrator, graph_state.get()); + GraphPodH2d graph_h2d; + graph_h2d.api = api; + // Both hook and arena live on the run-owned GraphHostState, and the guard + // clears the hook on every exit path — an early return below must not + // leave the orchestrator able to call into graph_h2d after it dies. + struct GraphUploadScope { + GraphHostState &state; + ~GraphUploadScope() { graph_host_set_eager_upload(state, nullptr, nullptr); } + } graph_upload_scope{*graph_state}; + if (ensure_graph_pinned_pack(kGraphPinnedArenaBytes)) { + graph_host_set_pinned_arena(*graph_state, static_cast(g_graph_pack.host), g_graph_pack.cap); + } else { + LOG_WARN("host-orch: pinned Graph arena unavailable; POD H2D may stage"); + } + graph_host_set_eager_upload(*graph_state, &GraphPodH2d::eager_cb, &graph_h2d); + const int32_t block_dim = runtime->get_worker_count() / PLATFORM_CORES_PER_BLOCKDIM; if (block_dim < 1) { LOG_ERROR("host-orch: worker_count %d yields no clusters", runtime->get_worker_count()); @@ -612,6 +742,7 @@ int32_t run_host_orchestration( entry_points->entry(orch_l2); rt_scope_end(rt); rt_orchestration_done(rt); + graph_host_set_eager_upload(*graph_state, nullptr, nullptr); #if SIMPLER_ORCH_PROFILING // Per-sub-step cumulatives across this pass's submits. The accumulators only // exist in a SIMPLER_ORCH_PROFILING build (build_runtimes.py --profiling-orch 1), @@ -644,12 +775,11 @@ int32_t run_host_orchestration( // five markers, which must not be charged to the pass it measures. const int64_t t_graph_ns = bind_now_ns(); - uint64_t graph_bytes = 0; - if (!upload_graph_submissions(runtime, api, *graph_state, graph_bytes)) return -1; + if (!upload_leftover_graph_submissions(graph_h2d, *graph_state)) return -1; { char attrs[96]; - snprintf(attrs, sizeof(attrs), "count=%zu bytes=%" PRIu64, graph_host_upload_count(*graph_state), graph_bytes); - record_bind_phase(HostPhaseKind::BindGraphUpload, t_graph_ns, attrs, graph_bytes); + snprintf(attrs, sizeof(attrs), "count=%zu", graph_host_upload_count(*graph_state)); + record_bind_phase(HostPhaseKind::BindGraphUpload, t_graph_ns, attrs); } // total_tasks sizes the bounded per-segment H2D copies below; a value outside diff --git a/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h b/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h index 77598e3188..c0dffae60e 100644 --- a/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h +++ b/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h @@ -420,8 +420,11 @@ static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const C invoke(); if (!rt_graph_end()) invoke(); } else if (result.execute_block) { - // Un-cacheable at begin, or the Definition cache is full: ordinary path. - invoke(); + // Un-cacheable at begin, the Definition cache is full, or the runtime + // went fatal (e.g. the eager Graph POD upload failed after the outer + // task was published): ordinary path, except that a fatal runtime must + // not re-run the body — its run is already doomed. + if (!current_runtime()->ops->is_fatal(current_runtime())) invoke(); } // Cache hit: execute_block and recording are both false; the body is skipped. rt_graph_commit(); diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp index 0fe0181e5f..d078a6ba1d 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp @@ -31,6 +31,7 @@ #include #include #include +#include #include #include #include @@ -337,14 +338,47 @@ struct GraphRecording { struct GraphPendingUpload { PTO2TaskSlotState *outer_slot{nullptr}; std::vector image; + std::byte *pinned{nullptr}; + size_t bytes{0}; + bool h2d_done{false}; }; struct GraphHostState { std::unordered_map> definitions; std::unique_ptr recording; std::vector pending_uploads; + // Run-scoped eager-upload hook and pinned bump arena: registered per + // run_host_orchestration, so concurrent orchestrations each own their own + // uploader and bump cursor. + GraphHostEagerUploadFn eager_upload_fn{nullptr}; + void *eager_upload_ctx{nullptr}; + std::byte *pin_base{nullptr}; + size_t pin_cap{0}; + size_t pin_used{0}; }; +namespace { +constexpr size_t kPinnedBumpAlign = 64; + +std::byte *graph_host_pinned_bump(GraphHostState &state, size_t bytes) { + if (state.pin_base == nullptr || bytes == 0) return nullptr; + const size_t off = (state.pin_used + kPinnedBumpAlign - 1) & ~(kPinnedBumpAlign - 1); + if (off + bytes > state.pin_cap) return nullptr; + state.pin_used = off + bytes; + return state.pin_base + off; +} +} // namespace + +void graph_host_set_pinned_arena(GraphHostState &state, std::byte *base, size_t cap) { + state.pin_base = base; + state.pin_cap = cap; + state.pin_used = 0; +} + +std::byte *graph_host_pinned_base(const GraphHostState &state) { return state.pin_base; } + +size_t graph_host_pinned_used(const GraphHostState &state) { return state.pin_used; } + namespace { GraphHostState *graph_state_from(PTO2OrchestratorState *orch) { @@ -682,7 +716,11 @@ size_t graph_host_upload_count(const GraphHostState &state) { return state.pendi std::optional graph_host_upload(GraphHostState &state, size_t index) { if (index >= state.pending_uploads.size()) return std::nullopt; GraphPendingUpload &upload = state.pending_uploads[index]; - if (upload.outer_slot == nullptr || upload.image.empty()) return std::nullopt; + if (upload.outer_slot == nullptr) return std::nullopt; + if (upload.pinned != nullptr && upload.bytes > 0) { + return GraphHostUpload{upload.outer_slot, upload.pinned, upload.bytes}; + } + if (upload.image.empty()) return std::nullopt; return GraphHostUpload{upload.outer_slot, upload.image.data(), upload.image.size()}; } @@ -698,6 +736,19 @@ GraphHostDefinitionList graph_host_definitions(GraphHostState &state) { return list; } +void graph_host_set_eager_upload(GraphHostState &state, GraphHostEagerUploadFn fn, void *ctx) { + state.eager_upload_fn = fn; + state.eager_upload_ctx = ctx; +} + +bool graph_host_upload_h2d_done(const GraphHostState &state, size_t index) { + return index < state.pending_uploads.size() && state.pending_uploads[index].h2d_done; +} + +void graph_host_mark_upload_h2d_done(GraphHostState &state, size_t index) { + if (index < state.pending_uploads.size()) state.pending_uploads[index].h2d_done = true; +} + static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { uint32_t next = orch->fanin_seen_current_epoch + 1; if (next == 0) { @@ -1412,47 +1463,86 @@ void graph_reset_outer_payload(PTO2TaskPayload &payload) { payload.early_sync_drain_state.store(PTO2_EARLY_SYNC_DRAIN_NONE, std::memory_order_relaxed); } -bool graph_build_submission_image( - const std::vector &definition_image, const CoreTaskArgs &args, std::vector *submission_image +// Output pointers the caller does not need may be null. +static bool graph_submission_layout( + const std::vector &definition_image, const CoreTaskArgs &args, size_t *total_bytes, + size_t *tensors_offset, size_t *scalars_offset, size_t *scalar_bytes ) { - if (submission_image == nullptr || graph_definition(definition_image) == nullptr) return false; - const size_t tensors_offset = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphTensor)); + if (total_bytes == nullptr || graph_definition(definition_image) == nullptr) { + return false; + } + const size_t tensors = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphTensor)); const size_t tensor_bytes = static_cast(args.tensor_count()) * sizeof(GraphTensor); - if (tensors_offset > UINT32_MAX || tensors_offset > UINT32_MAX - tensor_bytes) { + if (tensors > UINT32_MAX || tensors > UINT32_MAX - tensor_bytes) { return false; } - const size_t tensors_end = tensors_offset + tensor_bytes; - const size_t scalar_bytes = static_cast(args.scalar_count()) * sizeof(uint64_t); - const size_t scalars_offset = args.scalar_count() == 0 ? 0 : PTO2_ALIGN_UP(tensors_end, alignof(uint64_t)); - const size_t total_bytes = args.scalar_count() == 0 ? tensors_end : scalars_offset + scalar_bytes; - if ((args.scalar_count() != 0 && (scalars_offset > UINT32_MAX || scalars_offset > UINT32_MAX - scalar_bytes)) || - total_bytes > UINT32_MAX) { + const size_t tensors_end = tensors + tensor_bytes; + const size_t bytes_of_scalars = static_cast(args.scalar_count()) * sizeof(uint64_t); + const size_t scalars = args.scalar_count() == 0 ? 0 : PTO2_ALIGN_UP(tensors_end, alignof(uint64_t)); + const size_t total = args.scalar_count() == 0 ? tensors_end : scalars + bytes_of_scalars; + if ((args.scalar_count() != 0 && (scalars > UINT32_MAX || scalars > UINT32_MAX - bytes_of_scalars)) || + total > UINT32_MAX) { return false; } - submission_image->assign(total_bytes, std::byte{0}); - auto *tensors = reinterpret_cast(submission_image->data() + tensors_offset); + if (tensors_offset != nullptr) *tensors_offset = tensors; + if (scalars_offset != nullptr) *scalars_offset = scalars; + if (scalar_bytes != nullptr) *scalar_bytes = bytes_of_scalars; + *total_bytes = total; + return true; +} + +static bool graph_fill_submission_image( + const std::vector &definition_image, const CoreTaskArgs &args, std::byte *dst, size_t dst_bytes +) { + size_t total_bytes = 0; + size_t tensors_offset = 0; + size_t scalars_offset = 0; + size_t scalar_bytes = 0; + if (dst == nullptr || + !graph_submission_layout( + definition_image, args, &total_bytes, &tensors_offset, &scalars_offset, &scalar_bytes + ) || + dst_bytes < total_bytes) { + return false; + } + std::memset(dst, 0, total_bytes); + auto *tensors = reinterpret_cast(dst + tensors_offset); for (int32_t i = 0; i < args.tensor_count(); ++i) tensors[i] = graph_tensor_pack(args.tensor(i).ref()); if (args.scalar_count() != 0) { - std::memcpy( - submission_image->data() + scalars_offset, args.scalar_data(), - static_cast(args.scalar_count()) * sizeof(uint64_t) - ); + std::memcpy(dst + scalars_offset, args.scalar_data(), scalar_bytes); } const GraphDefinition &definition = *graph_definition(definition_image); GraphSubmission submission{}; submission.graph_key = definition.full_key; submission.definition_hash = definition.content_hash; - submission.total_bytes = static_cast(submission_image->size()); + submission.total_bytes = static_cast(total_bytes); submission.tensors_offset = static_cast(tensors_offset); submission.tensor_count = static_cast(args.tensor_count()); submission.scalars_offset = static_cast(scalars_offset); submission.scalar_count = static_cast(args.scalar_count()); - std::memcpy(submission_image->data(), &submission, sizeof(submission)); + std::memcpy(dst, &submission, sizeof(submission)); return true; } +bool graph_build_submission_image( + const std::vector &definition_image, const CoreTaskArgs &args, std::vector *submission_image +) { + size_t total_bytes = 0; + size_t tensors_offset = 0; + size_t scalars_offset = 0; + size_t scalar_bytes = 0; + if (submission_image == nullptr || + !graph_submission_layout( + definition_image, args, &total_bytes, &tensors_offset, &scalars_offset, &scalar_bytes + )) { + return false; + } + submission_image->assign(total_bytes, std::byte{0}); + return graph_fill_submission_image(definition_image, args, submission_image->data(), submission_image->size()); +} + bool graph_submit_definition( PTO2OrchestratorState *orch, GraphHostState *state, const std::vector &definition_image, const CoreTaskArgs &args, PTO2TaskId *submitted_id @@ -1470,7 +1560,20 @@ bool graph_submit_definition( } GraphPendingUpload pending; - if (!graph_build_submission_image(definition_image, args, &pending.image)) return false; + size_t total_bytes = 0; + if (!graph_submission_layout(definition_image, args, &total_bytes, nullptr, nullptr, nullptr)) { + return false; + } + std::byte *pinned = graph_host_pinned_bump(*state, total_bytes); + if (pinned != nullptr) { + if (!graph_fill_submission_image(definition_image, args, pinned, total_bytes)) return false; + pending.pinned = pinned; + pending.bytes = total_bytes; + } else if (!graph_build_submission_image(definition_image, args, &pending.image)) { + return false; + } else { + pending.bytes = pending.image.size(); + } DepInputs boundary_inputs{ args.tensor_count(), args.tensor_data(), args.tag_data(), 0, nullptr, @@ -1521,6 +1624,21 @@ bool graph_submit_definition( pending.outer_slot = &slot; state->pending_uploads.push_back(std::move(pending)); + if (state->eager_upload_fn != nullptr) { + const size_t index = state->pending_uploads.size() - 1; + if (!state->eager_upload_fn(state->eager_upload_ctx, *state, index)) { + // The outer GRAPH task is already published (slot allocated, fanins + // wired, outputs registered), so the ordinary-path fallback the + // false return otherwise means would re-submit the body on top of a + // half-uploaded Graph task. Latch fatal so graph_begin aborts and + // the run fails loudly instead. + orch->report_fatal( + PTO2_ERROR_EXPLICIT_ORCH_FATAL, __FUNCTION__, "eager Graph POD H2D failed for key=%#llx", + static_cast(definition->full_key) + ); + return false; + } + } if (submitted_id != nullptr) *submitted_id = task_id; #if SIMPLER_DFX orch->tasks_submitted++; 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 ef0bad51e5..217adcafe1 100644 --- a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp @@ -461,27 +461,95 @@ static bool relocate_host_orch_image( return ok; } -bool upload_graph_submissions( - Runtime *runtime, const HostApi *api, GraphHostState &graph_state, uint64_t &uploaded_bytes -) { +// Retained pinned bump arena for Graph POD images. Orch writes each layer +// in place; upload_one synchronously H2Ds that POD (source is already pinned). +// The CANN entry points are resolved with dlopen so the same source builds on +// hosts without an Ascend toolkit (sim runners, CI); there the arena falls +// back to plain memory, which costs nothing because sim copies are memcpy. +struct GraphPinnedPack { + void *host = nullptr; + size_t cap = 0; + bool pinned = false; +}; + +static GraphPinnedPack g_graph_pack; +static std::mutex g_graph_pack_mu; +constexpr size_t kGraphPinnedArenaBytes = 16ull * 1024ull * 1024ull; + +namespace { +using AclrtMallocHostFn = int (*)(void **, size_t); +using AclrtFreeHostFn = int (*)(void *); + +AclrtMallocHostFn graph_aclrt_malloc_host() { + void *handle = dlopen("libascendcl.so", RTLD_NOW | RTLD_LOCAL); + if (handle == nullptr) return nullptr; + return reinterpret_cast(dlsym(handle, "aclrtMallocHost")); +} + +AclrtFreeHostFn graph_aclrt_free_host() { + void *handle = dlopen("libascendcl.so", RTLD_NOW | RTLD_LOCAL); + if (handle == nullptr) return nullptr; + return reinterpret_cast(dlsym(handle, "aclrtFreeHost")); +} +} // namespace + +static void free_graph_pack(GraphPinnedPack &pack) { + if (pack.host == nullptr) return; + if (pack.pinned) { + AclrtFreeHostFn free_host = graph_aclrt_free_host(); + if (free_host != nullptr) (void)free_host(pack.host); + } else { + delete[] static_cast(pack.host); + } + pack.host = nullptr; + pack.cap = 0; + pack.pinned = false; +} + +static bool ensure_graph_pinned_pack(size_t cap) { + std::lock_guard lock(g_graph_pack_mu); + if (g_graph_pack.cap >= cap && g_graph_pack.host != nullptr) { + return true; + } + free_graph_pack(g_graph_pack); + AclrtMallocHostFn malloc_host = graph_aclrt_malloc_host(); + if (malloc_host != nullptr) { + void *host = nullptr; + if (malloc_host(&host, cap) != 0 || host == nullptr) { + LOG_ERROR("host-orch: aclrtMallocHost(%zu) failed", cap); + return false; + } + std::memset(host, 0, cap); + g_graph_pack.host = host; + g_graph_pack.pinned = true; + } else { + g_graph_pack.host = new std::byte[cap]{}; + g_graph_pack.pinned = false; + } + g_graph_pack.cap = cap; + return true; +} + +struct GraphPodH2d { + const HostApi *api = nullptr; std::unordered_map occurrences; - uploaded_bytes = 0; - const size_t count = graph_host_upload_count(graph_state); - // Pass 1: upload each distinct Definition once as a shared device object - // ([GraphDefinitionHeader][Definition image]) keyed by content identity. - // Submissions reference the object's GM address, so this pass completing - // before any submission is uploaded is what makes the reference safe — - // the device boots only after both passes. - GraphHostDefinitionList definitions = graph_host_definitions(graph_state); + struct UploadedDefinition { void *device_object; // GM address; host must not dereference const GraphDefinition *host_view; // the host-side image the object was built from }; std::unordered_map definition_objects; - for (const GraphHostDefinition &entry : definitions.entries) { - if (entry.data == nullptr || entry.bytes < sizeof(GraphDefinition)) continue; + + // Upload one distinct Definition as a shared device object + // ([GraphDefinitionHeader][Definition image]) keyed by content identity. + // Submissions reference the object's GM address, so the object existing + // before any submission referencing it is uploaded is what makes the + // reference safe — the device boots only after both are done. + bool ensure_definition_object(const GraphHostDefinition &entry) { + if (entry.data == nullptr || entry.bytes < sizeof(GraphDefinition)) return false; const auto *definition = reinterpret_cast(entry.data); - if (definition->total_bytes != entry.bytes || definition->full_key != entry.full_key) continue; + if (definition->total_bytes != entry.bytes || definition->full_key != entry.full_key) return false; + if (definition_objects.count(definition->content_hash) != 0) return true; const size_t object_bytes = sizeof(GraphDefinitionHeader) + entry.bytes; void *object = api->acquire_graph_definition_buffer(entry.full_key, object_bytes, alignof(GraphDefinitionHeader)); @@ -507,11 +575,27 @@ bool upload_graph_submissions( return false; } definition_objects.emplace(definition->content_hash, UploadedDefinition{object, definition}); - uploaded_bytes += object_bytes; + return true; + } + + // Returns the Definition entry a submission references, or nullopt when the + // host state holds no valid Definition for it. + std::optional find_definition(const GraphHostState &graph_state, uint64_t content_hash) { + GraphHostDefinitionList definitions = graph_host_definitions(const_cast(graph_state)); + for (const GraphHostDefinition &entry : definitions.entries) { + if (entry.bytes < sizeof(GraphDefinition) || entry.data == nullptr) continue; + const auto *definition = reinterpret_cast(entry.data); + if (definition->content_hash == content_hash && definition->total_bytes == entry.bytes) { + return entry; + } + } + return std::nullopt; } - // Pass 2: per-submission execution storage + the small reference image. - for (size_t index = 0; index < count; ++index) { + bool upload_one(GraphHostState &graph_state, size_t index) { + if (graph_host_upload_h2d_done(graph_state, index)) { + return true; + } std::optional upload = graph_host_upload(graph_state, index); if (!upload.has_value() || upload->outer_slot == nullptr || upload->data == nullptr || upload->bytes < sizeof(GraphSubmission) || upload->outer_slot->task_kind != TaskKind::GRAPH || @@ -525,9 +609,19 @@ bool upload_graph_submissions( return false; } auto object_it = definition_objects.find(submission->definition_hash); - if (object_it == definition_objects.end() || object_it->second.device_object == nullptr) { - LOG_ERROR("host-orch: Graph submission has no uploaded Definition object"); - return false; + if (object_it == definition_objects.end()) { + // Eager uploads run during orch entry, before any batched pass, so + // the Definition object for this submission may not exist yet. + std::optional entry = find_definition(graph_state, submission->definition_hash); + if (!entry.has_value() || !ensure_definition_object(*entry)) { + LOG_ERROR("host-orch: Graph submission has no uploadable Definition"); + return false; + } + object_it = definition_objects.find(submission->definition_hash); + if (object_it == definition_objects.end()) { + LOG_ERROR("host-orch: Graph submission has no uploaded Definition object"); + return false; + } } // Capacities come from the host-side Definition image the device // object was built from; the GM object itself is never dereferenced @@ -548,10 +642,7 @@ bool upload_graph_submissions( submission->graph_key, occurrence, execution_bytes, alignof(GraphNodeStorage) ); if (execution_storage == nullptr) { - LOG_ERROR( - "host-orch: failed to retain %zu bytes for Graph execution key=%#llx occurrence=%u", execution_bytes, - static_cast(submission->graph_key), occurrence - ); + LOG_ERROR("host-orch: failed to retain Graph execution storage"); return false; } submission->definition_addr = reinterpret_cast(object_it->second.device_object); @@ -560,19 +651,41 @@ bool upload_graph_submissions( submission->local_execution = 0; submission->activation_gate = 0; - void *device_submission = api->device_malloc(upload->bytes); + // Retained runner-owned POD storage keyed by (graph_key, occurrence): + // reused across runs while capacity fits, released at Worker + // finalization — so it must not enter tensor_pairs_, which validate + // frees every round. + void *device_submission = api->acquire_graph_submission_buffer( + submission->graph_key, occurrence, upload->bytes, alignof(GraphSubmission) + ); if (device_submission == nullptr) { - LOG_ERROR("host-orch: failed to allocate %zu bytes for Graph submission", upload->bytes); + LOG_ERROR("host-orch: failed to retain Graph submission"); return false; } if (api->copy_to_device(device_submission, upload->data, upload->bytes) != 0) { LOG_ERROR("host-orch: failed to upload Graph submission POD image"); - api->device_free(device_submission); return false; } + upload->outer_slot->graph_context = device_submission; - runtime->tensor_pairs_.push_back({nullptr, device_submission, upload->bytes, false}); - uploaded_bytes += static_cast(upload->bytes); + graph_host_mark_upload_h2d_done(graph_state, index); + return true; + } + + static bool eager_cb(void *ctx, GraphHostState &state, size_t index) { + return static_cast(ctx)->upload_one(state, index); + } +}; + +static bool upload_leftover_graph_submissions(GraphPodH2d &h2d, GraphHostState &graph_state) { + const size_t count = graph_host_upload_count(graph_state); + for (size_t index = 0; index < count; ++index) { + if (graph_host_upload_h2d_done(graph_state, index)) { + continue; + } + if (!h2d.upload_one(graph_state, index)) { + return false; + } } return true; } @@ -630,6 +743,22 @@ int32_t run_host_orchestration( } GraphHostStateBinding graph_binding(rt->orchestrator, graph_state.get()); + GraphPodH2d graph_h2d; + graph_h2d.api = api; + // Both hook and arena live on the run-owned GraphHostState, and the guard + // clears the hook on every exit path — an early return below must not + // leave the orchestrator able to call into graph_h2d after it dies. + struct GraphUploadScope { + GraphHostState &state; + ~GraphUploadScope() { graph_host_set_eager_upload(state, nullptr, nullptr); } + } graph_upload_scope{*graph_state}; + if (ensure_graph_pinned_pack(kGraphPinnedArenaBytes)) { + graph_host_set_pinned_arena(*graph_state, static_cast(g_graph_pack.host), g_graph_pack.cap); + } else { + LOG_WARN("host-orch: pinned Graph arena unavailable; POD H2D may stage"); + } + graph_host_set_eager_upload(*graph_state, &GraphPodH2d::eager_cb, &graph_h2d); + // Install the ops table (host s_runtime_ops) and latch this run's cluster // counts. worker_count is published by DeviceRunner::prepare_launch_shape // before this bind, so the host orchestrator sees the same geometry the @@ -668,6 +797,7 @@ int32_t run_host_orchestration( entry_points->entry(orch_l2); rt_scope_end(rt); rt_orchestration_done(rt); + graph_host_set_eager_upload(*graph_state, nullptr, nullptr); #if SIMPLER_ORCH_PROFILING // Per-sub-step cumulatives across this pass's submits. The accumulators only // exist in a SIMPLER_ORCH_PROFILING build (build_runtimes.py --profiling-orch 1), @@ -700,12 +830,11 @@ int32_t run_host_orchestration( // five markers, which must not be charged to the pass it measures. const int64_t t_graph_ns = bind_now_ns(); - uint64_t graph_bytes = 0; - if (!upload_graph_submissions(runtime, api, *graph_state, graph_bytes)) return -1; + if (!upload_leftover_graph_submissions(graph_h2d, *graph_state)) return -1; { char attrs[96]; - snprintf(attrs, sizeof(attrs), "count=%zu bytes=%" PRIu64, graph_host_upload_count(*graph_state), graph_bytes); - record_bind_phase(HostPhaseKind::BindGraphUpload, t_graph_ns, attrs, graph_bytes); + snprintf(attrs, sizeof(attrs), "count=%zu", graph_host_upload_count(*graph_state)); + record_bind_phase(HostPhaseKind::BindGraphUpload, t_graph_ns, attrs); } // total_tasks sizes the bounded per-segment H2D copies below; a value outside 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 77598e3188..c0dffae60e 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 @@ -420,8 +420,11 @@ static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const C invoke(); if (!rt_graph_end()) invoke(); } else if (result.execute_block) { - // Un-cacheable at begin, or the Definition cache is full: ordinary path. - invoke(); + // Un-cacheable at begin, the Definition cache is full, or the runtime + // went fatal (e.g. the eager Graph POD upload failed after the outer + // task was published): ordinary path, except that a fatal runtime must + // not re-run the body — its run is already doomed. + if (!current_runtime()->ops->is_fatal(current_runtime())) invoke(); } // Cache hit: execute_block and recording are both false; the body is skipped. rt_graph_commit(); 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 0fe0181e5f..d078a6ba1d 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 @@ -31,6 +31,7 @@ #include #include #include +#include #include #include #include @@ -337,14 +338,47 @@ struct GraphRecording { struct GraphPendingUpload { PTO2TaskSlotState *outer_slot{nullptr}; std::vector image; + std::byte *pinned{nullptr}; + size_t bytes{0}; + bool h2d_done{false}; }; struct GraphHostState { std::unordered_map> definitions; std::unique_ptr recording; std::vector pending_uploads; + // Run-scoped eager-upload hook and pinned bump arena: registered per + // run_host_orchestration, so concurrent orchestrations each own their own + // uploader and bump cursor. + GraphHostEagerUploadFn eager_upload_fn{nullptr}; + void *eager_upload_ctx{nullptr}; + std::byte *pin_base{nullptr}; + size_t pin_cap{0}; + size_t pin_used{0}; }; +namespace { +constexpr size_t kPinnedBumpAlign = 64; + +std::byte *graph_host_pinned_bump(GraphHostState &state, size_t bytes) { + if (state.pin_base == nullptr || bytes == 0) return nullptr; + const size_t off = (state.pin_used + kPinnedBumpAlign - 1) & ~(kPinnedBumpAlign - 1); + if (off + bytes > state.pin_cap) return nullptr; + state.pin_used = off + bytes; + return state.pin_base + off; +} +} // namespace + +void graph_host_set_pinned_arena(GraphHostState &state, std::byte *base, size_t cap) { + state.pin_base = base; + state.pin_cap = cap; + state.pin_used = 0; +} + +std::byte *graph_host_pinned_base(const GraphHostState &state) { return state.pin_base; } + +size_t graph_host_pinned_used(const GraphHostState &state) { return state.pin_used; } + namespace { GraphHostState *graph_state_from(PTO2OrchestratorState *orch) { @@ -682,7 +716,11 @@ size_t graph_host_upload_count(const GraphHostState &state) { return state.pendi std::optional graph_host_upload(GraphHostState &state, size_t index) { if (index >= state.pending_uploads.size()) return std::nullopt; GraphPendingUpload &upload = state.pending_uploads[index]; - if (upload.outer_slot == nullptr || upload.image.empty()) return std::nullopt; + if (upload.outer_slot == nullptr) return std::nullopt; + if (upload.pinned != nullptr && upload.bytes > 0) { + return GraphHostUpload{upload.outer_slot, upload.pinned, upload.bytes}; + } + if (upload.image.empty()) return std::nullopt; return GraphHostUpload{upload.outer_slot, upload.image.data(), upload.image.size()}; } @@ -698,6 +736,19 @@ GraphHostDefinitionList graph_host_definitions(GraphHostState &state) { return list; } +void graph_host_set_eager_upload(GraphHostState &state, GraphHostEagerUploadFn fn, void *ctx) { + state.eager_upload_fn = fn; + state.eager_upload_ctx = ctx; +} + +bool graph_host_upload_h2d_done(const GraphHostState &state, size_t index) { + return index < state.pending_uploads.size() && state.pending_uploads[index].h2d_done; +} + +void graph_host_mark_upload_h2d_done(GraphHostState &state, size_t index) { + if (index < state.pending_uploads.size()) state.pending_uploads[index].h2d_done = true; +} + static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { uint32_t next = orch->fanin_seen_current_epoch + 1; if (next == 0) { @@ -1412,47 +1463,86 @@ void graph_reset_outer_payload(PTO2TaskPayload &payload) { payload.early_sync_drain_state.store(PTO2_EARLY_SYNC_DRAIN_NONE, std::memory_order_relaxed); } -bool graph_build_submission_image( - const std::vector &definition_image, const CoreTaskArgs &args, std::vector *submission_image +// Output pointers the caller does not need may be null. +static bool graph_submission_layout( + const std::vector &definition_image, const CoreTaskArgs &args, size_t *total_bytes, + size_t *tensors_offset, size_t *scalars_offset, size_t *scalar_bytes ) { - if (submission_image == nullptr || graph_definition(definition_image) == nullptr) return false; - const size_t tensors_offset = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphTensor)); + if (total_bytes == nullptr || graph_definition(definition_image) == nullptr) { + return false; + } + const size_t tensors = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphTensor)); const size_t tensor_bytes = static_cast(args.tensor_count()) * sizeof(GraphTensor); - if (tensors_offset > UINT32_MAX || tensors_offset > UINT32_MAX - tensor_bytes) { + if (tensors > UINT32_MAX || tensors > UINT32_MAX - tensor_bytes) { return false; } - const size_t tensors_end = tensors_offset + tensor_bytes; - const size_t scalar_bytes = static_cast(args.scalar_count()) * sizeof(uint64_t); - const size_t scalars_offset = args.scalar_count() == 0 ? 0 : PTO2_ALIGN_UP(tensors_end, alignof(uint64_t)); - const size_t total_bytes = args.scalar_count() == 0 ? tensors_end : scalars_offset + scalar_bytes; - if ((args.scalar_count() != 0 && (scalars_offset > UINT32_MAX || scalars_offset > UINT32_MAX - scalar_bytes)) || - total_bytes > UINT32_MAX) { + const size_t tensors_end = tensors + tensor_bytes; + const size_t bytes_of_scalars = static_cast(args.scalar_count()) * sizeof(uint64_t); + const size_t scalars = args.scalar_count() == 0 ? 0 : PTO2_ALIGN_UP(tensors_end, alignof(uint64_t)); + const size_t total = args.scalar_count() == 0 ? tensors_end : scalars + bytes_of_scalars; + if ((args.scalar_count() != 0 && (scalars > UINT32_MAX || scalars > UINT32_MAX - bytes_of_scalars)) || + total > UINT32_MAX) { return false; } - submission_image->assign(total_bytes, std::byte{0}); - auto *tensors = reinterpret_cast(submission_image->data() + tensors_offset); + if (tensors_offset != nullptr) *tensors_offset = tensors; + if (scalars_offset != nullptr) *scalars_offset = scalars; + if (scalar_bytes != nullptr) *scalar_bytes = bytes_of_scalars; + *total_bytes = total; + return true; +} + +static bool graph_fill_submission_image( + const std::vector &definition_image, const CoreTaskArgs &args, std::byte *dst, size_t dst_bytes +) { + size_t total_bytes = 0; + size_t tensors_offset = 0; + size_t scalars_offset = 0; + size_t scalar_bytes = 0; + if (dst == nullptr || + !graph_submission_layout( + definition_image, args, &total_bytes, &tensors_offset, &scalars_offset, &scalar_bytes + ) || + dst_bytes < total_bytes) { + return false; + } + std::memset(dst, 0, total_bytes); + auto *tensors = reinterpret_cast(dst + tensors_offset); for (int32_t i = 0; i < args.tensor_count(); ++i) tensors[i] = graph_tensor_pack(args.tensor(i).ref()); if (args.scalar_count() != 0) { - std::memcpy( - submission_image->data() + scalars_offset, args.scalar_data(), - static_cast(args.scalar_count()) * sizeof(uint64_t) - ); + std::memcpy(dst + scalars_offset, args.scalar_data(), scalar_bytes); } const GraphDefinition &definition = *graph_definition(definition_image); GraphSubmission submission{}; submission.graph_key = definition.full_key; submission.definition_hash = definition.content_hash; - submission.total_bytes = static_cast(submission_image->size()); + submission.total_bytes = static_cast(total_bytes); submission.tensors_offset = static_cast(tensors_offset); submission.tensor_count = static_cast(args.tensor_count()); submission.scalars_offset = static_cast(scalars_offset); submission.scalar_count = static_cast(args.scalar_count()); - std::memcpy(submission_image->data(), &submission, sizeof(submission)); + std::memcpy(dst, &submission, sizeof(submission)); return true; } +bool graph_build_submission_image( + const std::vector &definition_image, const CoreTaskArgs &args, std::vector *submission_image +) { + size_t total_bytes = 0; + size_t tensors_offset = 0; + size_t scalars_offset = 0; + size_t scalar_bytes = 0; + if (submission_image == nullptr || + !graph_submission_layout( + definition_image, args, &total_bytes, &tensors_offset, &scalars_offset, &scalar_bytes + )) { + return false; + } + submission_image->assign(total_bytes, std::byte{0}); + return graph_fill_submission_image(definition_image, args, submission_image->data(), submission_image->size()); +} + bool graph_submit_definition( PTO2OrchestratorState *orch, GraphHostState *state, const std::vector &definition_image, const CoreTaskArgs &args, PTO2TaskId *submitted_id @@ -1470,7 +1560,20 @@ bool graph_submit_definition( } GraphPendingUpload pending; - if (!graph_build_submission_image(definition_image, args, &pending.image)) return false; + size_t total_bytes = 0; + if (!graph_submission_layout(definition_image, args, &total_bytes, nullptr, nullptr, nullptr)) { + return false; + } + std::byte *pinned = graph_host_pinned_bump(*state, total_bytes); + if (pinned != nullptr) { + if (!graph_fill_submission_image(definition_image, args, pinned, total_bytes)) return false; + pending.pinned = pinned; + pending.bytes = total_bytes; + } else if (!graph_build_submission_image(definition_image, args, &pending.image)) { + return false; + } else { + pending.bytes = pending.image.size(); + } DepInputs boundary_inputs{ args.tensor_count(), args.tensor_data(), args.tag_data(), 0, nullptr, @@ -1521,6 +1624,21 @@ bool graph_submit_definition( pending.outer_slot = &slot; state->pending_uploads.push_back(std::move(pending)); + if (state->eager_upload_fn != nullptr) { + const size_t index = state->pending_uploads.size() - 1; + if (!state->eager_upload_fn(state->eager_upload_ctx, *state, index)) { + // The outer GRAPH task is already published (slot allocated, fanins + // wired, outputs registered), so the ordinary-path fallback the + // false return otherwise means would re-submit the body on top of a + // half-uploaded Graph task. Latch fatal so graph_begin aborts and + // the run fails loudly instead. + orch->report_fatal( + PTO2_ERROR_EXPLICIT_ORCH_FATAL, __FUNCTION__, "eager Graph POD H2D failed for key=%#llx", + static_cast(definition->full_key) + ); + return false; + } + } if (submitted_id != nullptr) *submitted_id = task_id; #if SIMPLER_DFX orch->tasks_submitted++; diff --git a/src/common/host_build_graph/graph_host_state.h b/src/common/host_build_graph/graph_host_state.h index 265f70907c..b3f3c032d8 100644 --- a/src/common/host_build_graph/graph_host_state.h +++ b/src/common/host_build_graph/graph_host_state.h @@ -34,6 +34,10 @@ struct GraphHostUpload { size_t bytes; }; +// Eager H2D hook signature: invoked right after each Graph POD image is +// appended during orch entry (compute-one-layer, copy-that-layer). +using GraphHostEagerUploadFn = bool (*)(void *ctx, GraphHostState &state, size_t index); + // The run's distinct Definition images (already deduplicated by the host-side // Definition cache), for upload as shared device objects ahead of submissions. struct GraphHostDefinition { @@ -50,3 +54,18 @@ GraphHostStatePtr make_graph_host_state(); size_t graph_host_upload_count(const GraphHostState &state); std::optional graph_host_upload(GraphHostState &state, size_t index); GraphHostDefinitionList graph_host_definitions(GraphHostState &state); + +// Optional eager H2D hook: invoked right after each Graph POD image is appended +// during orch entry (compute-one-layer, copy-that-layer). Registered on the +// run-owned GraphHostState (not process globals) so concurrent orchestrations +// cannot clobber each other's uploader; cleared before the state dies. +void graph_host_set_eager_upload(GraphHostState &state, GraphHostEagerUploadFn fn, void *ctx); +bool graph_host_upload_h2d_done(const GraphHostState &state, size_t index); +void graph_host_mark_upload_h2d_done(GraphHostState &state, size_t index); + +// Optional pinned bump arena for Graph POD images. Attached to the run-owned +// GraphHostState before orch entry so graph_submit_definition can write PODs in +// place. Not attached → fallback std::vector images. +void graph_host_set_pinned_arena(GraphHostState &state, std::byte *base, size_t cap); +std::byte *graph_host_pinned_base(const GraphHostState &state); +size_t graph_host_pinned_used(const GraphHostState &state); diff --git a/src/common/platform/include/common/host_api.h b/src/common/platform/include/common/host_api.h index 2a3b35e977..57513ea234 100644 --- a/src/common/platform/include/common/host_api.h +++ b/src/common/platform/include/common/host_api.h @@ -72,6 +72,15 @@ struct HostApiOps { void *(*acquire_graph_definition_buffer)( void *runner_ctx, uint32_t pipeline_slot, uint64_t key, size_t bytes, size_t alignment ); + // Runner-owned Graph submission POD storage, keyed by (graph_key, + // occurrence) exactly like acquire_graph_execution_buffer: one retained + // block per key, reused while capacity fits, all released at Worker + // finalization. Lets the eager per-layer H2D keep the device copy of the + // submission image alive across runs instead of per-run malloc/free. + void *(*acquire_graph_submission_buffer)( + void *runner_ctx, uint32_t pipeline_slot, uint64_t graph_key, uint32_t occurrence, size_t bytes, + size_t alignment + ); // Commit the three pooled regions (GM heap, runtime shared memory, and // prebuilt runtime arena) of the arena bank selected by this run, as three // independent device allocations. `runtime_arena_size == 0` skips the @@ -179,6 +188,13 @@ struct HostApi { if (ops_->acquire_graph_definition_buffer == nullptr) return nullptr; return ops_->acquire_graph_definition_buffer(runner_ctx_, pipeline_slot_, key, bytes, alignment); } + void * + acquire_graph_submission_buffer(uint64_t graph_key, uint32_t occurrence, size_t bytes, size_t alignment) const { + if (ops_->acquire_graph_submission_buffer == nullptr) return nullptr; + return ops_->acquire_graph_submission_buffer( + runner_ctx_, pipeline_slot_, graph_key, occurrence, bytes, alignment + ); + } int setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, size_t runtime_arena_size) const { return ops_->setup_static_arena(runner_ctx_, arena_bank_, gm_heap_size, gm_sm_size, runtime_arena_size); } diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index 4a169d7f2e..000b26c886 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -180,6 +180,18 @@ static void *acquire_graph_definition_buffer( } } +static void *acquire_graph_submission_buffer( + void *runner_ctx, uint32_t pipeline_slot, uint64_t graph_key, uint32_t occurrence, size_t bytes, size_t alignment +) { + if (runner_ctx == nullptr) return nullptr; + try { + return static_cast(runner_ctx) + ->acquire_graph_submission_buffer(pipeline_slot, graph_key, occurrence, bytes, alignment); + } catch (...) { + return nullptr; + } +} + static uint64_t upload_chip_callable_buffer_wrapper(void *runner_ctx, const void *callable) { if (runner_ctx == nullptr) return 0; try { @@ -295,6 +307,7 @@ static const HostApiOps g_host_api_ops = { .set_retained_temp_buffer = set_retained_temp_buffer, .acquire_graph_execution_buffer = acquire_graph_execution_buffer, .acquire_graph_definition_buffer = acquire_graph_definition_buffer, + .acquire_graph_submission_buffer = acquire_graph_submission_buffer, .setup_static_arena = setup_static_arena_wrapper, .acquire_pooled_gm_heap = acquire_pooled_gm_heap_wrapper, .acquire_pooled_gm_sm = acquire_pooled_gm_sm_wrapper, diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index e4c929ec36..fbbed57976 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -207,6 +207,42 @@ void *DeviceRunnerBase::acquire_graph_execution_buffer( return aligned_addr; } +void *DeviceRunnerBase::acquire_graph_submission_buffer( + uint32_t pipeline_slot, uint64_t graph_key, uint32_t occurrence, size_t bytes, size_t alignment +) { + if (pipeline_slot >= graph_submission_buffers_.size() || bytes == 0 || alignment == 0 || + (alignment & (alignment - 1)) != 0 || bytes > SIZE_MAX - (alignment - 1)) { + return nullptr; + } + std::vector &buffers = graph_submission_buffers_[pipeline_slot][graph_key]; + if (occurrence >= buffers.size()) buffers.resize(static_cast(occurrence) + 1); + RetainedGraphExecutionBuffer &buffer = buffers[occurrence]; + if (buffer.aligned_addr != nullptr && buffer.capacity >= bytes && + reinterpret_cast(buffer.aligned_addr) % alignment == 0) { + return buffer.aligned_addr; + } + + const size_t allocation_bytes = bytes + alignment - 1; + void *allocation = mem_alloc_.alloc(allocation_bytes); + if (allocation == nullptr) return nullptr; + const uintptr_t raw = reinterpret_cast(allocation); + if (raw > UINTPTR_MAX - (alignment - 1)) { + mem_alloc_.free(allocation); + return nullptr; + } + void *aligned_addr = reinterpret_cast((raw + alignment - 1) & ~(alignment - 1)); + if (device_memset(aligned_addr, 0, bytes) != 0) { + mem_alloc_.free(allocation); + return nullptr; + } + if (buffer.allocation != nullptr && mem_alloc_.free(buffer.allocation) != 0) { + mem_alloc_.free(allocation); + return nullptr; + } + buffer = RetainedGraphExecutionBuffer{allocation, aligned_addr, bytes}; + return aligned_addr; +} + void *DeviceRunnerBase::acquire_graph_definition_buffer( uint32_t pipeline_slot, uint64_t key, size_t bytes, size_t alignment ) { @@ -242,13 +278,19 @@ void *DeviceRunnerBase::acquire_graph_definition_buffer( } void DeviceRunnerBase::release_graph_execution_buffers() { - for (GraphExecutionBufferMap &by_key : graph_execution_buffers_) { + const auto release_occurrence_map = [](GraphExecutionBufferMap &by_key, MemoryAllocator &alloc) { for (auto &entry : by_key) { for (RetainedGraphExecutionBuffer &buffer : entry.second) { - if (buffer.allocation != nullptr) mem_alloc_.free(buffer.allocation); + if (buffer.allocation != nullptr) alloc.free(buffer.allocation); } } by_key.clear(); + }; + for (GraphExecutionBufferMap &by_key : graph_execution_buffers_) { + release_occurrence_map(by_key, mem_alloc_); + } + for (GraphExecutionBufferMap &by_key : graph_submission_buffers_) { + release_occurrence_map(by_key, mem_alloc_); } for (GraphDefinitionBufferMap &by_key : graph_definition_buffers_) { for (auto &entry : by_key) { diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index bb8d26b200..2fd7ed97e8 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -149,6 +149,9 @@ class DeviceRunnerBase { ); void * acquire_graph_definition_buffer(uint32_t pipeline_slot, uint64_t key, std::size_t bytes, std::size_t alignment); + void *acquire_graph_submission_buffer( + uint32_t pipeline_slot, uint64_t graph_key, uint32_t occurrence, std::size_t bytes, std::size_t alignment + ); void clear_temporary_buffer(); /** * Map a device buffer into the host address space and return a @@ -1083,6 +1086,9 @@ class DeviceRunnerBase { }; using GraphExecutionBufferMap = std::unordered_map>; std::array graph_execution_buffers_{}; + // Graph submission POD storage, one retained block per (pipeline slot, + // graph_key, occurrence) — see HostApi acquire_graph_submission_buffer. + std::array graph_submission_buffers_{}; // Graph Definition storage, one retained block per (pipeline slot, // definition key) — see HostApi acquire_graph_definition_buffer. Keyed by // content identity rather than occurrence: every submission of one run diff --git a/src/common/platform/sim/host/c_api_shared.cpp b/src/common/platform/sim/host/c_api_shared.cpp index ec4e6ef5d6..dd0f2952cc 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -166,6 +166,18 @@ static void *acquire_graph_definition_buffer( } } +static void *acquire_graph_submission_buffer( + void *runner_ctx, uint32_t pipeline_slot, uint64_t graph_key, uint32_t occurrence, size_t bytes, size_t alignment +) { + if (runner_ctx == nullptr) return nullptr; + try { + return static_cast(runner_ctx) + ->acquire_graph_submission_buffer(pipeline_slot, graph_key, occurrence, bytes, alignment); + } catch (...) { + return nullptr; + } +} + static uint64_t upload_chip_callable_buffer_wrapper(void *runner_ctx, const void *callable) { if (runner_ctx == nullptr) return 0; try { @@ -282,6 +294,7 @@ static const HostApiOps g_host_api_ops = { .set_retained_temp_buffer = set_retained_temp_buffer, .acquire_graph_execution_buffer = acquire_graph_execution_buffer, .acquire_graph_definition_buffer = acquire_graph_definition_buffer, + .acquire_graph_submission_buffer = acquire_graph_submission_buffer, .setup_static_arena = setup_static_arena_wrapper, .acquire_pooled_gm_heap = acquire_pooled_gm_heap_wrapper, .acquire_pooled_gm_sm = acquire_pooled_gm_sm_wrapper, diff --git a/src/common/platform/sim/host/device_runner_base.cpp b/src/common/platform/sim/host/device_runner_base.cpp index 7e8d55d854..e4be9f2d18 100644 --- a/src/common/platform/sim/host/device_runner_base.cpp +++ b/src/common/platform/sim/host/device_runner_base.cpp @@ -410,6 +410,42 @@ void *SimDeviceRunnerBase::acquire_graph_execution_buffer( return aligned_addr; } +void *SimDeviceRunnerBase::acquire_graph_submission_buffer( + uint32_t pipeline_slot, uint64_t graph_key, uint32_t occurrence, size_t bytes, size_t alignment +) { + if (pipeline_slot >= graph_submission_buffers_.size() || bytes == 0 || alignment == 0 || + (alignment & (alignment - 1)) != 0 || bytes > SIZE_MAX - (alignment - 1)) { + return nullptr; + } + std::vector &buffers = graph_submission_buffers_[pipeline_slot][graph_key]; + if (occurrence >= buffers.size()) buffers.resize(static_cast(occurrence) + 1); + RetainedGraphExecutionBuffer &buffer = buffers[occurrence]; + if (buffer.aligned_addr != nullptr && buffer.capacity >= bytes && + reinterpret_cast(buffer.aligned_addr) % alignment == 0) { + return buffer.aligned_addr; + } + + const size_t allocation_bytes = bytes + alignment - 1; + void *allocation = mem_alloc_.alloc(allocation_bytes); + if (allocation == nullptr) return nullptr; + const uintptr_t raw = reinterpret_cast(allocation); + if (raw > UINTPTR_MAX - (alignment - 1)) { + mem_alloc_.free(allocation); + return nullptr; + } + void *aligned_addr = reinterpret_cast((raw + alignment - 1) & ~(alignment - 1)); + if (device_memset(aligned_addr, 0, bytes) != 0) { + mem_alloc_.free(allocation); + return nullptr; + } + if (buffer.allocation != nullptr && mem_alloc_.free(buffer.allocation) != 0) { + mem_alloc_.free(allocation); + return nullptr; + } + buffer = RetainedGraphExecutionBuffer{allocation, aligned_addr, bytes}; + return aligned_addr; +} + void *SimDeviceRunnerBase::acquire_graph_definition_buffer( uint32_t pipeline_slot, uint64_t key, size_t bytes, size_t alignment ) { @@ -445,13 +481,19 @@ void *SimDeviceRunnerBase::acquire_graph_definition_buffer( } void SimDeviceRunnerBase::release_graph_execution_buffers() { - for (GraphExecutionBufferMap &by_key : graph_execution_buffers_) { + const auto release_occurrence_map = [](GraphExecutionBufferMap &by_key, MemoryAllocator &alloc) { for (auto &entry : by_key) { for (RetainedGraphExecutionBuffer &buffer : entry.second) { - if (buffer.allocation != nullptr) mem_alloc_.free(buffer.allocation); + if (buffer.allocation != nullptr) alloc.free(buffer.allocation); } } by_key.clear(); + }; + for (GraphExecutionBufferMap &by_key : graph_execution_buffers_) { + release_occurrence_map(by_key, mem_alloc_); + } + for (GraphExecutionBufferMap &by_key : graph_submission_buffers_) { + release_occurrence_map(by_key, mem_alloc_); } for (GraphDefinitionBufferMap &by_key : graph_definition_buffers_) { for (auto &entry : by_key) { diff --git a/src/common/platform/sim/host/device_runner_base.h b/src/common/platform/sim/host/device_runner_base.h index 8a6b84db71..097d928477 100644 --- a/src/common/platform/sim/host/device_runner_base.h +++ b/src/common/platform/sim/host/device_runner_base.h @@ -198,6 +198,9 @@ class SimDeviceRunnerBase { uint32_t pipeline_slot, uint64_t graph_key, uint32_t occurrence, size_t bytes, size_t alignment ); void *acquire_graph_definition_buffer(uint32_t pipeline_slot, uint64_t key, size_t bytes, size_t alignment); + void *acquire_graph_submission_buffer( + uint32_t pipeline_slot, uint64_t graph_key, uint32_t occurrence, size_t bytes, size_t alignment + ); void clear_temporary_buffer(); // On sim, allocate_tensor returns a plain host pointer, so the "device" @@ -344,6 +347,9 @@ class SimDeviceRunnerBase { }; using GraphExecutionBufferMap = std::unordered_map>; std::array graph_execution_buffers_{}; + // Graph submission POD storage, one retained block per (pipeline slot, + // graph_key, occurrence) — see HostApi acquire_graph_submission_buffer. + std::array graph_submission_buffers_{}; // Graph Definition storage, one retained block per (pipeline slot, // definition key) — see HostApi acquire_graph_definition_buffer. using GraphDefinitionBufferMap = std::unordered_map;