text: //===- 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/Affine/EDSC/Intrinsics.h"
  #include "mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.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/EDSC/Builders.h"
  #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.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::edsc;
  using namespace mlir::edsc::intrinsics;
  using namespace mlir::linalg;
  
  using edsc::op::operator+;
  
  static SmallVector<Value, 8> makeCanonicalAffineApplies(OpBuilder &b,
                                                          Location loc,
                                                          AffineMap map,
                                                          ArrayRef<Value> vals) {
    if (map.isEmpty())
      return {};
  
    assert(map.getNumInputs() == vals.size());
    SmallVector<Value, 8> res;
    res.reserve(map.getNumResults());
    auto dims = map.getNumDims();
    for (auto e : map.getResults()) {
      auto exprMap = AffineMap::get(dims, map.getNumSymbols(), e);
      SmallVector<Value, 4> operands(vals.begin(), vals.end());
      canonicalizeMapAndOperands(&exprMap, &operands);
      res.push_back(affine_apply(exprMap, operands));
    }
    return res;
  }
  
  static SmallVector<Value, 4> permuteIvs(ArrayRef<Value> ivs,
                                          Optional<AffineMap> permutation) {
    return permutation ? applyMapToValues(ScopedContext::getBuilderRef(),
                                          ScopedContext::getLocation(),
                                          permutation.getValue(), ivs)
                       : SmallVector<Value, 4>(ivs.begin(), ivs.end());
  }
  
  template <typename IndexedValueType, typename OpType>
  static void inlineRegionAndEmitStore(OpType op, ArrayRef<Value> indexedValues,
                                       ArrayRef<SmallVector<Value, 8>> indexing,
                                       ArrayRef<Value> outputBuffers) {
    assert(op->getNumRegions() == 1 && "Expected single region op");
    auto &b = ScopedContext::getBuilderRef();
    auto &block = op->getRegion(0).front();
    BlockAndValueMapping map;
    map.map(block.getArguments(), indexedValues);
    for (auto &op : block.without_terminator()) {
      assert(op.getNumRegions() == 0 && "expected a non-nested region");
      auto *newOp = b.clone(op, map);
      map.map(op.getResults(), newOp->getResults());
    }
  
    Operation &terminator = block.back();
    assert(isa<linalg::YieldOp>(terminator) &&
           "expected a yield op in the end of the region");
    for (unsigned i = 0, e = terminator.getNumOperands(); i < e; ++i) {
      IndexedValueType O(outputBuffers[i]);
      O(indexing[i]) = map.lookupOrDefault(terminator.getOperand(i));
    }
  }
  
  // Returns a pair that contains input indices and output indices of a
  // SingleInputPoolingOp `op`.
  struct InputAndOutputIndices {
    SmallVector<Value, 8> inputs;
    SmallVector<Value, 8> outputs;
  };
  template <typename SingleInputPoolingOp>
  static InputAndOutputIndices getInputAndOutputIndices(ArrayRef<Value> allIvs,
                                                        SingleInputPoolingOp op) {
    auto &b = ScopedContext::getBuilderRef();
    auto loc = ScopedContext::getLocation();
    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 IndexedValueType>
  static void emitScalarImplementation(ArrayRef<Value> allIvs,
                                       LinalgOp linalgOp) {
    assert(linalgOp.hasBufferSemantics() &&
           "expected linalg op with buffer semantics");
    auto &b = ScopedContext::getBuilderRef();
    auto loc = ScopedContext::getLocation();
    unsigned nInputs = linalgOp.getNumInputs();
    unsigned nOutputs = linalgOp.getNumOutputs();
    SmallVector<Value, 4> indexedValues;
    indexedValues.reserve(nInputs + nOutputs);
  
    auto allIvsPlusDims = SmallVector<Value, 4>(allIvs.begin(), allIvs.end());
  
    // TODO: Avoid the loads if the corresponding argument of the
    // region has no uses.
    // 1.a. Emit load from input views.
    for (unsigned i = 0; i < nInputs; ++i) {
      auto indexing = makeCanonicalAffineApplies(
          b, loc, linalgOp.getInputIndexingMap(i), allIvsPlusDims);
      // Passing through IndexedValueType emits the proper load operation.
      indexedValues.push_back(IndexedValueType(linalgOp.getInput(i))(indexing));
    }
    // 1.b. Emit load from output views.
    for (unsigned i = 0; i < nOutputs; ++i) {
      auto indexing = makeCanonicalAffineApplies(
          b, loc, linalgOp.getOutputIndexingMap(i), allIvsPlusDims);
      // Passing through IndexedValueType emits the proper load operation.
      indexedValues.push_back(
          IndexedValueType(linalgOp.getOutputBuffer(i))(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>, 8> indexing;
    SmallVector<Value, 8> outputBuffers;
    for (unsigned i = 0; i < nOutputs; ++i) {
      indexing.push_back(makeCanonicalAffineApplies(
          b, loc, linalgOp.getOutputIndexingMap(i), allIvsPlusDims));
      outputBuffers.push_back(linalgOp.getOutputBuffer(i));
    }
    inlineRegionAndEmitStore<IndexedValueType>(linalgOp, indexedValues, indexing,
                                               outputBuffers);
  }
  
  template <typename IndexedValueType>
  static void emitScalarImplementation(ArrayRef<Value> allIvs, CopyOp copyOp) {
    assert(copyOp.hasBufferSemantics() &&
           "expected linalg op with buffer semantics");
    auto nPar = copyOp.getNumParallelLoops();
    assert(nPar == allIvs.size());
    auto inputIvs =
        permuteIvs(allIvs.take_front(nPar), copyOp.inputPermutation());
    auto outputIvs =
        permuteIvs(allIvs.take_front(nPar), copyOp.outputPermutation());
    SmallVector<Value, 8> iivs(inputIvs.begin(), inputIvs.end());
    SmallVector<Value, 8> oivs(outputIvs.begin(), outputIvs.end());
    IndexedValueType O(copyOp.getOutputBuffer(0)), I(copyOp.getInput(0));
    // Emit the proper scalar assignment, whether we are dealing with a 0-D or
    // an n-D loop nest; with or without permutations.
    // clang-format off
      nPar > 0 ? O(oivs) = I(iivs) :
                 O() = I();
    // clang-format on
  }
  
  template <typename IndexedValueType>
  static void emitScalarImplementation(ArrayRef<Value> allIvs, FillOp fillOp) {
    assert(fillOp.hasBufferSemantics() &&
           "expected linalg op with buffer semantics");
    auto nPar = fillOp.getNumParallelLoops();
    assert(nPar == allIvs.size());
    auto ivs = SmallVector<Value, 4>(allIvs.begin(), allIvs.begin() + nPar);
    IndexedValueType O(fillOp.getOutputBuffer(0));
    // Emit the proper scalar assignment, whether we are dealing with a 0-D or
    // an n-D loop nest; with or without permutations.
    nPar > 0 ? O(ivs) = fillOp.value() : O() = fillOp.value();
  }
  
  // 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.
  template <typename IndexedValueType>
  Value getPaddedInput(Value input, ArrayRef<Value> indices,
                       ArrayRef<int> skipPadding, Value padValue) {
    // TODO: add a level of indirection to linalg.generic.
  
    IndexedValueType indexedInput(input);
  
    auto *context = ScopedContext::getContext();
    Value zeroIndex = std_constant_index(0);
    SmallVector<Value, 8> conds;
    SmallVector<Value, 8> 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;
      }
  
      using edsc::op::sge;
      using edsc::op::slt;
      using edsc::op::operator||;
      Value leftOutOfBound = slt(dim, zeroIndex);
      if (conds.empty())
        conds.push_back(leftOutOfBound);
      else
        conds.push_back(conds.back() || leftOutOfBound);
      Value rightBound = std_dim(input, idx);
      conds.push_back(conds.back() || (sge(dim, rightBound)));
  
      // When padding is involved, the indices will only be shifted to negative,
      // so having a max op is enough.
      auto maxMap = AffineMap::get(/*dimCount=*/1, 0,
                                   {getAffineDimExpr(/*position=*/0, context),
                                    getAffineConstantExpr(0, context)},
                                   context);
      clampedImIdx.push_back(affine_max(dim.getType(), maxMap, ValueRange{dim}));
    }
  
    Value readInput = indexedInput(clampedImIdx);
    return conds.empty() ? readInput
                         : (Value)std_select(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) {
    auto &b = ScopedContext::getBuilderRef();
    if (auto floatType = type.dyn_cast<FloatType>()) {
      return b.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 b.getIntegerAttr(intType, APInt::getSignedMinValue(width));
    }
    llvm_unreachable("Unsupported data type for PoolingMaxOp");
    return {};
  }
  
  template <>
  Attribute getPadValueAttr<PoolingMinOp>(Type type) {
    auto &b = ScopedContext::getBuilderRef();
    if (auto floatType = type.dyn_cast<FloatType>()) {
      return b.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 b.getIntegerAttr(intType, APInt::getSignedMaxValue(width));
    }
    llvm_unreachable("Unsupported data type for PoolingMinOp");
    return {};
  }
  
  template <>
  Attribute getPadValueAttr<PoolingSumOp>(Type type) {
    auto &b = ScopedContext::getBuilderRef();
    return b.getZeroAttr(type);
  }
  
  template <>
  Attribute getPadValueAttr<ConvOp>(Type type) {
    auto &b = ScopedContext::getBuilderRef();
    return b.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 IndexedValueType>
  static void emitScalarImplementation(ArrayRef<Value> allIvs, ConvOp convOp) {
    assert(convOp.hasBufferSemantics() &&
           "expected linalg op with buffer semantics");
    auto &b = ScopedContext::getBuilderRef();
    auto loc = ScopedContext::getLocation();
    auto mapsRange = convOp.indexing_maps().getAsRange<AffineMapAttr>();
    auto maps = llvm::to_vector<8>(
        llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
    SmallVector<Value, 8> fIdx(
        makeCanonicalAffineApplies(b, loc, maps[0], allIvs));
    SmallVector<Value, 8> imIdx(
        makeCanonicalAffineApplies(b, loc, maps[1], allIvs));
    SmallVector<Value, 8> oIdx(
        makeCanonicalAffineApplies(b, loc, maps[2], allIvs));
  
    IndexedValueType F(convOp.filter()), O(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 StdIndexedValue
    // when there is non-zero padding.
    if (hasPadding(convOp)) {
      Type type = convOp.input().getType().cast<MemRefType>().getElementType();
      Value padValue = std_constant(type, getPadValueAttr<ConvOp>(type));
      Value paddedInput = getPaddedInput<StdIndexedValue>(
          convOp.input(), imIdx,
          /* Only need to pad the window dimensions */
          {0, static_cast<int>(imIdx.size()) - 1}, padValue);
      O(oIdx) += F(fIdx) * paddedInput;
    } else {
      IndexedValueType I(convOp.input());
      O(oIdx) += F(fIdx) * I(imIdx);
    }
  }
  
  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 IndexedValueType, typename PoolingOp>
  static Value getPoolingInput(PoolingOp op, ArrayRef<Value> inputIndices) {
    if (hasPadding(op)) {
      Type type =
          op.input().getType().template cast<MemRefType>().getElementType();
      Value padValue = std_constant(type, getPadValueAttr<PoolingOp>(type));
      return getPaddedInput<StdIndexedValue>(op.input(), inputIndices,
                                             /*Pad every dimension*/ {},
                                             padValue);
    }
    IndexedValueType input(op.input());
    return input(inputIndices);
  }
  
  template <typename IndexedValueType, typename OpType>
  void emitPoolingMinMaxScalarImplementation(ArrayRef<Value> allIvs, OpType op) {
    InputAndOutputIndices indices = getInputAndOutputIndices(allIvs, op);
    // Emit scalar form.
    IndexedValueType output(op.output());
    Value lhs = output(indices.outputs);
    Value rhs = getPoolingInput<IndexedValueType>(op, indices.inputs);
    using edsc::op::sgt;
    using edsc::op::slt;
    Value value = std::is_same<OpType, PoolingMinOp>()
                      ? std_select(slt(lhs, rhs), lhs, rhs)
                      : std_select(sgt(lhs, rhs), lhs, rhs);
    output(indices.outputs) = value;
  }
  
  template <typename IndexedValueType>
  static void emitScalarImplementation(ArrayRef<Value> allIvs, PoolingMaxOp op) {
    emitPoolingMinMaxScalarImplementation<IndexedValueType, PoolingMaxOp>(allIvs,
                                                                          op);
  }
  
  template <typename IndexedValueType>
  static void emitScalarImplementation(ArrayRef<Value> allIvs, PoolingMinOp op) {
    emitPoolingMinMaxScalarImplementation<IndexedValueType, PoolingMinOp>(allIvs,
                                                                          op);
  }
  
  template <typename IndexedValueType>
  static void emitScalarImplementation(ArrayRef<Value> allIvs, PoolingSumOp op) {
    auto indices = getInputAndOutputIndices(allIvs, op);
    IndexedValueType output(op.output());
  
    // Emit scalar form.
    output(indices.outputs) +=
        getPoolingInput<IndexedValueType>(op, indices.inputs);
  }
  
  /// Emits the MLIR for the scalar part of the indexed 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 induction
  ///      variables and 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(%i, %j, %k, %11, %12, %13) :
  ///            (index, index, index, 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 IndexedValueType>
  static void emitScalarImplementation(ArrayRef<Value> allIvs,
                                       IndexedGenericOp indexedGenericOp) {
    assert(indexedGenericOp.hasBufferSemantics() &&
           "expected linalg op with buffer semantics");
    auto &b = ScopedContext::getBuilderRef();
    auto loc = ScopedContext::getLocation();
    unsigned nInputs = indexedGenericOp.getNumInputs();
    unsigned nOutputs = indexedGenericOp.getNumOutputs();
    unsigned nLoops = allIvs.size();
    SmallVector<Value, 4> indexedValues;
    indexedValues.reserve(nLoops + nInputs + nOutputs);
    for (unsigned i = 0; i < nLoops; ++i)
      indexedValues.push_back(allIvs[i]);
  
    // TODO: Avoid the loads if the corresponding argument of the
    // region has no uses.
    // 1.a. Emit load from input views.
    for (unsigned i = 0; i < nInputs; ++i) {
      auto indexing = makeCanonicalAffineApplies(
          b, loc, indexedGenericOp.getInputIndexingMap(i), allIvs);
      // Pass input i through IndexedValueType emits the proper load operation.
      indexedValues.push_back(
          IndexedValueType(indexedGenericOp.getInput(i))(indexing));
    }
    // 1.b. Emit load from output views.
    for (unsigned i = 0; i < nOutputs; ++i) {
      auto indexing = makeCanonicalAffineApplies(
          b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs);
      // Pass output i through IndexedValueType emits the proper load operation.
      indexedValues.push_back(
          IndexedValueType(indexedGenericOp.getOutputBuffer(i))(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>, 8> indexing;
    SmallVector<Value, 8> outputBuffers;
    for (unsigned i = 0; i < nOutputs; ++i) {
      indexing.push_back(makeCanonicalAffineApplies(
          b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs));
      outputBuffers.push_back(indexedGenericOp.getOutputBuffer(i));
    }
    inlineRegionAndEmitStore<IndexedValueType>(indexedGenericOp, indexedValues,
                                               indexing, outputBuffers);
  }
  
  template <typename LoopTy>
  static Optional<LinalgLoops> linalgOpToLoopsImpl(Operation *op,
                                                   OpBuilder &builder) {
    using IndexedValueTy = typename GenerateLoopNest<LoopTy>::IndexedValueTy;
  
    ScopedContext scope(builder, op->getLoc());
  
    // The flattened loopToOperandRangesMaps is expected to be an invertible
    // permutation map (which is asserted in the inverse calculation).
    auto linalgOp = cast<LinalgOp>(op);
    assert(linalgOp.hasBufferSemantics() &&
           "expected linalg op with buffer semantics");
    auto loopRanges = linalgOp.createLoopRanges(builder, op->getLoc());
    SmallVector<Value, 4> allIvs;
    GenerateLoopNest<LoopTy>::doit(
        loopRanges, /*iterInitArgs*/ {}, linalgOp.iterator_types().getValue(),
        [&](ValueRange ivs, ValueRange iterArgs) -> scf::ValueVector {
          assert(iterArgs.empty() && "unexpected iterArgs");
          allIvs.append(ivs.begin(), ivs.end());
          llvm::TypeSwitch<Operation *>(op)
              .Case<CopyOp, FillOp, ConvOp, PoolingMaxOp, PoolingMinOp,
                    PoolingSumOp, IndexedGenericOp, LinalgOp>([&](auto op) {
                emitScalarImplementation<IndexedValueTy>(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.
    llvm::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());
    return loops;
  }
  
  namespace {
  template <typename LoopType>
  class LinalgRewritePattern : public RewritePattern {
  public:
    LinalgRewritePattern() : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()) {}
  
    LogicalResult matchAndRewrite(Operation *op,
                                  PatternRewriter &rewriter) const override {
      if (!isa<LinalgOp>(op))
        return failure();
      if (!linalgOpToLoopsImpl<LoopType>(op, rewriter))
        return failure();
      rewriter.eraseOp(op);
      return success();
    }
  };
  
  struct FoldAffineOp;
  } // namespace
  
  template <typename LoopType>
  static void lowerLinalgToLoopsImpl(FuncOp funcOp, MLIRContext *context) {
    OwningRewritePatternList patterns;
    patterns.insert<LinalgRewritePattern<LoopType>>();
    DimOp::getCanonicalizationPatterns(patterns, context);
    AffineApplyOp::getCanonicalizationPatterns(patterns, context);
    patterns.insert<FoldAffineOp>(context);
    // Just apply the patterns greedily.
    applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
  }
  
  namespace {
  /// 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();
    }
  };
  
  struct LowerToAffineLoops
      : public LinalgLowerToAffineLoopsBase<LowerToAffineLoops> {
    void runOnFunction() override {
      lowerLinalgToLoopsImpl<AffineForOp>(getFunction(), &getContext());
    }
  };
  
  struct LowerToLoops : public LinalgLowerToLoopsBase<LowerToLoops> {
    void runOnFunction() override {
      lowerLinalgToLoopsImpl<scf::ForOp>(getFunction(), &getContext());
    }
  };
  
  struct LowerToParallelLoops
      : public LinalgLowerToParallelLoopsBase<LowerToParallelLoops> {
    void runOnFunction() override {
      lowerLinalgToLoopsImpl<scf::ParallelOp>(getFunction(), &getContext());
    }
  };
  } // namespace
  
  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 with the proper body for `op`.
  template <typename LoopTy>
  Optional<LinalgLoops> mlir::linalg::linalgLowerOpToLoops(OpBuilder &builder,
                                                           Operation *op) {
    return linalgOpToLoopsImpl<LoopTy>(op, builder);
  }
  
  template Optional<LinalgLoops>
  mlir::linalg::linalgLowerOpToLoops<AffineForOp>(OpBuilder &builder,
                                                  Operation *op);
  template Optional<LinalgLoops>
  mlir::linalg::linalgLowerOpToLoops<scf::ForOp>(OpBuilder &builder,
                                                 Operation *op);
  template Optional<LinalgLoops>
  mlir::linalg::linalgLowerOpToLoops<scf::ParallelOp>(OpBuilder &builder,
                                                      Operation *op);
  
  /// Emits a loop nest of `affine.for` with the proper body for `op`.
  LogicalResult mlir::linalg::linalgOpToAffineLoops(OpBuilder &builder,
                                                    Operation *op) {
    Optional<LinalgLoops> loops = linalgLowerOpToLoops<AffineForOp>(builder, op);
    return loops ? success() : failure();
  }
  
  /// Emits a loop nest of `scf.for` with the proper body for `op`.
  LogicalResult mlir::linalg::linalgOpToLoops(OpBuilder &builder, Operation *op) {
    Optional<LinalgLoops> loops = linalgLowerOpToLoops<scf::ForOp>(builder, op);
    return loops ? success() : failure();
  }
  
  /// Emits a loop nest of `scf.parallel` with the proper body for `op`.
  LogicalResult mlir::linalg::linalgOpToParallelLoops(OpBuilder &builder,
                                                      Operation *op) {
    Optional<LinalgLoops> loops =
        linalgLowerOpToLoops<scf::ParallelOp>(builder, op);
    return loops ? success() : failure();
  }
  