Change behavior of loopUnrollFull with unroll factor 1

Using loopUnrollFull with unroll factor 1 should promote the loop body as
opposed to doing nothing.

PiperOrigin-RevId: 214812126

diff --git a/mlir/lib/Transforms/LoopUnroll.cpp b/mlir/lib/Transforms/LoopUnroll.cpp
index cfbf5059659435fcefd05464d0f4fc8fbb89efa2..91944a9e51dc46c202334c0600fb24ffc0c3bc41 100644 (file)
--- a/mlir/lib/Transforms/LoopUnroll.cpp
+++ b/mlir/lib/Transforms/LoopUnroll.cpp
@@ -171,9 +171,14 @@ bool LoopUnroll::runOnForStmt(ForStmt *forStmt) {
 
 /// Unrolls this loop completely.
 bool mlir::loopUnrollFull(ForStmt *forStmt) {
-  Optional<uint64_t> tripCount = getConstantTripCount(*forStmt);
-  if (tripCount.hasValue())
-    return loopUnrollByFactor(forStmt, tripCount.getValue());
+  Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
+  if (mayBeConstantTripCount.hasValue()) {
+    uint64_t tripCount = mayBeConstantTripCount.getValue();
+    if (tripCount == 1) {
+      return promoteIfSingleIteration(forStmt);
+    }
+    return loopUnrollByFactor(forStmt, tripCount);
+  }
   return false;
 }
 

diff --git a/mlir/test/Transforms/unroll.mlir b/mlir/test/Transforms/unroll.mlir
index ee2c3d401608751138d359f6b1ed6a511b32089a..3948ae073abfb8c37dc8366127833c0f15047127 100644 (file)
--- a/mlir/test/Transforms/unroll.mlir
+++ b/mlir/test/Transforms/unroll.mlir
@@ -502,3 +502,15 @@ mlfunc @loop_nest_operand4(%N : affineint) {
   }
   return
 }
+
+// CHECK-LABEL: mlfunc @loop_nest_unroll_full() {
+mlfunc @loop_nest_unroll_full() {
+  // CHECK-NEXT: %0 = "foo"() : () -> i32
+  // CHECK-NEXT: %1 = "bar"() : () -> i32
+  // CHECK-NEXT:  return
+  for %i = 0 to 0 {
+    %x = "foo"() : () -> i32
+    %y = "bar"() : () -> i32
+  }
+  return
+} // CHECK }