From f82c10e167c3bb30e1c8c8e65efa7dae74155b72 Mon Sep 17 00:00:00 2001 From: ChaoZheng109 Date: Tue, 18 Aug 2026 00:58:41 -0700 Subject: [PATCH] Update: share Graph Definitions across HBG submissions Upload each distinct Graph Definition once per run and let Graph submissions reference the retained device object. Remove duplicate Definition storage from Graph executions while preserving content validation and affine reuse. Record the retained object length so malformed total_bytes values are rejected before hashing, remove the unused legacy Definition accessor, and use architecture-specific hints while waiting for peer verification. Add shared wire-format coverage, dirty-storage guardrails, and an investigation record that distinguishes the shared Definition upload from per-submission reference uploads. --- ...6-08-hbg-graph-definition-single-upload.md | 82 +++++++++ docs/investigations/README.md | 1 + .../host_build_graph/host/runtime_maker.cpp | 61 ++++++- .../orchestrator_core/pto_orchestrator.cpp | 20 ++- .../host_build_graph/host/runtime_maker.cpp | 61 ++++++- .../orchestrator_core/pto_orchestrator.cpp | 20 ++- .../host_build_graph/graph_execution.cpp | 130 ++++++++++---- src/common/host_build_graph/graph_execution.h | 68 ++++---- .../host_build_graph/graph_host_state.h | 15 ++ src/common/platform/include/common/host_api.h | 14 ++ .../platform/onboard/host/c_api_shared.cpp | 13 ++ .../onboard/host/device_runner_base.cpp | 43 +++++ .../onboard/host/device_runner_base.h | 8 + src/common/platform/sim/host/c_api_shared.cpp | 13 ++ .../platform/sim/host/device_runner_base.cpp | 40 +++++ .../platform/sim/host/device_runner_base.h | 5 + tests/ut/cpp/common/test_hbg_graph_cache.cpp | 165 +++++++++++++++--- 17 files changed, 649 insertions(+), 110 deletions(-) create mode 100644 docs/investigations/2026-08-hbg-graph-definition-single-upload.md diff --git a/docs/investigations/2026-08-hbg-graph-definition-single-upload.md b/docs/investigations/2026-08-hbg-graph-definition-single-upload.md new file mode 100644 index 0000000000..b06499a17e --- /dev/null +++ b/docs/investigations/2026-08-hbg-graph-definition-single-upload.md @@ -0,0 +1,82 @@ +# 2026-08 — hbg: upload Graph Definitions once as shared device objects + +## Question + +Breaking down the host side of `examples/a2a3/host_build_graph/qwen3_14b_decode` +(40 Graph submissions of one 277-node, 130,192-byte Definition) showed two +dominant costs outside pure orchestration: + +- the orchestration window rebuilt and zero-filled a 132,752-byte submission + image per layer — 98.1% of those bytes a byte-identical Definition copy + (KNOWN_ISSUES at the time); +- the upload stage shipped all 40 images: 5.31 MB across 40 + `device_malloc` + `rtMemcpy` round trips, a stable 14.4–15.9 ms per run. + +Both trace to one design choice: the Definition travels *inside* every +submission image. + +## Change + +`254f924e` (measured by `4d434174`/`b8095e39`, enabled by `f868ac52`): + +- each distinct Definition uploads **once** as a + `[GraphDefinitionHeader][Definition image]` object retained by + `acquire_graph_definition_buffer`, keyed by content identity; +- `GraphSubmission` carries `definition_addr` + `definition_hash` instead of + the inline image; a submission is now 2,568 bytes; +- device localize validates the shared object through a one-time verify gate + (first localizer FNV-hashes, peers spin on the state word) and binds + topology against the shared image in place — the per-occurrence embedded + Definition copy in execution storage is gone. + +## Result (qwen3, `SIMPLER_SKIP_DEVICE_RUN=1`, 5 serial runs each, median) + +| Stage | Before | After | Δ | +| ----- | ------ | ----- | - | +| orch image build (incl. zero-fill) | 931 µs (826) | 23.6 µs (10.0) | −97.5% | +| submission bytes | 5,310,080 | 232,944 | −95.6% | +| orch window total | 1,834 µs | 836 µs | −54% | +| **H2D upload time** | **14.66 ms** | **12.10 ms** | **−17%** | + +## Why the H2D time gain is far below the byte gain + +Bytes fell 95.6% but the upload *time* fell only 17%: the stage's cost was +never bandwidth-dominated. The effective rate is absurd on both sides — +0.36 GB/s before, 0.02 GB/s after — which is the signature of fixed +per-call costs dominating data movement. The optimized stage still makes 41 +allocation-and-copy pairs: one 130,192-byte shared Definition object and 40 +2,568-byte reference submissions. + +Each upload pays: + +1. **`rtMalloc` per object** — `MemoryAllocator::alloc` calls CANN + `rtMalloc(RT_MEMORY_HBM)` (a driver round trip) plus a mutex-guarded map + insert; the one Definition object and 40 submissions make 41 allocations + per run, freed again at teardown. +2. **`rtMemcpy` (sync, `RT_MEMCPY_HOST_TO_DEVICE`) per object** — each + call is a blocking submit-and-wait on the copy stream: host builds the + descriptor, pushes to the driver, and blocks for completion. Each 2.5 KB + reference submission never occupies the link long enough for bandwidth to + matter. The 12.1 ms / 41 ≈ 295 µs figure is the average across one large + and 40 small calls, not a uniform per-submission measurement, but it bounds + the scale of the driver and synchronization overhead dominating the stage. + +So the model is `time ≈ N × (malloc + memcpy latency) + bytes / BW`, and at +these sizes the first term dominates by two orders of magnitude. The byte +reduction could only remove the (already small) second term — consistent +with the measured 2.5 ms saving landing near 40 × 130 KB / 0.36 GB/s of +actual removed traffic. + +**The follow-up this points at**: batch the 40 reference submissions into one +packed upload (they are POD and contiguous by construction; one malloc, one +memcpy, per-device offsets), which attacks the *N* in the dominant term. +Expected to take the stage from ~12 ms to well under 1 ms. Not done in this +change to keep the wire-v2 commit reviewable and independently revertible. + +## Notes + +- The first-cut device verify gate returned "busy-looking" nulls to peer + submissions and surfaced as `sched_error_code=5 INVALID_ARGS`; the fix is + the spin-wait on `verify_state` (dispatch-path legal: spin, no sleep). +- `graph record` (245–687 µs) remains per-run and untouched — the per-run + Definition cache discard is a separate, still-open item (KNOWN_ISSUES). diff --git a/docs/investigations/README.md b/docs/investigations/README.md index da7b6058cc..b9fa49f7dd 100644 --- a/docs/investigations/README.md +++ b/docs/investigations/README.md @@ -84,6 +84,7 @@ that ...". Newest first. +- [2026-08 — hbg: uploading Graph Definitions once as shared device objects](2026-08-hbg-graph-definition-single-upload.md) — cut the per-replay 130 KB Definition re-serialization (image build 931→24 µs, orch total −54%), and measured the follow-on truth that the H2D stage is latency-bound: bytes fell 95.6% but time only 17% because fixed allocation and synchronous-copy costs dominate 41 calls (one 130 KB Definition plus 40 × 2.5 KB references) — the next lever is batching the reference submissions into one packed upload - [2026-07 — Why qwen3_14b_decode held a device for 406 s, and the four answers that were wrong](2026-07-qwen-scene-test-406s-decomposition.md) — root cause was torch thread oversubscription in goldens, not any of the big things: the reference walks 3584 tiny slice ops per layer and torch sizes its pool from the core count, so a 320-core host paid 6.35 s/layer against 1.05 s at 4. Capping to 8 (#1601) cut the golden 359 s → ~40 s and let the case rejoin the sweep. Records the measured decomposition (compile 59 s / fixture 13 s / golden 359 s / device tens of ms), how to split queue wait from card-held work via the `npu-lock` timestamps, and why golden caching, vectorisation, a nightly split and `skip_golden` were each dropped - [2026-07 — `sync_start` drain retry ABA across reusable barrier state](2026-07-sync-start-drain-retry-aba.md) — reproduced #1455 deterministically by suspending an old-attempt scheduler between the ack barrier and election; retained a generation-tagged O(log N) tree with fixed thread-0 coordination after packed-atomic, no-root-broadcast, and rotating-coordinator experiments; removed the invasive runtime hook after hook-free UT/ST coverage was established - [2026-07 — Host worker dispatch latency: where the remaining ~50 µs goes](2026-07-host-dispatch-latency-budget.md) — measured & dropped: after #1499 a `Worker.run()` costs **43.3 µs fixed + 8.08 µs per task**, so the "94% of latency is our runtime, not the IPC" headline is an artifact of benchmarking a *single* task — at 256 tasks/run the fixed cost is 2%. Also settles #1498 against itself for latency: a pipe wake measures 4.4 µs one way against a 3.5 µs bare fork+shm round trip, so blocking *adds* more than the whole IPC floor to every dispatch; only the CPU-while-idle argument survives 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 c1f18893c4..e5c008b912 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -422,6 +422,50 @@ bool upload_graph_submissions( 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; + const auto *definition = reinterpret_cast(entry.data); + if (definition->total_bytes != entry.bytes || definition->full_key != entry.full_key) continue; + const size_t object_bytes = sizeof(GraphDefinitionHeader) + entry.bytes; + void *object = + api->acquire_graph_definition_buffer(entry.full_key, object_bytes, alignof(GraphDefinitionHeader)); + if (object == nullptr) { + LOG_ERROR( + "host-orch: failed to retain %zu bytes for Graph Definition key=%#llx", object_bytes, + static_cast(entry.full_key) + ); + return false; + } + std::vector staging(object_bytes, std::byte{0}); + auto *header = reinterpret_cast(staging.data()); + header->magic = GRAPH_DEFINITION_OBJECT_MAGIC; + header->verify_state.store( + static_cast(GraphDefinitionVerifyState::UPLOADED), std::memory_order_relaxed + ); + header->definition_bytes = static_cast(entry.bytes); + header->content_hash = definition->content_hash; + header->full_key = definition->full_key; + std::memcpy(staging.data() + sizeof(GraphDefinitionHeader), entry.data, entry.bytes); + if (api->copy_to_device(object, staging.data(), object_bytes) != 0) { + LOG_ERROR("host-orch: failed to upload Graph Definition object"); + return false; + } + definition_objects.emplace(definition->content_hash, UploadedDefinition{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 || @@ -435,13 +479,21 @@ bool upload_graph_submissions( LOG_ERROR("host-orch: Graph submission size does not match its POD image"); return false; } - const GraphDefinition *definition = graph_submission_definition(*submission); + 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; + } + // Capacities come from the host-side Definition image the device + // object was built from; the GM object itself is never dereferenced + // on the host. + const GraphDefinition *definition = object_it->second.host_view; size_t execution_bytes = 0; - if (definition == nullptr || definition->full_key != submission->graph_key || definition->task_count == 0 || - definition->task_count > GRAPH_MAX_NODES || + if (definition->task_count == 0 || definition->task_count > GRAPH_MAX_NODES || + definition->full_key != submission->graph_key || !graph_execution_storage_bytes( static_cast(definition->task_count), definition->tensor_arg_count, - definition->scalar_arg_count, definition->total_bytes, &execution_bytes + definition->scalar_arg_count, &execution_bytes )) { LOG_ERROR("host-orch: invalid Graph execution storage request"); return false; @@ -457,6 +509,7 @@ bool upload_graph_submissions( ); return false; } + submission->definition_addr = reinterpret_cast(object_it->second.device_object); submission->execution_storage = reinterpret_cast(execution_storage); submission->execution_storage_bytes = execution_bytes; submission->local_execution = 0; 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 90ac9be666..5a080dc6bd 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 @@ -686,6 +686,18 @@ std::optional graph_host_upload(GraphHostState &state, size_t i return GraphHostUpload{upload.outer_slot, upload.image.data(), upload.image.size()}; } +GraphHostDefinitionList graph_host_definitions(GraphHostState &state) { + GraphHostDefinitionList list; + list.entries.reserve(state.definitions.size()); + for (auto &[key, image] : state.definitions) { + const GraphDefinition *header = graph_definition(image); + if (header != nullptr && header->total_bytes == image.size()) { + list.entries.push_back(GraphHostDefinition{key, image.data(), image.size()}); + } + } + return list; +} + static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { uint32_t next = orch->fanin_seen_current_epoch + 1; if (next == 0) { @@ -1401,10 +1413,9 @@ bool graph_build_submission_image( const std::vector &definition_image, const CoreTaskArgs &args, std::vector *submission_image ) { if (submission_image == nullptr || graph_definition(definition_image) == nullptr) return false; - const size_t definition_offset = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphDefinition)); - const size_t tensors_offset = PTO2_ALIGN_UP(definition_offset + definition_image.size(), alignof(GraphTensor)); + const size_t tensors_offset = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphTensor)); const size_t tensor_bytes = static_cast(args.tensor_count()) * sizeof(GraphTensor); - if (definition_offset > UINT32_MAX || 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; @@ -1416,7 +1427,6 @@ bool graph_build_submission_image( return false; } submission_image->assign(total_bytes, std::byte{0}); - std::memcpy(submission_image->data() + definition_offset, definition_image.data(), definition_image.size()); auto *tensors = reinterpret_cast(submission_image->data() + tensors_offset); for (int32_t i = 0; i < args.tensor_count(); ++i) tensors[i] = graph_tensor_pack(args.tensor(i).ref()); @@ -1430,8 +1440,8 @@ bool graph_build_submission_image( 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.definition_offset = static_cast(definition_offset); submission.tensors_offset = static_cast(tensors_offset); submission.tensor_count = static_cast(args.tensor_count()); submission.scalars_offset = static_cast(scalars_offset); 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 7f73a015af..1a24509c68 100644 --- a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp @@ -467,6 +467,50 @@ bool upload_graph_submissions( 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; + const auto *definition = reinterpret_cast(entry.data); + if (definition->total_bytes != entry.bytes || definition->full_key != entry.full_key) continue; + const size_t object_bytes = sizeof(GraphDefinitionHeader) + entry.bytes; + void *object = + api->acquire_graph_definition_buffer(entry.full_key, object_bytes, alignof(GraphDefinitionHeader)); + if (object == nullptr) { + LOG_ERROR( + "host-orch: failed to retain %zu bytes for Graph Definition key=%#llx", object_bytes, + static_cast(entry.full_key) + ); + return false; + } + std::vector staging(object_bytes, std::byte{0}); + auto *header = reinterpret_cast(staging.data()); + header->magic = GRAPH_DEFINITION_OBJECT_MAGIC; + header->verify_state.store( + static_cast(GraphDefinitionVerifyState::UPLOADED), std::memory_order_relaxed + ); + header->definition_bytes = static_cast(entry.bytes); + header->content_hash = definition->content_hash; + header->full_key = definition->full_key; + std::memcpy(staging.data() + sizeof(GraphDefinitionHeader), entry.data, entry.bytes); + if (api->copy_to_device(object, staging.data(), object_bytes) != 0) { + LOG_ERROR("host-orch: failed to upload Graph Definition object"); + return false; + } + definition_objects.emplace(definition->content_hash, UploadedDefinition{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 || @@ -480,13 +524,21 @@ bool upload_graph_submissions( LOG_ERROR("host-orch: Graph submission size does not match its POD image"); return false; } - const GraphDefinition *definition = graph_submission_definition(*submission); + 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; + } + // Capacities come from the host-side Definition image the device + // object was built from; the GM object itself is never dereferenced + // on the host. + const GraphDefinition *definition = object_it->second.host_view; size_t execution_bytes = 0; - if (definition == nullptr || definition->full_key != submission->graph_key || definition->task_count == 0 || - definition->task_count > GRAPH_MAX_NODES || + if (definition->task_count == 0 || definition->task_count > GRAPH_MAX_NODES || + definition->full_key != submission->graph_key || !graph_execution_storage_bytes( static_cast(definition->task_count), definition->tensor_arg_count, - definition->scalar_arg_count, definition->total_bytes, &execution_bytes + definition->scalar_arg_count, &execution_bytes )) { LOG_ERROR("host-orch: invalid Graph execution storage request"); return false; @@ -502,6 +554,7 @@ bool upload_graph_submissions( ); return false; } + submission->definition_addr = reinterpret_cast(object_it->second.device_object); submission->execution_storage = reinterpret_cast(execution_storage); submission->execution_storage_bytes = execution_bytes; submission->local_execution = 0; 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 90ac9be666..5a080dc6bd 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 @@ -686,6 +686,18 @@ std::optional graph_host_upload(GraphHostState &state, size_t i return GraphHostUpload{upload.outer_slot, upload.image.data(), upload.image.size()}; } +GraphHostDefinitionList graph_host_definitions(GraphHostState &state) { + GraphHostDefinitionList list; + list.entries.reserve(state.definitions.size()); + for (auto &[key, image] : state.definitions) { + const GraphDefinition *header = graph_definition(image); + if (header != nullptr && header->total_bytes == image.size()) { + list.entries.push_back(GraphHostDefinition{key, image.data(), image.size()}); + } + } + return list; +} + static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { uint32_t next = orch->fanin_seen_current_epoch + 1; if (next == 0) { @@ -1401,10 +1413,9 @@ bool graph_build_submission_image( const std::vector &definition_image, const CoreTaskArgs &args, std::vector *submission_image ) { if (submission_image == nullptr || graph_definition(definition_image) == nullptr) return false; - const size_t definition_offset = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphDefinition)); - const size_t tensors_offset = PTO2_ALIGN_UP(definition_offset + definition_image.size(), alignof(GraphTensor)); + const size_t tensors_offset = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphTensor)); const size_t tensor_bytes = static_cast(args.tensor_count()) * sizeof(GraphTensor); - if (definition_offset > UINT32_MAX || 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; @@ -1416,7 +1427,6 @@ bool graph_build_submission_image( return false; } submission_image->assign(total_bytes, std::byte{0}); - std::memcpy(submission_image->data() + definition_offset, definition_image.data(), definition_image.size()); auto *tensors = reinterpret_cast(submission_image->data() + tensors_offset); for (int32_t i = 0; i < args.tensor_count(); ++i) tensors[i] = graph_tensor_pack(args.tensor(i).ref()); @@ -1430,8 +1440,8 @@ bool graph_build_submission_image( 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.definition_offset = static_cast(definition_offset); submission.tensors_offset = static_cast(tensors_offset); submission.tensor_count = static_cast(args.tensor_count()); submission.scalars_offset = static_cast(scalars_offset); diff --git a/src/common/host_build_graph/graph_execution.cpp b/src/common/host_build_graph/graph_execution.cpp index c29fbae0fd..713d90e16d 100644 --- a/src/common/host_build_graph/graph_execution.cpp +++ b/src/common/host_build_graph/graph_execution.cpp @@ -33,12 +33,10 @@ void reset_execution( size_t nodes_offset = 0; size_t tensor_patches_offset = 0; size_t scalar_patches_offset = 0; - size_t definition_offset = 0; size_t bytes = 0; if (!graph_execution_storage_layout( - execution->node_capacity, execution->tensor_patch_capacity, execution->scalar_patch_capacity, - execution->definition_capacity, &nodes_offset, &tensor_patches_offset, &scalar_patches_offset, - &definition_offset, &bytes + execution->node_capacity, execution->tensor_patch_capacity, execution->scalar_patch_capacity, &nodes_offset, + &tensor_patches_offset, &scalar_patches_offset, &bytes )) { return; } @@ -69,7 +67,6 @@ void reset_execution( reinterpret_cast(reinterpret_cast(execution) + tensor_patches_offset); execution->scalar_patches = reinterpret_cast(reinterpret_cast(execution) + scalar_patches_offset); - execution->definition_storage = reinterpret_cast(execution) + definition_offset; if (!execution->definition_affine_reuse) { execution->definition = nullptr; execution->fanin_offsets = nullptr; @@ -83,7 +80,7 @@ void reset_execution( bool reusable_execution_header_valid(const GraphExecution &execution, size_t storage_bytes) { if (execution.storage_magic != GRAPH_EXECUTION_STORAGE_MAGIC || execution.node_capacity <= 0 || - execution.node_capacity > static_cast(GRAPH_MAX_NODES) || execution.definition_capacity == 0 || + execution.node_capacity > static_cast(GRAPH_MAX_NODES) || execution.tensor_patch_capacity > GRAPH_MAX_NODES * MAX_TENSOR_ARGS || execution.scalar_patch_capacity > GRAPH_MAX_NODES * MAX_SCALAR_ARGS || execution.materialized_tensor_patch_count > execution.tensor_patch_capacity || @@ -97,25 +94,23 @@ bool reusable_execution_header_valid(const GraphExecution &execution, size_t sto size_t expected_bytes = 0; return graph_execution_storage_bytes( execution.node_capacity, execution.tensor_patch_capacity, execution.scalar_patch_capacity, - execution.definition_capacity, &expected_bytes + &expected_bytes ) && execution.allocation_bytes == expected_bytes && expected_bytes <= storage_bytes; } GraphExecution *acquire_host_execution_storage( GraphSubmission &submission, int32_t node_count, uint64_t graph_key, uint64_t definition_hash, - uint32_t tensor_patch_count, uint32_t scalar_patch_count, size_t definition_bytes + uint32_t tensor_patch_count, uint32_t scalar_patch_count ) { - if (node_count <= 0 || node_count > static_cast(GRAPH_MAX_NODES) || definition_bytes == 0 || - submission.execution_storage == 0 || submission.execution_storage_bytes > SIZE_MAX || + if (node_count <= 0 || node_count > static_cast(GRAPH_MAX_NODES) || submission.execution_storage == 0 || + submission.execution_storage_bytes > SIZE_MAX || submission.execution_storage % alignof(GraphNodeStorage) != 0) { return nullptr; } const size_t storage_bytes = static_cast(submission.execution_storage_bytes); size_t required_bytes = 0; - if (!graph_execution_storage_bytes( - node_count, tensor_patch_count, scalar_patch_count, definition_bytes, &required_bytes - ) || + if (!graph_execution_storage_bytes(node_count, tensor_patch_count, scalar_patch_count, &required_bytes) || required_bytes > storage_bytes) { return nullptr; } @@ -128,15 +123,13 @@ GraphExecution *acquire_host_execution_storage( if (has_existing_execution && !valid_header) return nullptr; const bool capacities_fit = valid_header && execution->node_capacity >= node_count && execution->tensor_patch_capacity >= tensor_patch_count && - execution->scalar_patch_capacity >= scalar_patch_count && - execution->definition_capacity >= definition_bytes; + execution->scalar_patch_capacity >= scalar_patch_count; if (!capacities_fit) { if (valid_header) destroy_execution_nodes(execution); execution = new (execution) GraphExecution{}; execution->node_capacity = node_count; execution->tensor_patch_capacity = tensor_patch_count; execution->scalar_patch_capacity = scalar_patch_count; - execution->definition_capacity = definition_bytes; execution->allocation_bytes = required_bytes; } reset_execution(execution, node_count, tensor_patch_count, scalar_patch_count, graph_key, definition_hash); @@ -243,18 +236,68 @@ bool bind_graph_topology(GraphExecution &execution) { return true; } -bool graph_definition_hash_matches(const GraphDefinition &definition) { - if (definition.content_hash == 0 || definition.total_bytes < sizeof(GraphDefinition)) return false; +bool graph_definition_hash_matches(const GraphDefinition &definition, uint32_t definition_bytes) { + if (definition_bytes < sizeof(GraphDefinition) || definition.total_bytes != definition_bytes || + definition.content_hash == 0) { + return false; + } constexpr size_t HASH_OFFSET = offsetof(GraphDefinition, content_hash); constexpr size_t HASH_END = HASH_OFFSET + sizeof(GraphDefinition::content_hash); const auto *bytes = reinterpret_cast(&definition); uint64_t hash = graph_hash_bytes(1469598103934665603ULL, bytes, HASH_OFFSET); const uint64_t zero_hash = 0; hash = graph_hash_bytes(hash, &zero_hash, sizeof(zero_hash)); - hash = graph_hash_bytes(hash, bytes + HASH_END, definition.total_bytes - HASH_END); + hash = graph_hash_bytes(hash, bytes + HASH_END, definition_bytes - HASH_END); return hash == definition.content_hash; } +// One-time integrity gate for a shared Definition object. The first localizer +// wins the UPLOADED->VERIFYING CAS, hashes the image once, and publishes +// VERIFIED/INVALID. Localizers of the other submissions sharing this object +// spin on the state word — never re-hashing, never failing while a peer is +// mid-verify (the verify is bounded by the image size, so the spin is short). +GraphDefinition *graph_definition_object_verified(GraphDefinitionHeader &header) { + if (header.magic != GRAPH_DEFINITION_OBJECT_MAGIC) return nullptr; + if (header.definition_bytes < sizeof(GraphDefinition)) { + header.verify_state.store( + static_cast(GraphDefinitionVerifyState::INVALID), std::memory_order_release + ); + return nullptr; + } + auto *definition = reinterpret_cast(&header + 1); + uint32_t observed = header.verify_state.load(std::memory_order_acquire); + if (observed == static_cast(GraphDefinitionVerifyState::VERIFIED)) { + return definition; + } + if (observed == static_cast(GraphDefinitionVerifyState::INVALID)) return nullptr; + uint32_t expected = static_cast(GraphDefinitionVerifyState::UPLOADED); + if (!header.verify_state.compare_exchange_strong( + expected, static_cast(GraphDefinitionVerifyState::VERIFYING), std::memory_order_acq_rel, + std::memory_order_acquire + )) { + // A peer is verifying; wait for its verdict rather than racing it. + while (true) { + observed = header.verify_state.load(std::memory_order_acquire); + if (observed == static_cast(GraphDefinitionVerifyState::VERIFIED)) { + return definition; + } + if (observed == static_cast(GraphDefinitionVerifyState::INVALID)) return nullptr; +#if defined(__aarch64__) + __asm__ volatile("yield"); +#elif defined(__x86_64__) + __builtin_ia32_pause(); +#endif + } + } + const bool matched = header.content_hash == definition->content_hash && header.full_key == definition->full_key && + graph_definition_hash_matches(*definition, header.definition_bytes); + header.verify_state.store( + static_cast(matched ? GraphDefinitionVerifyState::VERIFIED : GraphDefinitionVerifyState::INVALID), + std::memory_order_release + ); + return matched ? definition : nullptr; +} + } // namespace GraphExecution *graph_execution_localize(PTO2TaskSlotState &outer_slot) { @@ -263,26 +306,27 @@ GraphExecution *graph_execution_localize(PTO2TaskSlotState &outer_slot) { if (GraphExecution *existing = graph_submission_local_execution(*submission)) return existing; if (graph_submission_execution_initializing(*submission)) return nullptr; - const GraphDefinition *definition = graph_submission_definition(*submission); + // The Definition lives in its own shared GM object; only its header is + // inspected here. The content hash gate runs inside the submission CAS + // below, so exactly one localizer verifies a shared object — concurrent + // localizers of other submissions see INITIALIZING and retry as BUSY + // rather than racing the verify state. + if (submission->definition_addr == 0 || submission->definition_addr % alignof(GraphDefinitionHeader) != 0) { + return nullptr; + } + auto *definition_header = + reinterpret_cast(static_cast(submission->definition_addr)); + if (definition_header->magic != GRAPH_DEFINITION_OBJECT_MAGIC) return nullptr; const GraphTensor *boundary_tensors = graph_submission_tensors(*submission); const uint64_t *boundary_scalars = graph_submission_scalars(*submission); const size_t boundary_tensor_end = static_cast(submission->tensors_offset) + static_cast(submission->tensor_count) * sizeof(GraphTensor); - if (definition == nullptr || definition->total_bytes == 0 || definition->task_count == 0 || - definition->task_count > GRAPH_MAX_NODES || - definition->total_bytes > submission->total_bytes - submission->definition_offset || - submission->graph_key != definition->full_key || submission->tensor_count != definition->boundary_count || - submission->scalar_count != definition->boundary_scalar_count || boundary_tensors == nullptr || + if (boundary_tensors == nullptr || outer_slot.task == nullptr || outer_slot.task->packed_buffer_base == nullptr || + outer_slot.task->packed_buffer_end == nullptr || (submission->scalar_count != 0 && boundary_scalars == nullptr) || - (submission->scalar_count != 0 && submission->scalars_offset < boundary_tensor_end) || - submission->tensors_offset < submission->definition_offset + definition->total_bytes || - !graph_definition_hash_matches(*definition) || outer_slot.task == nullptr || - outer_slot.task->packed_buffer_base == nullptr || outer_slot.task->packed_buffer_end == nullptr) { + (submission->scalar_count != 0 && submission->scalars_offset < boundary_tensor_end)) { return nullptr; } - const uintptr_t outer_base = reinterpret_cast(outer_slot.task->packed_buffer_base); - const uintptr_t outer_end = reinterpret_cast(outer_slot.task->packed_buffer_end); - if (outer_end < outer_base || definition->required_heap > outer_end - outer_base) return nullptr; for (uint32_t i = 0; i < submission->tensor_count; ++i) { if (!graph_tensor_wire_valid(boundary_tensors[i])) return nullptr; } @@ -297,9 +341,24 @@ GraphExecution *graph_execution_localize(PTO2TaskSlotState &outer_slot) { nullptr; } + const GraphDefinition *definition = graph_definition_object_verified(*definition_header); + if (definition == nullptr || definition->total_bytes == 0 || definition->task_count == 0 || + definition->task_count > GRAPH_MAX_NODES || submission->definition_hash != definition->content_hash || + submission->graph_key != definition->full_key || submission->tensor_count != definition->boundary_count || + submission->scalar_count != definition->boundary_scalar_count) { + __atomic_store_n(&submission->local_execution, 0, __ATOMIC_RELEASE); + return nullptr; + } + const uintptr_t outer_base = reinterpret_cast(outer_slot.task->packed_buffer_base); + const uintptr_t outer_end = reinterpret_cast(outer_slot.task->packed_buffer_end); + if (outer_end < outer_base || definition->required_heap > outer_end - outer_base) { + __atomic_store_n(&submission->local_execution, 0, __ATOMIC_RELEASE); + return nullptr; + } + GraphExecution *execution = acquire_host_execution_storage( *submission, static_cast(definition->task_count), submission->graph_key, definition->content_hash, - definition->tensor_arg_count, definition->scalar_arg_count, definition->total_bytes + definition->tensor_arg_count, definition->scalar_arg_count ); if (execution == nullptr) { __atomic_store_n(&submission->local_execution, 0, __ATOMIC_RELEASE); @@ -307,8 +366,9 @@ GraphExecution *graph_execution_localize(PTO2TaskSlotState &outer_slot) { } if (!execution->definition_affine_reuse) { - std::memcpy(execution->definition_storage, definition, definition->total_bytes); - execution->definition = static_cast(execution->definition_storage); + // The Definition is read in place from the shared GM object; no + // per-occurrence copy is taken. + execution->definition = definition; if (!bind_graph_topology(*execution)) { execution->retired_nodes.store(execution->node_count, std::memory_order_relaxed); execution->state.store(GraphExecutionState::COMPLETED, std::memory_order_release); diff --git a/src/common/host_build_graph/graph_execution.h b/src/common/host_build_graph/graph_execution.h index 9c3390f432..98e05a5919 100644 --- a/src/common/host_build_graph/graph_execution.h +++ b/src/common/host_build_graph/graph_execution.h @@ -101,6 +101,29 @@ struct GraphBoundarySignature { uint8_t reserved; }; +// Header prefixing each device-resident Definition object. The definition +// buffer uploaded by the host is [GraphDefinitionHeader][GraphDefinition image]; +// verify_state gates the one-time integrity hash so submissions sharing one +// Definition each pay only a state load. +inline constexpr uint64_t GRAPH_DEFINITION_OBJECT_MAGIC = 0x4752415048455844ULL; + +enum class GraphDefinitionVerifyState : uint32_t { + UPLOADED = 0, + VERIFYING = 1, + VERIFIED = 2, + INVALID = 3, +}; + +struct GraphDefinitionHeader { + uint64_t magic; + std::atomic verify_state; + uint32_t definition_bytes; + uint64_t content_hash; + uint64_t full_key; +}; + +static_assert(std::is_standard_layout_v); + struct GraphDefinition { uint64_t full_key; uint64_t content_hash; @@ -132,9 +155,13 @@ struct GraphSubmission { uint64_t execution_storage; uint64_t execution_storage_bytes; uint64_t local_execution; + // Device GM address of the shared Definition object this replay references + // (an integer-typed absolute address per the wire rules) plus the content + // hash the host computed for it. + uint64_t definition_addr; + uint64_t definition_hash; uint32_t activation_gate; uint32_t total_bytes; - uint32_t definition_offset; uint32_t tensors_offset; uint32_t tensor_count; uint32_t scalars_offset; @@ -238,22 +265,6 @@ inline GraphSubmission *graph_submission_from_slot(PTO2TaskSlotState &slot) { return slot.task_kind == TaskKind::GRAPH ? static_cast(slot.graph_context) : nullptr; } -inline const GraphDefinition *graph_submission_definition(const GraphSubmission &submission) { - if (submission.definition_offset == 0 || submission.definition_offset % alignof(GraphDefinition) != 0 || - submission.definition_offset > submission.total_bytes || - sizeof(GraphDefinition) > submission.total_bytes - submission.definition_offset) { - return nullptr; - } - const auto *definition = reinterpret_cast( - reinterpret_cast(&submission) + submission.definition_offset - ); - if (definition->total_bytes < sizeof(GraphDefinition) || - definition->total_bytes > submission.total_bytes - submission.definition_offset) { - return nullptr; - } - return definition; -} - inline bool graph_submission_wire_size_valid(const GraphSubmission &submission, size_t available_bytes) { return available_bytes >= sizeof(GraphSubmission) && submission.total_bytes == available_bytes; } @@ -361,7 +372,6 @@ struct GraphExecution { uint32_t materialized_scalar_patches{0}; uint32_t materialized_scalar_patch_count{0}; size_t allocation_bytes{0}; - size_t definition_capacity{0}; uint64_t graph_key{0}; uint64_t definition_hash{0}; uint64_t materialized_graph_key{0}; @@ -373,7 +383,6 @@ struct GraphExecution { GraphNodeStorage *node_storage{nullptr}; GraphTensorAddressPatch *tensor_patches{nullptr}; GraphScalarPatch *scalar_patches{nullptr}; - void *definition_storage{nullptr}; const GraphDefinition *definition{nullptr}; const uint32_t *fanin_offsets{nullptr}; const uint16_t *fanin_indices{nullptr}; @@ -389,12 +398,11 @@ static_assert(std::is_trivially_destructible_v); static_assert(std::is_trivially_destructible_v); inline bool graph_execution_storage_layout( - int32_t node_capacity, uint32_t tensor_patch_capacity, uint32_t scalar_patch_capacity, size_t definition_capacity, - size_t *nodes_offset, size_t *tensor_patches_offset, size_t *scalar_patches_offset, size_t *definition_offset, - size_t *storage_bytes + int32_t node_capacity, uint32_t tensor_patch_capacity, uint32_t scalar_patch_capacity, size_t *nodes_offset, + size_t *tensor_patches_offset, size_t *scalar_patches_offset, size_t *storage_bytes ) { if (nodes_offset == nullptr || tensor_patches_offset == nullptr || scalar_patches_offset == nullptr || - definition_offset == nullptr || storage_bytes == nullptr || node_capacity <= 0 || + storage_bytes == nullptr || node_capacity <= 0 || static_cast(node_capacity) > SIZE_MAX / sizeof(GraphNodeStorage) || tensor_patch_capacity > GRAPH_MAX_NODES * MAX_TENSOR_ARGS || scalar_patch_capacity > GRAPH_MAX_NODES * MAX_SCALAR_ARGS) { @@ -415,25 +423,21 @@ inline bool graph_execution_storage_layout( !checked_align_up( *tensor_patches_offset + tensor_patches_bytes, alignof(GraphScalarPatch), scalar_patches_offset ) || - *scalar_patches_offset > SIZE_MAX - scalar_patches_bytes || - !checked_align_up(*scalar_patches_offset + scalar_patches_bytes, alignof(GraphDefinition), definition_offset) || - *definition_offset > SIZE_MAX - definition_capacity) { + *scalar_patches_offset > SIZE_MAX - scalar_patches_bytes) { return false; } - return checked_align_up(*definition_offset + definition_capacity, alignof(GraphNodeStorage), storage_bytes); + return checked_align_up(*scalar_patches_offset + scalar_patches_bytes, alignof(GraphNodeStorage), storage_bytes); } inline bool graph_execution_storage_bytes( - int32_t node_capacity, uint32_t tensor_patch_capacity, uint32_t scalar_patch_capacity, size_t definition_capacity, - size_t *storage_bytes + int32_t node_capacity, uint32_t tensor_patch_capacity, uint32_t scalar_patch_capacity, size_t *storage_bytes ) { size_t nodes_offset = 0; size_t tensor_patches_offset = 0; size_t scalar_patches_offset = 0; - size_t definition_offset = 0; return graph_execution_storage_layout( - node_capacity, tensor_patch_capacity, scalar_patch_capacity, definition_capacity, &nodes_offset, - &tensor_patches_offset, &scalar_patches_offset, &definition_offset, storage_bytes + node_capacity, tensor_patch_capacity, scalar_patch_capacity, &nodes_offset, &tensor_patches_offset, + &scalar_patches_offset, storage_bytes ); } diff --git a/src/common/host_build_graph/graph_host_state.h b/src/common/host_build_graph/graph_host_state.h index 4d869b1c4d..265f70907c 100644 --- a/src/common/host_build_graph/graph_host_state.h +++ b/src/common/host_build_graph/graph_host_state.h @@ -12,8 +12,10 @@ #pragma once #include +#include #include #include +#include struct PTO2TaskSlotState; struct GraphHostState; @@ -32,6 +34,19 @@ struct GraphHostUpload { size_t bytes; }; +// 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 { + uint64_t full_key; + const std::byte *data; + size_t bytes; +}; + +struct GraphHostDefinitionList { + std::vector entries; +}; + 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); diff --git a/src/common/platform/include/common/host_api.h b/src/common/platform/include/common/host_api.h index 7319153e5b..2a3b35e977 100644 --- a/src/common/platform/include/common/host_api.h +++ b/src/common/platform/include/common/host_api.h @@ -62,6 +62,16 @@ struct HostApiOps { void *runner_ctx, uint32_t pipeline_slot, uint64_t graph_key, uint32_t occurrence, size_t bytes, size_t alignment ); + // Runner-owned Graph Definition storage, keyed by the Definition's content + // identity (full_key, content_hash, total_bytes folded into one key by the + // caller). Same grow-only retention contract as + // acquire_graph_execution_buffer: one block per key, reused while capacity + // fits, a growth request replaces the entry, all blocks released at Worker + // finalization. Lets every submission of one run reference a single + // device-resident Definition instead of each carrying a full copy. + void *(*acquire_graph_definition_buffer)( + void *runner_ctx, uint32_t pipeline_slot, uint64_t key, 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 @@ -165,6 +175,10 @@ struct HostApi { runner_ctx_, pipeline_slot_, graph_key, occurrence, bytes, alignment ); } + void *acquire_graph_definition_buffer(uint64_t key, size_t bytes, size_t alignment) const { + if (ops_->acquire_graph_definition_buffer == nullptr) return nullptr; + return ops_->acquire_graph_definition_buffer(runner_ctx_, pipeline_slot_, key, 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 8daaea45a7..265a2ba21f 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -168,6 +168,18 @@ static void *acquire_graph_execution_buffer( } } +static void *acquire_graph_definition_buffer( + void *runner_ctx, uint32_t pipeline_slot, uint64_t key, size_t bytes, size_t alignment +) { + if (runner_ctx == nullptr) return nullptr; + try { + return static_cast(runner_ctx) + ->acquire_graph_definition_buffer(pipeline_slot, key, 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 = { .get_retained_temp_buffer = get_retained_temp_buffer, .set_retained_temp_buffer = set_retained_temp_buffer, .acquire_graph_execution_buffer = acquire_graph_execution_buffer, + .acquire_graph_definition_buffer = acquire_graph_definition_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 3cc97d49bb..e4c929ec36 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -207,6 +207,40 @@ void *DeviceRunnerBase::acquire_graph_execution_buffer( return aligned_addr; } +void *DeviceRunnerBase::acquire_graph_definition_buffer( + uint32_t pipeline_slot, uint64_t key, size_t bytes, size_t alignment +) { + if (pipeline_slot >= graph_definition_buffers_.size() || bytes == 0 || alignment == 0 || + (alignment & (alignment - 1)) != 0 || bytes > SIZE_MAX - (alignment - 1)) { + return nullptr; + } + RetainedGraphExecutionBuffer &buffer = graph_definition_buffers_[pipeline_slot][key]; + 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::release_graph_execution_buffers() { for (GraphExecutionBufferMap &by_key : graph_execution_buffers_) { for (auto &entry : by_key) { @@ -216,12 +250,21 @@ void DeviceRunnerBase::release_graph_execution_buffers() { } by_key.clear(); } + for (GraphDefinitionBufferMap &by_key : graph_definition_buffers_) { + for (auto &entry : by_key) { + if (entry.second.allocation != nullptr) mem_alloc_.free(entry.second.allocation); + } + by_key.clear(); + } } void DeviceRunnerBase::abandon_graph_execution_buffers() { for (GraphExecutionBufferMap &by_key : graph_execution_buffers_) { by_key.clear(); } + for (GraphDefinitionBufferMap &by_key : graph_definition_buffers_) { + by_key.clear(); + } } void DeviceRunnerBase::clear_temporary_buffer() { diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index 050ed4fdf4..bb8d26b200 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -147,6 +147,8 @@ class DeviceRunnerBase { void *acquire_graph_execution_buffer( uint32_t pipeline_slot, uint64_t graph_key, uint32_t occurrence, std::size_t bytes, std::size_t alignment ); + void * + acquire_graph_definition_buffer(uint32_t pipeline_slot, uint64_t key, std::size_t bytes, std::size_t alignment); void clear_temporary_buffer(); /** * Map a device buffer into the host address space and return a @@ -1081,6 +1083,12 @@ class DeviceRunnerBase { }; using GraphExecutionBufferMap = std::unordered_map>; std::array graph_execution_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 + // references the same device-resident Definition. + using GraphDefinitionBufferMap = std::unordered_map; + std::array graph_definition_buffers_{}; // One independently committed set of the three pooled device regions. A // run reaches its set through the arena bank its lease selects, so diff --git a/src/common/platform/sim/host/c_api_shared.cpp b/src/common/platform/sim/host/c_api_shared.cpp index 39ee8e3846..0a1f9d98fe 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -154,6 +154,18 @@ static void *acquire_graph_execution_buffer( } } +static void *acquire_graph_definition_buffer( + void *runner_ctx, uint32_t pipeline_slot, uint64_t key, size_t bytes, size_t alignment +) { + if (runner_ctx == nullptr) return nullptr; + try { + return static_cast(runner_ctx) + ->acquire_graph_definition_buffer(pipeline_slot, key, 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 +281,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_execution_buffer = acquire_graph_execution_buffer, + .acquire_graph_definition_buffer = acquire_graph_definition_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 781e4bc484..7e8d55d854 100644 --- a/src/common/platform/sim/host/device_runner_base.cpp +++ b/src/common/platform/sim/host/device_runner_base.cpp @@ -410,6 +410,40 @@ void *SimDeviceRunnerBase::acquire_graph_execution_buffer( return aligned_addr; } +void *SimDeviceRunnerBase::acquire_graph_definition_buffer( + uint32_t pipeline_slot, uint64_t key, size_t bytes, size_t alignment +) { + if (pipeline_slot >= graph_definition_buffers_.size() || bytes == 0 || alignment == 0 || + (alignment & (alignment - 1)) != 0 || bytes > SIZE_MAX - (alignment - 1)) { + return nullptr; + } + RetainedGraphExecutionBuffer &buffer = graph_definition_buffers_[pipeline_slot][key]; + 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::release_graph_execution_buffers() { for (GraphExecutionBufferMap &by_key : graph_execution_buffers_) { for (auto &entry : by_key) { @@ -419,6 +453,12 @@ void SimDeviceRunnerBase::release_graph_execution_buffers() { } by_key.clear(); } + for (GraphDefinitionBufferMap &by_key : graph_definition_buffers_) { + for (auto &entry : by_key) { + if (entry.second.allocation != nullptr) mem_alloc_.free(entry.second.allocation); + } + by_key.clear(); + } } void SimDeviceRunnerBase::clear_temporary_buffer() { diff --git a/src/common/platform/sim/host/device_runner_base.h b/src/common/platform/sim/host/device_runner_base.h index fb7758b9f6..8a6b84db71 100644 --- a/src/common/platform/sim/host/device_runner_base.h +++ b/src/common/platform/sim/host/device_runner_base.h @@ -197,6 +197,7 @@ class SimDeviceRunnerBase { void *acquire_graph_execution_buffer( 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 clear_temporary_buffer(); // On sim, allocate_tensor returns a plain host pointer, so the "device" @@ -343,6 +344,10 @@ class SimDeviceRunnerBase { }; using GraphExecutionBufferMap = std::unordered_map>; std::array graph_execution_buffers_{}; + // Graph Definition storage, one retained block per (pipeline slot, + // definition key) — see HostApi acquire_graph_definition_buffer. + using GraphDefinitionBufferMap = std::unordered_map; + std::array graph_definition_buffers_{}; // Each arena bank backs the three pooled regions (PTO2 GM heap / PTO2 // shared memory / trb prebuilt runtime arena) for one pipeline slot. They diff --git a/tests/ut/cpp/common/test_hbg_graph_cache.cpp b/tests/ut/cpp/common/test_hbg_graph_cache.cpp index 86b55c6ef4..0a9121f162 100644 --- a/tests/ut/cpp/common/test_hbg_graph_cache.cpp +++ b/tests/ut/cpp/common/test_hbg_graph_cache.cpp @@ -120,12 +120,9 @@ std::vector make_test_submission( uint64_t graph_key, uint64_t boundary_address, uint64_t boundary_scalar, uint64_t execution_storage, size_t execution_storage_bytes ) { - const std::vector definition = make_test_definition(graph_key, boundary_address); - const size_t definition_offset = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphDefinition)); - const size_t tensors_offset = PTO2_ALIGN_UP(definition_offset + definition.size(), alignof(GraphTensor)); + const size_t tensors_offset = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphTensor)); const size_t scalars_offset = PTO2_ALIGN_UP(tensors_offset + sizeof(GraphTensor), alignof(uint64_t)); std::vector image(scalars_offset + sizeof(uint64_t)); - std::memcpy(image.data() + definition_offset, definition.data(), definition.size()); const GraphTensor boundary = make_test_tensor(boundary_address); std::memcpy(image.data() + tensors_offset, &boundary, sizeof(boundary)); std::memcpy(image.data() + scalars_offset, &boundary_scalar, sizeof(boundary_scalar)); @@ -135,7 +132,6 @@ std::vector make_test_submission( submission.execution_storage = execution_storage; submission.execution_storage_bytes = execution_storage_bytes; submission.total_bytes = static_cast(image.size()); - submission.definition_offset = static_cast(definition_offset); submission.tensors_offset = static_cast(tensors_offset); submission.tensor_count = 1; submission.scalars_offset = static_cast(scalars_offset); @@ -144,12 +140,49 @@ std::vector make_test_submission( return image; } +// A Definition device object exactly as upload_graph_submissions builds it: +// [GraphDefinitionHeader][Definition image]. +class TestDefinitionObject { +public: + explicit TestDefinitionObject(const std::vector &definition, uint32_t retained_definition_bytes = 0) { + const size_t object_bytes = sizeof(GraphDefinitionHeader) + definition.size(); + data_ = ::operator new(object_bytes, std::align_val_t(alignof(GraphDefinitionHeader))); + std::memset(data_, 0, object_bytes); + auto *header = static_cast(data_); + const auto *def = reinterpret_cast(definition.data()); + header->magic = GRAPH_DEFINITION_OBJECT_MAGIC; + header->verify_state.store( + static_cast(GraphDefinitionVerifyState::UPLOADED), std::memory_order_relaxed + ); + header->definition_bytes = + retained_definition_bytes == 0 ? static_cast(definition.size()) : retained_definition_bytes; + header->content_hash = def->content_hash; + header->full_key = def->full_key; + std::memcpy( + static_cast(data_) + sizeof(GraphDefinitionHeader), definition.data(), definition.size() + ); + } + + ~TestDefinitionObject() { ::operator delete(data_, std::align_val_t(alignof(GraphDefinitionHeader))); } + + uint64_t address() const { return reinterpret_cast(data_); } + uint64_t hash() const { return static_cast(data_)->content_hash; } + GraphDefinitionVerifyState verify_state() const { + return static_cast( + static_cast(data_)->verify_state.load(std::memory_order_acquire) + ); + } + +private: + void *data_{nullptr}; +}; + class AlignedStorage { public: - explicit AlignedStorage(size_t bytes) : + explicit AlignedStorage(size_t bytes, uint8_t fill = 0) : bytes_(bytes) { data_ = ::operator new(bytes, std::align_val_t(alignof(GraphNodeStorage))); - std::memset(data_, 0, bytes); + std::memset(data_, fill, bytes); } ~AlignedStorage() { ::operator delete(data_, std::align_val_t(alignof(GraphNodeStorage))); } @@ -222,35 +255,32 @@ TEST(GraphScalarProvenance, MutableAccessInvalidatesForwardedSource) { TEST(GraphExecutionStorage, ComputesAlignedExactSize) { constexpr int32_t NODE_COUNT = 7; - constexpr size_t DEFINITION_BYTES = 321; constexpr uint32_t TENSOR_PATCH_COUNT = 11; constexpr uint32_t SCALAR_PATCH_COUNT = 5; size_t nodes_offset = 0; size_t tensor_patches_offset = 0; size_t scalar_patches_offset = 0; - size_t definition_offset = 0; size_t storage_bytes = 0; ASSERT_TRUE(graph_execution_storage_layout( - NODE_COUNT, TENSOR_PATCH_COUNT, SCALAR_PATCH_COUNT, DEFINITION_BYTES, &nodes_offset, &tensor_patches_offset, - &scalar_patches_offset, &definition_offset, &storage_bytes + NODE_COUNT, TENSOR_PATCH_COUNT, SCALAR_PATCH_COUNT, &nodes_offset, &tensor_patches_offset, + &scalar_patches_offset, &storage_bytes )); EXPECT_EQ(nodes_offset % alignof(GraphNodeStorage), 0U); EXPECT_EQ(tensor_patches_offset % alignof(GraphTensorAddressPatch), 0U); EXPECT_EQ(scalar_patches_offset % alignof(GraphScalarPatch), 0U); - EXPECT_EQ(definition_offset % alignof(GraphDefinition), 0U); EXPECT_GE(tensor_patches_offset, nodes_offset + NODE_COUNT * sizeof(GraphNodeStorage)); EXPECT_GE(scalar_patches_offset, tensor_patches_offset + TENSOR_PATCH_COUNT * sizeof(GraphTensorAddressPatch)); - EXPECT_GE(definition_offset, scalar_patches_offset + SCALAR_PATCH_COUNT * sizeof(GraphScalarPatch)); - EXPECT_GE(storage_bytes, definition_offset + DEFINITION_BYTES); + EXPECT_GE(storage_bytes, scalar_patches_offset + SCALAR_PATCH_COUNT * sizeof(GraphScalarPatch)); EXPECT_EQ(storage_bytes % alignof(GraphNodeStorage), 0U); } TEST(GraphExecutionStorage, RejectsInvalidCapacity) { size_t storage_bytes = 0; - EXPECT_FALSE(graph_execution_storage_bytes(0, 0, 0, sizeof(GraphDefinition), &storage_bytes)); - EXPECT_FALSE(graph_execution_storage_bytes(1, 0, 0, SIZE_MAX, &storage_bytes)); + EXPECT_FALSE(graph_execution_storage_bytes(0, 0, 0, &storage_bytes)); + EXPECT_FALSE(graph_execution_storage_bytes(-1, 0, 0, &storage_bytes)); + EXPECT_FALSE(graph_execution_storage_bytes(1, GRAPH_MAX_NODES * MAX_TENSOR_ARGS + 1, 0, &storage_bytes)); } TEST(GraphExecutionReplay, AffineHitRefreshesOnlyDynamicFields) { @@ -262,14 +292,19 @@ TEST(GraphExecutionReplay, AffineHitRefreshesOnlyDynamicFields) { const std::vector definition = make_test_definition(GRAPH_KEY_VALUE, reinterpret_cast(first_boundary.data())); + const TestDefinitionObject definition_object(definition); size_t execution_bytes = 0; - ASSERT_TRUE(graph_execution_storage_bytes(2, 2, 2, definition.size(), &execution_bytes)); - AlignedStorage execution_storage(execution_bytes); + ASSERT_TRUE(graph_execution_storage_bytes(2, 2, 2, &execution_bytes)); + // 0xAA fill stands in for the unzeroed GM a device-side allocation + // returns: nothing in localize/materialize may depend on a zeroed block. + AlignedStorage execution_storage(execution_bytes, 0xAA); std::vector submission_image = make_test_submission( GRAPH_KEY_VALUE, reinterpret_cast(first_boundary.data()), 17, reinterpret_cast(execution_storage.data()), execution_storage.size() ); auto &submission = *reinterpret_cast(submission_image.data()); + submission.definition_addr = definition_object.address(); + submission.definition_hash = definition_object.hash(); PTO2TaskDescriptor outer_task{}; outer_task.task_id = PTO2TaskId::make(1, 7); @@ -347,14 +382,36 @@ TEST(GraphExecutionReplay, AffineHitRefreshesOnlyDynamicFields) { EXPECT_EQ(execution->materialized_tensor_patches, 1U); } -TEST(GraphSubmissionWire, RejectsDefinitionBeyondSubmission) { - constexpr uint64_t GRAPH_KEY_VALUE = 0x3456; - std::vector image = make_test_submission(GRAPH_KEY_VALUE, 0x1000, 17, 0x2000, 4096); - auto &submission = *reinterpret_cast(image.data()); - auto *definition = reinterpret_cast(image.data() + submission.definition_offset); - definition->total_bytes = submission.total_bytes - submission.definition_offset + 1; +TEST(GraphDefinitionObject, RejectsDefinitionBeyondRetainedBytes) { + constexpr uint64_t GRAPH_KEY_VALUE = 0x4567; + std::array heap{}; + std::array boundary{}; + const std::vector definition = + make_test_definition(GRAPH_KEY_VALUE, reinterpret_cast(boundary.data())); + ASSERT_GT(definition.size(), sizeof(GraphDefinition)); + const TestDefinitionObject definition_object(definition, sizeof(GraphDefinition)); + size_t execution_bytes = 0; + ASSERT_TRUE(graph_execution_storage_bytes(2, 2, 2, &execution_bytes)); + AlignedStorage execution_storage(execution_bytes); + std::vector submission_image = make_test_submission( + GRAPH_KEY_VALUE, reinterpret_cast(boundary.data()), 17, + reinterpret_cast(execution_storage.data()), execution_storage.size() + ); + auto &submission = *reinterpret_cast(submission_image.data()); + submission.definition_addr = definition_object.address(); + submission.definition_hash = definition_object.hash(); - EXPECT_EQ(graph_submission_definition(submission), nullptr); + PTO2TaskDescriptor outer_task{}; + outer_task.task_id = PTO2TaskId::make(1, 7); + outer_task.packed_buffer_base = heap.data(); + outer_task.packed_buffer_end = heap.data() + heap.size(); + PTO2TaskSlotState outer_slot{}; + outer_slot.task_kind = TaskKind::GRAPH; + outer_slot.task = &outer_task; + outer_slot.graph_context = &submission; + + EXPECT_EQ(graph_execution_localize(outer_slot), nullptr); + EXPECT_EQ(definition_object.verify_state(), GraphDefinitionVerifyState::INVALID); } TEST(GraphSubmissionWire, RequiresExactAvailableSize) { @@ -485,3 +542,61 @@ TEST(GraphExecutionProgress, InternalNodeResolutionIsNotAHostCompletion) { EXPECT_EQ(result.resolved, 1); EXPECT_EQ(result.completed, 0); } + +// Device-side execution storage is not guaranteed to be zero-initialized. +// localize + materialize must produce a fully valid execution from arbitrary +// initial bytes. +TEST(GraphExecutionMaterialize, DirtyStorageYieldsValidExecution) { + constexpr uint64_t GRAPH_KEY_VALUE = 0x9753; + std::array heap{}; + std::array boundary{}; + + const std::vector definition = + make_test_definition(GRAPH_KEY_VALUE, reinterpret_cast(boundary.data())); + const TestDefinitionObject definition_object(definition); + size_t execution_bytes = 0; + ASSERT_TRUE(graph_execution_storage_bytes(2, 2, 2, &execution_bytes)); + AlignedStorage execution_storage(execution_bytes, 0xAA); + std::vector submission_image = make_test_submission( + GRAPH_KEY_VALUE, reinterpret_cast(boundary.data()), 17, + reinterpret_cast(execution_storage.data()), execution_storage.size() + ); + auto &submission = *reinterpret_cast(submission_image.data()); + submission.definition_addr = definition_object.address(); + submission.definition_hash = definition_object.hash(); + + PTO2TaskDescriptor outer_task{}; + outer_task.task_id = PTO2TaskId::make(1, 5); + outer_task.packed_buffer_base = heap.data(); + outer_task.packed_buffer_end = heap.data() + heap.size(); + PTO2TaskSlotState outer_slot{}; + outer_slot.task_kind = TaskKind::GRAPH; + outer_slot.task = &outer_task; + outer_slot.graph_context = &submission; + + GraphExecution *execution = graph_execution_localize(outer_slot); + ASSERT_NE(execution, nullptr); + EXPECT_EQ(graph_execution_materialize_slice(outer_slot, *execution, 2), GraphMaterializeResult::PREPARED); + + // Every observable scheduling field must be a materialize-written value, + // not the 0xAA fill: state machine, counters, atomics, and the patch + // tables all start from values only the device side wrote. + for (int32_t i = 0; i < execution->node_count; ++i) { + const GraphNodeStorage &node = execution->node_storage[i]; + ASSERT_EQ(node.slot.task_state.load(std::memory_order_relaxed), PTO2_TASK_PENDING); + ASSERT_EQ(node.slot.task_kind, TaskKind::GRAPH_NODE); + ASSERT_EQ(node.slot.completed_subtasks.load(std::memory_order_relaxed), 0); + ASSERT_EQ(node.payload.dispatch_fanin.load(std::memory_order_relaxed), 0); + ASSERT_EQ(node.payload.tensor_count, 1); + ASSERT_EQ(node.payload.scalar_count, 1); + // make_test_definition assigns node i the heap offset 64*i, so the + // packed window starts at outer_base + 64*i, not at outer_base. + ASSERT_EQ(node.task.packed_buffer_base, static_cast(heap.data() + static_cast(i) * 64)); + } + EXPECT_EQ(execution->materialized_nodes, execution->node_count); + for (uint32_t i = 0; i < execution->materialized_tensor_patches; ++i) { + // A patch address of 0xAAAAAAAAAAAAAAAA would mean the fill leaked + // through into a field the scheduler later dereferences. + EXPECT_NE(execution->tensor_patches[i].address_offset, 0xAAAAAAAAAAAAAAAAULL); + } +}