MLIR  19.0.0git
BufferResultsToOutParams.cpp
Go to the documentation of this file.
1 //===- BufferResultsToOutParams.cpp - Calling convention conversion -------===//
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 
10 
13 #include "mlir/IR/Operation.h"
14 #include "mlir/Pass/Pass.h"
15 
16 namespace mlir {
17 namespace bufferization {
18 #define GEN_PASS_DEF_BUFFERRESULTSTOOUTPARAMS
19 #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
20 } // namespace bufferization
21 } // namespace mlir
22 
23 using namespace mlir;
25 
26 /// Return `true` if the given MemRef type has a fully dynamic layout.
27 static bool hasFullyDynamicLayoutMap(MemRefType type) {
28  int64_t offset;
30  if (failed(getStridesAndOffset(type, strides, offset)))
31  return false;
32  if (!llvm::all_of(strides, ShapedType::isDynamic))
33  return false;
34  if (!ShapedType::isDynamic(offset))
35  return false;
36  return true;
37 }
38 
39 /// Return `true` if the given MemRef type has a static identity layout (i.e.,
40 /// no layout).
41 static bool hasStaticIdentityLayout(MemRefType type) {
42  return type.getLayout().isIdentity();
43 }
44 
45 // Updates the func op and entry block.
46 //
47 // Any args appended to the entry block are added to `appendedEntryArgs`.
48 // If `addResultAttribute` is true, adds the unit attribute `bufferize.result`
49 // to each newly created function argument.
50 static LogicalResult
51 updateFuncOp(func::FuncOp func,
52  SmallVectorImpl<BlockArgument> &appendedEntryArgs,
53  bool addResultAttribute) {
54  auto functionType = func.getFunctionType();
55 
56  // Collect information about the results will become appended arguments.
57  SmallVector<Type, 6> erasedResultTypes;
58  BitVector erasedResultIndices(functionType.getNumResults());
59  for (const auto &resultType : llvm::enumerate(functionType.getResults())) {
60  if (auto memrefType = dyn_cast<MemRefType>(resultType.value())) {
61  if (!hasStaticIdentityLayout(memrefType) &&
62  !hasFullyDynamicLayoutMap(memrefType)) {
63  // Only buffers with static identity layout can be allocated. These can
64  // be casted to memrefs with fully dynamic layout map. Other layout maps
65  // are not supported.
66  return func->emitError()
67  << "cannot create out param for result with unsupported layout";
68  }
69  erasedResultIndices.set(resultType.index());
70  erasedResultTypes.push_back(memrefType);
71  }
72  }
73 
74  // Add the new arguments to the function type.
75  auto newArgTypes = llvm::to_vector<6>(
76  llvm::concat<const Type>(functionType.getInputs(), erasedResultTypes));
77  auto newFunctionType = FunctionType::get(func.getContext(), newArgTypes,
78  functionType.getResults());
79  func.setType(newFunctionType);
80 
81  // Transfer the result attributes to arg attributes.
82  auto erasedIndicesIt = erasedResultIndices.set_bits_begin();
83  for (int i = 0, e = erasedResultTypes.size(); i < e; ++i, ++erasedIndicesIt) {
84  func.setArgAttrs(functionType.getNumInputs() + i,
85  func.getResultAttrs(*erasedIndicesIt));
86  if (addResultAttribute)
87  func.setArgAttr(functionType.getNumInputs() + i,
88  StringAttr::get(func.getContext(), "bufferize.result"),
89  UnitAttr::get(func.getContext()));
90  }
91 
92  // Erase the results.
93  func.eraseResults(erasedResultIndices);
94 
95  // Add the new arguments to the entry block if the function is not external.
96  if (func.isExternal())
97  return success();
98  Location loc = func.getLoc();
99  for (Type type : erasedResultTypes)
100  appendedEntryArgs.push_back(func.front().addArgument(type, loc));
101 
102  return success();
103 }
104 
105 // Updates all ReturnOps in the scope of the given func::FuncOp by either
106 // keeping them as return values or copying the associated buffer contents into
107 // the given out-params.
108 static LogicalResult updateReturnOps(func::FuncOp func,
109  ArrayRef<BlockArgument> appendedEntryArgs,
110  MemCpyFn memCpyFn) {
111  auto res = func.walk([&](func::ReturnOp op) {
112  SmallVector<Value, 6> copyIntoOutParams;
113  SmallVector<Value, 6> keepAsReturnOperands;
114  for (Value operand : op.getOperands()) {
115  if (isa<MemRefType>(operand.getType()))
116  copyIntoOutParams.push_back(operand);
117  else
118  keepAsReturnOperands.push_back(operand);
119  }
120  OpBuilder builder(op);
121  for (auto t : llvm::zip(copyIntoOutParams, appendedEntryArgs)) {
122  if (failed(
123  memCpyFn(builder, op.getLoc(), std::get<0>(t), std::get<1>(t))))
124  return WalkResult::interrupt();
125  }
126  builder.create<func::ReturnOp>(op.getLoc(), keepAsReturnOperands);
127  op.erase();
128  return WalkResult::advance();
129  });
130  return failure(res.wasInterrupted());
131 }
132 
133 // Updates all CallOps in the scope of the given ModuleOp by allocating
134 // temporary buffers for newly introduced out params.
135 static LogicalResult
136 updateCalls(ModuleOp module,
138  bool didFail = false;
139  SymbolTable symtab(module);
140  module.walk([&](func::CallOp op) {
141  auto callee = symtab.lookup<func::FuncOp>(op.getCallee());
142  if (!callee) {
143  op.emitError() << "cannot find callee '" << op.getCallee() << "' in "
144  << "symbol table";
145  didFail = true;
146  return;
147  }
148  if (!options.filterFn(&callee))
149  return;
150  SmallVector<Value, 6> replaceWithNewCallResults;
151  SmallVector<Value, 6> replaceWithOutParams;
152  for (OpResult result : op.getResults()) {
153  if (isa<MemRefType>(result.getType()))
154  replaceWithOutParams.push_back(result);
155  else
156  replaceWithNewCallResults.push_back(result);
157  }
158  SmallVector<Value, 6> outParams;
159  OpBuilder builder(op);
160  for (Value memref : replaceWithOutParams) {
161  if (!cast<MemRefType>(memref.getType()).hasStaticShape()) {
162  op.emitError()
163  << "cannot create out param for dynamically shaped result";
164  didFail = true;
165  return;
166  }
167  auto memrefType = cast<MemRefType>(memref.getType());
168  auto allocType =
169  MemRefType::get(memrefType.getShape(), memrefType.getElementType(),
170  AffineMap(), memrefType.getMemorySpace());
171  Value outParam = builder.create<memref::AllocOp>(op.getLoc(), allocType);
172  if (!hasStaticIdentityLayout(memrefType)) {
173  // Layout maps are already checked in `updateFuncOp`.
174  assert(hasFullyDynamicLayoutMap(memrefType) &&
175  "layout map not supported");
176  outParam =
177  builder.create<memref::CastOp>(op.getLoc(), memrefType, outParam);
178  }
179  memref.replaceAllUsesWith(outParam);
180  outParams.push_back(outParam);
181  }
182 
183  auto newOperands = llvm::to_vector<6>(op.getOperands());
184  newOperands.append(outParams.begin(), outParams.end());
185  auto newResultTypes = llvm::to_vector<6>(llvm::map_range(
186  replaceWithNewCallResults, [](Value v) { return v.getType(); }));
187  auto newCall = builder.create<func::CallOp>(op.getLoc(), op.getCalleeAttr(),
188  newResultTypes, newOperands);
189  for (auto t : llvm::zip(replaceWithNewCallResults, newCall.getResults()))
190  std::get<0>(t).replaceAllUsesWith(std::get<1>(t));
191  op.erase();
192  });
193 
194  return failure(didFail);
195 }
196 
198  ModuleOp module,
200  for (auto func : module.getOps<func::FuncOp>()) {
201  if (!options.filterFn(&func))
202  continue;
203  SmallVector<BlockArgument, 6> appendedEntryArgs;
204  if (failed(
205  updateFuncOp(func, appendedEntryArgs, options.addResultAttribute)))
206  return failure();
207  if (func.isExternal())
208  continue;
209  auto defaultMemCpyFn = [](OpBuilder &builder, Location loc, Value from,
210  Value to) {
211  builder.create<memref::CopyOp>(loc, from, to);
212  return success();
213  };
214  if (failed(updateReturnOps(func, appendedEntryArgs,
215  options.memCpyFn.value_or(defaultMemCpyFn)))) {
216  return failure();
217  }
218  }
219  if (failed(updateCalls(module, options)))
220  return failure();
221  return success();
222 }
223 
224 namespace {
225 struct BufferResultsToOutParamsPass
226  : bufferization::impl::BufferResultsToOutParamsBase<
227  BufferResultsToOutParamsPass> {
228  explicit BufferResultsToOutParamsPass(
230  : options(options) {}
231 
232  void runOnOperation() override {
233  // Convert from pass options in tablegen to BufferResultsToOutParamsOpts.
234  if (addResultAttribute)
235  options.addResultAttribute = true;
236 
238  options)))
239  return signalPassFailure();
240  }
241 
242 private:
244 };
245 } // namespace
246 
249  return std::make_unique<BufferResultsToOutParamsPass>(options);
250 }
static LogicalResult updateReturnOps(func::FuncOp func, ArrayRef< BlockArgument > appendedEntryArgs, MemCpyFn memCpyFn)
bufferization::BufferResultsToOutParamsOpts::MemCpyFn MemCpyFn
static LogicalResult updateFuncOp(func::FuncOp func, SmallVectorImpl< BlockArgument > &appendedEntryArgs, bool addResultAttribute)
static bool hasStaticIdentityLayout(MemRefType type)
Return true if the given MemRef type has a static identity layout (i.e., no layout).
static LogicalResult updateCalls(ModuleOp module, const bufferization::BufferResultsToOutParamsOpts &options)
static bool hasFullyDynamicLayoutMap(MemRefType type)
Return true if the given MemRef type has a fully dynamic layout.
static llvm::ManagedStatic< PassManagerOptions > options
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:47
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
This class helps build Operations.
Definition: Builders.h:209
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:464
This is a value defined by a result of an operation.
Definition: Value.h:453
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:268
operand_range getOperands()
Returns an iterator on the underlying Value's.
Definition: Operation.h:373
void replaceAllUsesWith(ValuesT &&values)
Replace all uses of results of this operation with the provided 'values'.
Definition: Operation.h:272
result_range getResults()
Definition: Operation.h:410
void erase()
Remove this operation from its parent block and delete it.
Definition: Operation.cpp:539
This class allows for representing and managing the symbol table used by operations with the 'SymbolT...
Definition: SymbolTable.h:24
Operation * lookup(StringRef name) const
Look up a symbol with the specified name, returning null if no such name exists.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:125
static WalkResult advance()
Definition: Visitors.h:52
static WalkResult interrupt()
Definition: Visitors.h:51
LogicalResult promoteBufferResultsToOutParams(ModuleOp module, const BufferResultsToOutParamsOpts &options)
Replace buffers that are returned from a function with an out parameter.
std::unique_ptr< Pass > createBufferResultsToOutParamsPass(const BufferResultsToOutParamsOpts &options={})
Creates a pass that converts memref function results to out-params.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
LogicalResult getStridesAndOffset(MemRefType t, SmallVectorImpl< int64_t > &strides, int64_t &offset)
Returns the strides of the MemRef if the layout map is in strided form.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
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
std::function< LogicalResult(OpBuilder &, Location, Value, Value)> MemCpyFn
Memcpy function: Generate a memcpy between two memrefs.
Definition: Passes.h:154