Refactor: extract finalize_common() (stacked on #913)#917
Conversation
|
Warning Review limit reached
More reviews will be available in 5 minutes and 45 seconds. Learn how PR review limits work. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
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 DeviceRunner implementations for the a2a3 and a5 platforms by consolidating shared diagnostics collectors, enablement setters, and the common finalization routine into the DeviceRunnerBase class. This refactoring also resolves a latent use-after-free bug in a2a3 by ensuring the device_wall buffer is freed before the allocator is finalized. The review feedback correctly identifies that finalize_common() fails to capture and propagate error codes from stream destruction and kernel argument finalization as documented, and provides a code suggestion to address this issue.
12c94a1 to
41913e0
Compare
…order Move the byte-identical body of `finalize()` from a2a3 and a5 onboard DeviceRunners into a shared `DeviceRunnerBase::finalize_common()`. Each arch's `finalize()` shrinks to a thin wrapper handling only its arch-specific bits (-32 net lines, identical behavior modulo one ordering normalization noted below). `finalize_common()` does the shared body: - Destroy persistent AICPU/AICore streams (no pre-sync — supported teardown path even on error-state streams). - `kernel_args_.finalize_device_args()` + `aicore_bin_handle_` / `binaries_loaded_` reset. - Free + clear `chip_callable_buffers_`, `orch_so_dedup_`, `callables_` (with hbg dlclose), `aicpu_seen_callable_ids_`, `aicpu_dlopen_total_`. - Release the three pooled arenas. - Free `device_wall_dev_ptr_` BEFORE `mem_alloc_.finalize()`. - `mem_alloc_.finalize()`. - Reset `block_dim_`, `worker_count_`, `aicore_kernel_binary_`. Each arch's `finalize()`: - a2a3: early return + attach + `finalize_collectors()` (4 shared + dep_gen) + `finalize_common()` + cached arena sizes reset + ACL-or-rt device reset + `device_id_` reset. - a5: early return + attach + inline 4-collector finalize + `finalize_common()` + cached arena sizes reset + `rtDeviceReset` + `device_id_` reset. Normalization: device_wall_dev_ptr_ free order changes on a2a3 — it previously freed AFTER `mem_alloc_.finalize()` AND after the device reset, routing the free through an already-finalized allocator on a torn-down device context (a latent no-op / UAF). It now matches a5's ordering: before `mem_alloc_.finalize()`. No observable behavior change in tests (the path was effectively dead), but this fixes a real ordering bug the refactor surfaced. Both arches built clean. a2a3 onboard smoke 8/8 in 15s.
Summary
Moves the byte-identical body of
finalize()from both onboardDeviceRunners into a sharedDeviceRunnerBase::finalize_common(). Each arch'sfinalize()shrinks to a thin wrapper covering its arch-specific bits. Net -32 lines across the four touched files.finalize_common()body (in base)rtStreamDestroyis the supported teardown path even on error-state streams).kernel_args_.finalize_device_args()+ resetaicore_bin_handle_/binaries_loaded_.chip_callable_buffers_,orch_so_dedup_,callables_(with hbg dlclose),aicpu_seen_callable_ids_,aicpu_dlopen_total_.device_wall_dev_ptr_beforemem_alloc_.finalize().mem_alloc_.finalize().block_dim_,worker_count_,aicore_kernel_binary_.Arch-specific
finalize()wrappersfinalize_collectors()(4 shared + dep_gen) →finalize_common()→ cached arena sizes reset → ACL-or-rt device reset →device_id_reset.finalize_common()→ cached arena sizes reset →rtDeviceReset→device_id_reset.One real behavior normalization (call out)
a2a3's
device_wall_dev_ptr_free order changes: it previously freed aftermem_alloc_.finalize()and after the device reset, which routed the free through an already-finalized allocator on a torn-down device context — a latent no-op or UAF. It now matches a5's longstanding ordering (free beforemem_alloc_.finalize()). The path was effectively dead before this PR (tests pass either way), but this fixes the ordering bug the refactor surfaced.Verification
dummy_task,alternating_matmul_add,prepared_callablesuite) — 8/8 passed in 15s.Test plan