Fix: harden HBG Graph Execution across A2/A3 and A5

docs/troubleshooting/device-error-codes.md

@@ -29,7 +29,9 @@ grep -E "orch_error_code=|sched_error_code=|sub_class=|error detail:" <run log>
 
 ## How an error reaches you
 
-The `tensormap_and_ringbuffer` runtime runs orchestration and scheduling on the
+The `host_build_graph` runtime runs orchestration on the host, transfers the
+prepared image to the AICPU, and runs scheduling there. The
+`tensormap_and_ringbuffer` runtime runs both orchestration and scheduling on the
 AICPU. On a fatal condition the runtime **latches** a code into the shared-memory
 header; the host reads it back in `validate_runtime_impl` and prints the lines
 above.
@@ -77,6 +79,7 @@ layer to go looking in, which is what these columns are for.
 | 101 | ASYNC_COMPLETION_INVALID | kernel (async) |
 | 102 | ASYNC_WAIT_OVERFLOW | kernel (async) |
 | 103 | ASYNC_REGISTRATION_FAILED | runtime-internal |
+| 104 | READY_QUEUE_OVERFLOW | runtime-internal / config |
 
 ### SCHEDULER_TIMEOUT sub-classes
 
@@ -139,9 +142,9 @@ always *fallout* — scroll up for the first failure on that device. For classif
 
 ## Minimal reproductions
 
-Each code has a live, minimal trigger in `tests/st/runtime_fatal_codes/`. These are
-the fastest way to see what a code looks like, and the shape to copy when you
-suspect one:
+Each publicly triggerable code listed below has a live, minimal trigger in
+`tests/st/runtime_fatal_codes/`. These are the fastest way to see what a code
+looks like, and the shape to copy when you suspect one:
 
 | Code | How the ST provokes it | Fixture |
 | ---- | ---------------------- | ------- |
@@ -166,7 +169,7 @@ enforces coverage. Edit those and the log carries the new code correctly:
 | ---- | ----- |
 | runtime code names / descriptions / hints | `src/common/runtime_status/error_names.h` |
 | host-side CANN names / descriptions / hints | `src/common/platform/include/host/acl_error_names.h` |
-| `SCHEDULER_TIMEOUT` sub-class labels | `src/{arch}/runtime/tensormap_and_ringbuffer/common/pto_runtime_status.h` |
+| `SCHEDULER_TIMEOUT` sub-class labels | `src/{arch}/runtime/{host_build_graph,tensormap_and_ringbuffer}/common/pto_runtime_status.h` |
 | completeness test | `tests/ut/cpp/common/test_error_code_names.cpp` |
 
 **This page does not need updating for a new code** — deliberately. The tables
@@ -176,7 +179,7 @@ time, so there is nothing to drift out of sync.
 
 ## References
 
-- Code definitions: `src/{arch}/runtime/tensormap_and_ringbuffer/common/pto_runtime_status.h`
+- Code definitions: `src/{arch}/runtime/{host_build_graph,tensormap_and_ringbuffer}/common/pto_runtime_status.h`
 - Host print site: `.../host/runtime_maker.cpp` (`validate_runtime_impl`)
 - Sub-class logic: `.../runtime/scheduler/scheduler_cold_path.cpp` (`classify_stall_reason`)
 - End-to-end negative tests: `tests/st/runtime_fatal_codes/`

src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp

@@ -406,12 +406,17 @@ bool upload_graph_submissions(Runtime *runtime, const HostApi *api, GraphHostSta
     const size_t count = graph_host_upload_count(graph_state);
     for (size_t index = 0; index < count; ++index) {
         std::optional<GraphHostUpload> upload = graph_host_upload(graph_state, index);
-        if (!upload.has_value() || upload->outer_slot->task_kind != TaskKind::GRAPH ||
+        if (!upload.has_value() || upload->outer_slot == nullptr || upload->data == nullptr ||
+            upload->bytes < sizeof(GraphSubmission) || upload->outer_slot->task_kind != TaskKind::GRAPH ||
             upload->outer_slot->task == nullptr) {
             LOG_ERROR("host-orch: invalid pending Graph POD image");
             return false;
         }
         auto *submission = reinterpret_cast<GraphSubmission *>(upload->data);
+        if (!graph_submission_wire_size_valid(*submission, upload->bytes)) {
+            LOG_ERROR("host-orch: Graph submission size does not match its POD image");
+            return false;
+        }
         const GraphDefinition *definition = graph_submission_definition(*submission);
         size_t execution_bytes = 0;
         if (definition == nullptr || definition->full_key != submission->graph_key || definition->task_count == 0 ||

src/a2a3/runtime/host_build_graph/runtime/graph_execution.h

@@ -243,9 +243,18 @@ inline const GraphDefinition *graph_submission_definition(const GraphSubmission
         sizeof(GraphDefinition) > submission.total_bytes - submission.definition_offset) {
         return nullptr;
     }
-    return reinterpret_cast<const GraphDefinition *>(
+    const auto *definition = reinterpret_cast<const GraphDefinition *>(
         reinterpret_cast<const uint8_t *>(&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;
 }
 
 inline const GraphTensor *graph_submission_tensors(const GraphSubmission &submission) {
@@ -366,6 +375,8 @@ struct GraphExecution {
     uint32_t boundary_scalar_count{0};
 };
 
+static_assert(offsetof(GraphExecution, storage_magic) == 0);
+static_assert(sizeof(GraphExecution::storage_magic) == sizeof(uint64_t));
 static_assert(std::is_trivially_destructible_v<GraphNodeStorage>);
 static_assert(std::is_trivially_destructible_v<GraphExecution>);
 
@@ -442,7 +453,7 @@ inline void graph_execution_retire_node(GraphExecution &execution) {
 inline bool graph_submission_signal(GraphSubmission &submission, uint32_t bit) {
     constexpr uint32_t BOTH = 0x3;
     uint32_t observed = __atomic_fetch_or(&submission.activation_gate, bit, __ATOMIC_ACQ_REL);
-    return (observed | bit) == BOTH;
+    return observed != BOTH && (observed | bit) == BOTH;
 }
 
 inline GraphExecution *graph_submission_local_execution(GraphSubmission &submission) {

src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h

@@ -9,8 +9,7 @@
  * -----------------------------------------------------------------------------------------------------------
  */
 
-#ifndef PTO_ASYNC_WAIT_H
-#define PTO_ASYNC_WAIT_H
+#pragma once
 
 #include <atomic>
 #include <cstddef>
@@ -124,7 +123,8 @@ struct AsyncWaitEntry {
 };
 
 struct AsyncPollResult {
-    int32_t completed{0};
+    int32_t completed{0};  // Host-submitted stream tasks completed.
+    int32_t resolved{0};   // All task completions, including internal Graph nodes.
     int32_t error_code{PTO2_ERROR_NONE};
     PTO2TaskSlotState *failed_slot_state{nullptr};
 };
@@ -174,6 +174,8 @@ struct AsyncWaitList {
     struct DrainCompletionSink {
         PTO2SchedulerState *sched{nullptr};
         int32_t inline_completed{0};
+        int32_t inline_resolved{0};
+        int32_t error_code{PTO2_ERROR_NONE};
 #if SIMPLER_SCHED_PROFILING
         int32_t thread_idx{0};
 #endif
@@ -236,7 +238,10 @@ struct AsyncWaitList {
                     // conditions => NotDeferred. Complete it inline when the
                     // sink allows; otherwise fall back to the entry-store path.
                     if (sink.can_inline_complete()) {
-                        (void)try_inline_complete_locked(sink, *slot_state_ptr);
+                        if (!try_inline_complete_locked(sink, *slot_state_ptr)) {
+                            error_code = sink.error_code;
+                            return drained;
+                        }
                         continue;
                     }
                     if (count >= MAX_ASYNC_WAITS) {
@@ -295,5 +300,3 @@ struct AsyncWaitList {
 #endif
     );
 };
-
-#endif  // PTO_ASYNC_WAIT_H

src/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cpp

@@ -565,6 +565,10 @@ GraphMaterializeResult graph_execution_materialize_slice(
             }
         } else {
             for (int32_t j = 0; j < payload.tensor_count; ++j) {
+                if (execution.materialized_tensor_patches >= execution.materialized_tensor_patch_count) {
+                    execution.materialize_busy.store(0, std::memory_order_release);
+                    return GraphMaterializeResult::INVALID;
+                }
                 const GraphTensorAddressPatch &patch =
                     execution.tensor_patches[execution.materialized_tensor_patches++];
                 if (patch.source == static_cast<uint8_t>(GraphTensorAddressSource::BOUNDARY)) {

src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h

@@ -488,6 +488,25 @@ struct PTO2SchedulerState {
     // dummy_ready_queue and are retired inline; a ready sync_start cohort goes to
     // the per-shape ready_sync_queues[] (drained as Tier-0); everything else to
     // ready_queues[].
+    void latch_ready_queue_overflow(int32_t thread_idx = -1) {
+        int32_t expected = PTO2_ERROR_NONE;
+        const bool latched = sm_header->sched_error_code.compare_exchange_strong(
+            expected, PTO2_ERROR_READY_QUEUE_OVERFLOW, std::memory_order_acq_rel, std::memory_order_acquire
+        );
+        if (latched && thread_idx >= 0) {
+            sm_header->sched_error_thread.store(thread_idx, std::memory_order_release);
+        }
+        if (thread_idx >= 0 && thread_idx < 32) {
+            sm_header->sched_error_bitmap.fetch_or(1U << static_cast<uint32_t>(thread_idx), std::memory_order_acq_rel);
+        }
+    }
+
+    bool push_graph_prepare(PTO2TaskSlotState *slot_state, uint64_t task_id, int32_t thread_idx) {
+        if (graph_prepare_queue.push_tagged(slot_state, task_id)) return true;
+        latch_ready_queue_overflow(thread_idx);
+        return false;
+    }
+
     void push_ready_routed(PTO2TaskSlotState *slot_state) {
         bool pushed;
         if (slot_state->task_kind == TaskKind::GRAPH) {
@@ -511,10 +530,7 @@ struct PTO2SchedulerState {
         // forward-progress timeout). The graph_ready push is checked identically
         // so a graph task cannot be dropped either.
         if (!pushed) {
-            int32_t expected = PTO2_ERROR_NONE;
-            sm_header->sched_error_code.compare_exchange_strong(
-                expected, PTO2_ERROR_READY_QUEUE_OVERFLOW, std::memory_order_acq_rel, std::memory_order_acquire
-            );
+            latch_ready_queue_overflow();
         }
     }
 
@@ -959,6 +975,7 @@ struct PTO2SchedulerState {
     struct TaskCompletionOutcome {
         uint32_t fanout_edges{0};
         int32_t stream_tasks_completed{0};
+        int32_t error_code{PTO2_ERROR_NONE};
     };
 
     TaskCompletionOutcome complete_task(
@@ -981,11 +998,21 @@ struct PTO2SchedulerState {
         }
 
         GraphExecution *execution = graph_execution_from_slot(slot_state);
-        if (execution == nullptr || execution->definition == nullptr || execution->nodes == nullptr) return outcome;
+        if (execution == nullptr || execution->definition == nullptr || execution->nodes == nullptr ||
+            execution->state.load(std::memory_order_acquire) != GraphExecutionState::ACTIVE) {
+            outcome.error_code = PTO2_ERROR_INVALID_ARGS;
+            return outcome;
+        }
         const int32_t saved_node_index = slot_state.graph_node_index;
-        if (saved_node_index < 0) return outcome;
+        if (saved_node_index < 0) {
+            outcome.error_code = PTO2_ERROR_INVALID_ARGS;
+            return outcome;
+        }
         const uint32_t node_index = static_cast<uint32_t>(saved_node_index);
-        if (node_index >= static_cast<uint32_t>(execution->node_count)) return outcome;
+        if (node_index >= static_cast<uint32_t>(execution->node_count)) {
+            outcome.error_code = PTO2_ERROR_INVALID_ARGS;
+            return outcome;
+        }
 
         // Publish completion before closing the wake list. A consumer that
         // loses registration to the sentinel acquires this state when it
@@ -1113,7 +1140,12 @@ AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &si
 #else
     PTO2SchedulerState::TaskCompletionOutcome outcome = sink.sched->complete_task(slot_state);
 #endif
+    if (outcome.error_code != PTO2_ERROR_NONE) {
+        sink.error_code = outcome.error_code;
+        return false;
+    }
     sink.inline_completed += outcome.stream_tasks_completed;
+    sink.inline_resolved++;
     return true;
 }
 
@@ -1142,6 +1174,7 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete(
         return result;
     }
     result.completed += sink.inline_completed;
+    result.resolved += sink.inline_resolved;
 
     for (int32_t i = count - 1; i >= 0; --i) {
         AsyncWaitEntry &entry = entries[i];
@@ -1178,8 +1211,15 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete(
 #else
             PTO2SchedulerState::TaskCompletionOutcome outcome = sched->complete_task(*entry.slot_state);
 #endif
+            if (outcome.error_code != PTO2_ERROR_NONE) {
+                result.error_code = outcome.error_code;
+                result.failed_slot_state = entry.slot_state;
+                unlock();
+                return result;
+            }
             // Polling: completion is fully published inline; no deferred release.
             result.completed += outcome.stream_tasks_completed;
+            result.resolved++;
 
             int32_t last = count - 1;
             if (i != last) entries[i] = entries[last];

src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp

@@ -45,6 +45,13 @@ static void latch_scheduler_error(PTO2SharedMemoryHeader *header, int32_t thread
     }
 }
 
+void SchedulerContext::fail_scheduler(Runtime *runtime, int32_t thread_idx, int32_t error_code) {
+    latch_scheduler_error(sched_ == nullptr ? nullptr : sched_->sm_header, thread_idx, error_code);
+    if (!completed_.exchange(true, std::memory_order_acq_rel)) {
+        emergency_shutdown(runtime);
+    }
+}
+
 LoopAction SchedulerContext::handle_orchestrator_exit(
     int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime, int32_t &task_count
 ) {
@@ -1140,9 +1147,7 @@ void SchedulerContext::classify_partition(int32_t thread_idx, int32_t nthreads)
         }
         PTO2TaskSlotState &slot = ring.get_slot_state_by_task_id(id);
         if (slot.task_kind == TaskKind::GRAPH) {
-            while (!sched_->graph_prepare_queue.push_tagged(&slot, slot.task->task_id.raw)) {
-                SPIN_WAIT_HINT();
-            }
+            if (!sched_->push_graph_prepare(&slot, slot.task->task_id.raw, thread_idx)) return;
         }
         int32_t state = sched_->classify_fanin_state(&slot);
         if (state < 0) {

src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h

@@ -8,8 +8,7 @@
  * See LICENSE in the root of the software repository for the full text of the License.
  * -----------------------------------------------------------------------------------------------------------
  */
-#ifndef SCHEDULER_CONTEXT_H
-#define SCHEDULER_CONTEXT_H
+#pragma once
 
 #include "aicpu/device_phase_aicpu.h"
 #include "aicpu/platform_regs.h"
@@ -272,6 +271,8 @@ class SchedulerContext {
     // deinit their AICore register blocks. Idempotent.
     void emergency_shutdown(Runtime *runtime);
 
+    __attribute__((noinline, cold)) void fail_scheduler(Runtime *runtime, int32_t thread_idx, int32_t error_code);
+
     // =========================================================================
     // Dispatch (scheduler_dispatch.cpp)
     // =========================================================================
@@ -575,5 +576,3 @@ class SchedulerContext {
         return func_id_to_addr_[func_id];
     }
 };
-
-#endif  // SCHEDULER_CONTEXT_H