[Refactor] Simplify logic in the CompleteBufferFragment#1226
[Refactor] Simplify logic in the CompleteBufferFragment#1226LeiWang1999 merged 3 commits intotile-ai:mainfrom
CompleteBufferFragment#1226Conversation
WalkthroughRemoved a non-bijective pre-check and fallback in CompleteBufferFragment (src/op/parallel.cc), now directly inverting the bijective indices to compute replication extents; changed a persistent LOG(WARNING) to DLOG(WARNING) in LayoutNode::InverseWithLevel (src/layout/layout.cc). Changes
Sequence Diagram(s)sequenceDiagram
participant Caller as CompleteBufferFragment
participant Layout as Layout(loop_vars_)
participant Inv as ind_inv
participant Rep as RepExtentCalc
Note over Caller,Layout: OLD FLOW (with pre-check & fallback)
Caller->>Layout: compute bijective_indice
Layout-->>Caller: bijective_indice
Caller->>Caller: construct layout_before_inv
alt layout_before_inv not bijective
Caller->>Rep: fallback to no-rep inversion
Rep-->>Caller: fallback result
else bijective
Caller->>Layout: inverse Layout(loop_vars_, bijective_indice)
Layout-->>Inv: ind_inv
Caller->>Rep: derive indice_rep_extent, compute thd_b
Rep-->>Caller: rep results
end
sequenceDiagram
participant Caller as CompleteBufferFragment
participant Layout as Layout(loop_vars_)
participant Inv as ind_inv
participant Rep as RepExtentCalc
Note over Caller,Layout: NEW FLOW (direct inversion)
Caller->>Layout: compute bijective_indice
Layout-->>Caller: bijective_indice
Caller->>Layout: inverse Layout(loop_vars_, bijective_indice)
Layout-->>Inv: ind_inv
Caller->>Rep: derive indice_rep_extent, compute thd_b
Rep-->>Caller: rep results
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
🔇 Additional comments (1)
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 |
|
👋 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! 🚀 |
![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
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/op/parallel.cc(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
src/op/parallel.cc (1)
src/layout/layout.cc (2)
Layout(58-70)Layout(72-75)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Quick Lint
| std::nullopt) | ||
| ->CondenseReplicateVar(); | ||
| } | ||
| Layout ind_inv = Layout(loop_vars_, bijective_indice)->Inverse(); |
There was a problem hiding this comment.
🧩 Analysis chain
Verify that bijective mapping is guaranteed before unconditional inversion.
The removal of the pre-check for bijective mapping and the fallback logic means that Layout::Inverse() is now called unconditionally on line 652. If bijective_indice (created by appending rep_b from DivideUnusedIterators) is not actually bijective, this could lead to runtime failures or incorrect layout computations.
Potential risks:
Layout::Inverse()may fail or produce incorrect results for non-bijective mappings- Error messages may be less clear without the pre-check
- Silent correctness issues if the inverse is computed incorrectly
Please verify that DivideUnusedIterators guarantees a bijective mapping in all cases. Run the following script to check for existing test coverage of edge cases:
Additionally, consider:
- Adding explicit checks or assertions that validate bijectivity before inversion
- Documenting why the unconditional inversion is safe (if it is proven to be)
- Adding test cases that exercise non-trivial index patterns to ensure correctness
🏁 Script executed:
#!/bin/bash
# Description: Check for test coverage of CompleteBufferFragment and bijective mapping edge cases
# Search for tests related to CompleteBufferFragment
echo "=== Tests for CompleteBufferFragment ==="
rg -n "CompleteBufferFragment" --type=cpp --type=py
# Search for tests involving DivideUnusedIterators
echo -e "\n=== Tests for DivideUnusedIterators ==="
rg -n "DivideUnusedIterators" --type=cpp --type=py
# Search for bijective-related tests in parallel operations
echo -e "\n=== Bijective mapping tests ==="
rg -nC3 "bijective.*parallel|parallel.*bijective" --type=cpp --type=py -g '!*.md'
# Look for any comments or documentation about bijective guarantees
echo -e "\n=== Documentation about bijective guarantees ==="
rg -nC2 "DivideUnusedIterators" --type=cpp -A5 -B5Length of output: 4531
🏁 Script executed:
#!/bin/bash
# Check the implementation of DivideUnusedIterators and MakeFlattenedExpression
echo "=== DivideUnusedIterators implementation ==="
sed -n '122,160p' src/layout/utils.cc
echo -e "\n=== MakeFlattenedExpression signature ==="
rg -A10 "^PrimExpr MakeFlattenedExpression" src/layout/utils.cc
echo -e "\n=== Layout::Inverse() function ==="
rg -n "Inverse()" src/layout/layout.h
echo -e "\n=== Layout::Inverse() implementation ==="
rg -A20 "Layout.*::Inverse()" src/layout/layout.cc
echo -e "\n=== Check for bijective-related comments or checks ==="
rg -nC3 "bijective" src/op/parallel.cc
echo -e "\n=== Look for any error handling or assertions around Inverse ==="
rg -B5 -A5 "->Inverse()" src/op/parallel.ccLength of output: 995
🏁 Script executed:
#!/bin/bash
# Check the implementation of DivideUnusedIterators and MakeFlattenedExpression
echo "=== DivideUnusedIterators implementation ==="
sed -n '122,160p' src/layout/utils.cc
echo -e "\n=== MakeFlattenedExpression signature ==="
rg -A10 "^PrimExpr MakeFlattenedExpression" src/layout/utils.cc
echo -e "\n=== Layout::Inverse() function ==="
rg -n "Inverse()" src/layout/layout.h
echo -e "\n=== Layout::Inverse() implementation ==="
rg -A20 "Layout.*::Inverse()" src/layout/layout.cc
echo -e "\n=== Check for bijective-related comments or checks ==="
rg -nC3 "bijective" src/op/parallel.cc
echo -e "\n=== Look for error handling or assertions around Inverse in parallel.cc ==="
rg -B5 -A5 "Inverse" src/op/parallel.cc | head -50Length of output: 9042
🏁 Script executed:
#!/bin/bash
# Check the git history of the file to see what was removed
echo "=== Git history of line 652 changes ==="
git log --oneline -n 20 src/op/parallel.cc
echo -e "\n=== Check if InverseWithLevel has parameters for handling non-bijective ==="
rg -A30 "LayoutNode::InverseWithLevel()" src/layout/layout.cc | head -40
echo -e "\n=== Check what level parameter means in DetectIterMap ==="
rg -B3 "IterMapLevel" src/layout/layout.cc | head -30
echo -e "\n=== Check if there's any old fallback code in git ==="
git log -p --follow -S "pre-check\|fallback" src/op/parallel.cc | head -100Length of output: 3546
🏁 Script executed:
#!/bin/bash
# Get the full diff of the most relevant commit
echo "=== Full diff of commit 2957afc ==="
git show 2957afc -- src/op/parallel.cc | head -200
echo -e "\n=== Look for InverseWithLevel usage in the same file ==="
rg -n "InverseWithLevel\|Inverse()" src/op/parallel.ccLength of output: 4193
Remove unconditional Inverse() call or restore bijectivity pre-check and fallback handling.
The current code at line 652 removes critical safety mechanisms that were previously in place. The git history (commit 2957afc) shows the code previously included:
- Pre-check for bijectivity:
CanProveEqual(in_prod, out_prod)to verify the combined mapping was bijective - Fallback path: If the check failed, it would use
Layout(loop_vars_, indice_map_[buffer])->Inverse()withoutrep_b
Without these safeguards, calling Inverse() unconditionally can fail if bijective_indice is not actually bijective (for example, if DivideUnusedIterators output combined with the original indices creates a non-bijective mapping). The 2D path (lines 626–632) validates bijectivity explicitly before inversion; the fallback path should too.
Options:
- Restore the cardinality pre-check and fallback handling
- Or add explicit error handling around the
Inverse()call with recovery logic - Or document why bijectivity is now guaranteed
… to DLOG for symbolic layout fallback.
* fix * Fix logging level in LayoutNode::InverseWithLevel method from WARNING to DLOG for symbolic layout fallback. * lint fix --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2 participants
This pull request simplifies the logic in the
CompleteBufferFragmentfunction insrc/op/parallel.ccby removing a pre-check for bijective mapping and its fallback logic. The function now directly computes the inverse layout without checking for non-bijective mappings, streamlining the code.Code simplification:
rep_b, along with the associated fallback logic for non-bijective cases. The function now always computes the inverse layout directly.Summary by CodeRabbit
Refactor
Chore