[Language] Expose T.get_warp_idx_sync and T.shuffle_elect for efficient thread election

[LeiWang1999]

This pull request introduces new thread and warp identification intrinsics for CUDA backends, making it easier to query thread indices, warp indices, and warp group indices in TileLang programs. The changes span the TileLang Python API, CUDA code generation, and the underlying C++ operators and device intrinsics.

New CUDA thread and warp index intrinsics:

Python API additions:

• Added four new helper functions to tilelang/language/builtin.py: get_lane_idx, get_warp_idx_sync, get_warp_idx, and get_warp_group_idx, each returning the corresponding thread/warp/group index as a PrimExpr. These functions provide clear documentation and usage examples, and are lowered to new CUDA intrinsics.

C++ operator and codegen support:

• Registered new TIR built-in operators for each intrinsic in src/op/builtin.cc and declared them in src/op/builtin.h, marking them as pure (no side effects) and requiring zero inputs. [1] [2]
• Updated the CUDA codegen visitor in src/target/codegen_cuda.cc to emit calls to the new CUDA helpers for these operators, ensuring correct code generation.

Device-side CUDA implementations:

• Implemented the new helpers in src/tl_templates/cuda/intrin.h, each forwarding to the corresponding CUTLASS canonical index function, providing accurate thread/warp/group identification on supported hardware. [1] [2]

Summary by CodeRabbit

• New Features

  • Added built-ins to query thread hierarchy: get_lane_idx, get_warp_idx_sync, get_warp_idx, get_warp_group_idx.
  • Added shuffle_elect(thread_extent) to elect a single lane in a logical thread group.

• Bug Fixes / Enhancements

  • Support negative indices in JIT adapter list handling (-len..-1 now valid).

• Tests

  • Added CUDA-targeted tests validating lane/warp/group index helpers and shuffle_elect.

[coderabbitai]

Walkthrough

Adds four new TileLang (TL) intrinsics for lane/warp queries (get_lane_idx, get_warp_idx_sync, get_warp_idx, get_warp_group_idx) across C++ builtin registration, header declarations, CUDA codegen emission, CUDA device intrinsics, Python wrappers, tests, and a small index-handling tweak in a JIT adapter.

Changes

Cohort / File(s)	Summary
TL builtin op registration src/op/builtin.cc	Registers four zero-input TL ops: get_lane_idx, get_warp_idx_sync, get_warp_idx, get_warp_group_idx, all with TCallEffectKind = kPure.
TL builtin op declarations src/op/builtin.h	Adds TVM_DLL const Op & accessors for the four new intrinsics in tvm::tl.
CUDA codegen handlers src/target/codegen_cuda.cc	Extends CallNode handling to emit tl::get_lane_idx(...), tl::get_warp_idx_sync(...), tl::get_warp_idx(...), and tl::get_warp_group_idx(...) branches with arg-count checks.
CUDA device intrinsics src/tl_templates/cuda/intrin.h	Introduces inline implementations: tl::get_lane_idx, tl::get_warp_idx_sync, tl::get_warp_idx, tl::get_warp_group_idx and internal helpers/defaults; adjusts namespace boundary.
Python TL builtins tilelang/language/builtin.py	Adds Python wrappers lowering to new TL intrinsics: get_lane_idx, get_warp_idx_sync, get_warp_idx, get_warp_group_idx, and shuffle_elect(thread_extent).
Tests testing/python/language/test_tilelang_language_get_warp_info.py	New CUDA-only tests and runners validating lane id, warp idx (sync/non-sync), warp group idx, and shuffle_elect behavior across configurations.
JIT adapter index handling tilelang/jit/adapter/base.py	Adjusts negative-index handling in _legalize_result_idx: allows -len(params) and converts negative indices to positive equivalents in-place.

Sequence Diagram(s)

sequenceDiagram
    autonumber
    participant Py as Python API (tilelang.language.builtin)
    participant TL as TL Op Registry
    participant IR as Lowering / CallNode
    participant CG as CUDA Codegen (codegen_cuda.cc)
    participant CU as CUDA Device Helpers (tl_templates/cuda/intrin.h)

    Py->>TL: call get_lane_idx()/get_warp_idx(_sync)/get_warp_group_idx()
    TL->>IR: produce Call to tl.* builtin (0 args, kPure)
    IR->>CG: lower/visit CallNode for tl.* intrinsic
    CG->>CU: emit tl::get_* device call (with optional params)
    CU-->>Py: runtime returns lane/warp/group index value

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• [CI][AMD] Add AMD GPU CI and fix some related bugs #694 — Adds other TL builtin intrinsics in src/op/builtin.cc, similar registration pattern and potential conflicts in op ordering.
• [Bugfix] Add missing FP8 header include #752 — Modifies CUDA intrinsics/header layout and warp-related helpers; likely intersects with intrin.h changes here.
• [Refactor] Refactor Pass InjectFenceProxy and expose some warp group primitives in frontend #977 — Adds TL/CUDA intrinsics and codegen entries touching the same files (builtin.*, codegen_cuda.cc, intrin.h), relevant for integration and overlap.

Suggested reviewers

• LeiWang1999

Poem

A hop, a patch, in CUDA light,
I count my lanes from left to right.
Four new hops, all tidy and neat,
Sync or free — the warps we meet.
Carrots compiled, the threads all greet. 🥕🐇

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ 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 highlights two of the newly exposed intrinsics but omits other key additions such as get_lane_idx, get_warp_idx, and get_warp_group_idx, making it only partially representative of the overall changeset.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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

[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]

Actionable comments posted: 0

🧹 Nitpick comments (1)

tilelang/language/builtin.py (1)

350-371: Document the compile-time constant requirement.

The thread_extent parameter is used as a C++ template argument in the codegen (tl::tl_shuffle_elect<thread_extent>()), which requires a compile-time constant. However, the current documentation doesn't explicitly state this constraint.

Consider adding a note to clarify this requirement:

 def shuffle_elect(thread_extent: int) -> PrimExpr:
     """Elect exactly one lane within a logical thread group.
 
     Parameters
     ----------
     thread_extent : int
         Size (in threads) of the group in which a single lane should be elected.
         Passing 0 elects a single lane in the entire thread block.
+        Must be a compile-time constant (integer literal).

Alternatively, you could add runtime validation:

def shuffle_elect(thread_extent: int) -> PrimExpr:
    """..."""
    if not isinstance(thread_extent, int):
        raise TypeError(f"thread_extent must be an integer literal, got {type(thread_extent)}")
    return tir.call_intrin("bool", tir.op.Op.get("tl.tl_shuffle_elect"), thread_extent)

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 4a229dd and 9bffa4a.

📒 Files selected for processing (5)

• src/op/builtin.cc (1 hunks)
• src/op/builtin.h (1 hunks)
• src/target/codegen_cuda.cc (1 hunks)
• src/tl_templates/cuda/intrin.h (2 hunks)
• tilelang/language/builtin.py (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (5)

tilelang/language/builtin.py (1)

tilelang/language/tir/op.py (1)

• call_intrin (119-144)

src/op/builtin.cc (1)

tilelang/language/builtin.py (4)

• get_lane_idx (283-297)
• get_warp_idx_sync (300-314)
• get_warp_idx (317-330)
• get_warp_group_idx (333-347)

src/op/builtin.h (1)

tilelang/language/builtin.py (4)

• get_lane_idx (283-297)
• get_warp_idx_sync (300-314)
• get_warp_idx (317-330)
• get_warp_group_idx (333-347)

src/tl_templates/cuda/intrin.h (2)

src/tl_templates/cuda/gemm_sm90.h (1)

• tl (231-385)

tilelang/language/builtin.py (4)

• get_lane_idx (283-297)
• get_warp_idx_sync (300-314)
• get_warp_idx (317-330)
• get_warp_group_idx (333-347)

src/target/codegen_cuda.cc (1)

tilelang/language/builtin.py (4)

• get_lane_idx (283-297)
• get_warp_idx_sync (300-314)
• get_warp_idx (317-330)
• get_warp_group_idx (333-347)

⏰ 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: format-check

🔇 Additional comments (5)

src/tl_templates/cuda/intrin.h (1)

13-21: LGTM! Well-structured additions.

The four new index helpers are correctly placed outside the SM90+ conditional block, making them available across all CUDA architectures. The delegation to CUTLASS canonical functions is appropriate, and the namespace structure is preserved correctly.

src/op/builtin.h (1)

361-391: LGTM!

The four new intrinsic declarations follow the existing pattern and are well-documented. The positioning between warpgroup_wait() and wait_wgmma() is logical and maintains consistency with the file structure.

src/target/codegen_cuda.cc (1)

1971-1978: LGTM!

The CUDA codegen correctly emits the tl::get_*() calls for the new intrinsics. The implementation is straightforward and consistent with other zero-argument intrinsic patterns in the codebase.

src/op/builtin.cc (1)

221-239: LGTM! Marking intrinsics as pure is appropriate.

All four intrinsics are correctly registered with zero inputs and kPure effect kind. While get_warp_idx_sync uses __shfl_sync under the hood (which has implicit warp synchronization), marking it as kPure is reasonable because:

• The function is deterministic (same inputs → same outputs)
• It doesn't modify observable state or memory
• The warp-level synchronization is implicit in CUDA's execution model
• This aligns with how similar warp-level operations are typically classified

tilelang/language/builtin.py (1)

283-347: LGTM! Well-documented index query functions.

The four new index query functions (get_lane_idx, get_warp_idx_sync, get_warp_idx, get_warp_group_idx) are well-implemented with comprehensive docstrings, clear examples, and detailed implementation notes explaining the underlying CUDA/CUTLASS mechanisms.

[LyricZhao]

Some CUDA kernels treat, e.g. 5/6 warps, as a warp group. So I suggest the index calculation can be more general. For example, give warp size (AMD 64 threads per warp) and warp group size with an arch-specific default value.

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tilelang/jit/adapter/base.py (1)

33-39: Fix off-by-one bounds checks and error messages in _legalize_result_idx

• In the int branch (around line 24), change the upper bound check to >= len(params) and update the message:

  -    if result_idx > len(params) or result_idx < -len(params):
  +    if result_idx >= len(params) or result_idx < -len(params):
  ...
  -        f"result_idx should be an integer between {-len(params) - 1} and {len(params) - 1}"
  +        f"result_idx should be an integer between {-len(params)} and {len(params) - 1}"

• In the list branch (around line 33), adjust only the message:

  -        f"result_idx should be an integer between {-len(params) - 1} and {len(params) - 1}"
  +        f"result_idx should be an integer between {-len(params)} and {len(params) - 1}"

🧹 Nitpick comments (3)

src/tl_templates/cuda/intrin.h (1)

51-60: Consider parameterizing warp size in tl_shuffle_elect for consistency

Elsewhere you allow non-32 warp sizes (HIP/AMD 64). tl_shuffle_elect still divides by 32. Consider using detail::default_warp_size() (or a template param) instead of 32 to keep behavior consistent across vendors.

tilelang/language/builtin.py (1)

407-429: Enforce thread_extent be a compile-time constant
Add a check in shuffle_elect (tilelang/language/builtin.py) to require thread_extent is a tir.IntImm (e.g. assert isinstance(thread_extent, tir.IntImm)), so non-literal arguments error out before C++ codegen.

testing/python/language/test_tilelang_language_get_warp_info.py (1)

160-208: LGTM: Comprehensive test coverage.

The test suite covers both default configurations and custom parameter overrides for all five intrinsics. The @tilelang.testing.requires_cuda decorator correctly ensures tests only run on CUDA-capable hardware.

Consider adding brief docstrings to document what each test validates, especially for tests like test_shuffle_elect_block_leader where the behavior being tested (electing block leader with thread_extent=0) may not be immediately obvious.

Example:

@tilelang.testing.requires_cuda
def test_shuffle_elect_block_leader():
    """Test shuffle_elect with thread_extent=0 elects a single block-wide leader."""
    run_shuffle_elect(num_threads=128, thread_extent=0)

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 9bffa4a and cebc6ff.

📒 Files selected for processing (7)

• src/op/builtin.cc (1 hunks)
• src/op/builtin.h (1 hunks)
• src/target/codegen_cuda.cc (1 hunks)
• src/tl_templates/cuda/intrin.h (2 hunks)
• testing/python/language/test_tilelang_language_get_warp_info.py (1 hunks)
• tilelang/jit/adapter/base.py (1 hunks)
• tilelang/language/builtin.py (2 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• src/op/builtin.cc

🧰 Additional context used

🧬 Code graph analysis (6)

tilelang/jit/adapter/base.py (1)

tilelang/jit/kernel.py (1)

• params (466-467)

src/op/builtin.h (1)

tilelang/language/builtin.py (4)

• get_lane_idx (297-319)
• get_warp_idx_sync (322-343)
• get_warp_idx (346-367)
• get_warp_group_idx (370-404)

src/tl_templates/cuda/intrin.h (1)

tilelang/language/builtin.py (4)

• get_lane_idx (297-319)
• get_warp_idx_sync (322-343)
• get_warp_idx (346-367)
• get_warp_group_idx (370-404)

testing/python/language/test_tilelang_language_get_warp_info.py (3)

tilelang/utils/target.py (1)

• check_hip_availability (33-43)

tilelang/language/builtin.py (5)

• get_lane_idx (297-319)
• get_warp_idx_sync (322-343)
• get_warp_idx (346-367)
• get_warp_group_idx (370-404)
• shuffle_elect (407-428)

tilelang/jit/adapter/base.py (1)

• get_kernel_source (51-52)

src/target/codegen_cuda.cc (1)

tilelang/language/builtin.py (4)

• get_lane_idx (297-319)
• get_warp_idx_sync (322-343)
• get_warp_idx (346-367)
• get_warp_group_idx (370-404)

tilelang/language/builtin.py (2)

src/op/builtin.h (1)

• tvm (13-506)

tilelang/language/tir/op.py (1)

• call_intrin (119-144)

🪛 Ruff (0.13.3)

testing/python/language/test_tilelang_language_get_warp_info.py

135-135: Avoid specifying long messages outside the exception class

(TRY003)

---

143-143: Avoid specifying long messages outside the exception class

(TRY003)

tilelang/language/builtin.py

25-25: Avoid specifying long messages outside the exception class

(TRY003)

---

401-402: Avoid specifying long messages outside the exception class

(TRY003)

🔇 Additional comments (13)

src/tl_templates/cuda/intrin.h (1)

13-60: Defaults and helpers look good

Arch-aware defaults, argument sanitization, and linear index math are fine.

src/op/builtin.h (1)

361-392: LGTM: Op declarations match codegen and Python API

New builtin op declarations align with emitters and wrappers.

tilelang/language/builtin.py (5)

14-26: LGTM: argument normalization helper

Accepts int/PrimExpr, returns IntImm or None; straightforward.

---

297-320: get_lane_idx wrapper is correct

Lowers with optional warp_size; docs match behavior.

---

322-344: get_warp_idx_sync wrapper is correct

Matches intended lowering; optional warp_size supported.

---

346-368: get_warp_idx wrapper is correct

Consistent with codegen/device side.

---

370-405: get_warp_group_idx wrapper enforces valid args; good

Argument gating is sensible; lowering matches codegen expectation.

testing/python/language/test_tilelang_language_get_warp_info.py (6)

1-10: LGTM: Clean imports and architecture-aware setup.

The imports are well-organized, and the HIP availability check enables cross-platform testing. The default warps per group value (4) aligns with NVIDIA's typical warp-group configuration.

---

12-22: LGTM: Architecture-aware parameter resolution.

The helper functions correctly provide architecture-specific defaults (32 for NVIDIA, 64 for AMD), allowing tests to run across different GPU vendors while supporting explicit overrides.

---

24-86: LGTM: Well-structured kernel builders.

The JIT-decorated kernel builders follow a consistent pattern and correctly exercise each intrinsic (T.get_lane_idx, T.get_warp_idx_sync, T.get_warp_idx, T.get_warp_group_idx, T.shuffle_elect). The parameter forwarding is correct.

---

89-119: LGTM: Reference calculations are correct.

The reference tensors for lane and warp index tests correctly implement the expected semantics:

• Lane index: thread_id % warp_size
• Warp index (sync and non-sync): thread_id // warp_size

The use of torch.testing.assert_close for validation is appropriate.

---

122-138: LGTM: Warp group index calculation is correct.

The reference calculation thread_id // (warp_size * warps_per_group) correctly implements the expected warp group indexing semantics. The defensive check for threads_per_group <= 0 is reasonable, though unlikely to trigger in normal usage.

---

210-211: LGTM: Standard test runner setup.

The main block correctly invokes the test framework, and the commented debug line is useful for manual testing during development.

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

Add missing include for tl intrinsics to avoid undefined references

Emission for tl::get_lane_idx / tl::get_warp_idx(_sync) / tl::get_warp_group_idx looks correct and arity checks are fine. However, there’s no include for the tl intrinsics header; this can fail to compile when these symbols aren’t already brought in indirectly.

Include the header alongside other tl headers in Finish():

   decl_stream << "#include <tl_templates/cuda/gemm.h>\n";
   if (enable_sparse_gemm_) {
     decl_stream << "#include <tl_templates/cuda/gemm_sp.h>\n";
   }
   decl_stream << "#include <tl_templates/cuda/copy.h>\n";
   decl_stream << "#include <tl_templates/cuda/reduce.h>\n";
   decl_stream << "#include <tl_templates/cuda/ldsm.h>\n";
   decl_stream << "#include <tl_templates/cuda/threadblock_swizzle.h>\n";
   decl_stream << "#include <tl_templates/cuda/debug.h>\n";
+  decl_stream << "#include <tl_templates/cuda/intrin.h>\n";
   decl_stream << "#ifdef ENABLE_BF16\n";

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

Suggested change

	} else if (op->op.same_as(tl::get_lane_idx())) {
	ICHECK_LE(op->args.size(), 1)
	<< "tl.get_lane_idx expects at most one argument <warp_size>.";
	os << "tl::get_lane_idx(";
	if (!op->args.empty()) {
	os << PrintExpr(op->args[0]);
	}
	os << ")";
	} else if (op->op.same_as(tl::get_warp_idx_sync())) {
	ICHECK_LE(op->args.size(), 1)
	<< "tl.get_warp_idx_sync expects at most one argument <warp_size>.";
	os << "tl::get_warp_idx_sync(";
	if (!op->args.empty()) {
	os << PrintExpr(op->args[0]);
	}
	os << ")";
	} else if (op->op.same_as(tl::get_warp_idx())) {
	ICHECK_LE(op->args.size(), 1)
	<< "tl.get_warp_idx expects at most one argument <warp_size>.";
	os << "tl::get_warp_idx(";
	if (!op->args.empty()) {
	os << PrintExpr(op->args[0]);
	}
	os << ")";
	} else if (op->op.same_as(tl::get_warp_group_idx())) {
	ICHECK_LE(op->args.size(), 2)
	<< "tl.get_warp_group_idx expects <warp_size, warps_per_group>.";
	os << "tl::get_warp_group_idx(";
	for (size_t i = 0; i < op->args.size(); ++i) {
	if (i != 0) {
	os << ", ";
	}
	os << PrintExpr(op->args[i]);
	}
	os << ")";
	// In Finish(), alongside the other TL headers:
	decl_stream << "#include <tl_templates/cuda/gemm.h>\n";
	if (enable_sparse_gemm_) {
	decl_stream << "#include <tl_templates/cuda/gemm_sp.h>\n";
	}
	decl_stream << "#include <tl_templates/cuda/copy.h>\n";
	decl_stream << "#include <tl_templates/cuda/reduce.h>\n";
	decl_stream << "#include <tl_templates/cuda/ldsm.h>\n";
	decl_stream << "#include <tl_templates/cuda/threadblock_swizzle.h>\n";
	decl_stream << "#include <tl_templates/cuda/debug.h>\n";
	decl_stream << "#include <tl_templates/cuda/intrin.h>\n";
	decl_stream << "#ifdef ENABLE_BF16\n";

🤖 Prompt for AI Agents

In src/target/codegen_cuda.cc around lines 1971 to 2005, the code emits calls to
tl::get_lane_idx, tl::get_warp_idx(_sync) and tl::get_warp_group_idx but the
translation layer intrinsics header isn’t being included in Finish(), which can
cause undefined reference/compile errors; add an #include for the tl intrinsics
header (the header that declares
tl::get_lane_idx/get_warp_idx/get_warp_group_idx) alongside the other tl headers
in the Finish() function so the symbols are declared when these calls are
emitted.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Remove unreachable dead code.

Lines 153-154 are unreachable because:

• Line 142 raises an exception if thread_extent < 0
• Line 149 handles thread_extent == 0
• Line 151 handles thread_extent > 0

The else branch can never execute.

Apply this diff to remove the dead code:

     if thread_extent == 0:
         mask = indices == 0
     elif thread_extent > 0:
         mask = (indices % thread_extent) == 0
-    else:
-        mask = torch.zeros_like(indices, dtype=torch.bool)
     ref = mask.to(dtype=A.dtype, device=A.device)

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

Suggested change

	def run_shuffle_elect(num_threads: int = 128, thread_extent: int = 64):
	if thread_extent < 0:
	raise ValueError("thread_extent must be non-negative.")
	kernel = _shuffle_elect_kernel(num_threads, thread_extent)
	A = kernel()
	print(kernel.get_kernel_source())
	print(A)
	indices = torch.arange(num_threads, device=A.device, dtype=torch.int64)
	if thread_extent == 0:
	mask = indices == 0
	elif thread_extent > 0:
	mask = (indices % thread_extent) == 0
	else:
	mask = torch.zeros_like(indices, dtype=torch.bool)
	ref = mask.to(dtype=A.dtype, device=A.device)
	torch.testing.assert_close(A.cpu(), ref.cpu())
	return A
	def run_shuffle_elect(num_threads: int = 128, thread_extent: int = 64):
	if thread_extent < 0:
	raise ValueError("thread_extent must be non-negative.")
	kernel = _shuffle_elect_kernel(num_threads, thread_extent)
	A = kernel()
	print(kernel.get_kernel_source())
	print(A)
	indices = torch.arange(num_threads, device=A.device, dtype=torch.int64)
	if thread_extent == 0:
	mask = indices == 0
	elif thread_extent > 0:
	mask = (indices % thread_extent) == 0
	ref = mask.to(dtype=A.dtype, device=A.device)
	torch.testing.assert_close(A.cpu(), ref.cpu())
	return A

🧰 Tools

🪛 Ruff (0.13.3)

143-143: Avoid specifying long messages outside the exception class

(TRY003)

🤖 Prompt for AI Agents

In testing/python/language/test_tilelang_language_get_warp_info.py around lines
141 to 157, remove the unreachable final else branch (lines 153-154) that sets
mask to torch.zeros_like(...) because thread_extent < 0 is already prevented by
the early ValueError and the other branches handle ==0 and >0; simply delete
that else block so mask is only set in the existing thread_extent == 0 and
thread_extent > 0 branches (or convert the second branch to an else if you
prefer), leaving mask always defined before building ref.

[LeiWang1999]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

tilelang/src/tl_templates/cuda/intrin.h

Lines 100 to 103 in 65b77aa

	return __shfl_sync(0xffffffff, // full warp mask
	(threadIdx.x / 32) %
	(thread_extent / 32), // warp index within group
	0 // take the value from lane 0

Handle non-multiple warp sizes in tl_shuffle_elect

The general branch of tl_shuffle_elect divides by (thread_extent / 32) without validating that thread_extent is at least one warp and divisible by 32. When a caller passes any thread_extent in the range 1–31, the expression instantiates the template with a zero denominator, causing a compile-time error. For thread extents such as 48 or 96, integer division truncates to 1 or 3 warps, so the code elects one lane per physical warp instead of one per 48- or 96-thread group, yielding two or three leaders per logical group. The Python API exposes shuffle_elect without documenting these constraints, so valid-looking inputs produce either build failures or incorrect leaders.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.