SCF/Transforms/LoopPipelining.cpp

//===- LoopPipelining.cpp - Code to perform loop software pipelining-------===//
//
// 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 software pipelining
//
//===----------------------------------------------------------------------===//

#include "PassDetail.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/Dialect/SCF/Transforms.h"
#include "mlir/Dialect/SCF/Utils.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Support/MathExtras.h"

using namespace mlir;
using namespace mlir::scf;

namespace {

/// Helper to keep internal information during pipelining transformation.
struct LoopPipelinerInternal {
  /// Coarse liverange information for ops used across stages.
  struct LiverangeInfo {
    unsigned lastUseStage = 0;
    unsigned defStage = 0;
  };

protected:
  ForOp forOp;
  unsigned maxStage = 0;
  DenseMap<Operation *, unsigned> stages;
  std::vector<Operation *> opOrder;
  int64_t ub;
  int64_t lb;
  int64_t step;

  // When peeling the kernel we generate several version of each value for
  // different stage of the prologue. This map tracks the mapping between
  // original Values in the loop and the different versions
  // peeled from the loop.
  DenseMap<Value, llvm::SmallVector<Value>> valueMapping;

  /// Assign a value to `valueMapping`, this means `val` represents the version
  /// `idx` of `key` in the epilogue.
  void setValueMapping(Value key, Value el, int64_t idx);

public:
  /// Initalize the information for the given `op`, return true if it
  /// satisfies the pre-condition to apply pipelining.
  bool initializeLoopInfo(ForOp op, const PipeliningOption &options);
  /// Emits the prologue, this creates `maxStage - 1` part which will contain
  /// operations from stages [0; i], where i is the part index.
  void emitPrologue(PatternRewriter &rewriter);
  /// Gather liverange information for Values that are used in a different stage
  /// than its definition.
  llvm::MapVector<Value, LiverangeInfo> analyzeCrossStageValues();
  scf::ForOp createKernelLoop(
      const llvm::MapVector<Value, LiverangeInfo> &crossStageValues,
      PatternRewriter &rewriter,
      llvm::DenseMap<std::pair<Value, unsigned>, unsigned> &loopArgMap);
  /// Emits the pipelined kernel. This clones loop operations following user
  /// order and remaps operands defined in a different stage as their use.
  void createKernel(
      scf::ForOp newForOp,
      const llvm::MapVector<Value, LiverangeInfo> &crossStageValues,
      const llvm::DenseMap<std::pair<Value, unsigned>, unsigned> &loopArgMap,
      PatternRewriter &rewriter);
  /// Emits the epilogue, this creates `maxStage - 1` part which will contain
  /// operations from stages [i; maxStage], where i is the part index.
  void emitEpilogue(PatternRewriter &rewriter);
};

bool LoopPipelinerInternal::initializeLoopInfo(
    ForOp op, const PipeliningOption &options) {
  forOp = op;
  auto upperBoundCst = forOp.upperBound().getDefiningOp<ConstantIndexOp>();
  auto lowerBoundCst = forOp.lowerBound().getDefiningOp<ConstantIndexOp>();
  auto stepCst = forOp.step().getDefiningOp<ConstantIndexOp>();
  if (!upperBoundCst || !lowerBoundCst || !stepCst)
    return false;
  ub = upperBoundCst.getValue();
  lb = lowerBoundCst.getValue();
  step = stepCst.getValue();
  int64_t numIteration = ceilDiv(ub - lb, step);
  std::vector<std::pair<Operation *, unsigned>> schedule;
  options.getScheduleFn(forOp, schedule);
  if (schedule.empty())
    return false;

  opOrder.reserve(schedule.size());
  for (auto &opSchedule : schedule) {
    maxStage = std::max(maxStage, opSchedule.second);
    stages[opSchedule.first] = opSchedule.second;
    opOrder.push_back(opSchedule.first);
  }
  if (numIteration <= maxStage)
    return false;

  // All operations need to have a stage.
  if (forOp
          .walk([this](Operation *op) {
            if (op != forOp.getOperation() && !isa<scf::YieldOp>(op) &&
                stages.find(op) == stages.end())
              return WalkResult::interrupt();
            return WalkResult::advance();
          })
          .wasInterrupted())
    return false;

  // TODO: Add support for loop with operands.
  if (forOp.getNumIterOperands() > 0)
    return false;

  return true;
}

void LoopPipelinerInternal::emitPrologue(PatternRewriter &rewriter) {
  for (int64_t i = 0; i < maxStage; i++) {
    // special handling for induction variable as the increment is implicit.
    Value iv = rewriter.create<ConstantIndexOp>(forOp.getLoc(), lb + i);
    setValueMapping(forOp.getInductionVar(), iv, i);
    for (Operation *op : opOrder) {
      if (stages[op] > i)
        continue;
      Operation *newOp = rewriter.clone(*op);
      for (unsigned opIdx = 0; opIdx < op->getNumOperands(); opIdx++) {
        auto it = valueMapping.find(op->getOperand(opIdx));
        if (it != valueMapping.end())
          newOp->setOperand(opIdx, it->second[i - stages[op]]);
      }
      for (unsigned destId : llvm::seq(unsigned(0), op->getNumResults())) {
        setValueMapping(op->getResult(destId), newOp->getResult(destId),
                        i - stages[op]);
      }
    }
  }
}

llvm::MapVector<Value, LoopPipelinerInternal::LiverangeInfo>
LoopPipelinerInternal::analyzeCrossStageValues() {
  llvm::MapVector<Value, LoopPipelinerInternal::LiverangeInfo> crossStageValues;
  for (Operation *op : opOrder) {
    unsigned stage = stages[op];
    for (OpOperand &operand : op->getOpOperands()) {
      Operation *def = operand.get().getDefiningOp();
      if (!def)
        continue;
      auto defStage = stages.find(def);
      if (defStage == stages.end() || defStage->second == stage)
        continue;
      assert(stage > defStage->second);
      LiverangeInfo &info = crossStageValues[operand.get()];
      info.defStage = defStage->second;
      info.lastUseStage = std::max(info.lastUseStage, stage);
    }
  }
  return crossStageValues;
}

scf::ForOp LoopPipelinerInternal::createKernelLoop(
    const llvm::MapVector<Value, LoopPipelinerInternal::LiverangeInfo>
        &crossStageValues,
    PatternRewriter &rewriter,
    llvm::DenseMap<std::pair<Value, unsigned>, unsigned> &loopArgMap) {
  // Creates the list of initial values associated to values used across
  // stages. The initial values come from the prologue created above.
  // Keep track of the kernel argument associated to each version of the
  // values passed to the kernel.
  auto newLoopArg = llvm::to_vector<8>(forOp.getIterOperands());
  for (auto escape : crossStageValues) {
    LiverangeInfo &info = escape.second;
    Value value = escape.first;
    for (unsigned stageIdx = 0; stageIdx < info.lastUseStage - info.defStage;
         stageIdx++) {
      Value valueVersion =
          valueMapping[value][maxStage - info.lastUseStage + stageIdx];
      assert(valueVersion);
      newLoopArg.push_back(valueVersion);
      loopArgMap[std::make_pair(value, info.lastUseStage - info.defStage -
                                           stageIdx)] = newLoopArg.size() - 1;
    }
  }

  // Create the new kernel loop. Since we need to peel `numStages - 1`
  // iteration we change the upper bound to remove those iterations.
  Value newUb =
      rewriter.create<ConstantIndexOp>(forOp.getLoc(), ub - maxStage * step);
  auto newForOp = rewriter.create<scf::ForOp>(
      forOp.getLoc(), forOp.lowerBound(), newUb, forOp.step(), newLoopArg);
  return newForOp;
}

void LoopPipelinerInternal::createKernel(
    scf::ForOp newForOp,
    const llvm::MapVector<Value, LoopPipelinerInternal::LiverangeInfo>
        &crossStageValues,
    const llvm::DenseMap<std::pair<Value, unsigned>, unsigned> &loopArgMap,
    PatternRewriter &rewriter) {
  valueMapping.clear();

  // Create the kernel, we clone instruction based on the order given by
  // user and remap operands coming from a previous stages.
  rewriter.setInsertionPoint(newForOp.getBody(), newForOp.getBody()->begin());
  BlockAndValueMapping mapping;
  mapping.map(forOp.getInductionVar(), newForOp.getInductionVar());
  for (Operation *op : opOrder) {
    int64_t useStage = stages[op];
    auto *newOp = rewriter.clone(*op, mapping);
    for (OpOperand &operand : op->getOpOperands()) {
      // Special case for the induction variable uses. We replace it with a
      // version incremented based on the stage where it is used.
      if (operand.get() == forOp.getInductionVar()) {
        rewriter.setInsertionPoint(newOp);
        Value offset = rewriter.create<ConstantIndexOp>(
            forOp.getLoc(), (maxStage - stages[op]) * step);
        Value iv = rewriter.create<AddIOp>(forOp.getLoc(),
                                           newForOp.getInductionVar(), offset);
        newOp->setOperand(operand.getOperandNumber(), iv);
        rewriter.setInsertionPointAfter(newOp);
        continue;
      }
      // For operands defined in a previous stage we need to remap it to use
      // the correct region argument. We look for the right version of the
      // Value based on the stage where it is used.
      Operation *def = operand.get().getDefiningOp();
      if (!def)
        continue;
      auto stageDef = stages.find(def);
      if (stageDef == stages.end() || stageDef->second == useStage)
        continue;
      auto remap = loopArgMap.find(
          std::make_pair(operand.get(), useStage - stageDef->second));
      assert(remap != loopArgMap.end());
      newOp->setOperand(operand.getOperandNumber(),
                        newForOp.getRegionIterArgs()[remap->second]);
    }
  }

  // Collect the Values that need to be returned by the forOp. For each
  // value we need to have `LastUseStage - DefStage` number of versions
  // returned.
  // We create a mapping between original values and the associated loop
  // returned values that will be needed by the epilogue.
  llvm::SmallVector<Value> yieldOperands;
  for (auto &it : crossStageValues) {
    int64_t version = maxStage - it.second.lastUseStage + 1;
    unsigned numVersionReturned = it.second.lastUseStage - it.second.defStage;
    // add the original verstion to yield ops.
    // If there is a liverange spanning across more than 2 stages we need to add
    // extra arg.
    for (unsigned i = 1; i < numVersionReturned; i++) {
      setValueMapping(it.first, newForOp->getResult(yieldOperands.size()),
                      version++);
      yieldOperands.push_back(
          newForOp.getBody()->getArguments()[yieldOperands.size() + 1 +
                                             newForOp.getNumInductionVars()]);
    }
    setValueMapping(it.first, newForOp->getResult(yieldOperands.size()),
                    version++);
    yieldOperands.push_back(mapping.lookupOrDefault(it.first));
  }
  rewriter.create<scf::YieldOp>(forOp.getLoc(), yieldOperands);
}

void LoopPipelinerInternal::emitEpilogue(PatternRewriter &rewriter) {
  // Emit different versions of the induction variable. They will be
  // removed by dead code if not used.
  for (int64_t i = 0; i < maxStage; i++) {
    Value newlastIter = rewriter.create<ConstantIndexOp>(
        forOp.getLoc(), lb + step * ((((ub - 1) - lb) / step) - i));
    setValueMapping(forOp.getInductionVar(), newlastIter, maxStage - i);
  }
  // Emit `maxStage - 1` epilogue part that includes operations fro stages
  // [i; maxStage].
  for (int64_t i = 1; i <= maxStage; i++) {
    for (Operation *op : opOrder) {
      if (stages[op] < i)
        continue;
      Operation *newOp = rewriter.clone(*op);
      for (unsigned opIdx = 0; opIdx < op->getNumOperands(); opIdx++) {
        auto it = valueMapping.find(op->getOperand(opIdx));
        if (it != valueMapping.end()) {
          Value v = it->second[maxStage - stages[op] + i];
          assert(v);
          newOp->setOperand(opIdx, v);
        }
      }
      for (unsigned destId : llvm::seq(unsigned(0), op->getNumResults())) {
        setValueMapping(op->getResult(destId), newOp->getResult(destId),
                        maxStage - stages[op] + i);
      }
    }
  }
}

void LoopPipelinerInternal::setValueMapping(Value key, Value el, int64_t idx) {
  auto it = valueMapping.find(key);
  // If the value is not in the map yet add a vector big enough to store all
  // versions.
  if (it == valueMapping.end())
    it =
        valueMapping
            .insert(std::make_pair(key, llvm::SmallVector<Value>(maxStage + 1)))
            .first;
  it->second[idx] = el;
}

/// Generate a pipelined version of the scf.for loop based on the schedule given
/// as option. This applies the mechanical transformation of changing the loop
/// and generating the prologue/epilogue for the pipelining and doesn't make any
/// decision regarding the schedule.
/// Based on the option the loop is split into several stages.
/// The transformation assumes that the scheduling given by user is valid.
/// For example if we break a loop into 3 stages named S0, S1, S2 we would
/// generate the following code with the number in parenthesis the iteration
/// index:
/// S0(0)                        // Prologue
/// S0(1) S1(0)                  // Prologue
/// scf.for %I = %C0 to %N - 2 {
///  S0(I+2) S1(I+1) S2(I)       // Pipelined kernel
/// }
/// S1(N) S2(N-1)                // Epilogue
/// S2(N)                        // Epilogue
struct ForLoopPipelining : public OpRewritePattern<ForOp> {
  ForLoopPipelining(const PipeliningOption &options, MLIRContext *context)
      : OpRewritePattern<ForOp>(context), options(options) {}
  LogicalResult matchAndRewrite(ForOp forOp,
                                PatternRewriter &rewriter) const override {

    LoopPipelinerInternal pipeliner;
    if (!pipeliner.initializeLoopInfo(forOp, options))
      return failure();

    // 1. Emit prologue.
    pipeliner.emitPrologue(rewriter);

    // 2. Track values used across stages. When a value cross stages it will
    // need to be passed as loop iteration arguments.
    // We first collect the values that are used in a different stage than where
    // they are defined.
    llvm::MapVector<Value, LoopPipelinerInternal::LiverangeInfo>
        crossStageValues = pipeliner.analyzeCrossStageValues();

    // Mapping between original loop values used cross stage and the block
    // arguments associated after pipelining. A Value may map to several
    // arguments if its liverange spans across more than 2 stages.
    llvm::DenseMap<std::pair<Value, unsigned>, unsigned> loopArgMap;
    // 3. Create the new kernel loop and return the block arguments mapping.
    ForOp newForOp =
        pipeliner.createKernelLoop(crossStageValues, rewriter, loopArgMap);
    // Create the kernel block, order ops based on user choice and remap
    // operands.
    pipeliner.createKernel(newForOp, crossStageValues, loopArgMap, rewriter);

    // 4. Emit the epilogue after the new forOp.
    rewriter.setInsertionPointAfter(newForOp);
    pipeliner.emitEpilogue(rewriter);

    // 5. Erase the original loop and replace the uses with the epilogue output.
    if (forOp->getNumResults() > 0)
      rewriter.replaceOp(
          forOp, newForOp.getResults().take_front(forOp->getNumResults()));
    else
      rewriter.eraseOp(forOp);

    return success();
  }

protected:
  PipeliningOption options;
};

} // namespace

void mlir::scf::populateSCFLoopPipeliningPatterns(
    RewritePatternSet &patterns, const PipeliningOption &options) {
  patterns.add<ForLoopPipelining>(options, patterns.getContext());
}