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); + } +}