[mlir][bufferize][NFC] Decouple dropping of equivalent return values from bufferization

[springerm]

Summary

This simplifies the bufferization itself and is in preparation of connecting with the sparse compiler.

Depends On D126813

[springerm]

fix cmake build

[springerm]

update

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

 
 /// Options for analysis-enabled bufferization.
 struct OneShotBufferizationOptions : public BufferizationOptions {
   OneShotBufferizationOptions() = default;
 
   /// Specifies whether returning newly allocated memrefs should be allowed.
   /// Otherwise, a pass failure is triggered.
   bool allowReturnAllocs = false;
-
-  /// Specifies whether buffer return values that are equivalent to a FuncOp
-  /// bbArg should be dropped.
-  bool dropEquivalentFuncResults = true;
 };
 
 /// The BufferizationAliasInfo class maintains a list of buffer aliases and
 /// equivalence classes to support bufferization.
 class BufferizationAliasInfo {
 public:
   explicit BufferizationAliasInfo(Operation *rootOp);
 

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h

 
 /// Creates a pass that converts memref function results to out-params.
 std::unique_ptr<Pass> createBufferResultsToOutParamsPass();
 
 /// Replace buffers that are returned from a function with an out parameter.
 /// Also update all call sites.
 LogicalResult promoteBufferResultsToOutParams(ModuleOp module);
 
+/// Creates a pass that drops memref function results that are equivalent to a
+/// function argument.
+std::unique_ptr<Pass> createDropEquivalentBufferResultsPass();
+
+/// Drop all memref function results that are equivalent to a function argument.
+LogicalResult dropEquivalentBufferResults(ModuleOp module);
+
 /// Creates a pass that finalizes a partial bufferization by removing remaining
 /// bufferization.to_tensor and bufferization.to_memref operations.
 std::unique_ptr<OperationPass<func::FuncOp>> createFinalizingBufferizePass();
 
 /// Create a pass that bufferizes all ops that implement BufferizableOpInterface
 /// with One-Shot Bufferize.
 std::unique_ptr<Pass> createOneShotBufferizePass();
 

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

def FinalizingBufferize : Pass<"finalizing-bufferize", "func::FuncOp"> {
   let constructor = "mlir::bufferization::createFinalizingBufferizePass()";
 }
 
 def BufferizationBufferize : Pass<"bufferization-bufferize", "func::FuncOp"> {
   let summary = "Bufferize the `bufferization` dialect";
   let constructor = "mlir::bufferization::createBufferizationBufferizePass()";
 }
 
+def DropEquivalentBufferResults : Pass<"drop-equivalent-buffer-results", "ModuleOp">  {
+  let summary = "Remove MemRef return values that are equivalent to a bbArg";
+  let description = [{
+    This pass removes MemRef return values from functions if they are equivalent
+    to a function bbArg. In that case, the return value is redundant and the
+    respective CallOp operand can be used at the call site.
+
+    Note: If a bbArg buffer is not returned directly but casted to beforehand,
+    the buffer is still considered equivalent.
+  }];
+  let constructor = "mlir::bufferization::createDropEquivalentBufferResultsPass()";
+  let dependentDialects = ["memref::MemRefDialect"];
+}
+
 def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
   let summary = "One-Shot Bufferize";
   let description = [{
     This pass bufferizes all ops that implement `BufferizableOpInterface`. It
     first performs an inplacability analysis on SSA use-def chains of tensor
     values to determine which OpOperands may bufferize in-place, i.e., without
     inserting a buffer copy. It then rewrites the IR, inserting a buffer
     allocation and copy for each OpOperand that was decided to bufferize
let description = [{
     bufferize out-of-place. This is useful for understanding why One-Shot
     Bufferize chose to insert a certain buffer copy.
 
     `bufferize-function-boundaries` is an experimental flag for bufferizing
     `FuncOp`, `ReturnOp` and `CallOp`. This feature is still under development
     and supports only simple cases at the moment. In particular:
 
     * Recursive or circular function call graphs are not supported.
-    * When a returned tensor can be proven to be equivalent to a tensor function
-      argument, the return value disappears. Instead, the buffer of the tensor
-      argument is modified in-place.
-    * Returning non-equivalent tensors is forbidden by default and must be
-      explicitly activated with `allow-return-allocs`.
     * External functions (without bodies) that return a tensor are not
       supported.
     * Function with multiple blocks or multiple ReturnOps are not supported.
     * Layout maps on function signatures can be controlled with a separate
       `function-boundary-type-conversion` option, which is similar to
       `unknown-type-conversion` but supports an additional `infer-layout-map`
       option. `fully-dynamic-layout-map` and `identity-layout-map` ensure that
       function signatures bufferize to easily predictable types, potentially at
let description = [{
 
     One-Shot Bufferize gathers additional information during the analysis phase
     when function boundary bufferization is activated. E.g., whether a function
     argument is read/written and which returned values are aliasing/equivalent.
     For debugging purposes, such information can be printed with
     `test-analysis-only`.
   }];
   let options = [
-    Option<"dropEquivalentFuncResults", "drop-equivalent-func-results", "bool",
-           /*default=*/"true",
-           "Drop buffer return values that are equivalent to a FuncOp arg.">,
     Option<"allowReturnAllocs", "allow-return-allocs", "bool",
             /*default=*/"false",
            "Allows returning/yielding new allocations from a block.">,
     Option<"allowUnknownOps", "allow-unknown-ops", "bool",
            /*default=*/"false",
            "Allows unknown (not bufferizable) ops in the input IR.">,
     Option<"alwaysAliasingWithDest", "always-aliasing-with-dest", "bool",
             /*default=*/"true",

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

void getDependentDialects(DialectRegistry &registry) const override {
     registerAllocationOpInterfaceExternalModels(registry);
   }
 
   void runOnOperation() override {
     OneShotBufferizationOptions opt;
     if (!options) {
       // Make new bufferization options if none were provided when creating the
       // pass.
-      opt.dropEquivalentFuncResults = dropEquivalentFuncResults;
       opt.allowReturnAllocs = allowReturnAllocs;
       opt.allowUnknownOps = allowUnknownOps;
       opt.alwaysAliasingWithDest = alwaysAliasingWithDest;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.createDeallocs = createDeallocs;
       opt.functionBoundaryTypeConversion =
           parseLayoutMapOption(functionBoundaryTypeConversion);
       opt.printConflicts = printConflicts;

mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRBufferizationTransforms
   AllocTensorElimination.cpp
   Bufferize.cpp
   BufferDeallocation.cpp
   BufferOptimizations.cpp
   BufferResultsToOutParams.cpp
   BufferUtils.cpp
+  DropEquivalentBufferResults.cpp
   FuncBufferizableOpInterfaceImpl.cpp
   OneShotAnalysis.cpp
   OneShotModuleBufferize.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
 
   DEPENDS
   MLIRBufferizationPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRBufferization
   MLIRControlFlowInterfaces
   MLIRFunc
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRMemRef
   MLIRPass
   MLIRTensor
   MLIRTransforms
 )
+

mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp

+//===- DropEquivalentBufferResults.cpp - Calling convention conversion ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass drops return values from functions if they are equivalent to one of
+// their arguments. E.g.:
+//
+// ```
+// func.func @foo(%m : memref<?xf32>) -> (memref<?xf32>) {
+//   return %m : memref<?xf32>
+// }
+// ```
+//
+// This functions is rewritten to:
+//
+// ```
+// func.func @foo(%m : memref<?xf32>) {
+//   return
+// }
+// ```
+//
+// All call sites are updated accordingly. If a function returns a cast of a
+// function argument, it is also considered equivalent. A cast is inserted at
+// the call site in that case.
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/Pass.h"
+
+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;
+  for (Block &b : funcOp.getBody()) {
+    if (auto candidateOp = dyn_cast<func::ReturnOp>(b.getTerminator())) {
+      if (returnOp)
+        return nullptr;
+      returnOp = candidateOp;
+    }
+  }
+  return returnOp;
+}
+
+/// Return the func::FuncOp called by `callOp`.
+static func::FuncOp getCalledFunction(CallOpInterface callOp) {
+  SymbolRefAttr sym = callOp.getCallableForCallee().dyn_cast<SymbolRefAttr>();
+  if (!sym)
+    return nullptr;
+  return dyn_cast_or_null<func::FuncOp>(
+      SymbolTable::lookupNearestSymbolFrom(callOp, sym));
+}
+
+LogicalResult
+mlir::bufferization::dropEquivalentBufferResults(ModuleOp module) {
+  IRRewriter rewriter(module.getContext());
+
+  for (auto funcOp : module.getOps<func::FuncOp>()) {
+    if (funcOp.isExternal())
+      continue;
+    func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
+    // TODO: Support functions with multiple blocks.
+    if (!returnOp)
+      continue;
+
+    // Compute erased results.
+    SmallVector<Value> newReturnValues;
+    BitVector erasedResultIndices(funcOp.getFunctionType().getNumResults());
+    DenseMap<int64_t, int64_t> resultToArgs;
+    for (const auto &it : llvm::enumerate(returnOp.operands())) {
+      bool erased = false;
+      for (BlockArgument bbArg : funcOp.getArguments()) {
+        Value val = it.value();
+        while (auto castOp = val.getDefiningOp<memref::CastOp>())
+          val = castOp.source();
+
+        if (val == bbArg) {
+          resultToArgs[it.index()] = bbArg.getArgNumber();
+          erased = true;
+          break;
+        }
+      }
+
+      if (erased) {
+        erasedResultIndices.set(it.index());
+      } else {
+        newReturnValues.push_back(it.value());
+      }
+    }
+
+    // Update function.
+    funcOp.eraseResults(erasedResultIndices);
+    returnOp.operandsMutable().assign(newReturnValues);
+
+    // Update function calls.
+    module.walk([&](func::CallOp callOp) {
+      if (getCalledFunction(callOp) != funcOp)
+        return WalkResult::skip();
+
+      rewriter.setInsertionPoint(callOp);
+      auto newCallOp = rewriter.create<func::CallOp>(callOp.getLoc(), funcOp,
+                                                     callOp.operands());
+      SmallVector<Value> newResults;
+      int64_t nextResult = 0;
+      for (int64_t i = 0; i < callOp.getNumResults(); ++i) {
+        if (!resultToArgs.count(i)) {
+          // This result was not erased.
+          newResults.push_back(newCallOp.getResult(nextResult++));
+          continue;
+        }
+
+        // This result was erased.
+        Value replacement = callOp.getOperand(resultToArgs[i]);
+        Type expectedType = callOp.getResult(i).getType();
+        if (replacement.getType() != expectedType) {
+          // A cast must be inserted at the call site.
+          replacement = rewriter.create<memref::CastOp>(
+              callOp.getLoc(), expectedType, replacement);
+        }
+        newResults.push_back(replacement);
+      }
+      rewriter.replaceOp(callOp, newResults);
+      return WalkResult::advance();
+    });
+  }
+
+  return success();
+}
+
+namespace {
+struct DropEquivalentBufferResultsPass
+    : DropEquivalentBufferResultsBase<DropEquivalentBufferResultsPass> {
+  void runOnOperation() override {
+    if (failed(bufferization::dropEquivalentBufferResults(getOperation())))
+      return signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass>
+mlir::bufferization::createDropEquivalentBufferResultsPass() {
+  return std::make_unique<DropEquivalentBufferResultsPass>();
+}

mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp

getAliasingOpOperand(Operation *op, OpResult opResult,
     SmallVector<OpOperand *> result;
     for (int64_t bbArgIdx : aliasingFuncArgs)
       result.push_back(&callOp->getOpOperand(bbArgIdx));
     return result;
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
+    func::CallOp callOp = cast<func::CallOp>(op);
+    FuncOp funcOp = getCalledFunction(callOp);
+    assert(funcOp && "expected CallOp to a FuncOp");
+    const FuncAnalysisState &funcState = getFuncAnalysisState(state);
+    if (getFuncOpAnalysisState(state, funcOp) !=
+        FuncOpAnalysisState::Analyzed) {
+      // Function not analyzed yet. The conservative answer is "None".
+      return BufferRelation::None;
+    }
+
+    Optional<int64_t> maybeEquiv =
+        getEquivalentFuncArgIdx(funcOp, funcState, opResult.getResultNumber());
+    if (maybeEquiv.hasValue()) {
+#ifndef NDEBUG
+      SmallVector<OpOperand *> aliasingOpOperands =
+          getAliasingOpOperand(op, opResult, state);
+      assert(aliasingOpOperands.size() == 1 &&
+             "expected exactly 1 aliasing OpOperand");
+      assert(aliasingOpOperands.front()->getOperandNumber() ==
+                 maybeEquiv.getValue() &&
+             "inconsistent analysis state");
+#endif
       return BufferRelation::Equivalent;
     }
+    return BufferRelation::None;
+  }
 
   /// All function arguments are writable. It is the responsibility of the
   /// CallOp to insert buffer copies where necessary.
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     unsigned numResults = callOp.getNumResults();
     unsigned numOperands = callOp->getNumOperands();
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
     FunctionType funcType = funcOp.getFunctionType();
-    const FuncAnalysisState &funcState =
-        getFuncAnalysisState(state.getAnalysisState());
-    const OneShotBufferizationOptions &options =
-        static_cast<const OneShotBufferizationOptions &>(state.getOptions());
 
     // Result types of the bufferized CallOp.
     SmallVector<Type> resultTypes;
     // Replacement values for the existing CallOp. These are usually the results
     // of the bufferized CallOp, unless a tensor result folds onto an operand.
     SmallVector<Value> replacementValues(numResults, Value());
     // For non-tensor results: A mapping from return val indices of the old
     // CallOp to return val indices of the bufferized CallOp.
     SmallVector<Optional<unsigned>> retValMapping(numResults, None);
     // Operands of the bufferized CallOp.
     SmallVector<Value> newOperands(numOperands, Value());
 
-    // Based on previously gathered equivalence information, we know if a
-    // tensor result folds onto an operand. These are the only tensor value
-    // results that are supported at the moment.
-    //
-    // For tensors return values that do not fold onto an operand, additional
-    // work is needed (TODO) to either:
-    // * hoist a result into an inplaceable operand or
-    // * devise a better representation to truly return a buffer.
-    //
-    // Note: If a function has no body, no equivalence information is
-    // available. Consequently, a tensor return value cannot be proven to fold
-    // onto a FuncOp bbArg, so calls to such functions are not bufferizable at
-    // the moment.
-
-    // 1. Compute the result types of the new CallOp. Tensor results that are
-    // equivalent to a FuncOp bbArg are no longer returned.
+    // 1. Compute the result types of the new CallOp.
     for (const auto &it : llvm::enumerate(callOp.getResultTypes())) {
       unsigned returnValIdx = it.index();
       Type returnType = it.value();
       if (!returnType.isa<TensorType>()) {
         // Non-tensor values are returned.
         retValMapping[returnValIdx] = resultTypes.size();
         resultTypes.push_back(returnType);
         continue;
       }
 
-      if (options.dropEquivalentFuncResults) {
-        if (Optional<int64_t> bbArgIdx =
-                getEquivalentFuncArgIdx(funcOp, funcState, returnValIdx)) {
-          // Return operands that are equivalent to some bbArg, are not
-          // returned.
-          FailureOr<Value> bufferOrFailure =
-              state.getBuffer(rewriter, callOp->getOpOperand(*bbArgIdx));
-          if (failed(bufferOrFailure))
-            return failure();
-          replacementValues[returnValIdx] = *bufferOrFailure;
-          newOperands[*bbArgIdx] = *bufferOrFailure;
-          continue;
-        }
-      }
-
-      if (!options.allowReturnAllocs)
-        return callOp->emitError(
-            "call to FuncOp that returns non-equivalent tensors not supported");
-
-      // Returning a memref. This memref is not equivalent to any bbArg. It is
-      // likely a newly allocated buffer. We may want to hoist such allocations
-      // to the call site in the future.
+      // Returning a memref.
       retValMapping[returnValIdx] = resultTypes.size();
       resultTypes.push_back(funcType.getResult(resultTypes.size()));
     }
 
     // 2. Rewrite tensor operands as memrefs based on `bufferizedFuncType`.
     for (OpOperand &opOperand : callOp->getOpOperands()) {
       unsigned idx = opOperand.getOperandNumber();
       Value tensorOperand = opOperand.get();
 
       // Non-tensor operands are just copied.
       if (!tensorOperand.getType().isa<TensorType>()) {
         newOperands[idx] = tensorOperand;
         continue;
       }
 
-      // Retrieve buffers for tensor operands. Tensor operand buffers, who's
-      // corresponding FuncOp bbArgs are equivalent to a returned tensor, were
-      // already stored in `newOperands` during Step 1.
+      // Retrieve buffers for tensor operands.
       Value buffer = newOperands[idx];
       if (!buffer) {
         FailureOr<Value> bufferOrFailure = state.getBuffer(rewriter, opOperand);
         if (failed(bufferOrFailure))
           return failure();
         buffer = *bufferOrFailure;
       }
 
for (OpOperand &opOperand : callOp->getOpOperands()) {
       }
       newOperands[idx] = buffer;
     }
 
     // 3. Create the new CallOp.
     Operation *newCallOp = rewriter.create<func::CallOp>(
         callOp.getLoc(), funcOp.getSymName(), resultTypes, newOperands);
     newCallOp->setAttrs(callOp->getAttrs());
-    // Get replacement values for non-tensor / non-equivalent results.
+    // Get replacement values.
     for (unsigned i = 0; i < replacementValues.size(); ++i) {
       if (replacementValues[i])
         continue;
       replacementValues[i] = newCallOp->getResult(*retValMapping[i]);
     }
 
     // 4. Replace the old op with the new op.
     replaceOpWithBufferizedValues(rewriter, callOp, replacementValues);

 
 struct FuncOpInterface
     : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
   /// Rewrite function bbArgs and return values into buffer form. This function
   /// bufferizes the function signature and the ReturnOp. When the entire
   /// function body has been bufferized, function return types can be switched
   /// to more concise memref types as part of `foldMemRefCasts`.
   ///
-  /// When a tensor function argument is known to be equivalent to a tensor
-  /// result, it is dropped from the return values.
-  ///
   /// All function bbArgs are writable unless they are explicitly marked as
   /// read-only. Callers must insert copies when needed.
-  ///
-  /// Note: Returning a memref is possible, but corresponding CallOp
-  /// bufferizations fail unless `allowReturnAllocs`.
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto funcOp = cast<FuncOp>(op);
     FunctionType funcType = funcOp.getFunctionType();
-    const FuncAnalysisState &funcState =
-        getFuncAnalysisState(state.getAnalysisState());
     const OneShotBufferizationOptions &options =
         static_cast<const OneShotBufferizationOptions &>(state.getOptions());
 
     // Construct the bufferized function type.
     SmallVector<Type> argTypes;
     for (const auto &it : llvm::enumerate(funcType.getInputs())) {
       Type argType = it.value();
       if (auto tensorType = argType.dyn_cast<TensorType>()) {
for (OpOperand &returnOperand : returnOp->getOpOperands()) {
       rewriter.setInsertionPoint(returnOp);
 
       // If not a tensor type just forward it.
       if (!tensorType) {
         returnValues.push_back(returnVal);
         continue;
       }
 
-      // If return operand is equivalent to some bbArg, no need to return it.
-      if (options.dropEquivalentFuncResults) {
-        if (Optional<int64_t> equivBbArgIdx = getEquivalentFuncArgIdx(
-                funcOp, funcState, returnOperand.getOperandNumber())) {
-          // TODO: Use memref type with fully dynamic layout map and add folder
-          // for memref.cast + memref.copy.
-          Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
-              loc, getMemRefType(tensorType, options), returnVal);
-          BlockArgument equivBbArg = funcOp.getArgument(*equivBbArgIdx);
-          // Note: This copy will fold away. It must be inserted here to ensure
-          // that `returnVal` still has at least one use and does not fold away.
-          if (failed(
-                  options.createMemCpy(rewriter, loc, toMemrefOp, equivBbArg)))
-            return funcOp->emitError("could not generate copy for bbArg");
-          continue;
-        }
-      }
-
       BaseMemRefType resultType;
       if (options.functionBoundaryTypeConversion ==
           BufferizationOptions::LayoutMapOption::IdentityLayoutMap) {
         resultType = getMemRefTypeWithStaticIdentityLayout(tensorType);
       } else {
         // Note: If `InferLayoutMap`, cast are later folded away.
         resultType = getMemRefTypeWithFullyDynamicLayout(tensorType);
       }

mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp

 // gathered and stored in `FuncAnalysisState`.
 //
 // * `aliasingFuncOpBBArgsAnalysis` determines the equivalent/aliasing bbArgs
 // for
 //   each tensor return value (if any).
 // * `funcOpBbArgReadWriteAnalysis` determines whether or not a tensor bbArg is
 //   read/written.
 //
-// Only tensors that are equivalent to some FuncOp bbArg may be returned.
-// Bufferization currently fails if other tensors (in particular tensors that
-// bufferize out-of-place and result in a new buffer allocation) are returned.
-// In the future, such allocations could be hoisted to the caller.
-//
-// Example: `foo` fails bufferization because %0 is not equivalent to any bbArg.
-// ```
-// func @foo() -> tensor<?xf32> {
-//   %0 = bufferization.alloc_tensor(...) : tensor<?xf32>
-//   return %0 : tensor<?xf32>
-// }
-// ```
-//
 // Module Bufferization implements the following calling convention.
 //
 // * In the absence of conflicts within a FuncOp, the FuncOp's bbArgs may always
 //   be written to in-place.
 // * If a tensor operand of a CallOp is read after the CallOp, the operand of
 //   the CallOp must bufferize out-of-place.
 //
 // Example: The tensor.insert op bufferizes in-place because it is allowed to
for (func::FuncOp funcOp : orderedFuncOps) {
     if (failed(bufferizeOp(funcOp, bufferizationState)))
       return failure();
     // Change buffer return types to more precise layout maps.
     if (options.functionBoundaryTypeConversion ==
         BufferizationOptions::LayoutMapOption::InferLayoutMap)
       foldMemRefCasts(funcOp);
   }
 
-  // Check result.
-  for (func::FuncOp funcOp : orderedFuncOps) {
-    if (!options.allowReturnAllocs &&
-        llvm::any_of(funcOp.getFunctionType().getResults(), [](Type t) {
-          return t.isa<MemRefType, UnrankedMemRefType>();
-        })) {
-      funcOp->emitError("memref return type is unsupported");
-      return failure();
-    }
-  }
-
   // Post-pass cleanup of function argument attributes.
   moduleOp.walk([&](func::FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationAttributes(bbArg);
   });
 
   return success();
 }

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-allow-return-allocs.mlir

-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs" -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs drop-equivalent-func-results=false" -split-input-file | FileCheck %s --check-prefix=EQUIV
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs" -drop-equivalent-buffer-results -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=NO-DROP
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map" -split-input-file -o /dev/null

 
 func.func @return_arg(%A: tensor<?xf32>) -> tensor<?xf32> {
   func.return %A : tensor<?xf32>
 }
 // CHECK-LABEL: func @return_arg
 // CHECK-SAME:      %[[A:.*]]: memref<?xf32
 //  CHECK-NOT:    return %[[A]]
 
-// EQUIV-LABEL: func @return_arg
-// EQUIV-SAME:      %[[A:.*]]: memref<?xf32
-//      EQUIV:    return %[[A]]
+// NO-DROP-LABEL: func @return_arg
+//  NO-DROP-SAME:     %[[A:.*]]: memref<?xf32
+//       No_DROP:   return %[[A]]

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-out-params.mlir

-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=fully-dynamic-layout-map" -buffer-results-to-out-params -buffer-deallocation -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=identity-layout-map" -buffer-results-to-out-params -buffer-deallocation -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT
-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=infer-layout-map" -buffer-deallocation -split-input-file | FileCheck %s --check-prefix=CHECK-BASELINE
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=fully-dynamic-layout-map" -drop-equivalent-buffer-results -buffer-results-to-out-params -buffer-deallocation -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=identity-layout-map" -drop-equivalent-buffer-results -buffer-results-to-out-params -buffer-deallocation -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=infer-layout-map" -drop-equivalent-buffer-results -buffer-deallocation -split-input-file | FileCheck %s --check-prefix=CHECK-BASELINE
 
 // Note: function-boundary-type-conversion=infer-layout-map with
 // promote-buffer-results-to-out-params is an unsupported combination.
 
 // Note: This bufferization is not very efficient yet, but it works.
 
 // CHECK: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 // CHECK-LABEL: func @callee(

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir

 // Note: Default is function-boundary-type-conversion=infer-layout-map
-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs" -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs" -drop-equivalent-buffer-results -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
func.func @call_func_with_non_tensor_return(
   // CHECK: %[[alloc:.*]] = memref.alloc
   // CHECK-DAG: memref.copy %[[arg0]], %[[alloc]]
   // CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]
   // CHECK: %[[call:.*]] = call @inner_func(%[[casted]])
   %0, %1 = call @inner_func(%t0) : (tensor<?xf32>) -> (tensor<?xf32>, f32)
 
   // Note: The tensor return value cannot fold away because the CallOp
   // bufferized out-of-place.
-  // CHECK: return %[[call]], %[[alloc]] : f32, memref<?xf32>
+  // CHECK: return %[[call]], %[[casted]] : f32, memref<?xf32
   return %1, %0 : f32, tensor<?xf32>
 }
 
 // -----
 
 // A chain of function calls. The last function f0 is potentially writing to the
 // buffer. This becomes a problem when bufferizing main and a copy must be
 // inserted then. (No copies in the other functions.)

 
 // CHECK-LABEL: func @main(
 //  CHECK-SAME:     %[[t:.*]]: memref<?xf32
 //       CHECK:   %[[alloc:.*]] = memref.alloc
 //   CHECK-NOT:   copy
 //       CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
 //   CHECK-NOT:   copy
 //       CHECK:   call @does_not_read(%[[casted]])
-//       CHECK:   %[[r:.*]] = memref.load %[[alloc]]
+//       CHECK:   %[[r:.*]] = memref.load %[[casted]]
 //       CHECK:   memref.dealloc %[[alloc]]
 func.func @main(%t: tensor<?xf32> {bufferization.writable = false}) -> f32 {
   %0 = call @does_not_read(%t) : (tensor<?xf32>) -> (tensor<?xf32>)
   %idx = arith.constant 4 : index
   %r = tensor.extract %0[%idx] : tensor<?xf32>
   return %r : f32
 }
 
func.func @main() {
   %AA = linalg.fill ins(%v1 : f32) outs(%A : tensor<64xf32>) -> tensor<64xf32>
   %BB = linalg.fill ins(%v2 : f32) outs(%B : tensor<64xf32>) -> tensor<64xf32>
   %CC = linalg.fill ins(%v0 : f32) outs(%C : tensor<f32>) -> tensor<f32>
 
   // CHECK-NEXT:   call @init_and_dot(%[[cA]], %[[cB]], %[[cC]])
   %res = call @init_and_dot(%AA, %BB, %CC) :
     (tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32>
 
-  // CHECK-NEXT:   %[[dC:.*]] = memref.cast %[[C]] : memref<f32> to memref<*xf32>
+  // CHECK-NEXT:   %[[dC:.*]] = memref.cast %[[cC]] : memref<f32, {{.*}}> to memref<*xf32>
   %res2 = tensor.cast %res: tensor<f32> to tensor<*xf32>
 
   // CHECK-NEXT:   call @printMemrefF32(%[[dC]]) : (memref<*xf32>) -> ()
   call @printMemrefF32(%res2) : (tensor<*xf32>) -> ()
 
   // CHECK-DAG:   memref.dealloc %[[A]] : memref<64xf32>
   // CHECK-DAG:   memref.dealloc %[[B]] : memref<64xf32>
   // CHECK-DAG:   memref.dealloc %[[C]] : memref<f32>

 }
 
 // CHECK-LABEL: func @equivalent_func_arg(
 //  CHECK-SAME:     %[[arg0:.*]]: memref<?xf32
 func.func @equivalent_func_arg(%t0: tensor<?xf32> {bufferization.writable = true},
                                %c0: index, %c10: index, %c1: index) -> tensor<?xf32> {
   // CHECK-NOT: alloc
   // CHECK-NOT: copy
+  // CHECK: scf.for {{.*}} iter_args(%[[t1:.*]] = %[[arg0]])
   %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {
-    // CHECK: call @inner_func(%[[arg0]])
+    // CHECK: call @inner_func(%[[t1]])
     %3 = func.call @inner_func(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+    // CHECK: scf.yield %[[t1]]
     scf.yield %3 : tensor<?xf32>
   }
   return %1: tensor<?xf32>
 }
 
 // -----
 
 // inner_func_2 modifies the bbArg, but the loop yields the original value. A

mlir/test/Dialect/Linalg/one-shot-bufferize.mlir

-// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -buffer-loop-hoisting -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -buffer-loop-hoisting -drop-equivalent-buffer-results -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
-// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -drop-equivalent-buffer-results -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
 
 // TODO: Some test cases from this file should be moved to other dialects.
 
 // CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
 // CHECK-LABEL: func @fill_inplace(
 //  CHECK-SAME:   %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
 // CHECK-NO-LAYOUT-MAP-LABEL: func @fill_inplace(%{{.*}}: memref<?xf32>) {

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

-// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -buffer-deallocation -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -drop-equivalent-buffer-results -buffer-deallocation -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries" -buffer-deallocation -split-input-file -o /dev/null

mlir/test/Dialect/Tensor/one-shot-bufferize.mlir

-// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -drop-equivalent-buffer-results -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null
 
 // Test bufferization using memref types that have no layout map.
 // RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file -o /dev/null