llvm/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp

//===- 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/Arith/Utils/Utils.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"
#include "llvm/Support/Debug.h"

usingnamespacemlir;
usingnamespacemlir::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) { … }

/// 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) { … }

static Operation *movePaddingToFillOrGenericOp(RewriterBase &rewriter,
                                               Location loc, PadOp padOp,
                                               Value dest) { … }

static SmallVector<Value> reifyOrComputeDynamicSizes(OpBuilder &b,
                                                     Value value) { … }

static Value
createAllocationForTensor(RewriterBase &rewriter, Location loc, Value value,
                          const linalg::BufferizeToAllocationOptions &options,
                          Attribute memorySpace = { … }

Value linalg::bufferizeToAllocation(
    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
    PadOp padOp, Attribute memorySpace, Operation *insertionPoint) { … }

Value linalg::bufferizeToAllocation(
    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
    vector::MaskOp maskOp, Attribute memorySpace, Operation *insertionPoint) { … }

Value linalg::bufferizeToAllocation(
    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
    bufferization::AllocTensorOp allocTensorOp, Attribute memorySpace,
    Operation *insertionPoint) { … }

/// Lower tensor.from_elements to a sequence of chained tensor.insert.
FailureOr<Operation *> mlir::linalg::rewriteInDestinationPassingStyle(
    RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp) { … }

/// Lower tensor.generate to linalg.generic.
FailureOr<Operation *>
mlir::linalg::rewriteInDestinationPassingStyle(RewriterBase &rewriter,
                                               tensor::GenerateOp generateOp) { … }

/// Lower tensor.pad to linalg.generic + tensor.insert_slice.
FailureOr<Operation *>
mlir::linalg::rewriteInDestinationPassingStyle(RewriterBase &rewriter,
                                               tensor::PadOp padOp) { … }

Value linalg::bufferizeToAllocation(
    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
    Operation *op, Attribute memorySpace, Operation *insertionPoint) { … }

namespace {

template <typename OpTy>
LogicalResult rewriteOpInDestinationPassingStyle(OpTy op,
                                                 PatternRewriter &rewriter) { … }

} // namespace

void linalg::populateConvertToDestinationStylePatterns(
    RewritePatternSet &patterns) { … }