[BugFix] Fix LayoutInference divide-by-zero on non-power-of-two broadcast#2469
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthrough
ChangesLayoutInference divide-by-zero fix and tests
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related issues
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![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 `@testing/python/transform/test_tilelang_transform_layout_inference.py`:
- Around line 225-236: Update test_column_broadcast_fragment_values to include
the original failing CUDA repro widths 24 and 40 in the block_n parametrization,
in addition to the existing 32, 48, and 96 cases. Keep the rest of the test
logic unchanged so the regression coverage in this tilelang transform layout
inference test exercises the reported non-power-of-two shapes.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: c103e9b1-f761-4c8a-ac18-b0dcef9c4d03
📒 Files selected for processing (2)
src/layout/utils.cctesting/python/transform/test_tilelang_transform_layout_inference.py
|
@regression-perf |
Performance Regression Test ReportTriggered by: @LeiWang1999 Results
Artifacts
|
The I ran a controlled local A/B test using the repository's perf regression entry point ( Test environment: NVIDIA TITAN RTX (Turing, sm_75), driver 550.144.03, CUDA 12.4, PyTorch 2.6.0+cu124.
This PR (10 runs, ms):
|
|
@regression-perf |
Performance Regression Test ReportTriggered by: @LeiWang1999 Results
Artifacts
|
2 participants
Summary
LayoutInferencecan crash with an internalDivide by zerowhen lowering a valid column-broadcast fragment access whose tile width is not a power of two.Root cause
get_unused_iterstracks the portion of an iterator already covered by existingIterSplitExprs using:For non-power-of-two tile widths, iter split normalization may over-approximate the covered span. For example, an iterator with extent
48may be represented by a split span that covers up to64.The old logic only treated exact equality as full coverage:
So when the covered span became larger than the mark extent, it was treated as incomplete coverage.
get_unused_itersthen synthesized a leftover split with zero extent, which later normalized into a divide-by-zero.
Fix
Keep the split-chain progression intact, but treat proven over-coverage as full coverage at the final leftover check:
If the consumed span is equal to or greater than the original iterator extent, there is no remaining iterator space to synthesize, so
get_unused_itersdoes not emit a zero-extent leftover split.The fix also simplifies the newly advanced
expected_lower_factorbefore the final coverage check.Tested
testing/python/transform/test_tilelang_transform_layout_inference.py:block_n = 24, 40, 48, 64, 96block_n = 32, 48, 96, checking that the column broadcast matchesD * 2testing/python/transform/test_tilelang_transform_layout_inference.pytesting/python/layoutFixes #2394
Summary
LayoutInferencecompile-time crash in leftover-iterator handling by treating proven over-coverage as full coverage, preventing zero-extent leftover splits that could lead to a divide-by-zero during lowering of valid non-power-of-two column-broadcast fragments (issue#2394).block_nvalues (including newly added24and40), and verified that numerical output matches the reference broadcasted result.C++ style / lint notes
src/layout/utils.cc, so it is relevant todocs/developer_guide/cpp_style.md.