[Mlir-commits] [mlir] [MLIR][SCF] Fix empty loop range in mlir::loopUnrollByFactor. (PR #237094)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 29 22:10:25 PDT 2026
https://github.com/sun-jacobi created https://github.com/llvm/llvm-project/pull/237094
## TL; DR
This patches fixes emtry range such as $[42, 0; 4)$ in `mlir::loopUnrollByFactor`
According to the SCF documentation, the precondition is
> Lower bound, upper bound, and step are interpreted as signed integers by default, or as unsigned integers if the unsignedCmp unit attribute is present. The step is required to be strictly positive.
However, even though in the case that lower bound, upper bound, and step are all positive, we can construct an empty range such as $[42, 0; 4)$, which leads to wrong behavior as reported in triton-lang/triton/issues/12525
>From 8443a4df1332ee245afa8c79557e6646f6023a7a Mon Sep 17 00:00:00 2001
From: sun-jacobi <sun1011jacobi at gmail.com>
Date: Fri, 30 Oct 2026 13:45:42 +0900
Subject: [PATCH] [MLIR][SCF] Fix empty loop range in mlir::loopUnrollByFactor.
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 4 +-
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 15 ++++-
mlir/test/Dialect/SCF/loop-unroll.mlir | 75 ++++++++++++++++++++-
mlir/test/Transforms/scf-loop-unroll.mlir | 5 +-
4 files changed, 91 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index 4cc7e62e6fd585..a4b3f0ee802360 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -128,8 +128,8 @@ splitForOpAtPoint(RewriterBase &rewriter, scf::ForOp forOp, Value splitPoint);
/// Unrolls this for operation by the specified unroll factor. Returns the
/// unrolled main loop and the epilogue loop, if the loop is unrolled. Otherwise
/// returns failure if the loop cannot be unrolled either due to restrictions or
-/// due to invalid unroll factors. Requires positive loop bounds and step. If
-/// specified, annotates the Ops in each unrolled iteration by applying
+/// due to invalid unroll factors. Requires positive step. If specified,
+// annotates the Ops in each unrolled iteration by applying
/// `annotateFn`.
/// If `shouldPromoteIfSingleIteration` is true, the function will promote the
/// loop body up if this has turned into a single iteration loop.
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 7ba0a796c86315..2054e160d56555 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -513,6 +513,17 @@ FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
auto upperBound = forOp.getUpperBound();
Value diff =
arith::SubIOp::create(boundsBuilder, loc, upperBound, lowerBound);
+
+ Value zero = arith::ConstantOp::create(
+ boundsBuilder, loc, boundsBuilder.getZeroAttr(diff.getType()));
+ if (forOp.getUnsignedCmp()) {
+ Value nonEmpty =
+ arith::CmpIOp::create(boundsBuilder, loc, arith::CmpIPredicate::ult,
+ lowerBound, upperBound);
+ diff = arith::SelectOp::create(boundsBuilder, loc, nonEmpty, diff, zero);
+ } else
+ diff = arith::MaxSIOp::create(boundsBuilder, loc, diff, zero);
+
Value tripCount = ceilDivPositive(boundsBuilder, loc, diff, step);
Value unrollFactorCst = arith::ConstantOp::create(
boundsBuilder, loc,
@@ -522,11 +533,11 @@ FailureOr<UnrolledLoopInfo> mlir::loopUnrollByFactor(
// Compute tripCountEvenMultiple = tripCount - (tripCount % unrollFactor)
Value tripCountEvenMultiple =
arith::SubIOp::create(boundsBuilder, loc, tripCount, tripCountRem);
- // Compute upperBoundUnrolled = lowerBound + tripCountEvenMultiple * step
+ // Compute upperBoundUnrolled = UpperBound + tripCountEvenMultiple * step
upperBoundUnrolled = arith::AddIOp::create(
boundsBuilder, loc, lowerBound,
arith::MulIOp::create(boundsBuilder, loc, tripCountEvenMultiple, step));
- // Scale 'step' by 'unrollFactor'.
+ // Compute stepUnrolled = step * unrollFactor
stepUnrolled =
arith::MulIOp::create(boundsBuilder, loc, step, unrollFactorCst);
}
diff --git a/mlir/test/Dialect/SCF/loop-unroll.mlir b/mlir/test/Dialect/SCF/loop-unroll.mlir
index 91decbde075c79..842862edecfaf3 100644
--- a/mlir/test/Dialect/SCF/loop-unroll.mlir
+++ b/mlir/test/Dialect/SCF/loop-unroll.mlir
@@ -24,9 +24,11 @@ func.func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
// UNROLL-BY-2-SAME: %[[MEM:.*3]]: memref<?xf32>
//
// UNROLL-BY-2-DAG: %[[V0:.*]] = arith.subi %[[UB]], %[[LB]] : index
+// UNROLL-BY-2-DAG: %[[C0:.*]] = arith.constant 0 : index
+// UNROLL-BY-2-DAG: %[[V0_NONNEG:.*]] = arith.maxsi %[[V0]], %[[C0]] : index
// UNROLL-BY-2-DAG: %[[C1:.*]] = arith.constant 1 : index
// UNROLL-BY-2-DAG: %[[V1:.*]] = arith.subi %[[STEP]], %[[C1]] : index
-// UNROLL-BY-2-DAG: %[[V2:.*]] = arith.addi %[[V0]], %[[V1]] : index
+// UNROLL-BY-2-DAG: %[[V2:.*]] = arith.addi %[[V0_NONNEG]], %[[V1]] : index
// Compute trip count in V3.
// UNROLL-BY-2-DAG: %[[V3:.*]] = arith.divui %[[V2]], %[[STEP]] : index
// Store unroll factor in C2.
@@ -57,9 +59,11 @@ func.func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
// UNROLL-BY-3-SAME: %[[MEM:.*3]]: memref<?xf32>
//
// UNROLL-BY-3-DAG: %[[V0:.*]] = arith.subi %[[UB]], %[[LB]] : index
+// UNROLL-BY-3-DAG: %[[C0:.*]] = arith.constant 0 : index
+// UNROLL-BY-3-DAG: %[[V0_NONNEG:.*]] = arith.maxsi %[[V0]], %[[C0]] : index
// UNROLL-BY-3-DAG: %[[C1:.*]] = arith.constant 1 : index
// UNROLL-BY-3-DAG: %[[V1:.*]] = arith.subi %[[STEP]], %[[C1]] : index
-// UNROLL-BY-3-DAG: %[[V2:.*]] = arith.addi %[[V0]], %[[V1]] : index
+// UNROLL-BY-3-DAG: %[[V2:.*]] = arith.addi %[[V0_NONNEG]], %[[V1]] : index
// Compute trip count in V3.
// UNROLL-BY-3-DAG: %[[V3:.*]] = arith.divui %[[V2]], %[[STEP]] : index
// Store unroll factor in C3.
@@ -156,6 +160,73 @@ func.func @dynamic_loop_unroll_inner_by_2(
// UNROLL-INNER-BY-2-NEXT: }
// UNROLL-INNER-BY-2-NEXT: return
+func.func @dynamic_unsigned_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
+ %arg3: memref<?xf32>) {
+ %0 = arith.constant 7.0 : f32
+ scf.for unsigned %i0 = %arg0 to %arg1 step %arg2 {
+ memref.store %0, %arg3[%i0] : memref<?xf32>
+ }
+ return
+}
+// UNROLL-BY-2-LABEL: func @dynamic_unsigned_loop_unroll
+// UNROLL-BY-2-SAME: %[[LB:.*0]]: index,
+// UNROLL-BY-2-SAME: %[[UB:.*1]]: index,
+// UNROLL-BY-2-SAME: %[[STEP:.*2]]: index,
+// UNROLL-BY-2-SAME: %[[MEM:.*3]]: memref<?xf32>
+//
+// UNROLL-BY-2-DAG: %[[DIFF:.*]] = arith.subi %[[UB]], %[[LB]] : index
+// UNROLL-BY-2-DAG: %[[C0:.*]] = arith.constant 0 : index
+// UNROLL-BY-2-DAG: %[[NONEMPTY:.*]] = arith.cmpi ult, %[[LB]], %[[UB]] : index
+// UNROLL-BY-2-DAG: %[[NONNEG:.*]] = arith.select %[[NONEMPTY]], %[[DIFF]], %[[C0]] : index
+// UNROLL-BY-2-DAG: %[[C1:.*]] = arith.constant 1 : index
+// UNROLL-BY-2-DAG: %[[STEP_MINUS_ONE:.*]] = arith.subi %[[STEP]], %[[C1]] : index
+// UNROLL-BY-2-DAG: %[[SUM:.*]] = arith.addi %[[NONNEG]], %[[STEP_MINUS_ONE]] : index
+// UNROLL-BY-2-DAG: %[[TRIP_COUNT:.*]] = arith.divui %[[SUM]], %[[STEP]] : index
+// UNROLL-BY-2-DAG: %[[C2:.*]] = arith.constant 2 : index
+// UNROLL-BY-2-DAG: %[[REM:.*]] = arith.remsi %[[TRIP_COUNT]], %[[C2]] : index
+// UNROLL-BY-2-DAG: %[[EVEN:.*]] = arith.subi %[[TRIP_COUNT]], %[[REM]] : index
+// UNROLL-BY-2-DAG: %[[UB_OFFSET:.*]] = arith.muli %[[EVEN]], %[[STEP]] : index
+// UNROLL-BY-2-DAG: %[[UNROLLED_UB:.*]] = arith.addi %[[LB]], %[[UB_OFFSET]] : index
+// UNROLL-BY-2-DAG: %[[UNROLLED_STEP:.*]] = arith.muli %[[STEP]], %[[C2]] : index
+// UNROLL-BY-2: scf.for unsigned %[[IV:.*]] = %[[LB]] to %[[UNROLLED_UB]] step %[[UNROLLED_STEP]] {
+// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>
+// UNROLL-BY-2: memref.store
+// UNROLL-BY-2-NEXT: }
+// UNROLL-BY-2-NEXT: scf.for unsigned %[[IV:.*]] = %[[UNROLLED_UB]] to %[[UB]] step %[[STEP]] {
+// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>
+// UNROLL-BY-2-NEXT: }
+// UNROLL-BY-2-NEXT: return
+
+// UNROLL-BY-3-LABEL: func @dynamic_unsigned_loop_unroll
+// UNROLL-BY-3-SAME: %[[LB:.*0]]: index,
+// UNROLL-BY-3-SAME: %[[UB:.*1]]: index,
+// UNROLL-BY-3-SAME: %[[STEP:.*2]]: index,
+// UNROLL-BY-3-SAME: %[[MEM:.*3]]: memref<?xf32>
+//
+// UNROLL-BY-3-DAG: %[[DIFF:.*]] = arith.subi %[[UB]], %[[LB]] : index
+// UNROLL-BY-3-DAG: %[[C0:.*]] = arith.constant 0 : index
+// UNROLL-BY-3-DAG: %[[NONEMPTY:.*]] = arith.cmpi ult, %[[LB]], %[[UB]] : index
+// UNROLL-BY-3-DAG: %[[NONNEG:.*]] = arith.select %[[NONEMPTY]], %[[DIFF]], %[[C0]] : index
+// UNROLL-BY-3-DAG: %[[C1:.*]] = arith.constant 1 : index
+// UNROLL-BY-3-DAG: %[[STEP_MINUS_ONE:.*]] = arith.subi %[[STEP]], %[[C1]] : index
+// UNROLL-BY-3-DAG: %[[SUM:.*]] = arith.addi %[[NONNEG]], %[[STEP_MINUS_ONE]] : index
+// UNROLL-BY-3-DAG: %[[TRIP_COUNT:.*]] = arith.divui %[[SUM]], %[[STEP]] : index
+// UNROLL-BY-3-DAG: %[[C3:.*]] = arith.constant 3 : index
+// UNROLL-BY-3-DAG: %[[REM:.*]] = arith.remsi %[[TRIP_COUNT]], %[[C3]] : index
+// UNROLL-BY-3-DAG: %[[EVEN:.*]] = arith.subi %[[TRIP_COUNT]], %[[REM]] : index
+// UNROLL-BY-3-DAG: %[[UB_OFFSET:.*]] = arith.muli %[[EVEN]], %[[STEP]] : index
+// UNROLL-BY-3-DAG: %[[UNROLLED_UB:.*]] = arith.addi %[[LB]], %[[UB_OFFSET]] : index
+// UNROLL-BY-3-DAG: %[[UNROLLED_STEP:.*]] = arith.muli %[[STEP]], %[[C3]] : index
+// UNROLL-BY-3: scf.for unsigned %[[IV:.*]] = %[[LB]] to %[[UNROLLED_UB]] step %[[UNROLLED_STEP]] {
+// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>
+// UNROLL-BY-3: memref.store
+// UNROLL-BY-3: memref.store
+// UNROLL-BY-3-NEXT: }
+// UNROLL-BY-3-NEXT: scf.for unsigned %[[IV:.*]] = %[[UNROLLED_UB]] to %[[UB]] step %[[STEP]] {
+// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>
+// UNROLL-BY-3-NEXT: }
+// UNROLL-BY-3-NEXT: return
+
// Test that no epilogue clean-up loop is generated because the trip count is
// a multiple of the unroll factor.
func.func @static_loop_unroll_by_2(%arg0 : memref<?xf32>) {
diff --git a/mlir/test/Transforms/scf-loop-unroll.mlir b/mlir/test/Transforms/scf-loop-unroll.mlir
index db96c659c49fbc..e78c9c9d2ae972 100644
--- a/mlir/test/Transforms/scf-loop-unroll.mlir
+++ b/mlir/test/Transforms/scf-loop-unroll.mlir
@@ -36,8 +36,9 @@ func.func @scf_loop_unroll_double_symbolic_ub(%arg0 : f32, %arg1 : f32, %n : ind
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
- // CHECK-NEXT: %[[REM:.*]] = arith.remsi %[[N]], %[[C3]]
- // CHECK-NEXT: %[[UB:.*]] = arith.subi %[[N]], %[[REM]]
+ // CHECK-NEXT: %[[NONNEG:.*]] = arith.maxsi %[[N]], %[[C0]]
+ // CHECK-NEXT: %[[REM:.*]] = arith.remsi %[[NONNEG]], %[[C3]]
+ // CHECK-NEXT: %[[UB:.*]] = arith.subi %[[NONNEG]], %[[REM]]
// CHECK-NEXT: %[[SUM:.*]]:2 = scf.for {{.*}} = %[[C0]] to %[[UB]] step %[[C3]] iter_args
// CHECK: }
// CHECK-NEXT: %[[SUM1:.*]]:2 = scf.for {{.*}} = %[[UB]] to %[[N]] step %[[C1]] iter_args(%[[V1:.*]] = %[[SUM]]#0, %[[V2:.*]] = %[[SUM]]#1)
More information about the Mlir-commits
mailing list