[Mlir-commits] [mlir] [mlir][bufferize] Make drop-equivalent-buffer-results support mult blocks (PR #163388)

[Matthias Springer]

================
@@ -41,18 +41,38 @@ namespace bufferization {
 
 using namespace mlir;
 
-/// Return the unique ReturnOp that terminates `funcOp`.
-/// Return nullptr if there is no such unique ReturnOp.
-static func::ReturnOp getAssumedUniqueReturnOp(func::FuncOp funcOp) {
-  func::ReturnOp returnOp;
+/// Get all the ReturnOp in the funcOp.
+static SmallVector<func::ReturnOp> getReturnOps(func::FuncOp funcOp) {
+  SmallVector<func::ReturnOp> returnOps;
   for (Block &b : funcOp.getBody()) {
     if (auto candidateOp = dyn_cast<func::ReturnOp>(b.getTerminator())) {
-      if (returnOp)
-        return nullptr;
-      returnOp = candidateOp;
+      returnOps.push_back(candidateOp);
     }
   }
-  return returnOp;
+  return returnOps;
+}
+
+/// Get the values at the same position in the `returnOps`.
+static SmallVector<Value>
+getReturnOpsOperandInPos(ArrayRef<func::ReturnOp> returnOps, size_t pos) {
+  SmallVector<Value> operands;
+  for (func::ReturnOp returnOp : returnOps) {
+    operands.push_back(returnOp.getOperand(pos));
+  }
+  return operands;
+}
+
+/// Check if the value in operands is equal to the argument.
----------------
matthias-springer wrote:

nit: `Check if all given values are the same buffer as the block argument (modulo cast ops).`

https://github.com/llvm/llvm-project/pull/163388