[mlir] Introduce allowMemrefFunctionResults for the helper operation converters of buffer placement

[dfki-ehna]

Summary

This parameter gives the developers the freedom to choose their desired function signature conversion for preparing their functions for buffer placement. It is introduced for BufferAssignmentFuncOpConverter, and also for BufferAssignmentReturnOpConverter, and BufferAssignmentCallOpConverter to adapt the return and call operations with the selected function signature conversion. If the parameter is set, buffer placement won't also deallocate the returned buffers.
Here are the examples if it is set and unset:

Input IR:

func test(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
  return %arg0 : tensor<4x8xf32>
}

func caller(%arg0: tensor<4x8xf32>) {
  %0 = call @test(%arg0) : (tensor<4x8xf32>) -> (tensor<4x8xf32>)
  return
}

Transformed IR (allowMemrefFunctionResults = true):

func test(%arg0: memref<4x8xf32>) -> memref<4x8xf32> {
  return %arg0 : memref<4x8xf32>
}

func caller(%arg0: memref<4x8xf32>) {
  %0 = call @test(%arg0) : (memref<4x8xf32>) -> (memref<4x8xf32>)
  return
}

Transformed IR (allowMemrefFunctionResults = false):

func test(%arg0: memref<4x8xf32>, %result: memref<4x8xf32) {
  copy(%arg0, %result)
  return
}

func caller(%arg0: memref<4x8xf32>) {
  %0 = alloc() : memref<4x8xf32>
  call @test(%arg0, %0) : (memref<4x8xf32>, memref<4x8xf32>) -> ()
  return
}

[pifon2a]

Before I start reviewing deeper. I there some special reason to use template argument to flip the flag instead of using pass options?
https://github.com/llvm/llvm-project/blob/0bfd70bdad7e4ac22d96503fa78a5dd55d4b430e/mlir/include/mlir/Dialect/SCF/Passes.td#L29

we would get a command line flag like in
https://github.com/llvm/llvm-project/blob/0bfd70bdad7e4ac22d96503fa78a5dd55d4b430e/mlir/test/Dialect/SCF/parallel-loop-tiling.mlir#L1

and the code might become clearer.

mlir/test/lib/Transforms/TestBufferPlacement.cpp

38

Maybe move

using BAFuncOpConverter =
     BufferAssignmentFuncOpConverter<allowMemrefFunctionResults>;
 using BAReturnOpConverter =
     BufferAssignmentReturnOpConverter<ReturnOp, ReturnOp, linalg::CopyOp,
                                       allowMemrefFunctionResults>;
 using BACallOpConverter =
     BufferAssignmentCallOpConverter<allowMemrefFunctionResults>;

closer to insert<> call. In that case, just remove BA part:

using FuncOpConverter =
     BufferAssignmentFuncOpConverter<allowMemrefFunctionResults>;
 using ReturnOpConverter =
     BufferAssignmentReturnOpConverter<ReturnOp, ReturnOp, linalg::CopyOp,
                                       allowMemrefFunctionResults>;
 using CallOpConverter =
     BufferAssignmentCallOpConverter<allowMemrefFunctionResults>;

165

just out of curiosity:
will replacing Super:: with this-> work?

[rriddle]

mlir/include/mlir/Transforms/BufferPlacement.h

117

Return the error directly.

[pifon2a]

mlir/include/mlir/Transforms/BufferPlacement.h

105

let's not have default values for allowMemrefFunctionResults.

161

It might make sense to hide this behaviour behind some populatePatterns func template, so that the users won't add this template with different allowMemrefFunctionResults values.

222

though->through

mlir/test/Transforms/buffer-placement-preparation-allowed-memref-results.mlir

90

nit: reduce indentation so that it fits?

mlir/test/lib/Transforms/TestBufferPlacement.cpp

188

"test-buffer-placement-preparation-with-allowed-memref-results"? to make it consistent with the pass above.

mlir/tools/mlir-opt/mlir-opt.cpp

44

sort?

[dfki-ehna]

Address the comments.

[dfki-ehna]

Correct the position of registerTestPreparationPassWithAllowedMemrefResults.

[dfki-ehna]

Thanks for suggesting to use the Pass option. I made TestBufferPlacementPreparationPass a template struct due to the fact that registerTestBufferPlacementPreparationPass and registerTestPreparationPassWithAllowedMemrefResults are two passes for only test purposes which shouldn't be publicly seen.

mlir/include/mlir/Transforms/BufferPlacement.h

161

populateWithBufferAssignmentOpConversionPatternswas introduced. All three operation conversion patterns were moved to detail namespace in case if somebody needs to call them selectively.

mlir/test/lib/Transforms/TestBufferPlacement.cpp

38

Replaced with populateWithBufferAssignmentOpConversionPatterns

165

Replaced.

[pifon2a]

mlir/lib/Dialect/Linalg/Transforms/TensorsToBuffers.cpp

111

/*allowMemrefFunctionResults=*/false

mlir/include/mlir/Transforms/BufferPlacement.h

 
 #ifndef MLIR_TRANSFORMS_BUFFERPLACEMENT_H
 #define MLIR_TRANSFORMS_BUFFERPLACEMENT_H
 
 #include "mlir/Analysis/Liveness.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dominance.h"
+#include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 
 /// Prepares a buffer placement phase. It can place (user-defined) alloc
 /// nodes. This simplifies the integration of the actual buffer-placement
 /// pass. Sample usage:

 public:
   BufferAssignmentTypeConverter();
 
   /// A helper function to check if `type` has been converted from non-memref
   /// type to memref.
   static bool isConvertedMemref(Type type, Type before);
 };
 
-/// Converts the signature of the function using the type converter. It adds an
-/// extra argument for each function result type which is going to be a memref
-/// type after type conversion. The other function result types remain
-/// unchanged. `BufferAssignmentTypeConverter` is a helper `TypeConverter` for
-/// this purpose.
-class FunctionAndBlockSignatureConverter
+/// Converts the signature of the function based on whether the function is
+/// allowed to return memref typed results or not using
+/// `allowMemrefFunctionResults` parameter. If this option is false, then it
+/// adds an extra function argument as an output buffer for each function result
+/// which is going to be a memref type only after type conversion. The
+/// other function result types remain unchanged. If
+/// `allowMemrefFunctionResults` is true, the types are converted in place.
+/// Any changes in function signature need to be applied
+/// to return and caller operations. `BufferAssignmentReturnOpConverter` and
+/// `BufferAssignmentCallOpConverter` are two helper function that match the
+/// return and caller operation with the new function signature. Furthermore,
+/// `BufferAssignmentTypeConverter` is a helper `TypeConverter` for converting
+/// tensor typed values to memref typed ones.
+template <bool allowMemrefFunctionResults = false>
+class BufferAssignmentFuncOpConverter
     : public BufferAssignmentOpConversionPattern<FuncOp> {
 public:
   using BufferAssignmentOpConversionPattern<
       FuncOp>::BufferAssignmentOpConversionPattern;
 
   /// Performs the actual signature rewriting step.
   LogicalResult
-  matchAndRewrite(FuncOp funcOp, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const final;
+  matchAndRewrite(mlir::FuncOp funcOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    if (!converter) {
+      funcOp.emitError("The type converter has not been defined for "
+                       "BufferAssignmentFuncOpConverter");
+      return failure();
+    }
+    auto funcType = funcOp.getType();
+
+    // Convert function arguments using the provided TypeConverter.
+    TypeConverter::SignatureConversion conversion(funcType.getNumInputs());
+    for (auto argType : llvm::enumerate(funcType.getInputs()))
+      conversion.addInputs(argType.index(),
+                           converter->convertType(argType.value()));
+
+    // If allowMemrefFunctionResults is false and a function result type is not
+    // a memref but it would be a memref after type conversion, a new argument
+    // should be appended to the function arguments list for this result.
+    // Otherwise, it remains unchanged as a function result.
+    SmallVector<Type, 2> newResultTypes;
+    newResultTypes.reserve(funcOp.getNumResults());
+    for (Type resType : funcType.getResults()) {
+      Type convertedType = converter->convertType(resType);
+      if (!allowMemrefFunctionResults &&
+          BufferAssignmentTypeConverter::isConvertedMemref(convertedType,
+                                                           resType))
+        conversion.addInputs(convertedType);
+      else
+        newResultTypes.push_back(convertedType);
+    }
+
+    // Update the signature of the function.
+    rewriter.updateRootInPlace(funcOp, [&] {
+      funcOp.setType(rewriter.getFunctionType(conversion.getConvertedTypes(),
+                                              newResultTypes));
+      rewriter.applySignatureConversion(&funcOp.getBody(), conversion);
+    });
+    return success();
+  }
 };
 
 /// Rewrites the `ReturnOp` to conform with the changed function signature.
-/// Operands that correspond to return values that have been rewritten from
-/// tensor results to memref arguments are dropped. In their place, a
-/// corresponding copy operation from the operand to the new function argument
-/// is inserted.
+/// if allowMemrefFunctionResults is false, operands that correspond to return
+/// values and have been rewritten from illegal typed results to memref
+/// arguments are dropped. In their place, a corresponding copy operation from
+/// the operand to the output function argument is inserted. Otherwise, the
+/// memref typed operands are returned.
+/// Note: If this pattern rewriter is used with BufferAssignmentFuncOpConverter,
+/// allowMemrefFunctionResults must be set/unset for both.
 template <typename ReturnOpSourceTy, typename ReturnOpTargetTy,
-          typename CopyOpTy>
+          typename CopyOpTy, bool allowMemrefFunctionResults = false>
 class BufferAssignmentReturnOpConverter
     : public BufferAssignmentOpConversionPattern<ReturnOpSourceTy> {
 public:
   using BufferAssignmentOpConversionPattern<
       ReturnOpSourceTy>::BufferAssignmentOpConversionPattern;
 
   /// Performs the actual return-op conversion step.
   LogicalResult
   matchAndRewrite(ReturnOpSourceTy returnOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
+    // If the memref typed results can be returned as function results, the new
+    // `ReturnOp` should only return the type converted operands.
+    if (allowMemrefFunctionResults) {
+      rewriter.replaceOpWithNewOp<ReturnOpTargetTy>(returnOp, operands);
+      return success();
+    }
+
     // Split the operands by their kinds whether they are converted memref or
     // not.
     SmallVector<Value, 2> needCopyOperands, newOperands;
     unsigned operandsSize = operands.size();
     needCopyOperands.reserve(operandsSize);
     newOperands.reserve(operandsSize);
     for (auto operand : llvm::enumerate(operands))
       if (BufferAssignmentTypeConverter::isConvertedMemref(
matchAndRewrite(ReturnOpSourceTy returnOp, ArrayRef<Value> operands,
     }
 
     // Insert the new target Return operation.
     rewriter.replaceOpWithNewOp<ReturnOpTargetTy>(returnOp, newOperands);
     return success();
   }
 };
 
-/// Converts `CallOp` to match its operands and results with the
-/// the callee after rewriting the callee with
-/// FunctionAndBlockSignatureConverter.
+/// Rewrites the `CallOp` to match its operands and results with the signature
+/// of the callee after rewriting the callee with
+/// BufferAssignmentFuncOpConverter. If allowMemrefFunctionResults is false, a
+/// buffer is allocated as an output buffer only for each memref typed result
+/// that has been rewritten. The new allocated buffer is passed though the
+/// operands list of the new `CallOp`.
+/// Note: If this pattern rewriter is used with BufferAssignmentFuncOpConverter,
+/// allowMemrefFunctionResults must be set/unset for both.
+template <bool allowMemrefFunctionResults = false>
 class BufferAssignmentCallOpConverter
     : public BufferAssignmentOpConversionPattern<CallOp> {
 public:
   using BufferAssignmentOpConversionPattern<
       CallOp>::BufferAssignmentOpConversionPattern;
 
-  /// Performs the actual `CallOp` conversion step.
   LogicalResult
   matchAndRewrite(CallOp callOp, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const final;
+                  ConversionPatternRewriter &rewriter) const final {
+    Location loc = callOp.getLoc();
+
+    // If the memref typed results can be returned as function results, there is
+    // no need to create output buffers. It is only required to convert the type
+    // of operands and results in place for creating the new `CallOp`.
+    if (allowMemrefFunctionResults) {
+      SmallVector<Type, 2> resultTypes;
+      resultTypes.reserve(callOp.getNumResults());
+      for (Type type : callOp.getResultTypes())
+        resultTypes.push_back(converter->convertType(type));
+      rewriter.replaceOpWithNewOp<CallOp>(callOp, callOp.getCallee(),
+                                          resultTypes, operands);
+      return success();
+    }
+
+    SmallVector<Value, 2> newOperands, replacingValues;
+    SmallVector<Type, 2> newResultTypes;
+    unsigned numResults = callOp.getNumResults();
+    newOperands.reserve(numResults + operands.size());
+    newOperands.append(operands.begin(), operands.end());
+    newResultTypes.reserve(numResults);
+    replacingValues.reserve(numResults);
+
+    // For each memref result of `CallOp` which has not been a memref before
+    // the type conversion, a new buffer is allocated and passed to the operands
+    // list of the new `CallOp`. Otherwise, it remains as a caller result.
+    for (Value result : callOp.getResults()) {
+      Type currType = result.getType();
+      Type newType = converter->convertType(result.getType());
+      if (BufferAssignmentTypeConverter::isConvertedMemref(newType, currType)) {
+        OpBuilder::InsertionGuard guard(rewriter);
+        rewriter.restoreInsertionPoint(bufferAssignment->computeAllocPosition(
+            result.dyn_cast<OpResult>()));
+        Value alloc =
+            rewriter.create<AllocOp>(loc, newType.dyn_cast<MemRefType>());
+        newOperands.push_back(alloc);
+        replacingValues.push_back(alloc);
+      } else {
+        newResultTypes.push_back(currType);
+
+        // No replacing is required.
+        replacingValues.push_back(nullptr);
+      }
+    }
+
+    // Creating the new `CallOp`.
+    rewriter.create<CallOp>(loc, callOp.getCallee(), newResultTypes,
+                            newOperands);
+
+    // Replacing the results of the old `CallOp`.
+    rewriter.replaceOp(callOp, replacingValues);
+    return success();
+  }
 };
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_BUFFERPLACEMENT_H

mlir/lib/Dialect/Linalg/Transforms/TensorsToBuffers.cpp

 };
 
 /// Populate the given list with patterns to convert Linalg operations on
 /// tensors to buffers.
 static void populateConvertLinalgOnTensorsToBuffersPattern(
     MLIRContext *context, BufferAssignmentPlacer *placer,
     TypeConverter *converter, OwningRewritePatternList *patterns) {
   // clang-format off
-  patterns->insert<FunctionAndBlockSignatureConverter,
+  patterns->insert<BufferAssignmentFuncOpConverter<>,
                    GenericOpConverter,
                    ReturnOpConverter>(context, placer, converter);
   // clang-format on
 }
 
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
 struct ConvertLinalgOnTensorsToBuffers
     : public LinalgOnTensorsToBuffersBase<ConvertLinalgOnTensorsToBuffers> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);

mlir/lib/Transforms/BufferPlacement.cpp

 // be invalid with respect to program semantics. The only thing that is
 // currently missing is a high-level loop analysis that allows us to move allocs
 // and deallocs outside of the loop blocks. Furthermore, it doesn't also accept
 // functions which return buffers already.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/BufferPlacement.h"
-#include "mlir/IR/Function.h"
-#include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 
 namespace {
 
 //===----------------------------------------------------------------------===//

 /// Computes the actual position to place allocs for the given value.
 OpBuilder::InsertPoint
 BufferAssignmentPlacer::computeAllocPosition(OpResult result) {
   Operation *owner = result.getOwner();
   return OpBuilder::InsertPoint(owner->getBlock(), Block::iterator(owner));
 }
 
 //===----------------------------------------------------------------------===//
-// FunctionAndBlockSignatureConverter
-//===----------------------------------------------------------------------===//
-
-// Performs the actual signature rewriting step.
-LogicalResult FunctionAndBlockSignatureConverter::matchAndRewrite(
-    FuncOp funcOp, ArrayRef<Value> operands,
-    ConversionPatternRewriter &rewriter) const {
-  if (!converter) {
-    funcOp.emitError("The type converter has not been defined for "
-                     "FunctionAndBlockSignatureConverter");
-    return failure();
-  }
-  auto funcType = funcOp.getType();
-
-  // Convert function arguments using the provided TypeConverter.
-  TypeConverter::SignatureConversion conversion(funcType.getNumInputs());
-  for (auto argType : llvm::enumerate(funcType.getInputs()))
-    conversion.addInputs(argType.index(),
-                         converter->convertType(argType.value()));
-
-  // If a function result type is not a memref but it would be a memref after
-  // type conversion, a new argument should be appended to the function
-  // arguments list for this result. Otherwise, it remains unchanged as a
-  // function result.
-  SmallVector<Type, 2> newResultTypes;
-  newResultTypes.reserve(funcOp.getNumResults());
-  for (Type resType : funcType.getResults()) {
-    Type convertedType = converter->convertType(resType);
-    if (BufferAssignmentTypeConverter::isConvertedMemref(convertedType,
-                                                         resType))
-      conversion.addInputs(convertedType);
-    else
-      newResultTypes.push_back(convertedType);
-  }
-
-  // Update the signature of the function.
-  rewriter.updateRootInPlace(funcOp, [&] {
-    funcOp.setType(rewriter.getFunctionType(conversion.getConvertedTypes(),
-                                            newResultTypes));
-    rewriter.applySignatureConversion(&funcOp.getBody(), conversion);
-  });
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
-// BufferAssignmentCallOpConverter
-//===----------------------------------------------------------------------===//
-
-// Performs `CallOp` conversion to match its operands and results with the
-// signature of the callee after rewriting the callee with
-// FunctionAndBlockSignatureConverter.
-LogicalResult BufferAssignmentCallOpConverter::matchAndRewrite(
-    CallOp callOp, ArrayRef<Value> operands,
-    ConversionPatternRewriter &rewriter) const {
-
-  Location loc = callOp.getLoc();
-  SmallVector<Value, 2> newOperands, replacingValues;
-  SmallVector<Type, 2> newResultTypes;
-  unsigned numResults = callOp.getNumResults();
-  newOperands.reserve(numResults + operands.size());
-  newOperands.append(operands.begin(), operands.end());
-  newResultTypes.reserve(numResults);
-  replacingValues.reserve(numResults);
-
-  // For each memref result of `CallOp` which has not been a memref before type
-  // conversion, a new buffer is allocated and passed to the operands list of
-  // the new `CallOp`. Otherwise, it remains as a caller result.
-  for (Value result : callOp.getResults()) {
-    Type currType = result.getType();
-    Type newType = converter->convertType(result.getType());
-    if (BufferAssignmentTypeConverter::isConvertedMemref(newType, currType)) {
-      OpBuilder::InsertionGuard guard(rewriter);
-      rewriter.restoreInsertionPoint(
-          bufferAssignment->computeAllocPosition(result.dyn_cast<OpResult>()));
-      Value alloc =
-          rewriter.create<AllocOp>(loc, newType.dyn_cast<MemRefType>());
-      newOperands.push_back(alloc);
-      replacingValues.push_back(alloc);
-    } else {
-      newResultTypes.push_back(currType);
-
-      // No replacing is required.
-      replacingValues.push_back(nullptr);
-    }
-  }
-
-  // Creating the new `CallOp`.
-  rewriter.create<CallOp>(loc, callOp.getCallee(), newResultTypes, newOperands);
-
-  // Replacing the results of the old `CallOp`.
-  rewriter.replaceOp(callOp, replacingValues);
-
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
 // BufferAssignmentTypeConverter
 //===----------------------------------------------------------------------===//
 
 /// Registers conversions into BufferAssignmentTypeConverter
 BufferAssignmentTypeConverter::BufferAssignmentTypeConverter() {
   // Keep all types unchanged.
   addConversion([](Type type) { return type; });
   // A type conversion that converts ranked-tensor type to memref type.

mlir/test/Transforms/buffer-placement-preparation-allowed-memref-results.mlir

+// RUN: mlir-opt -test-preparation-with-allowed-memref-results -split-input-file %s | FileCheck %s -dump-input-on-failure
+
+// Since allowMemrefEscaping is on for Buffer Placement in this test pass, all
+// tensor typed function results are converted to memref and remain as function
+// results. All memref typed function results will escape from the deallocation
+// phase of Buffer Placement.
+
+// CHECK-LABEL: func @void_function_signature_conversion
+func @void_function_signature_conversion(%arg0: tensor<4x8xf32>) {
+    return
+}
+// CHECK: ({{.*}}: memref<4x8xf32>)
+
+// -----
+
+#map0 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: func @complex_signature_conversion
+func @complex_signature_conversion(%arg0: tensor<5xf32>, %arg1: memref<10xf32>, %arg2: i1, %arg3: f16) -> (i1, tensor<5xf32>, memref<10xf32>, memref<15xf32>, f16) {
+  %0 = alloc() : memref<15xf32>
+  %1 = linalg.generic {
+          args_in = 1 : i64,
+          args_out = 1 : i64,
+          indexing_maps = [#map0, #map0],
+          iterator_types = ["parallel"]
+        } %arg0 {
+        ^bb0(%gen1_arg0: f32):
+          %tmp1 = exp %gen1_arg0 : f32
+          linalg.yield %tmp1 : f32
+        }: tensor<5xf32> -> tensor<5xf32>
+  return %arg2, %1, %arg1, %0, %arg3 : i1, tensor<5xf32>, memref<10xf32>, memref<15xf32>, f16
+}
+//      CHECK: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>, %[[ARG2:.*]]: i1, %[[ARG3:.*]]: f16)
+// CHECK-SAME: (i1, memref<5xf32>, memref<10xf32>, memref<15xf32>, f16)
+//      CHECK: %[[FIRST_ALLOC:.*]] = alloc()
+//      CHECK: %[[LINALG_ALLOC:.*]] = alloc()
+//      CHECK: return %[[ARG2]], %[[LINALG_ALLOC]], %[[ARG1]], %[[FIRST_ALLOC]], %[[ARG3]]
+
+// -----
+
+// CHECK-LABEL: func @no_signature_conversion_is_needed
+func @no_signature_conversion_is_needed(%arg0: memref<4x8xf32>) {
+  return
+}
+// CHECK: ({{.*}}: memref<4x8xf32>)
+
+// -----
+
+// CHECK-LABEL: func @no_signature_conversion_is_needed
+func @no_signature_conversion_is_needed(%arg0: i1, %arg1: f16) -> (i1, f16){
+  return %arg0, %arg1 : i1, f16
+}
+// CHECK: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: f16) -> (i1, f16)
+// CHECK: return %[[ARG0]], %[[ARG1]]
+
+// -----
+
+// CHECK-LABEL: func @simple_signature_conversion
+func @simple_signature_conversion(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
+  return %arg0 : tensor<4x8xf32>
+}
+//      CHECK: (%[[ARG0:.*]]: [[TYPE:.*]]<[[RANK:.*]]>) -> [[TYPE]]<[[RANK]]>
+// CHECK-NEXT: return %[[ARG0]]
+
+// -----
+
+// CHECK-LABEL: func @func_and_block_signature_conversion
+func @func_and_block_signature_conversion(%arg0 : tensor<2xf32>, %cond : i1, %arg1: tensor<4x4xf32>) -> tensor<4x4xf32>{
+    cond_br %cond, ^bb1, ^bb2
+  ^bb1:
+    br ^exit(%arg0 : tensor<2xf32>)
+  ^bb2:
+    br ^exit(%arg0 : tensor<2xf32>)
+  ^exit(%arg2: tensor<2xf32>):
+    return %arg1 : tensor<4x4xf32>
+}
+//      CHECK: (%[[ARG0:.*]]: [[ARG0_TYPE:.*]], %[[COND:.*]]: i1, %[[ARG1:.*]]: [[ARG1_TYPE:.*]]) -> [[RESULT_TYPE:.*]]
+//      CHECK: br ^[[EXIT_BLOCK:.*]](%[[ARG0]] : [[ARG0_TYPE]])
+//      CHECK: br ^[[EXIT_BLOCK]](%[[ARG0]] : [[ARG0_TYPE]])
+//      CHECK: ^[[EXIT_BLOCK]](%{{.*}}: [[ARG0_TYPE]])
+// CHECK-NEXT:  return %[[ARG1]]
+
+// -----
+
+// CHECK-LABEL: func @callee
+func @callee(%arg1: tensor<5xf32>) -> (tensor<5xf32>, memref<2xf32>) {
+  %buff = alloc() : memref<2xf32>
+  return %arg1, %buff : tensor<5xf32>, memref<2xf32>
+}
+//      CHECK: (%[[CALLEE_ARG:.*]]: memref<5xf32>) -> (memref<5xf32>, memref<2xf32>)
+//      CHECK: %[[ALLOC:.*]] = alloc()
+//      CHECK: return %[[CALLEE_ARG]], %[[ALLOC]]
+
+// CHECK-LABEL: func @caller
+func @caller(%arg0: tensor<5xf32>) -> tensor<5xf32> {
+  %x:2 = call @callee(%arg0) : (tensor<5xf32>) -> (tensor<5xf32>, memref<2xf32>)
+  %y:2 = call @callee(%x#0) : (tensor<5xf32>) -> (tensor<5xf32>, memref<2xf32>)
+  return %y#0 : tensor<5xf32>
+}
+// CHECK: (%[[CALLER_ARG:.*]]: memref<5xf32>) -> memref<5xf32>
+// CHECK: %[[X:.*]]:2 = call @callee(%[[CALLER_ARG]])
+// CHECK: %[[Y:.*]]:2 = call @callee(%[[X]]#0)
+// CHECK: return %[[Y]]#0
+
+
+
+
+

mlir/test/Transforms/buffer-placement-preparation.mlir

 // CHECK-NEXT: %[[ALLOC6:.*]] = alloc()
 // CHECK-NEXT: linalg.generic {{.*}} %[[ARG0]], %[[ALLOC6]]
 //      CHECK: %[[ALLOC7:.*]] = alloc()
 // CHECK-NEXT: linalg.generic {{.*}} %[[ALLOC6]], %[[ALLOC7]]
 
 // -----
 
 // Test case: Checking BufferAssignmentCallOpConverter and
-// FunctionAndBlockSignatureConverter and BufferAssignmentReturnOpConverter all
+// BufferAssignmentFuncOpConverter and BufferAssignmentReturnOpConverter all
 // together. The signature of `callee` after signature conversion would be:
 
 // func @callee(%arg0: memref<5xf32>,%arg1: memref<5xf32>) -> ()
 
 // The operands and results of caller and return operations must be matched
 // respectively.
 
 #map0 = affine_map<(d0) -> (d0)>

 // CHECK: %[[SECOND_ALLOC:.*]] = alloc()
 // CHECK: call @callee(%[[FIRST_ALLOC]], %[[SECOND_ALLOC]])
 // CHECK: linalg.copy(%[[SECOND_ALLOC]], %[[CALLER_RESULT]])
 // CHECK: return
 
 // -----
 
 // Test case: Checking BufferAssignmentCallOpConverter and
-// FunctionAndBlockSignatureConverter and BufferAssignmentReturnOpConverter all
+// BufferAssignmentFuncOpConverter and BufferAssignmentReturnOpConverter all
 // together on functions that also have memref typed results. The signature of
 // `callee` after signature conversion would be:
 
 // func @callee(%arg0: memref<5xf32>,%arg1: memref<5xf32>)-> memref<2xf32>
 
 // where %arg0 is the input and %arg1 is the output buffer and the original memref
 // type result remain as the function result. Then, the rewriter should match the
 // caller's signature with the callee. Thus, two buffers will be allocated instead

mlir/test/lib/Transforms/TestBufferPlacement.cpp

 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/BufferPlacement.h"
 
 using namespace mlir;
 
 namespace {
-/// This pass tests the computeAllocPosition helper method and two provided
-/// operation converters, FunctionAndBlockSignatureConverter and
-/// BufferAssignmentReturnOpConverter. Furthermore, this pass converts linalg
-/// operations on tensors to linalg operations on buffers to prepare them for
-/// the BufferPlacement pass that can be applied afterwards.
+/// This pass tests the computeAllocPosition helper method and buffer assignment
+/// operation converters. Furthermore, this pass converts linalg operations on
+/// tensors to linalg operations on buffers to prepare them for the
+/// BufferPlacement pass that can be applied afterwards.
+/// `allowMemrefFunctionResults` informs the buffer placement to allow functions
+/// that have memref typed results. Buffer assignment operation converters will
+/// be adapted respectively. It will also allow memref typed results to escape
+/// from the deallocation.
+template <bool allowMemrefFunctionResults = false>
 struct TestBufferPlacementPreparationPass
-    : mlir::PassWrapper<TestBufferPlacementPreparationPass,
-                        OperationPass<ModuleOp>> {
+    : mlir::PassWrapper<
+          TestBufferPlacementPreparationPass<allowMemrefFunctionResults>,
+  using Super = TestBufferPlacementPreparationPass<allowMemrefFunctionResults>;
+  using BAFuncOpConverter =
+      BufferAssignmentFuncOpConverter<allowMemrefFunctionResults>;
+  using BAReturnOpConverter =
+      BufferAssignmentReturnOpConverter<ReturnOp, ReturnOp, linalg::CopyOp,
+                                        allowMemrefFunctionResults>;
+  using BACallOpConverter =
+      BufferAssignmentCallOpConverter<allowMemrefFunctionResults>;
 
-  /// Converts tensor-type generic linalg operations to memref ones using buffer
-  /// assignment.
+  /// Converts tensor-type generic linalg operations to memref ones using
+  /// buffer assignment.
   class GenericOpConverter
       : public BufferAssignmentOpConversionPattern<linalg::GenericOp> {
   public:
     using BufferAssignmentOpConversionPattern<
         linalg::GenericOp>::BufferAssignmentOpConversionPattern;
 
     LogicalResult
     matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
public:
     }
   };
 
   void populateTensorLinalgToBufferLinalgConversionPattern(
       MLIRContext *context, BufferAssignmentPlacer *placer,
       TypeConverter *converter, OwningRewritePatternList *patterns) {
     // clang-format off
     patterns->insert<
-                   BufferAssignmentCallOpConverter,
-                   FunctionAndBlockSignatureConverter,
-                   GenericOpConverter,
-                   BufferAssignmentReturnOpConverter<
-                      ReturnOp, ReturnOp, linalg::CopyOp>
+                  BACallOpConverter,
+                  BAFuncOpConverter,
+                  BAReturnOpConverter,
+                  GenericOpConverter
     >(context, placer, converter);
     // clang-format on
   }
 
   void runOnOperation() override {
-    MLIRContext &context = getContext();
+    MLIRContext &context = Super::getContext();
     ConversionTarget target(context);
     BufferAssignmentTypeConverter converter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<StandardOpsDialect>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
     auto isIllegalType = [&](Type type) { return !converter.isLegal(type); };
void runOnOperation() override {
     });
 
     // Mark the function whose arguments are in tensor-type illegal.
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
       return converter.isSignatureLegal(funcOp.getType());
     });
 
     // Walk over all the functions to apply buffer assignment.
-    getOperation().walk([&](FuncOp function) -> WalkResult {
+    Super::getOperation().walk([&](FuncOp function) -> WalkResult {
       OwningRewritePatternList patterns;
       BufferAssignmentPlacer placer(function);
       populateTensorLinalgToBufferLinalgConversionPattern(
           &context, &placer, &converter, &patterns);
 
       // Applying full conversion
       return applyFullConversion(function, target, patterns, &converter);
     });
   };
 };
 } // end anonymous namespace
 
 namespace mlir {
 void registerTestBufferPlacementPreparationPass() {
-  PassRegistration<TestBufferPlacementPreparationPass>(
+  PassRegistration<TestBufferPlacementPreparationPass<false>>(
       "test-buffer-placement-preparation",
       "Tests buffer placement helper methods including its "
       "operation-conversion patterns");
 }
-} // end namespace mlir
+
+void registerTestPreparationPassWithAllowedMemrefResults() {
+  PassRegistration<TestBufferPlacementPreparationPass<true>>(
+      "test-preparation-with-allowed-memref-results",
+      "Tests the helper operation converters of buffer placement for allowing "
+      "functions to have memref typed results.");
+}
+} // end namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerPatternsTestPass();
 void registerPrintOpAvailabilityPass();
 void registerSideEffectTestPasses();
 void registerSimpleParametricTilingPass();
 void registerSymbolTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAllReduceLoweringPass();
 void registerTestAffineLoopUnswitchingPass();
+void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestBufferPlacementPreparationPass();
 void registerTestLoopPermutationPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDominancePass();
 void registerTestFunc();
void registerTestPasses() {
   registerTestCallGraphPass();
   registerTestConstantFold();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToHsacoPass();
 #endif
+  registerTestPreparationPassWithAllowedMemrefResults();
   registerTestBufferPlacementPreparationPass();
   registerTestDominancePass();
   registerTestFunc();
   registerTestGpuMemoryPromotionPass();
   registerTestLinalgTransforms();
   registerTestLivenessPass();
   registerTestLoopFusion();
   registerTestLoopMappingPass();