fix: preserve singleton broadcast dims in SplitVectorKernel

[ndleslx]

Summary

• preserve singleton split-axis dimensions in SplitVectorKernel instead of halving them unconditionally
• only adjust offsets and tracked split metadata when the split is actually applied
• add a regression test for [1, 128] broadcast tile loads used by col_expand_mul under SplitMode.UP_DOWN

Repro

This fixes the failure where split AIV kernels hit:

SplitVectorKernel requires an even split dimension, got 1

for broadcast tile loads such as Tile[[1, 128], ...] on the split axis.

Testing

• pip install -e ".[dev]"
• python -m pytest tests/ut/ir/transforms/test_split_vector_kernel.py -k singleton_broadcast -v

Closes #975

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 45e7c288-5c6e-4227-a38a-ebf7b90ee8fd

📥 Commits

Reviewing files that changed from the base of the PR and between 53d9f93 and 8574ef1.

📒 Files selected for processing (2)

• src/ir/transforms/split_vector_kernel_pass.cpp
• tests/ut/ir/transforms/test_split_vector_kernel.py

---

📝 Walkthrough

Walkthrough

Modified the split_vector_kernel_pass.cpp to preserve singleton dimensions on the split axis instead of unconditionally halving them. Introduced ComputeSplitDimInfo() function returning both the computed dimension and a split_applied flag, gating offset adjustments and tile tracking on actual splits. Added regression test validating singleton broadcast tile preservation.

Changes

Cohort / File(s)	Summary
Split Vector Kernel Pass Core Logic src/ir/transforms/split_vector_kernel_pass.cpp	Introduced ComputeSplitDimInfo() function that returns both computed dimension expression and split_applied flag. Updated HalveTileShape() and HalveTupleElement() to accept preserve_singleton parameter. Modified ProcessStmt() to conditionally apply offset adjustments and tile-variable tracking only when split_applied is true, while preserving singleton dimensions (e.g., [1, N]) in broadcast tiles.
Regression Test tests/ut/ir/transforms/test_split_vector_kernel.py	Added test_load_preserves_singleton_broadcast_dim test validating that split_vector_kernel() preserves singleton broadcast tiles under SplitMode.UP_DOWN, ensuring split dimension (e.g., [1, 128]) remains unchanged while normal tiles (e.g., [16, 128] → [8, 128]) are correctly halved.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• fix(runtime): stabilize mixed-kernel split AIV flow #951: Modifies split_vector_kernel_pass.cpp split handling with tracked half-dim info and conditional tile tracking/offset updates.
• feat(codegen): Enable a2a3 up/down split cross-core pipe codegen #869: Directly modifies same file internals for AIV/AIC split handling and halving logic in ProcessFunction/ProcessStmt.
• fix(codegen): share split AIV kernels on a2a3 #892: Adjusts tile dims, offsets, and subblock-dependent logic computation in same pass's AIV/AIC split lowering.

Suggested reviewers

• lyfne123

Poem

🐰 A tile so small, with just one in its way,
No need to halve it—let it stay!
While siblings split right down the seam,
The broadcast singleton achieves its dream. ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 27.27% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: preserving singleton broadcast dimensions in SplitVectorKernel instead of halving them unconditionally.
Description check	✅ Passed	The description is well-organized with Summary, Repro, and Testing sections that directly relate to the changeset and issue being fixed.
Linked Issues check	✅ Passed	The code changes fully implement the requirements from issue #975: preserving singleton dimensions, conditional offset/metadata adjustments, and adding a regression test for [1,128] broadcast tiles.
Out of Scope Changes check	✅ Passed	All changes are scoped to fixing singleton broadcast dimension handling in SplitVectorKernel; no unrelated modifications were introduced.

_{✏️ 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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to preserve singleton dimensions during vector kernel splitting for AIV operations by replacing the ComputeHalfDimSize function with ComputeSplitDimInfo. This allows the transformation to skip splitting for dimensions of size 1 when the preserve_singleton flag is enabled, preventing potential value errors. The logic was updated for tile.load, tile.full, and tile.create operations, and a new unit test was added to verify the behavior. Feedback points out that the tile.tpop_from_aic operation was not fully updated to use this new logic, which could lead to crashes or incorrect offset adjustments for singleton tiles.

[gemini-code-assist]

The tile.tpop_from_aic operation handling is incomplete and will still crash when encountering a singleton dimension on the split axis.

1. This line calls ComputeSplitDimInfo with the default preserve_singleton=false, which will throw a ValueError for dimensions of size 1.
2. Additionally, tile_vars should only be updated if the split was actually applied (split_applied == true). Otherwise, subsequent tile.store operations will incorrectly adjust offsets for the non-split singleton tile.
3. The helper RebuildTpopWithHalvedShape (called at line 296) also needs to be updated to support and pass the preserve_singleton flag to HalveTileShape to avoid a similar crash during type reconstruction.

References

1. When an AIV operation produces a TileType, ensure that any shape-related arguments within the Call itself are also updated (e.g., halved) to maintain type consistency and prevent failures in subsequent passes or codegen.