[Triton] Integrate reorder-loop-loads into LoopUnroll pass

[isaname]

Background

In Triton kernels, load operations and compute operations are interleaved after loop unrolling. On some hardware, this interleaving prevents memory access latency from being effectively hidden — because each load's result is immediately consumed by the computation that follows it, leaving the pipeline underutilized.

The ideal approach is to hoist (cluster) the load operations from each unrolled iteration to the front, so that multiple loads can be in-flight simultaneously, thereby exploiting memory system bandwidth and hiding latency.

Design

1. Reorder logic integrated into LoopUnroll

For loops with tt.reorder attribute:

• Full unroll: loads and their transitive deps are clustered to the front of the unrolled op range in the parent block.
• Partial unroll: same clustering is applied within the expanded loop body.

Key helpers: willFullyUnroll() (static trip count check), collectDepsInBlock() (transitive dep collection), reorderLoadsInRange() (load-first reordering preserving relative order within each group).

2. User interface

Added reorder=False parameter to tl.range. When reorder=True with loop_unroll_factor set, the tt.reorder attribute is attached to the scf.for, triggering load clustering during unroll.

for i in tl.range(0, N, loop_unroll_factor=4, reorder=True):
    a = tl.load(a_ptr + i * BLOCK)
    b = tl.load(b_ptr + i * BLOCK)
    c = a + b
    tl.store(c_ptr + i * BLOCK, c)

Performance

fused_inv_rope_fp8_quant

num_tokens	num_heads	n_groups	heads_per_group	tma_aligned_scales	Δ%(pipeline)	Δ%(reorder)
256	128	8	16	False	+9.41%	+12.89%
256	128	8	16	True	+11.37%	+13.76%
512	128	8	16	False	+8.04%	+12.49%
512	128	8	16	True	+12.26%	+12.76%
1024	128	8	16	False	+11.11%	+17.39%
1024	128	8	16	True	+14.46%	+16.27%
2048	128	8	16	False	+13.58%	+18.02%
2048	128	8	16	True	+16.42%	+17.92%
4096	128	8	16	False	+12.46%	+17.03%
4096	128	8	16	True	+15.78%	+17.88%

[CLAassistant]

All committers have signed the CLA.

[zhzhcookie]

Please refer to PR 763 to add the C++ macro switch.
For every change to the Python code, add the comment # flagtree reorder-loop-loads.

[sunnycase]

Thanks for working on this optimization. Could you please avoid modifying the upstream triton.language.range / tl.range primitive? Since reorder is a FlagTree/TLE-specific extension, please introduce tle.range(..., reorder=False) with reorder as a boolean parameter, and keep the existing Triton primitive unchanged.

Would it also be possible to move as much of the backend-specific implementation as possible out of upstream Triton-owned directories and into the TLE subtree? If a hook in an upstream-owned file is unavoidable, please keep it minimal and place the TLE-specific code behind the appropriate TLE build guard. This will minimize divergence from upstream and make ownership of the extension explicit.

Please update the relevant call sites and tests to use tle.range. Thank you.