153 lines · cpp
1//===- LoopUnroll.cpp - Code to perform loop unrolling --------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements loop unrolling.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Affine/Passes.h"14 15#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"16#include "mlir/Dialect/Affine/IR/AffineOps.h"17#include "mlir/Dialect/Affine/LoopUtils.h"18#include "llvm/Support/CommandLine.h"19#include <optional>20 21namespace mlir {22namespace affine {23#define GEN_PASS_DEF_AFFINELOOPUNROLL24#include "mlir/Dialect/Affine/Passes.h.inc"25} // namespace affine26} // namespace mlir27 28#define DEBUG_TYPE "affine-loop-unroll"29 30using namespace mlir;31using namespace mlir::affine;32 33namespace {34 35// TODO: this is really a test pass and should be moved out of dialect36// transforms.37 38/// Loop unrolling pass. Unrolls all innermost loops unless full unrolling and a39/// full unroll threshold was specified, in which case, fully unrolls all loops40/// with trip count less than the specified threshold. The latter is for testing41/// purposes, especially for testing outer loop unrolling.42struct LoopUnroll : public affine::impl::AffineLoopUnrollBase<LoopUnroll> {43 // Callback to obtain unroll factors; if this has a callable target, takes44 // precedence over command-line argument or passed argument.45 const std::function<unsigned(AffineForOp)> getUnrollFactor;46 47 LoopUnroll() : getUnrollFactor(nullptr) {}48 LoopUnroll(const LoopUnroll &other) = default;49 explicit LoopUnroll(50 std::optional<unsigned> unrollFactor = std::nullopt,51 bool unrollUpToFactor = false,52 const std::function<unsigned(AffineForOp)> &getUnrollFactor = nullptr)53 : getUnrollFactor(getUnrollFactor) {54 if (unrollFactor)55 this->unrollFactor = *unrollFactor;56 this->unrollUpToFactor = unrollUpToFactor;57 }58 59 void runOnOperation() override;60 61 /// Unroll this for op. Returns failure if nothing was done.62 LogicalResult runOnAffineForOp(AffineForOp forOp);63};64} // namespace65 66/// Returns true if no other affine.for ops are nested within `op`.67static bool isInnermostAffineForOp(AffineForOp op) {68 return !op.getBody()69 ->walk([&](AffineForOp nestedForOp) {70 return WalkResult::interrupt();71 })72 .wasInterrupted();73}74 75/// Gathers loops that have no affine.for's nested within.76static void gatherInnermostLoops(FunctionOpInterface f,77 SmallVectorImpl<AffineForOp> &loops) {78 f.walk([&](AffineForOp forOp) {79 if (isInnermostAffineForOp(forOp))80 loops.push_back(forOp);81 });82}83 84void LoopUnroll::runOnOperation() {85 if (!(unrollFactor.getValue() > 0 || unrollFactor.getValue() == -1)) {86 emitError(UnknownLoc::get(&getContext()),87 "Invalid option: 'unroll-factor' should be greater than 0 or "88 "equal to -1");89 return signalPassFailure();90 }91 FunctionOpInterface func = getOperation();92 if (func.isExternal())93 return;94 95 if (unrollFactor.getValue() == -1 && unrollFullThreshold.hasValue()) {96 // Store short loops as we walk.97 SmallVector<AffineForOp, 4> loops;98 99 // Gathers all loops with trip count <= minTripCount. Do a post order walk100 // so that loops are gathered from innermost to outermost (or else101 // unrolling an outer one may delete gathered inner ones).102 getOperation().walk([&](AffineForOp forOp) {103 std::optional<uint64_t> tripCount = getConstantTripCount(forOp);104 if (tripCount && *tripCount <= unrollFullThreshold)105 loops.push_back(forOp);106 });107 for (auto forOp : loops)108 (void)loopUnrollFull(forOp);109 return;110 }111 112 // If the call back is provided, we will recurse until no loops are found.113 SmallVector<AffineForOp, 4> loops;114 for (unsigned i = 0; i < numRepetitions || getUnrollFactor; i++) {115 loops.clear();116 gatherInnermostLoops(func, loops);117 if (loops.empty())118 break;119 bool unrolled = false;120 for (auto forOp : loops)121 unrolled |= succeeded(runOnAffineForOp(forOp));122 if (!unrolled)123 // Break out if nothing was unrolled.124 break;125 }126}127 128/// Unrolls a 'affine.for' op. Returns success if the loop was unrolled,129/// failure otherwise. The default unroll factor is 4.130LogicalResult LoopUnroll::runOnAffineForOp(AffineForOp forOp) {131 // Use the function callback if one was provided.132 if (getUnrollFactor)133 return loopUnrollByFactor(forOp, getUnrollFactor(forOp),134 /*annotateFn=*/nullptr, cleanUpUnroll);135 // Unroll completely if full loop unroll was specified.136 if (unrollFactor.getValue() == -1)137 return loopUnrollFull(forOp);138 // Otherwise, unroll by the given unroll factor.139 if (unrollUpToFactor)140 return loopUnrollUpToFactor(forOp, unrollFactor);141 return loopUnrollByFactor(forOp, unrollFactor, /*annotateFn=*/nullptr,142 cleanUpUnroll);143}144 145std::unique_ptr<InterfacePass<FunctionOpInterface>>146mlir::affine::createLoopUnrollPass(147 int unrollFactor, bool unrollUpToFactor,148 const std::function<unsigned(AffineForOp)> &getUnrollFactor) {149 return std::make_unique<LoopUnroll>(150 unrollFactor == -1 ? std::nullopt : std::optional<unsigned>(unrollFactor),151 unrollUpToFactor, getUnrollFactor);152}153