From e71efaceba867cd733c7ac85e638d348eab12447 Mon Sep 17 00:00:00 2001 From: oyrh <3456310189@qq.com> Date: Thu, 9 Jul 2026 11:39:29 +0000 Subject: [PATCH 1/4] [Triton] Integrate reorder-loop-loads into LoopUnroll pass --- lib/Dialect/Triton/Transforms/LoopUnroll.cpp | 133 +++++++++++++++++++ python/triton/compiler/code_generator.py | 4 + python/triton/language/core.py | 4 +- 3 files changed, 140 insertions(+), 1 deletion(-) diff --git a/lib/Dialect/Triton/Transforms/LoopUnroll.cpp b/lib/Dialect/Triton/Transforms/LoopUnroll.cpp index 294dff873..9abb0690e 100644 --- a/lib/Dialect/Triton/Transforms/LoopUnroll.cpp +++ b/lib/Dialect/Triton/Transforms/LoopUnroll.cpp @@ -8,6 +8,7 @@ #include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "triton/Dialect/Triton/IR/Dialect.h" #include "triton/Dialect/Triton/Transforms/Passes.h" +#include "llvm/ADT/SetVector.h" #include "llvm/Support/Debug.h" namespace mlir::triton { @@ -30,8 +31,101 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { return 1; } + // Returns true if unrolling by the given factor will fully eliminate the + // loop. + bool willFullyUnroll(scf::ForOp forOp, int64_t unrollFactor) { + IntegerAttr lbAttr, ubAttr, stepAttr; + if (!matchPattern(forOp.getLowerBound(), m_Constant(&lbAttr)) || + !matchPattern(forOp.getUpperBound(), m_Constant(&ubAttr)) || + !matchPattern(forOp.getStep(), m_Constant(&stepAttr))) + return false; + int64_t lb = lbAttr.getInt(); + int64_t ub = ubAttr.getInt(); + int64_t step = stepAttr.getInt(); + if (step <= 0 || ub <= lb) + return false; + int64_t tripCount = llvm::divideCeil(ub - lb, step); + return tripCount > 0 && tripCount <= unrollFactor; + } + + bool hasReorderAttr(scf::ForOp forOp) { + if (auto reorderAttr = forOp->getAttrOfType(reorderAttrName)) + return reorderAttr.getValue(); + return false; + } + + // Collect all transitive in-block dependencies of `op`. + void collectDepsInBlock(Operation *op, Block *block, + llvm::SetVector &deps) { + for (Value operand : op->getOperands()) { + auto *defOp = operand.getDefiningOp(); + if (!defOp || defOp->getBlock() != block) + continue; + if (deps.insert(defOp)) { + collectDepsInBlock(defOp, block, deps); + } + } + } + + // Reorder loads in a range of operations within a block. Moves all load ops + // and their transitive dependencies before other ops in the range, preserving + // relative order within each group. + void reorderLoadsInRange(Block *block, Operation *beginOp, Operation *endOp) { + // Collect all ops in the range [beginOp, endOp) + SmallVector rangeOps; + for (auto it = Block::iterator(beginOp); &*it != endOp; ++it) { + rangeOps.push_back(&*it); + } + + if (rangeOps.empty()) + return; + + // Identify loads and their dependency cluster + SmallVector loads; + llvm::SetVector loadCluster; + + for (Operation *op : rangeOps) { + if (isa(op)) + loads.push_back(op); + } + + if (loads.empty()) { + LDBG("No loads found in unrolled range, skipping reorder"); + return; + } + + for (Operation *load : loads) { + loadCluster.insert(load); + collectDepsInBlock(load, block, loadCluster); // * 收集load元素的依赖 + } + + LDBG("Full-unroll reorder: load cluster size: " + << loadCluster.size() << " (loads: " << loads.size() << ")"); + + // Build reordered sequence: + // 1. Load cluster ops (deps + loads) in original relative order + // 2. Remaining ops in original relative order + SmallVector reordered; + for (Operation *op : rangeOps) { + if (loadCluster.contains(op)) + reordered.push_back(op); + } + for (Operation *op : rangeOps) { + if (!loadCluster.contains(op)) + reordered.push_back(op); + } + + // Apply the new ordering by moving ops before endOp + for (Operation *op : reordered) { + op->moveBefore(endOp); + } + + LDBG("Full-unroll reorder complete"); + } + const char *loopUnrollFactorAttrName = "tt.loop_unroll_factor"; const char *pipelineStagesAttrName = "tt.num_stages"; + const char *reorderAttrName = "tt.reorder"; public: void runOnOperation() override { @@ -46,9 +140,48 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { auto ctx = getOperation()->getContext(); for (auto loop : loops) { auto unrollFactor = getUnrollFactorOrDefault(loop); + bool needsReorder = hasReorderAttr(loop); + bool fullyUnrolls = willFullyUnroll(loop, unrollFactor); + loop->removeAttr(loopUnrollFactorAttrName); + if (needsReorder) + loop->removeAttr(reorderAttrName); + + // Record position before unrolling so we can identify unrolled ops + // if the loop gets fully eliminated. + Block *parentBlock = loop->getBlock(); + Operation *opBeforeLoop = loop->getPrevNode(); // may be nullptr + LDBG("Unrolling loop by " << unrollFactor << " times\n" << loop); auto resultLoops = loopUnrollByFactor(loop, unrollFactor); + + if (failed(resultLoops)) + continue; + if (needsReorder) { + if (fullyUnrolls) { + // Full unroll: the loop has been completely eliminated. + // The unrolled ops are now in parentBlock between opBeforeLoop and + // the block terminator. + Operation *rangeBegin = opBeforeLoop ? opBeforeLoop->getNextNode() + : &parentBlock->front(); + Operation *rangeEnd = parentBlock->getTerminator(); + reorderLoadsInRange(parentBlock, rangeBegin, rangeEnd); + } else { + // Partial unroll: the loop still exists, body has been replicated. + // Reorder loads within the loop body. + // After partial unroll the original loop is updated in-place, so + // `loop` is still valid but its body has been expanded. + // However loopUnrollByFactor may return a new main loop op when + // there's an epilogue. Use the main loop if available. + scf::ForOp mainLoop = + resultLoops->mainLoopOp ? *resultLoops->mainLoopOp : loop; + if (mainLoop) { + Block *body = mainLoop.getBody(); + reorderLoadsInRange(body, &body->front(), body->getTerminator()); + } + } + } + // Do not pipeline the epilog loop. if (succeeded(resultLoops) && resultLoops->epilogueLoopOp) { (*resultLoops->epilogueLoopOp) diff --git a/python/triton/compiler/code_generator.py b/python/triton/compiler/code_generator.py index e79845f51..e318502d4 100644 --- a/python/triton/compiler/code_generator.py +++ b/python/triton/compiler/code_generator.py @@ -1244,6 +1244,7 @@ def visit_For(self, node): flatten = False warp_specialize = False disable_licm = False + reorder = False # flagtree tle try: from ..experimental.tle import language as tle @@ -1265,6 +1266,7 @@ def visit_For(self, node): flatten = iterator.flatten warp_specialize = iterator.warp_specialize disable_licm = iterator.disable_licm + reorder = getattr(iterator, 'reorder', False) elif IteratorClass is range: # visit iterator arguments # note: only `range` iterator is supported now @@ -1329,6 +1331,8 @@ def visit_For(self, node): for_op.set_attr("tt.warp_specialize", self.builder.get_unit_attr()) if disable_licm: for_op.set_attr("llvm.loop_annotation", self.builder.get_disable_loop_licm_attr()) + if reorder and _unwrap_if_constexpr(loop_unroll_factor) is not None: + for_op.set_attr("tt.reorder", self.builder.get_bool_attr(True)) self.scf_stack.append(node) for_op_body = for_op.get_body(0) diff --git a/python/triton/language/core.py b/python/triton/language/core.py index 439671c3a..a1c2cd13e 100644 --- a/python/triton/language/core.py +++ b/python/triton/language/core.py @@ -3287,7 +3287,8 @@ def kernel(...): """ def __init__(self, arg1, arg2=None, step=None, num_stages=None, loop_unroll_factor=None, - disallow_acc_multi_buffer=False, flatten=False, warp_specialize=False, disable_licm=False): + disallow_acc_multi_buffer=False, flatten=False, warp_specialize=False, disable_licm=False, + reorder=False): if step is None: self.step = constexpr(1) else: @@ -3304,6 +3305,7 @@ def __init__(self, arg1, arg2=None, step=None, num_stages=None, loop_unroll_fact self.flatten = flatten self.warp_specialize = warp_specialize self.disable_licm = disable_licm + self.reorder = reorder def __iter__(self): raise RuntimeError("tl.range can only be used in @triton.jit'd functions") From 3c584e4f3e555e21a4a73f4827d945c3da64d2ca Mon Sep 17 00:00:00 2001 From: oyrh <3456310189@qq.com> Date: Mon, 13 Jul 2026 02:58:11 +0000 Subject: [PATCH 2/4] [Tutorial] Add fused inv_rope fp8_quant reorder tutorial --- .../12-fused_inv_rope_fp8_quant_reorder.py | 240 ++++++++++++++++++ 1 file changed, 240 insertions(+) create mode 100644 python/tutorials/12-fused_inv_rope_fp8_quant_reorder.py diff --git a/python/tutorials/12-fused_inv_rope_fp8_quant_reorder.py b/python/tutorials/12-fused_inv_rope_fp8_quant_reorder.py new file mode 100644 index 000000000..204c15798 --- /dev/null +++ b/python/tutorials/12-fused_inv_rope_fp8_quant_reorder.py @@ -0,0 +1,240 @@ +import logging +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from flag_gems.runtime import torch_device_fn +from flag_gems.utils.device_info import get_device_capability + +if torch_device_fn.is_available() and get_device_capability() >= (9, 0): + SUPPORTED_FP8_DTYPE = torch.float8_e4m3fn +else: + SUPPORTED_FP8_DTYPE = torch.float32 + +logger = logging.getLogger(__name__) + + +def _get_tma_aligned_size(size: int, align: int) -> int: + return ((size + align - 1) // align) * align + + +@triton.jit +def _fused_inv_rope_fp8_quant_per_head( + o_ptr, + positions_ptr, + cos_sin_cache_ptr, + fp8_ptr, + scale_ptr, + num_tokens, + heads_per_group: tl.constexpr, + o_stride_token, + o_stride_head, + cache_stride_pos, + fp8_stride_group, + fp8_stride_token, + scale_stride_group, + scale_stride_k, + fp8_max: tl.constexpr, + eps: tl.constexpr, + QUANT_GROUP_SIZE: tl.constexpr, + CHUNKS_PER_HEAD: tl.constexpr, + ROPE_START: tl.constexpr, + HALF_ROPE: tl.constexpr, + TMA_ALIGNED_SCALES: tl.constexpr, + HEADS_PER_PROGRAM: tl.constexpr, +): + pid_token = tl.program_id(0).to(tl.int64) + pid_gh_base = tl.program_id(1).to(tl.int64) * HEADS_PER_PROGRAM + + HEAD_DIM: tl.constexpr = CHUNKS_PER_HEAD * QUANT_GROUP_SIZE + + if pid_token >= num_tokens: + # Zero-fill scales for padding tokens across all heads in this program. + for i in tl.static_range(0, HEADS_PER_PROGRAM): + global_head = pid_gh_base + i + g = global_head // heads_per_group + head_in_group = global_head % heads_per_group + qb_start = head_in_group * CHUNKS_PER_HEAD + if TMA_ALIGNED_SCALES: + scale_addr = (scale_ptr + g * scale_stride_group + pid_token + head_in_group * scale_stride_k) + tl.store(scale_addr, tl.zeros((), dtype=tl.int32)) + else: + block_offsets = tl.arange(0, CHUNKS_PER_HEAD) + qb_indices = qb_start + block_offsets + scale_addrs = (scale_ptr + g * scale_stride_group + pid_token + qb_indices * scale_stride_k) + tl.store(scale_addrs, tl.zeros((CHUNKS_PER_HEAD, ), dtype=tl.float32)) + return + + offsets = tl.arange(0, HEAD_DIM) + rope_abs_start: tl.constexpr = (CHUNKS_PER_HEAD - 1) * QUANT_GROUP_SIZE + ROPE_START + pos = tl.load(positions_ptr + pid_token) + cache_base = cos_sin_cache_ptr + pos * cache_stride_pos + is_rope = offsets >= rope_abs_start + rope_local = offsets - rope_abs_start + cs_idx = tl.maximum(rope_local >> 1, 0) + is_even = (rope_local & 1) == 0 + + # Load cos/sin once — shared across all heads for the same token. + cos_v = tl.load(cache_base + cs_idx, mask=is_rope, other=1.0) + sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope, other=0.0) + + for i in tl.range(0, HEADS_PER_PROGRAM, loop_unroll_factor=HEADS_PER_PROGRAM, reorder=True): + global_head = pid_gh_base + i + g = global_head // heads_per_group + head_in_group = global_head % heads_per_group + qb_start = head_in_group * CHUNKS_PER_HEAD + + input_base = o_ptr + pid_token * o_stride_token + global_head * o_stride_head + + x = tl.load(input_base + offsets).to(tl.float32) + + x_partner = tl.load(input_base + (offsets ^ 1), mask=is_rope, other=0.0).to(tl.float32) + x_add = x * cos_v + x_partner * sin_v + x_sub = x * cos_v - x_partner * sin_v + rotated = tl.where(is_even, x_add, x_sub) + x = tl.where(is_rope, rotated, x) + + x_2d = tl.reshape(tl.abs(x), (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE)) + block_absmax = tl.maximum(tl.max(x_2d, axis=1), eps) + scales = block_absmax * (1.0 / fp8_max) + if TMA_ALIGNED_SCALES: + scales = tl.math.exp2(tl.ceil(tl.log2(tl.maximum(tl.abs(scales), 1e-10)))) + + scales_exp = tl.reshape( + tl.broadcast_to( + tl.reshape(scales, (CHUNKS_PER_HEAD, 1)), + (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE), + ), + (HEAD_DIM, ), + ) + x_quant = tl.clamp(x / scales_exp, -fp8_max, fp8_max).to(tl.float8e4nv) + + fp8_base = (fp8_ptr + g * fp8_stride_group + pid_token * fp8_stride_token + qb_start * QUANT_GROUP_SIZE) + tl.store(fp8_base + offsets, x_quant) + + block_offsets = tl.arange(0, CHUNKS_PER_HEAD) + qb_indices = qb_start + block_offsets + if TMA_ALIGNED_SCALES: + scale_bits = scales.to(tl.int32, bitcast=True) + ue8m0_bytes = (scale_bits >> 23) & 0xFF + packed_val = tl.sum(ue8m0_bytes << (block_offsets * 8)) + scale_addr = (scale_ptr + g * scale_stride_group + pid_token + head_in_group * scale_stride_k) + tl.store(scale_addr, packed_val) + else: + scale_addrs = (scale_ptr + g * scale_stride_group + pid_token + qb_indices * scale_stride_k) + tl.store(scale_addrs, scales) + + +def fused_inv_rope_fp8_quant_reorder( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + n_groups: int, + heads_per_group: int, + nope_dim: int = 448, + rope_dim: int = 64, + quant_group_size: int = 128, + eps: float = 1e-10, + dtype: Optional[torch.dtype] = None, + tma_aligned_scales: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Triton draft of DeepSeek-V4 fused inverse-RoPE + FP8 group quant. + + Uses tl.range with reorder=True to cluster loads across multiple heads, + improving memory latency hiding via software pipelining after unroll. + + Args: + o: [num_tokens, num_heads, head_dim] + positions: [num_tokens] + cos_sin_cache: [max_position, rope_dim] laid out as cos || sin + + Returns: + o_fp8: [num_tokens, n_groups, heads_per_group * head_dim] + o_scale: [num_tokens, n_groups, num_scale_blocks] or packed UE8M0 view + """ + logger.debug("GEMS FUSED INV ROPE FP8 QUANT REORDER") + fp8_dtype = SUPPORTED_FP8_DTYPE if dtype is None else dtype + assert fp8_dtype == torch.float8_e4m3fn, "only torch.float8_e4m3fn is supported" + assert o.ndim == 3, "`o` must be [num_tokens, num_heads, head_dim]" + assert positions.ndim == 1, "`positions` must be 1D" + assert cos_sin_cache.ndim == 2, "`cos_sin_cache` must be 2D" + assert o.stride(-1) == 1, "head_dim must be contiguous" + assert positions.shape[0] == o.shape[0], "positions and o token count mismatch" + + num_tokens, num_heads, head_dim = o.shape + assert num_heads == n_groups * heads_per_group + assert head_dim == nope_dim + rope_dim + assert head_dim % quant_group_size == 0 + assert nope_dim % quant_group_size == (quant_group_size - rope_dim) + assert rope_dim % 2 == 0 + assert cos_sin_cache.shape[-1] == rope_dim + assert cos_sin_cache.dtype == torch.float32 + + chunks_per_head = head_dim // quant_group_size + if tma_aligned_scales: + assert (chunks_per_head <= 4), "packed UE8M0 path currently expects at most 4 scale blocks per head" + + d = heads_per_group * head_dim + num_scale_blocks = d // quant_group_size + tma_aligned_t = _get_tma_aligned_size(num_tokens, 4) + + if tma_aligned_scales: + scale_inner = (num_scale_blocks + 3) // 4 + scale_dtype = torch.int32 + else: + scale_inner = num_scale_blocks + scale_dtype = torch.float32 + + finfo = torch.finfo(fp8_dtype) + fp8_q = torch.empty((n_groups, num_tokens, d), dtype=fp8_dtype, device=o.device) + scale = torch.empty( + n_groups * scale_inner * tma_aligned_t, + dtype=scale_dtype, + device=o.device, + ).as_strided( + (n_groups, num_tokens, scale_inner), + (scale_inner * tma_aligned_t, 1, tma_aligned_t), + ) + + # Determine how many heads each program handles. + # Use smaller batches (2 or 4) to control register pressure, as each head + # holds HEAD_DIM × multiple arrays simultaneously. + total_heads = n_groups * heads_per_group + heads_per_program = 1 + for candidate in [4, 2]: + if total_heads % candidate == 0: + heads_per_program = candidate + break + + grid = (tma_aligned_t, total_heads // heads_per_program) + _fused_inv_rope_fp8_quant_per_head[grid]( + o, + positions, + cos_sin_cache, + fp8_q, + scale, + num_tokens, + heads_per_group=heads_per_group, + o_stride_token=o.stride(0), + o_stride_head=o.stride(1), + cache_stride_pos=cos_sin_cache.stride(0), + fp8_stride_group=fp8_q.stride(0), + fp8_stride_token=fp8_q.stride(1), + scale_stride_group=scale.stride(0), + scale_stride_k=scale.stride(2), + fp8_max=finfo.max, + eps=eps, + QUANT_GROUP_SIZE=quant_group_size, + CHUNKS_PER_HEAD=chunks_per_head, + ROPE_START=nope_dim % quant_group_size, + HALF_ROPE=rope_dim // 2, + TMA_ALIGNED_SCALES=tma_aligned_scales, + HEADS_PER_PROGRAM=heads_per_program, + num_warps=1, + num_stages=1, + ) + + return fp8_q.transpose(0, 1), scale.transpose(0, 1) From 8e63d2d64afeb06c357640ee9ade3bfc005e7b64 Mon Sep 17 00:00:00 2001 From: oyrh <3456310189@qq.com> Date: Wed, 15 Jul 2026 03:25:42 +0000 Subject: [PATCH 3/4] [Triton] Add compile-time toggle for reorder-loop-loads feature --- CMakeLists.txt | 1 + lib/Dialect/Triton/Transforms/LoopUnroll.cpp | 17 +++++++++++++++-- python/triton/compiler/code_generator.py | 8 ++++---- python/triton/language/core.py | 4 ++-- 4 files changed, 22 insertions(+), 8 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 52b0bee20..bd09b0904 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -35,6 +35,7 @@ endif() if(NOT FLAGTREE_BACKEND) add_definitions(-D__NVIDIA__) add_definitions(-D__AMD__) + add_definitions(-D__FLAGTREE_REORDER_LOOP_LOADS__) elseif(FLAGTREE_BACKEND STREQUAL "iluvatar") add_definitions(-D__ILUVATAR__) set(FLAGTREE_TLE OFF) diff --git a/lib/Dialect/Triton/Transforms/LoopUnroll.cpp b/lib/Dialect/Triton/Transforms/LoopUnroll.cpp index 9abb0690e..2e8cec6bb 100644 --- a/lib/Dialect/Triton/Transforms/LoopUnroll.cpp +++ b/lib/Dialect/Triton/Transforms/LoopUnroll.cpp @@ -31,6 +31,7 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { return 1; } +#ifdef __FLAGTREE_REORDER_LOOP_LOADS__ // Returns true if unrolling by the given factor will fully eliminate the // loop. bool willFullyUnroll(scf::ForOp forOp, int64_t unrollFactor) { @@ -45,7 +46,10 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { if (step <= 0 || ub <= lb) return false; int64_t tripCount = llvm::divideCeil(ub - lb, step); - return tripCount > 0 && tripCount <= unrollFactor; + // The main loop after partial unroll has tripCount / unrollFactor + // iterations. If that is <= 1, MLIR will eliminate it by inlining the + // body, so from our perspective the main loop is fully unrolled. + return tripCount > 0 && tripCount / unrollFactor <= 1; } bool hasReorderAttr(scf::ForOp forOp) { @@ -96,7 +100,7 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { for (Operation *load : loads) { loadCluster.insert(load); - collectDepsInBlock(load, block, loadCluster); // * 收集load元素的依赖 + collectDepsInBlock(load, block, loadCluster); } LDBG("Full-unroll reorder: load cluster size: " @@ -122,10 +126,13 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { LDBG("Full-unroll reorder complete"); } +#endif // __FLAGTREE_REORDER_LOOP_LOADS__ const char *loopUnrollFactorAttrName = "tt.loop_unroll_factor"; const char *pipelineStagesAttrName = "tt.num_stages"; +#ifdef __FLAGTREE_REORDER_LOOP_LOADS__ const char *reorderAttrName = "tt.reorder"; +#endif // __FLAGTREE_REORDER_LOOP_LOADS__ public: void runOnOperation() override { @@ -140,10 +147,13 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { auto ctx = getOperation()->getContext(); for (auto loop : loops) { auto unrollFactor = getUnrollFactorOrDefault(loop); +#ifdef __FLAGTREE_REORDER_LOOP_LOADS__ bool needsReorder = hasReorderAttr(loop); bool fullyUnrolls = willFullyUnroll(loop, unrollFactor); +#endif // __FLAGTREE_REORDER_LOOP_LOADS__ loop->removeAttr(loopUnrollFactorAttrName); +#ifdef __FLAGTREE_REORDER_LOOP_LOADS__ if (needsReorder) loop->removeAttr(reorderAttrName); @@ -151,12 +161,14 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { // if the loop gets fully eliminated. Block *parentBlock = loop->getBlock(); Operation *opBeforeLoop = loop->getPrevNode(); // may be nullptr +#endif // __FLAGTREE_REORDER_LOOP_LOADS__ LDBG("Unrolling loop by " << unrollFactor << " times\n" << loop); auto resultLoops = loopUnrollByFactor(loop, unrollFactor); if (failed(resultLoops)) continue; +#ifdef __FLAGTREE_REORDER_LOOP_LOADS__ if (needsReorder) { if (fullyUnrolls) { // Full unroll: the loop has been completely eliminated. @@ -181,6 +193,7 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { } } } +#endif // __FLAGTREE_REORDER_LOOP_LOADS__ // Do not pipeline the epilog loop. if (succeeded(resultLoops) && resultLoops->epilogueLoopOp) { diff --git a/python/triton/compiler/code_generator.py b/python/triton/compiler/code_generator.py index e318502d4..47f74d17b 100644 --- a/python/triton/compiler/code_generator.py +++ b/python/triton/compiler/code_generator.py @@ -1244,7 +1244,7 @@ def visit_For(self, node): flatten = False warp_specialize = False disable_licm = False - reorder = False + reorder = False # flagtree reorder-loop-loads # flagtree tle try: from ..experimental.tle import language as tle @@ -1266,7 +1266,7 @@ def visit_For(self, node): flatten = iterator.flatten warp_specialize = iterator.warp_specialize disable_licm = iterator.disable_licm - reorder = getattr(iterator, 'reorder', False) + reorder = getattr(iterator, 'reorder', False) # flagtree reorder-loop-loads elif IteratorClass is range: # visit iterator arguments # note: only `range` iterator is supported now @@ -1331,8 +1331,8 @@ def visit_For(self, node): for_op.set_attr("tt.warp_specialize", self.builder.get_unit_attr()) if disable_licm: for_op.set_attr("llvm.loop_annotation", self.builder.get_disable_loop_licm_attr()) - if reorder and _unwrap_if_constexpr(loop_unroll_factor) is not None: - for_op.set_attr("tt.reorder", self.builder.get_bool_attr(True)) + if reorder and _unwrap_if_constexpr(loop_unroll_factor) is not None: # flagtree reorder-loop-loads + for_op.set_attr("tt.reorder", self.builder.get_bool_attr(True)) # flagtree reorder-loop-loads self.scf_stack.append(node) for_op_body = for_op.get_body(0) diff --git a/python/triton/language/core.py b/python/triton/language/core.py index a1c2cd13e..9dcae2e67 100644 --- a/python/triton/language/core.py +++ b/python/triton/language/core.py @@ -3288,7 +3288,7 @@ def kernel(...): def __init__(self, arg1, arg2=None, step=None, num_stages=None, loop_unroll_factor=None, disallow_acc_multi_buffer=False, flatten=False, warp_specialize=False, disable_licm=False, - reorder=False): + reorder=False): # flagtree reorder-loop-loads if step is None: self.step = constexpr(1) else: @@ -3305,7 +3305,7 @@ def __init__(self, arg1, arg2=None, step=None, num_stages=None, loop_unroll_fact self.flatten = flatten self.warp_specialize = warp_specialize self.disable_licm = disable_licm - self.reorder = reorder + self.reorder = reorder # flagtree reorder-loop-loads def __iter__(self): raise RuntimeError("tl.range can only be used in @triton.jit'd functions") From 26cebfa91c26b5a800d39b8918441d4f8495fb39 Mon Sep 17 00:00:00 2001 From: oyrh <3456310189@qq.com> Date: Wed, 15 Jul 2026 08:06:12 +0000 Subject: [PATCH 4/4] Style: fix formatting issue --- lib/Dialect/Triton/Transforms/LoopUnroll.cpp | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/lib/Dialect/Triton/Transforms/LoopUnroll.cpp b/lib/Dialect/Triton/Transforms/LoopUnroll.cpp index 2e8cec6bb..39a11627d 100644 --- a/lib/Dialect/Triton/Transforms/LoopUnroll.cpp +++ b/lib/Dialect/Triton/Transforms/LoopUnroll.cpp @@ -8,7 +8,9 @@ #include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "triton/Dialect/Triton/IR/Dialect.h" #include "triton/Dialect/Triton/Transforms/Passes.h" +#ifdef __FLAGTREE_REORDER_LOOP_LOADS__ #include "llvm/ADT/SetVector.h" +#endif // __FLAGTREE_REORDER_LOOP_LOADS__ #include "llvm/Support/Debug.h" namespace mlir::triton { @@ -151,7 +153,6 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { bool needsReorder = hasReorderAttr(loop); bool fullyUnrolls = willFullyUnroll(loop, unrollFactor); #endif // __FLAGTREE_REORDER_LOOP_LOADS__ - loop->removeAttr(loopUnrollFactorAttrName); #ifdef __FLAGTREE_REORDER_LOOP_LOADS__ if (needsReorder) @@ -162,13 +163,11 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { Block *parentBlock = loop->getBlock(); Operation *opBeforeLoop = loop->getPrevNode(); // may be nullptr #endif // __FLAGTREE_REORDER_LOOP_LOADS__ - LDBG("Unrolling loop by " << unrollFactor << " times\n" << loop); auto resultLoops = loopUnrollByFactor(loop, unrollFactor); - +#ifdef __FLAGTREE_REORDER_LOOP_LOADS__ if (failed(resultLoops)) continue; -#ifdef __FLAGTREE_REORDER_LOOP_LOADS__ if (needsReorder) { if (fullyUnrolls) { // Full unroll: the loop has been completely eliminated. @@ -194,7 +193,6 @@ class LoopUnrollPass : public impl::TritonLoopUnrollBase { } } #endif // __FLAGTREE_REORDER_LOOP_LOADS__ - // Do not pipeline the epilog loop. if (succeeded(resultLoops) && resultLoops->epilogueLoopOp) { (*resultLoops->epilogueLoopOp)