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 3bf60517d4..4b74caf8be 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,22 +417,237 @@ static bool relocate_host_orch_image( return ok; } +// Retained pinned bump arena for Graph submission POD images. Orch writes each +// layer in place; after entry() one H2D copies the used prefix into a device +// blob reused across binds. The buffer is owned by the DeviceRunner through +// HostApi acquire_pinned_host_buffer: the platform allocates it with +// aclrtMallocHost (sim: aligned host memory, where it costs nothing because +// sim copies are memcpy), reuses it while capacity fits, and releases it in +// finalize while the device context is still alive — a pinned mapping freed +// only at process exit races the next process's chip bring-up on the card. +constexpr size_t kGraphPinnedArenaBytes = 16ull * 1024ull * 1024ull; + +struct GraphPodH2d { + const HostApi *api = nullptr; + + struct RetainedBuf { + void *ptr = nullptr; + size_t bytes = 0; + }; + + static std::mutex &buf_mu() { + static std::mutex mu; + return mu; + } + static std::unordered_map &retained_subs() { + // One entry per (graph, occurrence) ever submitted on this fallback + // path, never evicted: process-lifetime and unbounded by design. The + // bulk-blob path above is the steady state; this map only grows when + // PODs fall out of the pinned bump, and each entry is one device + // allocation reused across binds. + static std::unordered_map m; + return m; + } + static RetainedBuf &retained_blob() { + static RetainedBuf b; + return b; + } + + void *acquire_sized(RetainedBuf &slot, size_t bytes) { + if (slot.ptr != nullptr && slot.bytes >= bytes) { + return slot.ptr; + } + if (slot.ptr != nullptr) { + api->device_free(slot.ptr); + slot.ptr = nullptr; + slot.bytes = 0; + } + slot.ptr = api->device_malloc(bytes); + if (slot.ptr == nullptr) { + return nullptr; + } + slot.bytes = bytes; + return slot.ptr; + } + + void *acquire_submission(uint64_t graph_key, uint32_t occurrence, size_t bytes) { + // FNV-1a over the full graph_key and the occurrence. A shift-based + // packing ((graph_key << 32) ^ occurrence) discards graph_key's upper + // 32 bits, so two graphs agreeing in the low half would collide on one + // retained buffer. + uint64_t key = 1469598103934665603ull; + for (int i = 0; i < 8; ++i) { + key ^= (graph_key >> (8 * i)) & 0xff; + key *= 1099511628211ull; + } + for (int i = 0; i < 4; ++i) { + key ^= (occurrence >> (8 * i)) & 0xff; + key *= 1099511628211ull; + } + std::scoped_lock lock(buf_mu()); + return acquire_sized(retained_subs()[key], bytes); + } + + void *acquire_blob(size_t bytes) { + std::scoped_lock lock(buf_mu()); + return acquire_sized(retained_blob(), bytes); + } + + // Wire the fields the device reads into a pending submission POD. Returns + // nullopt when the upload is invalid. + std::optional take_upload(GraphHostState &graph_state, size_t index) { + if (graph_host_upload_h2d_done(graph_state, index)) { + return std::nullopt; + } + 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 || + upload->outer_slot->task == nullptr) { + LOG_ERROR("host-orch: invalid pending Graph POD image"); + return std::nullopt; + } + auto *submission = reinterpret_cast(upload->data); + if (!graph_submission_wire_size_valid(*submission, upload->bytes)) { + LOG_ERROR("host-orch: Graph submission size does not match its POD image"); + return std::nullopt; + } + return upload; + } + + // Definition objects are uploaded by content identity before any + // submission referencing them, so the device-side reference is safe. + bool wire_definition(GraphSubmission *submission, const GraphDefinition *definition, void *device_object) { + if (definition->task_count == 0 || definition->task_count > GRAPH_MAX_NODES || + definition->full_key != submission->graph_key || definition->execution_storage_bytes == 0) { + LOG_ERROR("host-orch: invalid Graph Definition for submission"); + return false; + } + // Checked against the host-side Definition image the device object was + // built from; the GM object itself is never dereferenced on the host. + // Execution storage needs no retained buffer: it is the tail of the + // outer task's own heap allocation, which graph_submit_definition sized + // to required_heap + execution_storage_bytes. + submission->definition_addr = reinterpret_cast(device_object); + submission->local_execution = 0; + submission->activation_gate = 0; + return true; + } + + bool upload_per_layer( + GraphHostState &graph_state, + const std::unordered_map> &definition_objects + ) { + 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; + } + std::optional upload = take_upload(graph_state, index); + if (!upload.has_value()) { + return false; + } + auto *submission = reinterpret_cast(upload->data); + auto object_it = definition_objects.find(submission->definition_hash); + if (object_it == definition_objects.end() || object_it->second.first == nullptr) { + LOG_ERROR("host-orch: Graph submission has no uploaded Definition object"); + return false; + } + if (!wire_definition(submission, object_it->second.second, object_it->second.first)) { + return false; + } + void *device_submission = + acquire_submission(submission->graph_key, static_cast(index), upload->bytes); + if (device_submission == nullptr) { + LOG_ERROR("host-orch: failed to allocate 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"); + return false; + } + upload->outer_slot->graph_context = device_submission; + // Device POD storage is retained across binds; it must not enter + // tensor_pairs_, which validate frees every round. + graph_host_mark_upload_h2d_done(graph_state, index); + } + return true; + } + + static bool all_pinned_in_bump(GraphHostState &graph_state, std::byte *base, size_t used) { + const size_t count = graph_host_upload_count(graph_state); + if (count == 0 || base == nullptr || used == 0) { + return false; + } + for (size_t index = 0; index < count; ++index) { + std::optional upload = graph_host_upload(graph_state, index); + if (!upload.has_value() || upload->data == nullptr || upload->bytes == 0) { + return false; + } + if (upload->data < base || upload->data + upload->bytes > base + used) { + return false; + } + } + return true; + } + + bool upload_all( + GraphHostState &graph_state, + const std::unordered_map> &definition_objects + ) { + std::byte *base = graph_host_pinned_base(); + const size_t used = graph_host_pinned_used(); + if (!all_pinned_in_bump(graph_state, base, used)) { + return upload_per_layer(graph_state, definition_objects); + } + const size_t count = graph_host_upload_count(graph_state); + for (size_t index = 0; index < count; ++index) { + std::optional upload = take_upload(graph_state, index); + if (!upload.has_value()) { + return false; + } + auto *submission = reinterpret_cast(upload->data); + auto object_it = definition_objects.find(submission->definition_hash); + if (object_it == definition_objects.end() || object_it->second.first == nullptr) { + LOG_ERROR("host-orch: Graph submission has no uploaded Definition object"); + return false; + } + if (!wire_definition(submission, object_it->second.second, object_it->second.first)) { + return false; + } + } + void *device_blob = acquire_blob(used); + if (device_blob == nullptr) { + LOG_ERROR("host-orch: failed to allocate packed Graph POD blob (%zu bytes)", used); + return false; + } + if (api->copy_to_device(device_blob, base, used) != 0) { + LOG_ERROR("host-orch: failed to upload packed Graph POD blob (%zu bytes, %zu pods)", used, count); + return false; + } + auto *dev_bytes = static_cast(device_blob); + for (size_t index = 0; index < count; ++index) { + std::optional upload = graph_host_upload(graph_state, index); + if (!upload.has_value()) { + return false; + } + upload->outer_slot->graph_context = dev_bytes + static_cast(upload->data - base); + graph_host_mark_upload_h2d_done(graph_state, index); + } + return true; + } +}; + +// Upload each distinct Definition once as a shared device object +// ([GraphDefinitionHeader][Definition image]) keyed by content identity, then +// hand the submissions to the per-run GraphPodH2d. Definitions first: +// submissions reference the object's GM address, and the device boots only +// after both phases complete. bool upload_graph_submissions( - Runtime *runtime, const HostApi *api, GraphHostState &graph_state, uint64_t &uploaded_bytes + Runtime *runtime, const HostApi *api, GraphHostState &graph_state, uint64_t &uploaded_bytes, GraphPodH2d &h2d ) { 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. + std::unordered_map> definition_objects; 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; const auto *definition = reinterpret_cast(entry.data); @@ -460,59 +676,10 @@ bool upload_graph_submissions( LOG_ERROR("host-orch: failed to upload Graph Definition object"); return false; } - definition_objects.emplace(definition->content_hash, UploadedDefinition{object, definition}); + definition_objects.emplace(definition->content_hash, std::make_pair(object, definition)); uploaded_bytes += object_bytes; } - - // Pass 2: per-submission execution storage + the small reference image. - for (size_t index = 0; index < count; ++index) { - 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 || - upload->outer_slot->task == nullptr) { - LOG_ERROR("host-orch: invalid pending Graph POD image"); - return false; - } - auto *submission = reinterpret_cast(upload->data); - if (!graph_submission_wire_size_valid(*submission, upload->bytes)) { - LOG_ERROR("host-orch: Graph submission size does not match its POD image"); - 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; - } - // Checked against the host-side Definition image the device object was - // built from; the GM object itself is never dereferenced on the host. - // Execution storage needs no retained buffer: it is the tail of the - // outer task's own heap allocation, which graph_submit_definition sized - // to required_heap + execution_storage_bytes. - const GraphDefinition *definition = object_it->second.host_view; - if (definition->task_count == 0 || definition->task_count > GRAPH_MAX_NODES || - definition->full_key != submission->graph_key || definition->execution_storage_bytes == 0) { - LOG_ERROR("host-orch: invalid Graph Definition for submission"); - return false; - } - submission->definition_addr = reinterpret_cast(object_it->second.device_object); - submission->local_execution = 0; - submission->activation_gate = 0; - - void *device_submission = api->device_malloc(upload->bytes); - if (device_submission == nullptr) { - LOG_ERROR("host-orch: failed to allocate %zu bytes for Graph submission", upload->bytes); - 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); - } - return true; + return h2d.upload_all(graph_state, definition_objects); } struct GraphHostStateBinding { @@ -566,6 +733,23 @@ int32_t run_host_orchestration( } GraphHostStateBinding graph_binding(rt->orchestrator, graph_state.get()); + GraphPodH2d graph_h2d; + graph_h2d.api = api; + // The arena guard clears the orchestrator's pointer to the pinned bump on + // every exit path — an early return below must not leave the orchestrator + // allocating PODs from a bump owned by this run. + struct PinnedArenaScope { + ~PinnedArenaScope() { graph_host_clear_pinned_arena(); } + } pinned_arena_scope; + // The buffer is the runner's retained slot, so acquiring it here is a map + // lookup once it has settled at 16 MB, not a fresh aclrtMallocHost per bind. + if (void *pinned = api->acquire_pinned_host_buffer(kGraphPinnedArenaBytes, kGraphPinnedBumpAlign); + pinned != nullptr) { + graph_host_set_pinned_arena(static_cast(pinned), kGraphPinnedArenaBytes); + } else { + LOG_WARN("host-orch: pinned Graph arena unavailable; POD H2D may stage"); + } + 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()); @@ -631,10 +815,13 @@ int32_t run_host_orchestration( 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_graph_submissions(runtime, api, *graph_state, graph_bytes, graph_h2d)) return -1; { char attrs[96]; - snprintf(attrs, sizeof(attrs), "count=%zu bytes=%" PRIu64, graph_host_upload_count(*graph_state), graph_bytes); + snprintf( + attrs, sizeof(attrs), "count=%zu bytes=%" PRIu64, graph_host_upload_count(*graph_state), + graph_host_pinned_used() + graph_bytes + ); record_bind_phase(HostPhaseKind::BindGraphUpload, t_graph_ns, attrs, graph_bytes); } 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 fd0f081513..f497967335 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,6 +338,9 @@ 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 { @@ -345,6 +349,36 @@ struct GraphHostState { std::vector pending_uploads; }; +namespace { +std::byte *g_pin_base = nullptr; +size_t g_pin_cap = 0; +size_t g_pin_used = 0; + +std::byte *graph_host_pinned_bump(size_t bytes) { + if (g_pin_base == nullptr || bytes == 0) return nullptr; + const size_t off = (g_pin_used + kGraphPinnedBumpAlign - 1) & ~(kGraphPinnedBumpAlign - 1); + if (off + bytes > g_pin_cap) return nullptr; + g_pin_used = off + bytes; + return g_pin_base + off; +} +} // namespace + +void graph_host_set_pinned_arena(std::byte *base, size_t cap) { + g_pin_base = base; + g_pin_cap = cap; + g_pin_used = 0; +} + +void graph_host_clear_pinned_arena() { + g_pin_base = nullptr; + g_pin_cap = 0; + g_pin_used = 0; +} + +std::byte *graph_host_pinned_base() { return g_pin_base; } + +size_t graph_host_pinned_used() { return g_pin_used; } + namespace { GraphHostState *graph_state_from(PTO2OrchestratorState *orch) { @@ -688,10 +722,22 @@ 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()}; } +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; +} + GraphHostDefinitionList graph_host_definitions(GraphHostState &state) { GraphHostDefinitionList list; list.entries.reserve(state.definitions.size()); @@ -1418,47 +1464,82 @@ 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 +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 || tensors_offset == nullptr || scalars_offset == nullptr || scalar_bytes == nullptr || + graph_definition(definition_image) == nullptr) { + return false; + } + *tensors_offset = 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_offset > UINT32_MAX || *tensors_offset > 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_offset + tensor_bytes; + *scalar_bytes = static_cast(args.scalar_count()) * sizeof(uint64_t); + *scalars_offset = args.scalar_count() == 0 ? 0 : PTO2_ALIGN_UP(tensors_end, alignof(uint64_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) { return false; } - submission_image->assign(total_bytes, std::byte{0}); - auto *tensors = reinterpret_cast(submission_image->data() + tensors_offset); + 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 @@ -1484,7 +1565,25 @@ bool graph_submit_definition( } GraphPendingUpload pending; - if (!graph_build_submission_image(definition_image, args, &pending.image)) return false; + size_t total_bytes = 0; + size_t tensors_offset = 0; + size_t scalars_offset = 0; + size_t scalar_bytes = 0; + if (!graph_submission_layout( + definition_image, args, &total_bytes, &tensors_offset, &scalars_offset, &scalar_bytes + )) { + return false; + } + std::byte *pinned = graph_host_pinned_bump(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, 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 7a224b4343..7f312765ff 100644 --- a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp @@ -44,6 +44,7 @@ #include #include #include +#include #include #include #include @@ -461,22 +462,237 @@ static bool relocate_host_orch_image( return ok; } +// Retained pinned bump arena for Graph submission POD images. Orch writes each +// layer in place; after entry() one H2D copies the used prefix into a device +// blob reused across binds. The buffer is owned by the DeviceRunner through +// HostApi acquire_pinned_host_buffer: the platform allocates it with +// aclrtMallocHost (sim: aligned host memory, where it costs nothing because +// sim copies are memcpy), reuses it while capacity fits, and releases it in +// finalize while the device context is still alive — a pinned mapping freed +// only at process exit races the next process's chip bring-up on the card. +constexpr size_t kGraphPinnedArenaBytes = 16ull * 1024ull * 1024ull; + +struct GraphPodH2d { + const HostApi *api = nullptr; + + struct RetainedBuf { + void *ptr = nullptr; + size_t bytes = 0; + }; + + static std::mutex &buf_mu() { + static std::mutex mu; + return mu; + } + static std::unordered_map &retained_subs() { + // One entry per (graph, occurrence) ever submitted on this fallback + // path, never evicted: process-lifetime and unbounded by design. The + // bulk-blob path above is the steady state; this map only grows when + // PODs fall out of the pinned bump, and each entry is one device + // allocation reused across binds. + static std::unordered_map m; + return m; + } + static RetainedBuf &retained_blob() { + static RetainedBuf b; + return b; + } + + void *acquire_sized(RetainedBuf &slot, size_t bytes) { + if (slot.ptr != nullptr && slot.bytes >= bytes) { + return slot.ptr; + } + if (slot.ptr != nullptr) { + api->device_free(slot.ptr); + slot.ptr = nullptr; + slot.bytes = 0; + } + slot.ptr = api->device_malloc(bytes); + if (slot.ptr == nullptr) { + return nullptr; + } + slot.bytes = bytes; + return slot.ptr; + } + + void *acquire_submission(uint64_t graph_key, uint32_t occurrence, size_t bytes) { + // FNV-1a over the full graph_key and the occurrence. A shift-based + // packing ((graph_key << 32) ^ occurrence) discards graph_key's upper + // 32 bits, so two graphs agreeing in the low half would collide on one + // retained buffer. + uint64_t key = 1469598103934665603ull; + for (int i = 0; i < 8; ++i) { + key ^= (graph_key >> (8 * i)) & 0xff; + key *= 1099511628211ull; + } + for (int i = 0; i < 4; ++i) { + key ^= (occurrence >> (8 * i)) & 0xff; + key *= 1099511628211ull; + } + std::scoped_lock lock(buf_mu()); + return acquire_sized(retained_subs()[key], bytes); + } + + void *acquire_blob(size_t bytes) { + std::scoped_lock lock(buf_mu()); + return acquire_sized(retained_blob(), bytes); + } + + // Wire the fields the device reads into a pending submission POD. Returns + // nullopt when the upload is invalid. + std::optional take_upload(GraphHostState &graph_state, size_t index) { + if (graph_host_upload_h2d_done(graph_state, index)) { + return std::nullopt; + } + 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 || + upload->outer_slot->task == nullptr) { + LOG_ERROR("host-orch: invalid pending Graph POD image"); + return std::nullopt; + } + auto *submission = reinterpret_cast(upload->data); + if (!graph_submission_wire_size_valid(*submission, upload->bytes)) { + LOG_ERROR("host-orch: Graph submission size does not match its POD image"); + return std::nullopt; + } + return upload; + } + + // Definition objects are uploaded by content identity before any + // submission referencing them, so the device-side reference is safe. + bool wire_definition(GraphSubmission *submission, const GraphDefinition *definition, void *device_object) { + if (definition->task_count == 0 || definition->task_count > GRAPH_MAX_NODES || + definition->full_key != submission->graph_key || definition->execution_storage_bytes == 0) { + LOG_ERROR("host-orch: invalid Graph Definition for submission"); + return false; + } + // Checked against the host-side Definition image the device object was + // built from; the GM object itself is never dereferenced on the host. + // Execution storage needs no retained buffer: it is the tail of the + // outer task's own heap allocation, which graph_submit_definition sized + // to required_heap + execution_storage_bytes. + submission->definition_addr = reinterpret_cast(device_object); + submission->local_execution = 0; + submission->activation_gate = 0; + return true; + } + + bool upload_per_layer( + GraphHostState &graph_state, + const std::unordered_map> &definition_objects + ) { + 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; + } + std::optional upload = take_upload(graph_state, index); + if (!upload.has_value()) { + return false; + } + auto *submission = reinterpret_cast(upload->data); + auto object_it = definition_objects.find(submission->definition_hash); + if (object_it == definition_objects.end() || object_it->second.first == nullptr) { + LOG_ERROR("host-orch: Graph submission has no uploaded Definition object"); + return false; + } + if (!wire_definition(submission, object_it->second.second, object_it->second.first)) { + return false; + } + void *device_submission = + acquire_submission(submission->graph_key, static_cast(index), upload->bytes); + if (device_submission == nullptr) { + LOG_ERROR("host-orch: failed to allocate 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"); + return false; + } + upload->outer_slot->graph_context = device_submission; + // Device POD storage is retained across binds; it must not enter + // tensor_pairs_, which validate frees every round. + graph_host_mark_upload_h2d_done(graph_state, index); + } + return true; + } + + static bool all_pinned_in_bump(GraphHostState &graph_state, std::byte *base, size_t used) { + const size_t count = graph_host_upload_count(graph_state); + if (count == 0 || base == nullptr || used == 0) { + return false; + } + for (size_t index = 0; index < count; ++index) { + std::optional upload = graph_host_upload(graph_state, index); + if (!upload.has_value() || upload->data == nullptr || upload->bytes == 0) { + return false; + } + if (upload->data < base || upload->data + upload->bytes > base + used) { + return false; + } + } + return true; + } + + bool upload_all( + GraphHostState &graph_state, + const std::unordered_map> &definition_objects + ) { + std::byte *base = graph_host_pinned_base(); + const size_t used = graph_host_pinned_used(); + if (!all_pinned_in_bump(graph_state, base, used)) { + return upload_per_layer(graph_state, definition_objects); + } + const size_t count = graph_host_upload_count(graph_state); + for (size_t index = 0; index < count; ++index) { + std::optional upload = take_upload(graph_state, index); + if (!upload.has_value()) { + return false; + } + auto *submission = reinterpret_cast(upload->data); + auto object_it = definition_objects.find(submission->definition_hash); + if (object_it == definition_objects.end() || object_it->second.first == nullptr) { + LOG_ERROR("host-orch: Graph submission has no uploaded Definition object"); + return false; + } + if (!wire_definition(submission, object_it->second.second, object_it->second.first)) { + return false; + } + } + void *device_blob = acquire_blob(used); + if (device_blob == nullptr) { + LOG_ERROR("host-orch: failed to allocate packed Graph POD blob (%zu bytes)", used); + return false; + } + if (api->copy_to_device(device_blob, base, used) != 0) { + LOG_ERROR("host-orch: failed to upload packed Graph POD blob (%zu bytes, %zu pods)", used, count); + return false; + } + auto *dev_bytes = static_cast(device_blob); + for (size_t index = 0; index < count; ++index) { + std::optional upload = graph_host_upload(graph_state, index); + if (!upload.has_value()) { + return false; + } + upload->outer_slot->graph_context = dev_bytes + static_cast(upload->data - base); + graph_host_mark_upload_h2d_done(graph_state, index); + } + return true; + } +}; + +// Upload each distinct Definition once as a shared device object +// ([GraphDefinitionHeader][Definition image]) keyed by content identity, then +// hand the submissions to the per-run GraphPodH2d. Definitions first: +// submissions reference the object's GM address, and the device boots only +// after both phases complete. bool upload_graph_submissions( - Runtime *runtime, const HostApi *api, GraphHostState &graph_state, uint64_t &uploaded_bytes + Runtime *runtime, const HostApi *api, GraphHostState &graph_state, uint64_t &uploaded_bytes, GraphPodH2d &h2d ) { 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. + std::unordered_map> definition_objects; 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; const auto *definition = reinterpret_cast(entry.data); @@ -505,59 +721,10 @@ bool upload_graph_submissions( LOG_ERROR("host-orch: failed to upload Graph Definition object"); return false; } - definition_objects.emplace(definition->content_hash, UploadedDefinition{object, definition}); + definition_objects.emplace(definition->content_hash, std::make_pair(object, definition)); uploaded_bytes += object_bytes; } - - // Pass 2: per-submission execution storage + the small reference image. - for (size_t index = 0; index < count; ++index) { - 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 || - upload->outer_slot->task == nullptr) { - LOG_ERROR("host-orch: invalid pending Graph POD image"); - return false; - } - auto *submission = reinterpret_cast(upload->data); - if (!graph_submission_wire_size_valid(*submission, upload->bytes)) { - LOG_ERROR("host-orch: Graph submission size does not match its POD image"); - 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; - } - // Checked against the host-side Definition image the device object was - // built from; the GM object itself is never dereferenced on the host. - // Execution storage needs no retained buffer: it is the tail of the - // outer task's own heap allocation, which graph_submit_definition sized - // to required_heap + execution_storage_bytes. - const GraphDefinition *definition = object_it->second.host_view; - if (definition->task_count == 0 || definition->task_count > GRAPH_MAX_NODES || - definition->full_key != submission->graph_key || definition->execution_storage_bytes == 0) { - LOG_ERROR("host-orch: invalid Graph Definition for submission"); - return false; - } - submission->definition_addr = reinterpret_cast(object_it->second.device_object); - submission->local_execution = 0; - submission->activation_gate = 0; - - void *device_submission = api->device_malloc(upload->bytes); - if (device_submission == nullptr) { - LOG_ERROR("host-orch: failed to allocate %zu bytes for Graph submission", upload->bytes); - 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); - } - return true; + return h2d.upload_all(graph_state, definition_objects); } struct GraphHostStateBinding { @@ -613,6 +780,23 @@ int32_t run_host_orchestration( } GraphHostStateBinding graph_binding(rt->orchestrator, graph_state.get()); + GraphPodH2d graph_h2d; + graph_h2d.api = api; + // The arena guard clears the orchestrator's pointer to the pinned bump on + // every exit path — an early return below must not leave the orchestrator + // allocating PODs from a bump owned by this run. The buffer itself is the + // runner's retained slot, so acquiring it here is a map lookup once it has + // settled at 16 MB, not a fresh aclrtMallocHost per bind. + struct PinnedArenaScope { + ~PinnedArenaScope() { graph_host_clear_pinned_arena(); } + } pinned_arena_scope; + if (void *pinned = api->acquire_pinned_host_buffer(kGraphPinnedArenaBytes, kGraphPinnedBumpAlign); + pinned != nullptr) { + graph_host_set_pinned_arena(static_cast(pinned), kGraphPinnedArenaBytes); + } else { + LOG_WARN("host-orch: pinned Graph arena unavailable; POD H2D may stage"); + } + // 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 @@ -687,10 +871,13 @@ int32_t run_host_orchestration( 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_graph_submissions(runtime, api, *graph_state, graph_bytes, graph_h2d)) return -1; { char attrs[96]; - snprintf(attrs, sizeof(attrs), "count=%zu bytes=%" PRIu64, graph_host_upload_count(*graph_state), graph_bytes); + snprintf( + attrs, sizeof(attrs), "count=%zu bytes=%" PRIu64, graph_host_upload_count(*graph_state), + graph_host_pinned_used() + graph_bytes + ); record_bind_phase(HostPhaseKind::BindGraphUpload, t_graph_ns, attrs, graph_bytes); } 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 fd0f081513..f497967335 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,6 +338,9 @@ 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 { @@ -345,6 +349,36 @@ struct GraphHostState { std::vector pending_uploads; }; +namespace { +std::byte *g_pin_base = nullptr; +size_t g_pin_cap = 0; +size_t g_pin_used = 0; + +std::byte *graph_host_pinned_bump(size_t bytes) { + if (g_pin_base == nullptr || bytes == 0) return nullptr; + const size_t off = (g_pin_used + kGraphPinnedBumpAlign - 1) & ~(kGraphPinnedBumpAlign - 1); + if (off + bytes > g_pin_cap) return nullptr; + g_pin_used = off + bytes; + return g_pin_base + off; +} +} // namespace + +void graph_host_set_pinned_arena(std::byte *base, size_t cap) { + g_pin_base = base; + g_pin_cap = cap; + g_pin_used = 0; +} + +void graph_host_clear_pinned_arena() { + g_pin_base = nullptr; + g_pin_cap = 0; + g_pin_used = 0; +} + +std::byte *graph_host_pinned_base() { return g_pin_base; } + +size_t graph_host_pinned_used() { return g_pin_used; } + namespace { GraphHostState *graph_state_from(PTO2OrchestratorState *orch) { @@ -688,10 +722,22 @@ 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()}; } +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; +} + GraphHostDefinitionList graph_host_definitions(GraphHostState &state) { GraphHostDefinitionList list; list.entries.reserve(state.definitions.size()); @@ -1418,47 +1464,82 @@ 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 +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 || tensors_offset == nullptr || scalars_offset == nullptr || scalar_bytes == nullptr || + graph_definition(definition_image) == nullptr) { + return false; + } + *tensors_offset = 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_offset > UINT32_MAX || *tensors_offset > 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_offset + tensor_bytes; + *scalar_bytes = static_cast(args.scalar_count()) * sizeof(uint64_t); + *scalars_offset = args.scalar_count() == 0 ? 0 : PTO2_ALIGN_UP(tensors_end, alignof(uint64_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) { return false; } - submission_image->assign(total_bytes, std::byte{0}); - auto *tensors = reinterpret_cast(submission_image->data() + tensors_offset); + 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 @@ -1484,7 +1565,25 @@ bool graph_submit_definition( } GraphPendingUpload pending; - if (!graph_build_submission_image(definition_image, args, &pending.image)) return false; + size_t total_bytes = 0; + size_t tensors_offset = 0; + size_t scalars_offset = 0; + size_t scalar_bytes = 0; + if (!graph_submission_layout( + definition_image, args, &total_bytes, &tensors_offset, &scalars_offset, &scalar_bytes + )) { + return false; + } + std::byte *pinned = graph_host_pinned_bump(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, diff --git a/src/common/host_build_graph/graph_host_state.h b/src/common/host_build_graph/graph_host_state.h index 265f70907c..5f5bbd4a06 100644 --- a/src/common/host_build_graph/graph_host_state.h +++ b/src/common/host_build_graph/graph_host_state.h @@ -49,4 +49,17 @@ struct GraphHostDefinitionList { GraphHostStatePtr make_graph_host_state(); size_t graph_host_upload_count(const GraphHostState &state); std::optional graph_host_upload(GraphHostState &state, size_t index); +bool graph_host_upload_h2d_done(const GraphHostState &state, size_t index); +void graph_host_mark_upload_h2d_done(GraphHostState &state, size_t index); GraphHostDefinitionList graph_host_definitions(GraphHostState &state); + +// Optional pinned bump arena for Graph submission POD images. Set by the host +// runtime before orch entry so graph_submit_definition can write each POD in +// place; unset means the fallback std::vector images. The base handed to +// graph_host_set_pinned_arena must satisfy kGraphPinnedBumpAlign — the +// orchestrator aligns offsets relative to it without re-aligning the base. +inline constexpr size_t kGraphPinnedBumpAlign = 64; +void graph_host_set_pinned_arena(std::byte *base, size_t cap); +void graph_host_clear_pinned_arena(); +std::byte *graph_host_pinned_base(); +size_t graph_host_pinned_used(); diff --git a/src/common/platform/include/common/host_api.h b/src/common/platform/include/common/host_api.h index d79459779e..77f26bf998 100644 --- a/src/common/platform/include/common/host_api.h +++ b/src/common/platform/include/common/host_api.h @@ -62,6 +62,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 pinned host buffer for DMA staging (hbg's Graph submission + // POD arena). One {addr, size} slot per pipeline slot, grow-only: a larger + // request replaces the block, an equal-or-smaller one reuses it. The + // returned base honors `alignment` (power of two). Onboard backs it with + // aclrtMallocHost and releases it in finalize_common() while the device + // context is still alive — a pinned mapping outliving aclFinalize is what + // wedges the next process's chip bring-up on the same card. Sim backs it + // with aligned host memory (sim copies are memcpy). + void *(*acquire_pinned_host_buffer)(void *runner_ctx, uint32_t pipeline_slot, 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 @@ -162,6 +171,10 @@ 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_pinned_host_buffer(size_t bytes, size_t alignment) const { + if (ops_->acquire_pinned_host_buffer == nullptr) return nullptr; + return ops_->acquire_pinned_host_buffer(runner_ctx_, pipeline_slot_, 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 df609176c6..8932f16049 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -168,6 +168,15 @@ static void *acquire_graph_definition_buffer( } } +static void *acquire_pinned_host_buffer(void *runner_ctx, uint32_t pipeline_slot, size_t bytes, size_t alignment) { + if (runner_ctx == nullptr) return nullptr; + try { + return static_cast(runner_ctx)->acquire_pinned_host_buffer(pipeline_slot, 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 +291,7 @@ static const HostApiOps g_host_api_ops = { .get_retained_temp_buffer = get_retained_temp_buffer, .set_retained_temp_buffer = set_retained_temp_buffer, .acquire_graph_definition_buffer = acquire_graph_definition_buffer, + .acquire_pinned_host_buffer = acquire_pinned_host_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 e3a75da59d..d8b2752554 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -214,6 +214,46 @@ void DeviceRunnerBase::release_graph_definition_buffers() { } } +void *DeviceRunnerBase::acquire_pinned_host_buffer(uint32_t pipeline_slot, std::size_t bytes, std::size_t alignment) { + if (pipeline_slot >= pinned_host_buffers_.size() || bytes == 0 || alignment == 0 || + (alignment & (alignment - 1)) != 0 || bytes > SIZE_MAX - (alignment - 1)) { + return nullptr; + } + PinnedHostBuffer &buffer = pinned_host_buffers_[pipeline_slot]; + 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 = nullptr; + if (aclrtMallocHost(&allocation, allocation_bytes) != ACL_SUCCESS || allocation == nullptr) { + LOG_ERROR("aclrtMallocHost(%zu) for pinned staging failed", allocation_bytes); + return nullptr; + } + const uintptr_t raw = reinterpret_cast(allocation); + if (raw > UINTPTR_MAX - (alignment - 1)) { + (void)aclrtFreeHost(allocation); + return nullptr; + } + void *aligned_addr = reinterpret_cast((raw + alignment - 1) & ~(alignment - 1)); + if (buffer.allocation != nullptr && aclrtFreeHost(buffer.allocation) != ACL_SUCCESS) { + (void)aclrtFreeHost(allocation); + return nullptr; + } + buffer = PinnedHostBuffer{allocation, aligned_addr, bytes}; + return aligned_addr; +} + +void DeviceRunnerBase::release_pinned_host_buffers() { + for (PinnedHostBuffer &buffer : pinned_host_buffers_) { + if (buffer.allocation != nullptr) { + (void)aclrtFreeHost(buffer.allocation); + } + buffer = PinnedHostBuffer{}; + } +} + void DeviceRunnerBase::abandon_graph_definition_buffers() { for (GraphDefinitionBufferMap &by_key : graph_definition_buffers_) { by_key.clear(); @@ -1290,9 +1330,15 @@ int DeviceRunnerBase::finalize_common_impl(bool abandon_device_resources) { abandon_graph_definition_buffers(); retained_temp_addrs_.fill(nullptr); retained_temp_sizes_.fill(0); + // The pinned staging blocks are host mappings, not device resources — + // they can and must be released even on the fatal path: a force reset + // does not invalidate a pinned mapping, and leaving one behind keeps a + // driver-side registration alive past aclFinalize. + release_pinned_host_buffers(); } else { release_graph_definition_buffers(); clear_temporary_buffer(); + release_pinned_host_buffers(); } // Free the device-phase/task-timing buffer (allocated lazily in run()) while diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index 9cb66e2ba3..c82f436c2f 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -146,6 +146,12 @@ class DeviceRunnerBase { void set_retained_temp_buffer(uint32_t pipeline_slot, void *addr, std::size_t size); void * acquire_graph_definition_buffer(uint32_t pipeline_slot, uint64_t key, std::size_t bytes, std::size_t alignment); + /** + * Retained pinned host buffer per pipeline slot (see HostApi + * acquire_pinned_host_buffer). Grow-only; released in finalize_common() + * while the device context is alive. `alignment` must be a power of two. + */ + void *acquire_pinned_host_buffer(uint32_t pipeline_slot, std::size_t bytes, std::size_t alignment); void clear_temporary_buffer(); /** * Map a device buffer into the host address space and return a @@ -914,6 +920,8 @@ class DeviceRunnerBase { */ int finalize_common(); void release_graph_definition_buffers(); + /** aclrtFreeHost every retained pinned host staging block (idempotent). */ + void release_pinned_host_buffers(); /** * Drop the retained graph-definition buffers without freeing them. @@ -1073,6 +1081,18 @@ class DeviceRunnerBase { // the grow/pack logic lives in trb bind. std::array retained_temp_addrs_{}; std::array retained_temp_sizes_{}; + // Pinned host staging block per pipeline slot — the raw aclrtMallocHost + // allocation plus the aligned base handed out (allocation is padded by + // `alignment - 1` so the aligned base always exists). Freed in + // finalize_common() before the device reset; a pinned mapping released + // only at process exit races the next process's chip bring-up on the + // same card. + struct PinnedHostBuffer { + void *allocation{nullptr}; + void *aligned_addr{nullptr}; + std::size_t capacity{0}; + }; + std::array pinned_host_buffers_{}; // One retained device block: the raw allocation plus the aligned address // handed out. Backs the Graph Definition cache below. struct RetainedGraphBuffer { diff --git a/src/common/platform/sim/host/c_api_shared.cpp b/src/common/platform/sim/host/c_api_shared.cpp index 26eeb64286..7c06c34b0f 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -154,6 +154,16 @@ static void *acquire_graph_definition_buffer( } } +static void *acquire_pinned_host_buffer(void *runner_ctx, uint32_t pipeline_slot, size_t bytes, size_t alignment) { + if (runner_ctx == nullptr) return nullptr; + try { + return static_cast(runner_ctx) + ->acquire_pinned_host_buffer(pipeline_slot, 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 { @@ -269,6 +279,7 @@ static const HostApiOps g_host_api_ops = { .get_retained_temp_buffer = get_retained_temp_buffer, .set_retained_temp_buffer = set_retained_temp_buffer, .acquire_graph_definition_buffer = acquire_graph_definition_buffer, + .acquire_pinned_host_buffer = acquire_pinned_host_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 9868c4feaa..d8c970d8f0 100644 --- a/src/common/platform/sim/host/device_runner_base.cpp +++ b/src/common/platform/sim/host/device_runner_base.cpp @@ -426,6 +426,21 @@ void SimDeviceRunnerBase::clear_temporary_buffer() { } } +void *SimDeviceRunnerBase::acquire_pinned_host_buffer(uint32_t pipeline_slot, size_t bytes, size_t alignment) { + // Sim device memory is host memory, so "pinned" staging is an aligned + // plain block through the same allocator the copies already use. The key + // is distinct from every Graph Definition key so the staging block never + // aliases a Definition buffer on the same slot; it is freed with the rest + // of the map by release_graph_definition_buffers() in finalize(). + constexpr uint64_t kPinnedHostBufferKey = 0x70696e6e6564ull; // "pinned" + return acquire_graph_definition_buffer(pipeline_slot, kPinnedHostBufferKey, bytes, alignment); +} + +void SimDeviceRunnerBase::release_pinned_host_buffers() { + // No dedicated state: the staging block lives in graph_definition_buffers_ + // and is freed with that map in finalize(). +} + int SimDeviceRunnerBase::stamp_orch_so(Runtime &runtime, int32_t cid) { // Registered-callable flow only: the orch SO was already delivered to the // sim AICPU at launch_device_register time. A run just needs the active diff --git a/src/common/platform/sim/host/device_runner_base.h b/src/common/platform/sim/host/device_runner_base.h index cda535c5be..b454faa70b 100644 --- a/src/common/platform/sim/host/device_runner_base.h +++ b/src/common/platform/sim/host/device_runner_base.h @@ -195,6 +195,8 @@ class SimDeviceRunnerBase { void get_retained_temp_buffer(uint32_t pipeline_slot, void **addr, size_t *size); void set_retained_temp_buffer(uint32_t pipeline_slot, void *addr, size_t size); void *acquire_graph_definition_buffer(uint32_t pipeline_slot, uint64_t key, size_t bytes, size_t alignment); + void *acquire_pinned_host_buffer(uint32_t pipeline_slot, size_t bytes, size_t alignment); + void release_pinned_host_buffers(); void clear_temporary_buffer(); // On sim, allocate_tensor returns a plain host pointer, so the "device"