Skip to content

Refactor: extract AICPU scheduler into SchedulerContext class (a2a3 + a5) - #654

Merged
jvjhfhg merged 10 commits into
hw-native-sys:mainfrom
poursoul:refactor-aicpu-sched
Apr 22, 2026
Merged

jvjhfhg merged 10 commits into
hw-native-sys:mainfrom
poursoul:refactor-aicpu-sched

Conversation

@poursoul

@poursoul poursoul commented Apr 22, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Extract the AICPU scheduler logic from aicpu_executor.cpp into a dedicated runtime/scheduler/ subdirectory with a new SchedulerContext class. Applied consistently to both a2a3 and a5 architectures.

  • Extract SchedulerContext class: split completion, cold-path, and dispatch methods out of AicpuExecutor into separate .cpp files, shrinking aicpu_executor.cpp from ~2800 lines to ~1200–1500 lines
  • Create runtime/scheduler/ directory: houses scheduler_types.h (type definitions & constants), scheduler_context.h (class declaration), scheduler_completion.cpp, scheduler_cold_path.cpp, scheduler_dispatch.cpp, plus relocated pto_scheduler.h/.cpp
  • Create runtime/shared/ directory: isolate host build dependencies to only runtime.cpp and pto_shared_memory.cpp, preventing device-only code from compiling into the host .so
  • Integrate PMU profiling: adapt upstream PMU sampling (Support: add a2a3 PMU profiling path #639) into SchedulerContext (per-task recording, init, buffer flush)
  • Sync documentation: update function names and file path references in device_log_profiling.md, profiling_levels.md, tensor-dump.md, and dynamic-linking.md

Pure structural refactoring — no behavioral change.

Changes

Area Change
src/{a2a3,a5}/.../runtime/scheduler/ New SchedulerContext class with 5 header/source files
src/{a2a3,a5}/.../aicpu/aicpu_executor.cpp Remove ~1700 lines of inline scheduler methods, replace with sched_ctx_ member calls
src/{a2a3,a5}/.../runtime/shared/ Relocate runtime.cpp and pto_shared_memory.cpp
src/{a2a3,a5}/.../build_config.py Host source_dirs changed to "runtime/shared"
src/{a2a3,a5}/platform/*/CMakeLists.txt GLOB → GLOB_RECURSE, add missing copyright headers
docs/ 4 docs updated with new paths and function references

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/a2a3/platform/onboard/aicore/CMakeLists.txt
- 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
@poursoul
poursoul force-pushed the refactor-aicpu-sched branch from 8ed62e3 to 41af1f0 Compare April 22, 2026 10:09
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
@poursoul
poursoul force-pushed the refactor-aicpu-sched branch from 41af1f0 to 8758241 Compare April 22, 2026 10:29
@poursoul poursoul changed the title Refactor aicpu sched Refactor: extract AICPU scheduler into SchedulerContext class (a2a3 + a5) Apr 22, 2026
@jvjhfhg
jvjhfhg merged commit f79c3bf into hw-native-sys:main Apr 22, 2026
14 checks passed
nalinaly pushed a commit to nalinaly/simpler that referenced this pull request Jul 31, 2026
… 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants