refactor(ir): extract wrapper-call finders into shared utility + drop InferFunctionCoreType body-walking#1321
Merged
Hzfengsy merged 4 commits intohw-native-sys:mainfrom May 9, 2026
Conversation
InnerCallFinder / CalleeFinder were duplicated three times across
codegen and analysis: two inside orchestration codegen
(FindWrapperInnerCall, FindGroupCallees) and one inside
orchestration_analysis.cpp's ComputeGroupEffectiveDirections. Each
walked a wrapper body, resolved GlobalVar callees via
Program::GetFunction, and collected matches by a slightly different
predicate.
Pull the scaffold out into wrapper_call_utils.{h,cpp} under
include/pypto/ir/transforms/utils/, exposing three named functions:
- FindFirstInnerCall(wrapper, program) -- first matching call
- FindGroupCallees(group_func, program) -- AIC/AIV name + first inner
call (priority AIC > AIV > InCore)
- CollectInnerCalls(wrapper, program) -- all non-Orchestration,
non-Opaque inner calls
Each user calls the named function it wants; the shared CallVisitor
under the hood handles tree traversal, GlobalVar resolution, and early
termination via a std::function callback.
Pure refactor, no behavior change; all 4439 unit tests pass.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughExtract inner-call discovery into a new IR utility module (header + implementation), refactor orchestration_codegen and orchestration_analysis to use the utilities, add the new source to the build, and update tests to mark many kernel/task functions as AIV instead of InCore. ChangesInner-Call Discovery Utility Extraction
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 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}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the codebase by extracting redundant IR visitor logic into a new utility module, wrapper_call_utils. It introduces centralized functions—FindFirstInnerCall, FindGroupCallees, and CollectInnerCalls—to handle the identification of inner function calls and group callees. These utilities replace local visitor implementations in orchestration_analysis.cpp and orchestration_codegen.cpp, improving code maintainability and reuse. I have no feedback to provide.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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/ir/transforms/utils/wrapper_call_utils.cpp`:
- Around line 73-90: The lambda passed to CallVisitor currently sets
info.inner_call only on the first match which allows lower-priority InCore/AIV
to win over later AIC; change the selection to enforce AIC > AIV > InCore by
updating info based on callee->func_type_: when you see an AIC, always
set/override info.aic_name, info.inner_call and info.inner_callee; when you see
an AIV, set/override info.aiv_name and set inner_call/inner_callee only if the
current inner is not an AIC; when you see an InCore, set inner_call/inner_callee
only if neither AIC nor AIV have been recorded. Keep using the existing
CallPtr/FunctionPtr and info.* fields (aic_name, aiv_name, inner_call,
inner_callee) to implement this priority logic.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 6390931e-374c-43b9-8df7-028806dfef4e
📒 Files selected for processing (5)
CMakeLists.txtinclude/pypto/ir/transforms/utils/wrapper_call_utils.hsrc/codegen/orchestration/orchestration_analysis.cppsrc/codegen/orchestration/orchestration_codegen.cppsrc/ir/transforms/utils/wrapper_call_utils.cpp
InferFunctionCoreType used to fall through to a 38-line tile-memory-space walk when the function arrived as plain InCore. That branch is reached only when a caller bypasses the IR pipeline (so ExpandMixedKernel never got to convert InCore → AIC/AIV). The only callers that bypass the pipeline are tests that hand-build user-style IR and call codegen.generate_orchestration directly. Convert those tests to declare their kernels with the appropriate @pl.function(type=pl.FunctionType.AIV) explicitly — all 32 affected kernels in tests/ut/codegen/test_orchestration_codegen.py operate on Vec memory only (pl.load + pl.add + pl.store, no matmul / Mat / Left / Right), so AIV is the correct concrete type. Tests that already run the pipeline (pm.run_passes / PassManager.get_strategy / direct calls to expand_mixed_kernel / outline_incore_scopes) keep InCore because the pipeline correctly converts them. InferFunctionCoreType is now an exact-match FunctionType lookup with INTERNAL_UNREACHABLE_SPAN on anything other than AIC / AIV — codegen contracts that only AIC / AIV reach this path, with a clear error pointing maintainers at the two recovery options if it ever fires. Migration of catalog item hw-native-sys#10 (Phase 2). Pure refactor in the codegen file (-38 / +6 LoC); test changes are mechanical type swaps.
lyfne123
changed the title
- wrapper_call_utils.h / .cpp: clarify FindGroupCallees doc and inline comment. The original CalleeFinder used first-match-in-source-order for inner_call (no AIC > AIV > InCore priority); my refactor copied that behavior verbatim, but the new docstring incorrectly claimed a type priority. Update both header and impl to describe the actual behavior — first-match wins, with the practical observation that Group bodies emitted by ExpandMixedKernel place AIC before AIV in source order, so the AIC call wins naturally. CodeRabbit flagged this as Major / Potential issue. Pure docstring + comment fix; behavior unchanged from main. If a real type-priority contract is desired, that's a separate, behavior-changing follow-up.
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/codegen/orchestration/orchestration_codegen.cpp`:
- Around line 60-66: The doc comment refers to a stale caller name
"GenerateUserFunctionCall"; update the comment to reference the actual caller
"GenerateFunctionCallCode" (and keep "GenerateSpmdCallCode") so the contract
described after ExpandMixedKernel is accurate; locate the comment block near the
top of orchestration_codegen.cpp (surrounding the comment mentioning
ExpandMixedKernel) and replace the stale symbol with GenerateFunctionCallCode
and adjust any surrounding wording if necessary to remain consistent with the
current callers.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 192cba72-c2b2-4170-ad52-f286a58c04bc
📒 Files selected for processing (4)
include/pypto/ir/transforms/utils/wrapper_call_utils.hsrc/codegen/orchestration/orchestration_codegen.cppsrc/ir/transforms/utils/wrapper_call_utils.cpptests/ut/codegen/test_orchestration_codegen.py
🚧 Files skipped from review as they are similar to previous changes (2)
- include/pypto/ir/transforms/utils/wrapper_call_utils.h
- src/ir/transforms/utils/wrapper_call_utils.cpp
Doc comment in InferFunctionCoreType referred to GenerateUserFunctionCall; the actual caller in this file is GenerateFunctionCallCode (line 1262). Updated to match.
Hzfengsy
approved these changes
May 9, 2026
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
2 participants
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.
Summary
Phase 2 of the codegen-analysis migration. Two related cleanups in one PR:
Commit 1 — #12: extract wrapper-call finders into shared IR utility
Pull three duplicated visitor classes (`InnerCallFinder` x2 + `CalleeFinder`) out of `src/codegen/orchestration/` into a shared IR-layer utility `include/pypto/ir/transforms/utils/wrapper_call_utils.h`, exposing three named functions:
A small `CallVisitor` scaffold inside the cpp handles tree traversal + `GlobalVar` resolution; each named function plugs in a callback. Pure refactor.
Commit 2 — #10: drop `InferFunctionCoreType` body-walking fallback
`InferFunctionCoreType` used to fall through to a 38-line tile-memory-space walk when the function arrived as plain `InCore`. That branch was reached only when a caller bypassed the IR pipeline (so `ExpandMixedKernel` never got to convert `InCore` → `AIC`/`AIV`). The only such callers are 23 tests in `test_orchestration_codegen.py` that hand-build user-style IR.
Decision: explicitly mark those test kernels with their concrete core type. All 32 affected kernels operate on `Vec` memory only (`pl.load` + `pl.add` + `pl.store`, no matmul / Mat / Left / Right), so `pl.FunctionType.AIV` is the correct type. Tests that already run the pipeline (`pm.run_passes` / `PassManager.get_strategy` / direct `expand_mixed_kernel` / `outline_incore_scopes` calls) keep `InCore` so the pipeline can convert them.
`InferFunctionCoreType` is now an exact-match `FunctionType` lookup with `INTERNAL_UNREACHABLE_SPAN` on anything else, with a clear error message pointing maintainers at the two recovery options if it ever fires.
Why one PR (and not two)
Both cleanups touch `src/codegen/orchestration/orchestration_codegen.cpp` and aim at the same goal (remove embedded analysis from codegen). Reviewing them together makes the codegen contract clearer in one place. Each commit stands alone if you want to land them separately.
Test plan
Stats