[Bugfix] Fix threadIdx variable lookup by thread_tag instead of position in ThreadSync

[LeiWang1999]

Added condition expression visiting in constr_visitor.h to ensure buffer accesses in if-conditions are properly recorded before applying guards

Summary by CodeRabbit

• Bug Fixes

  • Ensured buffer accesses inside conditional expressions are captured correctly by visiting conditions before applying branch guards.

• Tests

  • Added test coverage validating thread-synchronization hoisting with shared memory in loops and non-uniform conditional branches.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

Caution

Review failed

The pull request is closed.

📝 Walkthrough

Walkthrough

The visitor for if-then-else expressions now visits the condition expression unguarded before applying branch guards, ensuring condition buffer accesses are recorded; a test was added to validate thread-sync hoisting for non-uniform ifs inside loops using shared (thread) memory.

Changes

Cohort / File(s)	Summary
Constraint Visitor Modification src/transform/common/constr_visitor.h	VisitIfThenElseExpr now visits the condition unguarded before applying guards to true/false branches, so condition buffer accesses are pushed to the constraint stack.
Thread Sync Test Addition testing/python/transform/test_tilelang_transform_thread_sync.py	Added test_sync_hoist_non_uniform_if_in_loop_with_shared_memory (appears duplicated) to assert storage sync is hoisted before non-uniform if inside a loop using shared memory.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• [Refactor] Move ConstrVisitor to src/transform/common/constr_visitor.h for reuse #1622: Changes to ConstrVisitor's VisitIfThenElseExpr implementation affecting condition traversal logic.
• [Bugfix] Improves the accuracy of dependency analysis in the storage access  #1205: Visitor logic updates for IfThenElse nodes to include memory accesses from condition expressions for sync/insertion.
• [Bugfix] Fix threadIdx variable lookup by thread_tag instead of position in ThreadSync #1829: Contains similar change to visit condition unguarded and adds a ThreadSync hoisting test.

Suggested reviewers

• kurisu6912
• SiriusNEO
• silentCoder-dev

Poem

🐰 I nibble guards and peek the road ahead,
I note each buffer where my paws have tread,
A sync I tuck before branches diverge,
Hoisted and tidy—now threads safely merge. 🥕

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main bugfix: correcting threadIdx variable lookup methodology in ThreadSync, which is evidenced by the refactoring in thread_storage_sync.cc switching from optional tag-based lookup to direct position-based indexing.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch fix-thread-sync-threadidx-lookup

Warning

Review ran into problems

🔥 Problems

Git: Failed to clone repository. Please run the @coderabbitai full review command to re-trigger a full review. If the issue persists, set path_filters to include or exclude specific files.

---

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.}

[Copilot]

Pull request overview

Fixes ThreadSync-related guard/sync placement by ensuring buffer accesses in if-conditions are visited/recorded prior to applying guards, and adds a regression test for hoisting sync before a non-uniform if inside a loop.

Changes:

• Update ConstrVisitor to visit IfThenElse condition expressions before creating/applying guards.
• Add a CUDA test covering sync hoisting with shared memory reads in a loop/conditional.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
testing/python/transform/test_tilelang_transform_thread_sync.py	Adds a regression test ensuring shared-memory sync is hoisted before a non-uniform if inside a loop.
src/transform/common/constr_visitor.h	Ensures condition expressions are visited unguarded so their buffer accesses are recorded before guard application.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.