[Refactor] Recycle T.unroll(explicit=True) for early explicit unrolling

[Yongqi-Zhuo]

This PR refactors the pragma_unroll_explicit behavior: by default, it is off now. Nothing in the compiler needs explicit unrolling to function correctly. And cicc will unroll the code for you with #pragma unroll, which does not need explicit unrolling at all. So I decide to leverage T.unroll(explicit=True) for program pre-processing to facilitate pipeline planning: the software pipeliner and warp-specializer want to see IR nodes flattened for easier scheduling. So I added an UnrollLoop pass before the scheduling in the pass pipeline.

Summary

• Disabled explicit unrolling by default.
• Reused T.unroll(explicit=True) during preprocessing.
• Added UnrollLoop before pipeline planning.
• Updated loop annotation handling and unroll validation.
• Updated CUDA pipeline and unroll tests.
• Adjusted TMA tests for emitted loops.

C++ style / lint notes

• The PR does not modify rules in docs/developer_guide/cpp_style.md.
• The C++ API Style Audit (warning only) CI step remains relevant.
• No correctness or build issues are identified.
• TLCPP003/TLCPP004 findings remain advisory.

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

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 949e9e72-72d6-4d4c-855d-95c0946ebf95

📥 Commits

Reviewing files that changed from the base of the PR and between 9d5f81b and 87a3bc9.

📒 Files selected for processing (11)

• maint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.py
• src/backend/common/op/reduce.h
• src/transform/loop_partition.cc
• src/transform/unroll_loop.cc
• testing/python/language/test_tilelang_language_atomic.py
• testing/python/layout/test_tilelang_cute.py
• testing/python/transform/test_tilelang_transform_unroll_loop.py
• tilelang/cuda/pipeline.py
• tilelang/language/loop.py
• tilelang/language/tir/ir.py
• tilelang/tools/pass_visualizer/core.py

💤 Files with no reviewable changes (1)

• src/transform/loop_partition.cc

---

📝 Walkthrough

Walkthrough

Changes

Explicit loop unrolling

Layer / File(s)	Summary
Per-loop unroll semantics tilelang/language/..., src/transform/..., testing/python/transform/*	Unroll annotations now determine explicit unrolling per loop. Tests cover explicit, dynamic, and factor-annotated loops.
Pipeline preprocessing integration tilelang/cuda/pipeline.py, tilelang/tools/pass_visualizer/core.py	The pipeline runs UnrollLoop and Simplify before PipelinePlanning.
Lowering behavior and generated-source validation src/backend/common/op/reduce.h, maint/gemm/..., testing/python/language/*, testing/python/layout/*	Lowering paths stop disabling explicit unrolling. TMA source checks now validate generated loops with regular expressions.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant LoopAPI
  participant UnrollLoop
  participant CUDAPipeline
  participant PipelinePlanning
  LoopAPI->>UnrollLoop: provide loop annotations
  UnrollLoop->>UnrollLoop: resolve explicitness per loop
  CUDAPipeline->>UnrollLoop: preprocess explicit loops
  UnrollLoop->>CUDAPipeline: return expanded loops
  CUDAPipeline->>PipelinePlanning: submit simplified scheduling units

Loading

Possibly related PRs

• tile-ai/tilelang#2449: Updates the pass visualizer pipeline represented by build_pass_stages.
• tile-ai/tilelang#2718: Also changes loop partitioning and explicit-unroll behavior.
• tile-ai/tilelang#2836: Also changes explicit loop unrolling and pipeline preprocessing.

Suggested reviewers: leiwang1999, siriusneo

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes reusing explicit unrolling for early preprocessing, which matches the main objectives and changes.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[LeiWang1999]

@regression-perf

[github-actions]

Performance Regression Test Report

Triggered by: @LeiWang1999
Workflow run: https://github.com/tile-ai/tilelang/actions/runs/30878834122

Results

File	Original Latency	Current Latency	Speedup
example_topk	0.0426851	0.056502	0.755461
example_dequant_gemm_bf16_mxfp4_hopper	0.252317	0.255508	0.987511
sparse_mla_bwd	0.136669	0.138283	0.988328
example_gemm	0.0148016	0.0149256	0.991695
example_gqa_decode	0.0304913	0.0306849	0.993693
example_warp_specialize_gemm_copy_0_gemm_1	0.0235397	0.0236692	0.994528
sparse_mla_fwd	0.0529198	0.0532022	0.994693
example_mha_bwd_bshd	0.0138476	0.0139192	0.994856
example_gqa_bwd_tma_reduce_varlen	0.0277781	0.0279195	0.994935
fp8_lighting_indexer	0.0119068	0.0119585	0.995676
example_gqa_fwd_bshd	0.0296676	0.0297918	0.995833
example_mha_fwd_varlen	0.0205631	0.0206221	0.997138
example_tilelang_nsa_fwd	0.00405489	0.00406372	0.997827
example_tilelang_block_sparse_attn	0.0057211	0.00573333	0.997867
example_tilelang_gemm_splitk_vectorize_atomicadd	0.584247	0.585352	0.998112
example_dequant_gemm_bf16_fp4_hopper	0.268608	0.269113	0.998124
example_convolution	0.584262	0.585341	0.998157
example_mha_sink_fwd_bhsd	0.00982929	0.00984398	0.998507
example_mha_sink_bwd_bhsd	0.040819	0.0408598	0.999001
topk_selector	0.0272658	0.0272837	0.999344
example_warp_specialize_gemm_softpipe_stage2	0.015469	0.0154777	0.99944
example_mha_bwd_bhsd	0.0139969	0.0140027	0.999584
example_dequant_gemv_fp16xint4	0.0173423	0.0173488	0.999626
example_mha_sink_bwd_bhsd_sliding_window	0.0261233	0.0261324	0.999653
example_gemm_intrinsics	0.020162	0.0201678	0.999711
example_elementwise_add	0.0690938	0.0691059	0.999825
example_gqa_bwd	0.0287892	0.0287938	0.999842
example_blocksparse_gemm	0.0116757	0.0116774	0.999853
sparse_mla_fwd_pipelined	0.0343762	0.0343738	1.00007
example_dequant_gemm_w4a8	2.67897	2.67874	1.00008
example_convolution_autotune	0.591207	0.591083	1.00021
example_tilelang_gemm_fp8_2xAcc	0.0678591	0.0678407	1.00027
example_gqa_sink_bwd_bhsd_sliding_window	0.0152786	0.0152686	1.00066
example_tilelang_sparse_gqa_decode_varlen_indice	0.0108152	0.0108081	1.00066
example_tilelang_gemm_splitk	0.58909	0.588678	1.0007
example_fusedmoe_tilelang	0.0765446	0.0764861	1.00076
example_linear_attn_bwd	0.0971181	0.0970412	1.00079
example_dynamic	0.388655	0.3883	1.00091
example_mha_fwd_bshd	0.0148068	0.0147933	1.00092
block_sparse_attn_tilelang	0.00619208	0.00618628	1.00094
example_mhc_post	0.0656767	0.0656107	1.00101
example_mhc_pre	0.115449	0.115308	1.00122
example_tilelang_gemm_fp8	0.171223	0.171003	1.00129
example_mha_sink_fwd_bhsd_sliding_window	0.00973786	0.00972509	1.00131
example_tilelang_sparse_gqa_decode_varlen_mask	0.0282558	0.0282046	1.00182
example_mha_inference	0.0329875	0.0329247	1.00191
example_gqa_sink_bwd_bhsd	0.0250996	0.0250438	1.00223
example_per_token_cast_to_fp8	0.00432852	0.00431863	1.00229
example_mha_fwd_bhsd	0.00689009	0.00687368	1.00239
example_group_per_split_token_cast_to_fp8	0.00561904	0.00560517	1.00247
example_tilelang_nsa_decode	0.00418619	0.00417446	1.00281
example_warp_specialize_gemm_barrierpipe_stage2	0.0248498	0.0247075	1.00576
example_vertical_slash_sparse_attn	0.135681	0.134883	1.00592
example_linear_attn_fwd	0.0229264	0.0227856	1.00618
example_gemv	0.149076	0.148081	1.00672
example_mla_decode	0.29921	0.296348	1.00966
example_dequant_gemm_fp4_hopper	0.542422	0.533364	1.01698
example_warp_specialize_gemm_copy_1_gemm_0	0.0161919	0.0154961	1.0449

Artifacts

• regression_result.png (speedup plot) is attached as a workflow artifact. Download it from the workflow run page above.

[Yongqi-Zhuo]

example_topk regression confirmed to be noise cc @LeiWang1999