MLIR  17.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_DROPEQUIVALENTBUFFERRESULTS
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 = callOp.getCallableForCallee().dyn_cast<SymbolRefAttr>();
63  if (!sym)
64  return nullptr;
65  return dyn_cast_or_null<func::FuncOp>(
67 }
68 
71  IRRewriter rewriter(module.getContext());
72 
73  for (auto funcOp : module.getOps<func::FuncOp>()) {
74  if (funcOp.isExternal())
75  continue;
76  func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
77  // TODO: Support functions with multiple blocks.
78  if (!returnOp)
79  continue;
80 
81  // Compute erased results.
82  SmallVector<Value> newReturnValues;
83  BitVector erasedResultIndices(funcOp.getFunctionType().getNumResults());
84  DenseMap<int64_t, int64_t> resultToArgs;
85  for (const auto &it : llvm::enumerate(returnOp.getOperands())) {
86  bool erased = false;
87  for (BlockArgument bbArg : funcOp.getArguments()) {
88  Value val = it.value();
89  while (auto castOp = val.getDefiningOp<memref::CastOp>())
90  val = castOp.getSource();
91 
92  if (val == bbArg) {
93  resultToArgs[it.index()] = bbArg.getArgNumber();
94  erased = true;
95  break;
96  }
97  }
98 
99  if (erased) {
100  erasedResultIndices.set(it.index());
101  } else {
102  newReturnValues.push_back(it.value());
103  }
104  }
105 
106  // Update function.
107  funcOp.eraseResults(erasedResultIndices);
108  returnOp.getOperandsMutable().assign(newReturnValues);
109 
110  // Update function calls.
111  module.walk([&](func::CallOp callOp) {
112  if (getCalledFunction(callOp) != funcOp)
113  return WalkResult::skip();
114 
115  rewriter.setInsertionPoint(callOp);
116  auto newCallOp = rewriter.create<func::CallOp>(callOp.getLoc(), funcOp,
117  callOp.getOperands());
118  SmallVector<Value> newResults;
119  int64_t nextResult = 0;
120  for (int64_t i = 0; i < callOp.getNumResults(); ++i) {
121  if (!resultToArgs.count(i)) {
122  // This result was not erased.
123  newResults.push_back(newCallOp.getResult(nextResult++));
124  continue;
125  }
126 
127  // This result was erased.
128  Value replacement = callOp.getOperand(resultToArgs[i]);
129  Type expectedType = callOp.getResult(i).getType();
130  if (replacement.getType() != expectedType) {
131  // A cast must be inserted at the call site.
132  replacement = rewriter.create<memref::CastOp>(
133  callOp.getLoc(), expectedType, replacement);
134  }
135  newResults.push_back(replacement);
136  }
137  rewriter.replaceOp(callOp, newResults);
138  return WalkResult::advance();
139  });
140  }
141 
142  return success();
143 }
144 
145 namespace {
146 struct DropEquivalentBufferResultsPass
147  : bufferization::impl::DropEquivalentBufferResultsBase<
148  DropEquivalentBufferResultsPass> {
149  void runOnOperation() override {
151  return signalPassFailure();
152  }
153 };
154 } // namespace
155 
156 std::unique_ptr<Pass>
158  return std::make_unique<DropEquivalentBufferResultsPass>();
159 }
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:304
Block represents an ordered list of Operations.
Definition: Block.h:30
This class coordinates rewriting a piece of IR outside of a pattern rewrite, providing a way to keep ...
Definition: PatternMatch.h:605
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:351
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:422
virtual void replaceOp(Operation *op, ValueRange newValues)
This method replaces the results of the operation with the specified list of values.
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:93
Type getType() const
Return the type of this value.
Definition: Value.h:122
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
static WalkResult skip()
Definition: Visitors.h:53
static WalkResult advance()
Definition: Visitors.h:52
std::unique_ptr< Pass > createDropEquivalentBufferResultsPass()
Creates a pass that drops memref function results that are equivalent to a function argument.
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:223
Include the generated interface declarations.
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