[mlir] Add a flag to allow equivalent results.

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

struct OneShotBufferizationOptions : public BufferizationOptions {
   }
 
   /// Registered post analysis steps.
   PostAnalysisStepList postAnalysisSteps;
 
   /// 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.td

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.fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;

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

for (const auto &it : llvm::enumerate(callOp.getResultTypes())) {
       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.
struct FuncOpInterface
   /// 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 BufferizationOptions &options = state.getOptions();
+    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>()) {
         argTypes.push_back(
             getBufferizedFunctionArgType(funcOp, it.index(), options));
for (OpOperand &returnOperand : returnOp->getOpOperands()) {
 
       // If not a tensor type just forward it.
       if (!returnVal.getType().isa<RankedTensorType>()) {
         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())) {
           rewriter.setInsertionPoint(returnOp);
           Location loc = returnOp.getLoc();
           Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
-            loc, getMemRefType(returnVal.getType().cast<TensorType>(), options),
+              loc,
+              getMemRefType(returnVal.getType().cast<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(
                   createMemCpy(rewriter, loc, toMemrefOp, equivBbArg, options)))
             return funcOp->emitError("could not generate copy for bbArg");
           continue;
         }
+      }
 
       returnValues.push_back(*state.getBuffer(rewriter, returnOperand));
     }
 
     // 3. Rewrite the terminator without the in-place bufferizable values.
     returnOp.operandsMutable().assign(returnValues);
 
     // 4. Rewrite the FuncOp type to buffer form.

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 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 fully-dynamic-layout-maps=0" -split-input-file -o /dev/null

 func.func @main(%t: tensor<?xf32>, %sz: index, %idx: index) -> (f32, f32) {
   %cst = arith.constant 1.0 : f32
   %0 = call @return_slice(%t, %sz) : (tensor<?xf32>, index) -> (tensor<?xf32>)
   %filled = linalg.fill ins(%cst : f32) outs(%t : tensor<?xf32>) -> tensor<?xf32>
   %r1 = tensor.extract %0[%idx] : tensor<?xf32>
   %r2 = tensor.extract %filled[%idx] : tensor<?xf32>
   return %r1, %r2 : f32, f32
 }
+
+// -----
+
+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]]