[Bugfix] Fix missing reg alloc in custom warp specialization

[chengyupku]

This pull request makes targeted changes to the SetMaxNRegInjector logic in annotate_warp_group_reg_alloc.cc, refining how register allocation annotations are handled. The main focus is to adjust when and how set_max_nreg calls are injected and removed, specifically tightening the conditions under which register hints are applied.

Refinements to register allocation annotation:

• The check for removing set_max_nreg calls in VisitStmt_ is now limited to only no_set_max_nreg, meaning original set_max_nreg calls are no longer removed at this stage.

• The logic for injecting register hints has changed: register hints are now only injected if both dec_reg and inc_reg are exactly zero and there is no SIMT copy detected, simplifying and restricting the conditions for annotation. The register numbers are now hardcoded to 240 and 24 for inc_reg and dec_reg respectively.

Summary by CodeRabbit

• Refactor
  • Internal optimization adjustments to register allocation handling. No user-visible changes or new features introduced.

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

Walkthrough

Modification to warp group register allocation annotation logic that restricts set_max_nreg handling to only the no_set_max_nreg path, tightens the injection condition to require exactly zero register adjustments and no SIMT copy detection, and replaces computed register hint constants with fixed values.

Changes

Cohort / File(s)	Summary
Register allocation injection refinement src/transform/annotate_warp_group_reg_alloc.cc	Removed set_max_nreg handling in Evaluate visitor; tightened injection guard from dec_reg >= 0 && inc_reg >= 0 to dec_reg == 0 && inc_reg == 0 plus SIMT copy absence check; fixed inc_reg_num to 240 and dec_reg_num to 24 instead of deriving from hints

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

The change modifies register allocation logic with stricter conditions and fixed constants, requiring review context on SIMT copy detection semantics, register adjustment implications, and verification that the hardcoded values (240, 24) are appropriate replacements for dynamic computation.

Possibly related PRs

• [Refactor] Use has_simt_copy to decide whether to insert set_max_nreg #982: Modifies the same file (annotate_warp_group_reg_alloc.cc) and introduces SimtCopyDetector for SIMT-copy-aware set_max_nreg handling, directly complementing this PR's stricter injection guards.

Poem

🐰 Register hints simplified, their dance now bound,
No SIMT copies crossing hallowed ground,
Two-forty and twenty-four, constants now stand still,
Where injection once computed, now precision fills.
Tighter gates, simpler paths, allocation refined. ✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.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 pull request title "[Bugfix] Fix missing reg alloc in custom warp specialization" is clearly related to the changeset, which modifies register allocation annotation behavior in the warp group specialization code. The title correctly identifies this as a bugfix and references the core domain area (register allocation in warp specialization). However, the phrase "missing reg alloc" is somewhat imprecise relative to the actual changes, which involve tightening injection conditions, restricting what gets removed, and hardcoding constants—refinements to existing functionality rather than additions of missing features. While the title could be more specific about the nature of the fix, it communicates a meaningful change and is not misleading or off-topic.

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

[coderabbitai]

Actionable comments posted: 1

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between bc37ea6 and f60e3fe.

📒 Files selected for processing (1)

• src/transform/annotate_warp_group_reg_alloc.cc (2 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

src/transform/annotate_warp_group_reg_alloc.cc (1)

tilelang/language/builtin.py (1)

• no_set_max_nreg (174-177)

⏰ 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

🔇 Additional comments (1)

src/transform/annotate_warp_group_reg_alloc.cc (1)

135-145: Fix the misleading comment; the behavioral change appears intentional.

The comment at line 135 is indeed misleading. It states "Only inject if we have valid register hints" but the condition on line 138 actually injects only when both dec_reg and inc_reg are exactly zero—meaning no explicit set_max_nreg() calls were found in the code.

Based on the code structure:

• nreg_ is initialized to {0, 0} (lines 56-57)
• The SetMaxNRegCollector populates nreg_ from explicit set_max_nreg() calls in user code
• When both values remain zero, it means the user provided no hints, so hardcoded defaults (240, 24) are injected
• This behavior appears intentional: only inject defaults when no user-provided hints exist

Apply this fix to clarify the comment:

-      // Only inject if we have valid register hints and no SIMT copy
+      // Only inject default register hints when no explicit hints were provided and no SIMT copy is detected
       bool has_simt_copy = SimtCopyDetector::Detect(producer_body);

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Update the misleading comment.

The comment on lines 99-100 claims to "Remove the original set_max_nreg calls" but the code now only removes no_set_max_nreg() calls. Original set_max_nreg() calls are preserved in this version.

Apply this diff to correct the comment:

       if (call->op.same_as(no_set_max_nreg())) {
-        // Remove the original set_max_nreg calls as they will be re-inserted
-        // at appropriate locations
+        // Remove no_set_max_nreg() calls to allow injection of default hints
+        // if no explicit register hints were provided
         return Evaluate(0);
       }

📝 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

	if (call->op.same_as(no_set_max_nreg())) {
	// Remove the original set_max_nreg calls as they will be re-inserted
	// at appropriate locations
	return Evaluate(0);
	if (call->op.same_as(no_set_max_nreg())) {
	// Remove no_set_max_nreg() calls to allow injection of default hints
	// if no explicit register hints were provided
	return Evaluate(0);

🤖 Prompt for AI Agents

In src/transform/annotate_warp_group_reg_alloc.cc around lines 98 to 101, the
comment incorrectly states that "Remove the original set_max_nreg calls" while
the code only removes no_set_max_nreg() calls; update the comment to accurately
state that only no_set_max_nreg() calls are removed (or similar wording) so it
reflects the current behavior, e.g., change the comment to indicate removal of
no_set_max_nreg() calls which will be re-inserted later.