Repository navigation
Refactor: split per-device invariants into a one-shot AICPU init entry - #1201
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughReplaces the single ChangesAICPU init + register_callable ABI refactor
Sequence Diagram(s)sequenceDiagram
participant Host as DeviceRunnerBase (host)
participant Loader as LoadAicpuOp
participant AICPU as AICPU SO (device)
participant Executor as AicpuExecutor
rect rgba(70, 130, 180, 0.5)
note over Host,AICPU: Device initialization (once per device)
Host->>Loader: ensure_binaries_loaded()
Host->>Loader: launch_aicpu_payload(InitArgs)
Loader->>AICPU: simpler_aicpu_init(InitArgs)
AICPU->>AICPU: latch device_id, log_level, log_info_v into globals
end
rect rgba(60, 179, 113, 0.5)
note over Host,Executor: Callable registration (once per callable)
Host->>Host: aicpu_register_callable(callable_id)
Host->>Host: build RegisterCallableArgs from CallableState
Host->>Loader: launch_aicpu_payload(RegisterCallableArgs)
Loader->>AICPU: simpler_aicpu_register_callable(RegisterCallableArgs)
AICPU->>Executor: aicpu_register_callable(RegisterCallableArgs)
Executor->>Executor: ensure_orch_so_loaded_core(so_addr, so_size, symbols)
Executor->>Executor: dlopen + dlsym → populate orch_so_table_[callable_id]
Host->>Host: commit_aicpu_callable_load(callable_id)
end
rect rgba(255, 165, 0, 0.5)
note over Host,AICPU: Per-task execution
Host->>Loader: launch_aicpu_kernel(KernelArgs)
Loader->>AICPU: simpler_aicpu_exec(KernelArgs)
AICPU->>AICPU: set_platform_regs only (invariants already latched)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 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 |
There was a problem hiding this comment.
Code Review
This pull request refactors the AICPU kernel launch mechanism by separating per-device one-shot invariants into a new InitArgs structure initialized via simpler_aicpu_init, and replacing the old prewarm path with a more efficient simpler_aicpu_register_callable entry that uses a smaller RegisterCallableArgs payload. Feedback focuses on guarding against potential nullptr values passed to snprintf in the a2a3 and a5 device runners, and adding defensive validation checks for launch arguments in LoadAicpuOp::AicpuKernelLaunch.
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.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/common/platform/onboard/host/device_runner_base.cpp`:
- Around line 358-388: The teardown path leaves the one-shot guard stale, so a
reused DeviceRunnerBase instance can skip the AICPU init after cleanup and miss
relatching the globals. Update DeviceRunnerBase::finalize_common() to clear
aicpu_init_launched_ alongside the other teardown state (like binaries_loaded_)
so ensure_aicpu_init_launched() will invoke launch_aicpu_payload again after
reload.
🪄 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: c205be38-6d08-4883-b531-4d468f4343f7
📒 Files selected for processing (24)
docs/aicpu-kernel-launch-mechanisms.mddocs/callable-identity-registration.mdsrc/a2a3/platform/include/common/kernel_args.hsrc/a2a3/platform/onboard/aicpu/kernel.cppsrc/a2a3/platform/sim/host/device_runner.cppsrc/a2a3/platform/sim/host/device_runner.hsrc/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cppsrc/a2a3/runtime/host_build_graph/runtime/runtime.hsrc/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cppsrc/a5/platform/include/common/kernel_args.hsrc/a5/platform/onboard/aicpu/kernel.cppsrc/a5/platform/sim/host/device_runner.cppsrc/a5/platform/sim/host/device_runner.hsrc/a5/runtime/host_build_graph/aicpu/aicpu_executor.cppsrc/a5/runtime/host_build_graph/runtime/runtime.hsrc/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cppsrc/common/aicpu_loader/host/load_aicpu_op.cppsrc/common/aicpu_loader/host/load_aicpu_op.hsrc/common/platform/onboard/host/c_api_shared.cppsrc/common/platform/onboard/host/device_runner_base.cppsrc/common/platform/onboard/host/device_runner_base.hsrc/common/platform/sim/host/c_api_shared.cppsrc/common/platform/sim/host/device_runner_base.cppsrc/common/platform/sim/host/device_runner_base.h
3eba236 to
298d96f
Compare
Per-device-fixed fields rode on every per-task KernelArgs launch even
though their values never change across runs. Lift them into a new
worker-init-time entry and slim the callable-registration payload.
- Add InitArgs { device_id, log_level, log_info_v } and a new
simpler_aicpu_init entry, launched once per device from
ensure_device_initialized(). It latches these into the resident AICPU
SO globals, so exec / register_callable no longer re-push them. Remove
the corresponding setters from simpler_aicpu_exec and drop the fields
from KernelArgs (regs / pmu_reg_addrs / ffts stay per-run; a5 AICore
still reads regs off the device KernelArgs copy).
- Rename the prewarm path to register_callable end to end
(simpler_aicpu_prewarm_callable -> simpler_aicpu_register_callable,
aicpu_prewarm_callable -> aicpu_register_callable, host
prewarm_callable -> aicpu_register_callable, KernelNames::PrewarmName
-> RegisterCallableName, add InitName).
- Introduce RegisterCallableArgs carrying only the orch-SO descriptor
extracted from Runtime (callable id, dev SO addr/size, entry/config
symbol names) so the register path no longer H2D's a full Runtime.
Refactor ensure_orch_so_loaded into a core taking those values; the
run-path Runtime* wrapper and the register entry both delegate to it.
- Generalize the AICPU launch path to void*+size so the three entries
(exec / init / register_callable) share one mechanism; register all
three symbols in the inner-SO JSON.
- hbg variants are no-ops (host-side orchestration); add the two missing
dev_orch_so getters to hbg Runtime so the shared sim path compiles.
Validated: a2a3+a5 build clean; sim register/exec paths pass on both
runtimes; full a2a3 onboard sweep passes (HBG + TMARB).
…rwarding + KernelArgs cruft (#1207) The device register-callable entry was a platform-layer forwarding shell (simpler_aicpu_register_callable -> runtime aicpu_register_callable) that every runtime had to satisfy by link-time strong symbol, forcing host_build_graph to carry a no-op it never uses. Move the entry to where the capability lives and let the runtime advertise it, so the common/platform layers hold no runtime-specific symbol knowledge. - The TMARB runtime now exports simpler_aicpu_register_callable(void*) directly (was the platform shell forwarding to aicpu_register_callable). The platform kernel.cpp forwarding shell and its extern decl are deleted (a2a3 + a5). - host_build_graph no longer defines the no-op: it simply does not export the symbol. Its inner SO links, JSON-registers, and loads without it. - Each runtime's host part reports the AICPU entry symbols it exports beyond the base {exec, init} via a new runtime_extra_aicpu_symbols() — TMARB returns register_callable, hbg returns none. The common AICPU loader builds its JSON-registration + handle-resolution set from base + runtime-reported extras, so it no longer hardcodes any runtime-specific symbol. The sim runners dlsym the new exported name (optional; absent on hbg). KernelArgs cleanup (leftover from the InitArgs split in #1201): - Remove the dead a2a3 KernelArgs::device_id (orch device id is latched once via InitArgs/simpler_aicpu_init; no per-run reader remained). - Reorder both arches' KernelArgs so all uint64_t precede the uint32_t tail, which removes the explicit _pad alignment fillers entirely. runtime_args@0 / regs@8 stay offset-locked (static_asserts hold). - Fix stale comments that still described log config / device id as travelling on per-run KernelArgs (platform_regs.h, device_log.{h,cpp}, host_log.{h,cpp}, device_runner_base.h init_runtime_args_with_metadata). Validated: a2a3+a5 build clean; TMARB sim register->run passes; a2a3 onboard passes for both runtimes — hbg loads and runs without exporting simpler_aicpu_register_callable, TMARB register path intact.
…tail Make the host/device boundary of the tensormap_and_ringbuffer (trb) Runtime explicit in the type system instead of resting on a field-ordering convention, and drop the register-only fields that the RegisterCallableArgs hand-off (hw-native-sys#1201/hw-native-sys#1203/hw-native-sys#1207) already superseded. Device-read fields move into a named first member `DeviceRuntimeLaunchDesc dev` (offset 0); host-only state stays in the tail and is no longer uploaded. The H2D rtMemcpy and the AICPU cache_invalidate narrow from sizeof(Runtime) (~34KB) to sizeof(dev). - runtime.h (a2a3+a5): define DeviceRuntimeLaunchDesc, nest as `dev`, add static_asserts (offsetof==0, is_standard_layout_v, is_trivially_copyable_v, sizeof%64==0) and inline accessors so the shared platform layer compiles against both trb and hbg Runtimes; trb-only device code reads runtime->dev.X directly. The struct is alignas(64) so sizeof stays a cache-line multiple and cache_invalidate_range(sizeof(dev)) never rounds into a neighbouring line. - shared/runtime.cpp: ctor + getters/setters forward into dev.*; define runtime_device_copy_size() = sizeof(DeviceRuntimeLaunchDesc). hbg defines it = sizeof(Runtime) (uploads the whole object). device_runner_helpers.cpp uses it for both the alloc and the rtMemcpy length, staying runtime-agnostic. - Drop the register-only fields (dev_orch_so_addr_/size_, device_orch_func_name_/config_name_, register_new_callable_id_) from trb Runtime and the parity copies from hbg. The sim register path builds RegisterCallableArgs straight from CallableState (mirroring onboard) instead of round-tripping through a throwaway Runtime; the dead post-run register_new_callable_id() commit gate is removed; set_active_callable_id loses its is_new param and stamp_orch_so collapses to setting the cid. - test_runtime_orch_so.cpp now pins the RegisterCallableArgs POD shape. - Drop the unused RUNTIME_MAX_ORCH_SO_SIZE macro (no users; dead pre-PR). Verified: a2a3+a5 sim+onboard builds (trb+hbg); a2a3/a5 sim prepared_callable + orch_so_cache; a2a3 onboard prepared_callable (dlopen-count exactly-once) + orch_so_cache + dummy_task + mixed_example; hbg a2a3 onboard regression. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…tail (#1216) Make the host/device boundary of the tensormap_and_ringbuffer (trb) Runtime explicit in the type system instead of resting on a field-ordering convention, and drop the register-only fields that the RegisterCallableArgs hand-off (#1201/#1203/#1207) already superseded. Device-read fields move into a named first member `DeviceRuntimeLaunchDesc dev` (offset 0); host-only state stays in the tail and is no longer uploaded. The H2D rtMemcpy and the AICPU cache_invalidate narrow from sizeof(Runtime) (~34KB) to sizeof(dev). - runtime.h (a2a3+a5): define DeviceRuntimeLaunchDesc, nest as `dev`, add static_asserts (offsetof==0, is_standard_layout_v, is_trivially_copyable_v, sizeof%64==0) and inline accessors so the shared platform layer compiles against both trb and hbg Runtimes; trb-only device code reads runtime->dev.X directly. The struct is alignas(64) so sizeof stays a cache-line multiple and cache_invalidate_range(sizeof(dev)) never rounds into a neighbouring line. - shared/runtime.cpp: ctor + getters/setters forward into dev.*; define runtime_device_copy_size() = sizeof(DeviceRuntimeLaunchDesc). hbg defines it = sizeof(Runtime) (uploads the whole object). device_runner_helpers.cpp uses it for both the alloc and the rtMemcpy length, staying runtime-agnostic. - Drop the register-only fields (dev_orch_so_addr_/size_, device_orch_func_name_/config_name_, register_new_callable_id_) from trb Runtime and the parity copies from hbg. The sim register path builds RegisterCallableArgs straight from CallableState (mirroring onboard) instead of round-tripping through a throwaway Runtime; the dead post-run register_new_callable_id() commit gate is removed; set_active_callable_id loses its is_new param and stamp_orch_so collapses to setting the cid. - test_runtime_orch_so.cpp now pins the RegisterCallableArgs POD shape. - Drop the unused RUNTIME_MAX_ORCH_SO_SIZE macro (no users; dead pre-PR). Verified: a2a3+a5 sim+onboard builds (trb+hbg); a2a3/a5 sim prepared_callable + orch_so_cache; a2a3 onboard prepared_callable (dlopen-count exactly-once) + orch_so_cache + dummy_task + mixed_example; hbg a2a3 onboard regression.
hw-native-sys#1201) Per-device-fixed fields rode on every per-task KernelArgs launch even though their values never change across runs. Lift them into a new worker-init-time entry and slim the callable-registration payload. - Add InitArgs { device_id, log_level, log_info_v } and a new simpler_aicpu_init entry, launched once per device from ensure_device_initialized(). It latches these into the resident AICPU SO globals, so exec / register_callable no longer re-push them. Remove the corresponding setters from simpler_aicpu_exec and drop the fields from KernelArgs (regs / pmu_reg_addrs / ffts stay per-run; a5 AICore still reads regs off the device KernelArgs copy). - Rename the prewarm path to register_callable end to end (simpler_aicpu_prewarm_callable -> simpler_aicpu_register_callable, aicpu_prewarm_callable -> aicpu_register_callable, host prewarm_callable -> aicpu_register_callable, KernelNames::PrewarmName -> RegisterCallableName, add InitName). - Introduce RegisterCallableArgs carrying only the orch-SO descriptor extracted from Runtime (callable id, dev SO addr/size, entry/config symbol names) so the register path no longer H2D's a full Runtime. Refactor ensure_orch_so_loaded into a core taking those values; the run-path Runtime* wrapper and the register entry both delegate to it. - Generalize the AICPU launch path to void*+size so the three entries (exec / init / register_callable) share one mechanism; register all three symbols in the inner-SO JSON. - hbg variants are no-ops (host-side orchestration); add the two missing dev_orch_so getters to hbg Runtime so the shared sim path compiles. Validated: a2a3+a5 build clean; sim register/exec paths pass on both runtimes; full a2a3 onboard sweep passes (HBG + TMARB).
Summary
Per-device-fixed values rode on every per-task
KernelArgslaunch eventhough they never change across runs. This lifts them into a single
worker-init-time AICPU entry and slims the callable-registration payload so it
no longer ships a full
Runtime.InitArgs+simpler_aicpu_init—{ device_id, log_level, log_info_v }launched once per device from
ensure_device_initialized(). It latches theseinto the resident AICPU SO globals, so
exec/register_callableno longerre-push them. The corresponding setters are removed from
simpler_aicpu_execand the fields dropped from
KernelArgs.regs/pmu_reg_addrs/fftsstay per-run (a5 AICore still reads
regsoff the deviceKernelArgscopy).prewarm→register_callableend to end —simpler_aicpu_prewarm_callable→simpler_aicpu_register_callable,aicpu_prewarm_callable→aicpu_register_callable, hostprewarm_callable→
aicpu_register_callable,KernelNames::PrewarmName→RegisterCallableName, plus newInitName.RegisterCallableArgscarrying only the orch-SO descriptor extractedfrom
Runtime(callable id, dev SO addr/size, entry/config symbol names), sothe register path no longer H2D's a full
Runtime.ensure_orch_so_loadedis refactored into a core taking those values; the run-path
Runtime*wrapper and the register entry both delegate to it.
void* + sizeso the three entries(exec / init / register_callable) share one mechanism; all three symbols are
registered in the inner-SO JSON.
hbgvariants are no-ops (host-side orchestration); the two missingdev_orch_sogetters are added tohbgRuntimeso the shared sim pathcompiles.
Testing