90 lines · cpp
1//===- LoopUnrollAndJam.cpp - Code to perform loop unroll and jam ---------===//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 unroll and jam. Unroll and jam is a transformation10// that improves locality, in particular, register reuse, while also improving11// operation level parallelism. The example below shows what it does in nearly12// the general case. Loop unroll and jam currently works if the bounds of the13// loops inner to the loop being unroll-jammed do not depend on the latter.14//15// Before After unroll and jam of i by factor 2:16//17// for i, step = 218// for i S1(i);19// S1; S2(i);20// S2; S1(i+1);21// for j S2(i+1);22// S3; for j23// S4; S3(i, j);24// S5; S4(i, j);25// S6; S3(i+1, j)26// S4(i+1, j)27// S5(i);28// S6(i);29// S5(i+1);30// S6(i+1);31//32// Note: 'if/else' blocks are not jammed. So, if there are loops inside if33// op's, bodies of those loops will not be jammed.34//===----------------------------------------------------------------------===//35 36#include "mlir/Dialect/Affine/Passes.h"37 38#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"39#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"40#include "mlir/Dialect/Affine/IR/AffineOps.h"41#include "mlir/Dialect/Affine/LoopUtils.h"42#include "llvm/Support/CommandLine.h"43#include <optional>44 45namespace mlir {46namespace affine {47#define GEN_PASS_DEF_AFFINELOOPUNROLLANDJAM48#include "mlir/Dialect/Affine/Passes.h.inc"49} // namespace affine50} // namespace mlir51 52#define DEBUG_TYPE "affine-loop-unroll-jam"53 54using namespace mlir;55using namespace mlir::affine;56 57namespace {58/// Loop unroll jam pass. Currently, this just unroll jams the first59/// outer loop in a Function.60struct LoopUnrollAndJam61 : public affine::impl::AffineLoopUnrollAndJamBase<LoopUnrollAndJam> {62 explicit LoopUnrollAndJam(63 std::optional<unsigned> unrollJamFactor = std::nullopt) {64 if (unrollJamFactor)65 this->unrollJamFactor = *unrollJamFactor;66 }67 68 void runOnOperation() override;69};70} // namespace71 72std::unique_ptr<InterfacePass<FunctionOpInterface>>73mlir::affine::createLoopUnrollAndJamPass(int unrollJamFactor) {74 return std::make_unique<LoopUnrollAndJam>(75 unrollJamFactor == -1 ? std::nullopt76 : std::optional<unsigned>(unrollJamFactor));77}78 79void LoopUnrollAndJam::runOnOperation() {80 if (getOperation().isExternal())81 return;82 83 // Currently, just the outermost loop from the first loop nest is84 // unroll-and-jammed by this pass. However, runOnAffineForOp can be called on85 // any for operation.86 auto &entryBlock = getOperation().front();87 if (auto forOp = dyn_cast<AffineForOp>(entryBlock.front()))88 (void)loopUnrollJamByFactor(forOp, unrollJamFactor);89}90