Move definitions of lopoUnroll* functions to LoopUtils.cpp.

These functions are declared in Transforms/LoopUtils.h (included to the
Transforms/Utils library) but were defined in the loop unrolling pass in
Transforms/LoopUnroll.cpp.  As a result, targets depending only on
TransformUtils library but not on Transforms could get link errors.  Move the
definitions to Transforms/Utils/LoopUtils.cpp where they should actually live.
This does not modify any code.

PiperOrigin-RevId: 221508882

diff --git a/mlir/lib/Transforms/LoopUnroll.cpp b/mlir/lib/Transforms/LoopUnroll.cpp
index dc8b2f10269481d78484fb4eba42af5d00943f1a..23dff377f7154498c2606c3113e91457ed8660c2 100644 (file)
--- a/mlir/lib/Transforms/LoopUnroll.cpp
+++ b/mlir/lib/Transforms/LoopUnroll.cpp
@@ -169,123 +169,6 @@ bool LoopUnroll::runOnForStmt(ForStmt *forStmt) {
   return loopUnrollByFactor(forStmt, 4);
 }
 
-/// Unrolls this loop completely.
-bool mlir::loopUnrollFull(ForStmt *forStmt) {
-  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;
-}
-
-/// Unrolls and jams this loop by the specified factor or by the trip count (if
-/// constant) whichever is lower.
-bool mlir::loopUnrollUpToFactor(ForStmt *forStmt, uint64_t unrollFactor) {
-  Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
-
-  if (mayBeConstantTripCount.hasValue() &&
-      mayBeConstantTripCount.getValue() < unrollFactor)
-    return loopUnrollByFactor(forStmt, mayBeConstantTripCount.getValue());
-  return loopUnrollByFactor(forStmt, unrollFactor);
-}
-
-/// Unrolls this loop by the specified factor. Returns true if the loop
-/// is successfully unrolled.
-bool mlir::loopUnrollByFactor(ForStmt *forStmt, uint64_t unrollFactor) {
-  assert(unrollFactor >= 1 && "unroll factor should be >= 1");
-
-  if (unrollFactor == 1 || forStmt->getStatements().empty())
-    return false;
-
-  auto lbMap = forStmt->getLowerBoundMap();
-  auto ubMap = forStmt->getUpperBoundMap();
-
-  // Loops with max/min expressions won't be unrolled here (the output can't be
-  // expressed as an MLFunction in the general case). However, the right way to
-  // do such unrolling for an MLFunction would be to specialize the loop for the
-  // 'hotspot' case and unroll that hotspot.
-  if (lbMap.getNumResults() != 1 || ubMap.getNumResults() != 1)
-    return false;
-
-  // Same operand list for lower and upper bound for now.
-  // TODO(bondhugula): handle bounds with different operand lists.
-  if (!forStmt->matchingBoundOperandList())
-    return false;
-
-  Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
-
-  // If the trip count is lower than the unroll factor, no unrolled body.
-  // TODO(bondhugula): option to specify cleanup loop unrolling.
-  if (mayBeConstantTripCount.hasValue() &&
-      mayBeConstantTripCount.getValue() < unrollFactor)
-    return false;
-
-  // Generate the cleanup loop if trip count isn't a multiple of unrollFactor.
-  if (getLargestDivisorOfTripCount(*forStmt) % unrollFactor != 0) {
-    DenseMap<const MLValue *, MLValue *> operandMap;
-    MLFuncBuilder builder(forStmt->getBlock(), ++StmtBlock::iterator(forStmt));
-    auto *cleanupForStmt = cast<ForStmt>(builder.clone(*forStmt, operandMap));
-    auto clLbMap = getCleanupLoopLowerBound(*forStmt, unrollFactor, &builder);
-    assert(clLbMap &&
-           "cleanup loop lower bound map for single result bound maps can "
-           "always be determined");
-    cleanupForStmt->setLowerBoundMap(clLbMap);
-    // Promote the loop body up if this has turned into a single iteration loop.
-    promoteIfSingleIteration(cleanupForStmt);
-
-    // Adjust upper bound.
-    auto unrolledUbMap =
-        getUnrolledLoopUpperBound(*forStmt, unrollFactor, &builder);
-    assert(unrolledUbMap &&
-           "upper bound map can alwayys be determined for an unrolled loop "
-           "with single result bounds");
-    forStmt->setUpperBoundMap(unrolledUbMap);
-  }
-
-  // Scale the step of loop being unrolled by unroll factor.
-  int64_t step = forStmt->getStep();
-  forStmt->setStep(step * unrollFactor);
-
-  // Builder to insert unrolled bodies right after the last statement in the
-  // body of 'forStmt'.
-  MLFuncBuilder builder(forStmt, StmtBlock::iterator(forStmt->end()));
-
-  // Keep a pointer to the last statement in the original block so that we know
-  // what to clone (since we are doing this in-place).
-  StmtBlock::iterator srcBlockEnd = std::prev(forStmt->end());
-
-  // Unroll the contents of 'forStmt' (append unrollFactor-1 additional copies).
-  for (unsigned i = 1; i < unrollFactor; i++) {
-    DenseMap<const MLValue *, MLValue *> operandMap;
-
-    // If the induction variable is used, create a remapping to the value for
-    // this unrolled instance.
-    if (!forStmt->use_empty()) {
-      // iv' = iv + 1/2/3...unrollFactor-1;
-      auto d0 = builder.getAffineDimExpr(0);
-      auto bumpMap = builder.getAffineMap(1, 0, {d0 + i * step}, {});
-      auto *ivUnroll =
-          builder.create<AffineApplyOp>(forStmt->getLoc(), bumpMap, forStmt)
-              ->getResult(0);
-      operandMap[forStmt] = cast<MLValue>(ivUnroll);
-    }
-
-    // Clone the original body of 'forStmt'.
-    for (auto it = forStmt->begin(); it != std::next(srcBlockEnd); it++) {
-      builder.clone(*it, operandMap);
-    }
-  }
-
-  // Promote the loop body up if this has turned into a single iteration loop.
-  promoteIfSingleIteration(forStmt);
-
-  return true;
-}
-
 FunctionPass *mlir::createLoopUnrollPass(int unrollFactor, int unrollFull) {
   return new LoopUnroll(unrollFactor == -1 ? None
                                            : Optional<unsigned>(unrollFactor),

diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index 2603ea8d806551e5b3d9ceb24302aabf5465e03a..55cb64b8ba9f48c31d0bc0d8e1abacd5a96e4afb 100644 (file)
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -314,3 +314,120 @@ UtilResult mlir::stmtBodySkew(ForStmt *forStmt, ArrayRef<uint64_t> delays,
 
   return UtilResult::Success;
 }
+
+/// Unrolls this loop completely.
+bool mlir::loopUnrollFull(ForStmt *forStmt) {
+  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;
+}
+
+/// Unrolls and jams this loop by the specified factor or by the trip count (if
+/// constant) whichever is lower.
+bool mlir::loopUnrollUpToFactor(ForStmt *forStmt, uint64_t unrollFactor) {
+  Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
+
+  if (mayBeConstantTripCount.hasValue() &&
+      mayBeConstantTripCount.getValue() < unrollFactor)
+    return loopUnrollByFactor(forStmt, mayBeConstantTripCount.getValue());
+  return loopUnrollByFactor(forStmt, unrollFactor);
+}
+
+/// Unrolls this loop by the specified factor. Returns true if the loop
+/// is successfully unrolled.
+bool mlir::loopUnrollByFactor(ForStmt *forStmt, uint64_t unrollFactor) {
+  assert(unrollFactor >= 1 && "unroll factor should be >= 1");
+
+  if (unrollFactor == 1 || forStmt->getStatements().empty())
+    return false;
+
+  auto lbMap = forStmt->getLowerBoundMap();
+  auto ubMap = forStmt->getUpperBoundMap();
+
+  // Loops with max/min expressions won't be unrolled here (the output can't be
+  // expressed as an MLFunction in the general case). However, the right way to
+  // do such unrolling for an MLFunction would be to specialize the loop for the
+  // 'hotspot' case and unroll that hotspot.
+  if (lbMap.getNumResults() != 1 || ubMap.getNumResults() != 1)
+    return false;
+
+  // Same operand list for lower and upper bound for now.
+  // TODO(bondhugula): handle bounds with different operand lists.
+  if (!forStmt->matchingBoundOperandList())
+    return false;
+
+  Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(*forStmt);
+
+  // If the trip count is lower than the unroll factor, no unrolled body.
+  // TODO(bondhugula): option to specify cleanup loop unrolling.
+  if (mayBeConstantTripCount.hasValue() &&
+      mayBeConstantTripCount.getValue() < unrollFactor)
+    return false;
+
+  // Generate the cleanup loop if trip count isn't a multiple of unrollFactor.
+  if (getLargestDivisorOfTripCount(*forStmt) % unrollFactor != 0) {
+    DenseMap<const MLValue *, MLValue *> operandMap;
+    MLFuncBuilder builder(forStmt->getBlock(), ++StmtBlock::iterator(forStmt));
+    auto *cleanupForStmt = cast<ForStmt>(builder.clone(*forStmt, operandMap));
+    auto clLbMap = getCleanupLoopLowerBound(*forStmt, unrollFactor, &builder);
+    assert(clLbMap &&
+           "cleanup loop lower bound map for single result bound maps can "
+           "always be determined");
+    cleanupForStmt->setLowerBoundMap(clLbMap);
+    // Promote the loop body up if this has turned into a single iteration loop.
+    promoteIfSingleIteration(cleanupForStmt);
+
+    // Adjust upper bound.
+    auto unrolledUbMap =
+        getUnrolledLoopUpperBound(*forStmt, unrollFactor, &builder);
+    assert(unrolledUbMap &&
+           "upper bound map can alwayys be determined for an unrolled loop "
+           "with single result bounds");
+    forStmt->setUpperBoundMap(unrolledUbMap);
+  }
+
+  // Scale the step of loop being unrolled by unroll factor.
+  int64_t step = forStmt->getStep();
+  forStmt->setStep(step * unrollFactor);
+
+  // Builder to insert unrolled bodies right after the last statement in the
+  // body of 'forStmt'.
+  MLFuncBuilder builder(forStmt, StmtBlock::iterator(forStmt->end()));
+
+  // Keep a pointer to the last statement in the original block so that we know
+  // what to clone (since we are doing this in-place).
+  StmtBlock::iterator srcBlockEnd = std::prev(forStmt->end());
+
+  // Unroll the contents of 'forStmt' (append unrollFactor-1 additional copies).
+  for (unsigned i = 1; i < unrollFactor; i++) {
+    DenseMap<const MLValue *, MLValue *> operandMap;
+
+    // If the induction variable is used, create a remapping to the value for
+    // this unrolled instance.
+    if (!forStmt->use_empty()) {
+      // iv' = iv + 1/2/3...unrollFactor-1;
+      auto d0 = builder.getAffineDimExpr(0);
+      auto bumpMap = builder.getAffineMap(1, 0, {d0 + i * step}, {});
+      auto *ivUnroll =
+          builder.create<AffineApplyOp>(forStmt->getLoc(), bumpMap, forStmt)
+              ->getResult(0);
+      operandMap[forStmt] = cast<MLValue>(ivUnroll);
+    }
+
+    // Clone the original body of 'forStmt'.
+    for (auto it = forStmt->begin(); it != std::next(srcBlockEnd); it++) {
+      builder.clone(*it, operandMap);
+    }
+  }
+
+  // Promote the loop body up if this has turned into a single iteration loop.
+  promoteIfSingleIteration(forStmt);
+
+  return true;
+}