[Optimization] Replace wiring with polling-based task readiness test (~17% median device speedup)#1137
Open
SergioMartin86 wants to merge 8 commits into
Open
[Optimization] Replace wiring with polling-based task readiness test (~17% median device speedup)#1137SergioMartin86 wants to merge 8 commits into
SergioMartin86 wants to merge 8 commits into
Conversation
📝 WalkthroughWalkthroughThis PR redesigns the PTO2 runtime/scheduler polling architecture by moving orchestrator, scheduler, ring-buffer, shared-memory, and runtime lifecycle implementations from ChangesPTO2 Runtime/Scheduler Polling Redesign
Sequence Diagram(s)sequenceDiagram
participant AicpuExecutor
participant OrchestratorThread as Orchestrator Thread
participant SchedulerThread as Scheduler Thread(s)
participant PTO2OrchestratorState as Orchestrator State
participant PTO2SchedulerState as Scheduler State
AicpuExecutor->>OrchestratorThread: run() - orchestrator path
OrchestratorThread->>OrchestratorThread: dep_gen_aicpu_init()
OrchestratorThread->>OrchestratorThread: sm_handle->init(task_window_size, heap_size)
OrchestratorThread->>AicpuExecutor: signal runtime_init_ready_
SchedulerThread->>AicpuExecutor: wait on runtime_init_ready_
OrchestratorThread->>PTO2OrchestratorState: rt_scope_begin / orch_func() / rt_scope_end
PTO2OrchestratorState->>PTO2SchedulerState: push slot_state via wiring SpscQueue
OrchestratorThread->>OrchestratorThread: rt_orchestration_done()
OrchestratorThread->>OrchestratorThread: dep_gen_aicpu_flush()
OrchestratorThread->>PTO2SchedulerState: on_orchestration_done()
SchedulerThread->>PTO2SchedulerState: resolve_and_dispatch()
PTO2SchedulerState->>PTO2SchedulerState: drain_wiring_queue → classify_fanin_state
PTO2SchedulerState->>PTO2SchedulerState: on_mixed_task_complete → publish completion_flags
PTO2SchedulerState->>PTO2SchedulerState: advance completed_watermark
SchedulerThread->>AicpuExecutor: return rc (last thread calls runtime_destroy + deinit)
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the runtime by inlining orchestrator, scheduler, and ring buffer logic into header files to support a polling-based completion design. The code review identified several critical issues: potential truncation and undefined behavior in the CPU affinity gate (g_cpumask) for CPU IDs of 16 or greater; a lack of zero-initialization for completion_flags in shared memory, which could cause premature slot reclamation and memory corruption; missing bounds checks on dep_task_id.ring() before indexing ring headers; insufficient capacity in the pending FIFO (pending_capacity) under heavy multi-ring workloads; relaxed memory ordering in register_wake that could lead to memory visibility violations; and potential undefined behavior in __builtin_ctz if new_num_buckets is zero.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| static std::atomic<int32_t> s_reported{0}; | ||
| static std::atomic<int32_t> s_gate_init{0}; | ||
| static std::atomic<int32_t> s_gate_ready{0}; | ||
| static std::atomic<uint16_t> g_cpumask{0}; |
There was a problem hiding this comment.
Using uint16_t for g_cpumask can lead to truncation and deadlocks if sched_getcpu() returns a CPU ID of 16 or greater, as 1 << cpu will be truncated to 0 when OR'd into g_cpumask. Additionally, 1 << cpu is undefined behavior if cpu >= 31 because 1 is a 32-bit signed integer. It is safer to use uint64_t for g_cpumask and 1ULL << cpu for the bitwise operations, along with __builtin_popcountll.
Suggested change
| static std::atomic<uint16_t> g_cpumask{0}; | |
| static std::atomic<uint64_t> g_cpumask{0}; |
Comment on lines
+54
to
+61
| g_cpumask.fetch_or(1 << cpu, std::memory_order_relaxed); | ||
|
|
||
| struct ClusterInfo { | ||
| int32_t count{0}; | ||
| int32_t tids[MAX_GATE_THREADS]; | ||
| }; | ||
| ClusterInfo clusters[MAX_CLUSTERS]; | ||
| // Barrier wait until all the spawned threads are here before choosing which ones will be used. | ||
| while(__builtin_popcount(g_cpumask) < total_launched) {} | ||
|
|
||
| for (int32_t tid = 0; tid < total_launched; ++tid) { | ||
| int32_t c = s_thread_cpu[tid]; | ||
| if (c < 0) continue; | ||
| int32_t cluster_id = c / CPUS_PER_CLUSTER; | ||
| if (cluster_id < 0 || cluster_id >= MAX_CLUSTERS) continue; | ||
| ClusterInfo &info = clusters[cluster_id]; | ||
| if (info.count < MAX_GATE_THREADS) info.tids[info.count++] = tid; | ||
| } | ||
| // Choose the thread based on reverse bit order (highest cpu id to lowest) | ||
| // This assures that all the threads lie in the same NUMA domain | ||
| int how_many_on_top = __builtin_popcount(g_cpumask >> cpu); |
There was a problem hiding this comment.
To support g_cpumask as a uint64_t and prevent undefined behavior or truncation for CPU IDs >= 16, use 1ULL << cpu and __builtin_popcountll.
Suggested change
| g_cpumask.fetch_or(1 << cpu, std::memory_order_relaxed); | |
| struct ClusterInfo { | |
| int32_t count{0}; | |
| int32_t tids[MAX_GATE_THREADS]; | |
| }; | |
| ClusterInfo clusters[MAX_CLUSTERS]; | |
| // Barrier wait until all the spawned threads are here before choosing which ones will be used. | |
| while(__builtin_popcount(g_cpumask) < total_launched) {} | |
| for (int32_t tid = 0; tid < total_launched; ++tid) { | |
| int32_t c = s_thread_cpu[tid]; | |
| if (c < 0) continue; | |
| int32_t cluster_id = c / CPUS_PER_CLUSTER; | |
| if (cluster_id < 0 || cluster_id >= MAX_CLUSTERS) continue; | |
| ClusterInfo &info = clusters[cluster_id]; | |
| if (info.count < MAX_GATE_THREADS) info.tids[info.count++] = tid; | |
| } | |
| // Choose the thread based on reverse bit order (highest cpu id to lowest) | |
| // This assures that all the threads lie in the same NUMA domain | |
| int how_many_on_top = __builtin_popcount(g_cpumask >> cpu); | |
| g_cpumask.fetch_or(1ULL << cpu, std::memory_order_relaxed); | |
| // Barrier wait until all the spawned threads are here before choosing which ones will be used. | |
| while(__builtin_popcountll(g_cpumask) < total_launched) {} | |
| // Choose the thread based on reverse bit order (highest cpu id to lowest) | |
| // This assures that all the threads lie in the same NUMA domain | |
| int how_many_on_top = __builtin_popcountll(g_cpumask >> cpu); |
Comment on lines
+211
to
+217
| for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) | ||
| {} | ||
| } | ||
| bool validate() | ||
| { | ||
| if (!sm_base) return false; | ||
| if (!header) return false; |
There was a problem hiding this comment.
The completion_flags array in shared memory is not zero-initialized during init_header_per_ring. Since shared memory is not guaranteed to be zero-initialized on the device, completion_flags can contain garbage non-zero values. If a task completes out of order, the watermark-advancing loop in on_mixed_task_complete will read these garbage non-zero values for uncompleted intermediate slots and prematurely advance completed_watermark past them. This will lead to premature slot reclamation and severe memory corruption. completion_flags must be explicitly zero-initialized (e.g., using memset) during init_header_per_ring.
Suggested change
| for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) | |
| {} | |
| } | |
| bool validate() | |
| { | |
| if (!sm_base) return false; | |
| if (!header) return false; | |
| for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) | |
| { | |
| header->rings[r].fc.init(); | |
| // -1 = "no task completed yet"; first task to complete (local_id 0) | |
| // will advance the watermark to 0. | |
| header->rings[r].completed_watermark.store(-1, std::memory_order_relaxed); | |
| memset(header->rings[r].completion_flags, 0, task_window_sizes[r] * sizeof(std::atomic<uint8_t>)); | |
| } |
References
- Do not assume that allocated shared memory or device memory is zero-initialized. Always explicitly initialize all fields (such as thread/core counts and mapping arrays) to prevent garbage values from causing segmentation faults or undefined behavior.
Comment on lines
+560
to
+564
| if (!dep_task_id.is_valid()) | ||
| { | ||
| orch->report_fatal(PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "Arg.set_dependencies(...) requires valid task ids"); | ||
| return result; | ||
| } |
There was a problem hiding this comment.
The dep_task_id.ring() value is used to index orch->sm_header->rings on line 565 without validating that it is within bounds. Since dep_task_id.ring() can return a signed integer, we must perform bounds checks (checking both if the index is negative and if it exceeds the maximum count) to prevent undefined behavior, out-of-bounds memory access, or signed-to-unsigned conversion overflows. A check should be added to ensure dep_task_id.ring() >= 0 && dep_task_id.ring() < PTO2_MAX_RING_DEPTH.
if (!dep_task_id.is_valid() || dep_task_id.ring() < 0 || dep_task_id.ring() >= PTO2_MAX_RING_DEPTH)
{
orch->report_fatal(PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "Arg.set_dependencies(...) requires valid task ids");
return result;
}References
- Ensure that index-based accessors using signed integers perform bounds checks (e.g., checking if the index is negative or exceeds the count) to prevent undefined behavior, out-of-bounds memory access, or signed-to-unsigned conversion overflows, maintaining consistency with other accessors in the codebase.
| PTO2SchedulerLayout layout{}; | ||
| layout.ready_queue_capacity = PTO2_READY_QUEUE_SIZE; | ||
| layout.spsc_capacity = PTO2_WRIRING_QUEUE_SIZE; | ||
| layout.pending_capacity = PTO2_TASK_WINDOW_SIZE; // bounded by per-ring slot window |
There was a problem hiding this comment.
The pending_capacity is set to PTO2_TASK_WINDOW_SIZE (16384), but the pending FIFO is shared across all PTO2_MAX_RING_DEPTH (4) rings. Since the total number of concurrent in-flight tasks across all rings can be up to PTO2_SCOPE_TASKS_CAP (65536), the pending FIFO can overflow under heavy multi-ring workloads. Because pending_push_back does not perform any overflow checks, this will silently overwrite pending tasks, leading to deadlocks or corruption. pending_capacity should be increased to PTO2_SCOPE_TASKS_CAP to guarantee safety.
Suggested change
| layout.pending_capacity = PTO2_TASK_WINDOW_SIZE; // bounded by per-ring slot window | |
| layout.pending_capacity = PTO2_SCOPE_TASKS_CAP; // bounded by total in-flight slot budget |
Comment on lines
+508
to
+519
| PTO2TaskSlotState *expected = producer->wake_list_head.load(std::memory_order_relaxed); | ||
| while (true) | ||
| { | ||
| if (expected == WAKE_LIST_SENTINEL) | ||
| { | ||
| // Producer already completed and drained its wake list. The | ||
| // last unmet fanin is now satisfied; push consumer to ready. | ||
| push_ready_routed(consumer); | ||
| return; | ||
| } | ||
| } | ||
| // This pre-staged consumer was just released by its doorbell — it starts | ||
| // running NOW, so propagate dispatch_fanin to ITS consumers (auto-chain, | ||
| // knob A). Defer it via the sink so it runs after the whole fanout walk: | ||
| // doing it inline here would delay the doorbells of later consumers in the | ||
| // same producer's fanout. Fallback to inline if no sink / sink full. | ||
| if (sink == nullptr || !sink->push(&slot_state)) { | ||
| propagate_dispatch_fanin(slot_state); | ||
| } | ||
| // No explicit removal from the cross-thread queue: a still-queued entry for | ||
| // this consumer is now DISPATCHED and is dropped when a peer pops it. | ||
| // Fully pre-staged => skip the ready queue. Partially staged SPMD consumer => | ||
| // fall through so the caller pushes C; dispatch resumes from next_block_idx. | ||
| return slot_state.next_block_idx.load(std::memory_order_seq_cst) >= slot_state.logical_block_num; | ||
| } | ||
|
|
||
| bool release_fanin_and_check_ready( | ||
| PTO2TaskSlotState &slot_state, PTO2LocalReadyBuffer *local_bufs = nullptr, SpecReleaseSink *sink = nullptr | ||
| ) { | ||
| // Atomically increment fanin_refcount and check if all producers are done | ||
| // ACQ_REL on fanin_refcount already synchronizes with the orchestrator's | ||
| // init release, making fanin_count visible — plain load suffices. | ||
| int32_t new_refcount = slot_state.fanin_refcount.fetch_add(1, std::memory_order_acq_rel) + 1; | ||
|
|
||
| if (new_refcount == slot_state.fanin_count) { | ||
| // Speculative early-dispatch: pre-staged tasks are released by doorbell | ||
| // here, skipping the ready-queue round-trip entirely. | ||
| if (try_speculative_release(slot_state, sink)) return true; | ||
| // Local-first: try per-CoreType thread-local buffer before global queue | ||
| // Route by active_mask: AIC-containing tasks → buf[0], AIV-only → buf[1] | ||
| // DUMMY shape is out of range for local_bufs (sized PTO2_NUM_RESOURCE_SHAPES); | ||
| // dummy slots bypass the local fast path and go straight to dummy_ready_queue. | ||
| PTO2ResourceShape shape = slot_state.active_mask.to_shape(); | ||
| if (shape == PTO2ResourceShape::DUMMY) { | ||
| dummy_ready_queue.push(&slot_state); | ||
| } else if (!local_bufs || !local_bufs[static_cast<int32_t>(shape)].try_push(&slot_state)) { | ||
| ready_queues[static_cast<int32_t>(shape)].push(&slot_state); | ||
| consumer->next_in_wake_list = expected; | ||
| if (producer->wake_list_head.compare_exchange_weak(expected, consumer, std::memory_order_acq_rel, std::memory_order_relaxed)) |
There was a problem hiding this comment.
In register_wake, the initial load of wake_list_head and the failure memory order of compare_exchange_weak both use std::memory_order_relaxed. If expected == WAKE_LIST_SENTINEL is observed, the consumer is immediately routed to the ready queue. However, because of the relaxed memory order, there is no acquire synchronization with the producer's completion writes. This can allow the consumer to start executing on a worker core while seeing stale producer outputs (memory visibility violation). To ensure correctness, the initial load and the CAS failure order should both use std::memory_order_acquire.
Suggested change
| PTO2TaskSlotState *expected = producer->wake_list_head.load(std::memory_order_relaxed); | |
| while (true) | |
| { | |
| if (expected == WAKE_LIST_SENTINEL) | |
| { | |
| // Producer already completed and drained its wake list. The | |
| // last unmet fanin is now satisfied; push consumer to ready. | |
| push_ready_routed(consumer); | |
| return; | |
| } | |
| } | |
| // This pre-staged consumer was just released by its doorbell — it starts | |
| // running NOW, so propagate dispatch_fanin to ITS consumers (auto-chain, | |
| // knob A). Defer it via the sink so it runs after the whole fanout walk: | |
| // doing it inline here would delay the doorbells of later consumers in the | |
| // same producer's fanout. Fallback to inline if no sink / sink full. | |
| if (sink == nullptr || !sink->push(&slot_state)) { | |
| propagate_dispatch_fanin(slot_state); | |
| } | |
| // No explicit removal from the cross-thread queue: a still-queued entry for | |
| // this consumer is now DISPATCHED and is dropped when a peer pops it. | |
| // Fully pre-staged => skip the ready queue. Partially staged SPMD consumer => | |
| // fall through so the caller pushes C; dispatch resumes from next_block_idx. | |
| return slot_state.next_block_idx.load(std::memory_order_seq_cst) >= slot_state.logical_block_num; | |
| } | |
| bool release_fanin_and_check_ready( | |
| PTO2TaskSlotState &slot_state, PTO2LocalReadyBuffer *local_bufs = nullptr, SpecReleaseSink *sink = nullptr | |
| ) { | |
| // Atomically increment fanin_refcount and check if all producers are done | |
| // ACQ_REL on fanin_refcount already synchronizes with the orchestrator's | |
| // init release, making fanin_count visible — plain load suffices. | |
| int32_t new_refcount = slot_state.fanin_refcount.fetch_add(1, std::memory_order_acq_rel) + 1; | |
| if (new_refcount == slot_state.fanin_count) { | |
| // Speculative early-dispatch: pre-staged tasks are released by doorbell | |
| // here, skipping the ready-queue round-trip entirely. | |
| if (try_speculative_release(slot_state, sink)) return true; | |
| // Local-first: try per-CoreType thread-local buffer before global queue | |
| // Route by active_mask: AIC-containing tasks → buf[0], AIV-only → buf[1] | |
| // DUMMY shape is out of range for local_bufs (sized PTO2_NUM_RESOURCE_SHAPES); | |
| // dummy slots bypass the local fast path and go straight to dummy_ready_queue. | |
| PTO2ResourceShape shape = slot_state.active_mask.to_shape(); | |
| if (shape == PTO2ResourceShape::DUMMY) { | |
| dummy_ready_queue.push(&slot_state); | |
| } else if (!local_bufs || !local_bufs[static_cast<int32_t>(shape)].try_push(&slot_state)) { | |
| ready_queues[static_cast<int32_t>(shape)].push(&slot_state); | |
| consumer->next_in_wake_list = expected; | |
| if (producer->wake_list_head.compare_exchange_weak(expected, consumer, std::memory_order_acq_rel, std::memory_order_relaxed)) | |
| PTO2TaskSlotState *expected = producer->wake_list_head.load(std::memory_order_acquire); | |
| while (true) | |
| { | |
| if (expected == WAKE_LIST_SENTINEL) | |
| { | |
| // Producer already completed and drained its wake list. The | |
| // last unmet fanin is now satisfied; push consumer to ready. | |
| push_ready_routed(consumer); | |
| return; | |
| } | |
| consumer->next_in_wake_list = expected; | |
| if (producer->wake_list_head.compare_exchange_weak(expected, consumer, std::memory_order_acq_rel, std::memory_order_acquire)) |
| static PTO2TensorMapLayout reserve_layout(DeviceArena &arena, int32_t new_num_buckets, int32_t new_pool_size, const int32_t new_task_window_sizes[PTO2_MAX_RING_DEPTH]) | ||
| { | ||
| // num_buckets must be a power of two for the hash truncation to work. | ||
| always_assert((new_num_buckets & (new_num_buckets - 1)) == 0); |
There was a problem hiding this comment.
If new_num_buckets is 0, the assertion (new_num_buckets & (new_num_buckets - 1)) == 0 will still pass, but it will lead to undefined behavior in __builtin_ctz(num_buckets) inside the hash function. The assertion should be updated to ensure new_num_buckets > 0.
Suggested change
| always_assert((new_num_buckets & (new_num_buckets - 1)) == 0); | |
| always_assert(new_num_buckets > 0 && (new_num_buckets & (new_num_buckets - 1)) == 0); |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
🟠 Major comments (18)
src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp-547-548 (1)
547-548: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRun
deinit()exactly once.Line 548 lets every thread that observes
finished_ == truecalldeinit(). Once the last worker flipsfinished_in Lines 480-495, slower threads can concurrently resetsched_ctx_, atomics, and dep-gen state, so teardown becomes nondeterministic. Gatedeinit()behind a one-shot exchange/counter and keep it on the actual last-finishing thread.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp` around lines 547 - 548, The teardown path in AICPUExecutor currently lets every thread that sees finished_ true call deinit(), which can race and reset shared state multiple times. Update the cleanup logic in aicpu_executor.cpp around the finished_ check in AICPUExecutor::deinit/runtime teardown so only the actual last-finishing thread performs deinit exactly once, using a one-shot exchange/counter or equivalent guard tied to the finished_ transition.src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp-239-277 (1)
239-277: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winClamp env overrides before heap-size multiplication.
parse_env_uint64()accepts anyuint64_t, then Line 262 multiplieseff_heap_sizebyPTO2_MAX_RING_DEPTHwith no overflow check. A largePTO2_RING_HEAPvalue can wraptotal_heap_size, under-allocate GM memory, and still leaveruntime_init_data_from_layout(..., eff_heap_size)configuring each ring as if the full heap were available. That turns a bad env setting into out-of-bounds device accesses. Reject values above the downstream multiplication/type limits before storing them onruntime.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp` around lines 239 - 277, Clamp the PTO2_RING_HEAP override before it is stored on runtime and before computing total_heap_size in runtime_maker.cpp, since parse_env_uint64() can accept values that later overflow the eff_heap_size * PTO2_MAX_RING_DEPTH multiplication. Add a bounds check alongside the existing PTO2_RING_DEP_POOL INT32_MAX guard, reject oversized heap values with LOG_ERROR, and make sure runtime_init_data_from_layout and related sizing logic only see validated heap sizes so the runtime cannot under-allocate the GM arena.src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp-211-217 (1)
211-217: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftPreserve tensor direction metadata for copy-back.
Line 217 now records every non-child tensor the same way, but Lines 434-445 later copy back every recorded pair and use
graph_output_ptrfor the first one. If the first tensor is an input, the packed output buffer gets written into the caller’s input storage. Please keep output/input identity inTensorPair(or record output indices separately) and only apply copy-back /graph_output_ptrto true outputs.Also applies to: 434-445
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp` around lines 211 - 217, The tensor copy-back path in runtime_maker.cpp is losing direction metadata when `runtime->tensor_pairs_` is populated in the tensor setup loop, so the later copy-back logic can mistakenly treat the first recorded pair as the graph output even if it is an input. Update `TensorPair` (or add separate output-index tracking) in the code around `runtime->tensor_pairs_.push_back` and the copy-back block that uses `graph_output_ptr` so only true output tensors are copied back and the caller’s input storage is never overwritten.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_shared_memory.h-184-188 (1)
184-188: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winValidate
task_window_sizebefore deriving a mask.
task_window_mask = task_window_size - 1only works for non-zero powers of two; otherwiseget_*_by_task_id()aliases task IDs to the wrong slots.Proposed fix
bool init(void *sm_base_arg, uint64_t sm_size_arg, uint64_t task_window_size, uint64_t heap_size) { if (!sm_base_arg || sm_size_arg == 0) return false; + if (task_window_size == 0 || (task_window_size & (task_window_size - 1)) != 0) return false; if (sm_size_arg < calculate_size(task_window_size)) return false;Also applies to: 256-257
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_shared_memory.h` around lines 184 - 188, Validate task_window_size in pto_shared_memory::init before any mask is derived or used, since task_window_mask = task_window_size - 1 only works for non-zero powers of two. Add an early guard in init to reject zero or non-power-of-two values, and ensure the get_*_by_task_id() helpers continue to rely on a valid task_window_mask so task IDs do not alias to incorrect slots.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h-356-380 (1)
356-380: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftAdvance allocator-visible reclamation for inline-completed allocations.
alloc_tensors()advances onlycompleted_watermark, butPTO2TaskAllocator::alloc()gates slot reuse onfc.last_task_alive. Since inline allocations never enteron_mixed_task_complete(), an alloc-only burst can fill the task window and spin untilPTO2_ALLOC_SPIN_LIMITeven though every slot is completed. After publishing the watermark, also trigger the same guardedadvance_ring_pointers()path used by scheduler completions.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h` around lines 356 - 380, The inline-completion path in the PTO orchestrator currently updates only completed_watermark, but allocator reuse in PTO2TaskAllocator::alloc still depends on fc.last_task_alive, so alloc-only bursts can stall even when every slot is done. In the inline-completed branch inside the completion handling logic, after advancing the per-ring completed_watermark, also invoke the same guarded advance_ring_pointers() path used by normal scheduler completions so allocator-visible reclamation progresses for inline allocations.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h-236-242 (1)
236-242: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse the latching helper for allocator deadlock errors.
report_deadlock()overwritesorch_error_codewithstore(), unlikeorch_mark_fatal()and the newlatch_pool_error()helper. This can hide the first/root fatal error if the allocator later times out while unwinding.Suggested fix
void report_deadlock(bool heap_blocked) { - if (error_code_ptr_) - { - int32_t code = heap_blocked ? PTO2_ERROR_HEAP_RING_DEADLOCK : PTO2_ERROR_FLOW_CONTROL_DEADLOCK; - error_code_ptr_->store(code, std::memory_order_release); - } + int32_t code = heap_blocked ? PTO2_ERROR_HEAP_RING_DEADLOCK : PTO2_ERROR_FLOW_CONTROL_DEADLOCK; + latch_pool_error(error_code_ptr_, code); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h` around lines 236 - 242, The deadlock reporting path in report_deadlock() is overwriting orch_error_code directly with store(), which can replace the original fatal error after unwinding starts. Update report_deadlock() to use the same latching behavior as orch_mark_fatal() and the new latch_pool_error() helper so allocator deadlock errors only set the code once and preserve the first/root failure.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h-507-520 (1)
507-520: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winReturn failure when scope tracking overflows.
scope_tasks_push()can latch a fatal error, butprepare_task()ignores that and returns success, so the current task can still be initialized/enqueued after the scope buffer is saturated. Make the push helper returnbooland stop task preparation on failure.Suggested fix
-inline void scope_tasks_push(PTO2OrchestratorState *orch, PTO2TaskSlotState *task_slot_state); +inline bool scope_tasks_push(PTO2OrchestratorState *orch, PTO2TaskSlotState *task_slot_state); ... - scope_tasks_push(orch, out->slot_state); + if (!scope_tasks_push(orch, out->slot_state)) return false; ... -inline void scope_tasks_push(PTO2OrchestratorState *orch, PTO2TaskSlotState *task_slot_state) +inline bool scope_tasks_push(PTO2OrchestratorState *orch, PTO2TaskSlotState *task_slot_state) { if (orch->scope_tasks_size >= orch->scope_tasks_capacity) { orch->report_fatal(PTO2_ERROR_SCOPE_TASKS_OVERFLOW, __FUNCTION__, "scope_tasks buffer saturated at %d entries (all rings full)", orch->scope_tasks_capacity); - return; + return false; } orch->scope_tasks[orch->scope_tasks_size++] = task_slot_state; + return true; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h` around lines 507 - 520, The scope tracking overflow path in scope_tasks_push is only latching a fatal error, but prepare_task still continues and returns success, so task initialization can proceed after the buffer is saturated. Change scope_tasks_push in PTO2OrchestratorState flow to return bool, have it report failure when orch->scope_tasks_size reaches orch->scope_tasks_capacity, and update prepare_task to check that return value and stop early by returning false if pushing the task slot state fails.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h-216-235 (1)
216-235: 🩺 Stability & Availability | 🟠 MajorMake scope bounds checks fatal in release builds.
begin_scope()andend_scope()still rely onassert(), but the release build path usesNDEBUG, so overflow/underflow can reach thescope_beginsaccess and post-decrement pop unchecked. Switch both toreport_fatal(PTO2_ERROR_INVALID_ARGS, ...)and return before mutating the stack.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h` around lines 216 - 235, The scope stack bounds checks in begin_scope() and end_scope() are only enforced by assert(), so release builds can still overflow or underflow into scope_begins access and stack mutation. Update the PTO2Orchestrator scope handling to use report_fatal(PTO2_ERROR_INVALID_ARGS, __FUNCTION__, ...) for both invalid stack states, and return immediately before incrementing scope_stack_top or reading scope_begins. Keep the existing in_manual_scope() and manual_begin_depth logic intact, but make the failure path fatal in all builds.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h-107-115 (1)
107-115: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winValidate
task_window_sizebefore casting and using mask-based slots.Line 114 truncates
uint64_t task_window_sizetoint32_t, and downstream code indexes withlocal_id & (task_window_size - 1). Reject zero, non-power-of-two, and values aboveINT32_MAXhere to avoid bad arena sizes and slot aliasing.Proposed fix
inline PTO2RuntimeArenaLayout runtime_reserve_layout(DeviceArena &arena, uint64_t task_window_size, int32_t dep_pool_capacity) { + always_assert(task_window_size > 0); + always_assert((task_window_size & (task_window_size - 1)) == 0); + always_assert(task_window_size <= static_cast<uint64_t>(INT32_MAX)); + PTO2RuntimeArenaLayout layout{};🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h` around lines 107 - 115, `runtime_reserve_layout` currently stores `task_window_size` into `PTO2RuntimeArenaLayout` and then uses it for mask-based slot indexing, so validate the input before any cast or layout use. In `runtime_reserve_layout`, reject zero, non-power-of-two values, and any value above `INT32_MAX` for `task_window_size`, and ensure the checks happen before filling `task_window_sizes` or assigning to the layout so downstream `local_id & (task_window_size - 1)` logic remains safe.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h-352-365 (1)
352-365: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winEnforce TensorMap mask invariants in the layout contract.
new_num_buckets == 0passes the current check but later makes__builtin_ctz(num_buckets)undefined. Also validate eachnew_task_window_sizes[r]as positive power-of-two becauseget_task_local_id_slot, cleanup, and task-entry linking all use& (size - 1).Proposed fix
// num_buckets must be a power of two for the hash truncation to work. - always_assert((new_num_buckets & (new_num_buckets - 1)) == 0); + always_assert(new_num_buckets > 0); + always_assert((new_num_buckets & (new_num_buckets - 1)) == 0); + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) + { + always_assert(new_task_window_sizes[r] > 0); + always_assert((new_task_window_sizes[r] & (new_task_window_sizes[r] - 1)) == 0); + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h` around lines 352 - 365, `reserve_layout` in `PTO2TensorMapLayout` only checks that `new_num_buckets` is a power of two, but it still allows zero and does not validate `new_task_window_sizes`, which later breaks `__builtin_ctz` and the mask-based slot logic used by `get_task_local_id_slot` and task-entry cleanup/linking. Tighten the layout contract in `reserve_layout` by asserting `new_num_buckets > 0` and that every `new_task_window_sizes[r]` is also a positive power of two before reserving any buffers, keeping the invariants enforced at the point where the layout is constructed.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h-126-153 (1)
126-153: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winInitialize the shared-memory handle before wiring it.
runtime_init_data_from_layoutzeroessm_wrap, andruntime_wire_arena_pointersexposes it viart->sm_handle, but the handle fields remain null/zero. Any code that readsrt->sm_handle->header,sm_base, orsm_sizewill observe an uninitialized handle despite a validsm_dev_base.Proposed fix
-inline PTO2Runtime *runtime_init_data_from_layout(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2RuntimeMode mode, void *sm_dev_base, uint64_t, void *gm_heap_dev_base, uint64_t heap_size) +inline PTO2Runtime *runtime_init_data_from_layout(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2RuntimeMode mode, void *sm_dev_base, uint64_t sm_size, void *gm_heap_dev_base, uint64_t heap_size) { @@ auto *sm_wrap = static_cast<PTO2SharedMemoryHandle *>(arena.region_ptr(layout.off_sm_handle)); memset(sm_wrap, 0, sizeof(*sm_wrap)); + sm_wrap->sm_base = sm_dev_base; + sm_wrap->sm_size = sm_size; + sm_wrap->header = reinterpret_cast<PTO2SharedMemoryHeader *>(sm_dev_base); + sm_wrap->is_owner = false;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h` around lines 126 - 153, Initialize the shared-memory handle before exposing it through runtime_wire_arena_pointers: runtime_init_data_from_layout currently zeros the PTO2SharedMemoryHandle at layout.off_sm_handle but never sets its fields, so populate that handle with sm_dev_base and the corresponding size/header information after the memset and before returning. Use the existing symbols runtime_init_data_from_layout, runtime_wire_arena_pointers, and PTO2SharedMemoryHandle to locate the fix, and ensure rt->sm_handle points to a fully initialized handle rather than a zeroed placeholder.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h-322-327 (1)
322-327: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMask
pending_occupied_for idle AIV dispatch too.Line 325 is now the only IDLE-path branch that ignores
pending_occupied_.CoreExecStatestill trackspending_slot_state/pending_reg_task_idfor every core, so an idle AIV core can still already have a staged payload. Returning it here lets the scheduler stage a second task onto the same core and clobber the first pending assignment.Suggested fix
BitStates get_idle_core_offset_states(PTO2ResourceShape shape) const { if (shape == PTO2ResourceShape::AIC) return get_valid_cluster_offset_states(shape) & ~(pending_occupied_ & aic_mask_); - if (shape == PTO2ResourceShape::AIV) return core_states_ & aiv_mask_; + if (shape == PTO2ResourceShape::AIV) return (core_states_ & aiv_mask_) & ~(pending_occupied_ & aiv_mask_); return get_valid_cluster_offset_states(shape); // MIX: cluster-level }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h` around lines 322 - 327, The idle-core selection in get_idle_core_offset_states currently masks pending_occupied_ for AIC only, but not for AIV, which can expose cores that already have staged work. Update the PTO2ResourceShape::AIV branch in get_idle_core_offset_states so it also excludes pending_occupied_ using the same pending-state mask logic as the AIC path, while keeping the MIX/cluster-level path unchanged.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.h-45-48 (1)
45-48: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
defer_error()no longer trips the scheduler's count-first error path.Line 47 only writes
completion_error_code. Downstream,complete_slot_task()readsdeferred_slab->countfirst and never inspectserror_codewhen that count is zero, so async registration failures from callers likeregister_pto_async_event()can be dropped silently. Publish the error through the same “non-empty slab” contract the scheduler already consumes.Suggested fix
inline __aicore__ void defer_error(AsyncCtx &ctx, int32_t error_code) { - if (ctx.task_token.is_valid() && ctx.completion_error_code != nullptr) *ctx.completion_error_code = error_code; + if (ctx.task_token.is_invalid() || ctx.completion_error_code == nullptr || ctx.completion_count == nullptr) return; + *ctx.completion_error_code = error_code; + if (*ctx.completion_count == 0) *ctx.completion_count = 1; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.h` around lines 45 - 48, The async error path in defer_error() only sets completion_error_code, so the scheduler’s count-first flow in complete_slot_task() can skip the failure when deferred_slab->count is zero. Update defer_error() to publish the error using the same non-empty slab contract the scheduler already checks, so callers like register_pto_async_event() still surface registration failures. Keep the fix localized around AsyncCtx, deferred_slab, and the existing completion bookkeeping so complete_slot_task() sees the error without relying on error_code alone.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h-271-275 (1)
271-275: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAvoid the unlocked
async_wait_list.countread.Line 272 reads
countoutsideAsyncWaitList::busy, whilepoll_and_complete()mutates it under the lock. Multiple scheduler threads can race here; either always callpoll_and_complete()when the mailbox exists, or expose an atomic/locked “has waits” check.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h` around lines 271 - 275, The scheduler logic in AsyncWaitList/AsyncPollResult currently reads async_wait_list.count outside the AsyncWaitList::busy lock, which can race with poll_and_complete() mutations. Update the check in scheduler_context so it does not touch count directly; either always invoke poll_and_complete() whenever rt_->aicore_mailbox exists, or add a locked/atomic “has waits” helper on AsyncWaitList and use that instead.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h-616-652 (1)
616-652: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winValidate the discovered AIC/AIV cluster shape before indexing.
Line 648 and Line 649 assume two AIV workers per AIC worker. If handshake discovery returns an incomplete topology, this reads uninitialized/out-of-range worker IDs. Return failure before assignment when
aiv_count_ < 2 * aic_count_.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h` around lines 616 - 652, In assign_cores_to_threads(), add a topology validation step before using aic_worker_ids_ and aiv_worker_ids_ so the discovered cluster shape is confirmed to be 1 AIC plus 2 AIV per cluster. Specifically, check that aiv_count_ is at least 2 * aic_count_ before the loop that calls core_trackers_[t].set_cluster(), and return false early if the handshake-discovered topology is incomplete. This should be placed near the existing cluster_count/max_clusters_per_thread setup in scheduler_context.h.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h-1120-1172 (1)
1120-1172: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGuard async forwarding when the mailbox is unavailable.
Line 1120 permits
mailbox == nullptr, but Line 1156 and Line 1171 dereference it if a task registered deferred completions. Latch a scheduler error or fail the task before forwarding when the mailbox is missing.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h` around lines 1120 - 1172, The deferred-completion path in the scheduler context can still reach mailbox forwarding even when AICoreCompletionMailbox is null, which risks a null dereference. In the completion handling around the deferred_slab processing and the mixed_complete handoff, add a guard on mailbox availability and, if it is missing, set the appropriate sched_error_code and complete/abort the task before any try_push_condition or try_push_normal_done calls. Use the existing symbols AICoreCompletionMailbox, slot_state.any_subtask_deferred, sched_->sm_header->sched_error_code, and mailbox forwarding calls to localize the fix.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h-386-403 (1)
386-403: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winDo not skip retirement when the advance lock is busy.
Line 676 returns if another completer owns
advance_lock; if that owner loaded an older watermark at Line 388, the newer watermark from this completion may never retire slots until a future completion happens. Retry after the lock is released or add a dirty/needs-advance flag.Also applies to: 674-680
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h` around lines 386 - 403, The retirement path in advance_ring_pointers() can be skipped when advance_lock is already held, which leaves newer completed_watermark values unapplied until some later completion happens. Update the completer logic around advance_lock so a failed try-lock does not just return; instead retry after the lock becomes available or set a dirty/needs-advance flag that guarantees advance_ring_pointers() runs again once the lock is released. Use advance_ring_pointers, completed_watermark, and advance_lock to locate the affected completer code.src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h-132-145 (1)
132-145: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winInitialize per-core contexts for
active_sched_threads_.Line 132 skips this setup when
sched_thread_num_ == 0, butassign_cores_to_threads()falls back to one active scheduler thread. That leaves AIV1sub_block_idas zero in single-thread mode.Proposed fix
- for (int32_t t = 0; t < sched_thread_num_; t++) + for (int32_t t = 0; t < active_sched_threads_; t++)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h` around lines 132 - 145, The per-core sub-block initialization currently loops over sched_thread_num_, so it is skipped when no scheduler threads are configured even though assign_cores_to_threads() still creates one active scheduler thread. Update the initialization in scheduler_context setup to iterate over active_sched_threads_ (or the actual assigned scheduler thread count) and keep using CoreTracker and payload_per_core_ so AIV1 gets the correct sub_block_id in single-thread mode as well.
🟡 Minor comments (4)
src/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.md-51-51 (1)
51-51: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winReplace
MAX_IDLE_ITERATIONSwording with current scheduler threshold names.The current scheduler constants shown for this codepath are
STALL_LOG_INTERVAL/FATAL_ERROR_CHECK_INTERVAL(and note removal of the old max-idle symbol), so this line documents a non-current trigger name.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.md` at line 51, Update the profiling_levels.md wording to use the current scheduler threshold symbols instead of the removed MAX_IDLE_ITERATIONS reference. In the section describing stall detection and dump, replace the old trigger name with the scheduler’s current constants, STALL_LOG_INTERVAL and FATAL_ERROR_CHECK_INTERVAL, and keep the wording aligned with the codepath that uses these thresholds.src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md-541-541 (1)
541-541: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFix SchedulerContext header path in doc text.
Both lines now point to
runtime/scheduler_context.h, but the provided runtime file location isruntime/scheduler/scheduler_context.h; this will mislead readers navigating code locations.Also applies to: 555-555
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md` at line 541, The doc text points to the wrong SchedulerContext header path; update the references in RUNTIME_LOGIC.md to use the actual runtime/scheduler/scheduler_context.h location instead of runtime/scheduler_context.h. Make sure both mentions in the SchedulerContext/AicpuExecutor description are corrected so readers can navigate to the right header.src/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.md-454-454 (1)
454-454: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winCorrect the Scheduler profiling code-location path.
runtime/scheduler_context.happears inconsistent with the provided file locationruntime/scheduler/scheduler_context.h.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.md` at line 454, The Scheduler profiling code-location reference is inconsistent and should point to the actual scheduler header path used by the runtime. Update the profiling documentation entry in profiling_levels.md to reference scheduler_context.h under the scheduler subdirectory, and make sure the path matches the corresponding scheduler_context symbol/location used elsewhere in the runtime docs.src/common/platform/onboard/aicpu/platform_aicpu_affinity.cpp-57-57 (1)
57-57: 📐 Maintainability & Code Quality | 🟡 MinorFormat
src/common/platform/onboard/aicpu/platform_aicpu_affinity.cpp— line 57 still useswhile(spacing; run clang-format on this file.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/common/platform/onboard/aicpu/platform_aicpu_affinity.cpp` at line 57, The formatting in platform_aicpu_affinity.cpp still violates the project style because the busy-wait uses while( without the expected spacing; run clang-format on the file and ensure the loop around g_cpumask in the affinity logic is reformatted consistently with the rest of the codebase.Source: Coding guidelines
🧹 Nitpick comments (2)
src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h (1)
85-88: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAvoid adding public sizing override knobs without a measured need.
These public
Runtimefields create a new configurable surface. Please either keep them as internal/effective sizing state or document the concrete measured workload that requires user-facing overrides. Based on learnings, “Any new configurable surface must be backed by a concrete, measured workload that demonstrates the need.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h` around lines 85 - 88, The new public sizing override fields on Runtime add user-facing configuration without a demonstrated need. Either remove these overrides from the public Runtime interface in runtime.h and keep sizing as internal/effective state, or add clear documentation and only expose them if there is a concrete measured workload that justifies the knobs; update the Runtime definition and any related initialization code accordingly.Source: Learnings
src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.h (1)
50-138: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftPlease stop mirroring these async helper bodies across headers.
defer_flush_range,get_async_ctx,async_ctx_is_deferred,send_notification, and the token-save path are duplicated across multiple runtime headers in the supplied graph context. That makes contract fixes like thedefer_error()change easy to miss in one copy. A single internal helper header would be much safer.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.h` around lines 50 - 138, The async helper bodies are duplicated across runtime headers, so contract changes can drift between copies. Refactor `defer_flush_range`, `defer_flush`, `get_async_ctx`, `async_ctx_is_deferred`, `register_completion_condition`, `send_notification`, and `save_expected_notification_counter` into a single internal helper header and make the existing headers include it instead of mirroring the implementations. Keep the public function names and behavior unchanged, but ensure all async runtime paths use the shared definitions so fixes like `defer_error()` land everywhere consistently.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h`:
- Around line 595-610: Cross-ring producer reclamation is currently not
protected because the fanin update loop in pto_orchestrator.h skips
last_consumer_local_id for producers whose ring_id differs from self_ring,
leaving those slots eligible for early reuse. Update the scheduling/reclamation
logic around the fanin_builder.slots loop and PTO2TaskSlotState so cross-ring
consumers add an explicit lifetime pin or refcount that keeps the producer slot
and packed buffer alive until the consumer finishes, rather than relying only on
same-ring last_consumer_local_id updates.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_shared_memory.h`:
- Around line 253-263: The per-ring layout calculation in
setup_pointers_per_ring is missing the completion_flags block, so later rings
end up with incorrect offsets. Update the offset walk that assigns each ring’s
task_descriptors_offset to also advance past completion_flags after
task_slot_states, matching the layout used by setup_pointers_per_ring. Apply the
same fix anywhere the ring layout is recomputed so every ring’s offsets stay
aligned and ring_id > 0 no longer points into the previous ring’s flags.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h`:
- Around line 416-460: The pending FIFO in PendingState can overwrite unread
entries because pending_push_back writes unconditionally and the buffer is only
sized to a single PTO2_TASK_WINDOW_SIZE even though pending work may come from
multiple rings. Update the PendingState/PendingState::pending_push_back path to
either allocate the FIFO for the aggregate pending window or reject/backpressure
pushes when pending_count() reaches pending_cap, and make sure the
initialization/allocation site uses the same capacity logic so pending_tail_idx
never wraps onto live entries.
- Around line 113-143: The ready-queue enqueue path is inconsistent:
`PTO2ReadyQueue::push_batch` can spin indefinitely when the bounded queue is
full, while the `push()` call sites around the scheduler enqueue path ignore
failure and may silently drop tasks. Update `push_batch`, `push()`, and the
scheduler enqueue logic to use one backpressure strategy consistently, either
returning a failure/overflow signal or blocking in a controlled way, and make
the callers handle that result instead of discarding it.
In `@src/common/platform/onboard/aicpu/platform_aicpu_affinity.cpp`:
- Around line 26-67: The arrival synchronization in platform_aicpu_affinity_gate
is incorrectly using g_cpumask as the barrier, which can deadlock when multiple
threads land on the same CPU and can also assign duplicate tl_filter_exec_idx
values. Replace the popcount-based wait with a separate per-launch arrival
counter in platform_aicpu_affinity_gate, and classify threads using per-thread
CPU slots or a first-wins mechanism so each survivor gets a unique exec
index/role. Keep tl_filter_exec_idx as the downstream slot contract used by
platform_aicpu_affinity_thread_idx and reset the per-launch state only after all
threads have consumed their result.
---
Major comments:
In `@src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp`:
- Around line 547-548: The teardown path in AICPUExecutor currently lets every
thread that sees finished_ true call deinit(), which can race and reset shared
state multiple times. Update the cleanup logic in aicpu_executor.cpp around the
finished_ check in AICPUExecutor::deinit/runtime teardown so only the actual
last-finishing thread performs deinit exactly once, using a one-shot
exchange/counter or equivalent guard tied to the finished_ transition.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp`:
- Around line 239-277: Clamp the PTO2_RING_HEAP override before it is stored on
runtime and before computing total_heap_size in runtime_maker.cpp, since
parse_env_uint64() can accept values that later overflow the eff_heap_size *
PTO2_MAX_RING_DEPTH multiplication. Add a bounds check alongside the existing
PTO2_RING_DEP_POOL INT32_MAX guard, reject oversized heap values with LOG_ERROR,
and make sure runtime_init_data_from_layout and related sizing logic only see
validated heap sizes so the runtime cannot under-allocate the GM arena.
- Around line 211-217: The tensor copy-back path in runtime_maker.cpp is losing
direction metadata when `runtime->tensor_pairs_` is populated in the tensor
setup loop, so the later copy-back logic can mistakenly treat the first recorded
pair as the graph output even if it is an input. Update `TensorPair` (or add
separate output-index tracking) in the code around
`runtime->tensor_pairs_.push_back` and the copy-back block that uses
`graph_output_ptr` so only true output tensors are copied back and the caller’s
input storage is never overwritten.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.h`:
- Around line 45-48: The async error path in defer_error() only sets
completion_error_code, so the scheduler’s count-first flow in
complete_slot_task() can skip the failure when deferred_slab->count is zero.
Update defer_error() to publish the error using the same non-empty slab contract
the scheduler already checks, so callers like register_pto_async_event() still
surface registration failures. Keep the fix localized around AsyncCtx,
deferred_slab, and the existing completion bookkeeping so complete_slot_task()
sees the error without relying on error_code alone.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h`:
- Around line 356-380: The inline-completion path in the PTO orchestrator
currently updates only completed_watermark, but allocator reuse in
PTO2TaskAllocator::alloc still depends on fc.last_task_alive, so alloc-only
bursts can stall even when every slot is done. In the inline-completed branch
inside the completion handling logic, after advancing the per-ring
completed_watermark, also invoke the same guarded advance_ring_pointers() path
used by normal scheduler completions so allocator-visible reclamation progresses
for inline allocations.
- Around line 507-520: The scope tracking overflow path in scope_tasks_push is
only latching a fatal error, but prepare_task still continues and returns
success, so task initialization can proceed after the buffer is saturated.
Change scope_tasks_push in PTO2OrchestratorState flow to return bool, have it
report failure when orch->scope_tasks_size reaches orch->scope_tasks_capacity,
and update prepare_task to check that return value and stop early by returning
false if pushing the task slot state fails.
- Around line 216-235: The scope stack bounds checks in begin_scope() and
end_scope() are only enforced by assert(), so release builds can still overflow
or underflow into scope_begins access and stack mutation. Update the
PTO2Orchestrator scope handling to use report_fatal(PTO2_ERROR_INVALID_ARGS,
__FUNCTION__, ...) for both invalid stack states, and return immediately before
incrementing scope_stack_top or reading scope_begins. Keep the existing
in_manual_scope() and manual_begin_depth logic intact, but make the failure path
fatal in all builds.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h`:
- Around line 236-242: The deadlock reporting path in report_deadlock() is
overwriting orch_error_code directly with store(), which can replace the
original fatal error after unwinding starts. Update report_deadlock() to use the
same latching behavior as orch_mark_fatal() and the new latch_pool_error()
helper so allocator deadlock errors only set the code once and preserve the
first/root failure.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h`:
- Around line 107-115: `runtime_reserve_layout` currently stores
`task_window_size` into `PTO2RuntimeArenaLayout` and then uses it for mask-based
slot indexing, so validate the input before any cast or layout use. In
`runtime_reserve_layout`, reject zero, non-power-of-two values, and any value
above `INT32_MAX` for `task_window_size`, and ensure the checks happen before
filling `task_window_sizes` or assigning to the layout so downstream `local_id &
(task_window_size - 1)` logic remains safe.
- Around line 126-153: Initialize the shared-memory handle before exposing it
through runtime_wire_arena_pointers: runtime_init_data_from_layout currently
zeros the PTO2SharedMemoryHandle at layout.off_sm_handle but never sets its
fields, so populate that handle with sm_dev_base and the corresponding
size/header information after the memset and before returning. Use the existing
symbols runtime_init_data_from_layout, runtime_wire_arena_pointers, and
PTO2SharedMemoryHandle to locate the fix, and ensure rt->sm_handle points to a
fully initialized handle rather than a zeroed placeholder.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_shared_memory.h`:
- Around line 184-188: Validate task_window_size in pto_shared_memory::init
before any mask is derived or used, since task_window_mask = task_window_size -
1 only works for non-zero powers of two. Add an early guard in init to reject
zero or non-power-of-two values, and ensure the get_*_by_task_id() helpers
continue to rely on a valid task_window_mask so task IDs do not alias to
incorrect slots.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h`:
- Around line 352-365: `reserve_layout` in `PTO2TensorMapLayout` only checks
that `new_num_buckets` is a power of two, but it still allows zero and does not
validate `new_task_window_sizes`, which later breaks `__builtin_ctz` and the
mask-based slot logic used by `get_task_local_id_slot` and task-entry
cleanup/linking. Tighten the layout contract in `reserve_layout` by asserting
`new_num_buckets > 0` and that every `new_task_window_sizes[r]` is also a
positive power of two before reserving any buffers, keeping the invariants
enforced at the point where the layout is constructed.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h`:
- Around line 386-403: The retirement path in advance_ring_pointers() can be
skipped when advance_lock is already held, which leaves newer
completed_watermark values unapplied until some later completion happens. Update
the completer logic around advance_lock so a failed try-lock does not just
return; instead retry after the lock becomes available or set a
dirty/needs-advance flag that guarantees advance_ring_pointers() runs again once
the lock is released. Use advance_ring_pointers, completed_watermark, and
advance_lock to locate the affected completer code.
In
`@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h`:
- Around line 271-275: The scheduler logic in AsyncWaitList/AsyncPollResult
currently reads async_wait_list.count outside the AsyncWaitList::busy lock,
which can race with poll_and_complete() mutations. Update the check in
scheduler_context so it does not touch count directly; either always invoke
poll_and_complete() whenever rt_->aicore_mailbox exists, or add a locked/atomic
“has waits” helper on AsyncWaitList and use that instead.
- Around line 616-652: In assign_cores_to_threads(), add a topology validation
step before using aic_worker_ids_ and aiv_worker_ids_ so the discovered cluster
shape is confirmed to be 1 AIC plus 2 AIV per cluster. Specifically, check that
aiv_count_ is at least 2 * aic_count_ before the loop that calls
core_trackers_[t].set_cluster(), and return false early if the
handshake-discovered topology is incomplete. This should be placed near the
existing cluster_count/max_clusters_per_thread setup in scheduler_context.h.
- Around line 1120-1172: The deferred-completion path in the scheduler context
can still reach mailbox forwarding even when AICoreCompletionMailbox is null,
which risks a null dereference. In the completion handling around the
deferred_slab processing and the mixed_complete handoff, add a guard on mailbox
availability and, if it is missing, set the appropriate sched_error_code and
complete/abort the task before any try_push_condition or try_push_normal_done
calls. Use the existing symbols AICoreCompletionMailbox,
slot_state.any_subtask_deferred, sched_->sm_header->sched_error_code, and
mailbox forwarding calls to localize the fix.
- Around line 132-145: The per-core sub-block initialization currently loops
over sched_thread_num_, so it is skipped when no scheduler threads are
configured even though assign_cores_to_threads() still creates one active
scheduler thread. Update the initialization in scheduler_context setup to
iterate over active_sched_threads_ (or the actual assigned scheduler thread
count) and keep using CoreTracker and payload_per_core_ so AIV1 gets the correct
sub_block_id in single-thread mode as well.
In
`@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h`:
- Around line 322-327: The idle-core selection in get_idle_core_offset_states
currently masks pending_occupied_ for AIC only, but not for AIV, which can
expose cores that already have staged work. Update the PTO2ResourceShape::AIV
branch in get_idle_core_offset_states so it also excludes pending_occupied_
using the same pending-state mask logic as the AIC path, while keeping the
MIX/cluster-level path unchanged.
---
Minor comments:
In `@src/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.md`:
- Line 51: Update the profiling_levels.md wording to use the current scheduler
threshold symbols instead of the removed MAX_IDLE_ITERATIONS reference. In the
section describing stall detection and dump, replace the old trigger name with
the scheduler’s current constants, STALL_LOG_INTERVAL and
FATAL_ERROR_CHECK_INTERVAL, and keep the wording aligned with the codepath that
uses these thresholds.
- Line 454: The Scheduler profiling code-location reference is inconsistent and
should point to the actual scheduler header path used by the runtime. Update the
profiling documentation entry in profiling_levels.md to reference
scheduler_context.h under the scheduler subdirectory, and make sure the path
matches the corresponding scheduler_context symbol/location used elsewhere in
the runtime docs.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md`:
- Line 541: The doc text points to the wrong SchedulerContext header path;
update the references in RUNTIME_LOGIC.md to use the actual
runtime/scheduler/scheduler_context.h location instead of
runtime/scheduler_context.h. Make sure both mentions in the
SchedulerContext/AicpuExecutor description are corrected so readers can navigate
to the right header.
In `@src/common/platform/onboard/aicpu/platform_aicpu_affinity.cpp`:
- Line 57: The formatting in platform_aicpu_affinity.cpp still violates the
project style because the busy-wait uses while( without the expected spacing;
run clang-format on the file and ensure the loop around g_cpumask in the
affinity logic is reformatted consistently with the rest of the codebase.
---
Nitpick comments:
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.h`:
- Around line 50-138: The async helper bodies are duplicated across runtime
headers, so contract changes can drift between copies. Refactor
`defer_flush_range`, `defer_flush`, `get_async_ctx`, `async_ctx_is_deferred`,
`register_completion_condition`, `send_notification`, and
`save_expected_notification_counter` into a single internal helper header and
make the existing headers include it instead of mirroring the implementations.
Keep the public function names and behavior unchanged, but ensure all async
runtime paths use the shared definitions so fixes like `defer_error()` land
everywhere consistently.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h`:
- Around line 85-88: The new public sizing override fields on Runtime add
user-facing configuration without a demonstrated need. Either remove these
overrides from the public Runtime interface in runtime.h and keep sizing as
internal/effective state, or add clear documentation and only expose them if
there is a concrete measured workload that justifies the knobs; update the
Runtime definition and any related initialization code accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 6b696aae-e13f-45a8-ab9b-e4ccfdbf9955
📒 Files selected for processing (44)
src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/common/intrinsic.hsrc/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.mdsrc/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.mdsrc/a2a3/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.mdsrc/a2a3/runtime/tensormap_and_ringbuffer/docs/device_log_profiling.mdsrc/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.mdsrc/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/orchestration/common.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.hsrc/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_orchestration_api.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox_types.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/backend/sdma/sdma_completion_kernel.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/backend/sdma/sdma_completion_scheduler.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto2_dispatch_payload.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_completion_token.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_shared_memory.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_shared_memory.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/runtime.cppsrc/common/platform/onboard/aicpu/platform_aicpu_affinity.cpp
Comment on lines
+595
to
+610
| // Push this consumer's local_id into each producer's last_consumer high- | ||
| // water-mark, replacing the per-completion fanout_refcount notification. | ||
| // Reclamation gates on the per-ring completed_watermark reaching this | ||
| // value. Only update for same-ring fanin: cross-ring consumers live in a | ||
| // different local_id space, so their id is meaningless to the producer's | ||
| // ring's watermark. Cross-ring producer slots reclaim on scope_end / | ||
| // ring wrap instead — acceptable since cross-ring fanin (e.g. | ||
| // alloc_tensors output) is sparse. | ||
| const uint8_t self_ring = task_id.ring(); | ||
| const int32_t self_local = static_cast<int32_t>(task_id.local()); | ||
| for (int32_t i = 0; i < fanin_builder.count; i++) | ||
| { | ||
| PTO2TaskSlotState *prod = fanin_builder.slots[i]; | ||
| if (prod->ring_id != self_ring) continue; | ||
| if (self_local > prod->last_consumer_local_id) prod->last_consumer_local_id = self_local; | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | 🏗️ Heavy lift
Do not reclaim cross-ring producers before their consumers finish.
Line 608 skips last_consumer_local_id updates for cross-ring fanin, leaving the producer seeded to “self”. The scheduler reclamation path can then retire/reuse that producer’s slot and packed buffer as soon as the producer completes, even if a consumer on another ring has not run yet. Add a cross-ring lifetime pin/refcount or otherwise block producer reclamation until cross-ring consumers complete.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h` around
lines 595 - 610, Cross-ring producer reclamation is currently not protected
because the fanin update loop in pto_orchestrator.h skips last_consumer_local_id
for producers whose ring_id differs from self_ring, leaving those slots eligible
for early reuse. Update the scheduling/reclamation logic around the
fanin_builder.slots loop and PTO2TaskSlotState so cross-ring consumers add an
explicit lifetime pin or refcount that keeps the producer slot and packed buffer
alive until the consumer finishes, rather than relying only on same-ring
last_consumer_local_id updates.
Comment on lines
+253
to
+263
| uint64_t offset = PTO2_ALIGN_UP(sizeof(PTO2SharedMemoryHeader), PTO2_ALIGN_SIZE); | ||
| for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) | ||
| { | ||
| header->rings[r].task_window_size = task_window_sizes[r]; | ||
| header->rings[r].task_window_mask = static_cast<int32_t>(task_window_sizes[r] - 1); | ||
| header->rings[r].heap_size = heap_sizes[r]; | ||
| header->rings[r].task_descriptors_offset = offset; | ||
| offset += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskDescriptor), PTO2_ALIGN_SIZE); | ||
| offset += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskPayload), PTO2_ALIGN_SIZE); | ||
| offset += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskSlotState), PTO2_ALIGN_SIZE); | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win
Include completion_flags in every per-ring layout walk.
setup_pointers_per_ring() wires a completion-flag block after each ring’s slot states, but these offset walkers skip it. For ring_id > 0, descriptor offsets point into the previous ring’s flags.
Proposed fix
offset += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskDescriptor), PTO2_ALIGN_SIZE);
offset += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskPayload), PTO2_ALIGN_SIZE);
offset += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskSlotState), PTO2_ALIGN_SIZE);
+offset += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(std::atomic<uint8_t>), PTO2_ALIGN_SIZE); p += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskDescriptor), PTO2_ALIGN_SIZE);
p += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskPayload), PTO2_ALIGN_SIZE);
p += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskSlotState), PTO2_ALIGN_SIZE);
+p += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(std::atomic<uint8_t>), PTO2_ALIGN_SIZE);Also applies to: 341-351
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_shared_memory.h` around
lines 253 - 263, The per-ring layout calculation in setup_pointers_per_ring is
missing the completion_flags block, so later rings end up with incorrect
offsets. Update the offset walk that assigns each ring’s task_descriptors_offset
to also advance past completion_flags after task_slot_states, matching the
layout used by setup_pointers_per_ring. Apply the same fix anywhere the ring
layout is recomputed so every ring’s offsets stay aligned and ring_id > 0 no
longer points into the previous ring’s flags.
Comment on lines
+113
to
143
| void push_batch(PTO2TaskSlotState **items, int count) | ||
| { | ||
| if (count == 0) return; | ||
|
|
||
| uint64_t pos; | ||
| while (true) { | ||
| while (true) | ||
| { | ||
| pos = enqueue_pos.load(std::memory_order_relaxed); | ||
| bool ready = true; | ||
| for (int i = 0; i < count; i++) { | ||
| for (int i = 0; i < count; i++) | ||
| { | ||
| PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; | ||
| int64_t seq = slot->sequence.load(std::memory_order_acquire); | ||
| int64_t diff = seq - static_cast<int64_t>(pos + i); | ||
| if (diff != 0) { | ||
| if (diff != 0) | ||
| { | ||
| ready = false; | ||
| break; | ||
| } | ||
| } | ||
| if (!ready) { | ||
| continue; | ||
| } | ||
| if (enqueue_pos.compare_exchange_weak( | ||
| pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed | ||
| )) { | ||
| break; | ||
| } | ||
| if (!ready) continue; | ||
| if (enqueue_pos.compare_exchange_weak(pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed)) break; | ||
| } | ||
|
|
||
| for (int i = 0; i < count; i++) { | ||
| for (int i = 0; i < count; i++) | ||
| { | ||
| PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; | ||
| slot->slot_state = items[i]; | ||
| slot->sequence.store(static_cast<int64_t>(pos + i + 1), std::memory_order_release); | ||
| } | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
Handle ready-queue saturation consistently.
Line 113 can spin forever waiting for a full bounded queue, while Line 451/Line 452 ignore push() failure and can silently drop a ready task. Please make the enqueue path return/handle backpressure consistently instead of mixing spin-until-space and ignored failure.
Also applies to: 448-453
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h`
around lines 113 - 143, The ready-queue enqueue path is inconsistent:
`PTO2ReadyQueue::push_batch` can spin indefinitely when the bounded queue is
full, while the `push()` call sites around the scheduler enqueue path ignore
failure and may silently drop tasks. Update `push_batch`, `push()`, and the
scheduler enqueue logic to use one backpressure strategy consistently, either
returning a failure/overflow signal or blocking in a controlled way, and make
the callers handle that result instead of discarding it.
Comment on lines
416
to
460
| // Thread 0 exclusive: circular FIFO of tasks awaiting fanin readiness. | ||
| // SPSC queue receives slot_states from the orchestrator; thread 0 drains | ||
| // them into the pending ring and polls fanin readiness. Storing the FIFO | ||
| // out of band (instead of intrusively in PTO2TaskSlotState) keeps the | ||
| // task struct free of scheduler-private state. | ||
| struct alignas(64) PendingState | ||
| { | ||
| static constexpr int BACKOFF_LIMIT = 32; | ||
|
|
||
| // --- Thread 0 exclusive: local batch buffer + backoff --- | ||
| int batch_count = 0; | ||
| int batch_index = 0; | ||
| int backoff_counter = 0; | ||
| PTO2TaskSlotState *batch[BATCH_SIZE]; | ||
| static constexpr int DRAIN_BATCH = 30; | ||
| static constexpr int POLL_MAX_PER_ITER = 128; | ||
|
|
||
| // --- Thread 0 exclusive --- | ||
| PTO2TaskSlotState **pending_buf{nullptr}; // capacity slots, arena-owned | ||
| uint32_t pending_cap{0}; | ||
| uint32_t pending_mask{0}; | ||
| uint32_t pending_head_idx{0}; // next pop | ||
| uint32_t pending_tail_idx{0}; // next push | ||
| int backoff_counter{0}; | ||
| PTO2TaskSlotState *drain_buf[DRAIN_BATCH]; | ||
|
|
||
| // --- SPSC queue: orchestrator (push) ↔ thread 0 (pop) --- | ||
| PTO2SpscQueue queue; | ||
|
|
||
| // --- Orchestrator write, thread 0 read --- | ||
| alignas(64) std::atomic<bool> orch_needs_drain{false}; | ||
| } wiring; | ||
|
|
||
| static_assert( | ||
| offsetof(WiringState, queue) == 256, "WiringState: batch region must be exactly 4 cache lines before queue" | ||
| ); | ||
| static_assert(sizeof(WiringState) == 640, "WiringState must be exactly 10 cache lines (640B)"); | ||
| uint32_t pending_count() const { return pending_tail_idx - pending_head_idx; } | ||
| bool pending_empty() const { return pending_tail_idx == pending_head_idx; } | ||
| } wiring; | ||
|
|
||
| alignas(64) AsyncWaitList async_wait_list; | ||
|
|
||
| // Statistics (cold path, isolated from hot-path fields) | ||
| #if PTO2_SCHED_PROFILING | ||
| alignas(64) std::atomic<int64_t> tasks_completed; | ||
| std::atomic<int64_t> tasks_consumed; | ||
| #endif | ||
| // ========================================================================= | ||
| // Inline hot-path methods | ||
| // ========================================================================= | ||
|
|
||
| /** | ||
| * Drain wiring queue: pop submitted tasks and wire their fanout edges. | ||
| * Called by scheduler thread 0 each loop iteration. Sets fanin_count, | ||
| * acquires fanout_lock per producer, allocates dep_pool entries, and | ||
| * pushes ready tasks to the appropriate ready queue. | ||
| * | ||
| * @return Number of tasks wired this call. | ||
| */ | ||
|
|
||
| int drain_wiring_queue(bool force_drain = false) { | ||
| int wired = 0; | ||
|
|
||
| // Refill local batch buffer when exhausted. | ||
| if (wiring.batch_index >= wiring.batch_count) { | ||
| // Backoff: defer pop when queue holds fewer than a full batch, | ||
| // unless force_drain, orch_needs_drain, or backoff limit reached. | ||
| if (!force_drain && wiring.queue.size() < WiringState::BATCH_SIZE) { | ||
| if (!wiring.orch_needs_drain.load(std::memory_order_acquire) && | ||
| wiring.backoff_counter < WiringState::BACKOFF_LIMIT) { | ||
| wiring.backoff_counter++; | ||
| return 0; | ||
| } | ||
| } | ||
| wiring.backoff_counter = 0; | ||
| wiring.batch_count = wiring.queue.pop_batch(wiring.batch, WiringState::BATCH_SIZE); | ||
| wiring.batch_index = 0; | ||
| if (wiring.batch_count == 0) return 0; | ||
| } | ||
|
|
||
| // Process tasks from local buffer in strict FIFO order. | ||
| while (wiring.batch_index < wiring.batch_count) { | ||
| PTO2TaskSlotState *ws = wiring.batch[wiring.batch_index]; | ||
| int ring_id = ws->ring_id; | ||
| auto &rss = ring_sched_states[ring_id]; | ||
| int32_t wfanin = ws->payload->fanin_actual_count; | ||
|
|
||
| if (wfanin > 0 && rss.dep_pool.available() < wfanin) { | ||
| rss.dep_pool.reclaim(*rss.ring, rss.last_task_alive); | ||
| if (rss.dep_pool.available() < wfanin) { | ||
| #if PTO2_PROFILING | ||
| if (is_scope_stats_enabled()) { | ||
| rss.publish_dep_pool_snapshot(); | ||
| } | ||
| #endif | ||
| break; // not enough dep_pool space — keep remainder for next call | ||
| } | ||
| } | ||
|
|
||
| wiring.batch_index++; | ||
| wire_task(rss, ws, wfanin); | ||
| wired++; | ||
| } | ||
|
|
||
| return wired; | ||
| } | ||
|
|
||
| // Route a ready slot to the right global queue. Dummy tasks (empty | ||
| // active_mask) live in dummy_ready_queue; everything else goes to the | ||
| // per-shape ready_queues[]. Used by paths that do not have a thread-local | ||
| // ready buffer (e.g. wiring). See push_ready_routed_local for the | ||
| // dispatch-time fast path. | ||
| void push_ready_routed(PTO2TaskSlotState *slot_state) { | ||
| void push_ready_routed(PTO2TaskSlotState *slot_state) | ||
| { | ||
| PTO2ResourceShape shape = slot_state->active_mask.to_shape(); | ||
| if (shape == PTO2ResourceShape::DUMMY) { | ||
| dummy_ready_queue.push(slot_state); | ||
| } else { | ||
| ready_queues[static_cast<int32_t>(shape)].push(slot_state); | ||
| } | ||
| if (shape == PTO2ResourceShape::DUMMY) dummy_ready_queue.push(slot_state); | ||
| else ready_queues[static_cast<int32_t>(shape)].push(slot_state); | ||
| } | ||
|
|
||
| /** | ||
| * Wire fanout edges for a single task. Sets fanin_count, acquires each | ||
| * producer's fanout_lock, allocates dep_pool entries for live producers, | ||
| * pushes the task to the ready queue once its fanin refcount is satisfied. | ||
| */ | ||
| void wire_task(RingSchedState &rss, PTO2TaskSlotState *ws, int32_t wfanin) { | ||
| PTO2TaskPayload *wp = ws->payload; | ||
| ws->fanin_count = wfanin + 1; | ||
|
|
||
| if (wfanin != 0) { | ||
| int32_t early_finished = 0; | ||
| for_each_fanin_slot_state(*wp, [&](PTO2TaskSlotState *producer) { | ||
| producer->lock_fanout(); | ||
| int32_t pstate = producer->task_state.load(std::memory_order_acquire); | ||
| if (pstate >= PTO2_TASK_COMPLETED) { | ||
| early_finished++; | ||
| } else { | ||
| producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, ws); | ||
| } | ||
| producer->unlock_fanout(); | ||
| }); | ||
|
|
||
| // Seed dispatch_fanin with producers already complete at wiring | ||
| // time (e.g. buffer-creator tasks that finished before this | ||
| // consumer entered the graph). Such producers never dispatch at | ||
| // runtime, so they can never bump dispatch_fanin via the fanout | ||
| // walk; without this seed the candidate compare | ||
| // (dispatch_fanin == fanin_actual_count) would be unreachable | ||
| // whenever any producer is pre-completed. Mirrors the | ||
| // early_finished seed that ready_fanin gets via init_rc. | ||
| if (early_finished != 0) { | ||
| wp->dispatch_fanin.fetch_add(early_finished, std::memory_order_acq_rel); | ||
| } | ||
|
|
||
| int32_t init_rc = early_finished + 1; | ||
| int32_t new_rc = ws->fanin_refcount.fetch_add(init_rc, std::memory_order_acq_rel) + init_rc; | ||
| if (new_rc >= ws->fanin_count) { | ||
| push_ready_routed(ws); | ||
| } | ||
| } else { | ||
| ws->fanin_refcount.fetch_add(1, std::memory_order_acq_rel); | ||
| push_ready_routed(ws); | ||
| } | ||
|
|
||
| ws->dep_pool_mark = rss.dep_pool.top; | ||
| #if PTO2_PROFILING | ||
| if (is_scope_stats_enabled()) { | ||
| rss.publish_dep_pool_snapshot(); | ||
| } | ||
| #endif | ||
| // Append slot to the tail of the pending FIFO. | ||
| void pending_push_back(PTO2TaskSlotState *s) | ||
| { | ||
| wiring.pending_buf[wiring.pending_tail_idx & wiring.pending_mask] = s; | ||
| wiring.pending_tail_idx++; | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | 🏗️ Heavy lift
Prevent pending FIFO overwrite when aggregate pending work exceeds the ring window.
Line 456 writes without checking pending_count() < pending_cap, and Line 690 sizes the FIFO to a single PTO2_TASK_WINDOW_SIZE. With multi-ring pending tasks, pending_tail_idx can wrap onto unread entries and lose tasks. Size this for the aggregate window or fail/backpressure before writing.
Also applies to: 535-542, 685-695
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h`
around lines 416 - 460, The pending FIFO in PendingState can overwrite unread
entries because pending_push_back writes unconditionally and the buffer is only
sized to a single PTO2_TASK_WINDOW_SIZE even though pending work may come from
multiple rings. Update the PendingState/PendingState::pending_push_back path to
either allocate the FIFO for the aggregate pending window or reject/backpressure
pushes when pending_count() reaches pending_cap, and make sure the
initialization/allocation site uses the same capacity logic so pending_tail_idx
never wraps onto live entries.
Comment on lines
+26
to
+67
| static std::atomic<uint16_t> g_cpumask{0}; | ||
|
|
||
| static int32_t s_thread_cpu[MAX_GATE_THREADS]; | ||
| static bool s_thread_survive[MAX_GATE_THREADS]; | ||
|
|
||
| // Per-thread exec/slot index set by BOTH gates (legacy a2a3 + filter a5): | ||
| // -1 = dropped, otherwise this thread's index among the launched threads. | ||
| // Read via platform_aicpu_affinity_thread_idx(); used as the run-wall buffer | ||
| // slot and (filter gate) the sched/orch role id. | ||
| static thread_local int32_t tl_exec_idx = -1; | ||
|
|
||
| static inline int32_t popcount64(uint64_t v) { return __builtin_popcountll(static_cast<unsigned long long>(v)); } | ||
| // Per-thread exec/slot index published by both gates. -1 = dropped, otherwise | ||
| // the survivor's role id used downstream (e.g. run-wall buffer slot in | ||
| // kernel.cpp). Declared at file scope here so the legacy a2a3 gate (above the | ||
| // filter gate's own definition further down) can write it. | ||
| static thread_local int32_t tl_filter_exec_idx; | ||
|
|
||
| /** | ||
| * This function determines which threads to use. | ||
| * | ||
| * It tries to use all the threads in the same NUMA domain (Both A2A3 and A5) | ||
| * | ||
| * @return: true if the thread is used, false if it gets dumped | ||
| */ | ||
| bool platform_aicpu_affinity_gate(int32_t logical_count, int32_t total_launched) { | ||
| tl_exec_idx = -1; | ||
| tl_filter_exec_idx = -1; | ||
| // This should be impossible... | ||
| // Going to return false to dump all the threads. | ||
| if (logical_count >= total_launched) { | ||
| // No over-launch (unreachable for a2a3, where logical < total always): | ||
| // every thread survives. Leave tl_exec_idx = -1 so no run-wall slot is | ||
| // claimed on this degenerate path — and, crucially, do NOT touch | ||
| // s_reported here: this early return has no barrier/cleanup to reset | ||
| // it, so incrementing it would leak state into the next launch | ||
| // (s_reported is reset only by the full path's last-thread cleanup). | ||
| return true; | ||
| } | ||
|
|
||
| // Assign thread index | ||
| int32_t idx = s_reported.fetch_add(1, std::memory_order_acq_rel); | ||
|
|
||
| // Report CPU | ||
| #if defined(__aarch64__) | ||
| int32_t cpu = sched_getcpu(); | ||
| #elif defined(__x86_64__) | ||
| int32_t cpu = sched_getcpu(); | ||
| #else | ||
| int32_t cpu = -1; | ||
| #endif | ||
|
|
||
| int32_t normalized_cpu = -1; | ||
| if (cpu >= 0) { | ||
| if (cpu < 63) { | ||
| s_cpumask.fetch_or(1ULL << cpu, std::memory_order_release); | ||
| } | ||
| normalized_cpu = cpu % AICPU_CORES_PER_CHIP; | ||
| } | ||
| if (idx < MAX_GATE_THREADS) { | ||
| s_thread_cpu[idx] = normalized_cpu; | ||
| LOG_ERROR("Illegal: logical_count=%d is greater or equal then total_launched=%d", logical_count, total_launched); | ||
| return false; | ||
| } | ||
|
|
||
| // Barrier: wait until all total_launched threads have reported | ||
| while (popcount64(s_cpumask.load(std::memory_order_acquire)) < total_launched && | ||
| s_reported.load(std::memory_order_acquire) < total_launched) {} | ||
| // Get current CPU ID | ||
| int cpu = sched_getcpu(); | ||
|
|
||
| // CAS winner does cluster classification | ||
| int32_t expected = 0; | ||
| if (s_gate_init.compare_exchange_strong(expected, 1, std::memory_order_acq_rel, std::memory_order_acquire)) { | ||
| // Initialize survive flags | ||
| for (int32_t i = 0; i < total_launched; ++i) { | ||
| s_thread_survive[i] = false; | ||
| } | ||
| // At to cpumask | ||
| g_cpumask.fetch_or(1 << cpu, std::memory_order_relaxed); | ||
|
|
||
| struct ClusterInfo { | ||
| int32_t count{0}; | ||
| int32_t tids[MAX_GATE_THREADS]; | ||
| }; | ||
| ClusterInfo clusters[MAX_CLUSTERS]; | ||
| // Barrier wait until all the spawned threads are here before choosing which ones will be used. | ||
| while(__builtin_popcount(g_cpumask) < total_launched) {} | ||
|
|
||
| for (int32_t tid = 0; tid < total_launched; ++tid) { | ||
| int32_t c = s_thread_cpu[tid]; | ||
| if (c < 0) continue; | ||
| int32_t cluster_id = c / CPUS_PER_CLUSTER; | ||
| if (cluster_id < 0 || cluster_id >= MAX_CLUSTERS) continue; | ||
| ClusterInfo &info = clusters[cluster_id]; | ||
| if (info.count < MAX_GATE_THREADS) info.tids[info.count++] = tid; | ||
| } | ||
| // Choose the thread based on reverse bit order (highest cpu id to lowest) | ||
| // This assures that all the threads lie in the same NUMA domain | ||
| int how_many_on_top = __builtin_popcount(g_cpumask >> cpu); | ||
| bool will_be_used = how_many_on_top <= logical_count ? true : false; | ||
| // Publish per-survivor exec idx (lowest surviving cpu → 0, highest → | ||
| // logical_count - 1). Read by kernel.cpp via | ||
| // platform_aicpu_affinity_thread_idx() to claim the run-wall buffer slot | ||
| // so device timing is captured. | ||
| tl_filter_exec_idx = will_be_used ? (logical_count - how_many_on_top) : -1; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
Do not use the CPU bitmask as the arrival barrier.
g_cpumask collapses duplicate sched_getcpu() landings, so popcount(g_cpumask) can remain below total_launched forever. Even if the hang is avoided, duplicate threads on a selected CPU can compute the same tl_filter_exec_idx, which downstream uses as the executor role/run-wall slot. Use a separate arrival counter plus per-thread CPU slots/first-wins classification, and reset per-launch state after all threads consume their result. The supplied downstream/header snippets show the exec index is a slot/role contract.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/common/platform/onboard/aicpu/platform_aicpu_affinity.cpp` around lines
26 - 67, The arrival synchronization in platform_aicpu_affinity_gate is
incorrectly using g_cpumask as the barrier, which can deadlock when multiple
threads land on the same CPU and can also assign duplicate tl_filter_exec_idx
values. Replace the popcount-based wait with a separate per-launch arrival
counter in platform_aicpu_affinity_gate, and classify threads using per-thread
CPU slots or a first-wins mechanism so each survivor gets a unique exec
index/role. Keep tl_filter_exec_idx as the downstream slot contract used by
platform_aicpu_affinity_thread_idx and reset the per-launch state only after all
threads have consumed their result.
SergioMartin86
changed the title
SergioMartin86
force-pushed
the
polling-pr-minimal
branch
from
3a1bc17 to
1ac2de8
Compare
SergioMartin86
changed the title
SergioMartin86
changed the title
SergioMartin86
changed the title
Consolidates the polling-redesign series (originally 6 commits on polling-pr-minimal at base fcc33bc) onto current upstream/main as a single squash. Resolves 24 file conflicts using the recipe in RECONCILIATION_NOTES.md: - Polling-redesigned files (pto_scheduler.h, ring_buffer, shared_memory, runtime, tensormap, runtime2_types, orchestrator.cpp, scheduler_*.cpp): take theirs (polling). - Upstream-evolved files with feature additions (aicpu_executor.cpp, runtime_maker.cpp, orchestration/, platform_aicpu_affinity.cpp): take ours (upstream/main). Cross-cutting adapter overloads added so upstream call sites compile against polling structs: - PTO2RuntimeArenaLayout: task_window_sizes[]/heap_sizes[]/dep_pool_capacities[] per-ring arrays (was single scalars). - runtime_reserve_layout: per-ring overload + single-size broadcast adapter. - runtime_init_data_from_layout: heap_sizes[] per-ring overload + scalar adapter. - runtime_destroy(rt, arena): 2-arg overload forwarding to single-arg. - PTO2SharedMemoryHandle::init_per_ring: forwards to init_header_per_ring. - PTO2OrchestratorState::l2_swimlane_level: L2SwimlaneLevel field for upstream aicpu_executor's orchestrator-to-scheduler bridge. - SchedulerContext::on_orchestration_done: thread_idx overload. Renames applied (commit 10a7680's surface): - bare `Arg` -> `L0TaskArgs` in pto_runtime2.h, pto_orchestrator.h, pto_dep_compute.h, pto_runtime2_types.h. - `TensorRef::ptr` -> `.ref()` (returns reference, callers take address). - `.create_info` -> `.create_info()` method call. dep_gen_aicpu_record_submit signature: pass args.launch_spec.block_num() (matches upstream a5's call). The needs_copy_back D2H optimization is already in upstream/main, so the polling-pr-minimal commit 3a1bc17 that restored it on the old base is unnecessary here and dropped. Compile-clean after this commit. Three further bug fixes follow.
The polling-side Runtime constructor initializes its own fields but not the upstream-added aicpu_allowed_cpus[] / aicpu_allowed_cpu_count / aicpu_launch_count. Host code populates them later from prepare_runtime_for_launch, so this is defensive — but garbage left in them between construction and host population would be visible to the AICPU side if the host path were ever skipped. Zero them in the same memset/clear block as the other fields.
Upstream aicpu_executor.cpp calls sched_ctx_.on_orchestration_done(runtime, rt, thread_idx, total_tasks) but my adapter received the args in reverse order. The result was total_tasks_ = thread_idx (0/1/2 instead of the real task count), so the scheduler thought it was done before the orchestrator finished — and the test hung in 507018 territory regardless. Fix puts thread_idx and total_tasks in the same positions as upstream. Still hangs after this fix — runtime hangs earlier than on_orchestration_done. No LOG_INFO_V0 output from polling kernel at all (even with --log-level v9). Working theory: macro wiring drift between polling-side scheduler_context.h and upstream's unified_log + orchestration_api log entry points. To diagnose further, would need to verify which log_info_v fn the macro resolves to in the built libaicpu_kernel.so.
setup_pointers(task_window_sizes[0]) broadcasts the first ring size to all rings, which is fine only when all rings have the same window. Use the canonical per-ring setup_pointers_per_ring to handle the general case. (Test workload happens to use uniform 16384 across all 4 rings, so this fix is correctness-improving but not the cause of the current hang.) Hang location now identified: orchestrator thread 3 hangs inside the loaded orch SO's (*p_func)(orch_args_cached_) call — i.e. inside the user-graph submit loop — between aicpu_executor.cpp's printed Ring sizes (line 487) and Orchestrator completed (line 685). Most likely candidates: prepare_task allocator wait, tensormap insert, or submit_task_common's last_consumer_local_id update path. Next session should add LOG_INFO_V0 inside submit_task_common to bracket which call hangs.
ROOT CAUSE of the runtime hang at vnl-main reconciliation. The squash-merge took polling-pr-minimal's version of pto_runtime2.h (had only `log_info_v` in PTO2RuntimeOps) but kept upstream's pto_orchestration_api.h (declares `log_error`, `log_warn`, `log_debug`, `log_info_v` in that order). When the orchestration .so called rt->ops->log_info_v(...) from inside the dlopen'd user-graph SO, the compiler resolved log_info_v's offset using the orch-side layout (after 6 fn ptrs + 3 logging fn ptrs). But the runtime had initialized rt->ops as s_runtime_ops using the polling-side layout (log_info_v right after report_fatal). The orch SO followed the wrong function pointer into a get_tensor_data/set_tensor_data slot, which jumped into corruption and silently hung the entire AICPU thread. Symptom: aicpu_executor.cpp reached "DIAG pre-p_func" then (*p_func) never returned. orch_diag_step on shared memory stayed at 30. AICore stream timed out at 507018. Fix: restore log_error / log_warn / log_debug fields (and their rt_log_error / rt_log_warn / rt_log_debug dispatcher implementations, populated in s_runtime_ops) before log_info_v. ABI now matches between both sides. Result: paged_attention Case1 PASSED on a 5-round run, dev 7. Avg Host 208 ms, Avg Device 31 ms. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The wake-list-only redesign made classify_fanin_state's decision terminal — each task is either routed to a ready queue (all fanins met) or registered on a producer's wake_list (first unmet), with no "leave it for next iter" outcome. With DRAIN_BATCH (30) below the old POLL_MAX_PER_ITER cap (128), the intermediate FIFO emptied within the same iter that filled it, paying push_back + pop_front overhead per task for no carry-over benefit. drain_wiring_queue now classifies + routes each drained task in-line. Removes pending_buf/cap/mask/head_idx/tail_idx state, pending_push_back /pending_pop_front/pending_count/pending_empty helpers, off_pending_buffer + pending_capacity layout fields, POLL_MAX_PER_ITER, and the per-iter PTO2_TASK_WINDOW_SIZE pointer array arena reservation. Net diff: −81 / +36. Smoke-tested paged_attention Case1 100 rounds PASS on dev 6; targeted A/B vs HEAD shows consistent small device improvement (−0.5% on paged_attention C1, −1.5% on paged_attention_manual_scope C1, −2.6% on alternating_matmul_add C1). Host time is too noisy on this shared box to claim a host win, but is neutral-or-better across the three samples.
Two consecutive loops over fanin_builder ran back-to-back per task:
the first updated each same-ring producer's last_consumer_local_id
high-water-mark, the second copied (local_id, ring_id) into the
payload's flat arrays. Fold into one loop.
Side benefit: read ring_id from the cache-warm fanin_builder.ring_ids
SOA slice (already populated by append_fanin_or_fail) instead of
dereferencing slot_state->ring_id. Cross-ring fanin iters now skip
the slot dereference entirely; only same-ring iters touch the
producer's slot_state cache line.
A/B on dev 6, 100 rounds trimmed-80:
alternating_matmul_add C1 — Host 165.2 → 148.7 ms (-10.0%),
Device 1.43 → 1.43 ms (flat).
paged_attention C1 — Host noisy across samples,
Device 31.03 → 30.85 ms (-0.6%).
Smaller tests see most of the host benefit (per-task host overhead
dominates their wall time); large device-bound tests see negligible
delta as expected.
SergioMartin86
force-pushed
the
polling-pr-minimal
branch
from
fa60918 to
a6fb764
Compare
In on_mixed_task_complete's wake-list walk, after reading next, the code was assigning waiter->next_in_wake_list = nullptr. The store has no observable effect: register_wake() unconditionally overwrites the field on every re-registration (before the CAS that publishes the consumer onto a producer's wake list), and reset_for_reuse() clears it on slot reuse. No reader exists between this point and the next overwrite/reset. Saves one store per waiter across every producer completion. Tiny absolute win (paged_attention ~5K-10K wake-list iters/round) but removes confusing-by-omission code: a reader could conclude the nullptr clear was load-bearing for an ordering or visibility invariant when it isn't. Smoke-tested paged_attention Case1 (5 rounds) on dev 6: PASS.
SergioMartin86
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Human Summary
This PR replaces the entire wiring logic (one where tasks notify each other on completion), by a polling-based approach where the scheduler constantly asks for readiness. This is an optimization because it severely reduces contention, pointer chasing, and initial wiring effort. As a consequence an ~17% median device speedup (host timing reflects no change within noise variance) is achieved across all tests used.
Two more optimizations were done here:
AI-generated summary follows.
AI Summary
TL;DR
Rewrites the per-completion consumer-notification path in the a2a3
tensormap_and_ringbufferscheduler. Producer completion is now one bytestore to a per-ring
completion_flags[]array, and consumers poll thebytes via a thread-0-managed pending FIFO. Removes the ~195K cross-slot
fanin_refcount.fetch_addatomics that dominatepaged_attentionruns.On 18 tests against
upstream/main(18a13a26), 100 rounds each,trimmed-80, dev 6, pto-isa
e722679b:(13 / 18 directly comparable — 3 tests fail on baseline only, 2 fail
on both. See "Baseline-only failures" below.)
Headline
paged_attention:Case1: device 38.92 → 31.29 ms (−19.6%),host 209.99 → 223.51 ms (+6.4%, within trial noise on this shared
box).
What changes
Each producer task no longer walks its consumer list to RMW each
consumer's
fanin_refcounton completion. Instead, per ring:completion_flags[]byte array in shared memory. Producercompletion sets one byte (release store) — replaces the per-consumer
atomic RMW.
fanin_local_ids[]+fanin_ring_ids[]onPTO2TaskPayload.Recorded at submit time; the consumer-side pending poll indexes the
producer's ring's bytes. The parallel
ring_ids[]is required becausenested
PTO2_SCOPEblocks can land producer and consumer in differentrings.
SPSC wiring queue into a pending FIFO, then
classify_fanin_statesweeps each task once: all met → ready queue; any unmet → register on
the first unmet producer's
wake_list_headand leave the FIFO(one-shot; the producer's completion drains its wake list and rescans).
completed_watermarkper ring (atomic, CAS-advanced throughconsecutive COMPLETED slots) replaces
fanout_refcountfor slotreclamation. Slots retire when the watermark passes the producer's
last_consumer_local_id.fanin_refcount,fanout_refcount,fanout_count,fanout_head,PTO2DepListEntrypool,try_speculative_release,release_fanin_and_check_ready,PTO2_TASK_CONSUMED, andon_scope_end's producer-release loop are removed.Reconciliation with upstream/main
This PR was originally written at base
fcc33bcb. Rebasing onto currentupstream/main(32 commits later) surfaced shape changes the pollingdesign has to absorb. The recipe is documented in
RECONCILIATION_NOTES.md; the high-level shape:git rebase upstream/main. Resolution:take theirs (polling) for files the redesign owns (
pto_scheduler.h,ring_buffer, shared_memory, runtime, tensormap, runtime2_types,
orchestrator.cpp, scheduler_*.cpp); take ours (upstream/main) for
files with upstream feature additions (
aicpu_executor.cpp,runtime_maker.cpp, orchestration/,platform_aicpu_affinity.cpp).Arg → L0TaskArgs,TensorRef::ptr → .ref(),.create_info → .create_info()propagated through the polling headers.runtime_reserve_layout,runtime_init_data_from_layout,runtime_destroy(rt, arena), andPTO2SharedMemoryHandle::init_per_ringso upstream call sites compileunchanged against the polling-side structs.
explicitly is now part of upstream — no separate gate change needed.
runtime_maker.cpp'sneeds_copy_backD2H skip is already present inupstream/main, so no change is needed here.
Bugs the rebase caught (each is its own commit in this PR)
Four small follow-up commits land on top of the squash. Each fixes a real
bug we hit while bringing the design forward; each could re-bite anyone
attempting the same rebase, so they're called out by name:
0c62a5fePTO2RuntimeOpsABI alignment. The polling-sidedefinition in
pto_runtime2.hdeclared onlylog_info_vunderLogging; the upstream
pto_orchestration_api.hdeclaredlog_error,log_warn,log_debug,log_info_vin that order. The orch SOcomputed
log_info_v's offset using one layout while the runtimeinitialized
s_runtime_opsusing the other, so the call ended up ina
get_tensor_dataslot and silently hung the AICPU thread(manifests as
aclrtSynchronizeStreamWithTimeout (AICPU) failed: 507018). Restoring the three missing log fields + theirrt_log_error / rt_log_warn / rt_log_debugdispatchers brings bothlayouts into agreement. This is the single most important bug to
carry forward — anyone rebasing the polling-side
pto_runtime2.hwithout it will hit the exact same hang.
0461340fon_orchestration_doneadapter arg order. Upstreamcalls
sched_ctx_.on_orchestration_done(runtime, rt, thread_idx, total_tasks)but the polling overload had the last two parametersreversed, so the scheduler's
total_tasks_was set tothread_idx(0/1/2). Symptom:
sched_error_code=100 PTO2_ERROR_SCHEDULER_TIMEOUT.47afba14per-ringsetup_pointersininit_per_ringadapter.Original used
setup_pointers(task_window_sizes[0])which broadcaststhe first ring's size to all rings. Correct on uniform workloads but
silently wrong on per-ring-sized configurations. One-line fix: call
setup_pointers_per_ring(task_window_sizes).dd621525Runtime ctor missing zero-init for new fields.Upstream-added
aicpu_allowed_cpus[],aicpu_allowed_cpu_count,aicpu_launch_countwere not initialized in the polling-sideconstructor. Host code populates them later, so this is defensive.
Further cleanups on top of the rebase
Three small follow-up commits land on top of the four reconciliation
fixes. Each is correctness-preserving and falls out of the wake-list-only
redesign — these are architectural simplifications, not new perf
optimisations:
c5e97adeDrop redundant pending FIFO from wiring-queue drainpath. The original design polled a pending FIFO so multi-fanin
consumers that hadn't been satisfied yet could roll to the next
iter. The wake-list-only redesign made
classify_fanin_state'sdecision terminal (route to ready OR register on first unmet
producer's wake list), so the FIFO emptied every iter and paid one
push_back+pop_frontper task for no carry-over benefit.Removes the FIFO, its buffer,
POLL_MAX_PER_ITER, and thePTO2_TASK_WINDOW_SIZEpointer-array arena reservation.Net −81 / +36 lines.
a6fb7642Mergefanin_builderloops insubmit_task_common.Two consecutive loops over
fanin_builderran back-to-back per task:the first updated each same-ring producer's
last_consumer_local_idhigh-water-mark, the second copied
(local_id, ring_id)into thepayload's flat arrays. Folds into one loop; the same-ring check now
reads from the cache-warm
fanin_builder.ring_idsSOA sliceinstead of dereferencing
slot_state->ring_id. Cross-ring faniniters now skip the slot-state cache-line load entirely.
c8e802dfDrop dead defensive store in wake-list drain. Inon_mixed_task_complete, after readingnext, the code wasassigning
waiter->next_in_wake_list = nullptr. The store has noobservable effect —
register_wake()unconditionally overwritesthe field on every re-registration before the CAS that publishes
the consumer onto a producer's wake list, and
reset_for_reuse()clears it on slot reuse. No reader exists between this point and
the next overwrite/reset.
Hot-path cost eliminated
upstream/main)fanin_refcount, runtry_speculative_releaseseq_cst CAS per ready consumer, check deadlock detector, pollorch_error_codefanout_headlinked list (withfanout_lock) + tensormap entry-pool back-pressure checkfanin_local_ids[]+fanin_ring_ids[], push to wiring SPSCDiff scope
44 runtime files + 1 host-side file modified in-place. No file
deletions, no file additions, no structural moves — every change is in
an existing upstream/main file. Reviewer focus areas:
runtime/scheduler/pto_scheduler.h::on_mixed_task_complete— the4-op completion path (byte store, wake_list exchange, watermark CAS,
ring advance try-lock).
runtime/scheduler/pto_scheduler.h::drain_wiring_queue— first-unmetregistration in the pending poll.
runtime/pto_orchestrator.h::submit_task_common— same-ringlast_consumer_local_idupdate + payload populate.runtime/pto_runtime2.h::PTO2RuntimeOps— ABI-aligned logging fields(see commit
0c62a5fe).Net line count is +4500 / −10500, dominated by removing the existing
notification machinery (fanout list pool, dep_list maintenance,
spec-dispatch release, scope-end release walk).
Benchmark sweep — a2a3 onboard, 100 rounds (trimmed to 80)
Method: branch-checkout in a single docker
.venv, rebuild on eachcheckout, run sequentially on dev 6 inside container
s00831018-aicpu,matched pto-isa
e722679b(the repo-pinned commit). All 18 tests runwith
--case <Case> --manual include --rounds 100 --skip-golden.Trimmed-mean of 80 (10 low + 10 high dropped).
Baseline:
upstream/main(18a13a26).Current: this PR (
c8e802df, vnl-main HEAD).alternating_matmul_addCase1paged_attentionCase1paged_attentionCase2paged_attentionCase3paged_attention_manual_scopeCase1paged_attention_unrollCase1paged_attention_unrollCase2paged_attention_unroll_4dimsCase1paged_attention_unroll_manual_scopeCase1paged_attention_unroll_manual_scopeCase2spmd_basicCase1spmd_multiblock_aivCase1spmd_multiblock_mixCase1Baseline-only failures (3/18) — robustness signal
Three tests fail on
upstream/mainbut pass on this PR underidentical run conditions:
batch_paged_attentionCase1,benchmark_bgemmCase0,multi_round_paged_attentionCase1. All hitaclrtSynchronizeStreamWithTimeout (AICPU) failed: 507018onupstream/mainand recover; the polling design runs them tocompletion. We aren't claiming "the PR fixes these tests" — the
mechanism is environment-specific — but it's worth noting that the
polling design tolerates this environment that the upstream scheduler
does not.
Pre-existing failures (2/18, on both)
spmd_paged_attentionCase1 and Case2 fail on both branches with thesame 507018 / 507046 pattern. Pre-existing flake unrelated to this PR.
Notes on the host regressions (
paged_attention_unroll*)Two
paged_attention_unrollvariants show host regressions of +18.3%and +30.3% with no other host signal of similar magnitude:
paged_attention_unrollCase1 (+18.3%)paged_attention_unroll_4dimsCase1 (+30.3%)These are small device-bound tests (~2 ms device per round) where
host wall is dominated by Python dispatch overhead and per-task
constant scheduler costs. The same family's other variants
(
paged_attention_unrollCase2,paged_attention_unroll_manual_scopeCase1/Case2) show host improvements of −2.5%, −18.4%, −7.6%
respectively, which argues the regressions are workload-shape and
shared-box-noise rather than a fundamental host regression. The
device delta is favourable on every variant (−6.8% to −13.5%).
paged_attentionCase1 / Case3 host is essentially flat (+6.4% /+3.1%) — within trial-to-trial host variance we've measured on this
shared box (±30 ms / single trial on a 200 ms host wall).
Tests excluded (8 of 26 cases)
These fail under direct
python3 test_*.pyinvocation on bothbranches — they are pytest-only tests and the benchmark script invokes
them as scripts. Not regressions; the same tests all pass under
pytest. Pre-existing benchmark-runner limitation, unrelated to this PR.dummy_task,fanin_lookup_perf,mixed_example,paged_attention_ringbuffer,qwen3_14b_decode,scalar_data_test,spmd_paged_attention_highperf,vector_exampleTest plan
python3 examples/a2a3/tensormap_and_ringbuffer/paged_attention/test_paged_attention.py --case Case1 --rounds 100 --skip-golden --manual include --platform a2a3 --device 6→ PASSupstream/main18a13a26, pto-isae722679bwith the rebased history)
1e-3). The one failure is
paged_attentionCase3, which fails thesame way on
upstream/main(max_diff ~0.09 vs ~0.12 on this PR —both about 100× over tolerance). Pre-existing kernel coverage
gap, not a runtime issue; see "Known pre-existing issues" below
for the root cause + fix shape.
paged_attentionCase1 (the case originally flagged in thereconciliation notes) now passes with max_diff ~5×10⁻⁶, a margin of
~200× under tolerance — almost certainly fixed by upstream's
pto-isa bump to
016396b5("fixes mask-form scatter") in09f1546f.Compatibility / risk
already existed on upstream/main.
dispatcher fields added to
PTO2RuntimeOps(which were already in theorch-side declaration; this PR makes the runtime-side definition
match, see commit
0c62a5fe).Known pre-existing issues (not addressed by this PR)
paged_attentionCase3 fails the default golden tolerance(rtol/atol 1e-3) with max_diff ~0.09–0.12 on both
upstream/mainand this PR. This is not a precision issue — it's a kernel
coverage gap. Case3 uses
head_dim=256, but the AIC kernels(
aic_qk_matmul.cpp,aic_pv_matmul.cpp) only specialise forhead_dim ∈ {16, 128}. The runtime dispatch falls through to thehead_dim=128template, so the matmul processes only the first128 of 256 head_dim elements per batch×head. 95% of output elements
violate, and the worst case shows
actual ≈ 0.002vsexpected ≈ 0.14. The fix is to addqk_matmul_impl<64, 256, 64>and
pv_matmul_impl<64, 64, 256>specialisations in the AICkernels and route Case3 to them. Out of scope for this runtime PR.
Benchmarks run with
--skip-goldenpartly to avoid this unrelatednoise and partly to measure pure runtime time (no allclose loop
overhead).