[BugFix] Fix two-instance modeling in ThreadSync cross-thread race checks - #2805
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! 🚀 |
📝 WalkthroughWalkthroughThe change adds fresh mutable-read handling and richer constraint-set operations, applies proof-based uniformity to synchronization hoisting, revises cross-thread conflict and parallel-loop race analysis, and expands regression tests for divergent and thread-private access patterns. ChangesThread synchronization analysis
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant TileLangThreadSyncPlanner
participant IsBlockUniformCondition
participant ConstrSet
participant Analyzer
TileLangThreadSyncPlanner->>IsBlockUniformCondition: prove if-condition uniformity
IsBlockUniformCondition->>ConstrSet: rename thread instances and merge constraints
ConstrSet->>Analyzer: populate predicate constraints
Analyzer-->>IsBlockUniformCondition: return agreement proof
IsBlockUniformCondition-->>TileLangThreadSyncPlanner: return uniformity result
TileLangThreadSyncPlanner->>TileLangThreadSyncPlanner: apply barrier hoisting decision
Possibly related PRs
Suggested reviewers: 🚥 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
🧹 Nitpick comments (7)
src/transform/thread_storage_sync.cc (2)
1533-1540: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value
threads.size() + idx - 3underflows when fewer than three thread dims are bound.
lhs.threads/rhs.threadsare indexed assize() + idx - 3; with fewer than three entries this is unsigned wraparound and an out-of-rangeArrayaccess. TileLang normally binds tx/ty/tz, so this mirrors the existing assumption inFindConflict, but anICHECK_GE(lhs.threads.size(), 3U)here would turn a potential OOB into a clear diagnostic.🤖 Prompt for 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. In `@src/transform/thread_storage_sync.cc` around lines 1533 - 1540, Add an ICHECK_GE(lhs.threads.size(), 3U) precondition before the indexing loop in the thread-variable synchronization logic, ensuring fewer than three bound thread dimensions produce a clear diagnostic instead of unsigned underflow. Keep the existing indexing and substitution behavior unchanged for valid inputs.
1014-1028: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSecond warning over-claims in the
requires_hoistcase.The unconditional "the race remains" message assumes both ends of the conflict sit inside the branch. That holds for the
depends_on_runtime && !is_block_uniformpath, but therequires_hoist && !proven_uniformpath hoists becauseThreadPartialSyncRewritercannot lower a partial barrier — there the hoist can be the correct fix, and telling users the race remains is misleading. Consider gating the second warning on the first predicate only.♻️ Suggested split
- if ((condition_prop.depends_on_runtime && !is_block_uniform) || - (condition_prop.requires_hoist && !proven_uniform)) { + bool divergent_runtime = + condition_prop.depends_on_runtime && !is_block_uniform; + bool unlowerable_partial = + condition_prop.requires_hoist && !proven_uniform; + if (divergent_runtime || unlowerable_partial) { LOG(WARNING) << "[ThreadSync] Hoisting sync from inside if to before if. " << "Condition is not safe for in-if sync: " << op->condition; - LOG(WARNING) << "[ThreadSync] The hoisted barrier no longer separates " + if (divergent_runtime) { + LOG(WARNING) << "[ThreadSync] The hoisted barrier no longer separates " "the accesses it was inserted for, as both ends of the " "conflict are inside the branch, so the race remains. " "Constrain the condition -- a T.assume on the " "parameters it reads -- to keep the barrier in place."; + }🤖 Prompt for 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. In `@src/transform/thread_storage_sync.cc` around lines 1014 - 1028, Gate the second warning about the hoisted barrier no longer separating accesses and the race remaining on the `condition_prop.depends_on_runtime && !is_block_uniform` predicate only. Keep the initial hoisting warning for both predicates, but do not emit the race-warning text when hoisting is caused solely by `condition_prop.requires_hoist && !proven_uniform`.src/transform/common/constr_visitor.h (2)
373-383: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueAsserts outside a
SeqStmtnow contribute no constraint.With the standalone
VisitStmt_(AssertStmtNode)removed, atirx::AssertStmtthat is the sole statement of a body (not wrapped in aSeqStmt) is silently ignored, so its premise is unavailable to downstream proofs. The direction is safe (fewer facts ⇒ more barriers), but it is an easy-to-miss hole; a fallback override that pushes the pure condition for the remainder of the enclosing scope, or a brief comment recording the limitation, would help.🤖 Prompt for 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. In `@src/transform/common/constr_visitor.h` around lines 373 - 383, Handle pure tirx::AssertStmt nodes that occur outside SeqStmt traversal so their conditions remain available to downstream proofs; update the AssertStmtNode handling in the statement visitor to push the condition for the enclosing scope, or document the intentional limitation if it cannot be supported. Preserve the existing side-effect purity check and mutable-state semantics.
261-267: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider rate-limiting this warning.
Mergeis reached fromFindConflict/PointerAccessIsDisjoint, which run once per access pair per statement, so a single mis-renamed bind would emit this warning O(n²) times per kernel.LOG_EVERY_N(WARNING, ...)(orDLOG) keeps the diagnostic without flooding the build log.🤖 Prompt for 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. In `@src/transform/common/constr_visitor.h` around lines 261 - 267, Rate-limit the conflicting-bind warning in ConstrSet::Merge, which can be reached repeatedly through FindConflict and PointerAccessIsDisjoint. Replace the unconditional LOG(WARNING) emission with the project’s supported rate-limited warning mechanism, such as LOG_EVERY_N, while preserving the existing diagnostic message and values.testing/python/transform/test_tilelang_transform_thread_sync.py (3)
1284-1285: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the placement this test is named for.
test_sync_may_stay_inside_block_uniform_guardonly checks that a barrier exists, so it would still pass if the barrier were hoisted out of the uniform guard — exactly the regression it is meant to pin.♻️ Add the position check
s = run_passes_script(func) assert 'T.tvm_storage_sync("shared")' in s, f"Expected a barrier for a cross-thread hazard:\n{s}" + sync_pos = s.index('T.tvm_storage_sync("shared")') + if_pos = s.index("if flag < 0") + assert sync_pos > if_pos, f"Barrier should stay inside the block-uniform guard:\n{s}"🤖 Prompt for 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. In `@testing/python/transform/test_tilelang_transform_thread_sync.py` around lines 1284 - 1285, Update test_sync_may_stay_inside_block_uniform_guard to assert that T.tvm_storage_sync("shared") remains inside the block’s uniform guard, not merely that it exists in the generated script. Inspect the generated string from run_passes_script(func) and verify the barrier’s position relative to the guard markers so hoisting it outside the guard fails the test.
590-593: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThree negative tests run
ThreadSyncoutside the shared harness. Each buildstvm.IRModule({"main": func})and calls the pass directly, skipping the cuda target attribute,AnnotateDeviceRegions, andSplitHostDevicethatrun_passesapplies — so these "no sync expected" assertions validate a different pass configuration than the rest of the file.
testing/python/transform/test_tilelang_transform_thread_sync.py#L590-L593: replace the manual module build and pass invocation withs = run_passes_script(func).testing/python/transform/test_tilelang_transform_thread_sync.py#L711-L714: same replacement withrun_passes_script(func).testing/python/transform/test_tilelang_transform_thread_sync.py#L769-L772: same replacement withrun_passes_script(func).🤖 Prompt for 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. In `@testing/python/transform/test_tilelang_transform_thread_sync.py` around lines 590 - 593, Update the three negative tests at testing/python/transform/test_tilelang_transform_thread_sync.py lines 590-593, 711-714, and 769-772 to replace manual IRModule construction and direct ThreadSync invocation with s = run_passes_script(func), preserving each existing no-sync assertion.
1116-1122: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExtend the test harness for non-standard pass orders.
InjectAssumesis exposed astilelang.transform.InjectAssumes, but this test bypassesrun_passesfor a specific pass ordering. If this ordering should be reusable, add support for an optional extra pass/list instead of duplicating the helper body.🤖 Prompt for 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. In `@testing/python/transform/test_tilelang_transform_thread_sync.py` around lines 1116 - 1122, The test’s manual pass sequence should be reusable through the existing run_passes helper. Extend run_passes to accept an optional extra pass or pass list, then invoke it for this non-standard ordering so InjectAssumes is applied without duplicating the helper’s setup logic.
🤖 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_thread_sync.py`:
- Around line 917-921: Update the placement assertion around run_passes_script
and the shared tvm_storage_sync marker to locate the specific divergent guard
text, matching the precise guard anchor used by the other placement tests,
instead of using the first generic "if " occurrence in the script.
---
Nitpick comments:
In `@src/transform/common/constr_visitor.h`:
- Around line 373-383: Handle pure tirx::AssertStmt nodes that occur outside
SeqStmt traversal so their conditions remain available to downstream proofs;
update the AssertStmtNode handling in the statement visitor to push the
condition for the enclosing scope, or document the intentional limitation if it
cannot be supported. Preserve the existing side-effect purity check and
mutable-state semantics.
- Around line 261-267: Rate-limit the conflicting-bind warning in
ConstrSet::Merge, which can be reached repeatedly through FindConflict and
PointerAccessIsDisjoint. Replace the unconditional LOG(WARNING) emission with
the project’s supported rate-limited warning mechanism, such as LOG_EVERY_N,
while preserving the existing diagnostic message and values.
In `@src/transform/thread_storage_sync.cc`:
- Around line 1533-1540: Add an ICHECK_GE(lhs.threads.size(), 3U) precondition
before the indexing loop in the thread-variable synchronization logic, ensuring
fewer than three bound thread dimensions produce a clear diagnostic instead of
unsigned underflow. Keep the existing indexing and substitution behavior
unchanged for valid inputs.
- Around line 1014-1028: Gate the second warning about the hoisted barrier no
longer separating accesses and the race remaining on the
`condition_prop.depends_on_runtime && !is_block_uniform` predicate only. Keep
the initial hoisting warning for both predicates, but do not emit the
race-warning text when hoisting is caused solely by
`condition_prop.requires_hoist && !proven_uniform`.
In `@testing/python/transform/test_tilelang_transform_thread_sync.py`:
- Around line 1284-1285: Update test_sync_may_stay_inside_block_uniform_guard to
assert that T.tvm_storage_sync("shared") remains inside the block’s uniform
guard, not merely that it exists in the generated script. Inspect the generated
string from run_passes_script(func) and verify the barrier’s position relative
to the guard markers so hoisting it outside the guard fails the test.
- Around line 590-593: Update the three negative tests at
testing/python/transform/test_tilelang_transform_thread_sync.py lines 590-593,
711-714, and 769-772 to replace manual IRModule construction and direct
ThreadSync invocation with s = run_passes_script(func), preserving each existing
no-sync assertion.
- Around line 1116-1122: The test’s manual pass sequence should be reusable
through the existing run_passes helper. Extend run_passes to accept an optional
extra pass or pass list, then invoke it for this non-standard ordering so
InjectAssumes is applied without duplicating the helper’s setup 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 0df4447a-223f-4608-b447-9e34c3efa7a0
📒 Files selected for processing (4)
src/transform/common/constr_visitor.hsrc/transform/thread_storage_sync.ccsrc/transform/verify_parallel_loop.cctesting/python/transform/test_tilelang_transform_thread_sync.py
| s = run_passes_script(func) | ||
| if 'T.tvm_storage_sync("shared")' in s: | ||
| sync_pos = s.index('T.tvm_storage_sync("shared")') | ||
| if_pos = s.index("if ") | ||
| assert sync_pos < if_pos, f"Barrier must be hoisted out of the divergent guard:\n{s}" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
s.index("if ") finds the first if in the whole script, not this guard.
If the printed module contains any earlier if (e.g. an assertion or a nested construct emitted by the harness), the position comparison checks the wrong anchor. The other placement tests already use the precise form; do the same here.
♻️ Anchor on the guard text
- s = run_passes_script(func)
- if 'T.tvm_storage_sync("shared")' in s:
- sync_pos = s.index('T.tvm_storage_sync("shared")')
- if_pos = s.index("if ")
- assert sync_pos < if_pos, f"Barrier must be hoisted out of the divergent guard:\n{s}"
+ s = run_passes_script(func)
+ if 'T.tvm_storage_sync("shared")' in s:
+ sync_pos = s.index('T.tvm_storage_sync("shared")')
+ if_pos = s.index("if bx * 4 + tx // 32 < n")
+ assert sync_pos < if_pos, f"Barrier must be hoisted out of the divergent guard:\n{s}"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| s = run_passes_script(func) | |
| if 'T.tvm_storage_sync("shared")' in s: | |
| sync_pos = s.index('T.tvm_storage_sync("shared")') | |
| if_pos = s.index("if ") | |
| assert sync_pos < if_pos, f"Barrier must be hoisted out of the divergent guard:\n{s}" | |
| s = run_passes_script(func) | |
| if 'T.tvm_storage_sync("shared")' in s: | |
| sync_pos = s.index('T.tvm_storage_sync("shared")') | |
| if_pos = s.index("if bx * 4 + tx // 32 < n") | |
| assert sync_pos < if_pos, f"Barrier must be hoisted out of the divergent guard:\n{s}" |
🤖 Prompt for 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.
In `@testing/python/transform/test_tilelang_transform_thread_sync.py` around lines
917 - 921, Update the placement assertion around run_passes_script and the
shared tvm_storage_sync marker to locate the specific divergent guard text,
matching the precise guard anchor used by the other placement tests, instead of
using the first generic "if " occurrence in the script.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/transform/thread_storage_sync.cc (1)
1503-1546: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winGuard the 3-thread indexing before renaming conflicts.
PointerAccessIsDisjointandFindConflictreadaccess.threads[threads.size() + idx - 3]foridx = 0..2, but new access entries are populated directly fromenv_threads_. If fewer than three rank-1 thread dims are active / reachable, this can trigger an OOB access. Scope the loops to the actual tail entries present (std::min({lhs.threads.unsigned_size(), rhs.threads.unsigned_size(), 3U})) or otherwise guarantee exactly three rank-1 thread vars before this logic runs.🤖 Prompt for 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. In `@src/transform/thread_storage_sync.cc` around lines 1503 - 1546, Guard the tail-thread indexing in PointerAccessIsDisjoint and FindConflict before accessing entries at threads.size() + idx - 3. Limit iteration to the available tail entries using the minimum of both thread-list sizes and three, or establish an equivalent precondition guaranteeing three rank-1 thread variables, so fewer active dimensions cannot cause out-of-bounds access.
🤖 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.
Outside diff comments:
In `@src/transform/thread_storage_sync.cc`:
- Around line 1503-1546: Guard the tail-thread indexing in
PointerAccessIsDisjoint and FindConflict before accessing entries at
threads.size() + idx - 3. Limit iteration to the available tail entries using
the minimum of both thread-list sizes and three, or establish an equivalent
precondition guaranteeing three rank-1 thread variables, so fewer active
dimensions cannot cause out-of-bounds access.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 4a36ef08-557d-4968-a925-eab3a1dcbf28
📒 Files selected for processing (2)
src/transform/common/constr_visitor.hsrc/transform/thread_storage_sync.cc
2 participants
Summary
FreshenMutableReads, improving constraint renaming/substitution/merging semantics, and refiningis_assume/assumption handling to avoid incorrect “two-instance” cross-thread equivalence checks.ifhoist candidates and refining condition/thread-variable property tracking.if, non-uniform if/else branches, loop cases), plus new assertions around barrier insertion and unorderable-hazard reporting with premise/assume/bind interactions.disjointnessto the spelling wordlist.C++ style / lint notes
docs/developer_guide/cpp_style.md.tvm::tl::FreshenMutableReadsandConstr::FreshenReads, and related constraint visitor/set behavior). These are warning-only and should be treated separately from correctness/build/test outcomes.