From eca540b808c545966d7364a7591286b1e79628ca Mon Sep 17 00:00:00 2001 From: TaoZQY Date: Thu, 6 Aug 2026 21:51:08 -0700 Subject: [PATCH] Add: support dynamic Graph boundary scalars - preserve boundary scalar provenance through direct and copied args - serialize current scalar bindings and refresh them during Graph replay - allow unused boundary scalars without creating node patches - document that host-derived plain scalar values become static Definition data - keep construction config in the cache key and reject mutable tracked scalars Refs #1712 --- .../host_build_graph/docs/GRAPH_EXECUTION.md | 68 ++++++---- .../host_build_graph/host/runtime_maker.cpp | 4 +- .../orchestration/pto_orchestration_api.h | 4 +- .../host_build_graph/runtime/graph_cache.h | 4 +- .../runtime/graph_execution.h | 75 +++++++++-- .../orchestrator_core/pto_orchestrator.cpp | 117 ++++++++++++++++-- .../host_build_graph/runtime/pto_types.h | 51 +++++++- .../runtime/scheduler/graph_execution.cpp | 95 +++++++++++--- .../orchestration/graph_execution_orch.cpp | 28 +++-- .../graph_execution/test_graph_execution.py | 10 +- tests/ut/cpp/a2a3/test_graph_cache.cpp | 92 ++++++++++++-- 11 files changed, 455 insertions(+), 93 deletions(-) diff --git a/src/a2a3/runtime/host_build_graph/docs/GRAPH_EXECUTION.md b/src/a2a3/runtime/host_build_graph/docs/GRAPH_EXECUTION.md index 1ad1b99b24..d5a91cd912 100644 --- a/src/a2a3/runtime/host_build_graph/docs/GRAPH_EXECUTION.md +++ b/src/a2a3/runtime/host_build_graph/docs/GRAPH_EXECUTION.md @@ -9,7 +9,7 @@ places one `GRAPH` task in the host task window. The device Scheduler expands the saved topology and dispatches its internal nodes; the Host Orchestrator does not submit those nodes again. -## Step-1 API +## API A Graph uses `CoreTaskArgs`, the existing incore argument type: @@ -27,7 +27,7 @@ void graph_function(const CoreTaskArgs &args, int variant) { CoreTaskArgs matmul_args; matmul_args.add_input(input, weight); matmul_args.add_output(intermediate); - matmul_args.add_scalar(uint32_t{16}); // fixed Definition data + matmul_args.copy_scalars_from(args, 0, 1); // current invocation's value TaskOutputTensors matmul = rt_submit_aic_task( variant == 0 ? FUNC_MATMUL : FUNC_MATMUL_TRANSPOSED, matmul_args @@ -69,19 +69,39 @@ the same key for different functions can select the wrong recorded topology. There are no public `GraphArgs`, `GraphBindings`, `Patch`, or `ScalarRef` types. The boundary is represented by `CoreTaskArgs`. -## Supported dynamic and static data +Boundary scalars are pass-through bindings. Forward them directly with +`node_args.add_scalar(args.scalar(i))` or `copy_scalars_from(args, i, count)` +so recording can retain their source indices. + +Ordinary C++ value transformations do not retain boundary provenance. Both +`node_args.add_scalar(args.scalar(i) + 1)` and copying `args.scalar(i)` into a +local arithmetic variable before calling `add_scalar` produce an ordinary +static node scalar. That value is stored in the Definition, and later cache +hits reuse the first invocation's value without a warning. The runtime cannot +distinguish such a derived value from an intentional static literal after the +C++ expression has produced a plain arithmetic value. Compute the derived value +before constructing the Graph boundary and pass it as another boundary scalar, +perform the transformation in a kernel, or use a construction parameter when +the value changes the Graph structure. + +Access through a non-const `scalar()` invalidates inherited boundary provenance +conservatively, because returning a mutable reference cannot distinguish a read +from a later write. A Graph containing such an invalidated binding is not +cached, which prevents replay from silently replacing the transformed value +with the unmodified boundary value. -Step 1 deliberately supports a narrow, safe contract: +## Supported dynamic and static data - Boundary ChipTensor addresses may change for every invocation. +- Boundary scalar values may change for every invocation. Their count is fixed + by the recorded boundary contract. Unused boundary scalars are allowed and do + not create internal scalar patches. - A Graph boundary contains at least one ChipTensor. - Construction parameters are part of Graph identity and may control the function's task count, kernel selection, or other structural choices. - Boundary ChipTensor shape, stride, dtype, size, direction, contiguity, and alias partition must match the first invocation. -- Scalars inside internal task args are fixed Definition data. -- Scalars in the boundary `CoreTaskArgs` are not cacheable yet. Such a call uses - the ordinary task-submit path. +- Internal task scalars with no boundary source are fixed Definition data. - Boundary storage is caller-owned. `INPUT`, `INOUT`, `OUTPUT_EXISTING`, and `NO_DEP` are supported. A boundary `TensorCreateInfo` tagged `OUTPUT` is not. - Early-resolve hints apply while recording the first invocation. Replayed @@ -116,7 +136,7 @@ void qwen_decoder_layer(const CoreTaskArgs &args) { CoreTaskArgs attention_args; attention_args.add_input(hidden, attention_weight); attention_args.add_output(attention_out); - attention_args.add_scalar(uint32_t{16}); // fixed model configuration + attention_args.copy_scalars_from(args, 0, 1); // dynamic token position TaskOutputTensors attention = rt_submit_aic_task(FUNC_ATTENTION, attention_args); @@ -138,7 +158,8 @@ void decode_three_layers( const std::array &hidden, const std::array &attention_weight, const std::array &mlp_weight, - const std::array &output + const std::array &output, + const std::array &token_position ) { for (std::size_t layer = 0; layer < hidden.size(); ++layer) { CoreTaskArgs args; @@ -148,15 +169,16 @@ void decode_three_layers( mlp_weight[layer] ); args.add_output(output[layer]); + args.add_scalar(token_position[layer]); submit_qwen_decoder_layer(args); } } ``` The first layer records ordinary task submissions. Layers two and three submit -one Graph task each when their ChipTensor metadata matches. A per-layer or -per-token scalar is not dynamic in step 1; use ordinary submission or a -different fixed Graph function/key until dynamic scalar support is added. +one Graph task each when their ChipTensor metadata and boundary scalar count +match. Each replay patches the current layer's `token_position`; its value is +not part of the Graph key. ## Definition @@ -179,7 +201,7 @@ Definition. It contains: - one packed-heap offset per node; - each node's ChipTensor source: `BOUNDARY_EXACT`, `BOUNDARY_VIEW`, `INTERNAL`, or `OWN_OUTPUT`; -- fixed scalar values; +- fixed scalar values plus boundary-scalar source indices; - fixed boundary signatures and alias representatives. The header also carries a content hash of the complete Definition image. The @@ -211,8 +233,8 @@ For a cache hit, the Host Orchestrator: 5. emits one outer `GRAPH` task; 6. stages the exact-size POD submission image for upload after orchestration; 7. asks the host runtime for an aligned execution block sized from the recorded - node count, tensor-address patch count, and Definition bytes, then writes - that device address into the submission wire image. + node count, Tensor-address and scalar patch capacities, and Definition + bytes, then writes that device address into the submission wire image. Internal nodes consume no ring task-window slots. Their descriptor, payload, and slot state are built in host-owned GM. The runtime retains one grow-only @@ -226,12 +248,14 @@ The `GraphSubmission` wire POD carries the aligned device address and usable byte capacity explicitly. The Scheduler validates both before placement- constructing `GraphExecution`; it never allocates execution storage from the AICPU process heap. A block whose prior Definition key and content hash match -retains the local Definition, static node fields, and the address patch table -generated during its first materialization. That graph-affine replay skips +retains the local Definition, static node fields, and the Tensor-address and +scalar patch tables generated during its first materialization. That +graph-affine replay skips topology binding, per-node count/offset validation, tensor-source -classification, tensor wire validation, static field stores, and scalar +classification, tensor wire validation, static field stores, and static scalar copies. It refreshes only task IDs, packed-buffer bases, boundary/internal -tensor addresses, scheduling state, dispatch atomics, and wake registrations. +tensor addresses, boundary scalar bindings, scheduling state, dispatch +atomics, and wake registrations. The retained blocks are addressed directly by `(pipeline slot, Graph key, occurrence index)`. Occurrence numbering restarts deterministically for every @@ -274,8 +298,8 @@ Internal dependency readiness borrows the completion-state polling idea, but dependency wiring remains an Orchestrator responsibility: - recording constructs both fanin and fanout CSR in the immutable Definition; -- first materialization builds static runnable node state plus a compact - boundary/internal address patch table; affine replay applies that table and +- first materialization builds static runnable node state plus compact Tensor + address and scalar patch tables; affine replay applies those tables and resets only dynamic runnable state; - materialization registers each non-root on one producer selected from its saved fanin CSR; @@ -315,7 +339,6 @@ error instead of leaving an already-submitted outer Graph unable to complete. These cases assert in debug builds and execute through the ordinary path in a release build: -- dynamic boundary scalars; - an empty Graph boundary; - variable ChipTensor shape or metadata; - changed boundary aliasing; @@ -325,6 +348,7 @@ release build: - cross-boundary explicit dependencies that are not represented by a boundary ChipTensor's creator; - an unclassifiable internal ChipTensor source; +- a boundary-derived scalar accessed through mutable `scalar()`; - more than 16 Definitions, 1024 internal nodes, or 32 boundary Tensors; - insufficient task-window or heap capacity detected before outer submission. 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 b057fe3b12..f23e6d647a 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -417,8 +417,8 @@ bool upload_graph_submissions(Runtime *runtime, const HostApi *api, GraphHostSta if (definition == nullptr || definition->full_key != submission->graph_key || definition->task_count == 0 || definition->task_count > GRAPH_MAX_NODES || !graph_execution_storage_bytes( - static_cast(definition->task_count), definition->tensor_arg_count, definition->total_bytes, - &execution_bytes + static_cast(definition->task_count), definition->tensor_arg_count, + definition->scalar_arg_count, definition->total_bytes, &execution_bytes )) { LOG_ERROR("host-orch: invalid Graph execution storage request"); return false; diff --git a/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h b/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h index f12d49150f..435d5d7427 100644 --- a/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h +++ b/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h @@ -391,10 +391,8 @@ template static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const CoreTaskArgs &args, Invoke invoke) { debug_assert(!args.has_error && "Graph boundary CoreTaskArgs construction failed"); debug_assert( - args.tensor_count() <= static_cast(GRAPH_MAX_TENSOR_ARGS) && - "Graph boundary exceeds the step-1 tensor limit" + args.tensor_count() <= static_cast(GRAPH_MAX_TENSOR_ARGS) && "Graph boundary exceeds the tensor limit" ); - debug_assert(args.scalar_count() == 0 && "Dynamic Graph boundary scalars are not supported in step 1"); debug_assert( args.explicit_dep_count() == 0 && "Explicit dependencies crossing the Graph boundary are not supported" ); diff --git a/src/a2a3/runtime/host_build_graph/runtime/graph_cache.h b/src/a2a3/runtime/host_build_graph/runtime/graph_cache.h index 7351bf6c8f..8455b25b1c 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/graph_cache.h +++ b/src/a2a3/runtime/host_build_graph/runtime/graph_cache.h @@ -46,10 +46,8 @@ constexpr uint64_t graph_const_hash_impl(const char *s, uint64_t h) { constexpr uint64_t GRAPH_KEY(const char *s) { return graph_const_hash_impl(s, 1469598103934665603ULL); } inline bool rt_graph_args_cacheable(const CoreTaskArgs &args) { - // Step 1 supports dynamic tensor addresses only. Kernel scalars are - // literals inside the Graph function and become immutable Definition data. if (args.has_error || args.tensor_count() <= 0 || - args.tensor_count() > static_cast(GRAPH_MAX_TENSOR_ARGS) || args.scalar_count() != 0) { + args.tensor_count() > static_cast(GRAPH_MAX_TENSOR_ARGS)) { return false; } for (int32_t i = 0; i < args.tensor_count(); ++i) { diff --git a/src/a2a3/runtime/host_build_graph/runtime/graph_execution.h b/src/a2a3/runtime/host_build_graph/runtime/graph_execution.h index 87c09d6db0..6b88f3f548 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/graph_execution.h +++ b/src/a2a3/runtime/host_build_graph/runtime/graph_execution.h @@ -62,6 +62,17 @@ struct GraphTensorSourceRef { uint64_t packed_offset; }; +enum class GraphScalarSource : uint8_t { + STATIC_VALUE = 0, + BOUNDARY = 1, +}; + +struct GraphScalarSourceRef { + uint16_t source_index; + uint8_t source; + uint8_t reserved; +}; + struct GraphNodeDefinition { int32_t kernel_id[PTO2_SUBTASK_SLOT_COUNT]; uint8_t active_mask; @@ -98,6 +109,7 @@ struct GraphDefinition { uint32_t edge_count; uint32_t root_count; uint32_t boundary_count; + uint32_t boundary_scalar_count; uint32_t tensor_arg_count; uint32_t scalar_arg_count; uint32_t off_fanout_offsets; @@ -110,6 +122,7 @@ struct GraphDefinition { uint32_t off_tensors; uint32_t off_tensor_sources; uint32_t off_scalars; + uint32_t off_scalar_sources; uint32_t off_boundary_signatures; }; @@ -123,13 +136,16 @@ struct GraphSubmission { uint32_t definition_offset; uint32_t tensors_offset; uint32_t tensor_count; - uint32_t reserved; + uint32_t scalars_offset; + uint32_t scalar_count; }; static_assert(std::is_trivially_copyable_v); static_assert(std::is_standard_layout_v); static_assert(std::is_trivially_copyable_v); static_assert(std::is_standard_layout_v); +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_standard_layout_v); static_assert(std::is_trivially_copyable_v); static_assert(std::is_standard_layout_v); static_assert(std::is_trivially_copyable_v); @@ -243,6 +259,18 @@ inline const GraphTensor *graph_submission_tensors(const GraphSubmission &submis ); } +inline const uint64_t *graph_submission_scalars(const GraphSubmission &submission) { + if (submission.scalar_count == 0) return nullptr; + if (submission.scalars_offset == 0 || submission.scalars_offset % alignof(uint64_t) != 0 || + submission.scalars_offset > submission.total_bytes || + submission.scalar_count > (submission.total_bytes - submission.scalars_offset) / sizeof(uint64_t)) { + return nullptr; + } + return reinterpret_cast( + reinterpret_cast(&submission) + submission.scalars_offset + ); +} + enum class GraphExecutionState : uint8_t { SUBMITTED = 0, MATERIALIZING = 1, @@ -277,6 +305,18 @@ static_assert(std::is_trivially_copyable_v); static_assert(std::is_standard_layout_v); static_assert(sizeof(GraphTensorAddressPatch) == 16); +struct GraphScalarPatch { + uint16_t node_index; + uint8_t node_scalar_index; + uint8_t boundary_scalar_index; +}; + +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_standard_layout_v); +static_assert(sizeof(GraphScalarPatch) == 4); +static_assert(GRAPH_MAX_NODES <= UINT16_MAX); +static_assert(MAX_SCALAR_ARGS <= UINT8_MAX); + struct alignas(64) GraphNodeStorage { PTO2TaskDescriptor task; PTO2TaskPayload payload; @@ -300,6 +340,9 @@ struct GraphExecution { uint32_t tensor_patch_capacity{0}; uint32_t materialized_tensor_patches{0}; uint32_t materialized_tensor_patch_count{0}; + uint32_t scalar_patch_capacity{0}; + 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}; @@ -312,25 +355,30 @@ struct GraphExecution { GraphNodeStorage *nodes{nullptr}; 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}; const GraphTensor *boundary_tensors{nullptr}; uint32_t boundary_tensor_count{0}; + const uint64_t *boundary_scalars{nullptr}; + uint32_t boundary_scalar_count{0}; }; 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, size_t definition_capacity, size_t *nodes_offset, - size_t *tensor_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 definition_capacity, + size_t *nodes_offset, size_t *tensor_patches_offset, size_t *scalar_patches_offset, size_t *definition_offset, + size_t *storage_bytes ) { - if (nodes_offset == nullptr || tensor_patches_offset == nullptr || definition_offset == nullptr || - storage_bytes == nullptr || node_capacity <= 0 || + if (nodes_offset == nullptr || tensor_patches_offset == nullptr || scalar_patches_offset == nullptr || + definition_offset == nullptr || storage_bytes == nullptr || node_capacity <= 0 || static_cast(node_capacity) > SIZE_MAX / sizeof(GraphNodeStorage) || - tensor_patch_capacity > GRAPH_MAX_NODES * MAX_TENSOR_ARGS) { + tensor_patch_capacity > GRAPH_MAX_NODES * MAX_TENSOR_ARGS || + scalar_patch_capacity > GRAPH_MAX_NODES * MAX_SCALAR_ARGS) { return false; } auto checked_align_up = [](size_t value, size_t alignment, size_t *result) { @@ -340,11 +388,16 @@ inline bool graph_execution_storage_layout( }; const size_t nodes_bytes = static_cast(node_capacity) * sizeof(GraphNodeStorage); const size_t tensor_patches_bytes = static_cast(tensor_patch_capacity) * sizeof(GraphTensorAddressPatch); + const size_t scalar_patches_bytes = static_cast(scalar_patch_capacity) * sizeof(GraphScalarPatch); if (!checked_align_up(sizeof(GraphExecution), alignof(GraphNodeStorage), nodes_offset) || *nodes_offset > SIZE_MAX - nodes_bytes || !checked_align_up(*nodes_offset + nodes_bytes, alignof(GraphTensorAddressPatch), tensor_patches_offset) || *tensor_patches_offset > SIZE_MAX - tensor_patches_bytes || - !checked_align_up(*tensor_patches_offset + tensor_patches_bytes, alignof(GraphDefinition), definition_offset) || + !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) { return false; } @@ -352,14 +405,16 @@ inline bool graph_execution_storage_layout( } inline bool graph_execution_storage_bytes( - int32_t node_capacity, uint32_t tensor_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 definition_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, definition_capacity, &nodes_offset, &tensor_patches_offset, - &definition_offset, storage_bytes + node_capacity, tensor_patch_capacity, scalar_patch_capacity, definition_capacity, &nodes_offset, + &tensor_patches_offset, &scalar_patches_offset, &definition_offset, storage_bytes ); } diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp index affa1cacdb..3cc90f0748 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 @@ -278,6 +278,17 @@ struct GraphRecordedTensorSourceRef { uint64_t packed_offset{0}; }; +enum class GraphRecordedScalarSource : uint8_t { + STATIC_VALUE, + BOUNDARY, + INVALIDATED_BOUNDARY, +}; + +struct GraphRecordedScalarSourceRef { + GraphRecordedScalarSource source{GraphRecordedScalarSource::STATIC_VALUE}; + size_t source_index{0}; +}; + struct GraphRecordedNode { std::array kernel_ids{}; ActiveMask active_mask{}; @@ -289,6 +300,7 @@ struct GraphRecordedNode { std::vector tensors; std::vector tensor_sources; std::vector scalars; + std::vector scalar_sources; std::vector internal_fanins; }; @@ -300,6 +312,7 @@ struct GraphRecording { std::vector current_fanins; std::vector boundary_tensors; std::vector boundary_types; + const CoreTaskArgs *boundary_args{nullptr}; std::vector nodes; }; @@ -361,6 +374,30 @@ bool graph_tensor_from_boundary( return false; } +GraphRecordedScalarSourceRef +graph_classify_scalar(const GraphRecording &recording, const CoreTaskArgs &args, int32_t scalar_index) { + if (recording.boundary_args == nullptr) return {}; + if (&args == recording.boundary_args && scalar_index < recording.boundary_args->scalar_count()) { + return GraphRecordedScalarSourceRef{GraphRecordedScalarSource::BOUNDARY, static_cast(scalar_index)}; + } + + const void *source = args.scalar_source(scalar_index); + const void *invalidated_source = args.invalidated_scalar_source(scalar_index); + if (source == nullptr && invalidated_source == nullptr) return {}; + for (int32_t i = 0; i < recording.boundary_args->scalar_count(); ++i) { + const void *boundary_source = static_cast(&recording.boundary_args->scalar(i)); + if (source == boundary_source) { + return GraphRecordedScalarSourceRef{GraphRecordedScalarSource::BOUNDARY, static_cast(i)}; + } + if (invalidated_source == boundary_source) { + return GraphRecordedScalarSourceRef{ + GraphRecordedScalarSource::INVALIDATED_BOUNDARY, static_cast(i) + }; + } + } + return {}; +} + bool graph_classify_tensor( const GraphRecording &recording, const GraphRecordedNode ¤t, int32_t task_index, const ChipTensor &tensor, GraphRecordedTensorSourceRef *source @@ -474,6 +511,11 @@ void graph_record_task( node.tensors.assign(payload.tensors, payload.tensors + payload.tensor_count); node.tensor_sources.resize(static_cast(payload.tensor_count)); node.scalars.assign(payload.scalars, payload.scalars + payload.scalar_count); + node.scalar_sources.resize(static_cast(payload.scalar_count)); + if (args.scalar_count() != payload.scalar_count) { + recording.unsupported = true; + return; + } for (int32_t i = 0; i < payload.tensor_count; ++i) { if (!graph_classify_tensor( recording, node, task_index, payload.tensors[i], &node.tensor_sources[static_cast(i)] @@ -482,6 +524,14 @@ void graph_record_task( return; } } + for (int32_t i = 0; i < payload.scalar_count; ++i) { + GraphRecordedScalarSourceRef source = graph_classify_scalar(recording, args, i); + if (source.source == GraphRecordedScalarSource::INVALIDATED_BOUNDARY) { + recording.unsupported = true; + return; + } + node.scalar_sources[static_cast(i)] = source; + } for (size_t producer : recording.current_fanins) { if (producer >= static_cast(task_index)) { recording.unsupported = true; @@ -543,10 +593,23 @@ std::optional graph_pack_tensor_source(const GraphRecorded return packed; } +std::optional graph_pack_scalar_source(const GraphRecordedScalarSourceRef &source) { + if (source.source == GraphRecordedScalarSource::INVALIDATED_BOUNDARY || source.source_index > UINT16_MAX) { + return std::nullopt; + } + + GraphScalarSourceRef packed{}; + packed.source = source.source == GraphRecordedScalarSource::BOUNDARY ? + static_cast(GraphScalarSource::BOUNDARY) : + static_cast(GraphScalarSource::STATIC_VALUE); + packed.source_index = static_cast(source.source_index); + return packed; +} + bool graph_build_definition(const GraphRecording &recording, std::vector *image) { if (image == nullptr || recording.unsupported || recording.nodes.empty() || recording.nodes.size() > GRAPH_MAX_NODES || recording.boundary_tensors.size() > UINT16_MAX || - recording.boundary_tensors.size() != recording.boundary_types.size() || + recording.boundary_tensors.size() != recording.boundary_types.size() || recording.boundary_args == nullptr || std::any_of(recording.boundary_tensors.begin(), recording.boundary_tensors.end(), [](const ChipTensor &tensor) { return tensor.ndims > MAX_TENSOR_DIMS; })) { @@ -562,6 +625,7 @@ bool graph_build_definition(const GraphRecording &recording, std::vector tensors; std::vector tensor_sources; std::vector scalars; + std::vector scalar_sources; uint64_t required_heap = 0; uint32_t edge_count = 0; @@ -571,9 +635,11 @@ bool graph_build_definition(const GraphRecording &recording, std::vector static_cast(INT32_MAX) || source.scalars.size() > static_cast(INT32_MAX) || source.internal_fanins.size() > UINT16_MAX || source.tensors.size() != source.tensor_sources.size() || + source.scalars.size() != source.scalar_sources.size() || tensors.size() > UINT32_MAX - source.tensors.size() || tensor_sources.size() > UINT32_MAX - source.tensor_sources.size() || scalars.size() > UINT32_MAX - source.scalars.size() || + scalar_sources.size() > UINT32_MAX - source.scalar_sources.size() || std::any_of(source.tensors.begin(), source.tensors.end(), [](const ChipTensor &tensor) { return tensor.ndims > MAX_TENSOR_DIMS; })) { @@ -611,7 +677,21 @@ bool graph_build_definition(const GraphRecording &recording, std::vector packed_source = + graph_pack_scalar_source(source.scalar_sources[scalar_index]); + if (!packed_source.has_value() || + (packed_source->source == static_cast(GraphScalarSource::BOUNDARY) && + packed_source->source_index >= recording.boundary_args->scalar_count())) { + return false; + } + scalar_sources.push_back(*packed_source); + scalars.push_back( + packed_source->source == static_cast(GraphScalarSource::BOUNDARY) ? + 0 : + source.scalars[scalar_index] + ); + } } std::vector fanout_offsets(recording.nodes.size() + 1, 0); @@ -649,6 +729,7 @@ bool graph_build_definition(const GraphRecording &recording, std::vector(roots.size()); definition.boundary_count = static_cast(signatures.size()); + definition.boundary_scalar_count = static_cast(recording.boundary_args->scalar_count()); definition.tensor_arg_count = static_cast(tensors.size()); definition.scalar_arg_count = static_cast(scalars.size()); definition.off_fanout_offsets = graph_append_section(image, fanout_offsets); @@ -661,12 +742,14 @@ bool graph_build_definition(const GraphRecording &recording, std::vector(definition.boundary_count)) { + if (args.scalar_count() != static_cast(definition.boundary_scalar_count) || + args.explicit_dep_count() != 0 || args.tensor_count() != static_cast(definition.boundary_count)) { LOG_WARN( - "[GraphExecution] fixed boundary contract mismatch: tensors=%d/%u scalars=%d explicit_deps=%u", - args.tensor_count(), definition.boundary_count, args.scalar_count(), args.explicit_dep_count() + "[GraphExecution] fixed boundary contract mismatch: tensors=%d/%u scalars=%d/%u explicit_deps=%u", + args.tensor_count(), definition.boundary_count, args.scalar_count(), definition.boundary_scalar_count, + args.explicit_dep_count() ); return false; } @@ -1426,11 +1510,25 @@ bool graph_build_submission_image( if (definition_offset > UINT32_MAX || tensors_offset > UINT32_MAX || tensors_offset > UINT32_MAX - tensor_bytes) { return false; } - submission_image->assign(tensors_offset + tensor_bytes, std::byte{0}); + const size_t tensors_end = tensors_offset + tensor_bytes; + const size_t scalar_bytes = static_cast(args.scalar_count()) * sizeof(uint64_t); + const size_t scalars_offset = args.scalar_count() == 0 ? 0 : PTO2_ALIGN_UP(tensors_end, alignof(uint64_t)); + const size_t total_bytes = args.scalar_count() == 0 ? tensors_end : scalars_offset + scalar_bytes; + if ((args.scalar_count() != 0 && (scalars_offset > UINT32_MAX || scalars_offset > UINT32_MAX - scalar_bytes)) || + total_bytes > UINT32_MAX) { + 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()); + if (args.scalar_count() != 0) { + std::memcpy( + submission_image->data() + scalars_offset, args.scalar_data(), + static_cast(args.scalar_count()) * sizeof(uint64_t) + ); + } const GraphDefinition &definition = *graph_definition(definition_image); GraphSubmission submission{}; @@ -1439,6 +1537,8 @@ bool graph_build_submission_image( 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); + submission.scalar_count = static_cast(args.scalar_count()); std::memcpy(submission_image->data(), &submission, sizeof(submission)); return true; } @@ -1526,7 +1626,6 @@ PTO2OrchestratorState::graph_begin(uint64_t graph_key, const CoreTaskArgs &args, GraphScopeResult result; GraphHostState *state = graph_state_from(orch); if (state == nullptr || !rt_graph_args_cacheable(args) || args.explicit_dep_count() != 0) { - debug_assert(args.scalar_count() == 0 && "Graph execution scalars are not supported in step 1"); debug_assert(args.explicit_dep_count() == 0 && "Graph boundary explicit dependencies are not supported"); return result; } @@ -1567,6 +1666,8 @@ PTO2OrchestratorState::graph_begin(uint64_t graph_key, const CoreTaskArgs &args, auto recording = std::make_unique(); recording->full_key = full_key; recording->start_local_task_id = orch->ring.task_allocator.active_count(); + args.anchor_scalar_sources(); + recording->boundary_args = &args; recording->boundary_tensors.reserve(static_cast(args.tensor_count())); recording->boundary_types.reserve(static_cast(args.tensor_count())); for (int32_t i = 0; i < args.tensor_count(); ++i) { diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_types.h index 039fac5ce3..492c311a77 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_types.h @@ -28,6 +28,8 @@ #include #include +#include +#include #include #include #include @@ -284,6 +286,8 @@ struct Arg : TaskArgsTpl { void clear() { Base::clear(); + scalar_sources_.fill(nullptr); + scalar_sources_invalidated_.fill(false); #if SIMPLER_DFX dump_arg_selection_.clear(); #endif @@ -446,6 +450,8 @@ struct Arg : TaskArgsTpl { return; } memcpy(&scalars_[scalar_count_], values, count * sizeof(uint64_t)); + std::fill_n(scalar_sources_.begin() + scalar_count_, count, nullptr); + std::fill_n(scalar_sources_invalidated_.begin() + scalar_count_, count, false); #if SIMPLER_DFX dump_arg_selection_.clear_scalar_metadata(scalar_count_, count); #endif @@ -481,6 +487,8 @@ struct Arg : TaskArgsTpl { dst[i] = static_cast(static_cast(values[i])); } #endif + std::fill_n(scalar_sources_.begin() + scalar_count_, count, nullptr); + std::fill_n(scalar_sources_invalidated_.begin() + scalar_count_, count, false); #if SIMPLER_DFX dump_arg_selection_.clear_scalar_metadata(scalar_count_, count); #endif @@ -501,12 +509,40 @@ struct Arg : TaskArgsTpl { return; } memcpy(&scalars_[scalar_count_], &src.scalars_[src_offset], count * sizeof(uint64_t)); + for (int i = 0; i < count; ++i) { + const int src_index = src_offset + i; + scalar_sources_[scalar_count_ + i] = src.scalar_sources_[src_index] != nullptr ? + src.scalar_sources_[src_index] : + static_cast(&src.scalars_[src_index]); + scalar_sources_invalidated_[scalar_count_ + i] = src.scalar_sources_invalidated_[src_index]; + } #if SIMPLER_DFX dump_arg_selection_.copy_scalar_dtypes_from(src.dump_arg_selection_, scalar_count_, src_offset, count); #endif scalar_count_ += count; } + const uint64_t &scalar(int32_t i) const { return scalars_[i]; } + uint64_t &scalar(int32_t i) { + // A mutable reference may escape and be written later. Conservatively + // invalidate inherited Graph-boundary provenance as soon as it is requested. + scalar_sources_invalidated_[i] = scalar_sources_[i] != nullptr; + return scalars_[i]; + } + const void *scalar_source(int32_t i) const { return scalar_sources_invalidated_[i] ? nullptr : scalar_sources_[i]; } + const void *invalidated_scalar_source(int32_t i) const { + return scalar_sources_invalidated_[i] ? scalar_sources_[i] : nullptr; + } + // Graph recording starts before its function runs. Anchor the host-only + // provenance at that point so every boundary slot has a unique identity, + // even when multiple slots were initialized from the same lvalue. + void anchor_scalar_sources() const { + for (int32_t i = 0; i < scalar_count_; ++i) { + scalar_sources_[i] = static_cast(&scalars_[i]); + scalar_sources_invalidated_[i] = false; + } + } + #if SIMPLER_DFX const uint8_t *scalar_dtypes() const { return dump_arg_selection_.scalar_dtypes(); } #else @@ -514,6 +550,12 @@ struct Arg : TaskArgsTpl { #endif private: + // In-process recording metadata only; it is never copied into a task payload + // or any host-device wire image. + mutable std::array scalar_sources_{}; + // Kept separately so invalidated boundary ancestry can reject Graph caching + // instead of being mistaken for an unrelated static scalar. + mutable std::array scalar_sources_invalidated_{}; // Caller-owned dependency array; lifetime must extend through submit. #if SIMPLER_DFX DumpArgSelection dump_arg_selection_; @@ -542,11 +584,14 @@ struct Arg : TaskArgsTpl { template void add_scalar_one(T &&value) { scalars_[scalar_count_] = to_u64(value); -#if SIMPLER_DFX - uintptr_t scalar_source_ptr = 0; + scalar_sources_[scalar_count_] = nullptr; + scalar_sources_invalidated_[scalar_count_] = false; if constexpr (std::is_lvalue_reference_v) { - scalar_source_ptr = reinterpret_cast(&value); + scalar_sources_[scalar_count_] = static_cast(&value); } +#if SIMPLER_DFX + uintptr_t scalar_source_ptr = 0; + scalar_source_ptr = reinterpret_cast(scalar_sources_[scalar_count_]); dump_arg_selection_.record_scalar_source( scalar_count_, scalar_source_ptr, dtype_of>>() ); diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cpp index 1b06758b3f..8c1b176932 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cpp @@ -27,16 +27,18 @@ void destroy_execution_nodes(GraphExecution *execution) { } void reset_execution( - GraphExecution *execution, int32_t node_count, uint32_t tensor_patch_count, uint64_t graph_key, - uint64_t definition_hash + GraphExecution *execution, int32_t node_count, uint32_t tensor_patch_count, uint32_t scalar_patch_count, + uint64_t graph_key, uint64_t definition_hash ) { 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->definition_capacity, &nodes_offset, - &tensor_patches_offset, &definition_offset, &bytes + 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 )) { return; } @@ -44,6 +46,7 @@ void reset_execution( execution->materialized_definition_hash == definition_hash && execution->materialized_node_count == node_count && execution->materialized_tensor_patch_count == tensor_patch_count && + execution->materialized_scalar_patch_count <= scalar_patch_count && execution->constructed_nodes >= node_count; execution->state.store(GraphExecutionState::SUBMITTED, std::memory_order_relaxed); execution->materialize_busy.store(0, std::memory_order_relaxed); @@ -52,6 +55,7 @@ void reset_execution( execution->node_count = node_count; execution->materialized_nodes = 0; execution->materialized_tensor_patches = 0; + execution->materialized_scalar_patches = 0; execution->allocation_bytes = bytes; execution->graph_key = graph_key; execution->definition_hash = definition_hash; @@ -61,6 +65,8 @@ void reset_execution( reinterpret_cast(reinterpret_cast(execution) + nodes_offset); execution->tensor_patches = 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; @@ -69,13 +75,17 @@ void reset_execution( } execution->boundary_tensors = nullptr; execution->boundary_tensor_count = 0; + execution->boundary_scalars = nullptr; + execution->boundary_scalar_count = 0; } 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.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 || + execution.materialized_scalar_patch_count > execution.scalar_patch_capacity || execution.constructed_nodes < 0 || execution.constructed_nodes > execution.node_capacity || execution.node_count <= 0 || execution.node_count > execution.node_capacity || execution.state.load(std::memory_order_acquire) != GraphExecutionState::COMPLETED || @@ -84,14 +94,15 @@ 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.definition_capacity, &expected_bytes + execution.node_capacity, execution.tensor_patch_capacity, execution.scalar_patch_capacity, + execution.definition_capacity, &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, size_t definition_bytes + uint32_t tensor_patch_count, uint32_t scalar_patch_count, size_t definition_bytes ) { if (node_count <= 0 || node_count > static_cast(GRAPH_MAX_NODES) || definition_bytes == 0 || submission.execution_storage == 0 || submission.execution_storage_bytes > SIZE_MAX || @@ -100,7 +111,9 @@ GraphExecution *acquire_host_execution_storage( } 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, definition_bytes, &required_bytes) || + if (!graph_execution_storage_bytes( + node_count, tensor_patch_count, scalar_patch_count, definition_bytes, &required_bytes + ) || required_bytes > storage_bytes) { return nullptr; } @@ -113,16 +126,18 @@ 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; 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, graph_key, definition_hash); + reset_execution(execution, node_count, tensor_patch_count, scalar_patch_count, graph_key, definition_hash); execution->storage_magic = GRAPH_EXECUTION_STORAGE_MAGIC; return execution; } @@ -145,6 +160,7 @@ void reset_graph_payload(PTO2TaskPayload &payload) { bool bind_graph_topology(GraphExecution &execution) { if (execution.definition == nullptr) return false; const GraphDefinition &definition = *execution.definition; + if (definition.boundary_scalar_count > MAX_SCALAR_ARGS) return false; const uint32_t *fanin_offsets = graph_definition_array(definition, definition.off_fanin_offsets, definition.task_count + 1); const uint16_t *fanin_indices = @@ -266,11 +282,16 @@ GraphExecution *graph_execution_localize(PTO2TaskSlotState &outer_slot) { const GraphDefinition *definition = graph_submission_definition(*submission); 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 || - boundary_tensors == nullptr || + submission->scalar_count != definition->boundary_scalar_count || boundary_tensors == 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) { @@ -295,7 +316,7 @@ GraphExecution *graph_execution_localize(PTO2TaskSlotState &outer_slot) { GraphExecution *execution = acquire_host_execution_storage( *submission, static_cast(definition->task_count), submission->graph_key, definition->content_hash, - definition->tensor_arg_count, definition->total_bytes + definition->tensor_arg_count, definition->scalar_arg_count, definition->total_bytes ); if (execution == nullptr) { __atomic_store_n(&submission->local_execution, 0, __ATOMIC_RELEASE); @@ -314,6 +335,8 @@ GraphExecution *graph_execution_localize(PTO2TaskSlotState &outer_slot) { } execution->boundary_tensors = boundary_tensors; execution->boundary_tensor_count = submission->tensor_count; + execution->boundary_scalars = boundary_scalars; + execution->boundary_scalar_count = submission->scalar_count; execution->outer_slot = &outer_slot; const uint64_t desired = static_cast(reinterpret_cast(execution)); @@ -327,7 +350,8 @@ GraphMaterializeResult graph_execution_materialize_slice( if (nodes_materialized != nullptr) *nodes_materialized = 0; if (outer_slot.task_kind != TaskKind::GRAPH || outer_slot.task == nullptr || outer_slot.task->packed_buffer_base == nullptr || max_nodes <= 0 || execution.definition == nullptr || - execution.node_storage == nullptr || execution.tensor_patches == nullptr) { + execution.node_storage == nullptr || execution.tensor_patches == nullptr || + execution.scalar_patches == nullptr) { return GraphMaterializeResult::INVALID; } @@ -362,6 +386,7 @@ GraphMaterializeResult graph_execution_materialize_slice( const GraphTensor *definition_tensors = nullptr; const GraphTensorSourceRef *tensor_sources = nullptr; const uint64_t *definition_scalars = nullptr; + const GraphScalarSourceRef *scalar_sources = nullptr; if (!affine_reuse) { nodes = graph_definition_array(definition, definition.off_nodes, definition.task_count); node_offsets = graph_definition_array(definition, definition.off_node_offsets, definition.task_count); @@ -378,9 +403,14 @@ GraphMaterializeResult graph_execution_materialize_slice( definition.scalar_arg_count == 0 ? nullptr : graph_definition_array(definition, definition.off_scalars, definition.scalar_arg_count); + scalar_sources = definition.scalar_arg_count == 0 ? + nullptr : + graph_definition_array( + definition, definition.off_scalar_sources, definition.scalar_arg_count + ); if (nodes == nullptr || node_offsets == nullptr || (definition.tensor_arg_count != 0 && (definition_tensors == nullptr || tensor_sources == nullptr)) || - (definition.scalar_arg_count != 0 && definition_scalars == nullptr)) { + (definition.scalar_arg_count != 0 && (definition_scalars == nullptr || scalar_sources == nullptr))) { execution.materialize_busy.store(0, std::memory_order_release); return GraphMaterializeResult::INVALID; } @@ -513,11 +543,25 @@ GraphMaterializeResult graph_execution_materialize_slice( execution.tensor_patches[execution.materialized_tensor_patches++] = patch; graph_tensor_unpack(rebound, &tensor); } - if (source.scalar_count > 0) { - std::memcpy( - payload.scalars, definition_scalars + source.scalar_offset, - static_cast(source.scalar_count) * sizeof(uint64_t) - ); + for (int32_t j = 0; j < source.scalar_count; ++j) { + const uint32_t scalar_index = source.scalar_offset + static_cast(j); + const GraphScalarSourceRef &ref = scalar_sources[scalar_index]; + if (ref.source == static_cast(GraphScalarSource::STATIC_VALUE)) { + payload.scalars[j] = definition_scalars[scalar_index]; + } else if (ref.source == static_cast(GraphScalarSource::BOUNDARY)) { + if (ref.source_index >= execution.boundary_scalar_count || execution.boundary_scalars == nullptr || + execution.materialized_scalar_patches >= execution.scalar_patch_capacity) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + payload.scalars[j] = execution.boundary_scalars[ref.source_index]; + execution.scalar_patches[execution.materialized_scalar_patches++] = GraphScalarPatch{ + static_cast(i), static_cast(j), static_cast(ref.source_index) + }; + } else { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } } } else { for (int32_t j = 0; j < payload.tensor_count; ++j) { @@ -529,6 +573,18 @@ GraphMaterializeResult graph_execution_materialize_slice( payload.tensors[j].buffer.addr = outer_base + patch.address_offset; } } + while (execution.materialized_scalar_patches < execution.materialized_scalar_patch_count) { + const GraphScalarPatch &patch = execution.scalar_patches[execution.materialized_scalar_patches]; + if (patch.node_index != static_cast(i)) break; + if (patch.node_scalar_index >= payload.scalar_count || + patch.boundary_scalar_index >= execution.boundary_scalar_count || + execution.boundary_scalars == nullptr) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + payload.scalars[patch.node_scalar_index] = execution.boundary_scalars[patch.boundary_scalar_index]; + execution.materialized_scalar_patches++; + } } reset_graph_payload(payload); if (!register_initial_graph_waiter(execution, i)) { @@ -550,6 +606,10 @@ GraphMaterializeResult graph_execution_materialize_slice( execution.materialize_busy.store(0, std::memory_order_release); return GraphMaterializeResult::INVALID; } + if (affine_reuse && execution.materialized_scalar_patches != execution.materialized_scalar_patch_count) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } // nodes is published before PREPARED. An activator that acquires the state // may therefore route saved roots without observing partially built nodes. @@ -558,6 +618,7 @@ GraphMaterializeResult graph_execution_materialize_slice( execution.materialized_definition_hash = execution.definition_hash; execution.materialized_node_count = execution.node_count; execution.materialized_tensor_patch_count = execution.materialized_tensor_patches; + execution.materialized_scalar_patch_count = execution.materialized_scalar_patches; execution.materialized_outer_base = outer_base; execution.state.store(GraphExecutionState::PREPARED, std::memory_order_release); execution.materialize_busy.store(0, std::memory_order_release); diff --git a/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_orch.cpp b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_orch.cpp index 27b586c94d..c6d1adf7a6 100644 --- a/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_orch.cpp +++ b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_orch.cpp @@ -21,11 +21,10 @@ namespace { -void layer(const CoreTaskArgs &args, float left_delta, float right_delta) { +void layer(const CoreTaskArgs &args, int variant) { const ChipTensor &a = args.tensor(0).ref(); const ChipTensor &b = args.tensor(1).ref(); const ChipTensor &output = args.tensor(2).ref(); - const float ndim_delta = b.ndims == 1 ? 0.0F : 2.0F; const std::array shape{a.shapes[0]}; TensorCreateInfo intermediate(shape.data(), static_cast(shape.size()), DataType::FLOAT32); @@ -45,7 +44,11 @@ void layer(const CoreTaskArgs &args, float left_delta, float right_delta) { CoreTaskArgs left_args; left_args.add_input(sum); left_args.add_output(intermediate); - left_args.add_scalar(left_delta + ndim_delta); + if (variant == 0) { + left_args.add_scalar(args.scalar(0)); + } else { + left_args.add_scalar(100.0F); + } left_args.set_allow_early_resolve(true); TaskOutputTensors left_outputs = rt_submit_aiv_task(FUNC_ADD_SCALAR, left_args); ChipTensor left = left_outputs.get_ref(0); @@ -53,7 +56,7 @@ void layer(const CoreTaskArgs &args, float left_delta, float right_delta) { CoreTaskArgs right_args; right_args.add_input(sum); right_args.add_output(intermediate); - right_args.add_scalar(right_delta + ndim_delta); + right_args.add_scalar(args.scalar(variant == 0 ? 1 : 0)); TaskOutputTensors right_outputs = rt_submit_aiv_task(FUNC_ADD_SCALAR, right_args); ChipTensor right = right_outputs.get_ref(0); @@ -63,7 +66,7 @@ void layer(const CoreTaskArgs &args, float left_delta, float right_delta) { rt_submit_aiv_task(FUNC_MUL, mul_args); } -void submit_layer(const CoreTaskArgs &args) { rt_submit_graph(&layer, args, 1.0F, 2.0F); } +void submit_layer(const CoreTaskArgs &args, int variant) { rt_submit_graph(&layer, args, variant); } } // namespace @@ -89,11 +92,20 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const Chip TaskOutputTensors seed_outputs = rt_submit_aiv_task(FUNC_ADD_SCALAR, seed_args); ChipTensor seeded_a = seed_outputs.get_ref(0); - for (int32_t output_index = 2; output_index < 5; ++output_index) { + const std::array left_deltas{1.0F, 3.0F, 5.0F}; + const std::array right_deltas{2.0F, 4.0F, 6.0F}; + const std::array unused_deltas{7.0F, 8.0F, 9.0F}; + const std::array variants{0, 1, 0}; + for (size_t layer = 0; layer < variants.size(); ++layer) { CoreTaskArgs layer_args; layer_args.add_input(seeded_a, b); - layer_args.add_output(args.tensor(output_index).ref()); - submit_layer(layer_args); // first call records; later calls submit one Graph task + layer_args.add_output(args.tensor(static_cast(layer) + 2).ref()); + if (variants[layer] == 0) { + layer_args.add_scalar(left_deltas[layer], right_deltas[layer], unused_deltas[layer]); + } else { + layer_args.add_scalar(right_deltas[layer]); + } + submit_layer(layer_args, variants[layer]); } } diff --git a/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution.py b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution.py index e53c547954..10545c1643 100644 --- a/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution.py +++ b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution.py @@ -7,7 +7,7 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- -"""Graph Execution records once and replays topology with dynamic CoreTaskArgs tensors.""" +"""Graph Execution replays dynamic tensor and scalar bindings across config-keyed definitions.""" import torch from simpler.task_interface import ArgDirection as D @@ -75,11 +75,9 @@ def generate_args(self, params): def compute_golden(self, args, params): base = args.a + args.b - ndim_delta = 0.0 if args.a.ndim == 1 else 2.0 - expected = (base + 1.0 + ndim_delta) * (base + 2.0 + ndim_delta) - args.output_1[:] = expected - args.output_3[:] = expected - args.output_5[:] = expected + args.output_1[:] = (base + 1.0) * (base + 2.0) + args.output_3[:] = (base + 100.0) * (base + 4.0) + args.output_5[:] = (base + 5.0) * (base + 6.0) if __name__ == "__main__": diff --git a/tests/ut/cpp/a2a3/test_graph_cache.cpp b/tests/ut/cpp/a2a3/test_graph_cache.cpp index 805c19b39e..9428c16817 100644 --- a/tests/ut/cpp/a2a3/test_graph_cache.cpp +++ b/tests/ut/cpp/a2a3/test_graph_cache.cpp @@ -73,7 +73,10 @@ std::vector make_test_definition(uint64_t graph_key, uint64_t boundar tensor_sources[0].source = static_cast(GraphTensorSource::BOUNDARY_EXACT); tensor_sources[1].source = static_cast(GraphTensorSource::INTERNAL); tensor_sources[1].packed_offset = 16; - std::vector scalars{17, 18}; + std::vector scalars{0, 18}; + std::vector scalar_sources(2); + scalar_sources[0].source = static_cast(GraphScalarSource::BOUNDARY); + scalar_sources[1].source = static_cast(GraphScalarSource::STATIC_VALUE); GraphDefinition definition{}; definition.full_key = graph_key; @@ -82,6 +85,7 @@ std::vector make_test_definition(uint64_t graph_key, uint64_t boundar definition.edge_count = 1; definition.root_count = 1; definition.boundary_count = 1; + definition.boundary_scalar_count = 1; definition.tensor_arg_count = 2; definition.scalar_arg_count = 2; definition.off_fanin_offsets = append_section(image, fanin_offsets); @@ -94,6 +98,7 @@ std::vector make_test_definition(uint64_t graph_key, uint64_t boundar definition.off_tensors = append_section(image, tensors); definition.off_tensor_sources = append_section(image, tensor_sources); definition.off_scalars = append_section(image, scalars); + definition.off_scalar_sources = append_section(image, scalar_sources); definition.total_bytes = static_cast(image.size()); std::memcpy(image.data(), &definition, sizeof(definition)); @@ -103,15 +108,18 @@ std::vector make_test_definition(uint64_t graph_key, uint64_t boundar } std::vector make_test_submission( - uint64_t graph_key, uint64_t boundary_address, uint64_t execution_storage, size_t execution_storage_bytes + 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)); - std::vector image(tensors_offset + sizeof(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)); GraphSubmission submission{}; submission.graph_key = graph_key; @@ -121,6 +129,8 @@ std::vector make_test_submission( 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); + submission.scalar_count = 1; std::memcpy(image.data(), &submission, sizeof(submission)); return image; } @@ -151,24 +161,78 @@ TEST(GraphCache, RejectsEmptyBoundary) { EXPECT_FALSE(rt_graph_args_cacheable(args)); } +TEST(GraphCache, AcceptsBoundaryScalars) { + std::array boundary{}; + const GraphTensor packed = make_test_tensor(reinterpret_cast(boundary.data())); + ChipTensor tensor{}; + graph_tensor_unpack(packed, &tensor); + + CoreTaskArgs args; + args.add_input(tensor); + args.add_scalar(uint32_t{17}); + + EXPECT_TRUE(rt_graph_args_cacheable(args)); +} + +TEST(GraphCache, ConfigValuesSelectDifferentDefinitions) { + constexpr uint64_t GRAPH_ID = 0x1234; + + EXPECT_NE(rt_graph_make_key(GRAPH_ID, 0), rt_graph_make_key(GRAPH_ID, 1)); + EXPECT_EQ(rt_graph_make_key(GRAPH_ID, 0), rt_graph_make_key(GRAPH_ID, 0)); +} + +TEST(GraphScalarProvenance, ForwardedScalarRetainsBoundarySource) { + uint32_t value = 17; + CoreTaskArgs boundary_args; + boundary_args.add_scalar(value, value); + boundary_args.anchor_scalar_sources(); + CoreTaskArgs forwarded_args; + forwarded_args.copy_scalars_from(boundary_args, 1, 1); + CoreTaskArgs node_args; + + node_args.copy_scalars_from(forwarded_args, 0, 1); + + EXPECT_EQ(node_args.scalar_source(0), static_cast(&std::as_const(boundary_args).scalar(1))); +} + +TEST(GraphScalarProvenance, MutableAccessInvalidatesForwardedSource) { + CoreTaskArgs boundary_args; + boundary_args.add_scalar(uint32_t{17}); + boundary_args.anchor_scalar_sources(); + CoreTaskArgs node_args; + node_args.copy_scalars_from(boundary_args, 0, 1); + ASSERT_NE(node_args.scalar_source(0), nullptr); + + node_args.scalar(0) = 18; + + EXPECT_EQ(node_args.scalar_source(0), nullptr); + EXPECT_EQ( + node_args.invalidated_scalar_source(0), static_cast(&std::as_const(boundary_args).scalar(0)) + ); +} + 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, DEFINITION_BYTES, &nodes_offset, &tensor_patches_offset, &definition_offset, - &storage_bytes + NODE_COUNT, TENSOR_PATCH_COUNT, SCALAR_PATCH_COUNT, DEFINITION_BYTES, &nodes_offset, &tensor_patches_offset, + &scalar_patches_offset, &definition_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(definition_offset, tensor_patches_offset + TENSOR_PATCH_COUNT * sizeof(GraphTensorAddressPatch)); + 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_EQ(storage_bytes % alignof(GraphNodeStorage), 0U); } @@ -176,8 +240,8 @@ TEST(GraphExecutionStorage, ComputesAlignedExactSize) { TEST(GraphExecutionStorage, RejectsInvalidCapacity) { size_t storage_bytes = 0; - EXPECT_FALSE(graph_execution_storage_bytes(0, 0, sizeof(GraphDefinition), &storage_bytes)); - EXPECT_FALSE(graph_execution_storage_bytes(1, 0, SIZE_MAX, &storage_bytes)); + 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)); } TEST(GraphExecutionReplay, AffineHitRefreshesOnlyDynamicFields) { @@ -190,10 +254,10 @@ TEST(GraphExecutionReplay, AffineHitRefreshesOnlyDynamicFields) { const std::vector definition = make_test_definition(GRAPH_KEY_VALUE, reinterpret_cast(first_boundary.data())); size_t execution_bytes = 0; - ASSERT_TRUE(graph_execution_storage_bytes(2, 2, definition.size(), &execution_bytes)); + ASSERT_TRUE(graph_execution_storage_bytes(2, 2, 2, definition.size(), &execution_bytes)); AlignedStorage execution_storage(execution_bytes); std::vector submission_image = make_test_submission( - GRAPH_KEY_VALUE, reinterpret_cast(first_boundary.data()), + GRAPH_KEY_VALUE, reinterpret_cast(first_boundary.data()), 17, reinterpret_cast(execution_storage.data()), execution_storage.size() ); auto &submission = *reinterpret_cast(submission_image.data()); @@ -213,6 +277,8 @@ TEST(GraphExecutionReplay, AffineHitRefreshesOnlyDynamicFields) { GraphNodeStorage &node = execution->node_storage[0]; ASSERT_EQ(node.payload.scalar_count, 1); ASSERT_EQ(node.payload.tensor_count, 1); + EXPECT_EQ(node.payload.scalars[0], 17U); + EXPECT_EQ(execution->node_storage[1].payload.scalars[0], 18U); graph_execution_mark_completed(*execution); execution->retired_nodes.store(2, std::memory_order_release); @@ -222,6 +288,8 @@ TEST(GraphExecutionReplay, AffineHitRefreshesOnlyDynamicFields) { outer_task.packed_buffer_end = second_heap.data() + second_heap.size(); auto *boundary = reinterpret_cast(submission_image.data() + submission.tensors_offset); boundary->buffer_addr = reinterpret_cast(second_boundary.data()); + auto *boundary_scalar = reinterpret_cast(submission_image.data() + submission.scalars_offset); + *boundary_scalar = 99; execution = graph_execution_localize(outer_slot); ASSERT_NE(execution, nullptr); @@ -233,6 +301,7 @@ TEST(GraphExecutionReplay, AffineHitRefreshesOnlyDynamicFields) { node.task.kernel_id[0] = 314; node.slot.active_mask = ActiveMask(3); node.payload.scalars[0] = 2718; + execution->node_storage[1].payload.scalars[0] = 31415; node.payload.tensors[0].version = 1618; node.slot.completed_subtasks.store(1, std::memory_order_relaxed); node.payload.dispatch_fanin.store(1, std::memory_order_relaxed); @@ -240,7 +309,8 @@ TEST(GraphExecutionReplay, AffineHitRefreshesOnlyDynamicFields) { EXPECT_EQ(graph_execution_materialize_slice(outer_slot, *execution, 2), GraphMaterializeResult::PREPARED); EXPECT_EQ(node.task.kernel_id[0], 314); EXPECT_EQ(node.slot.active_mask.raw(), 3); - EXPECT_EQ(node.payload.scalars[0], 2718U); + EXPECT_EQ(node.payload.scalars[0], 99U); + EXPECT_EQ(execution->node_storage[1].payload.scalars[0], 31415U); EXPECT_EQ(node.payload.tensors[0].version, 1618); EXPECT_EQ(node.task.task_id, PTO2TaskId::make(1, (8U << 10U))); EXPECT_EQ(node.task.packed_buffer_base, second_heap.data());