[Refactor] Use has_simt_copy to decide whether to insert set_max_nreg

[chengyupku]

This pull request introduces improvements to register allocation annotation and its integration with target-specific optimization passes. The main changes include adding SIMT copy detection to avoid incorrect register allocation and adjusting when the AnnotateWarpGroupRegAlloc pass is applied in the optimization pipeline.

Enhancements to register allocation annotation:

• Added a new SimtCopyDetector class in annotate_warp_group_reg_alloc.cc to detect the presence of SIMT copy operations and prevent register allocation annotation when such operations are present.
• Updated the injection logic in SetMaxNRegInjector to use the new SIMT copy detection, ensuring register allocation is only annotated when safe.

Optimization pipeline adjustments:

• Moved the application of AnnotateWarpGroupRegAlloc in phase.py to after InjectPTXAsyncCopy, and now only apply it if TMA and warp specialization are enabled for the target. [1] [2]

Code cleanup:

• Removed a call to T.no_set_max_nreg() in fp8_lighting_indexer.py as part of simplifying the kernel setup.

Summary by CodeRabbit

• Bug Fixes

  • Prevented over-aggressive register annotations by applying them only when TMA and warp specialization are enabled, reducing risk of performance regressions.
  • Updated an example kernel to rely on default register handling, improving stability across environments.

• Refactor

  • Improved detection of specialized memory copies to more accurately gate register optimization hints, leading to more predictable performance without changing user-facing behavior.

[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

Removes a kernel register directive in an example; adds SIMT copy detection (SimtCopyDetector) and integrates it to gate register-hint injection; adjusts OptimizeForTarget to conditionally invoke AnnotateWarpGroupRegAlloc after PTX async copy when warp specialization is allowed, removing its earlier unconditional call.

Changes

Cohort / File(s)	Summary of changes
Example kernel register directive removal examples/deepseek_v32/fp8_lighting_indexer.py	Deleted T.no_set_max_nreg() directive in mqa_attn_return_logits_kernel, reverting to default register allocation behavior.
SIMT copy detection and injector gating src/transform/annotate_warp_group_reg_alloc.cc	Added SimtCopyDetector (StmtExprVisitor) to detect non-global BufferStore operations; integrated Detect(producer_body) to gate SetMaxNRegInjector’s register-hint injection; preserved traversal/injection logic otherwise.
Optimize phase control-flow adjustment tilelang/engine/phase.py	Removed early unconditional AnnotateWarpGroupRegAlloc call; added conditional call after InjectPTXAsyncCopy when allow_tma_and_warp_specialized is true.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant P as OptimizeForTarget
  participant B as InjectTmaBarrier
  participant A as InjectPTXAsyncCopy
  participant W as AnnotateWarpGroupRegAlloc

  P->>B: InjectTmaBarrier()
  Note over P,B: Early call to AnnotateWarpGroupRegAlloc removed
  P->>A: InjectPTXAsyncCopy()
  alt allow_tma_and_warp_specialized == true
    P->>W: AnnotateWarpGroupRegAlloc()
    Note over P,W: Now gated and occurs later
  else
    Note over P: Skip warp-group reg alloc annotation
  end

Loading

sequenceDiagram
  autonumber
  participant I as SetMaxNRegInjector
  participant D as SimtCopyDetector
  participant PB as producer_body
  participant IR as IR Nodes

  I->>D: Detect(PB)
  D->>IR: Visit BufferStore nodes
  alt storage_scope != "global"
    D-->>I: SIMT copy present (true)
  else
    D-->>I: No SIMT copy (false)
  end
  alt detection == true
    I->>IR: Inject register hints
  else
    I->>IR: Do not inject
  end

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Poem

I twitch my whiskers at SIMT’s trace,
A detector hops through memory’s space.
Registers rest where hints are due,
And phases shift to pathways new.
With async winds and guarded gates—
I thump approval, then iterate. 🥕✨

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 title succinctly captures the main refactor by indicating that the PR uses has_simt_copy to gate set_max_nreg insertion, which reflects the core change to the register allocation logic without extraneous details.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

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

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 8a0de94 and 6a42e31.

📒 Files selected for processing (3)

• examples/deepseek_v32/fp8_lighting_indexer.py (0 hunks)
• src/transform/annotate_warp_group_reg_alloc.cc (2 hunks)
• tilelang/engine/phase.py (1 hunks)

💤 Files with no reviewable changes (1)

• examples/deepseek_v32/fp8_lighting_indexer.py

🧰 Additional context used

🧬 Code graph analysis (2)

src/transform/annotate_warp_group_reg_alloc.cc (1)

src/transform/warp_specialized_rewriter.cc (21)

• op (38-43)
• op (38-38)
• op (73-83)
• op (73-73)
• op (85-93)
• op (85-85)
• op (95-100)
• op (95-95)
• op (102-110)
• op (102-102)
• op (134-146)
• op (134-134)
• op (148-177)
• op (148-148)
• op (179-189)
• op (179-179)
• has_simt_copy_ (225-225)
• has_simt_copy_ (227-227)
• has_simt_copy_ (612-612)
• has_simt_copy (269-269)
• has_simt_copy (273-273)

tilelang/engine/phase.py (2)

src/transform/annotate_warp_group_reg_alloc.cc (2)

• AnnotateWarpGroupRegAlloc (179-185)
• AnnotateWarpGroupRegAlloc (179-179)

tilelang/transform/__init__.py (1)

• AnnotateWarpGroupRegAlloc (204-216)

⏰ 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). (2)

• GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
• GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.0 (on self-hosted-amd)

🔇 Additional comments (2)

tilelang/engine/phase.py (1)

208-209: LGTM! Conditional application aligns with warp specialization enablement.

The conditional invocation of AnnotateWarpGroupRegAlloc correctly gates the pass on allow_tma_and_warp_specialized, matching where the WarpSpecialized pass is applied (lines 133-137). Placing it after InjectPTXAsyncCopy also aligns with the comment indicating PTX async copies must be injected before register allocation annotation.

src/transform/annotate_warp_group_reg_alloc.cc (1)

137-137: Verify injection logic covers both producer and consumer
SimtCopyDetector::Detect is only invoked on producer_body. After updating the detector to eliminate false positives, manually confirm that set_max_nreg is injected when and only when each branch—producer and consumer—actually lacks SIMT copies.

---

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: 0

🧹 Nitpick comments (1)

src/transform/annotate_warp_group_reg_alloc.cc (1)

62-81: Consider using scope.rank for more robust scope checking.

Line 74 uses string comparison (scope.to_string() != "global") to detect non-global stores. The relevant code in warp_specialized_rewriter.cc (lines 147-176) uses scope.rank == StorageRank::kShared for similar checks, which is more type-safe and less fragile than string comparison.

Apply this diff to use rank-based comparison:

 private:
   void VisitStmt_(const BufferStoreNode *op) final {
     auto scope =
         runtime::StorageScope::Create(GetPtrStorageScope(op->buffer->data));
-    if (scope.to_string() != "global") {
+    if (scope.rank != runtime::StorageRank::kGlobal) {
       has_simt_copy_ = true;
     }
     StmtExprVisitor::VisitStmt_(op);
   }

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 117f2b8 and 26674b2.

📒 Files selected for processing (3)

• examples/deepseek_v32/fp8_lighting_indexer.py (0 hunks)
• src/transform/annotate_warp_group_reg_alloc.cc (2 hunks)
• tilelang/engine/phase.py (1 hunks)

💤 Files with no reviewable changes (1)

• examples/deepseek_v32/fp8_lighting_indexer.py

🧰 Additional context used

🧬 Code graph analysis (2)

tilelang/engine/phase.py (2)

src/transform/annotate_warp_group_reg_alloc.cc (2)

• AnnotateWarpGroupRegAlloc (179-185)
• AnnotateWarpGroupRegAlloc (179-179)

tilelang/transform/__init__.py (1)

• AnnotateWarpGroupRegAlloc (204-216)

src/transform/annotate_warp_group_reg_alloc.cc (1)

src/transform/warp_specialized_rewriter.cc (21)

• op (38-43)
• op (38-38)
• op (73-83)
• op (73-73)
• op (85-93)
• op (85-85)
• op (95-100)
• op (95-95)
• op (102-110)
• op (102-102)
• op (134-146)
• op (134-134)
• op (148-177)
• op (148-148)
• op (179-189)
• op (179-179)
• has_simt_copy_ (225-225)
• has_simt_copy_ (227-227)
• has_simt_copy_ (612-612)
• has_simt_copy (269-269)
• has_simt_copy (273-273)

⏰ 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: build-test-metal

🔇 Additional comments (3)

src/transform/annotate_warp_group_reg_alloc.cc (2)

137-137: LGTM! Proper integration of SIMT copy detection.

The replacement of placeholder logic with SimtCopyDetector::Detect(producer_body) correctly implements the conditional injection of register hints based on actual SIMT copy detection. This ensures that register allocation annotations are only applied when safe (i.e., when no SIMT copies are present).

---

62-81: Confirm broadened SIMT copy detection is intentional
The new SimtCopyDetector marks all non-global stores as SIMT copies, unlike existing rewriters that only flag producer stores from global loads—validate this wider scope won’t disable valid register allocation annotations or adjust logic/tests if needed.

tilelang/engine/phase.py (1)

207-208: LGTM! Properly gates register allocation annotation.

The conditional invocation of AnnotateWarpGroupRegAlloc ensures that register hints are only injected when TMA and warp specialization are enabled for the target. This aligns with the PR objective to make register allocation conditional and prevents unnecessary annotation on unsupported targets.

[chengyupku]

pre-commit.ci run

[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 cd821b9 and 937405a.

📒 Files selected for processing (5)

• .github/workflows/cuda-ci.yml (1 hunks)
• examples/deepseek_v32/fp8_lighting_indexer.py (0 hunks)
• requirements-test.txt (1 hunks)
• src/transform/annotate_warp_group_reg_alloc.cc (2 hunks)
• tilelang/engine/phase.py (1 hunks)

💤 Files with no reviewable changes (1)

• examples/deepseek_v32/fp8_lighting_indexer.py

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

• requirements-test.txt

🧰 Additional context used

🧬 Code graph analysis (2)

src/transform/annotate_warp_group_reg_alloc.cc (2)

src/transform/warp_specialized_rewriter.cc (21)

• op (38-43)
• op (38-38)
• op (73-83)
• op (73-73)
• op (85-93)
• op (85-85)
• op (95-100)
• op (95-95)
• op (102-110)
• op (102-102)
• op (134-146)
• op (134-134)
• op (148-177)
• op (148-148)
• op (179-189)
• op (179-179)
• has_simt_copy_ (225-225)
• has_simt_copy_ (227-227)
• has_simt_copy_ (612-612)
• has_simt_copy (269-269)
• has_simt_copy (273-273)

src/transform/warp_specialized_rewriter.h (1)

• VisitStmt_ (77-90)

tilelang/engine/phase.py (2)

src/transform/annotate_warp_group_reg_alloc.cc (2)

• AnnotateWarpGroupRegAlloc (179-185)
• AnnotateWarpGroupRegAlloc (179-179)

tilelang/transform/__init__.py (1)

• AnnotateWarpGroupRegAlloc (204-216)

⏰ 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). (2)

• GitHub Check: build-test-amd
• GitHub Check: format-check

🔇 Additional comments (1)

.github/workflows/cuda-ci.yml (1)

118-125: Single-worker pytest fits the GPU runners.

Running the example and test suites with -n 1 matches our single-GPU self-hosted machines and avoids the contention we were seeing. LGTM.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Over-broad SIMT copy detection disables register hints everywhere.

SimtCopyDetector now sets has_simt_copy_ for every BufferStore whose scope is anything other than "global". In warp-specialized producers, the stores almost always target shared memory—even when they're the result of plain computation rather than a SIMT copy from global. As a consequence, SetMaxNRegInjector will see has_simt_copy == true for nearly all producer branches and skip re-inserting set_max_nreg, effectively turning the pass off and regressing performance.

Please tighten the detection to match actual SIMT copies (e.g., shared stores fed solely by global loads, mirroring the logic already used in WarpSpecializedRewriter) instead of treating every non-global store as a copy.

🤖 Prompt for AI Agents

In src/transform/annotate_warp_group_reg_alloc.cc around lines 71-78, the
detector currently marks any non-"global" BufferStore as a SIMT copy; change it
to only mark has_simt_copy_ when the store writes to non-global storage AND the
stored value is composed solely of loads from global buffers (mirror the
WarpSpecializedRewriter logic). Implement a small recursive helper that walks
op->value and returns true only if every BufferLoad encountered has
GetPtrStorageScope(load->buffer->data).to_string() == "global" (allow other pure
arithmetic/casts between those loads), and use that helper in the VisitStmt_
check; otherwise do not set has_simt_copy_. Ensure you still call the base
StmtExprVisitor::VisitStmt_(op).

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

requirements-test.txt (1)

31-31: Move triton into a GPU-only requirements file and cap its version

• Remove triton>=3.4.0 from requirements-test.txt and add
  triton>=3.4,<3.5 to requirements-test-cuda.txt (installed only under contains(matrix.runner.toolkit, 'CUDA')).
• If you also want ROCm/Metal support, guard similarly in CI or create separate req files.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 937405a and 8a0de94.

📒 Files selected for processing (1)

• requirements-test.txt (1 hunks)

🧰 Additional context used

🪛 LanguageTool

requirements-test.txt

[grammar] ~31-~31: There might be a mistake here.
Context: ...sts scipy tabulate tornado triton>=3.4.0 wheel

(QB_NEW_EN)

⏰ 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: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)