[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:14:05 PDT 2026
https://github.com/sun-jacobi updated https://github.com/llvm/llvm-project/pull/237094
>From ceea5a0d76a3080cba8c0163cb02e183f5211b70 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