[mlir] Remove AppendToArgumentsList functionality from BufferizeTypeConverter.

[silvas]

Summary

This functionality is superceded by BufferResultsToOutParams pass (see
https://reviews.llvm.org/D90071) for users the require buffers to be
out-params. That pass should be run immediately after all tensors are gone from
the program (before buffer optimizations and deallocation insertion), such as
immediately after a "finalizing" bufferize pass.

The -test-finalizing-bufferize pass now defaults to what used to be the
allowMemrefFunctionResults=true flag. and the
finalizing-bufferize-allowed-memref-results.mlir file is moved
to test/Transforms/finalizing-bufferize.mlir.

[dfki-albo]

mlir/test/Transforms/finalizing-bufferize.mlir

15

nit: %arg3: f16 in a new line.

85–86

Check if the return yields the result type.

mlir/include/mlir/Transforms/Bufferize.h

 
 namespace mlir {
 
 /// A helper type converter class for using inside Buffer Assignment operation
 /// conversion patterns. The default constructor keeps all the types intact
 /// except for the ranked-tensor types which is converted to memref types.
 class BufferizeTypeConverter : public TypeConverter {
 public:
-  /// This enum is for showing how buffer placement operation converters should
-  /// conduct with certain result type after type conversion. This value can be
-  /// set/get for each specific type using setResultConversionKind or
-  /// getResultConversionKind.
-  enum ResultConversionKind { AppendToArgumentsList, KeepAsFunctionResult };
-
   BufferizeTypeConverter();
 
   /// This method tries to decompose a value of a certain type using provided
   /// decompose callback functions. If it is unable to do so, the original value
   /// is returned.
   void tryDecomposeValue(OpBuilder &, Location, Type, Value,
                          SmallVectorImpl<Value> &);
 
template <typename FnT,
             typename T = typename llvm::function_traits<FnT>::template arg_t<0>>
   void addDecomposeTypeConversion(FnT &&callback) {
     auto wrapper =
         wrapDecomposeTypeConversionCallback<T>(std::forward<FnT>(callback));
     decomposeTypeConversions.emplace_back(wrapper);
     addConversion(std::forward<FnT>(callback));
   }
 
-  /// This method returns ResultConversionKind for the mapping from `origin`
-  /// type to `input` type.
-  ResultConversionKind getResultConversionKind(Type origin, Type input);
-
-  /// This method registers ResultConversionKind for the mapping from type 'T'
-  /// to type 'U'.
-  template <typename T, typename U>
-  void setResultConversionKind(ResultConversionKind kind) {
-    assert((kind != AppendToArgumentsList ||
-            llvm::is_one_of<U, MemRefType, UnrankedMemRefType>::value) &&
-           "Only the memref typed values can be set to be appended to the "
-           "function argument list at the moment");
-    resultTypeConversions.emplace_back(
-        [=](Type origin, Type input) -> Optional<ResultConversionKind> {
-          if (origin.template isa<T>() && input.template isa<U>())
-            return kind;
-          return llvm::None;
-        });
-  }
-
 private:
   using DecomposeValueConversionCallFn = std::function<Optional<LogicalResult>(
       OpBuilder &, Location, Type, Value, SmallVectorImpl<Value> &)>;
 
   using DecomposeTypeConversionCallFn =
       std::function<Optional<LogicalResult>(Type, SmallVectorImpl<Type> &)>;
 
-  using ResultConversionKindFn =
-      std::function<Optional<ResultConversionKind>(Type, Type)>;
-
   /// Generate a wrapper for the given decompose value conversion callback.
   template <typename T, typename FnT>
   DecomposeValueConversionCallFn
   wrapDecomposeValueConversionCallback(FnT &&callback) {
     return [callback = std::forward<FnT>(callback)](
                OpBuilder &builder, Location loc, Type type, Value value,
                SmallVectorImpl<Value> &newValues) -> Optional<LogicalResult> {
       if (T derivedType = type.dyn_cast<T>())
return [callback = std::forward<FnT>(callback)](
                SmallVectorImpl<Type> &results) -> Optional<LogicalResult> {
       T derivedType = type.dyn_cast<T>();
       if (!derivedType)
         return llvm::None;
       return callback(derivedType, results);
     };
   }
 
-  SmallVector<ResultConversionKindFn, 2> resultTypeConversions;
   SmallVector<DecomposeValueConversionCallFn, 2> decomposeValueConversions;
   SmallVector<DecomposeTypeConversionCallFn, 2> decomposeTypeConversions;
 };
 
 /// Marks ops used by bufferization for type conversion materializations as
 /// "legal" in the given ConversionTarget.
 ///
 /// This function should be called by all bufferization passes using

 public:
   using BufferizeOpConversionPattern<
       ReturnOpSourceTy>::BufferizeOpConversionPattern;
 
   /// Performs the actual return-op conversion step.
   LogicalResult
   matchAndRewrite(ReturnOpSourceTy returnOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
-    Location loc = returnOp.getLoc();
-
-    // Split the operands depending on whether they need a copy operation or
-    // they remain as operands of the return operation. If an operand is
-    // decomposable and a decompose callback function has been provided by the
-    // user, it will be unpacked.
-    SmallVector<Value, 2> newOperands, needCopyOperands;
-    OpBuilder builder(returnOp);
-    for (auto operand : llvm::enumerate(operands)) {
-      SmallVector<Value, 2> values;
-      this->converter.tryDecomposeValue(builder, loc, operand.value().getType(),
-                                        operand.value(), values);
-      Type type = returnOp.getOperand(operand.index()).getType();
-      SmallVector<Type, 2> originTypes;
-      this->converter.tryDecomposeType(type, originTypes);
-      for (auto value : llvm::enumerate(values)) {
-        Type origin = originTypes[value.index()];
-        Type converted = value.value().getType();
-        auto kind = this->converter.getResultConversionKind(origin, converted);
-        if (kind == BufferizeTypeConverter::KeepAsFunctionResult)
-          newOperands.push_back(value.value());
-        else
-          // kind = BufferizeTypeConverter::AppendToArgumentsList
-          needCopyOperands.push_back(value.value());
-      }
-    }
-
-    // Insert Copy operations instead for the operands that have been removed
-    // from operand list and appended to the function arguments list.
-    Block &entryBlock = returnOp.getParentRegion()->front();
-    unsigned numFuncArgs = entryBlock.getNumArguments();
-    if (needCopyOperands.size() > numFuncArgs)
-      return returnOp.emitError(
-          "The number of operands that need Copy operations is more "
-          "than the number of target function arguments.");
-    unsigned destArgNum = numFuncArgs - needCopyOperands.size();
-    rewriter.setInsertionPoint(returnOp);
-    for (Value operand : needCopyOperands) {
-      rewriter.create<CopyOpTy>(loc, operand,
-                                entryBlock.getArgument(destArgNum));
-      ++destArgNum;
-    }
+    SmallVector<Value, 2> newOperands;
+    for (auto operand : operands)
+      this->converter.tryDecomposeValue(
+          rewriter, returnOp.getLoc(), operand.getType(), operand, newOperands);
     rewriter.replaceOpWithNewOp<ReturnOpTargetTy>(returnOp, newOperands);
     return success();
   }
 };
 
 /// Rewrites the `CallOp` to match its operands and results with the signature
 /// of the callee after rewriting the callee with
 /// BufferizeFuncOpConverter.

mlir/lib/Transforms/Bufferize.cpp

 void BufferizeTypeConverter::tryDecomposeType(Type type,
                                               SmallVectorImpl<Type> &types) {
   for (auto &conversion : decomposeTypeConversions)
     if (conversion(type, types))
       return;
   types.push_back(type);
 }
 
-/// This method returns ResultConversionKind for the input type.
-BufferizeTypeConverter::ResultConversionKind
-BufferizeTypeConverter::getResultConversionKind(Type origin, Type converted) {
-  for (auto &conversion : resultTypeConversions)
-    if (auto res = conversion(origin, converted))
-      return res.getValue();
-  return KeepAsFunctionResult;
-}
-
 void mlir::populateBufferizeMaterializationLegality(ConversionTarget &target) {
   target.addLegalOp<TensorLoadOp, TensorToMemrefOp>();
 }
 
 namespace {
 // In a finalizing bufferize conversion, we know that all tensors have been
 // converted to memrefs, thus, this op becomes an identity.
 class BufferizeTensorLoadOp : public OpConversionPattern<TensorLoadOp> {
LogicalResult BufferizeFuncOpConverter::matchAndRewrite(
   }
 
   // Convert the result types of the function.
   SmallVector<Type, 2> newResultTypes;
   newResultTypes.reserve(funcOp.getNumResults());
   for (Type resultType : funcType.getResults()) {
     SmallVector<Type, 2> originTypes;
     converter.tryDecomposeType(resultType, originTypes);
-    for (auto origin : originTypes) {
-      Type converted = converter.convertType(origin);
-      auto kind = converter.getResultConversionKind(origin, converted);
-      if (kind == BufferizeTypeConverter::AppendToArgumentsList) {
-        conversion.addInputs(converted);
-      } else {
-        assert(kind == BufferizeTypeConverter::KeepAsFunctionResult);
-        newResultTypes.push_back(converted);
-      }
-    }
+    for (auto origin : originTypes)
+      newResultTypes.push_back(converter.convertType(origin));
   }
 
   if (failed(rewriter.convertRegionTypes(&funcOp.getBody(), converter,
                                          &conversion)))
     return failure();
 
   // Update the signature of the function.
   rewriter.updateRootInPlace(funcOp, [&] {
     funcOp.setType(rewriter.getFunctionType(conversion.getConvertedTypes(),
                                             newResultTypes));
   });
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // BufferizeCallOpConverter
 //===----------------------------------------------------------------------===//
 
-namespace {
-// This class represents a mapping from a result to a list of values and some
-// results that have not yet constructed. Instead, the indices of these
-// results in the operation that will be constructed are known. They will be
-// replaced with the actual values when they are available. The order of
-// adding to this mapping is important.
-class CallOpResultMapping {
-public:
-  CallOpResultMapping() { order = 0; };
-
-  /// Add an available value to the mapping.
-  void addMapping(Value value) { toValuesMapping.push_back({order++, value}); }
-
-  /// Add the index of unavailble result value to the mapping.
-  void addMapping(unsigned index) {
-    toIndicesMapping.push_back({order++, index});
-  }
-
-  /// This method returns the mapping values list. The unknown result values
-  /// that only their indices are available are replaced with their values.
-  void getMappingValues(ValueRange valuesToReplaceIndices,
-                        SmallVectorImpl<Value> &values) {
-    // Append available values to the list.
-    SmallVector<std::pair<unsigned, Value>, 2> res(toValuesMapping.begin(),
-                                                   toValuesMapping.end());
-    // Replace the indices with the actual values.
-    for (const std::pair<unsigned, unsigned> &entry : toIndicesMapping) {
-      assert(entry.second < valuesToReplaceIndices.size() &&
-             "The value index is out of range.");
-      res.push_back({entry.first, valuesToReplaceIndices[entry.second]});
-    }
-    // Sort the values based on their adding orders.
-    llvm::sort(res, [](const std::pair<unsigned, Value> &v1,
-                       const std::pair<unsigned, Value> &v2) {
-      return v1.first < v2.first;
-    });
-    // Fill the values.
-    for (const std::pair<unsigned, Value> &entry : res)
-      values.push_back(entry.second);
-  }
-
-private:
-  /// Keeping the inserting order of mapping values.
-  int order;
-
-  /// Containing the mapping values with their inserting orders.
-  SmallVector<std::pair<unsigned, Value>, 2> toValuesMapping;
-
-  /// Containing the indices of result values with their inserting orders.
-  SmallVector<std::pair<unsigned, unsigned>, 2> toIndicesMapping;
-};
-} // namespace
-
 /// Performs the actual rewriting step.
 LogicalResult BufferizeCallOpConverter::matchAndRewrite(
     CallOp callOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
 
   Location loc = callOp.getLoc();
-  OpBuilder builder(callOp);
   SmallVector<Value, 2> newOperands;
 
   // TODO: if the CallOp references a FuncOp that only has a declaration (e.g.
   // to an externally defined symbol like an external library calls), only
   // convert if some special attribute is set.
   // This will allow more control of interop across ABI boundaries.
 
   // Create the operands list of the new `CallOp`. It unpacks the decomposable
   // values if a decompose callback function has been provided by the user.
-  for (auto operand : operands) {
-    SmallVector<Value, 2> values;
-    converter.tryDecomposeValue(builder, loc, operand.getType(), operand,
-                                values);
-    newOperands.append(values.begin(), values.end());
-  }
+  for (auto operand : operands)
+    converter.tryDecomposeValue(rewriter, loc, operand.getType(), operand,
+                                newOperands);
 
-  // Create the new result types for the new `CallOp` and a mapping from the old
-  // result to new value(s).
+  // Create the new result types for the new `CallOp` and track the indices in
+  // the new call op's results that correspond to the old call op's results.
   SmallVector<Type, 2> newResultTypes;
-  SmallVector<CallOpResultMapping, 4> mappings;
-  mappings.resize(callOp.getNumResults());
+  SmallVector<SmallVector<int, 2>, 4> expandedResultIndices;
+  expandedResultIndices.resize(callOp.getNumResults());
   for (auto result : llvm::enumerate(callOp.getResults())) {
     SmallVector<Type, 2> originTypes;
     converter.tryDecomposeType(result.value().getType(), originTypes);
-    auto &resultMapping = mappings[result.index()];
+    auto &resultMapping = expandedResultIndices[result.index()];
     for (Type origin : originTypes) {
       Type converted = converter.convertType(origin);
-      auto kind = converter.getResultConversionKind(origin, converted);
-      if (kind == BufferizeTypeConverter::KeepAsFunctionResult) {
-        newResultTypes.push_back(converted);
-        // The result value is not yet available. Its index is kept and it is
-        // replaced with the actual value of the new `CallOp` later.
-        resultMapping.addMapping(newResultTypes.size() - 1);
-      } else {
-        // kind = BufferizeTypeConverter::AppendToArgumentsList
-        MemRefType memref = converted.dyn_cast<MemRefType>();
-        if (!memref)
-          return callOp.emitError("Cannot allocate for a non-Memref type");
-        Value alloc = rewriter.create<AllocOp>(loc, memref);
-        newOperands.push_back(alloc);
-        resultMapping.addMapping(alloc);
-      }
+      resultMapping.push_back(newResultTypes.size() - 1);
     }
   }
 
   CallOp newCallOp = rewriter.create<CallOp>(loc, callOp.getCallee(),
                                              newResultTypes, newOperands);
 
   // Build a replacing value for each result to replace its uses. If a result
   // has multiple mapping values, it needs to be packed to a single value.
-  OpBuilder nextBuilder(callOp.getOperation()->getNextNode());
   SmallVector<Value, 2> replacedValues;
   replacedValues.reserve(callOp.getNumResults());
   for (unsigned i = 0, e = callOp.getNumResults(); i < e; ++i) {
-    SmallVector<Value, 2> valuesToPack;
-    mappings[i].getMappingValues(newCallOp.getResults(), valuesToPack);
+    auto valuesToPack = llvm::to_vector<6>(
+        llvm::map_range(expandedResultIndices[i],
+                        [&](int i) { return newCallOp.getResult(i); }));
     if (valuesToPack.empty()) {
       // No replacement is required.
       replacedValues.push_back(nullptr);
     } else if (valuesToPack.size() == 1) {
       replacedValues.push_back(valuesToPack.front());
     } else {
       // Values need to be packed using callback function. The same callback
       // that is used for materializeArgumentConversion is used for packing.
       Value packed = converter.materializeArgumentConversion(
-          nextBuilder, loc, callOp.getType(i), valuesToPack);
+          rewriter, loc, callOp.getType(i), valuesToPack);
       replacedValues.push_back(packed);
     }
   }
   rewriter.replaceOp(callOp, replacedValues);
   return success();
 }

mlir/test/Transforms/finalizing-bufferize-allowed-memref-results.mlir

-// RUN: mlir-opt -test-finalizing-bufferize-with-allowed-memref-results -split-input-file %s | FileCheck %s
-
-// Since allowMemrefEscaping is active for Bufferization 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 Bufferization.
-
-// CHECK-LABEL: func @void_function_signature_conversion
-func @void_function_signature_conversion(%arg0: tensor<4x8xf32>) {
-    return
-}
-// CHECK: ({{.*}}: memref<4x8xf32>)
-
-// -----
-
-// 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 = test.tensor_based in(%arg0 : 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>,
-// CHECK-SAME: %[[ARG2:.*]]: i1, %[[ARG3:.*]]: f16)
-// CHECK-SAME: (i1, memref<5xf32>, memref<10xf32>, memref<15xf32>, f16)
-//      CHECK: %[[FIRST_ALLOC:.*]] = alloc()
-//      CHECK: %[[TENSOR_ALLOC:.*]] = alloc()
-//      CHECK: return %[[ARG2]], %[[TENSOR_ALLOC]], %[[ARG1]], %[[FIRST_ALLOC]],
-// CHECK-SAME: %[[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_with_unranked_arg_and_result
-func @func_with_unranked_arg_and_result(%arg0: tensor<*xf32>) -> tensor<*xf32> {
-  return %arg0 : tensor<*xf32>
-}
-// CHECK-SAME: ([[ARG:%.*]]: memref<*xf32>) -> memref<*xf32>
-// CHECK-NEXT: return [[ARG]] : memref<*xf32>
-
-// -----
-
-// 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
-
-// -----
-
-// Test case: Testing BufferizeCallOpConverter to see if it matches with the
-// signature of the new signature of the callee function when there are tuple
-// typed args and results. BufferizeTypeConverter is set to flatten tuple typed
-// arguments. The tuple typed values should be decomposed and composed using
-// get_tuple_element and make_tuple operations of test dialect. Tensor types are
-// converted to Memref. Memref typed function results remain as function
-// results.
-
-// CHECK-LABEL: func @callee
-func @callee(%arg0: tuple<tensor<2xf32>,i1, tensor<5xf32>>) -> (tuple<tensor<2xf32>,i1, tensor<5xf32>>){
-  return %arg0 : tuple<tensor<2xf32>,i1, tensor<5xf32>>
-}
-// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1, %[[ARG2:.*]]: memref<5xf32>)
-// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>)
-// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) {index = 0 : i32}
-// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) {index = 1 : i32}
-// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[TUPLE]]) {index = 2 : i32}
-// CHECK-NEXT: return %[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]]
-
-// CHECK-LABEL: func @caller
-func @caller(%arg0: tuple<tensor<2xf32>,i1, tensor<5xf32>>) -> tuple<tensor<2xf32>,i1, tensor<5xf32>>{
-  %x0 = call @callee(%arg0) : (tuple<tensor<2xf32>,i1, tensor<5xf32>>) -> (tuple<tensor<2xf32>,i1, tensor<5xf32>>)
-  %y0 = call @callee(%x0) : (tuple<tensor<2xf32>,i1, tensor<5xf32>>) -> (tuple<tensor<2xf32>,i1, tensor<5xf32>>)
-  return %y0 : tuple<tensor<2xf32>,i1, tensor<5xf32>>
-}
-// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1, %[[ARG2:.*]]: memref<5xf32>)
-// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>)
-// CHECK-NEXT: %[[ARG_TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 0 : i32}
-// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 1 : i32}
-// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 2 : i32}
-// CHECK-NEXT: %[[CALLEE_RESULTS:.*]]:3 = call @callee(%[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]])
-// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) -> (memref<2xf32>, i1, memref<5xf32>)
-// CHECK-NEXT: %[[RESULT_TUPLE:.*]] = "test.make_tuple"(%[[CALLEE_RESULTS]]#0, %[[CALLEE_RESULTS]]#1, %[[CALLEE_RESULTS]]#2)
-// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 0 : i32}
-// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 1 : i32}
-// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 2 : i32}
-// CHECK-NEXT: %[[CALLEE_RESULTS:.*]]:3 = call @callee(%[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]])
-// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) -> (memref<2xf32>, i1, memref<5xf32>)
-// CHECK-NEXT: %[[RETURN_TUPLE:.*]] = "test.make_tuple"(%[[CALLEE_RESULTS]]#0, %[[CALLEE_RESULTS]]#1, %[[CALLEE_RESULTS]]#2)
-// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 0 : i32}
-// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 1 : i32}
-// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 2 : i32}
-// CHECK-NEXT: return %[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]]
-
-// -----
-
-// Test case: Testing BufferizeFuncOpConverter and
-// BufferizeReturnOpConverter to see if the return operation matches with the
-// new function signature when there are tuple typed args and results.
-// BufferizeTypeConverter is set to flatten tuple typed arguments. The tuple
-// typed values should be decomposed and composed using get_tuple_element and
-// make_tuple operations of test dialect. Tensor types are converted to Memref.
-// Memref typed function results remain as function results.
-
-// CHECK-LABEL: func @decompose_tuple_typed_function_args_and_results
-func @decompose_tuple_typed_function_args_and_results(%arg0: tuple<i1,f32>, %arg1: tensor<10xf32>, %arg2: tuple<i1, tensor<5xf32>>) -> (tuple<i1, tensor<5xf32>>, tensor<10xf32>, tuple<i1,f32>){
-  return %arg2, %arg1, %arg0 : tuple<i1, tensor<5xf32>>, tensor<10xf32>, tuple<i1,f32>
-}
-// CHECK-SAME: %[[ARG0:.*]]: i1, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: memref<10xf32>, %[[ARG3:.*]]: i1, %[[ARG4:.*]]: memref<5xf32>
-// CHECK-SAME: (i1, memref<5xf32>, memref<10xf32>, i1, f32)
-// CHECK-NEXT: %[[FIRST_TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]])
-// CHECK-NEXT: %[[SECOND_TUPLE:.*]] = "test.make_tuple"(%[[ARG3]], %[[ARG4]])
-// CHECK-NEXT: %[[SECOND_TUPLE_FIRST_ELEM:.*]]  = "test.get_tuple_element"(%[[SECOND_TUPLE]]) {index = 0 : i32}
-// CHECK-NEXT: %[[SECOND_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[SECOND_TUPLE]]) {index = 1 : i32}
-// CHECK-NEXT: %[[FIRST_TUPLE_FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[FIRST_TUPLE]]) {index = 0 : i32}
-// CHECK-NEXT: %[[FIRST_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[FIRST_TUPLE]]) {index = 1 : i32}
-// CHECK-NEXT: return %[[SECOND_TUPLE_FIRST_ELEM]], %[[SECOND_TUPLE_SECOND_ELEM]], %[[ARG2]], %[[FIRST_TUPLE_FIRST_ELEM]], %[[FIRST_TUPLE_SECOND_ELEM]]

mlir/test/Transforms/finalizing-bufferize.mlir

 // RUN: mlir-opt -test-finalizing-bufferize -split-input-file %s | FileCheck %s
 
-// CHECK-LABEL: func @func_signature_conversion
-func @func_signature_conversion(%arg0: tensor<4x8xf32>) {
+// CHECK-LABEL: func @void_function_signature_conversion
+func @void_function_signature_conversion(%arg0: tensor<4x8xf32>) {
     return
 }
-// CHECK: ({{.*}}: memref<4x8xf32>) {
+// CHECK: ({{.*}}: memref<4x8xf32>)
 
 // -----
 
-// Only tensor typed function result should be converted to memref and move to
-// the function arguments list. The other memref function results remain as
-// function results.
-
-// CHECK-LABEL: func @memref_in_function_results
-func @memref_in_function_results(%arg0: tensor<5xf32>, %arg1: memref<10xf32>)
-                            -> (tensor<5xf32>, memref<10xf32>, memref<15xf32>) {
+// 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 = test.tensor_based in(%arg0 : tensor<5xf32>) -> tensor<5xf32>
-  return %1, %arg1, %0 : tensor<5xf32>, memref<10xf32>, memref<15xf32>
+  return %arg2, %1, %arg1, %0, %arg3 :
+   i1, tensor<5xf32>, memref<10xf32>, memref<15xf32>, f16
 }
 //      CHECK: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>,
-// CHECK-SAME: %[[RESULT:.*]]: memref<5xf32>)
-// CHECK-SAME: (memref<10xf32>, memref<15xf32>)
+// CHECK-SAME: %[[ARG2:.*]]: i1, %[[ARG3:.*]]: f16)
+// CHECK-SAME: (i1, memref<5xf32>, memref<10xf32>, memref<15xf32>, f16)
 //      CHECK: %[[FIRST_ALLOC:.*]] = alloc()
 //      CHECK: %[[TENSOR_ALLOC:.*]] = alloc()
-//      CHECK: test.copy(%[[TENSOR_ALLOC]], %[[RESULT]])
-//      CHECK: return %[[ARG1]], %[[FIRST_ALLOC]]
+//      CHECK: return %[[ARG2]], %[[TENSOR_ALLOC]], %[[ARG1]], %[[FIRST_ALLOC]],
+// CHECK-SAME: %[[ARG3]]
 
 // -----
 
 // CHECK-LABEL: func @no_signature_conversion_is_needed
 func @no_signature_conversion_is_needed(%arg0: memref<4x8xf32>) {
   return
 }
-// CHECK: ({{.*}}: memref<4x8xf32>) {
+// 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 @complex_signature_conversion
-func @complex_signature_conversion(%arg0: tensor<4x8xf32>, %arg1: i1,
-                                   %arg2: tensor<5x5xf64>,%arg3: f16) ->
-                                   (i1, tensor<5x5xf64>, f16, tensor<4x8xf32>) {
-    return %arg1, %arg2, %arg3, %arg0 : i1, tensor<5x5xf64>, f16,
-           tensor<4x8xf32>
-}
-//      CHECK: (%[[ARG0:.*]]: memref<4x8xf32>, %[[ARG1:.*]]: i1
-// CHECK-SAME: %[[ARG2:.*]]: memref<5x5xf64>, %[[ARG3:.*]]: f16
-// CHECK-SAME: %[[RESULT1:.*]]: memref<5x5xf64>
-// CHECK-SAME: %[[RESULT2:.*]]: memref<4x8xf32>) -> (i1, f16) {
-// CHECK-NEXT: test.copy(%[[ARG2]], %[[RESULT1]])
-// CHECK-NEXT: test.copy(%[[ARG0]], %[[RESULT2]])
-// CHECK-NEXT: return %[[ARG1]], %[[ARG3]]
+// 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 @non_void_to_void_return_op_converter
-func @non_void_to_void_return_op_converter(%arg0: tensor<4x8xf32>)
-                                           -> tensor<4x8xf32> {
-  return %arg0 : tensor<4x8xf32>
+// CHECK-LABEL: func @func_with_unranked_arg_and_result
+func @func_with_unranked_arg_and_result(%arg0: tensor<*xf32>) -> tensor<*xf32> {
+  return %arg0 : tensor<*xf32>
 }
-//      CHECK: (%[[ARG0:.*]]: [[TYPE:.*]]<[[RANK:.*]]>,
-// CHECK-SAME: %[[RESULT:.*]]: [[TYPE]]<[[RANK]]>) {
-// CHECK-NEXT: test.copy(%[[ARG0]], %[[RESULT]])
-// CHECK-NEXT: return
+// CHECK-SAME: ([[ARG:%.*]]: memref<*xf32>) -> memref<*xf32>
+// CHECK-NEXT: return [[ARG]] : memref<*xf32>
 
 // -----
 
 // CHECK-LABEL: func @func_and_block_signature_conversion
-func @func_and_block_signature_conversion(%arg0 : tensor<2xf32>, %cond : i1,
-                                          %arg1: tensor<4x4xf32>)
-                                          -> tensor<4x4xf32>{
+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,
-// CHECK-SAME: %[[ARG1:.*]]: [[ARG1_TYPE:.*]],
-// CHECK-SAME: %[[RESULT:.*]]: [[RESULT_TYPE:.*]]) {
+//      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: test.copy(%[[ARG1]], %[[RESULT]])
-// CHECK-NEXT: return
-
-// -----
-
-// Test Case: Simple case for checking if BufferizePlacer creates AllocOps
-//            right before TensorBasedOp.
-
-// CHECK-LABEL: func @compute_allocs_position_simple
-func @compute_allocs_position_simple(%cond: i1, %arg0: tensor<2xf32>)
-                                     -> tensor<2xf32>{
-    %0 = test.tensor_based in(%arg0 : tensor<2xf32>) -> tensor<2xf32>
-    %1 = test.tensor_based in(%0 : tensor<2xf32>) -> tensor<2xf32>
-    return %1 : tensor<2xf32>
-}
-//      CHECK: (%{{.*}}: {{.*}}, %[[ARG0:.*]]: memref<2xf32>,
-// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
-// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}} out(%[[FIRST_ALLOC]]
-//      CHECK: %[[SECOND_ALLOC:.*]] = alloc()
-// CHECK-NEXT: test.buffer_based in(%[[FIRST_ALLOC]]{{.*}} out(%[[SECOND_ALLOC]]
-
-// -----
-
-// Test Case: if-else case for checking if BufferizePlacer creates AllocOps
-//            right before TensorBasedOp.
-
-// CHECK-LABEL: func @compute_allocs_position
-func @compute_allocs_position(%cond: i1, %arg0: tensor<2xf32>) -> tensor<2xf32>{
-    %0 = test.tensor_based in(%arg0 : tensor<2xf32>) -> tensor<2xf32>
-    %1 = test.tensor_based in(%0 : tensor<2xf32>) -> tensor<2xf32>
-    cond_br %cond, ^bb1(%arg0, %0: tensor<2xf32>, tensor<2xf32>),
-                   ^bb2(%0, %arg0: tensor<2xf32>, tensor<2xf32>)
-  ^bb1(%arg1 : tensor<2xf32>, %arg2 : tensor<2xf32>):
-    %2 = test.tensor_based in(%arg0 : tensor<2xf32>) -> tensor<2xf32>
-    %3 = test.tensor_based in(%2 : tensor<2xf32>) -> tensor<2xf32>
-    br ^exit(%arg1, %arg2 : tensor<2xf32>, tensor<2xf32>)
-  ^bb2(%arg3 : tensor<2xf32>, %arg4 : tensor<2xf32>):
-    %4 = test.tensor_based in(%arg0 : tensor<2xf32>) -> tensor<2xf32>
-    %5 = test.tensor_based in(%4 : tensor<2xf32>) -> tensor<2xf32>
-    br ^exit(%arg3, %arg4 : tensor<2xf32>, tensor<2xf32>)
-  ^exit(%arg5 : tensor<2xf32>, %arg6 : tensor<2xf32>):
-    %6 = test.tensor_based in(%arg0 : tensor<2xf32>)  -> tensor<2xf32>
-    %7 = test.tensor_based in(%6 : tensor<2xf32>) -> tensor<2xf32>
-    return %7 : tensor<2xf32>
-}
-//      CHECK: (%{{.*}}: {{.*}}, %[[ARG0:.*]]: memref<2xf32>,
-// CHECK-NEXT: %[[ALLOC0:.*]] = alloc()
-// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}} out(%[[ALLOC0]]
-//      CHECK: %[[ALLOC1:.*]] = alloc()
-// CHECK-NEXT: test.buffer_based in(%[[ALLOC0]]{{.*}} out(%[[ALLOC1]]
-//      CHECK: cond_br %{{.*}}, ^[[BB0:.*]]({{.*}}), ^[[BB1:.*]](
-// CHECK-NEXT: ^[[BB0]]
-// CHECK-NEXT: %[[ALLOC2:.*]] = alloc()
-// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}} out(%[[ALLOC2]]
-//      CHECK: %[[ALLOC3:.*]] = alloc()
-// CHECK-NEXT: test.buffer_based in(%[[ALLOC2]]{{.*}} out(%[[ALLOC3]]
-//      CHECK: br ^[[EXIT:.*]]({{.*}})
-// CHECK-NEXT: ^[[BB1]]
-// CHECK-NEXT: %[[ALLOC4:.*]] = alloc()
-// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}} out(%[[ALLOC4]]
-//      CHECK: %[[ALLOC5:.*]] = alloc()
-// CHECK-NEXT: test.buffer_based in(%[[ALLOC4]]{{.*}} out(%[[ALLOC5]]
-//      CHECK: br ^[[EXIT]]
-// CHECK-NEXT: ^[[EXIT]]
-// CHECK-NEXT: %[[ALLOC6:.*]] = alloc()
-// CHECK-NEXT: test.buffer_based in(%[[ARG0]]{{.*}} out(%[[ALLOC6]]
-//      CHECK: %[[ALLOC7:.*]] = alloc()
-// CHECK-NEXT: test.buffer_based in(%[[ALLOC6]]{{.*}} out(%[[ALLOC7]]
-
-// -----
-
-// Test case: Checking BufferizeCallOpConverter and
-// BufferizeFuncOpConverter and BufferizeReturnOpConverter 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.
-
-// CHECK-LABEL: func @callee
-func @callee(%arg1: tensor<5xf32>) -> tensor<5xf32> {
-  %0 = test.tensor_based in(%arg1 : tensor<5xf32>) -> tensor<5xf32>
-  return %0 : tensor<5xf32>
-}
-// CHECK: (%[[CALLEE_ARG:.*]]: memref<5xf32>,
-// CHECK-SAME: %[[CALLEE_RESULT:.*]]: memref<5xf32>)
-// CHECK: %[[ALLOC:.*]] = alloc()
-// CHECK: test.buffer_based
-// CHECK: test.copy(%[[ALLOC]], %[[CALLEE_RESULT]])
-// CHECK: return
-
-// CHECK-LABEL: func @caller
-func @caller(%arg0: tensor<5xf32>) -> tensor<5xf32> {
-  %x = call @callee(%arg0) : (tensor<5xf32>) -> tensor<5xf32>
-  %y = call @callee(%x) : (tensor<5xf32>) -> tensor<5xf32>
-  return %y : tensor<5xf32>
-}
-// CHECK: (%[[CALLER_ARG:.*]]: memref<5xf32>,
-// CHECK-SAME: %[[CALLER_RESULT:.*]]: memref<5xf32>)
-// CHECK: %[[FIRST_ALLOC:.*]] = alloc()
-// CHECK: call @callee(%[[CALLER_ARG]], %[[FIRST_ALLOC]])
-// CHECK: %[[SECOND_ALLOC:.*]] = alloc()
-// CHECK: call @callee(%[[FIRST_ALLOC]], %[[SECOND_ALLOC]])
-// CHECK: test.copy(%[[SECOND_ALLOC]], %[[CALLER_RESULT]])
-// CHECK: return
+// CHECK-NEXT:  return %[[ARG1]] : [[RESULT_TYPE]]
 
 // -----
 
-// Test case: Checking BufferizeCallOpConverter and
-// BufferizeFuncOpConverter and BufferizeReturnOpConverter 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 of %x0 and %y0 and they are passed to the callers' operands
-// list as the output buffers. %x1 and %y1 remain as callers' results.
-
 // 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>,
-// CHECK-SAME: %[[CALLEE_RESULT:.*]]: memref<5xf32>)
-// CHECK-SAME: memref<2xf32>
-//      CHECK: %[[ALLOC:.*]] = alloc()
-//      CHECK: test.copy(%[[CALLEE_ARG]], %[[CALLEE_RESULT]])
-//      CHECK: return %[[ALLOC]]
+// 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> {
-  %x0, %x1 = call @callee(%arg0) : (tensor<5xf32>)
-                                   -> (tensor<5xf32>, memref<2xf32>)
-  %y0, %y1 = call @callee(%x0) : (tensor<5xf32>)
-                                 -> (tensor<5xf32>, memref<2xf32>)
-  return %y0 : tensor<5xf32>
-}
-// CHECK: (%[[CALLER_ARG:.*]]: memref<5xf32>,
-// CHECK-SAME: %[[CALLER_RESULT:.*]]: memref<5xf32>)
-// CHECK: %[[X0:.*]] = alloc()
-// CHECK: %[[X1:.*]] = call @callee(%[[CALLER_ARG]], %[[X0]])
-// CHECK: %[[Y0:.*]] = alloc()
-// CHECK: %[[Y1:.*]] = call @callee(%[[X0]], %[[Y0]])
-// CHECK: test.copy(%[[Y0]], %[[CALLER_RESULT]])
-// CHECK: return
-
-// -----
-
-// CHECK-LABEL: func @func_with_unranked_arg
-func @func_with_unranked_arg(%arg0: tensor<*xf32>) {
-  return
-}
-// CHECK-SAME: ([[ARG:%.*]]: memref<*xf32>)
+  %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
 
 // -----
 
 // Test case: Testing BufferizeCallOpConverter to see if it matches with the
 // signature of the new signature of the callee function when there are tuple
-// typed args and results. BufferizeTypeConverter is set to flatten tuple
-// typed arguments. The tuple typed values should be decomposed and composed
-// using get_tuple_element and make_tuple operations of test dialect. Tensor
-// types are converted to Memref. Memref typed function results are appended to
-// the function arguments list.
+// typed args and results. BufferizeTypeConverter is set to flatten tuple typed
+// arguments. The tuple typed values should be decomposed and composed using
+// get_tuple_element and make_tuple operations of test dialect. Tensor types are
+// converted to Memref. Memref typed function results remain as function
+// results.
 
 // CHECK-LABEL: func @callee
-func @callee(%arg0: tuple<tensor<2xf32>,i1, tensor<5xf32>>)
-             -> (tuple<tensor<2xf32>,i1, tensor<5xf32>>){
+func @callee(%arg0: tuple<tensor<2xf32>,i1, tensor<5xf32>>) -> (tuple<tensor<2xf32>,i1, tensor<5xf32>>){
   return %arg0 : tuple<tensor<2xf32>,i1, tensor<5xf32>>
 }
-// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1,
-// CHECK-SAME: %[[ARG2:.*]]: memref<5xf32>, %[[RESULT0:.*]]: memref<2xf32>,
-// CHECK-SAME: %[[RESULT1:.*]]: memref<5xf32>) -> i1
-// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]],
-// CHECK-SAME: %[[ARG2]])
-// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 0 : i32}
-// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 1 : i32}
-// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 2 : i32}
-// CHECK-NEXT: test.copy(%[[FIRST_ELEM]], %[[RESULT0]])
-// CHECK-NEXT: test.copy(%[[THIRD_ELEM]], %[[RESULT1]])
-// CHECK-NEXT: return %[[SECOND_ELEM]]
+// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1, %[[ARG2:.*]]: memref<5xf32>)
+// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>)
+// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
+// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) {index = 0 : i32}
+// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]]) {index = 1 : i32}
+// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[TUPLE]]) {index = 2 : i32}
+// CHECK-NEXT: return %[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]]
 
 // CHECK-LABEL: func @caller
-func @caller(%arg0: tuple<tensor<2xf32>,i1, tensor<5xf32>>)
-             -> tuple<tensor<2xf32>,i1, tensor<5xf32>>{
-  %x0 = call @callee(%arg0) : (tuple<tensor<2xf32>,i1, tensor<5xf32>>)
-                              -> (tuple<tensor<2xf32>,i1, tensor<5xf32>>)
-  %y0 = call @callee(%x0) : (tuple<tensor<2xf32>,i1, tensor<5xf32>>)
-                            -> (tuple<tensor<2xf32>,i1, tensor<5xf32>>)
+func @caller(%arg0: tuple<tensor<2xf32>,i1, tensor<5xf32>>) -> tuple<tensor<2xf32>,i1, tensor<5xf32>>{
+  %x0 = call @callee(%arg0) : (tuple<tensor<2xf32>,i1, tensor<5xf32>>) -> (tuple<tensor<2xf32>,i1, tensor<5xf32>>)
+  %y0 = call @callee(%x0) : (tuple<tensor<2xf32>,i1, tensor<5xf32>>) -> (tuple<tensor<2xf32>,i1, tensor<5xf32>>)
   return %y0 : tuple<tensor<2xf32>,i1, tensor<5xf32>>
 }
-// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1,
-// CHECK-SAME: %[[ARG2:.*]]: memref<5xf32>, %[[RESULT0:.*]]: memref<2xf32>,
-// CHECK-SAME: %[[RESULT1:.*]]: memref<5xf32>) -> i1
-// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]],
-// CHECK-SAME: %[[ARG2]])
-// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 0 : i32}
-// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 1 : i32}
-// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 2 : i32}
-// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
-// CHECK-NEXT: %[[SECOND_ALLOC:.*]] = alloc()
-// CHECK-NEXT: %[[CALLEE_RESULT:.*]] = call @callee(%[[FIRST_ELEM]],
-// CHECK-SAME: %[[SECOND_ELEM]], %[[THIRD_ELEM]], %[[FIRST_ALLOC]],
-// CHECK-SAME: %[[SECOND_ALLOC]])
-// CHECK-SAME: (memref<2xf32>, i1,
-// CHECK-SAME: memref<5xf32>, memref<2xf32>, memref<5xf32>) -> i1
-// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[FIRST_ALLOC]],
-// CHECK-SAME: %[[CALLEE_RESULT]], %[[SECOND_ALLOC]])
-// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 0 : i32}
-// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 1 : i32}
-// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 2 : i32}
-// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
-// CHECK-NEXT: %[[SECOND_ALLOC:.*]] = alloc()
-// CHECK-NEXT: %[[CALLEE_RESULT:.*]] = call @callee(%[[FIRST_ELEM]],
-// CHECK-SAME: %[[SECOND_ELEM]], %[[THIRD_ELEM]], %[[FIRST_ALLOC]],
-// CHECK-SAME: %[[SECOND_ALLOC]])
-// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>, memref<2xf32>, memref<5xf32>)
-// CHECK-SAME: i1
-// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[FIRST_ALLOC]],
-// CHECK-SAME: %[[CALLEE_RESULT]], %[[SECOND_ALLOC]])
-// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 0 : i32}
-// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 1 : i32}
-// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[TUPLE]])
-// CHECK-SAME: {index = 2 : i32}
-// CHECK-NEXT: test.copy(%[[FIRST_ELEM]], %[[RESULT0]])
-// CHECK-NEXT: test.copy(%[[THIRD_ELEM]], %[[RESULT1]])
-// CHECK-NEXT: return %[[SECOND_ELEM]]
+// CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[ARG1:.*]]: i1, %[[ARG2:.*]]: memref<5xf32>)
+// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>)
+// CHECK-NEXT: %[[ARG_TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
+// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 0 : i32}
+// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 1 : i32}
+// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 2 : i32}
+// CHECK-NEXT: %[[CALLEE_RESULTS:.*]]:3 = call @callee(%[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]])
+// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) -> (memref<2xf32>, i1, memref<5xf32>)
+// CHECK-NEXT: %[[RESULT_TUPLE:.*]] = "test.make_tuple"(%[[CALLEE_RESULTS]]#0, %[[CALLEE_RESULTS]]#1, %[[CALLEE_RESULTS]]#2)
+// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 0 : i32}
+// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 1 : i32}
+// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[RESULT_TUPLE]]) {index = 2 : i32}
+// CHECK-NEXT: %[[CALLEE_RESULTS:.*]]:3 = call @callee(%[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]])
+// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>) -> (memref<2xf32>, i1, memref<5xf32>)
+// CHECK-NEXT: %[[RETURN_TUPLE:.*]] = "test.make_tuple"(%[[CALLEE_RESULTS]]#0, %[[CALLEE_RESULTS]]#1, %[[CALLEE_RESULTS]]#2)
+// CHECK-NEXT: %[[FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 0 : i32}
+// CHECK-NEXT: %[[SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 1 : i32}
+// CHECK-NEXT: %[[THIRD_ELEM:.*]]  = "test.get_tuple_element"(%[[RETURN_TUPLE]]) {index = 2 : i32}
+// CHECK-NEXT: return %[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]]
 
 // -----
 
-// Test case: Testing BufferizeFuncOpConverter and BufferizeReturnOpConverter
-// to see if the return operation matches with the new function signature when
-// there are tuple typed args and results. BufferizeTypeConverter is set to
-// flatten tuple typed arguments. The tuple typed values should be decomposed
-// and composed using get_tuple_element and make_tuple operations of test
-// dialect. Tensor types are converted to Memref. Memref typed function results
-// are appended to the function arguments list.
+// Test case: Testing BufferizeFuncOpConverter and
+// BufferizeReturnOpConverter to see if the return operation matches with the
+// new function signature when there are tuple typed args and results.
+// BufferizeTypeConverter is set to flatten tuple typed arguments. The tuple
+// typed values should be decomposed and composed using get_tuple_element and
+// make_tuple operations of test dialect. Tensor types are converted to Memref.
+// Memref typed function results remain as function results.
 
 // CHECK-LABEL: func @decompose_tuple_typed_function_args_and_results
-func @decompose_tuple_typed_function_args_and_results(%arg0: tuple<i1,f32>,
-                                                      %arg1: tensor<10xf32>,
-                                                      %arg2: tuple<i1,
-                                                             tensor<5xf32>>)
-                                                      -> (tuple<i1,
-                                                                tensor<5xf32>>,
-                                                      tensor<10xf32>,
-                                                            tuple<i1,f32>){
-  return %arg2, %arg1, %arg0 : tuple<i1, tensor<5xf32>>, tensor<10xf32>,
-                                    tuple<i1,f32>
-}
-// CHECK-SAME: %[[ARG0:.*]]: i1, %[[ARG1:.*]]: f32,
-// CHECK-SAME: %[[ARG2:.*]]: memref<10xf32>, %[[ARG3:.*]]: i1,
-// CHECK-SAME: %[[ARG4:.*]]: memref<5xf32>, %[[RESULT0:.*]]: memref<5xf32>,
-// CHECK-SAME: %[[RESULT1:.*]]: memref<10xf32>
-// CHECK-SAME: (i1, i1, f32)
+func @decompose_tuple_typed_function_args_and_results(%arg0: tuple<i1,f32>, %arg1: tensor<10xf32>, %arg2: tuple<i1, tensor<5xf32>>) -> (tuple<i1, tensor<5xf32>>, tensor<10xf32>, tuple<i1,f32>){
+  return %arg2, %arg1, %arg0 : tuple<i1, tensor<5xf32>>, tensor<10xf32>, tuple<i1,f32>
+}
+// CHECK-SAME: %[[ARG0:.*]]: i1, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: memref<10xf32>, %[[ARG3:.*]]: i1, %[[ARG4:.*]]: memref<5xf32>
+// CHECK-SAME: (i1, memref<5xf32>, memref<10xf32>, i1, f32)
 // CHECK-NEXT: %[[FIRST_TUPLE:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]])
 // CHECK-NEXT: %[[SECOND_TUPLE:.*]] = "test.make_tuple"(%[[ARG3]], %[[ARG4]])
-// CHECK-NEXT: %[[SECOND_TUPLE_FIRST_ELEM:.*]]  = "test.get_tuple_element"
-// CHECK-SAME: (%[[SECOND_TUPLE]]) {index = 0 : i32}
-// CHECK-NEXT: %[[SECOND_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element"
-// CHECK-SAME: (%[[SECOND_TUPLE]]) {index = 1 : i32}
-// CHECK-NEXT: %[[FIRST_TUPLE_FIRST_ELEM:.*]] = "test.get_tuple_element"
-// CHECK-SAME: (%[[FIRST_TUPLE]]) {index = 0 : i32}
-// CHECK-NEXT: %[[FIRST_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element"
-// CHECK-SAME: (%[[FIRST_TUPLE]]) {index = 1 : i32}
-// CHECK-NEXT: test.copy(%[[SECOND_TUPLE_SECOND_ELEM]], %[[RESULT0]])
-// CHECK-NEXT: test.copy(%[[ARG2]], %[[RESULT1]])
-// CHECK-NEXT: return %[[SECOND_TUPLE_FIRST_ELEM]], %[[FIRST_TUPLE_FIRST_ELEM]],
-// CHECK-SAME: %[[FIRST_TUPLE_SECOND_ELEM]]
+// CHECK-NEXT: %[[SECOND_TUPLE_FIRST_ELEM:.*]]  = "test.get_tuple_element"(%[[SECOND_TUPLE]]) {index = 0 : i32}
+// CHECK-NEXT: %[[SECOND_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[SECOND_TUPLE]]) {index = 1 : i32}
+// CHECK-NEXT: %[[FIRST_TUPLE_FIRST_ELEM:.*]] = "test.get_tuple_element"(%[[FIRST_TUPLE]]) {index = 0 : i32}
+// CHECK-NEXT: %[[FIRST_TUPLE_SECOND_ELEM:.*]] = "test.get_tuple_element"(%[[FIRST_TUPLE]]) {index = 1 : i32}
+// CHECK-NEXT: return %[[SECOND_TUPLE_FIRST_ELEM]], %[[SECOND_TUPLE_SECOND_ELEM]], %[[ARG2]], %[[FIRST_TUPLE_FIRST_ELEM]], %[[FIRST_TUPLE_SECOND_ELEM]]

mlir/test/lib/Transforms/TestFinalizingBufferize.cpp

 /// involves rewriting funcs (including block arguments of the contained
 /// region), calls, and returns. The unique property of finalizing bufferization
 /// passes is that they cannot be done via a local transformation with suitable
 /// materializations to ensure composability (as other bufferization passes do).
 /// For example, if a call is rewritten, the callee needs to be rewritten
 /// otherwise the IR will end up invalid. Thus, finalizing bufferization passes
 /// require an atomic change to the entire program (e.g. the whole module).
 ///
-/// `allowMemrefFunctionResults` informs the buffer finalization policy to allow
-/// functions that have memref typed results. Patterns involved with converting
-/// func/call/return respect the finalization policy to ensure a consistent
-/// atomic conversion of the entire module. `allowMemrefFunctionResults` also
-/// allows memref typed results to escape from the deallocation.
-///
 /// TODO: Split out BufferizeFinalizationPolicy from BufferizeTypeConverter.
-template <bool allowMemrefFunctionResults>
 struct TestFinalizingBufferizePass
-    : mlir::PassWrapper<TestFinalizingBufferizePass<allowMemrefFunctionResults>,
-                        OperationPass<ModuleOp>> {
+    : mlir::PassWrapper<TestFinalizingBufferizePass, OperationPass<ModuleOp>> {
 
   /// Converts tensor based test operations to buffer based ones using
   /// bufferize.
   class TensorBasedOpConverter
       : public BufferizeOpConversionPattern<test::TensorBasedOp> {
   public:
     using BufferizeOpConversionPattern<
         test::TensorBasedOp>::BufferizeOpConversionPattern;
target.addDynamicallyLegalOp<mlir::CallOp>(
         [&](mlir::CallOp callOp) { return converter.isLegal(callOp); });
 
     // Mark the function whose arguments are in tensor-type illegal.
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
       return converter.isSignatureLegal(funcOp.getType()) &&
              converter.isLegal(&funcOp.getBody());
     });
 
-    auto kind = allowMemrefFunctionResults
-                    ? BufferizeTypeConverter::KeepAsFunctionResult
-                    : BufferizeTypeConverter::AppendToArgumentsList;
-    converter.setResultConversionKind<RankedTensorType, MemRefType>(kind);
-    converter.setResultConversionKind<UnrankedTensorType, UnrankedMemRefType>(
-        kind);
-
     converter.addDecomposeTypeConversion(
         [](TupleType tupleType, SmallVectorImpl<Type> &types) {
           tupleType.getFlattenedTypes(types);
           return success();
         });
 
     converter.addArgumentMaterialization(
         [](OpBuilder &builder, TupleType resultType, ValueRange inputs,
if (failed(applyFullConversion(this->getOperation(), target,
       this->signalPassFailure();
   };
 };
 } // end anonymous namespace
 
 namespace mlir {
 namespace test {
 void registerTestFinalizingBufferizePass() {
-  PassRegistration<
-      TestFinalizingBufferizePass</*allowMemrefFunctionResults=*/false>>(
+  PassRegistration<TestFinalizingBufferizePass>(
       "test-finalizing-bufferize", "Tests finalizing bufferize conversions");
 }
-
-void registerTestPreparationPassWithAllowedMemrefResults() {
-  PassRegistration<
-      TestFinalizingBufferizePass</*allowMemrefFunctionResults=*/true>>(
-      "test-finalizing-bufferize-with-allowed-memref-results",
-      "Tests finalizing buffierize conversions, allowing functions to have "
-      "memref typed results.");
-}
 } // namespace test
 } // namespace mlir

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

#endif
   test::registerTestLinalgTransforms();
   test::registerTestLivenessPass();
   test::registerTestLoopFusion();
   test::registerTestLoopMappingPass();
   test::registerTestLoopUnrollingPass();
   test::registerTestMemRefDependenceCheck();
   test::registerTestMemRefStrideCalculation();
   test::registerTestOpaqueLoc();
-  test::registerTestPreparationPassWithAllowedMemrefResults();
   test::registerTestRecursiveTypesPass();
   test::registerTestSCFUtilsPass();
   test::registerTestVectorConversions();
 }
 #endif
 
 int main(int argc, char **argv) {
   registerAllPasses();