// Copyright 2020 The IREE Authors
//
// Licensed 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

//===- LinalgBufferizePass.cpp - Pass to bufferize Linalg on tensors ------===//
//
// The overall bufferizarion algorithm is summarized here. Each of the
// individual steps are explained in detail later.
//
// Problem statement:
//
// The bufferization in this file is intended for converting tensor-operations
// into memref-operations for ops within a dispatch region. The goal is to reuse
// the buffers provided as inputs/outputs by the hal layer as memrefs for each
// of the operations. If the transformation cannot reuse input/output buffer to
// store an intermediate tensor, an allocation is done. This allocation is
// typically meant to be to target scratchspace memory.
//
// The algorithm has two phases an analysis phase and a tranformation phase.
//
// - The analysis phase walks the function and organizes relevant tensors
//   (tensors that need to be converted to memrefs) into equivalence clases. Two
//   tensors are part of the same equivalence class if they can eventually be
//   mapped to the same memref. This allows determining which operations can use
//   the buffer provided for the outputs to compute the results in place.
// - The transformation phase walks the function again and inserts corresponding
//   memref operations. The tensor operations are still kept around since the
//   analysis driving the transformation is based on the tensor values.
//   - Converting tensor operations to memref operations when all operands use
//     either buffers that are inputs to the dispatch or are allocated
//     temporarily within the dispatch region can be achieved by a
//     straight-forward walk.
//   - Reusing memref for the result of the dispatch for operations is more
//     involved and explained below.
//
//===----------------------------------------------------------------------===//
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree/compiler/Codegen/Common/BufferizationAnalysis.h"
#include "iree/compiler/Codegen/PassDetail.h"
#include "iree/compiler/Codegen/Passes.h"
#include "iree/compiler/Codegen/Transforms/Transforms.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/Flow/IR/FlowTypes.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "llvm/ADT/EquivalenceClasses.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Vector/VectorOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/Value.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/Passes.h"

#define DEBUG_TYPE "iree-codegen-linalg-bufferize"

namespace mlir {
namespace iree_compiler {

//===----------------------------------------------------------------------===//
// Bufferization helper functions using BlockAndValueMapping.
//===----------------------------------------------------------------------===//

/// Returns the dynamic dimensions of a Value `v` that is assumed to be
/// ShapedType.
static SmallVector<Value, 4> getDynamicDims(OpBuilder &b, Location loc,
                                            Value v) {
  SmallVector<Value, 4> dynamicDims;
  for (auto shape : enumerate(v.getType().cast<ShapedType>().getShape())) {
    if (shape.value() == ShapedType::kDynamicSize) {
      dynamicDims.push_back(
          b.createOrFold<memref::DimOp>(loc, v, shape.index()));
    }
  }
  return dynamicDims;
}

/// Allocates a memref for the results of an operation. Uses the
/// `InferShapedTypeOpInterface` where possible to get the shape of the output
/// in terms of the shapes of the operands.
static Value allocateBufferForResult(OpBuilder &b, Operation *op,
                                     unsigned resultNum,
                                     WorkgroupMemoryAllocationFn allocationFn) {
  assert(op->getNumResults() > resultNum);
  RankedTensorType resultType =
      op->getResult(resultNum).getType().cast<RankedTensorType>();
  SmallVector<Value, 4> dynamicDims;

  // Get the shape of the result
  Location loc = op->getLoc();
  if (auto shapedOp = dyn_cast<ReifyRankedShapedTypeOpInterface>(op)) {
    ReifiedRankedShapedTypeDims resultShape;
    if (failed(shapedOp.reifyResultShapes(b, resultShape))) {
      return nullptr;
    }
    for (auto shape : enumerate(resultShape[resultNum])) {
      if (resultType.isDynamicDim(shape.index())) {
        dynamicDims.push_back(shape.value());
      }
    }
  } else if (auto loadOp = dyn_cast<IREE::Flow::DispatchTensorLoadOp>(op)) {
    dynamicDims = llvm::to_vector<4>(loadOp.sizes());
  } else if (auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
    dynamicDims = llvm::to_vector<4>(sliceOp.sizes());
  } else if (auto subTensorInsertOp = dyn_cast<tensor::InsertSliceOp>(op)) {
    dynamicDims = getDynamicDims(b, loc, subTensorInsertOp.dest());
  } else if (auto transferWriteOp = dyn_cast<vector::TransferWriteOp>(op)) {
    dynamicDims = getDynamicDims(b, loc, transferWriteOp.source());
  } else {
    dynamicDims = getDynamicDims(b, loc, op->getResult(resultNum));
  }

  // If its a static allocation hoist it all the way up at begining of the
  // function.
  if (dynamicDims.empty()) {
    auto funcOp = op->getParentOfType<FuncOp>();
    OpBuilder::InsertionGuard g(b);
    b.setInsertionPointToStart(&funcOp.front());
    return allocationFn(b, loc, resultType.getShape(),
                        resultType.getElementType(), dynamicDims);
  }
  return allocationFn(b, loc, resultType.getShape(),
                      resultType.getElementType(), dynamicDims);
}

template <typename TensorType>
static MemRefType getMemrefTypeForTensor(TensorType tensorType,
                                         Attribute layout = Attribute(),
                                         Attribute memorySpace = nullptr) {
  return MemRefType::get(tensorType.getShape(), tensorType.getElementType(),
                         layout, memorySpace);
}

/// Creates a subview operation given the `src`, `offsets`, `sizes` and
/// `strides`. Handles the corner case where the `offsets`, `sizes` and
/// `strides` are empty in which case just forward the `src` value.  If the
/// `resultRank` does not match the source rank, uses a rank-reduced subview.
static Value createSubviewOp(OpBuilder &b, Location loc, unsigned resultRank,
                             Value src, ArrayRef<OpFoldResult> offsets,
                             ArrayRef<OpFoldResult> sizes,
                             ArrayRef<OpFoldResult> strides) {
  if (offsets.empty() && sizes.empty() && strides.empty()) return src;
  MemRefType resultType;
  MemRefType srcType = src.getType().cast<MemRefType>();
  if (srcType.getRank() != resultRank) {
    resultType = memref::SubViewOp::inferRankReducedResultType(
                     resultRank, srcType, offsets, sizes, strides)
                     .cast<MemRefType>();
  }
  return b.create<memref::SubViewOp>(loc, resultType, src, offsets, sizes,
                                     strides);
}

//===----------------------------------------------------------------------===//
// There might be cases when the `value` stored into a
// `flow.dispatch.tensor.store` operation is obtained from operation that
// computes the value (say a `linalg` operation) through a series of `reshapes`,
// `cast` etc. When trying to reuse the buffer for the result passed in to the
// dispatch region for these operations, these operations need to be "replayed"
// in reverse so that the type of the buffer in the operation computing the
// value matches what is expected.
//
// For example,
// ```mlir
//   %buffer = hal.interface.binding.subspan .. : tensor<?xf32>
//   %result = linalg.matmul ins(%lhs, %rhs : tensor<?x?xf32>, tensor<?x?xf32>)
//       outs(%init : tensor<?x?xf32>) -> tensor<?x?xf32>
//   %value = linalg.tensor_reshape %result [affine_map<(d0, d1) -> (d0, d1)]
//       : tensor<?x?xf32> into tensor<?xf32>
//   flow.dispatch.tensor.store %value, %buffer[..] [..] [..]
// ```
//
// needs to be converted to
//
// ```mlir
//   %buffer = hal.interface.binding.subspan .. : memref<?xf32>
//   %result = subview %buffer[..] [..] [..] : memref<?xf32>
//   %value = linalg.reshape %result [affine_map<(d0, d1) -> (d0, d1)]
//       : memref<?xf32> into memref<?x?xf32>
//   linalg.matmul ins(%lhs, %rhs : memref<?x?xf32>, memref<?x?xf32>)
//       outs(%result : memref<?x?xf32>)
//   flow.dispatch.tensor.store %value, %buffer[..] [..] [..]
// ```
//
// ===----------------------------------------------------------------------===//

/// Returns the subview into the buffer that is supposed to be populated with
/// the `value` of the `flow.dispatch.tensor.store` operation. This can be used
/// to compute the results in place.
static Value getSubviewOpForTensorStoreOp(OpBuilder &b, Operation *storeOp,
                                          const BlockAndValueMapping &bvm) {
  SmallVector<Value, 4> operandsOfSubviewOp;
  auto op = cast<OffsetSizeAndStrideOpInterface>(storeOp);
  Value target, source;
  std::tie(source, target) =
      TypeSwitch<Operation *, std::tuple<Value, Value>>(op)
          .Case<IREE::Flow::DispatchTensorStoreOp>([&](auto storeOp) {
            return std::make_tuple(storeOp.value(), storeOp.target());
          })
          .Case<tensor::InsertSliceOp>([&](auto storeOp) {
            return std::make_tuple(storeOp.source(), storeOp.dest());
          })
          .Default([](Operation *) {
            return std::make_tuple<Value, Value>(nullptr, nullptr);
          });
  if (!target) return nullptr;

  // Clone the offset, size and stride values. They will be CSE-ed later.
  Operation *parentOp = storeOp->getParentOp();
  BlockAndValueMapping indexValMap;
  llvm::SetVector<Operation *> slice;
  auto cloneIndexValues = [&](ArrayRef<OpFoldResult> ofrs) {
    SmallVector<OpFoldResult> clonedVals;
    for (auto ofr : ofrs) {
      // Just copy the attributes.
      if (auto attr = ofr.dyn_cast<Attribute>()) {
        clonedVals.push_back(attr);
        continue;
      }
      Value val = ofr.get<Value>();
      // If it is a block argument use the same value.
      if (val.isa<BlockArgument>()) {
        clonedVals.push_back(val);
        continue;
      }
      // The slice of ops needed for index computation need to be cloned to
      // avoid use-def violations. If the value has been cloned already, reuse
      // that.
      if (auto lookupVal = indexValMap.lookupOrNull(val)) {
        clonedVals.push_back(lookupVal);
        continue;
      }
      slice.clear();
      getBackwardSlice(val, &slice, [&](Operation *sliceOp) {
        return sliceOp->getParentOp() == parentOp;
      });
      for (auto sliceOp : slice) {
        if (!indexValMap.contains(sliceOp->getResult(0))) {
          b.clone(*sliceOp, indexValMap);
        }
      }
      if (Operation *definingOp = val.getDefiningOp()) {
        b.clone(*definingOp, indexValMap);
      }
      clonedVals.push_back(indexValMap.lookup(val));
    }
    return clonedVals;
  };
  SmallVector<OpFoldResult> subViewOffsets, subViewSizes, subViewStrides;
  subViewOffsets = cloneIndexValues(op.getMixedOffsets());
  subViewSizes = cloneIndexValues(op.getMixedSizes());
  subViewStrides = cloneIndexValues(op.getMixedStrides());
  Value subview = createSubviewOp(
      b, op.getLoc(), source.getType().cast<ShapedType>().getRank(),
      bvm.lookup(target), subViewOffsets, subViewSizes, subViewStrides);
  return subview;
}

/// Gets the reverse of a
/// `linalg.tensor_expand_shape`/`linalg.tensor_collapse_shape` op to get a
/// memref type that can be used for in-place computation of the result of a
/// dispatch region.
template <typename TensorReshapeOpTy>
static Value getReverseOfReshapeOp(OpBuilder &b, TensorReshapeOpTy reshapeOp,
                                   Value resultBuffer) {
  auto memrefType = getMemrefTypeForTensor(
      reshapeOp.getSrcType(), {},
      resultBuffer.getType().cast<MemRefType>().getMemorySpace());
  using ReverseReshapeOpTy = typename std::conditional<
      std::is_same<TensorReshapeOpTy, linalg::TensorCollapseShapeOp>::value,
      memref::ExpandShapeOp, memref::CollapseShapeOp>::type;
  return b.create<ReverseReshapeOpTy>(reshapeOp.getLoc(), memrefType,
                                      resultBuffer, reshapeOp.reassociation());
}

/// Gets the reverse of a `tensor.cast` op to get a memref type that
/// can be used for in-place computation of the result of a disaptch region.
static Value getReverseOfCastOp(OpBuilder &b, tensor::CastOp castOp,
                                Value resultBuffer) {
  auto memrefType = getMemrefTypeForTensor(
      castOp.source().getType().cast<RankedTensorType>(),
      resultBuffer.getType().cast<MemRefType>().getLayout(),
      resultBuffer.getType().cast<MemRefType>().getMemorySpace());
  return b.create<memref::CastOp>(castOp.getLoc(), memrefType, resultBuffer);
}

/// Returns a tied result value give the operand. If no such result exists,
/// returns `nullptr`.
static Value getTiedResultForOperand(OpOperand &operand,
                                     const BufferizationPlan &plan) {
  for (Value result : operand.getOwner()->getResults()) {
    if (plan.isEquivalent(operand.get(), result)) {
      return result;
    }
  }
  return nullptr;
}

/// To perform updates directly into the result buffer, the uses need to be
/// walked to get to a value already mapped to a buffer or a
/// `flow.dispatch.tensor.store` operation. For each use, gets the tied result
/// and follow its uses. The traversed uses and thir tied results are returned
/// in `traversedUses`.
static Value walkUseToGetResultBuffer(
    OpBuilder &b, Value value, const BufferizationPlan &plan,
    const BlockAndValueMapping &bvm,
    SmallVectorImpl<std::pair<OpOperand *, Value>> &traversedUses) {
  Operation *user = nullptr;
  while (value.hasOneUse()) {
    OpOperand &use = *value.use_begin();
    user = use.getOwner();
    if (isa<IREE::Flow::DispatchTensorStoreOp, tensor::InsertSliceOp>(user)) {
      return getSubviewOpForTensorStoreOp(b, user, bvm);
    }
    value = getTiedResultForOperand(use, plan);
    if (!value) return nullptr;
    traversedUses.push_back(std::make_pair(&use, value));
    if (auto resultBuffer = bvm.lookupOrNull(value)) return resultBuffer;
  }
  return nullptr;
}

/// For an operation whose `resultValue` is the result of the dispatch region,
/// gets the buffer to use to compute the value in-place.
static Value getInplaceResultBuffer(OpBuilder &b, OpResult resultValue,
                                    const BufferizationPlan &plan,
                                    BlockAndValueMapping &bvm) {
  // Traverse the use-def chains to get the `flow.dispatch.tensor.store`
  // operation keeping track of all the traversed operations. Note that the
  // equivalence set construction should ensure that all operations traversed
  // here have a single use.
  SmallVector<std::pair<OpOperand *, Value>> traversedUses;
  Value resultBuffer =
      walkUseToGetResultBuffer(b, resultValue, plan, bvm, traversedUses);
  if (!resultBuffer) return nullptr;

  DEBUG_WITH_TYPE(DEBUG_TYPE, {
    llvm::dbgs() << "Pair :\n\tTensor :";
    resultValue.getOwner()->print(llvm::dbgs());
    llvm::dbgs() << "\nt\tMemref :";
    resultBuffer.print(llvm::dbgs());
    llvm::dbgs() << "\n";
  });

  // Now replay the instructions that are essentially doing type-conversion, in
  // reverse, to get the type needed for the operation computing the value.
  for (auto &it : llvm::reverse(traversedUses)) {
    Operation *op = it.first->getOwner();
    resultBuffer =
        TypeSwitch<Operation *, Value>(op)
            .Case<scf::IfOp, scf::ForOp, linalg::LinalgOp,
                  IREE::LinalgExt::LinalgExtOp, tensor::InsertSliceOp,
                  vector::TransferWriteOp>(
                [&](auto op) { return resultBuffer; })
            .Case<linalg::TensorCollapseShapeOp, linalg::TensorExpandShapeOp>(
                [&](auto reshapeOp) {
                  return getReverseOfReshapeOp(b, reshapeOp, resultBuffer);
                })
            .Case<tensor::CastOp>([&](tensor::CastOp castOp) {
              return getReverseOfCastOp(b, castOp, resultBuffer);
            })
            .Default([&](Operation *) { return nullptr; });
    // TODO(ravishankarm): Maybe this needs to be an error. This should not have
    // happened.
    if (!resultBuffer) return nullptr;
    Value tiedResult = it.second;
    bvm.map(tiedResult, resultBuffer);
    DEBUG_WITH_TYPE(DEBUG_TYPE, {
      llvm::dbgs() << "Pair :\n\tTensor result "
                   << tiedResult.cast<OpResult>().getResultNumber() << " :";
      op->print(llvm::dbgs());
      llvm::dbgs() << "\nt\tMemref :";
      resultBuffer.print(llvm::dbgs());
      llvm::dbgs() << "\n";
    });
  }
  return resultBuffer;
}

/// Converts a `tensor.cast` operation into a `memref.cast` operation with the
/// result aliasing the buffer for the operand.
static Value getAliasingBufferForResult(OpBuilder &b, tensor::CastOp castOp,
                                        BlockAndValueMapping &bvm) {
  Value inputBuffer = bvm.lookup(castOp.source());
  Value resultTensor = castOp.dest();
  auto outputType = getMemrefTypeForTensor(
      resultTensor.getType().cast<RankedTensorType>(), {},
      inputBuffer.getType().cast<MemRefType>().getMemorySpace());
  return b.create<memref::CastOp>(castOp.getLoc(), outputType, inputBuffer);
}

/// Returns the subview that indexes into the source of the interface buffer.
static Value getAliasingBufferForResult(OpBuilder &b,
                                        IREE::Flow::DispatchTensorLoadOp loadOp,
                                        BlockAndValueMapping &bvm) {
  Location loc = loadOp.getLoc();
  Value memref = bvm.lookup(loadOp.source());
  return createSubviewOp(b, loc,
                         loadOp.result().getType().cast<ShapedType>().getRank(),
                         memref, loadOp.getMixedOffsets(),
                         loadOp.getMixedSizes(), loadOp.getMixedStrides());
}

/// Converts a `linalg.tensor_collapse/expand_shape` operation to a
/// `linalg.collapse/expand_shape` operation with the result aliasing the buffer
/// for the operand.
template <typename TensorReshapeOpTy>
static Value getAliasingBufferForReshapeResult(OpBuilder &b,
                                               TensorReshapeOpTy op,
                                               BlockAndValueMapping &bvm) {
  Location loc = op.getLoc();
  Value srcTensor = op.src();
  RankedTensorType resultTensorType = op.getResultType();
  Value inputBuffer = bvm.lookup(srcTensor);

  // Create the reshape op.
  MemRefType inputBufferType = inputBuffer.getType().cast<MemRefType>();
  auto reshapeResultType = getMemrefTypeForTensor(
      resultTensorType, {}, inputBufferType.getMemorySpace());
  using ReshapeOpTy = typename std::conditional<
      std::is_same<TensorReshapeOpTy, linalg::TensorCollapseShapeOp>::value,
      memref::CollapseShapeOp, memref::ExpandShapeOp>::type;
  Value bufferReshape = b.create<ReshapeOpTy>(loc, reshapeResultType,
                                              inputBuffer, op.reassociation());
  return bufferReshape;
}

/// Converts a `subtensor` operation to a `subview` operation.
static Value getAliasingBufferForResult(OpBuilder &b, tensor::ExtractSliceOp op,
                                        BlockAndValueMapping &bvm) {
  Location loc = op.getLoc();
  Value srcTensor = op.source();
  Value inputBuffer = bvm.lookup(srcTensor);

  return createSubviewOp(b, loc, op.getType().getRank(), inputBuffer,
                         op.getMixedOffsets(), op.getMixedSizes(),
                         op.getMixedStrides());
}

/// Returns output buffers that aliases inputs.
static SmallVector<Value> getAliasingBuffersForResult(
    scf::ForOp scfFor, BlockAndValueMapping &bvm) {
  SmallVector<Value> aliasedBuffers(scfFor.results().size(), nullptr);
  for (int i = 0; i < scfFor.results().size(); ++i) {
    Value inputTensor = scfFor.initArgs()[i];
    if (!inputTensor.getType().isa<RankedTensorType>()) continue;
    Value inputBuffer = bvm.lookup(inputTensor);
    aliasedBuffers[i] = inputBuffer;
  }
  return aliasedBuffers;
}

/// Returns a `memref` for every result that aliases the buffer for one of its
/// operands. Returns the memref of the right shape/type based on the operation.
static SmallVector<Value, 4> getAliasingBuffersForResults(
    OpBuilder &b, Operation *op, BlockAndValueMapping &bvm) {
  return TypeSwitch<Operation *, SmallVector<Value, 4>>(op)
      .Case<IREE::Flow::DispatchTensorLoadOp, tensor::ExtractSliceOp,
            tensor::CastOp>([&](auto singleResultOp) -> SmallVector<Value, 4> {
        return {getAliasingBufferForResult(b, singleResultOp, bvm)};
      })
      .Case<linalg::TensorCollapseShapeOp, linalg::TensorExpandShapeOp>(
          [&](auto reshapeOp) -> SmallVector<Value, 4> {
            return {getAliasingBufferForReshapeResult(b, reshapeOp, bvm)};
          })
      .Case<scf::ForOp>([&](auto scfFor) -> SmallVector<Value> {
        return getAliasingBuffersForResult(scfFor, bvm);
      })
      .Default([&](Operation *op) -> SmallVector<Value, 4> {
        return SmallVector<Value, 4>(op->getNumResults(), nullptr);
      });
}

/// For a result value, gets the operand that is tied with it. If no such
/// operand exists, returns `nullptr`.
static Value getTiedOperandForResult(OpResult result,
                                     const BufferizationPlan &plan) {
  for (Value operand : result.getOwner()->getOperands()) {
    if (plan.isEquivalent(operand, result)) {
      return operand;
    }
  }
  return nullptr;
}
static bool hasTiedOperandForResult(OpResult result,
                                    const BufferizationPlan &plan) {
  return static_cast<bool>(getTiedOperandForResult(result, plan));
}

/// Computes the `memrefs` to use for the result of an operation based on
/// - If the result has a tied operand reuse the buffer for the tied operand (or
///   an alias of it) as the buffer for the result. The `alaisingBuffer` vector
///   is expected to be as large as the number of results.
/// - If the result has no tied operands, the corresponding position in the
///   `aliasingBuffer` list must be `nullptr`.
/// - If the result is in the same equivalence set as the result of the dispatch
///   region (i.e. `value` operand of a `flow.dispatch.tensor.store`) then
///   return an alias/view of the buffer passed into the dispatch region to
///   store the results.
/// - Lastly, allocate a temporary buffer for the result using the passed
///   allocation function.
static LogicalResult getOrAllocateResultBuffers(
    OpBuilder &b, Operation *op, ArrayRef<Value> aliasingBuffers,
    BlockAndValueMapping &bvm, BufferizationPlan &plan,
    WorkgroupMemoryAllocationFn allocationFn) {
  assert(aliasingBuffers.size() == op->getNumResults());
  auto results = op->getResults();
  for (auto result : llvm::enumerate(results)) {
    if (!result.value().getType().isa<RankedTensorType>() ||
        bvm.contains(result.value())) {
      continue;
    }
    Value buffer;

    if (aliasingBuffers[result.index()] &&
        hasTiedOperandForResult(result.value(), plan)) {
      buffer = aliasingBuffers[result.index()];
    }
    if (!buffer && plan.isInStoreSet(result.value())) {
      buffer = getInplaceResultBuffer(b, result.value(), plan, bvm);
    }
    if (!buffer) {
      buffer = allocateBufferForResult(b, op, result.index(), allocationFn);
    }
    if (!buffer) {
      return op->emitError("unable to get result buffer for op");
    }
    bvm.map(result.value(), buffer);
    DEBUG_WITH_TYPE(DEBUG_TYPE, {
      llvm::dbgs() << "Pair :\n\tTensor result " << result.index() << ":";
      op->print(llvm::dbgs());
      llvm::dbgs() << "\nt\tMemref :";
      buffer.print(llvm::dbgs());
      llvm::dbgs() << "\n";
    });
  }
  return success();
}

/// Convenience wrapper around core allocation function for the case where the
/// alias is the buffer for the result directly.
static LogicalResult getOrAllocateResultBuffers(
    OpBuilder &b, Operation *op, BlockAndValueMapping &bvm,
    BufferizationPlan &plan, WorkgroupMemoryAllocationFn allocationFn) {
  auto aliasingBuffers = llvm::to_vector<4>(
      llvm::map_range(op->getResults(), [&](OpResult result) {
        Value tiedOperand = getTiedOperandForResult(result, plan);
        return tiedOperand ? bvm.lookupOrNull(tiedOperand) : tiedOperand;
      }));
  return getOrAllocateResultBuffers(b, op, aliasingBuffers, bvm, plan,
                                    allocationFn);
}

/// Generic conversion pattern that matches any linalg::LinalgOp. This avoids
/// template instantiating one pattern for each linalg::LinalgOp. The method
/// expects all operands and results have already been mapped to memrefs.
template <typename OpTy>
static LogicalResult convertAnyLinalgOp(
    OpBuilder &b, OpTy op, BlockAndValueMapping &bvm, BufferizationPlan &plan,
    WorkgroupMemoryAllocationFn allocationFn) {
  // Skip linalg ops inserted by this pass.
  if (op.hasBufferSemantics()) return success();

  Location loc = op.getLoc();
  SmallVector<Value, 2> newInputBuffers;
  newInputBuffers.reserve(op.getNumInputs());
  for (OpOperand *opOperand : op.getInputOperands()) {
    if (op.isScalar(opOperand)) {
      newInputBuffers.push_back(opOperand->get());
      continue;
    }
    // For `linalg.poolin_*` ops, the input might be from a
    // `linalg.init_tensor`. In such cases, the `BlockAndValueMapping` wont have
    // a mapping for the buffer. Allocate a buffer for these.
    Value inputBuffer = bvm.lookupOrNull(opOperand->get());
    if (!inputBuffer) {
      OpResult definingOpResult = opOperand->get().dyn_cast<OpResult>();
      if (!definingOpResult) return failure();
      inputBuffer = allocateBufferForResult(b, definingOpResult.getOwner(),
                                            definingOpResult.getResultNumber(),
                                            allocationFn);
    }
    newInputBuffers.push_back(inputBuffer);
  }
  SmallVector<Value, 2> newOutputBuffers;
  auto results = op.getOperation()->getResults();
  auto outputs = op.getOutputOperands();
  for (auto it : llvm::zip(results, outputs)) {
    Value resultTensor = std::get<0>(it);
    Value resultBuffer = bvm.lookup(resultTensor);

    OpOperand *outOperand = std::get<1>(it);
    Value outTensor = outOperand->get();
    Value outBuffer = bvm.lookupOrNull(outTensor);
    if (outBuffer && !plan.isEquivalent(outTensor, resultTensor) &&
        op.payloadUsesValueFromOperand(outOperand)) {
      b.create<linalg::CopyOp>(loc, outBuffer, resultBuffer);
    }
    newOutputBuffers.push_back(resultBuffer);
  }

  SmallVector<Value, 8> newOperands(newInputBuffers.begin(),
                                    newInputBuffers.end());
  newOperands.append(newOutputBuffers.begin(), newOutputBuffers.end());
  op.clone(b, loc, {}, newOperands);
  return success();
}

/// Constants that return tensor types can be handled natively by the
/// backends. Here just provide a cast to memref to bridge the gap from tensors
/// to memrefs.
static LogicalResult convertConstantOp(OpBuilder &b,
                                       arith::ConstantOp constantOp,
                                       BlockAndValueMapping &bvm) {
  Value result = constantOp.getResult();
  assert(!bvm.lookupOrNull(result));
  RankedTensorType tensorType = result.getType().dyn_cast<RankedTensorType>();
  if (!tensorType) return success();
  OpBuilder::InsertionGuard g(b);
  b.setInsertionPointAfter(constantOp);
  auto memrefType = getMemrefTypeForTensor(tensorType);
  Value memref =
      b.create<memref::BufferCastOp>(constantOp.getLoc(), memrefType, result);
  bvm.map(result, memref);
  return success();
}

static LogicalResult convertDispatchTieShapeOp(
    OpBuilder &b, IREE::Flow::DispatchTieShapeOp shapeOp,
    BlockAndValueMapping &bvm) {
  if (Value v = bvm.lookupOrNull(shapeOp.operand())) {
    auto tieShapeOp = b.create<Shape::TieShapeOp>(shapeOp.getLoc(), v.getType(),
                                                  v, shapeOp.shape());
    bvm.map(shapeOp.getResult(), tieShapeOp.getResult());
  }
  return success();
}

/// Converts a `tensor.extract` operation into a `load`.
static LogicalResult convertTensorExtractOp(OpBuilder &b, tensor::ExtractOp op,
                                            const BlockAndValueMapping &bvm) {
  OpBuilder::InsertionGuard g(b);
  b.setInsertionPoint(op);
  Value inputBuffer = bvm.lookup(op.tensor());
  Value load =
      b.createOrFold<memref::LoadOp>(op.getLoc(), inputBuffer, op.indices());
  // Since the value is the scalar, and `bvm` is used to only track tensor ->
  // memref mappings, just replace the uses directly.
  op.result().replaceAllUsesWith(load);
  return success();
}

/// Converts a `flow.dispatch.tensor.store` operation to memrefs. If the `value`
/// and `target` are in the same equivalent set, then there is nothing to do. If
/// no create a subview into the result buffer and copy the `value`.
static LogicalResult convertInterfaceStoreTensorOp(
    OpBuilder &b, IREE::Flow::DispatchTensorStoreOp storeOp,
    BlockAndValueMapping &bvm, BufferizationPlan &plan) {
  if (plan.isEquivalent(storeOp.target(), storeOp.value())) {
    return success();
  }
  OpBuilder::InsertionGuard g(b);
  b.setInsertionPoint(storeOp);
  Value storeTo = bvm.lookup(storeOp.target());
  Value storeFrom = bvm.lookup(storeOp.value());

  Value subview = createSubviewOp(
      b, storeOp.getLoc(), storeFrom.getType().cast<ShapedType>().getRank(),
      storeTo, storeOp.getMixedOffsets(), storeOp.getMixedSizes(),
      storeOp.getMixedStrides());

  b.create<linalg::CopyOp>(storeOp->getLoc(), storeFrom, subview);
  return success();
}

/// Converts a `tensor.insert_slice` operation to buffers by
/// - Allocating a buffer for the result (if needed), and copying the
///   destination value into this buffer.
/// - Copying the source values into a subview of the result buffer.
static LogicalResult convertSubTensorInsertOp(OpBuilder &b,
                                              tensor::InsertSliceOp op,
                                              BlockAndValueMapping &bvm,
                                              BufferizationPlan &plan) {
  Location loc = op.getLoc();
  Value result = op.getResult();
  Value resultBuffer = bvm.lookup(result);

  // If `dest` and `result` are not equivalent, need a copy for that.
  Value dest = op.dest();
  if (!plan.isEquivalent(dest, result)) {
    Value destBuffer = bvm.lookup(dest);
    b.create<linalg::CopyOp>(loc, destBuffer, resultBuffer);
  }

  Value source = op.source();
  if (plan.isEquivalent(source, dest)) {
    return success();
  }

  // Copy from the source to the result subview.
  ShapedType sourceType = op.getSourceType();
  Value sourceBuffer = bvm.lookup(source);
  SmallVector<OpFoldResult> offsets = op.getMixedOffsets();
  SmallVector<OpFoldResult> sizes = op.getMixedSizes();
  SmallVector<OpFoldResult> strides = op.getMixedStrides();
  Value subViewOp = createSubviewOp(b, loc, sourceType.getRank(), resultBuffer,
                                    offsets, sizes, strides);
  b.create<linalg::CopyOp>(loc, sourceBuffer, subViewOp);
  return success();
}

/// Converts a `tensor.insert` operations into a `memref.store`.
static LogicalResult convertTensorInsertOp(OpBuilder &b, tensor::InsertOp op,
                                           BlockAndValueMapping &bvm,
                                           BufferizationPlan &plan) {
  Location loc = op.getLoc();
  Value result = op.result();
  Value resultBuffer = bvm.lookup(result);
  if (!plan.isEquivalent(op.dest(), result)) {
    Value destBuffer = bvm.lookup(op.dest());
    b.create<linalg::CopyOp>(loc, destBuffer, resultBuffer);
  }

  b.create<memref::StoreOp>(loc, op.scalar(), resultBuffer, op.indices());
  return success();
}

/// Converts a vector.transfer_write op to use memref operands for source.
static LogicalResult convertVectorTransferWriteOp(OpBuilder &b,
                                                  vector::TransferWriteOp op,
                                                  BlockAndValueMapping &bvm,
                                                  BufferizationPlan &plan) {
  Location loc = op.getLoc();
  Value result = op.result();
  RankedTensorType resultType = result.getType().dyn_cast<RankedTensorType>();
  if (!resultType) return success();
  Value resultBuffer = bvm.lookup(result);

  if (!plan.isEquivalent(op.source(), result) &&
      // If the source is linalg.init_tensor, then we don't care about the
      // initial value and can avoid the copy.
      !op.source().getDefiningOp<linalg::InitTensorOp>()) {
    Value destBuffer = bvm.lookup(op.source());
    b.create<linalg::CopyOp>(loc, destBuffer, resultBuffer);
  }

  // Create a new vector.transfer_write operation without a result value.
  b.create<vector::TransferWriteOp>(
      loc, op.vector(), resultBuffer, op.indices(), op.permutation_map(),
      op.mask(), op.in_bounds() ? *op.in_bounds() : ArrayAttr());
  return success();
}

static LogicalResult convertScfForOp(OpBuilder &b, scf::ForOp forOp,
                                     BlockAndValueMapping &bvm,
                                     BufferizationPlan &plan) {
  auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
  Location loc = forOp.getLoc();
  for (auto arg : llvm::enumerate(forOp.getRegionIterArgs())) {
    if (!arg.value().getType().isa<RankedTensorType>()) continue;
    Value resultTensor = forOp.getResult(arg.index());
    Value resultBuffer = bvm.lookup(resultTensor);
    OpOperand &initOperand = forOp.getOpOperandForRegionIterArg(arg.value());
    Value yieldOperand = yieldOp.getOperand(arg.index());
    bvm.map(arg.value(), resultBuffer);
    bvm.map(yieldOperand, resultBuffer);
    if (!plan.isEquivalent(arg.value(), initOperand.get())) {
      Value initBuffer = bvm.lookup(initOperand.get());
      b.create<linalg::CopyOp>(loc, initBuffer, resultBuffer);
    }
  }
  return success();
}

static LogicalResult convertScfIfOp(OpBuilder &b, scf::IfOp ifOp,
                                    BlockAndValueMapping &bvm,
                                    BufferizationPlan &plan) {
  auto thenYieldOp = ifOp.thenYield();
  auto elseYieldOp = ifOp.elseYield();
  for (auto result : llvm::enumerate(ifOp.getResults())) {
    Value resultValue = result.value();
    if (!resultValue.getType().isa<RankedTensorType>()) continue;
    Value resultBuffer = bvm.lookup(resultValue);
    Value thenYield = thenYieldOp->getOperand(result.index());
    Value elseYield = elseYieldOp->getOperand(result.index());
    bvm.map(thenYield, resultBuffer);
    bvm.map(elseYield, resultBuffer);
  }
  return success();
}

/// If the alias of the buffer for an input oeprand cannot be used for the
/// "tied" results, need to do an explicit copy of the memory pointed to by the
/// aliased buffer into the buffer assigned to the result.
static void copyFromAliasingBufferToResultBuffer(
    OpBuilder &b, Location loc, ArrayRef<Value> tiedOperands,
    ArrayRef<Value> tiedResults, ArrayRef<Value> aliasingBuffers,
    BlockAndValueMapping &bvm, BufferizationPlan &plan) {
  for (auto result : enumerate(tiedResults)) {
    Value operand = tiedOperands[result.index()];
    if (!plan.isEquivalent(result.value(), operand)) {
      b.create<linalg::CopyOp>(loc, aliasingBuffers[result.index()],
                               bvm.lookup(result.value()));
    }
  }
}

/// Returns the static/dynamic mixed sizes of the memref.
static SmallVector<OpFoldResult> getMemrefSizes(OpBuilder &b, Location loc,
                                                Value memref) {
  auto inputShape = memref.getType().cast<ShapedType>().getShape();
  SmallVector<OpFoldResult> sizeMixedValues;
  for (int64_t i = 0; i < inputShape.size(); ++i) {
    if (inputShape[i] == ShapedType::kDynamicSize) {
      Value dim = b.create<memref::DimOp>(loc, memref, i);
      sizeMixedValues.push_back(dim);
    } else {
      sizeMixedValues.push_back(b.getI64IntegerAttr(inputShape[i]));
    }
  }
  return sizeMixedValues;
}

static LogicalResult convertPadTensorOp(OpBuilder &b,
                                        linalg::PadTensorOp padTensorOp,
                                        BlockAndValueMapping &bvm) {
  auto inputTensor = padTensorOp.source();
  auto inputMemref = bvm.lookup(inputTensor);

  auto loc = padTensorOp.getLoc();

  auto resultPaddedBuffer = bvm.lookup(padTensorOp.result());

  // Get padding value and fill the result buffer.
  linalg::YieldOp yeildOp =
      *padTensorOp.region().getOps<linalg::YieldOp>().begin();
  Value paddingValue = yeildOp.values()[0];

  auto constOp = paddingValue.getDefiningOp<arith::ConstantOp>();
  if (!constOp) {
    return padTensorOp.emitError(
        "Converting linalg.pad_tensor with non-constant padding value");
  }
  if (constOp.value().isa<DenseElementsAttr>()) {
    return padTensorOp.emitError(
        "Converting linalg.pad_tensor with non-scalar constant padding "
        "value");
  }

  b.create<linalg::FillOp>(loc, paddingValue, resultPaddedBuffer);

  // Get the interior region.
  SmallVector<OpFoldResult> sizeMixedValues =
      getMemrefSizes(b, loc, inputMemref);
  SmallVector<OpFoldResult> strides(
      inputMemref.getType().cast<ShapedType>().getRank(),
      b.getI64IntegerAttr(1));

  auto resultSubView = b.create<memref::SubViewOp>(loc, resultPaddedBuffer,
                                                   padTensorOp.getMixedLowPad(),
                                                   sizeMixedValues, strides);
  // Copy to the interior region.
  b.create<linalg::CopyOp>(loc, inputMemref, resultSubView);
  return success();
}

namespace {
/// Pass to convert from tensor based ops to memref based ops.
class LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
 public:
  LinalgBufferizePass(WorkgroupMemoryAllocationFn fn) : allocationFn(fn) {}
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<IREE::Util::UtilDialect, linalg::LinalgDialect,
                    memref::MemRefDialect, scf::SCFDialect, StandardOpsDialect,
                    mlir::math::MathDialect, mlir::arith::ArithmeticDialect>();
  }
  void runOnOperation() override;

 private:
  WorkgroupMemoryAllocationFn allocationFn;
};
}  // namespace

void LinalgBufferizePass::runOnOperation() {
  BufferizationPlan plan;
  FuncOp funcOp = getOperation();
  if (failed(createTensorEquivalenceClasses(funcOp, plan))) {
    return signalPassFailure();
  }

  MLIRContext *context = &getContext();
  OpBuilder b(context);

  BlockAndValueMapping bvm;

  // First go over all hal.interface.binding.subspan ops and create counterparts
  // working with memrefs.
  funcOp.walk([&](IREE::HAL::InterfaceBindingSubspanOp op) {
    auto shapedType =
        op.getResult().getType().dyn_cast<IREE::Flow::DispatchTensorType>();
    if (!shapedType || !shapedType.hasRank()) return;
    OpBuilder::InsertionGuard g(b);
    b.setInsertionPoint(op);
    // Just change the result type of the InterfaceBindingSubspanOp to form
    // the base buffer.
    auto tensorType =
        op.result().getType().cast<IREE::Flow::DispatchTensorType>();
    auto memRefType = getMemrefTypeForTensor(tensorType);
    auto baseBuffer = b.create<IREE::HAL::InterfaceBindingSubspanOp>(
        op->getLoc(), memRefType, op.binding(), op.byte_offset(),
        op.byte_length(), op.dynamic_dims());
    bvm.map(op, baseBuffer);
  });

  // Visit all the operations that return `tensor`s and convert them to using
  // `memref`s.
  auto convertTensorProducingOps = [&](Operation *op) -> WalkResult {
    return TypeSwitch<Operation *, LogicalResult>(op)
        .Case<arith::ConstantOp>([&](arith::ConstantOp constantOp) {
          return convertConstantOp(b, constantOp, bvm);
        })
        .Case<IREE::Flow::DispatchTensorStoreOp>(
            [&](IREE::Flow::DispatchTensorStoreOp storeOp) {
              return convertInterfaceStoreTensorOp(b, storeOp, bvm, plan);
            })
        .Case<IREE::Flow::DispatchTieShapeOp>(
            [&](IREE::Flow::DispatchTieShapeOp shapeOp) {
              return convertDispatchTieShapeOp(b, shapeOp, bvm);
            })
        .Case<scf::ForOp>([&](scf::ForOp forOp) {
          if (failed(getOrAllocateResultBuffers(b, forOp, bvm, plan,
                                                allocationFn))) {
            return failure();
          }
          return convertScfForOp(b, forOp, bvm, plan);
        })
        .Case<scf::IfOp>([&](scf::IfOp ifOp) {
          if (failed(getOrAllocateResultBuffers(b, ifOp, bvm, plan,
                                                allocationFn))) {
            return failure();
          }
          return convertScfIfOp(b, ifOp, bvm, plan);
        })
        .Case<IREE::Flow::DispatchTensorLoadOp, linalg::TensorCollapseShapeOp,
              linalg::TensorExpandShapeOp, tensor::ExtractSliceOp,
              tensor::CastOp>([&](auto aliasingOp) {
          auto aliasingBuffers =
              getAliasingBuffersForResults(b, aliasingOp, bvm);
          if (failed(getOrAllocateResultBuffers(b, aliasingOp, aliasingBuffers,
                                                bvm, plan, allocationFn))) {
            return failure();
          }
          copyFromAliasingBufferToResultBuffer(
              b, aliasingOp->getLoc(), aliasingOp->getOperand(0),
              aliasingOp->getResult(0), aliasingBuffers, bvm, plan);
          return success();
        })
        .Case<linalg::PadTensorOp>([&](linalg::PadTensorOp padTensorOp) {
          if (failed(getOrAllocateResultBuffers(b, padTensorOp, bvm, plan,
                                                allocationFn))) {
            return failure();
          }
          return convertPadTensorOp(b, padTensorOp, bvm);
        })
        .Case<linalg::LinalgOp, IREE::LinalgExt::LinalgExtOp>([&](auto op) {
          if (failed(
                  getOrAllocateResultBuffers(b, op, bvm, plan, allocationFn))) {
            return failure();
          }
          return convertAnyLinalgOp(b, op, bvm, plan, allocationFn);
        })
        .Case<tensor::InsertSliceOp>(
            [&](tensor::InsertSliceOp subTensorInsertOp) {
              if (failed(getOrAllocateResultBuffers(b, subTensorInsertOp, bvm,
                                                    plan, allocationFn))) {
                return failure();
              }
              return convertSubTensorInsertOp(b, subTensorInsertOp, bvm, plan);
            })
        .Case<tensor::InsertOp>([&](tensor::InsertOp tensorInsertOp) {
          if (failed(getOrAllocateResultBuffers(b, tensorInsertOp, bvm, plan,
                                                allocationFn))) {
            return failure();
          }
          return convertTensorInsertOp(b, tensorInsertOp, bvm, plan);
        })
        .Case<vector::TransferWriteOp>(
            [&](vector::TransferWriteOp transferWriteOp) {
              if (failed(getOrAllocateResultBuffers(b, transferWriteOp, bvm,
                                                    plan, allocationFn))) {
                return failure();
              }
              return convertVectorTransferWriteOp(b, transferWriteOp, bvm,
                                                  plan);
            })
        .Default([&](Operation *op) { return success(); });
  };
  auto walkResult =
      funcOp.walk<WalkOrder::PreOrder>([&](Operation *op) -> WalkResult {
        b.setInsertionPoint(op);
        return convertTensorProducingOps(op);
      });
  if (walkResult.wasInterrupted()) {
    return signalPassFailure();
  }

  // Lastly visit the non-tensor return operations that still use `tensor`
  // values. These need to be updated to use the corresponding `memref` values,
  // but dont need to update the block-and-value mapping.
  auto convertNonTensorProducingOps = [&](Operation *op) -> LogicalResult {
    return TypeSwitch<Operation *, LogicalResult>(op)
        .Case<tensor::ExtractOp>([&](tensor::ExtractOp op) {
          return convertTensorExtractOp(b, op, bvm);
        })
        .Case<vector::TransferReadOp>([&](auto op) {
          for (unsigned i : llvm::seq<unsigned>(0, op->getNumOperands())) {
            Value operand = op->getOperand(i);
            if (operand.getType().isa<RankedTensorType>()) {
              Value remappedVal = bvm.lookupOrNull(operand);
              if (remappedVal) op->setOperand(i, remappedVal);
            }
          }
          return success();
        })
        .Case<tensor::DimOp>([&](tensor::DimOp dimOp) {
          Value operand = dimOp.source();
          Value remappedVal = bvm.lookupOrNull(operand);
          if (remappedVal) {
            Value newDimOp = b.create<memref::DimOp>(
                dimOp.getLoc(), remappedVal, dimOp.index());
            dimOp.replaceAllUsesWith(newDimOp);
          }
          return success();
        })
        .Case<scf::ForOp>([&](scf::ForOp forOp) {
          // To canonicalize the `scf.for` tensor result/operand/yield value
          // away, forward the init argument to the yeild of the loop.
          auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
          for (auto arg : llvm::enumerate(forOp.getIterOperands())) {
            if (!arg.value().getType().isa<RankedTensorType>()) continue;
            yieldOp.setOperand(arg.index(), arg.value());
          }
          return success();
        })
        .Case<IREE::Flow::DispatchTensorStoreOp>([&](auto op) {
          op.erase();
          return success();
        })
        .Default([&](Operation *op) { return success(); });
  };

  walkResult = funcOp.walk([&](Operation *op) -> WalkResult {
    b.setInsertionPoint(op);
    return convertNonTensorProducingOps(op);
  });
  if (walkResult.wasInterrupted()) {
    return signalPassFailure();
  }
}

static Value defaultAllocationFn(OpBuilder &builder, Location loc,
                                 ArrayRef<int64_t> staticShape,
                                 Type elementType,
                                 ArrayRef<Value> dynamicSizes) {
  auto allocationType = MemRefType::get(staticShape, elementType);
  return builder.create<memref::AllocOp>(loc, allocationType, dynamicSizes);
}

std::unique_ptr<OperationPass<FuncOp>> createLinalgBufferizePass(
    WorkgroupMemoryAllocationFn allocationFn) {
  return std::make_unique<LinalgBufferizePass>(
      allocationFn ? allocationFn : defaultAllocationFn);
}

void addLinalgBufferizePasses(OpPassManager &passManager,
                              WorkgroupMemoryAllocationFn allocationFn) {
  passManager.addNestedPass<FuncOp>(createLinalgBufferizePass(allocationFn));
  passManager.addPass(memref::createResolveShapedTypeResultDimsPass());
  passManager.addNestedPass<FuncOp>(createCanonicalizerPass());
  passManager.addNestedPass<FuncOp>(createCSEPass());
  passManager.addNestedPass<FuncOp>(createCleanupBufferAllocViewPass());
  // passManager.addPass(createBufferHoistingPass());
  // TODO(nicolasvasilache): bug in buffer loop hoisting with
  // dynamic_linalg_matmul_on_tensors_fuse_0.mlir
  // passManager.addPass(createBufferLoopHoistingPass());
}

}  // namespace iree_compiler
}  // namespace mlir
