Linalg/Transforms/Loops.cpp

//===- Loops.cpp - conversion from Linalg named and generic ops to loops --===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "PassDetail.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/SCF/Transforms.h"
#include "mlir/Dialect/StandardOps/Utils/Utils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/TypeSwitch.h"

using namespace mlir;
using namespace mlir::linalg;

static SmallVector<Value> makeCanonicalAffineApplies(OpBuilder &b, Location loc,
                                                     AffineMap map,
                                                     ArrayRef<Value> vals) {
  if (map.isEmpty())
    return {};

  assert(map.getNumInputs() == vals.size());
  SmallVector<Value> res;
  res.reserve(map.getNumResults());
  auto dims = map.getNumDims();
  for (auto e : map.getResults()) {
    auto exprMap = AffineMap::get(dims, map.getNumSymbols(), e);
    SmallVector<Value> operands(vals.begin(), vals.end());
    canonicalizeMapAndOperands(&exprMap, &operands);
    res.push_back(b.create<AffineApplyOp>(loc, exprMap, operands));
  }
  return res;
}

template <typename LoadOpTy, typename StoreOpTy, typename OpType>
static void inlineRegionAndEmitStore(OpBuilder &b, Location loc, OpType op,
                                     ArrayRef<Value> indexedValues,
                                     ArrayRef<SmallVector<Value>> indexing,
                                     ArrayRef<Value> outputBuffers) {
  auto &block = op->getRegion(0).front();
  BlockAndValueMapping map;
  map.map(block.getArguments(), indexedValues);
  for (auto &op : block.without_terminator()) {
    auto *newOp = b.clone(op, map);
    map.map(op.getResults(), newOp->getResults());
  }

  Operation *terminator = block.getTerminator();
  for (OpOperand &operand : terminator->getOpOperands()) {
    Value toStore = map.lookupOrDefault(operand.get());
    b.create<StoreOpTy>(loc, toStore, outputBuffers[operand.getOperandNumber()],
                        indexing[operand.getOperandNumber()]);
  }
}

// Returns a pair that contains input indices and output indices of a
// SingleInputPoolingOp `op`.
struct InputAndOutputIndices {
  SmallVector<Value> inputs;
  SmallVector<Value> outputs;
};
template <typename SingleInputPoolingOp>
static InputAndOutputIndices
getInputAndOutputIndices(OpBuilder &b, Location loc, ArrayRef<Value> allIvs,
                         SingleInputPoolingOp op) {
  auto mapsRange = op.indexing_maps().template getAsRange<AffineMapAttr>();
  auto maps = llvm::to_vector<8>(
      llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
  return InputAndOutputIndices{
      makeCanonicalAffineApplies(b, loc, maps[0], allIvs),
      makeCanonicalAffineApplies(b, loc, maps[2], allIvs)};
}

/// Emits the MLIR for the scalar part of the generic op by:
///   1. Emitting load ops for each input and output view in order. This is
///      achieved by applying the appropriate input or output map to the
///      enclosing induction variables.
///   2. Emitting a call to `op.fun()` that takes as arguments the scalars
///      from point 1. above.
///   3. Emitting store ops to store the results of 2. to the output
///      views.
///
/// An example output may resemble:
///
/// ```
///    scf.for %i = %c0 to %0 step %c1 {
///      scf.for %j = %c0 to %1 step %c1 {
///        scf.for %k = %c0 to %4 step %c1 {
///          %11 = load %arg0[%i, %j] :
///            memref<?x?xf32, stride_specification>
///          %12 = load %arg1[%i, %j, %k] :
///            memref<?x?x?xf32, stride_specification>
///          %13 = load %arg2[%i, %k, %j] :
///            memref<?x?x?xf32, stride_specification>
///          %14:2 = call @foo(%11, %12, %13) : (f32, f32, f32) -> (f32, f32)
///          store %14#0, %arg1[%i, %j, %k] :
///            memref<?x?x?Xf32, stride_specification>
///          store %14#1, %arg2[%i, %k, %j] :
///            memref<?x?x?Xf32, stride_specification>
///       }
///      }
///    }
/// ```
template <typename LoadOpTy, typename StoreOpTy>
static void emitScalarImplementation(OpBuilder &b, Location loc,
                                     ArrayRef<Value> allIvs,
                                     LinalgOp linalgOp) {
  assert(linalgOp.hasBufferSemantics() &&
         "expected linalg op with buffer semantics");
  SmallVector<Value> indexedValues;
  indexedValues.reserve(linalgOp.getNumInputsAndOutputs());

  auto allIvsPlusDims = SmallVector<Value>(allIvs.begin(), allIvs.end());

  // TODO: Avoid the loads if the corresponding argument of the
  // region has no uses.
  // 1.a. Emit load from input operand or for scalars access the operand itself.
  for (OpOperand *inputOperand : linalgOp.getInputOperands()) {
    if (linalgOp.isScalar(inputOperand)) {
      indexedValues.push_back(inputOperand->get());
      continue;
    }
    auto indexing = makeCanonicalAffineApplies(
        b, loc, linalgOp.getTiedIndexingMap(inputOperand), allIvsPlusDims);
    indexedValues.push_back(
        b.create<LoadOpTy>(loc, inputOperand->get(), indexing));
  }
  // 1.b. Emit load from output views.
  for (OpOperand *outputOperand : linalgOp.getOutputOperands()) {
    SmallVector<Value> indexing = makeCanonicalAffineApplies(
        b, loc, linalgOp.getTiedIndexingMap(outputOperand), allIvsPlusDims);
    indexedValues.push_back(
        b.create<LoadOpTy>(loc, outputOperand->get(), indexing));
  }

  // TODO: When a region inliner exists, use it.
  // 2. Inline region, currently only works for a single basic block.
  // 3. Emit store.
  SmallVector<SmallVector<Value>, 8> indexing;
  SmallVector<Value> outputBuffers;
  for (OpOperand *outputOperand : linalgOp.getOutputBufferOperands()) {
    indexing.push_back(makeCanonicalAffineApplies(
        b, loc, linalgOp.getTiedIndexingMap(outputOperand), allIvsPlusDims));
    outputBuffers.push_back(outputOperand->get());
  }
  inlineRegionAndEmitStore<LoadOpTy, StoreOpTy>(b, loc, linalgOp, indexedValues,
                                                indexing, outputBuffers);
}

// Create a padded view into the given `input` tensor using the 'indices'
// to access the tensor. `skipPadding` lists the dimensions for which no padding
// is needed e.g. the non-spatial dimensions for convolutions.
Value getPaddedInput(OpBuilder &b, Location loc, Value input,
                     ArrayRef<Value> indices, ArrayRef<int> skipPadding,
                     Value padValue) {
  Value zeroIndex = b.create<ConstantIndexOp>(loc, 0);
  SmallVector<Value> conds;
  SmallVector<Value> clampedImIdx;
  for (auto iter : llvm::enumerate(indices)) {
    int idx = iter.index();
    auto dim = iter.value();
    if (is_contained(skipPadding, idx)) {
      clampedImIdx.push_back(dim);
      continue;
    }

    Value leftOutOfBound =
        b.create<CmpIOp>(loc, CmpIPredicate::slt, dim, zeroIndex);
    if (conds.empty())
      conds.push_back(leftOutOfBound);
    else
      conds.push_back(b.create<OrOp>(loc, conds.back(), leftOutOfBound));
    Value rightBound = createOrFoldDimOp(b, loc, input, idx);
    Value rightOutOfBound =
        b.create<CmpIOp>(loc, CmpIPredicate::sge, dim, rightBound);
    conds.push_back(b.create<OrOp>(loc, conds.back(), rightOutOfBound));

    // When padding is involved, the indices will only be shifted to negative,
    // so having a max op is enough.
    MLIRContext *ctx = input.getContext();
    AffineExpr m = getAffineDimExpr(/*position=*/0, ctx),
               zero = getAffineConstantExpr(0, ctx);
    AffineMap maxMap =
        AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>>{{m, zero}})
            .front();
    clampedImIdx.push_back(b.create<AffineMaxOp>(loc, maxMap, ValueRange{dim}));
  }

  Value readInput = b.create<memref::LoadOp>(loc, input, clampedImIdx);
  if (conds.empty())
    return readInput;

  return b.create<SelectOp>(loc, conds.back(), padValue, readInput);
}

namespace {

/// The padding value for a given Op depends on the semantics of the Op.
/// The identity value for ConvOp and PoolingSumOp is 0, for PoolingMaxOp is
/// -inf or minInt and for PoolingMinOp is inf or maxInt.
template <typename OpType> Attribute getPadValueAttr(Type type) {
  llvm_unreachable("Unexpected op type for getPadValueAttr");
  return {};
}

template <> Attribute getPadValueAttr<PoolingMaxOp>(Type type) {
  if (auto floatType = type.dyn_cast<FloatType>()) {
    return OpBuilder(type.getContext())
        .getFloatAttr(floatType, APFloat::getInf(floatType.getFloatSemantics(),
                                                 /*Negative*/ true));
  }
  if (auto intType = type.dyn_cast<IntegerType>()) {
    unsigned width = intType.getWidth();
    // The select instruction used to lower the PoolingMin uses a signed
    // comparison, use a signed constant irrespective of the signedness of the
    // integer type.
    return OpBuilder(type.getContext())
        .getIntegerAttr(intType, APInt::getSignedMinValue(width));
  }
  llvm_unreachable("Unsupported data type for PoolingMaxOp");
  return {};
}

template <> Attribute getPadValueAttr<PoolingMinOp>(Type type) {
  if (auto floatType = type.dyn_cast<FloatType>()) {
    return OpBuilder(type.getContext())
        .getFloatAttr(floatType,
                      APFloat::getInf(floatType.getFloatSemantics()));
  }
  if (auto intType = type.dyn_cast<IntegerType>()) {
    unsigned width = intType.getWidth();
    // The select instruction used to lower the PoolingMin uses a signed
    // comparison, use a signed constant irrespective of the signedness of the
    // integer type.
    return OpBuilder(type.getContext())
        .getIntegerAttr(intType, APInt::getSignedMaxValue(width));
  }
  llvm_unreachable("Unsupported data type for PoolingMinOp");
  return {};
}

template <> Attribute getPadValueAttr<PoolingSumOp>(Type type) {
  return OpBuilder(type.getContext()).getZeroAttr(type);
}

template <> Attribute getPadValueAttr<ConvOp>(Type type) {
  return OpBuilder(type.getContext()).getZeroAttr(type);
}

} // namespace

/// Returns true is `convOp` has a non-zero padding.
static bool hasPadding(ConvOp convOp) {
  for (unsigned i = 0, e = convOp.getNumSpatialDimensions(); i < e; ++i) {
    if (convOp.getLowPad(i) > 0 || convOp.getHighPad(i) > 0)
      return true;
  }
  return false;
}

template <typename LoadOpTy, typename StoreOpTy>
static void emitScalarImplementation(OpBuilder &b, Location loc,
                                     ArrayRef<Value> allIvs, ConvOp convOp) {
  assert(convOp.hasBufferSemantics() &&
         "expected linalg op with buffer semantics");
  auto mapsRange = convOp.indexing_maps().getAsRange<AffineMapAttr>();
  auto maps = llvm::to_vector<8>(
      llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
  SmallVector<Value> fIdx(makeCanonicalAffineApplies(b, loc, maps[0], allIvs));
  SmallVector<Value> imIdx(makeCanonicalAffineApplies(b, loc, maps[1], allIvs));
  SmallVector<Value> oIdx(makeCanonicalAffineApplies(b, loc, maps[2], allIvs));

  Value filter = convOp.filter(), output = convOp.output();

  // Emit scalar form. Padded conv involves an affine.max in the memory access
  // which is not allowed by affine.load. Override to use an MemRefIndexedValue
  // when there is non-zero padding.
  if (hasPadding(convOp)) {
    Type type = convOp.input().getType().cast<MemRefType>().getElementType();
    Value padValue =
        b.create<ConstantOp>(loc, type, getPadValueAttr<ConvOp>(type));
    Value paddedInput =
        getPaddedInput(b, loc, convOp.input(), imIdx,
                       /* Only need to pad the window dimensions */
                       {0, static_cast<int>(imIdx.size()) - 1}, padValue);
    Value filterVal = b.create<LoadOpTy>(loc, filter, fIdx);
    Value mulVal = ArithBuilder(b, loc).mul(filterVal, paddedInput);
    Value outputVal = b.create<LoadOpTy>(loc, output, oIdx);
    Value addVal = ArithBuilder(b, loc).add(mulVal, outputVal);
    b.create<StoreOpTy>(loc, addVal, output, oIdx);
  } else {
    Value inputVal = b.create<LoadOpTy>(loc, convOp.input(), imIdx);
    Value filterVal = b.create<LoadOpTy>(loc, filter, fIdx);
    Value mulVal = ArithBuilder(b, loc).mul(filterVal, inputVal);
    Value outputVal = b.create<LoadOpTy>(loc, output, oIdx);
    Value addVal = ArithBuilder(b, loc).add(mulVal, outputVal);
    b.create<StoreOpTy>(loc, addVal, output, oIdx);
  }
}

template <typename PoolingOp> static bool hasPadding(PoolingOp poolingOp) {
  for (unsigned i = 0, e = poolingOp.getNumWindowLoops(); i < e; ++i) {
    if (poolingOp.getLowPad(i) > 0 || poolingOp.getHighPad(i) > 0)
      return true;
  }
  return false;
}

template <typename LoadOpTy, typename StoreOpTy, typename PoolingOp>
static Value getPoolingInput(OpBuilder &b, Location loc, PoolingOp op,
                             ArrayRef<Value> inputIndices) {
  if (hasPadding(op)) {
    Type type =
        op.input().getType().template cast<MemRefType>().getElementType();
    Value padValue =
        b.create<ConstantOp>(loc, type, getPadValueAttr<PoolingOp>(type));
    return getPaddedInput(b, loc, op.input(), inputIndices,
                          /*Pad every dimension*/ {}, padValue);
  }
  return b.create<LoadOpTy>(loc, op.input(), inputIndices);
}

template <typename LoadOpTy, typename StoreOpTy, typename OpType>
void emitPoolingMinMaxScalarImplementation(OpBuilder &b, Location loc,
                                           ArrayRef<Value> allIvs, OpType op) {
  InputAndOutputIndices indices = getInputAndOutputIndices(b, loc, allIvs, op);
  Value lhs = b.create<LoadOpTy>(loc, op.output(), indices.outputs);
  Value rhs = getPoolingInput<LoadOpTy, StoreOpTy>(b, loc, op, indices.inputs);
  Value value = llvm::TypeSwitch<Operation *, Value>(op)
                    .Case([&](PoolingMinOp poolingOp) {
                      return ArithBuilder(b, loc).select(
                          ArithBuilder(b, loc).slt(lhs, rhs), lhs, rhs);
                    })
                    .Case([&](PoolingMaxOp poolingOp) {
                      return ArithBuilder(b, loc).select(
                          ArithBuilder(b, loc).sgt(lhs, rhs), lhs, rhs);
                    })
                    .Default([&](auto) { return Value(); });
  b.create<StoreOpTy>(loc, value, op.output(), indices.outputs);
}

template <typename LoadOpTy, typename StoreOpTy>
static void emitScalarImplementation(OpBuilder &b, Location loc,
                                     ArrayRef<Value> allIvs, PoolingMaxOp op) {
  emitPoolingMinMaxScalarImplementation<LoadOpTy, StoreOpTy, PoolingMaxOp>(
      b, loc, allIvs, op);
}

template <typename LoadOpTy, typename StoreOpTy>
static void emitScalarImplementation(OpBuilder &b, Location loc,
                                     ArrayRef<Value> allIvs, PoolingMinOp op) {
  emitPoolingMinMaxScalarImplementation<LoadOpTy, StoreOpTy, PoolingMinOp>(
      b, loc, allIvs, op);
}

template <typename LoadOpTy, typename StoreOpTy>
static void emitScalarImplementation(OpBuilder &b, Location loc,
                                     ArrayRef<Value> allIvs, PoolingSumOp op) {
  auto indices = getInputAndOutputIndices(b, loc, allIvs, op);
  Value inputVal =
      getPoolingInput<LoadOpTy, StoreOpTy>(b, loc, op, indices.inputs);
  Value outputVal = b.create<LoadOpTy>(loc, op.output(), indices.outputs);
  Value added = ArithBuilder(b, loc).add(outputVal, inputVal);
  b.create<StoreOpTy>(loc, added, op.output(), indices.outputs);
}

/// Replace the index operations in the body of the loop nest by the matching
/// induction variables.
static void replaceIndexOpsByInductionVariables(LinalgOp linalgOp,
                                                PatternRewriter &rewriter,
                                                ArrayRef<Operation *> loopOps) {
  // Extract the induction variables of the loop nest from outer to inner.
  SmallVector<Value> allIvs;
  for (Operation *loopOp : loopOps) {
    llvm::TypeSwitch<Operation *>(loopOp)
        .Case([&](scf::ParallelOp parallelOp) {
          allIvs.append(parallelOp.getInductionVars().begin(),
                        parallelOp.getInductionVars().end());
        })
        .Case([&](scf::ForOp forOp) {
          allIvs.push_back(forOp.getInductionVar());
        })
        .Case([&](AffineForOp affineForOp) {
          allIvs.push_back(affineForOp.getInductionVar());
        })
        .Default([&](Operation *op) { assert(false && "unexpected op"); });
  }
  assert(linalgOp.getNumLoops() == allIvs.size() &&
         "expected the number of loops and induction variables to match");
  // Replace the index operations in the body of the innermost loop op.
  if (!loopOps.empty()) {
    LoopLikeOpInterface loopOp = loopOps.back();
    for (IndexOp indexOp :
         llvm::make_early_inc_range(loopOp.getLoopBody().getOps<IndexOp>()))
      rewriter.replaceOp(indexOp, allIvs[indexOp.dim()]);
  }
}

template <typename LoopTy>
static Optional<LinalgLoops> linalgOpToLoopsImpl(PatternRewriter &rewriter,
                                                 LinalgOp linalgOp) {
  using LoadOpTy =
      typename std::conditional<std::is_same<LoopTy, AffineForOp>::value,
                                AffineLoadOp, memref::LoadOp>::type;
  using StoreOpTy =
      typename std::conditional<std::is_same<LoopTy, AffineForOp>::value,
                                AffineStoreOp, memref::StoreOp>::type;

  // The flattened loopToOperandRangesMaps is expected to be an invertible
  // permutation map (which is asserted in the inverse calculation).
  assert(linalgOp.hasBufferSemantics() &&
         "expected linalg op with buffer semantics");

  auto loopRanges = linalgOp.createLoopRanges(rewriter, linalgOp.getLoc());
  auto iteratorTypes = llvm::to_vector<4>(linalgOp.iterator_types().getValue());

  SmallVector<Value> allIvs;
  GenerateLoopNest<LoopTy>::doit(
      rewriter, linalgOp.getLoc(), loopRanges, linalgOp, iteratorTypes,
      [&](OpBuilder &b, Location loc, ValueRange ivs,
          ValueRange operandValuesToUse) -> scf::ValueVector {
        assert(operandValuesToUse == linalgOp->getOperands() &&
               "expect operands are captured and not passed by loop argument");
        allIvs.append(ivs.begin(), ivs.end());
        llvm::TypeSwitch<Operation *>(linalgOp)
            .Case<ConvOp, PoolingMaxOp, PoolingMinOp, PoolingSumOp, LinalgOp>(
                [&](auto op) {
                  emitScalarImplementation<LoadOpTy, StoreOpTy>(b, loc, allIvs,
                                                                op);
                })
            .Default([&](Operation *op) { assert(false && "unexpected op"); });
        return scf::ValueVector{};
      });
  // Number of loop ops might be different from the number of ivs since some
  // loops like affine.parallel and scf.parallel have multiple ivs.
  SetVector<Operation *> loopSet;
  for (Value iv : allIvs) {
    if (!iv)
      return {};
    // The induction variable is a block argument of the entry block of the
    // loop operation.
    BlockArgument ivVal = iv.dyn_cast<BlockArgument>();
    if (!ivVal)
      return {};
    loopSet.insert(ivVal.getOwner()->getParentOp());
  }
  LinalgLoops loops(loopSet.begin(), loopSet.end());
  // Replace all index operations in the loop body.
  replaceIndexOpsByInductionVariables(linalgOp, rewriter, loops);
  return loops;
}

namespace {
template <typename LoopType>
class LinalgRewritePattern : public RewritePattern {
public:
  LinalgRewritePattern(MLIRContext *context)
      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}

  LogicalResult matchAndRewrite(Operation *op,
                                PatternRewriter &rewriter) const override {
    auto linalgOp = dyn_cast<LinalgOp>(op);
    if (!isa<LinalgOp>(op))
      return failure();
    if (!linalgOpToLoopsImpl<LoopType>(rewriter, linalgOp))
      return failure();
    rewriter.eraseOp(op);
    return success();
  }
};

/// Converts tiled_loop to SCF loop nests. All parallel dimensions are collected
/// into an scf.parallel loop and all sequential dimensions will result in the
/// nested scf.for loop nest. The pattern assumes that a tiled loop with
/// iterator_types ["reduction", "parallel", "reduction"] can be reordered. It
/// is true for the tiling that is currently suppported by Linalg.
struct TiledLoopToSCFPattern : public OpRewritePattern<TiledLoopOp> {
  using OpRewritePattern<TiledLoopOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TiledLoopOp tiledLoop,
                                PatternRewriter &rewriter) const override {
    // Fail conversion if the `tiled_loop` has not been bufferized.
    if (!tiledLoop.hasBufferSemantics())
      return failure();

    // Collect loop control parameters for parallel and sequential dimensions.
    SmallVector<Value, 3> seqLBs, seqUBs, seqSteps, seqIVs;
    SmallVector<Value, 3> parLBs, parUBs, parSteps, parIVs;
    for (auto en : llvm::enumerate(
             llvm::zip(tiledLoop.lowerBound(), tiledLoop.upperBound(),
                       tiledLoop.step(), tiledLoop.getInductionVars()))) {
      Value lb, ub, step, iv;
      std::tie(lb, ub, step, iv) = en.value();
      if (tiledLoop.isParallelDimension(en.index())) {
        parLBs.push_back(lb);
        parUBs.push_back(ub);
        parSteps.push_back(step);
        parIVs.push_back(iv);
      } else {
        seqLBs.push_back(lb);
        seqUBs.push_back(ub);
        seqSteps.push_back(step);
        seqIVs.push_back(iv);
      }
    }

    Location loc = tiledLoop.getLoc();
    auto generateForLoopNestAndCloneBody = [&](OpBuilder &builder, Location loc,
                                               ValueRange ivs) {
      BlockAndValueMapping bvm;
      bvm.map(parIVs, ivs);
      bvm.map(tiledLoop.getRegionInputArgs(), tiledLoop.inputs());
      bvm.map(tiledLoop.getRegionOutputArgs(), tiledLoop.outputs());

      // If not all dimensions of the tiled loop are parallel, an scf.for loop
      // nest is generated.
      if (!seqIVs.empty()) {
        scf::LoopNest nest =
            scf::buildLoopNest(builder, loc, seqLBs, seqUBs, seqSteps,
                               [&](OpBuilder &builder, Location loc,
                                   ValueRange ivs) { bvm.map(seqIVs, ivs); });
        builder.setInsertionPointToStart(nest.loops.back().getBody());
      }
      for (auto &op : tiledLoop.getBody()->without_terminator())
        builder.clone(op, bvm);
    };

    if (parIVs.empty())
      generateForLoopNestAndCloneBody(rewriter, loc, llvm::None);
    else
      rewriter.create<scf::ParallelOp>(loc, parLBs, parUBs, parSteps,
                                       generateForLoopNestAndCloneBody);
    rewriter.eraseOp(tiledLoop);
    return success();
  }
};

/// Local folding pattern for AffineApplyOp that we can apply greedily.
/// This replaces AffineApplyOp by the proper value in cases where the
/// associated map is trivial.
/// A trivial map here is defined as a map with a single result and either:
///   1. Zero operand + returns a single AffineConstantExpr
///   2. One operand + returns a single AffineDimExpr
///   3. One operand + returns a single AffineSymbolExpr
//
/// In the first case, the AffineApplyOp is replaced by a new constant. In the
/// other cases, it is replaced by its unique operand.
struct FoldAffineOp : public RewritePattern {
  FoldAffineOp(MLIRContext *context)
      : RewritePattern(AffineApplyOp::getOperationName(), 0, context) {}

  LogicalResult matchAndRewrite(Operation *op,
                                PatternRewriter &rewriter) const override {
    AffineApplyOp affineApplyOp = cast<AffineApplyOp>(op);
    auto map = affineApplyOp.getAffineMap();
    if (map.getNumResults() != 1 || map.getNumInputs() > 1)
      return failure();

    AffineExpr expr = map.getResult(0);
    if (map.getNumInputs() == 0) {
      if (auto val = expr.dyn_cast<AffineConstantExpr>()) {
        rewriter.replaceOpWithNewOp<ConstantIndexOp>(op, val.getValue());
        return success();
      }
      return failure();
    }
    if (expr.dyn_cast<AffineDimExpr>() || expr.dyn_cast<AffineSymbolExpr>()) {
      rewriter.replaceOp(op, op->getOperand(0));
      return success();
    }
    return failure();
  }
};

template <typename LoopType>
static void lowerLinalgToLoopsImpl(FuncOp funcOp) {
  MLIRContext *context = funcOp.getContext();
  RewritePatternSet patterns(context);
  patterns.add<LinalgRewritePattern<LoopType>>(context);
  memref::DimOp::getCanonicalizationPatterns(patterns, context);
  tensor::DimOp::getCanonicalizationPatterns(patterns, context);
  AffineApplyOp::getCanonicalizationPatterns(patterns, context);
  patterns.add<FoldAffineOp>(context);
  // Just apply the patterns greedily.
  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
}

struct LowerToAffineLoops
    : public LinalgLowerToAffineLoopsBase<LowerToAffineLoops> {
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<memref::MemRefDialect>();
  }
  void runOnFunction() override {
    lowerLinalgToLoopsImpl<AffineForOp>(getFunction());
  }
};

struct LowerToLoops : public LinalgLowerToLoopsBase<LowerToLoops> {
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<memref::MemRefDialect, scf::SCFDialect>();
  }
  void runOnFunction() override {
    lowerLinalgToLoopsImpl<scf::ForOp>(getFunction());
  }
};

struct LowerToParallelLoops
    : public LinalgLowerToParallelLoopsBase<LowerToParallelLoops> {
  void runOnFunction() override {
    lowerLinalgToLoopsImpl<scf::ParallelOp>(getFunction());
  }
};

struct LowerTiledLoopsToSCF
    : public LinalgLowerTiledLoopsToSCFBase<LowerTiledLoopsToSCF> {
  void runOnFunction() override {
    MLIRContext *context = &getContext();
    RewritePatternSet patterns(context);
    populateTiledLoopToSCFPattern(patterns);
    (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
  }
};
} // namespace

/// Rewrite a TiledLoopOp with bounds/step that potentially do not divide evenly
/// into two TiledLoopOps: One where the step divides the iteration space
/// evenly, followed another one for the last (partial) iteration (if any). This
/// function only rewrites the `idx`-th loop of the loop nest represented by
/// the TiledLoopOp. To peel the entire loop nest, this function must be called
/// multiple times.
///
/// This function rewrites the given TiledLoopOp in-place and creates a new
/// TiledLoopOp for the last iteration. It replaces all uses of the original
/// TiledLoopOp with the results of the newly generated one.
///
/// The newly generated TiledLoopOp is returned via `result`. The boundary
/// at which the loop is split (new upper bound) is returned via `splitBound`.
/// The return value indicates whether the TiledLoopOp was rewritten or not.
static LogicalResult peelTiledLoop(RewriterBase &b, TiledLoopOp loopOp,
                                   int64_t idx, TiledLoopOp &result,
                                   Value &splitBound) {
  Value lb = loopOp.lowerBound()[idx], ub = loopOp.upperBound()[idx],
        step = loopOp.step()[idx];
  auto ubInt = getConstantIntValue(ub);

  auto loc = loopOp.getLoc();
  AffineExpr exprLb, exprUb, exprStep;
  bindSymbols(b.getContext(), exprLb, exprUb, exprStep);
  // New upper bound: %ub - (%ub - %lb) mod %step
  auto modMap = AffineMap::get(0, 3, {exprUb - ((exprUb - exprLb) % exprStep)});
  SmallVector<Value> operands{lb, ub, step};
  mlir::canonicalizeMapAndOperands(&modMap, &operands);
  modMap = mlir::simplifyAffineMap(modMap);
  RewriterBase::InsertionGuard guard(b);
  b.setInsertionPoint(loopOp);
  splitBound = b.createOrFold<AffineApplyOp>(loc, modMap, operands);
  // No specialization necessary if step already divides upper bound evenly.
  if (splitBound == ub || (ubInt && ubInt == getConstantIntValue(splitBound)))
    return failure();

  // Create remainder loop.
  b.setInsertionPointAfter(loopOp);
  auto remainderLoop = cast<TiledLoopOp>(b.clone(*loopOp.getOperation()));
  loopOp.replaceAllUsesWith(remainderLoop->getResults());
  // Outputs: Take tensors from main loop's results. Take memrefs from main
  // loop's outputs.
  SmallVector<Value> remainderOutputs;
  for (unsigned o = 0, t = 0; o < loopOp.getNumOutputs(); ++o) {
    remainderOutputs.push_back(loopOp.outputs()[o].getType().isa<MemRefType>()
                                   ? loopOp.outputs()[o]
                                   : loopOp->getResult(t++));
  }
  remainderLoop.outputsMutable().assign(remainderOutputs);

  // Set new loop bounds.
  b.updateRootInPlace(loopOp, [&]() {
    SmallVector<Value> ubs = loopOp.upperBound();
    ubs[idx] = splitBound;
    loopOp.upperBoundMutable().assign(ubs);
  });
  SmallVector<Value> lbs = remainderLoop.lowerBound();
  lbs[idx] = splitBound;
  remainderLoop.lowerBoundMutable().assign(lbs);

  result = remainderLoop;
  return success();
}

template <typename OpTy, bool IsMin>
static void
rewriteAffineOpAfterPeeling(RewriterBase &rewriter, TiledLoopOp mainLoop,
                            TiledLoopOp remainderLoop, Value mainIv,
                            Value remainderIv, Value ub, Value step) {
  mainLoop.walk([&](OpTy affineOp) {
    AffineMap map = affineOp.getAffineMap();
    (void)scf::rewritePeeledMinMaxOp(rewriter, affineOp, map,
                                     affineOp.operands(), IsMin, mainIv, ub,
                                     step, /*insideLoop=*/true);
  });
  remainderLoop.walk([&](OpTy affineOp) {
    AffineMap map = affineOp.getAffineMap();
    (void)scf::rewritePeeledMinMaxOp(rewriter, affineOp, map,
                                     affineOp.operands(), IsMin, remainderIv,
                                     ub, step, /*insideLoop=*/false);
  });
}

LogicalResult mlir::linalg::peelAndCanonicalizeTiledLoop(RewriterBase &rewriter,
                                                         TiledLoopOp loopOp,
                                                         int64_t idx,
                                                         TiledLoopOp &result) {
  int64_t numLoops = loopOp.iterator_types().size();
  if (idx < 0 || numLoops <= idx)
    return failure();
  // Only parallel iterator supported.
  if (!isParallelIterator(loopOp.iterator_types()[idx]))
    return failure();

  Value ub = loopOp.upperBound()[idx];
  TiledLoopOp remainderLoop;
  Value splitBound;
  if (failed(peelTiledLoop(rewriter, loopOp, idx, remainderLoop, splitBound)))
    return failure();

  // Rewrite affine.min and affine.max ops.
  Value mainIv = loopOp.getInductionVars()[idx], step = loopOp.step()[idx],
        remainderIv = remainderLoop.getInductionVars()[idx];

  rewriteAffineOpAfterPeeling<AffineMinOp, /*IsMin=*/true>(
      rewriter, loopOp, remainderLoop, mainIv, remainderIv, ub, step);
  rewriteAffineOpAfterPeeling<AffineMaxOp, /*IsMin=*/false>(
      rewriter, loopOp, remainderLoop, mainIv, remainderIv, ub, step);

  result = remainderLoop;
  return success();
}

void mlir::linalg::populateTiledLoopToSCFPattern(RewritePatternSet &patterns) {
  patterns.add<TiledLoopToSCFPattern>(patterns.getContext());
}

std::unique_ptr<OperationPass<FuncOp>>
mlir::createConvertLinalgTiledLoopsToSCFPass() {
  return std::make_unique<LowerTiledLoopsToSCF>();
}

std::unique_ptr<OperationPass<FuncOp>> mlir::createConvertLinalgToLoopsPass() {
  return std::make_unique<LowerToLoops>();
}

std::unique_ptr<OperationPass<FuncOp>>
mlir::createConvertLinalgToParallelLoopsPass() {
  return std::make_unique<LowerToParallelLoops>();
}

std::unique_ptr<OperationPass<FuncOp>>
mlir::createConvertLinalgToAffineLoopsPass() {
  return std::make_unique<LowerToAffineLoops>();
}

/// Emits a loop nest of `affine.for` with the proper body for `linalgOp`.
Optional<LinalgLoops>
mlir::linalg::linalgOpToAffineLoops(PatternRewriter &rewriter,
                                    LinalgOp linalgOp) {
  return linalgOpToLoopsImpl<AffineForOp>(rewriter, linalgOp);
}

/// Emits a loop nest of `scf.for` with the proper body for `linalgOp`.
Optional<LinalgLoops> mlir::linalg::linalgOpToLoops(PatternRewriter &rewriter,
                                                    LinalgOp linalgOp) {
  return linalgOpToLoopsImpl<scf::ForOp>(rewriter, linalgOp);
}

/// Emits a loop nest of `scf.parallel` with the proper body for `linalgOp`.
Optional<LinalgLoops>
mlir::linalg::linalgOpToParallelLoops(PatternRewriter &rewriter,
                                      LinalgOp linalgOp) {
  return linalgOpToLoopsImpl<scf::ParallelOp>(rewriter, linalgOp);
}