Add an option for unrolling loops up to a factor.

[llitchev]

Summary

Currently, there is no option to allow for unrolling a loop up to a specific factor (specified by the user).
The code for doing that is there and there are benefits when unrolling is done to smaller loops (smaller than the factor specified).

[llitchev]

For some reason a code was removed in a diagnostic message. Added back.

[bondhugula]

mlir/include/mlir/Dialect/Affine/Passes.td

74

Format fix please.

mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp

48

This isn't properly indented I think - clang-format please.

130–132

Drop trivial braces.

mlir/test/Dialect/SCF/loop-unroll.mlir

252

Nit: drop extra line.

260

Nit: Doesn't make a difference for the test, but you could use affine.store.

271

tore -> store

[bondhugula]

mlir/test/Dialect/SCF/loop-unroll.mlir

266–272

Nit: Please indent these lines by one more shift.

[bondhugula]

Thanks for adding this option.

mlir/include/mlir/Dialect/Affine/Passes.h

 
 /// Creates a loop unrolling pass with the provided parameters.
 /// 'getUnrollFactor' is a function callback for clients to supply a function
 /// that computes an unroll factor - the callback takes precedence over unroll
 /// factors supplied through other means. If -1 is passed as the unrollFactor
 /// and no callback is provided, anything passed from the command-line (if at
 /// all) or the default unroll factor is used (LoopUnroll:kDefaultUnrollFactor).
 std::unique_ptr<OperationPass<FuncOp>> createLoopUnrollPass(
-    int unrollFactor = -1, bool unrollFull = false,
+    int unrollFactor = -1, bool unrollUpToFactor = false,
+    bool unrollFull = false,
     const std::function<unsigned(AffineForOp)> &getUnrollFactor = nullptr);
 
 /// Creates a loop unroll jam pass to unroll jam by the specified factor. A
 /// factor of -1 lets the pass use the default factor or the one on the command
 /// line if provided.
 std::unique_ptr<OperationPass<FuncOp>>
 createLoopUnrollAndJamPass(int unrollJamFactor = -1);
 

mlir/include/mlir/Dialect/Affine/Passes.td

 }
 
 def AffineLoopUnroll : FunctionPass<"affine-loop-unroll"> {
   let summary = "Unroll affine loops";
   let constructor = "mlir::createLoopUnrollPass()";
   let options = [
     Option<"unrollFactor", "unroll-factor", "unsigned", /*default=*/"4",
            "Use this unroll factor for all loops being unrolled">,
+    Option<"unrollUpToFactor", "unroll-up-to-factor", "bool", /*default=*/"false",
+           "Allow unroling up to the factor specicied">,
     Option<"unrollFull", "unroll-full", "bool", /*default=*/"false",
            "Fully unroll loops">,
     Option<"numRepetitions", "unroll-num-reps", "unsigned", /*default=*/"1",
            "Unroll innermost loops repeatedly this many times">,
     Option<"unrollFullThreshold", "unroll-full-threshold", "unsigned",
            /*default=*/"1",
            "Unroll all loops with trip count less than or equal to this">,
   ];

mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp

 //===- LoopUnroll.cpp - Code to perform loop unrolling --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements loop unrolling.
 //
 //===----------------------------------------------------------------------===//
-
 #include "PassDetail.h"
 #include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Transforms/LoopUtils.h"
struct LoopUnroll : public AffineLoopUnrollBase<LoopUnroll> {
   // precedence over command-line argument or passed argument.
   const std::function<unsigned(AffineForOp)> getUnrollFactor;
 
   LoopUnroll() : getUnrollFactor(nullptr) {}
   LoopUnroll(const LoopUnroll &other)
       : AffineLoopUnrollBase<LoopUnroll>(other),
         getUnrollFactor(other.getUnrollFactor) {}
   explicit LoopUnroll(
-      Optional<unsigned> unrollFactor = None, bool unrollFull = false,
+      Optional<unsigned> unrollFactor = None, bool unrollUpToFactor = false,
+      bool unrollFull = false,
       const std::function<unsigned(AffineForOp)> &getUnrollFactor = nullptr)
       : getUnrollFactor(getUnrollFactor) {
     if (unrollFactor)
       this->unrollFactor = *unrollFactor;
+    this->unrollUpToFactor = unrollUpToFactor;
     this->unrollFull = unrollFull;
   }
 
   void runOnFunction() override;
 
   /// Unroll this for op. Returns failure if nothing was done.
   LogicalResult runOnAffineForOp(AffineForOp forOp);
 };

 LogicalResult LoopUnroll::runOnAffineForOp(AffineForOp forOp) {
   // Use the function callback if one was provided.
   if (getUnrollFactor)
     return loopUnrollByFactor(forOp, getUnrollFactor(forOp));
   // Unroll completely if full loop unroll was specified.
   if (unrollFull)
     return loopUnrollFull(forOp);
   // Otherwise, unroll by the given unroll factor.
+  if (unrollUpToFactor) {
+    return loopUnrollUpToFactor(forOp, unrollFactor);
+  }
   return loopUnrollByFactor(forOp, unrollFactor);
 }
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLoopUnrollPass(
-    int unrollFactor, bool unrollFull,
+    int unrollFactor, bool unrollUpToFactor, bool unrollFull,
     const std::function<unsigned(AffineForOp)> &getUnrollFactor) {
   return std::make_unique<LoopUnroll>(
-      unrollFactor == -1 ? None : Optional<unsigned>(unrollFactor), unrollFull,
-      getUnrollFactor);
+      unrollFactor == -1 ? None : Optional<unsigned>(unrollFactor),
+      unrollUpToFactor, unrollFull, getUnrollFactor);
 }

mlir/lib/Transforms/Utils/LoopUtils.cpp

LogicalResult mlir::loopUnrollFull(AffineForOp forOp) {
   return failure();
 }
 
 /// Unrolls this loop by the specified factor or by the trip count (if constant)
 /// whichever is lower.
 LogicalResult mlir::loopUnrollUpToFactor(AffineForOp forOp,
                                          uint64_t unrollFactor) {
   Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
-
   if (mayBeConstantTripCount.hasValue() &&
       mayBeConstantTripCount.getValue() < unrollFactor)
     return loopUnrollByFactor(forOp, mayBeConstantTripCount.getValue());
   return loopUnrollByFactor(forOp, unrollFactor);
 }
 
 // Generates unrolled copies of AffineForOp or scf::ForOp 'loopBodyBlock', with
 // associated 'forOpIV' by 'unrollFactor', calling 'ivRemapFn' to remap

mlir/test/Dialect/SCF/loop-unroll.mlir

 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2' | FileCheck %s --check-prefix UNROLL-BY-2
 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=3' | FileCheck %s --check-prefix UNROLL-BY-3
 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=0' | FileCheck %s --check-prefix UNROLL-OUTER-BY-2
 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=1' | FileCheck %s --check-prefix UNROLL-INNER-BY-2
+// RUN: mlir-opt %s --affine-loop-unroll='unroll-factor=6 unroll-up-to-factor=true' | FileCheck %s --check-prefix UNROLL-UP-TO
 
 func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
                           %arg3: memref<?xf32>) {
   %0 = constant 7.0 : f32
   scf.for %i0 = %arg0 to %arg1 step %arg2 {
     store %0, %arg3[%i0] : memref<?xf32>
   }
   return

 //  UNROLL-BY-3-NEXT:    store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>
 //  UNROLL-BY-3-NEXT:    %[[C2_IV:.*]] = constant 2 : index
 //  UNROLL-BY-3-NEXT:    %[[V2:.*]] = muli %[[C1]], %[[C2_IV]] : index
 //  UNROLL-BY-3-NEXT:    %[[V3:.*]] = addi %[[IV]], %[[V2]] : index
 //  UNROLL-BY-3-NEXT:    store %{{.*}}, %[[MEM]][%[[V3]]] : memref<?xf32>
 //  UNROLL-BY-3-NEXT:  }
 //  UNROLL-BY-3-NEXT:  store %{{.*}}, %[[MEM]][%[[C9]]] : memref<?xf32>
 //  UNROLL-BY-3-NEXT:  return
+
+
+// Test unroll-up-to functionality.
+func @static_loop_unroll_up_to_factor(%arg0 : memref<?xf32>) {
+  %0 = constant 7.0 : f32
+  %lb = constant 0 : index
+  %ub = constant 2 : index
+  affine.for %i0 = %lb to %ub {
+    store %0, %arg0[%i0] : memref<?xf32>
+  }
+  return
+}
+// UNROLL-UP-TO-LABEL: func @static_loop_unroll_up_to_factor
+//  UNROLL-UP-TO-SAME:  %[[MEM:.*0]]: memref<?xf32>
+//  UNROLL-UP-TO-DAG:  %[[C0:.*]] = constant 0 : index
+//  UNROLL-UP-TO-DAG:  %[[C2:.*]] = constant 2 : index
+//  UNROLL-UP-TO-NEXT: %[[V0:.*]] = affine.apply {{.*}}
+//  UNROLL-UP-TO-NEXT: store %{{.*}}, %[[MEM]][%[[V0]]] : memref<?xf32>
+//  UNROLL-UP-TO-NEXT: %[[V1:.*]] = affine.apply {{.*}}
+//  UNROLL-UP-TO-NEXT: tore %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>
+//  UNROLL-UP-TO-NEXT: return

mlir/test/lib/Transforms/TestLoopUnrolling.cpp

void runOnFunction() override {
     });
     for (auto loop : loops) {
       loopUnrollByFactor(loop, unrollFactor);
     }
   }
   Option<uint64_t> unrollFactor{*this, "unroll-factor",
                                 llvm::cl::desc("Loop unroll factor."),
                                 llvm::cl::init(1)};
+  Option<bool> unrollUpToFactor{*this, "unroll-up-to-factor",
+                                llvm::cl::desc("Loop unroll up to factor."),
+                                llvm::cl::init(false)};
   Option<unsigned> loopDepth{*this, "loop-depth", llvm::cl::desc("Loop depth."),
                              llvm::cl::init(0)};
 };
 } // end namespace
 
 namespace mlir {
 void registerTestLoopUnrollingPass() {
   PassRegistration<TestLoopUnrollingPass>(
       "test-loop-unrolling", "Tests loop unrolling transformation");
 }
 } // namespace mlir