ConvertToDestinationStyle.cpp

Go to the documentation of this file.

//===- ConvertToDestinationStyle.cpp - Convert non-DPS to DPS ops ---------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file contains patterns to convert non-DPS ops to DPS ops. New
// tensor.empty ops are inserted as a destination. Such tensor.empty can be
// eliminated with "empty tensor elimination", allowing them to bufferize
// without an allocation (assuming there are no further conflicts).
//
//===----------------------------------------------------------------------===//
//
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "llvm/ADT/STLExtras.h"
 
using namespace mlir;
using namespace mlir::tensor;
 
// Implements backtracking to traverse indices of the output buffer while
// iterating over op.elements().
static Value createInserts(RewriterBase &rewriter, Location loc, int dim,
                           Value destination, ArrayRef<int64_t> shape,
                           ArrayRef<Value> constants,
                           OperandRange::iterator &elementIt,
                           SmallVectorImpl<Value> &indices) {
  if (dim == static_cast<int>(shape.size()) - 1) {
    for (int i = 0; i < shape.back(); ++i) {
      indices.back() = constants[i];
      destination = tensor::InsertOp::create(rewriter, loc, *elementIt,
                                             destination, indices);
      ++elementIt;
    }
    return destination;
  }
  for (int i = 0; i < shape[dim]; ++i) {
    indices[dim] = constants[i];
    destination = createInserts(rewriter, loc, dim + 1, destination, shape,
                                constants, elementIt, indices);
  }
  return destination;
}
 
/// Create a memcpy from the given source tensor to the given destination
/// memref. The copy op type can be specified in the `options`.
static void createMemcpy(OpBuilder &b, Location loc, Value tensorSource,
                         Value memrefDest,
                         const linalg::BufferizeToAllocationOptions &options) {
  auto tensorType = dyn_cast<RankedTensorType>(tensorSource.getType());
  assert(tensorType && "expected ranked tensor");
  assert(isa<MemRefType>(memrefDest.getType()) && "expected ranked memref");
 
  switch (options.memcpyOp) {
  case linalg::BufferizeToAllocationOptions::MemcpyOp::
      MaterializeInDestination: {
    // Note: This is the preferred way of memcpy'ing because no layout map
    // and/or memory space must be specified for the source.
    auto materializeOp = bufferization::MaterializeInDestinationOp::create(
        b, loc, tensorSource, memrefDest);
    materializeOp.setWritable(true);
  } break;
  case linalg::BufferizeToAllocationOptions::MemcpyOp::MemrefCopy: {
    // TODO: Support custom memory space on source.
    // We do not know the layout map of the source yet, so use a fully dynamic
    // layout for best compatibility.
    Value toBuffer = bufferization::ToBufferOp::create(
        b, loc, bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType),
        tensorSource, /*read_only=*/true);
    memref::CopyOp::create(b, loc, toBuffer, memrefDest);
  } break;
  case linalg::BufferizeToAllocationOptions::MemcpyOp::LinalgCopy: {
    // TODO: Support custom memory space on source.
    // We do not know the layout map of the source yet, so use a fully dynamic
    // layout for best compatibility.
    Value toBuffer = bufferization::ToBufferOp::create(
        b, loc, bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType),
        tensorSource, /*read_only=*/true);
    linalg::CopyOp::create(b, loc, toBuffer, memrefDest);
  } break;
  };
}
 
static Operation *movePaddingToFillOrGenericOp(RewriterBase &rewriter,
                                               Location loc, PadOp padOp,
                                               Value dest) {
  OpBuilder::InsertionGuard g(rewriter);
  RankedTensorType resultType = padOp.getResultType();
 
  // Examine the yielded value to decide if a linalg.generic is neede or a
  // linalg.fill is sufficient.
  Value yieldedValue =
      cast<tensor::YieldOp>(padOp.getBody()->getTerminator()).getValue();
  Attribute constYieldedValue;
  // Is the yielded value a bbArg defined outside of the PadOp?
  bool outsideBbArg =
      isa<BlockArgument>(yieldedValue) &&
      cast<BlockArgument>(yieldedValue).getOwner()->getParentOp() !=
          padOp.getOperation();
  // Is the yielded value an OpResult defined outside of the PadOp?
  bool outsideOpResult =
      isa<OpResult>(yieldedValue) &&
      yieldedValue.getDefiningOp()->getParentOp() != padOp.getOperation();
  bool invariantYieldedValue = outsideBbArg || outsideOpResult;
  if (matchPattern(yieldedValue, m_Constant(&constYieldedValue))) {
    // Padding with a constant: Create linalg.fill.
    Dialect *arithDialect =
        rewriter.getContext()->getLoadedDialect<arith::ArithDialect>();
    Value fillValue =
        arithDialect
            ->materializeConstant(rewriter, constYieldedValue,
                                  yieldedValue.getType(), yieldedValue.getLoc())
            ->getResult(0);
    auto fillOp = linalg::FillOp::create(rewriter, loc, ValueRange(fillValue),
                                         ValueRange(dest));
    return fillOp;
  }
 
  if (invariantYieldedValue) {
    // Padding with an invariant value.
    auto fillOp = linalg::FillOp::create(
        rewriter, loc, ValueRange(yieldedValue), ValueRange(dest));
    return fillOp;
  }
 
  // Create linalg.generic.
  SmallVector<utils::IteratorType> iteratorTypes(resultType.getRank(),
                                                 utils::IteratorType::parallel);
  SmallVector<AffineMap> indexingMaps(
      1, rewriter.getMultiDimIdentityMap(resultType.getRank()));
  auto genericOp = linalg::GenericOp::create(
      rewriter, loc, resultType, /*inputs=*/ValueRange(),
      /*outputs=*/ValueRange{dest}, /*indexingMaps=*/
      indexingMaps, iteratorTypes);
  Block *body = rewriter.createBlock(&genericOp->getRegion(0), {},
                                     resultType.getElementType(), loc);
  rewriter.setInsertionPointToStart(body);
  SmallVector<Value> bbArgReplacements;
  for (int64_t i = 0; i < resultType.getRank(); ++i)
    bbArgReplacements.push_back(linalg::IndexOp::create(rewriter, loc, i));
  rewriter.mergeBlocks(padOp.getBody(), body, bbArgReplacements);
 
  // Update terminator.
  auto yieldOp = cast<tensor::YieldOp>(body->getTerminator());
  rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOp.getValue());
  return genericOp;
}
 
static SmallVector<Value> reifyOrComputeDynamicSizes(OpBuilder &b,
                                                     Value value) {
  auto tensorType = cast<RankedTensorType>(value.getType());
  if (tensorType.hasStaticShape())
    return {};
 
  // Try to reify dynamic sizes.
  ReifiedRankedShapedTypeDims reifiedShape;
  if (isa<OpResult>(value) &&
      succeeded(reifyResultShapes(b, value.getDefiningOp(), reifiedShape))) {
    SmallVector<Value> dynSizes;
    for (int64_t i = 0; i < tensorType.getRank(); ++i) {
      if (tensorType.isDynamicDim(i))
        dynSizes.push_back(cast<Value>(
            reifiedShape[cast<OpResult>(value).getResultNumber()][i]));
    }
    return dynSizes;
  }
 
  // Create tensor.dim ops.
  SmallVector<Value> dynSizes;
  for (int64_t i = 0; i < tensorType.getRank(); ++i) {
    if (tensorType.isDynamicDim(i))
      dynSizes.push_back(
          DimOp::create(b, value.getLoc(), value,
                        arith::ConstantIndexOp::create(b, value.getLoc(), i)));
  }
  return dynSizes;
}
 
static Value
createAllocationForTensor(RewriterBase &rewriter, Location loc, Value value,
                          const linalg::BufferizeToAllocationOptions &options,
                          Attribute memorySpace = {}) {
  OpBuilder::InsertionGuard g(rewriter);
  auto tensorType = cast<RankedTensorType>(value.getType());
 
  // Create buffer allocation.
  auto memrefType =
      cast<MemRefType>(bufferization::getMemRefTypeWithStaticIdentityLayout(
          tensorType, memorySpace));
  SmallVector<Value> dynamicSizes = reifyOrComputeDynamicSizes(rewriter, value);
 
  Value alloc;
  if (options.allocOp ==
      linalg::BufferizeToAllocationOptions::AllocOp::MemrefAlloc) {
    alloc = memref::AllocOp::create(rewriter, loc, memrefType, dynamicSizes);
    if (options.emitDealloc) {
      // Place deallocation at the end of the block.
      rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator());
      memref::DeallocOp::create(rewriter, loc, alloc);
    }
  } else if (options.allocOp ==
             linalg::BufferizeToAllocationOptions::AllocOp::MemrefAlloca) {
    alloc = memref::AllocaOp::create(rewriter, loc, memrefType, dynamicSizes);
    // No dealloc is needed.
  }
 
  return alloc;
}
 
Value linalg::bufferizeToAllocation(
    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
    PadOp padOp, Attribute memorySpace, Operation *insertionPoint) {
  // tensor.pad does not have a destination operand.
  assert(!options.bufferizeDestinationOnly && "invalid options");
 
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(insertionPoint ? insertionPoint : padOp);
  Location loc = padOp.getLoc();
 
  // Create buffer allocation.
  Value alloc = createAllocationForTensor(rewriter, loc, padOp.getResult(),
                                          options, memorySpace);
  rewriter.setInsertionPoint(padOp);
 
  if (!padOp.hasZeroLowPad() || !padOp.hasZeroHighPad()) {
    // Create linalg.fill or linalg.generic. Not needed if there is no padding.
    Operation *fillOp =
        movePaddingToFillOrGenericOp(rewriter, loc, padOp, alloc);
    rewriter.setInsertionPointAfter(fillOp);
  }
 
  // Create memcpy.
  SmallVector<OpFoldResult> sizes =
      getMixedSizes(rewriter, loc, padOp.getSource());
  SmallVector<OpFoldResult> strides(padOp.getResultType().getRank(),
                                    rewriter.getIndexAttr(1));
  Value subview = memref::SubViewOp::create(
      rewriter, loc, alloc, /*offsets=*/padOp.getMixedLowPad(), sizes, strides);
  createMemcpy(rewriter, loc, padOp.getSource(), subview, options);
 
  // Create bufferization.to_tensor with "restrict" and "writable". The returned
  // tensor is a new buffer allocation, so it does not alias with any buffer.
  Value toTensorOp = bufferization::ToTensorOp::create(
      rewriter, loc, padOp.getResult().getType(), alloc, /*restrict=*/true,
      /*writable=*/true);
  rewriter.replaceOp(padOp, toTensorOp);
  return alloc;
}
 
Value linalg::bufferizeToAllocation(
    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
    vector::MaskOp maskOp, Attribute memorySpace, Operation *insertionPoint) {
  assert(llvm::range_size(maskOp.getMaskBlock()->without_terminator()) == 1 &&
         "expected single masked op");
  OpBuilder::InsertionGuard g(rewriter);
 
  // Should the bufferization options and state be function arguments?
  bufferization::BufferizationOptions bufferizationOptions;
  bufferization::BufferizationState bufferizationState;
 
  Operation *yieldOp = maskOp.getMaskRegion().front().getTerminator();
  assert(isa<vector::YieldOp>(yieldOp) && "expected yield op terminator");
 
  // Bufferize maskable op. By default, place the buffer allocation right before
  // the mask op.
  Value alloc = bufferizeToAllocation(
      rewriter, options, maskOp.getMaskableOp(), memorySpace,
      /*insertionPoint=*/insertionPoint ? insertionPoint : maskOp);
 
  if (options.bufferizeDestinationOnly)
    return alloc;
 
  // Bufferize terminator.
  rewriter.setInsertionPoint(yieldOp);
  if (failed(cast<bufferization::BufferizableOpInterface>(yieldOp).bufferize(
          rewriter, bufferizationOptions, bufferizationState)))
    return nullptr;
 
  // Erase dead to_tensor ops inside of the mask op. This is necessary because
  // there only be one op (apart from the terminator) inside the mask op.
  // TODO: Remove dead to_tensor ops more aggressively during bufferization.
  SmallVector<Operation *> toTensorOps;
  maskOp.walk([&](bufferization::ToTensorOp toTensorOp) {
    if (toTensorOp->getUses().empty())
      toTensorOps.push_back(toTensorOp.getOperation());
  });
  for (Operation *op : toTensorOps)
    rewriter.eraseOp(op);
 
  // Bufferize mask op.
  SmallVector<OpOperand *> resultUses;
  for (Value result : maskOp.getResults())
    if (isa<TensorType>(result.getType()))
      for (OpOperand &use : result.getUses())
        resultUses.push_back(&use);
  rewriter.setInsertionPoint(maskOp);
  if (failed(
          cast<bufferization::BufferizableOpInterface>(maskOp.getOperation())
              .bufferize(rewriter, bufferizationOptions, bufferizationState)))
    return nullptr;
 
  // Set "restrict" attribute, indicating that no other tensor aliases with
  // this tensor. That is because we just allocated a new buffer for the tensor.
  for (OpOperand *resultUse : resultUses) {
    auto toTensorOp =
        resultUse->get().getDefiningOp<bufferization::ToTensorOp>();
    assert(toTensorOp && "expected to_tensor op");
    rewriter.modifyOpInPlace(toTensorOp, [&]() {
      toTensorOp.setRestrict(true);
      toTensorOp.setWritable(true);
    });
  }
 
  return alloc;
}
 
Value linalg::bufferizeToAllocation(
    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
    bufferization::AllocTensorOp allocTensorOp, Attribute memorySpace,
    Operation *insertionPoint) {
  Location loc = allocTensorOp.getLoc();
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(insertionPoint ? insertionPoint : allocTensorOp);
  bufferization::BufferizationOptions bufferizationOptions;
 
  // Create buffer allocation.
  Value alloc = createAllocationForTensor(
      rewriter, loc, allocTensorOp.getResult(), options, memorySpace);
 
  // Create bufferization.to_tensor with "restrict" and "writable". The returned
  // tensor is a new buffer allocation, so it does not alias with any buffer.
  Value toTensorOp = bufferization::ToTensorOp::create(
      rewriter, loc, allocTensorOp.getResult().getType(), alloc,
      /*restrict=*/true,
      /*writable=*/true);
  rewriter.replaceOp(allocTensorOp, toTensorOp);
  return alloc;
}
 
/// Lower tensor.from_elements to a sequence of chained tensor.insert.
FailureOr<Operation *> mlir::linalg::rewriteInDestinationPassingStyle(
    RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp) {
  Location loc = fromElementsOp.getLoc();
  RankedTensorType tensorType =
      cast<RankedTensorType>(fromElementsOp.getType());
  auto shape = tensorType.getShape();
 
  // Create tensor.empty.
  auto emptyOp = EmptyOp::create(rewriter, loc, tensorType, ValueRange());
 
  // Case: tensor<elem_type>.
  if (shape.empty()) {
    Operation *res = rewriter.replaceOpWithNewOp<tensor::InsertOp>(
        fromElementsOp, fromElementsOp.getElements().front(),
        emptyOp.getResult(), ValueRange());
    return res;
  }
 
  // Create constants for the range of possible indices [0, max{shape_i}).
  auto maxDim = *llvm::max_element(shape);
  SmallVector<Value, 2> constants;
  constants.reserve(maxDim);
  for (int i = 0; i < maxDim; ++i)
    constants.push_back(arith::ConstantIndexOp::create(rewriter, loc, i));
 
  // Traverse all elements and create tensor.insert ops.
  auto elementIt = fromElementsOp.getElements().begin();
  SmallVector<Value, 2> indices(tensorType.getRank(), constants[0]);
  Value result = createInserts(rewriter, loc, /*dim=*/0, emptyOp.getResult(),
                               shape, constants, elementIt, indices);
 
  // Replace tensor.from_elements.
  rewriter.replaceOp(fromElementsOp, result);
  return result.getDefiningOp();
}
 
/// Lower tensor.generate to linalg.generic.
FailureOr<Operation *>
mlir::linalg::rewriteInDestinationPassingStyle(RewriterBase &rewriter,
                                               tensor::GenerateOp generateOp) {
  // Only ops with exactly one block are supported.
  if (!generateOp.getBody().hasOneBlock())
    return failure();
 
  Location loc = generateOp.getLoc();
  RankedTensorType tensorType = cast<RankedTensorType>(generateOp.getType());
 
  // Create tensor.empty.
  auto emptyOp = EmptyOp::create(rewriter, loc, tensorType,
                                 generateOp.getDynamicExtents());
 
  // Create linalg.generic.
  SmallVector<utils::IteratorType> iteratorTypes(tensorType.getRank(),
                                                 utils::IteratorType::parallel);
  SmallVector<AffineMap> indexingMaps(
      1, rewriter.getMultiDimIdentityMap(tensorType.getRank()));
  auto genericOp = linalg::GenericOp::create(
      rewriter, loc, tensorType, /*inputs=*/ValueRange(),
      /*outputs=*/ValueRange{emptyOp.getResult()}, /*indexingMaps=*/
      indexingMaps, iteratorTypes);
  Block *body = rewriter.createBlock(&genericOp->getRegion(0), {},
                                     tensorType.getElementType(), loc);
  rewriter.setInsertionPointToStart(body);
  SmallVector<Value> bbArgReplacements;
  for (int64_t i = 0; i < tensorType.getRank(); ++i)
    bbArgReplacements.push_back(linalg::IndexOp::create(rewriter, loc, i));
  rewriter.mergeBlocks(&generateOp.getBody().front(), body, bbArgReplacements);
 
  // Update terminator.
  auto yieldOp = cast<tensor::YieldOp>(body->getTerminator());
  rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOp.getValue());
 
  // Replace tensor.generate.
  rewriter.replaceOp(generateOp, genericOp->getResult(0));
  return genericOp.getOperation();
}
 
/// Lower tensor.pad to linalg.generic + tensor.insert_slice.
FailureOr<Operation *>
mlir::linalg::rewriteInDestinationPassingStyle(RewriterBase &rewriter,
                                               tensor::PadOp padOp) {
  // Only ops with exactly one block are supported.
  if (!padOp.getBodyRegion().hasOneBlock())
    return failure();
 
  // Create tensor.empty.
  Location loc = padOp.getLoc();
  RankedTensorType resultType = padOp.getResultType();
  ReifiedRankedShapedTypeDims reifiedShape;
  if (failed(reifyResultShapes(rewriter, padOp, reifiedShape)))
    return rewriter.notifyMatchFailure(
        padOp, "failed to reify tensor.pad op result shape");
  SmallVector<Value> dynamicSizes;
  for (int64_t i = 0; i < resultType.getRank(); ++i)
    if (resultType.isDynamicDim(i))
      dynamicSizes.push_back(cast<Value>(reifiedShape[0][i]));
 
  // If the `padOp` has a nofold attribute and all paddings are known to be 0,
  // explicitly insert a `linalg.copy`.
  if (padOp.getNofoldAttr() &&
      llvm::all_of(padOp.getMixedLowPad(), isZeroInteger) &&
      llvm::all_of(padOp.getMixedHighPad(), isZeroInteger)) {
    using bufferization::AllocTensorOp;
    Value allocated =
        AllocTensorOp::create(rewriter, loc, resultType, dynamicSizes);
    auto copyOp = rewriter.replaceOpWithNewOp<linalg::CopyOp>(
        padOp, padOp.getSource(), allocated);
    return copyOp.getOperation();
  }
 
  Value empty = EmptyOp::create(rewriter, loc, resultType, dynamicSizes);
  // Create linalg.fill or linalg.generic.
  Operation *fillOp = movePaddingToFillOrGenericOp(rewriter, loc, padOp, empty);
  rewriter.setInsertionPointAfter(fillOp);
 
  // Create tensor::InsertSliceOp.
  SmallVector<OpFoldResult> sliceSizes =
      getMixedSizes(rewriter, loc, padOp.getSource());
  SmallVector<OpFoldResult> sliceStrides(resultType.getRank(),
                                         rewriter.getIndexAttr(1));
  auto insertSliceOp = rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
      padOp, padOp.getSource(), fillOp->getResult(0),
      /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
  return insertSliceOp.getOperation();
}
 
Value linalg::bufferizeToAllocation(
    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
    Operation *op, Attribute memorySpace, Operation *insertionPoint) {
  using namespace bufferization;
 
  // Call specialized overload for certain ops.
  if (auto padOp = dyn_cast<tensor::PadOp>(op))
    return bufferizeToAllocation(rewriter, options, padOp, memorySpace);
  if (auto maskOp = dyn_cast<vector::MaskOp>(op))
    return bufferizeToAllocation(rewriter, options, maskOp, memorySpace);
  if (auto allocTensorOp = dyn_cast<bufferization::AllocTensorOp>(op))
    return bufferizeToAllocation(rewriter, options, allocTensorOp, memorySpace);
 
  // Only bufferizable ops are supported.
  auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
  if (!bufferizableOp)
    return nullptr;
 
  // Should the bufferization options and states be function arguments?
  BufferizationOptions bufferizationOptions;
  AnalysisState analysisState(bufferizationOptions);
  BufferizationState bufferizationState;
 
#ifndef NDEBUG
  if (!options.bufferizeDestinationOnly) {
    // Ops with nested tensor ops are not supported yet. At the moment, this
    // function just bufferizes the given op itself, but not its body.
    op->walk([&](Operation *nestedOp) {
      if (op == nestedOp)
        return;
      if (llvm::any_of(nestedOp->getOperands(),
                       [](Value v) { return isa<TensorType>(v.getType()); }))
        llvm_unreachable("ops with nested tensor ops are not supported yet");
      if (llvm::any_of(nestedOp->getResults(),
                       [](Value v) { return isa<TensorType>(v.getType()); }))
        llvm_unreachable("ops with nested tensor ops are not supported yet");
    });
  }
#endif // NDEBUG
 
  // Gather tensor results.
  SmallVector<OpResult> tensorResults;
  for (OpResult result : op->getResults()) {
    if (!isa<TensorType>(result.getType()))
      continue;
    // Unranked tensors are not supported
    if (!isa<RankedTensorType>(result.getType()))
      return nullptr;
    // Ops that bufferize to an allocation are not supported.
    if (bufferizableOp.bufferizesToAllocation(result))
      return nullptr;
    tensorResults.push_back(result);
  }
 
  // Gather all operands that should bufferize to a new allocation. I.e.,
  // bufferize out-of-place.
  SmallVector<OpOperand *> outOfPlaceOperands, resultUses;
  auto addOutOfPlaceOperand = [&](OpOperand *operand) {
    if (!llvm::is_contained(outOfPlaceOperands, operand))
      outOfPlaceOperands.push_back(operand);
  };
  for (OpResult result : tensorResults) {
    AliasingOpOperandList aliasingOperands =
        analysisState.getAliasingOpOperands(result);
    for (const AliasingOpOperand &operand : aliasingOperands) {
      addOutOfPlaceOperand(operand.opOperand);
      for (OpOperand &resultUse : result.getUses())
        resultUses.push_back(&resultUse);
    }
  }
  for (OpOperand &operand : op->getOpOperands()) {
    if (!analysisState.bufferizesToMemoryWrite(operand))
      continue;
    if (!isa<RankedTensorType>(operand.get().getType()))
      continue;
    addOutOfPlaceOperand(&operand);
  }
  // TODO: Support multiple buffers.
  if (outOfPlaceOperands.size() != 1)
    return nullptr;
 
  // Allocate buffers.
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(insertionPoint ? insertionPoint : op);
  SmallVector<Value> allocs;
  for (OpOperand *operand : outOfPlaceOperands) {
    Value alloc = createAllocationForTensor(
        rewriter, op->getLoc(), operand->get(), options, memorySpace);
    allocs.push_back(alloc);
    if (!analysisState.findDefinitions(operand).empty()) {
      // Initialize buffer with a copy of the operand data. Not needed if the
      // tensor is uninitialized.
      createMemcpy(rewriter, op->getLoc(), operand->get(), alloc, options);
    }
    rewriter.modifyOpInPlace(op, [&]() {
      auto toTensorOp = ToTensorOp::create(rewriter, op->getLoc(),
                                           operand->get().getType(), alloc);
      operand->set(toTensorOp);
      if (options.bufferizeDestinationOnly) {
        rewriter.modifyOpInPlace(toTensorOp, [&]() {
          toTensorOp.setRestrict(true);
          toTensorOp.setWritable(true);
        });
      }
    });
  }
 
  if (options.bufferizeDestinationOnly)
    return allocs.front();
 
  // Bufferize the op.
  rewriter.setInsertionPoint(op);
  if (failed(bufferizableOp.bufferize(rewriter, bufferizationOptions,
                                      bufferizationState)))
    return nullptr;
 
  // Set "restrict" attribute, indicating that no other tensor aliases with
  // this tensor. That is because we just allocated a new buffer for the tensor.
  for (OpOperand *resultUse : resultUses) {
    auto toTensorOp = resultUse->get().getDefiningOp<ToTensorOp>();
    assert(toTensorOp && "expected to_tensor op");
    rewriter.modifyOpInPlace(toTensorOp, [&]() {
      toTensorOp.setRestrict(true);
      toTensorOp.setWritable(true);
    });
  }
  return allocs.front();
}
 
namespace {
 
template <typename OpTy>
LogicalResult rewriteOpInDestinationPassingStyle(OpTy op,
                                                 PatternRewriter &rewriter) {
  return linalg::rewriteInDestinationPassingStyle(rewriter, op);
}
 
} // namespace
 
void linalg::populateConvertToDestinationStylePatterns(
    RewritePatternSet &patterns) {
  patterns.add(rewriteOpInDestinationPassingStyle<tensor::FromElementsOp>);
  patterns.add(rewriteOpInDestinationPassingStyle<tensor::GenerateOp>);
  patterns.add(rewriteOpInDestinationPassingStyle<tensor::PadOp>);
}