Refactor: extract AICPU scheduler into SchedulerContext class (a2a3 + a5) - #654
Conversation
There was a problem hiding this comment.
Code Review
This pull request refactors the AICPU runtime by extracting scheduler-related logic and state from the AicpuExecutor class into a new SchedulerContext class, with implementation files organized in a dedicated directory for better modularity. Documentation and build scripts are also updated to reflect these changes. The review feedback identifies several improvement opportunities: the use of recursive globbing in CMake could include unwanted files, hardcoded constants in scheduler_context.h should be moved to a shared header to avoid maintenance risks, a redundant query in the drain dispatch loop is inefficient, and an incorrect (void) cast exists for a variable that is actually used in the profiling path.
- Move pto_scheduler.h/.cpp into runtime/scheduler/ subdirectory - Create scheduler_types.h with CoreExecState, CoreTracker, SlotTransition, SchedProfilingCounters, SyncStartDrainState, profiling macros and scheduler constants - Update include paths to "scheduler/pto_scheduler.h" in 3 files - Switch CMake from GLOB to GLOB_RECURSE for source collection to support subdirectories without build_config.py changes - Remove duplicate type definitions from aicpu_executor.cpp - Add missing copyright headers to onboard CMakeLists.txt files Pure restructuring — no behavioral change.
- Create scheduler_context.h with full SchedulerContext class declaration (all method signatures for completion, cold path, and dispatch) - Create scheduler_completion.cpp implementing completion and drain methods as SchedulerContext members: decide_slot_transition, complete_slot_task, check_running_cores_for_completion, enter_drain_mode, count_global_available, drain_worker_dispatch, handle_drain_mode - Add #include of scheduler_context.h to aicpu_executor.cpp AicpuExecutor retains its original methods unchanged. SchedulerContext methods exist in parallel but are not yet called. The switchover happens in Phase 5-6.
- Move runtime.cpp and pto_shared_memory.cpp into runtime/shared/ (the only two runtime sources the host target actually needs) - Update build_config.py: host source_dirs uses "runtime/shared" instead of "runtime", eliminating compilation of device-only code (pto_orchestrator, pto_scheduler, pto_ring_buffer, pto_tensormap, pto_runtime2) in the host .so - Add RUNTIME_MAX_WORKER and RUNTIME_MAX_FUNC_ID guard macros to scheduler_context.h so it compiles without pulling in runtime.h Headers remain in runtime/ — only .cpp files are split.
Create scheduler_cold_path.cpp implementing cold-path helper methods as SchedulerContext members: handle_orchestrator_exit, handle_core_transition, check_idle_fatal_error, log_stall_diagnostics, handle_timeout_exit, log_profiling_summary. Accesses shared AicpuExecutor state through pointer members (completed_tasks_ptr_, orchestrator_done_ptr_, etc.) and uses emergency_shutdown_fn_ callback for core shutdown. AicpuExecutor retains its original methods unchanged; SchedulerContext methods exist in parallel but are not yet called.
Create scheduler_dispatch.cpp implementing dispatch helpers and the main scheduler loop as SchedulerContext members: - shape_name, get_dispatch_order, pop_ready_tasks_batch - build_payload, dispatch_subtask_to_core, dispatch_mix_block_to_cluster - dispatch_block, dispatch_shape - resolve_and_dispatch (the main scheduler entry point) Uses sched_-> for PTO2SchedulerState access and payload_per_core_ for dual-buffer dispatch payloads. AicpuExecutor retains its original methods unchanged; SchedulerContext methods exist in parallel but are not yet called.
- Replace AicpuExecutor's scheduler fields (core_exec_states_, core_trackers_, drain_state_, sched_perf_, s_pto2_payload_per_core) with single SchedulerContext sched_ctx_ member - Delete all old inline methods (cold-path, completion, dispatch, resolve_and_dispatch_pto2) from AicpuExecutor — ~1300 lines removed - Wire sched_ctx_ pointer members in init() and run() to connect to AicpuExecutor's shared atomic state - Switch run() to call sched_ctx_.resolve_and_dispatch() - Add core_assignments_ pointer to SchedulerContext for profiling flush - Set sched_ctx_.sched_ after rt creation (not in init() where rt may be null in device orchestration mode) aicpu_executor.cpp shrinks from ~2870 to ~1480 lines.
- device_log_profiling.md: resolve_and_dispatch_pto2 -> SchedulerContext::resolve_and_dispatch - profiling_levels.md: update scheduler profiling file paths from aicpu_executor.cpp to scheduler/ directory, fix markdownlint issues (fenced code block languages, table alignment) - tensor-dump.md: update dump_tensors_for_task call site references to scheduler_completion.cpp and scheduler_dispatch.cpp - dynamic-linking.md: document SchedulerContext sched_ctx_ member in AicpuExecutor, update core_trackers_ reference path
Apply the same scheduler refactoring to a5 as done in a2a3: - Create runtime/scheduler/ with SchedulerContext class and extracted methods (scheduler_types.h, scheduler_context.h, scheduler_completion.cpp, scheduler_cold_path.cpp, scheduler_dispatch.cpp) - Move pto_scheduler.h/.cpp into runtime/scheduler/ - Move runtime.cpp and pto_shared_memory.cpp into runtime/shared/ - Update build_config.py: host source_dirs uses "runtime/shared" - Switch CMake from GLOB to GLOB_RECURSE in all platform CMakeLists.txt - Update include paths to "scheduler/pto_scheduler.h" - Slim aicpu_executor.cpp from ~2854 to ~1244 lines - Add missing copyright headers to onboard CMakeLists.txt files a5-specific adaptations: - RUNTIME_MAX_WORKER=108 (vs a2a3's 72) in scheduler_context.h - CoreExecState profiling layout differs (no dispatch_count field) - LocalContext uses s_block_idx/s_block_num (vs block_idx/block_num) - No perf_aicpu_switch_buffer or perf_aicpu_flush_buffers calls
- device_log_profiling.md: resolve_and_dispatch_pto2 -> SchedulerContext::resolve_and_dispatch - profiling_levels.md: update scheduler profiling file paths to a5 scheduler/ directory, fix markdownlint format issues, replace src/a2a3 references with src/a5
8ed62e3 to
41af1f0
Compare
Merge upstream PMU profiling (hw-native-sys#639) into the extracted scheduler: - scheduler_completion.cpp: add pmu_aicpu_record_task() call in complete_slot_task for per-task PMU counter sampling - scheduler_dispatch.cpp: add pmu_aicpu_init() in one-time init, PTO2_DISABLE_DUAL_ISSUE conditional for single-phase dispatch, pmu_aicpu_flush_buffers() at dispatch loop exit - scheduler_context.h: add physical_core_ids_ and cores_total_num_ pointers for PMU MMIO base resolution - aicpu_executor.cpp: wire physical_core_ids_ in init() and record physical_core_id during handshake
41af1f0 to
8758241
Compare
… a5) (hw-native-sys#654) * Refactor: create runtime/scheduler/ and extract scheduler types - Move pto_scheduler.h/.cpp into runtime/scheduler/ subdirectory - Create scheduler_types.h with CoreExecState, CoreTracker, SlotTransition, SchedProfilingCounters, SyncStartDrainState, profiling macros and scheduler constants - Update include paths to "scheduler/pto_scheduler.h" in 3 files - Switch CMake from GLOB to GLOB_RECURSE for source collection to support subdirectories without build_config.py changes - Remove duplicate type definitions from aicpu_executor.cpp - Add missing copyright headers to onboard CMakeLists.txt files Pure restructuring — no behavioral change. * Refactor: add SchedulerContext class with completion and drain methods - Create scheduler_context.h with full SchedulerContext class declaration (all method signatures for completion, cold path, and dispatch) - Create scheduler_completion.cpp implementing completion and drain methods as SchedulerContext members: decide_slot_transition, complete_slot_task, check_running_cores_for_completion, enter_drain_mode, count_global_available, drain_worker_dispatch, handle_drain_mode - Add #include of scheduler_context.h to aicpu_executor.cpp AicpuExecutor retains its original methods unchanged. SchedulerContext methods exist in parallel but are not yet called. The switchover happens in Phase 5-6. * Refactor: slim host build deps with runtime/shared/ for sources - Move runtime.cpp and pto_shared_memory.cpp into runtime/shared/ (the only two runtime sources the host target actually needs) - Update build_config.py: host source_dirs uses "runtime/shared" instead of "runtime", eliminating compilation of device-only code (pto_orchestrator, pto_scheduler, pto_ring_buffer, pto_tensormap, pto_runtime2) in the host .so - Add RUNTIME_MAX_WORKER and RUNTIME_MAX_FUNC_ID guard macros to scheduler_context.h so it compiles without pulling in runtime.h Headers remain in runtime/ — only .cpp files are split. * Refactor: add SchedulerContext cold-path methods Create scheduler_cold_path.cpp implementing cold-path helper methods as SchedulerContext members: handle_orchestrator_exit, handle_core_transition, check_idle_fatal_error, log_stall_diagnostics, handle_timeout_exit, log_profiling_summary. Accesses shared AicpuExecutor state through pointer members (completed_tasks_ptr_, orchestrator_done_ptr_, etc.) and uses emergency_shutdown_fn_ callback for core shutdown. AicpuExecutor retains its original methods unchanged; SchedulerContext methods exist in parallel but are not yet called. * Refactor: add SchedulerContext dispatch methods and main loop Create scheduler_dispatch.cpp implementing dispatch helpers and the main scheduler loop as SchedulerContext members: - shape_name, get_dispatch_order, pop_ready_tasks_batch - build_payload, dispatch_subtask_to_core, dispatch_mix_block_to_cluster - dispatch_block, dispatch_shape - resolve_and_dispatch (the main scheduler entry point) Uses sched_-> for PTO2SchedulerState access and payload_per_core_ for dual-buffer dispatch payloads. AicpuExecutor retains its original methods unchanged; SchedulerContext methods exist in parallel but are not yet called. * Refactor: wire SchedulerContext into AicpuExecutor and remove old code - Replace AicpuExecutor's scheduler fields (core_exec_states_, core_trackers_, drain_state_, sched_perf_, s_pto2_payload_per_core) with single SchedulerContext sched_ctx_ member - Delete all old inline methods (cold-path, completion, dispatch, resolve_and_dispatch_pto2) from AicpuExecutor — ~1300 lines removed - Wire sched_ctx_ pointer members in init() and run() to connect to AicpuExecutor's shared atomic state - Switch run() to call sched_ctx_.resolve_and_dispatch() - Add core_assignments_ pointer to SchedulerContext for profiling flush - Set sched_ctx_.sched_ after rt creation (not in init() where rt may be null in device orchestration mode) aicpu_executor.cpp shrinks from ~2870 to ~1480 lines. * Update: sync docs with scheduler extraction refactoring - device_log_profiling.md: resolve_and_dispatch_pto2 -> SchedulerContext::resolve_and_dispatch - profiling_levels.md: update scheduler profiling file paths from aicpu_executor.cpp to scheduler/ directory, fix markdownlint issues (fenced code block languages, table alignment) - tensor-dump.md: update dump_tensors_for_task call site references to scheduler_completion.cpp and scheduler_dispatch.cpp - dynamic-linking.md: document SchedulerContext sched_ctx_ member in AicpuExecutor, update core_trackers_ reference path * Refactor: sync a5 scheduler extraction with a2a3 Apply the same scheduler refactoring to a5 as done in a2a3: - Create runtime/scheduler/ with SchedulerContext class and extracted methods (scheduler_types.h, scheduler_context.h, scheduler_completion.cpp, scheduler_cold_path.cpp, scheduler_dispatch.cpp) - Move pto_scheduler.h/.cpp into runtime/scheduler/ - Move runtime.cpp and pto_shared_memory.cpp into runtime/shared/ - Update build_config.py: host source_dirs uses "runtime/shared" - Switch CMake from GLOB to GLOB_RECURSE in all platform CMakeLists.txt - Update include paths to "scheduler/pto_scheduler.h" - Slim aicpu_executor.cpp from ~2854 to ~1244 lines - Add missing copyright headers to onboard CMakeLists.txt files a5-specific adaptations: - RUNTIME_MAX_WORKER=108 (vs a2a3's 72) in scheduler_context.h - CoreExecState profiling layout differs (no dispatch_count field) - LocalContext uses s_block_idx/s_block_num (vs block_idx/block_num) - No perf_aicpu_switch_buffer or perf_aicpu_flush_buffers calls * Update: sync a5 runtime docs with scheduler extraction refactoring - device_log_profiling.md: resolve_and_dispatch_pto2 -> SchedulerContext::resolve_and_dispatch - profiling_levels.md: update scheduler profiling file paths to a5 scheduler/ directory, fix markdownlint format issues, replace src/a2a3 references with src/a5 * Support: integrate PMU profiling into SchedulerContext Merge upstream PMU profiling (hw-native-sys#639) into the extracted scheduler: - scheduler_completion.cpp: add pmu_aicpu_record_task() call in complete_slot_task for per-task PMU counter sampling - scheduler_dispatch.cpp: add pmu_aicpu_init() in one-time init, PTO2_DISABLE_DUAL_ISSUE conditional for single-phase dispatch, pmu_aicpu_flush_buffers() at dispatch loop exit - scheduler_context.h: add physical_core_ids_ and cores_total_num_ pointers for PMU MMIO base resolution - aicpu_executor.cpp: wire physical_core_ids_ in init() and record physical_core_id during handshake
Summary
Extract the AICPU scheduler logic from
aicpu_executor.cppinto a dedicatedruntime/scheduler/subdirectory with a newSchedulerContextclass. Applied consistently to both a2a3 and a5 architectures..cppfiles, shrinking aicpu_executor.cpp from ~2800 lines to ~1200–1500 linesruntime/scheduler/directory: housesscheduler_types.h(type definitions & constants),scheduler_context.h(class declaration),scheduler_completion.cpp,scheduler_cold_path.cpp,scheduler_dispatch.cpp, plus relocatedpto_scheduler.h/.cppruntime/shared/directory: isolate host build dependencies to onlyruntime.cppandpto_shared_memory.cpp, preventing device-only code from compiling into the host .soPure structural refactoring — no behavioral change.
Changes
src/{a2a3,a5}/.../runtime/scheduler/src/{a2a3,a5}/.../aicpu/aicpu_executor.cppsched_ctx_member callssrc/{a2a3,a5}/.../runtime/shared/src/{a2a3,a5}/.../build_config.pysrc/{a2a3,a5}/platform/*/CMakeLists.txtdocs/