MLIR  18.0.0git
BufferizableOpInterfaceImpl.cpp
Go to the documentation of this file.
1 //===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
15 #include "mlir/IR/Dialect.h"
16 #include "mlir/IR/Operation.h"
18 
19 using namespace mlir;
20 using namespace linalg;
21 using namespace mlir::bufferization;
22 
23 namespace {
24 
25 /// Generic conversion for any DestinationStyleOpInterface on tensors.
26 static LogicalResult
27 bufferizeDestinationStyleOpInterface(RewriterBase &rewriter,
28  DestinationStyleOpInterface op,
30  // Take a guard before anything else.
31  OpBuilder::InsertionGuard g(rewriter);
32  rewriter.setInsertionPoint(op);
33 
34  // Nothing to do. This op is already bufferized.
35  if (op.hasBufferSemantics())
36  return success();
37 
38  // Ensure op has only tensors. Allow mixed tensor-buffer mode on a per-need
39  // basis.
40  if (!op.hasTensorSemantics())
41  return op->emitError() << "op does not have tensor semantics";
42 
43  // New input operands for the cloned op.
44  SmallVector<Value> newInputBuffers;
45  newInputBuffers.reserve(op.getNumDpsInputs());
46  for (OpOperand *opOperand : op.getDpsInputOperands()) {
47  if (op.isScalar(opOperand)) {
48  newInputBuffers.push_back(opOperand->get());
49  continue;
50  }
51  FailureOr<Value> buffer = getBuffer(rewriter, opOperand->get(), options);
52  if (failed(buffer))
53  return failure();
54  newInputBuffers.push_back(*buffer);
55  }
56 
57  // New output operands for the cloned op.
58  SmallVector<Value> newOutputBuffers;
59  for (OpResult opResult : op->getOpResults()) {
60  OpOperand *opOperand = op.getDpsInitOperand(opResult.getResultNumber());
61  FailureOr<Value> resultBuffer =
62  getBuffer(rewriter, opOperand->get(), options);
63  if (failed(resultBuffer))
64  return failure();
65  newOutputBuffers.push_back(*resultBuffer);
66  }
67 
68  // Merge input/output operands.
69  SmallVector<Value> newOperands = newInputBuffers;
70  newOperands.append(newOutputBuffers.begin(), newOutputBuffers.end());
71 
72  // Set insertion point now that potential alloc/dealloc are introduced.
73  rewriter.setInsertionPoint(op);
74  // Clone the op, but use the new operands. Move the existing block into the
75  // new op. Since the new op does not have any tensor results, it does not
76  // return anything.
77  assert(op->getNumRegions() == 1 && "expected that op has 1 region");
78  auto newOp = cast<DestinationStyleOpInterface>(cloneWithoutRegions(
79  rewriter, op, /*newResultTypes=*/TypeRange{}, newOperands));
80  rewriter.inlineRegionBefore(op->getRegion(0), newOp->getRegion(0),
81  newOp->getRegion(0).begin());
82 
83  // Replace the results of the old op with the new output buffers.
84  replaceOpWithBufferizedValues(rewriter, op, newOutputBuffers);
85 
86  return success();
87 }
88 
89 /// Bufferization of linalg.generic. Replace with a new linalg.generic that
90 /// operates entirely on memrefs.
91 template <typename OpTy>
92 struct LinalgOpInterface
93  : public DstBufferizableOpInterfaceExternalModel<LinalgOpInterface<OpTy>,
94  OpTy> {
95  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
96  const AnalysisState &state) const {
97  // Operand is read if it is used in the computation.
98  auto linalgOp = cast<linalg::LinalgOp>(op);
99  return linalgOp.payloadUsesValueFromOperand(&opOperand);
100  }
101 
102  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
103  const AnalysisState &state) const {
104  // Operand is written to if it is not an input/init.
105  auto dpsOp = cast<DestinationStyleOpInterface>(op);
106  return dpsOp.isDpsInit(&opOperand);
107  }
108 
109  bool bufferizesToElementwiseAccess(Operation *op, const AnalysisState &state,
110  ArrayRef<OpOperand *> opOperands) const {
111  auto linalgOp = cast<linalg::LinalgOp>(op);
112 
113  // All loops must be parallel.
114  if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops())
115  return false;
116 
117  // All index maps of tensors must be identity maps.
118  SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
119  assert(linalgOp->getNumOperands() == indexingMaps.size() &&
120  "unexpected number of indexing maps");
121  for (auto [operand, map] :
122  llvm::zip(linalgOp->getOpOperands(), indexingMaps)) {
123  // Non-tensors do not participate in bufferization, so they can be
124  // ignored.
125  if (!isa<RankedTensorType, MemRefType>(operand.get().getType()))
126  continue;
127  // Only consider operands in `opOperands`.
128  if (llvm::find(opOperands, &operand) == opOperands.end())
129  continue;
130  // TODO: This could be generalized to other indexing maps. (All indexing
131  // must be the same.)
132  if (!map.isIdentity())
133  return false;
134  }
135 
136  return true;
137  }
138 
139  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
140  const BufferizationOptions &options) const {
141  return bufferizeDestinationStyleOpInterface(
142  rewriter, cast<DestinationStyleOpInterface>(op), options);
143  }
144 };
145 
146 /// Helper structure that iterates over all LinalgOps in `OpTys` and registers
147 /// the `BufferizableOpInterface` with each of them.
148 template <typename... Ops>
149 struct LinalgOpInterfaceHelper {
150  static void registerOpInterface(MLIRContext *ctx) {
151  (Ops::template attachInterface<LinalgOpInterface<Ops>>(*ctx), ...);
152  }
153 };
154 } // namespace
155 
157  DialectRegistry &registry) {
158  registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
159  // Register all Linalg structured ops. `LinalgOp` is an interface and it is
160  // not possible to attach an external interface to an existing interface.
161  // Therefore, attach the `BufferizableOpInterface` to all ops one-by-one.
162  LinalgOpInterfaceHelper<
163 #define GET_OP_LIST
164 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
165  >::registerOpInterface(ctx);
166  });
167 }
static llvm::ManagedStatic< PassManagerOptions > options
Base class for generic analysis states.
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
void addExtension(std::unique_ptr< DialectExtensionBase > extension)
Add the given extension to the registry.
This class provides support for representing a failure result, or a valid value of type T.
Definition: LogicalResult.h:78
IRValueT get() const
Return the current value being used by this operand.
Definition: UseDefLists.h:160
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:333
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:383
This class represents an operand of an operation.
Definition: Value.h:263
This is a value defined by a result of an operation.
Definition: Value.h:453
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
unsigned getNumRegions()
Returns the number of regions held by this operation.
Definition: Operation.h:652
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:267
Region & getRegion(unsigned index)
Returns the region held by this operation at position 'index'.
Definition: Operation.h:665
result_range getOpResults()
Definition: Operation.h:415
iterator begin()
Definition: Region.h:55
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:399
virtual void inlineRegionBefore(Region &region, Region &parent, Region::iterator before)
Move the blocks that belong to "region" before the given position in another region "parent".
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:36
void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op, ValueRange values)
Replace an op with replacement values.
FailureOr< Value > getBuffer(RewriterBase &rewriter, Value value, const BufferizationOptions &options)
Lookup the buffer for the given value.
void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry)
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
Operation * cloneWithoutRegions(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:72
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
Options for BufferizableOpInterface-based bufferization.
Bufferizable ops that implement the DestinationStyleOpInterface can use this external model base clas...