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;