[Feature] T.sync_warp & T.shfl_sync support

[silentCoder-dev]

Support for #1598

Summary by CodeRabbit

• New Features

  • Introduced warp-level thread synchronization and broadcast shuffle operations to enable advanced CUDA kernel synchronization patterns and efficient inter-thread communication within warps.

• Tests

  • Added comprehensive test coverage validating warp-level synchronization and shuffle operations functionality in CUDA kernels.
  • Added test suite for broadcast value hoisting optimization pass applied during code generation and lowering.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[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

This PR introduces warp-level synchronization primitives (sync_warp, shfl_sync) to TileLang, implements a transformation pass to hoist constant broadcast values into temporary variables for optimization, and adds comprehensive test coverage for both features in CUDA kernels.

Changes

Cohort / File(s)	Summary
Language Builtins tilelang/language/builtin.py	Adds two new public synchronization functions: sync_warp(mask) emitting __syncwarp for warp-level synchronization, and shfl_sync(mask, value, srcLane, width) emitting __shfl_sync for warp-shuffle broadcast operations.
Transformation Pass tilelang/transform/hoist_broadcast_values.py	Introduces HoistBroadcastValuesMutator that extracts constant values from Broadcast operations into LetStmt-bound temporaries, with a top-level HoistBroadcastValues() pass factory function. Reduces broadcast operand duplication via IR transformation.
Transform Module Export tilelang/transform/__init__.py	Exports HoistBroadcastValues to make the transformation pass publicly accessible.
Lowering Pipeline tilelang/engine/lower.py	Integrates HoistBroadcastValues pass into device codegen pipelines (both with and without compile), applied after Simplify.
Warp Synchronization Tests testing/python/language/test_tilelang_language_warp_sync.py	Adds test_warp_sync() verifying __syncwarp codegen and correctness; adds test_shfl_sync() verifying __shfl_sync broadcast behavior across 32 threads with assertion on final values.
Broadcast Hoisting Tests testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py	Adds parameterized test_broadcast(dtype) over float8/float16 types that validates hoisting of broadcast constants via regex matching in generated kernel source.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related issues

• [Feature Request] T.sync_warp support #1598: Changes directly implement the sync_warp builtin feature requested in this issue for warp-level thread synchronization.

Possibly related PRs

• [AMD] fix bugs in warp shuffle #790: Related modifications to warp-shuffle and synchronization builtins in tilelang/language/builtin.py with overlapping scope in primitive operations.

Suggested reviewers

• LeiWang1999
• tzj-fxz

Poem

🐰 A warp of threads now sync as one,
With shuffle dance and broadcasts done,
Through hoisted values, constants fly,
Our kernels shine, optimized high! ✨

✨ Finishing touches

• 📝 Generate docstrings

---

📜 Recent review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between cfbc49b and 4d96988.

📒 Files selected for processing (6)

• testing/python/language/test_tilelang_language_warp_sync.py
• testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py
• tilelang/engine/lower.py
• tilelang/language/builtin.py
• tilelang/transform/__init__.py
• tilelang/transform/hoist_broadcast_values.py

---

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