MLIR  21.0.0git
DropEquivalentBufferResults.cpp
Go to the documentation of this file.
1 //===- DropEquivalentBufferResults.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 //
9 // This pass drops return values from functions if they are equivalent to one of
10 // their arguments. E.g.:
11 //
12 // ```
13 // func.func @foo(%m : memref<?xf32>) -> (memref<?xf32>) {
14 // return %m : memref<?xf32>
15 // }
16 // ```
17 //
18 // This functions is rewritten to:
19 //
20 // ```
21 // func.func @foo(%m : memref<?xf32>) {
22 // return
23 // }
24 // ```
25 //
26 // All call sites are updated accordingly. If a function returns a cast of a
27 // function argument, it is also considered equivalent. A cast is inserted at
28 // the call site in that case.
29 
31 
34 #include "mlir/IR/Operation.h"
35 #include "mlir/Pass/Pass.h"
36 
37 namespace mlir {
38 namespace bufferization {
39 #define GEN_PASS_DEF_DROPEQUIVALENTBUFFERRESULTSPASS
40 #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
41 } // namespace bufferization
42 } // namespace mlir
43 
44 using namespace mlir;
45 
46 /// Return the unique ReturnOp that terminates `funcOp`.
47 /// Return nullptr if there is no such unique ReturnOp.
48 static func::ReturnOp getAssumedUniqueReturnOp(func::FuncOp funcOp) {
49  func::ReturnOp returnOp;
50  for (Block &b : funcOp.getBody()) {
51  if (auto candidateOp = dyn_cast<func::ReturnOp>(b.getTerminator())) {
52  if (returnOp)
53  return nullptr;
54  returnOp = candidateOp;
55  }
56  }
57  return returnOp;
58 }
59 
60 /// Return the func::FuncOp called by `callOp`.
61 static func::FuncOp getCalledFunction(CallOpInterface callOp) {
62  SymbolRefAttr sym =
63  llvm::dyn_cast_if_present<SymbolRefAttr>(callOp.getCallableForCallee());
64  if (!sym)
65  return nullptr;
66  return dyn_cast_or_null<func::FuncOp>(
68 }
69 
70 LogicalResult
72  IRRewriter rewriter(module.getContext());
73 
75  // Collect the mapping of functions to their call sites.
76  module.walk([&](func::CallOp callOp) {
77  if (func::FuncOp calledFunc = getCalledFunction(callOp)) {
78  callerMap[calledFunc].insert(callOp);
79  }
80  });
81 
82  for (auto funcOp : module.getOps<func::FuncOp>()) {
83  if (funcOp.isExternal())
84  continue;
85  func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
86  // TODO: Support functions with multiple blocks.
87  if (!returnOp)
88  continue;
89 
90  // Compute erased results.
91  SmallVector<Value> newReturnValues;
92  BitVector erasedResultIndices(funcOp.getFunctionType().getNumResults());
93  DenseMap<int64_t, int64_t> resultToArgs;
94  for (const auto &it : llvm::enumerate(returnOp.getOperands())) {
95  bool erased = false;
96  for (BlockArgument bbArg : funcOp.getArguments()) {
97  Value val = it.value();
98  while (auto castOp = val.getDefiningOp<memref::CastOp>())
99  val = castOp.getSource();
100 
101  if (val == bbArg) {
102  resultToArgs[it.index()] = bbArg.getArgNumber();
103  erased = true;
104  break;
105  }
106  }
107 
108  if (erased) {
109  erasedResultIndices.set(it.index());
110  } else {
111  newReturnValues.push_back(it.value());
112  }
113  }
114 
115  // Update function.
116  if (failed(funcOp.eraseResults(erasedResultIndices)))
117  return failure();
118  returnOp.getOperandsMutable().assign(newReturnValues);
119 
120  // Update function calls.
121  for (func::CallOp callOp : callerMap[funcOp]) {
122  rewriter.setInsertionPoint(callOp);
123  auto newCallOp = rewriter.create<func::CallOp>(callOp.getLoc(), funcOp,
124  callOp.getOperands());
125  SmallVector<Value> newResults;
126  int64_t nextResult = 0;
127  for (int64_t i = 0; i < callOp.getNumResults(); ++i) {
128  if (!resultToArgs.count(i)) {
129  // This result was not erased.
130  newResults.push_back(newCallOp.getResult(nextResult++));
131  continue;
132  }
133 
134  // This result was erased.
135  Value replacement = callOp.getOperand(resultToArgs[i]);
136  Type expectedType = callOp.getResult(i).getType();
137  if (replacement.getType() != expectedType) {
138  // A cast must be inserted at the call site.
139  replacement = rewriter.create<memref::CastOp>(
140  callOp.getLoc(), expectedType, replacement);
141  }
142  newResults.push_back(replacement);
143  }
144  rewriter.replaceOp(callOp, newResults);
145  }
146  }
147 
148  return success();
149 }
150 
151 namespace {
152 struct DropEquivalentBufferResultsPass
153  : bufferization::impl::DropEquivalentBufferResultsPassBase<
154  DropEquivalentBufferResultsPass> {
155  void runOnOperation() override {
156  if (failed(bufferization::dropEquivalentBufferResults(getOperation())))
157  return signalPassFailure();
158  }
159 };
160 } // namespace
static func::FuncOp getCalledFunction(CallOpInterface callOp)
Return the func::FuncOp called by callOp.
static func::ReturnOp getAssumedUniqueReturnOp(func::FuncOp funcOp)
Return the unique ReturnOp that terminates funcOp.
This class represents an argument of a Block.
Definition: Value.h:295
Block represents an ordered list of Operations.
Definition: Block.h:33
This class coordinates rewriting a piece of IR outside of a pattern rewrite, providing a way to keep ...
Definition: PatternMatch.h:730
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:395
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:453
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
static Operation * lookupNearestSymbolFrom(Operation *from, StringAttr symbol)
Returns the operation registered with the given symbol name within the closest parent operation of,...
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:105
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
LogicalResult dropEquivalentBufferResults(ModuleOp module)
Drop all memref function results that are equivalent to a function argument.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:344
Include the generated interface declarations.