Don't fail if unable to promote loops during unrolling

[nimiwio]

Summary

When the unroll factor is 1, we should only fail "unrolling" when the trip count also is determined to be 1 and it is unable to be promoted.

[nimiwio]

Sorry didn't mean to send this for review quite yet, I still need to tweak this a bit, and add tests.

[nimiwio]

Add checks for trip count and tests

[bondhugula]

Is this ready for review?

mlir/test/Transforms/scf-loop-unroll.mlir

53

Use test.foo. (That will eventually make it easier to get rid of all unregistered ops.)

[nimiwio]

Use test.foo

[nimiwio]

In D115365#3186504, @bondhugula wrote:

Is this ready for review?

Yes this is ready now for review

[bondhugula]

mlir/test/Dialect/Affine/unroll.mlir

634

We don't need this test case - it's redundant and isn't testing anything more?

637

Please use test.foo.

641

You can drop the %{{.*}} = part.

[bondhugula]

mlir/test/Transforms/scf-loop-unroll.mlir

60–76

Likewise.

[nimiwio]

Address comments

[nimiwio]

Thank you for the review! Can this be landed?

mlir/lib/Transforms/Utils/LoopUtils.cpp

 
 /// Unrolls this loop by the specified factor. Returns success if the loop
 /// is successfully unrolled.
 LogicalResult mlir::loopUnrollByFactor(
     AffineForOp forOp, uint64_t unrollFactor,
     function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn) {
   assert(unrollFactor > 0 && "unroll factor should be positive");
 
-  if (unrollFactor == 1)
-    return promoteIfSingleIteration(forOp);
+  Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
+  if (unrollFactor == 1) {
+    if (mayBeConstantTripCount.hasValue() &&
+        mayBeConstantTripCount.getValue() == 1 &&
+        failed(promoteIfSingleIteration(forOp)))
+      return failure();
+    return success();
+  }
 
   // Nothing in the loop body other than the terminator.
   if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
     return success();
 
   // If the trip count is lower than the unroll factor, no unrolled body.
   // TODO: option to specify cleanup loop unrolling.
-  Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
   if (mayBeConstantTripCount.hasValue() &&
       mayBeConstantTripCount.getValue() < unrollFactor)
     return failure();
 
   // Generate the cleanup loop if trip count isn't a multiple of unrollFactor.
   if (getLargestDivisorOfTripCount(forOp) % unrollFactor != 0) {
     // Loops where the lower bound is a max expression or the upper bound is
     // a min expression and the trip count doesn't divide the unroll factor
LogicalResult mlir::loopUnrollByFactor(
   return success();
 }
 
 /// Unrolls 'forOp' by 'unrollFactor', returns success if the loop is unrolled.
 LogicalResult mlir::loopUnrollByFactor(
     scf::ForOp forOp, uint64_t unrollFactor,
     function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn) {
   assert(unrollFactor > 0 && "expected positive unroll factor");
-  if (unrollFactor == 1)
-    return promoteIfSingleIteration(forOp);
 
   // Return if the loop body is empty.
   if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
     return success();
 
   // Compute tripCount = ceilDiv((upperBound - lowerBound), step) and populate
   // 'upperBoundUnrolled' and 'stepUnrolled' for static and dynamic cases.
   OpBuilder boundsBuilder(forOp);
LogicalResult mlir::loopUnrollByFactor(
   if (lbCstOp && ubCstOp && stepCstOp) {
     // Constant loop bounds computation.
     int64_t lbCst = lbCstOp.value();
     int64_t ubCst = ubCstOp.value();
     int64_t stepCst = stepCstOp.value();
     assert(lbCst >= 0 && ubCst >= 0 && stepCst >= 0 &&
            "expected positive loop bounds and step");
     int64_t tripCount = mlir::ceilDiv(ubCst - lbCst, stepCst);
+
+    if (unrollFactor == 1) {
+      if (tripCount == 1 && failed(promoteIfSingleIteration(forOp)))
+        return failure();
+      return success();
+    }
+
     int64_t tripCountEvenMultiple = tripCount - (tripCount % unrollFactor);
     int64_t upperBoundUnrolledCst = lbCst + tripCountEvenMultiple * stepCst;
     assert(upperBoundUnrolledCst <= ubCst);
     int64_t stepUnrolledCst = stepCst * unrollFactor;
 
     // Create constant for 'upperBoundUnrolled' and set epilogue loop flag.
     generateEpilogueLoop = upperBoundUnrolledCst < ubCst;
     if (generateEpilogueLoop)
struct JamBlockGatherer {
   }
 };
 
 /// Unrolls and jams this loop by the specified factor.
 LogicalResult mlir::loopUnrollJamByFactor(AffineForOp forOp,
                                           uint64_t unrollJamFactor) {
   assert(unrollJamFactor > 0 && "unroll jam factor should be positive");
 
-  if (unrollJamFactor == 1)
-    return promoteIfSingleIteration(forOp);
+  Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
+  if (unrollJamFactor == 1) {
+    if (mayBeConstantTripCount.hasValue() &&
+        mayBeConstantTripCount.getValue() == 1 &&
+        failed(promoteIfSingleIteration(forOp)))
+      return failure();
+    return success();
+  }
 
   // Nothing in the loop body other than the terminator.
   if (llvm::hasSingleElement(forOp.getBody()->getOperations()))
     return success();
 
-  Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
   // If the trip count is lower than the unroll jam factor, no unroll jam.
   if (mayBeConstantTripCount.hasValue() &&
       mayBeConstantTripCount.getValue() < unrollJamFactor) {
     LLVM_DEBUG(llvm::dbgs() << "[failed] trip count < unroll-jam factor\n");
     return failure();
   }
 
   // If any control operand of any inner loop of `forOp` is defined within

mlir/test/Dialect/Affine/unroll.mlir

affine.for %i = 0 to 1 {
     %x = "foo"(%i) : (index) -> i32
   }
   return
 // UNROLL-BY-1-NEXT: %c0 = arith.constant 0 : index
 // UNROLL-BY-1-NEXT: %0 = "foo"(%c0) : (index) -> i32
 // UNROLL-BY-1-NEXT: return
 }
 
+// UNROLL-BY-1-LABEL: func @unroll_by_one_should_succeed_if_no_promotion()
+func @unroll_by_one_should_succeed_if_no_promotion() {
+  affine.for %i = 0 to 2 {
+    %x = "foo"(%i) : (index) -> i32
+  }
+  return
+// UNROLL-BY-1-NEXT:   affine.for %[[IV:.*]] = 0 to 2 {
+// UNROLL-BY-1-NEXT:     %{{.*}} = "foo"(%[[IV]]) : (index) -> i32
+// UNROLL-BY-1-NEXT:   }
+}
+
 // Test unrolling with affine.for iter_args.
 
 // UNROLL-BY-4-LABEL: loop_unroll_with_iter_args_and_cleanup
 func @loop_unroll_with_iter_args_and_cleanup(%arg0 : f32, %arg1 : f32, %n : index) -> (f32,f32) {
   %cf1 = arith.constant 1.0 : f32
   %cf2 = arith.constant 2.0 : f32
   %sum:2 = affine.for %iv = 0 to 10 iter_args(%i0 = %arg0, %i1 = %arg1) -> (f32, f32) {
     %sum0 = arith.addf %i0, %cf1 : f32

mlir/test/Transforms/scf-loop-unroll.mlir

 // RUN: mlir-opt %s --test-loop-unrolling="unroll-factor=3" -split-input-file -canonicalize | FileCheck %s
+// RUN: mlir-opt %s --test-loop-unrolling="unroll-factor=1" -split-input-file -canonicalize | FileCheck %s --check-prefix UNROLL-BY-1
 
 // CHECK-LABEL: scf_loop_unroll_single
 func @scf_loop_unroll_single(%arg0 : f32, %arg1 : f32) -> f32 {
   %from = arith.constant 0 : index
   %to = arith.constant 10 : index
   %step = arith.constant 1 : index
   %sum = scf.for %iv = %from to %to step %step iter_args(%sum_iter = %arg0) -> (f32) {
     %next = arith.addf %sum_iter, %arg1 : f32
func @scf_loop_unroll_double_symbolic_ub(%arg0 : f32, %arg1 : f32, %n : index) -> (f32,f32) {
   // CHECK-NEXT: %[[REM:.*]] = arith.remsi %[[N]], %[[C3]]
   // CHECK-NEXT: %[[UB:.*]] = arith.subi %[[N]], %[[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)
   // CHECK:      }
   // CHECK-NEXT: return %[[SUM1]]#0, %[[SUM1]]#1
 }
+
+// UNROLL-BY-1-LABEL: scf_loop_unroll_factor_1_promote
+func @scf_loop_unroll_factor_1_promote() -> () {
+  %step = arith.constant 1 : index
+  %lo = arith.constant 0 : index
+  %hi = arith.constant 1 : index
+  scf.for %i = %lo to %hi step %step {
+    %x = "test.foo"(%i) : (index) -> i32
+  }
+  return
+  // UNROLL-BY-1-NEXT: %[[C0:.*]] = arith.constant 0 : index
+  // UNROLL-BY-1-NEXT: %{{.*}} = "test.foo"(%[[C0]]) : (index) -> i32
+}
+
+// UNROLL-BY-1-LABEL: scf_loop_unroll_factor_1_no_promotion
+func @scf_loop_unroll_factor_1_no_promotion() -> () {
+  %step = arith.constant 1 : index
+  %lo = arith.constant 0 : index
+  %hi = arith.constant 2 : index
+  scf.for %i = %lo to %hi step %step {
+    %x = "test.foo"(%i) : (index) -> i32
+  }
+  return
+  // UNROLL-BY-1-DAG: %[[STEP:.*]] = arith.constant 1 : index
+  // UNROLL-BY-1-DAG: %[[LO:.*]] = arith.constant 0 : index
+  // UNROLL-BY-1-DAG: %[[HI:.*]] = arith.constant 2 : index
+  // UNROLL-BY-1-NEXT: scf.for %[[IV:.*]] = %[[LO]] to %[[HI]] step %[[STEP]] {
+  // UNROLL-BY-1-NEXT:  %{{.*}} = "test.foo"(%[[IV]]) : (index) -> i32
+  // UNROLL-BY-1-NEXT: }
+  // UNROLL-BY-1-NEXT: return
+}