fix(codegen): dispatch dual_aiv_dispatch AIV kernels on both vector lanes#2012
Open
luohuan19 wants to merge 3 commits into
Open
fix(codegen): dispatch dual_aiv_dispatch AIV kernels on both vector lanes#2012luohuan19 wants to merge 3 commits into
luohuan19 wants to merge 3 commits into
Conversation
Two parser-side asymmetries made python_print -> parse lossy for a function
carrying a pl.split_aiv region, so every pass over such IR failed the roundtrip
verifier (only ever noticed now, because the split_aiv pass tests run under an
empty PassContext with no instruments).
1. An explicit `attrs={"split": pl.SplitMode.NONE}` was dropped on parse while
the printer emits it. This hit exactly the functions a task-parallel
`pl.split_aiv(..., mode=NONE)` region produces: OutlineIncoreScopes derives a
representative function-level `split` mode of NONE from a single NONE region.
Store every mode instead. This is behaviour-preserving — Function::GetSplitMode()
maps `split == 0` to nullopt and is the sole reader of the attr, so a stored
NONE and an absent `split` are indistinguishable to every consumer.
2. A `pl.split_aiv` loop directly in an `@pl.function(type=pl.FunctionType.InCore)`
body was wrapped in a synthesized InCoreScopeStmt. That wrapper exists to keep a
bare top-level region eligible for OutlineIncoreScopes, but an InCore function is
already a core function — there is nothing left to outline. It is also the shape
post-outline IR prints as (outlining consumes the InCoreScopeStmt and leaves the
region directly in an InCore body), so the reparsed function grew a scope the
original did not have. Emit the region in place there too.
Both arms are guarded by new roundtrip tests.
…anes Fixes hw-native-sys#2006 A `pl.split_aiv` region is task-parallel: both AIV lanes run the body on disjoint work selected by `aiv_id`. SplitVectorKernel stamps such an AIV kernel `dual_aiv_dispatch`, but only the mixed-kernel Group submit honoured it. The three single-kernel dispatch paths — direct call, Spmd wrapper, and AIV-only Group — all built their submit through CoreTypeToSubmitPrefix(), which unconditionally yields `rt_submit_aiv_task` for VECTOR. `rt_submit_aiv_task` fills only the AIV0 slot of MixedKernels, so the runtime schedules a single-AIV task: one AIV core per block. The second lane never launches, and the lone lane that does reads a `get_sub_block_id()` the runtime leaves undefined for single-AIV tasks (runtime common/intrinsic.h), so its `aiv_id` is really its physical cluster position. Half the work is silently dropped — the destination rows come back as exact zeros, with no error, no NaN, and a block-parity pattern for which lane survives. A pure-vector kernel (no cube work) hits this on the Spmd-wrapper path. Route the three single-kernel paths through a shared BuildSingleKernelSubmitExpr(), which emits `MixedKernels{INVALID_KERNEL_ID, id, id}` + rt_submit_task for a `dual_aiv_dispatch` VECTOR kernel. Two active AIV slots make it a MIX-shape task, so the scheduler places both lanes of one cluster under the same block_idx and gives them sub_block_id 0 and 1; the cluster's AIC core is unused. No runtime change is needed. A plain vector kernel carries no `dual_aiv_dispatch` and keeps `rt_submit_aiv_task`, so it retains its full independent-AIV parallelism; mixed kernels are untouched. The existing NONE-region codegen ST compiled this exact shape but only asserted a .pto was emitted, which is why the defect shipped; it now also asserts the submit launches both lanes.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (8)
📝 WalkthroughWalkthroughThe change preserves ChangesDual-AIV Split-NONE dispatch
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant SplitAivParser
participant OrchestrationCodegen
participant MixedKernels
participant Runtime
SplitAivParser->>OrchestrationCodegen: preserve SplitMode.NONE and dual-AIV function shape
OrchestrationCodegen->>MixedKernels: configure AIC inactive and both AIV slots
OrchestrationCodegen->>Runtime: submit rt_submit_task(mixed_...)
Runtime->>MixedKernels: execute both AIV lanes
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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 |
![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 addresses an issue where dual-AIV vector kernels (e.g., split_aiv kernels) were incorrectly scheduled as single-AIV tasks, causing one of the lanes to never launch and silently dropping half the work. It fixes this by dispatching dual-AIV kernels as two-lane MixedKernels instead of using rt_submit_aiv_task. Additionally, it preserves SplitMode.NONE attributes during parsing to prevent lossy roundtrips, and ensures that split_aiv regions directly inside InCore function bodies do not synthesize spurious wrapper scopes. Feedback from the reviewer points out that checking only ir.FunctionType.InCore in ast_parser.py is too restrictive, as AIC and AIV are also core functions that do not require outlining, and suggests expanding the check to include them.
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.
- Widen the parser's core-function guard from InCore to the full IsInCoreType() set (InCore, AIC, AIV). An @pl.function(type=pl.FunctionType.AIV) body carrying a pl.split_aiv region is already a core function, but the InCore-only check sent it down the bare-top-level arm and synthesized a spurious InCoreScopeStmt — a scope the original does not have, which OutlineIncoreScopes would then try to outline out of an already-specialized kernel. Extracted the set as _CORE_FUNCTION_TYPES, mirroring the C++ IsInCoreType(). Guarded by a new roundtrip test. Reported by gemini-code-assist. - Fix two pyright reportOptionalMemberAccess errors in the roundtrip tests (Program.get_function() returns Function | None) — this was failing pre-commit in CI. - Correct the pass attribution in test_roundtrip_split_none_attr_survives: the function-level split attr is stamped by LowerAutoVectorSplit (the only stamper that stores NONE — SplitVectorKernel explicitly skips it), not by OutlineIncoreScopes, which never touches the attr. - Condense the new orchestration-codegen doc section (EN + zh-CN): the prose, code block and table stated the same fact three times, and the EN file had crossed the 700-line split threshold (701 -> 693). Cross-reference the MixedKernels table instead of repeating it. - Scope the ST's rt_submit_aiv_task negative assertion to the dual-AIV kernel's own id, so it stays correct if the fixture ever gains a plain vector kernel (which legitimately keeps rt_submit_aiv_task).
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
1 participant
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
A
pl.split_aivregion is task-parallel: both AIV lanes run the body on disjoint work selected byaiv_id.SplitVectorKernelstamps such an AIV kerneldual_aiv_dispatch, but only the mixed-kernel Group submit honoured it. The three single-kernel dispatch paths — direct call, Spmd wrapper, and AIV-only Group — all built their submit throughCoreTypeToSubmitPrefix(), which unconditionally yieldsrt_submit_aiv_taskfor VECTOR.rt_submit_aiv_taskfills only the AIV0 slot ofMixedKernels, so the runtime schedules a single-AIV task: one AIV core per block. The second lane never launches, and the lone lane that does reads aget_sub_block_id()the runtime leaves undefined for single-AIV tasks, so itsaiv_idis really its physical cluster position. Half the work is silently dropped — the destination rows come back as exact zeros, with no error, no NaN, and a block-parity pattern for which lane survives. A pure-vector kernel (no cube work) hits this on the Spmd-wrapper path.Changes
fix(language): makesplit_aivIR survive print → parse unchanged. Two parser-side asymmetries madepython_print→ parse lossy for a function carrying apl.split_aivregion, failing the roundtrip verifier on every pass over such IR:attrs={"split": pl.SplitMode.NONE}was dropped on parse while the printer emits it. Store every mode instead — behaviour-preserving, sinceFunction::GetSplitMode()mapssplit == 0tonulloptand is the attr's sole reader.pl.split_aivloop directly in an InCore function body was wrapped in a synthesizedInCoreScopeStmt. An InCore function is already a core function — nothing is left to outline — and it is the shape post-outline IR prints as. Emit the region in place there too.fix(codegen): dispatchdual_aiv_dispatchAIV kernels on both vector lanes. Route the three single-kernel paths through a sharedBuildSingleKernelSubmitExpr(), which emitsMixedKernels{INVALID_KERNEL_ID, id, id}+rt_submit_taskfor adual_aiv_dispatchVECTOR kernel. Two active AIV slots make it a MIX-shape task, so the scheduler places both lanes of one cluster under the sameblock_idxand gives themsub_block_id0 and 1; the cluster's AIC core is unused. No runtime change is needed.A plain vector kernel carries no
dual_aiv_dispatchand keepsrt_submit_aiv_task, so it retains its full independent-AIV parallelism. Mixed kernels are untouched.Testing
tests/ut/ir/printing/test_split_aiv_roundtrip.py) guard both parser arms.tests/ut/codegen/test_orchestration_codegen_more.py) covers thedual_aiv_dispatchsubmit shape..ptowas emitted — which is why the defect shipped. It now also asserts the submit launches both lanes.docs/en/dev/codegen/01-orchestration_codegen.md+ zh-cn).Related Issues
Fixes #2006