[mlir] Extend BufferAssignmentTypeConverter with result conversion callbacks

[dfki-ehna]

Summary

In this PR, the users of BufferPlacement can configure
BufferAssginmentTypeConverter. These new configurations would give the user more
freedom in the process of converting function signature, and return and call
operation conversions.

These are the new features:

• Accepting callback functions for decomposing types (i.e. 1 to N type conversion such as unpacking tuple types).
• Defining ResultConversionKind for specifying whether a function result with a certain type should be appended to the function arguments list or should be kept as function result. (Usage: converter.setResultConversionKind<MemRefType>(AppendToArgumentList))
• Accepting callback functions for composing or decomposing values (i.e. N to 1 and 1 to N value conversion).

Input:

func @test(%arg0: tuple<tensor<2xf32>,i1>) -> (tuple<tensor<2xf32>,i1>){
  return %arg0 : tuple<tensor<2xf32>,i1>
}

Output with [converter.setResultConversionKind<MemRefType>(AppendToArgumentList)]:

func @test(%arg0: memref<2xf32>, %arg1: i1, %arg2: memref<2xf32>) -> i1 {
  %0 = "test.make_tuple"(%arg0, %arg1) : (memref<2xf32>, i1) -> tuple<memref<2xf32>, i1>
  %1 = "test.get_tuple_element"(%0) {index = 0 : i32} : (tuple<memref<2xf32>, i1>) -> memref<2xf32>
  %2 = "test.get_tuple_element"(%0) {index = 1 : i32} : (tuple<memref<2xf32>, i1>) -> i1
  linalg.copy(%1, %arg2) : memref<2xf32>, memref<2xf32>
  return %2 : i1
}

Output with [converter.setResultConversionKind<MemRefType>(KeepAsFunctionResult)]:

func @test(%arg0: memref<2xf32>, %arg1: i1) -> (memref<2xf32>, i1) {
  %0 = "test.make_tuple"(%arg0, %arg1) : (memref<2xf32>, i1) -> tuple<memref<2xf32>, i1>
  %1 = "test.get_tuple_element"(%0) {index = 0 : i32} : (tuple<memref<2xf32>, i1>) -> memref<2xf32>
  %2 = "test.get_tuple_element"(%0) {index = 1 : i32} : (tuple<memref<2xf32>, i1>) -> i1
  return %1, %2 : memref<2xf32>, i1
}

[herhut]

Just some quick comments.

mlir/include/mlir/Transforms/BufferPlacement.h

93

Currently this could just be a map from Type to ResultConversionKind. Callbacks only help if the user can provide them and hence have some extra logic in them. Maybe by also passing the originating type?

219–243

This should be an assert, best in the super class' constructor.

257

Here reporting an error instead would be nicer. You could also append a note on the error showing the function signature,

mlir/lib/Transforms/BufferPlacement.cpp

872

If you had the original and new type available here, you could still differentiate between a rewritten memref and a newly created memref.

[pifon2a]

mlir/include/mlir/Transforms/BufferPlacement.h

93

I agree with Stephan. It would be nice to make it a map.
If we need more flexibility in future, then we would need to provide a method that allows to set a map from functor(Type) -> bool to ResultConversionKind so that users would what and how to convert. For now, a map would be sufficient.

[bondhugula]

The functionality added looks great! Thanks. I'll shortly submit a review.

[bondhugula]

mlir/include/mlir/Transforms/BufferPlacement.h

224

You can move this declaration to right under the comment below.

226

Nit: " ... depending on whether they need a ..."

254–255

int -> unsigned ?

[bondhugula]

mlir/lib/Transforms/BufferPlacement.cpp

734–735

Nit: Can fold these two into one line.

837

auto -> const std::pair<unsigned, Value> &

837

I think you have more copy here due to lack of a const ref.

841–843

Comments for the member variables please.

878

Avoid trailing comment - please move to next line.

895

Typo: it -> its

900

Use i = 0, e = ...; i < e; ... form to avoid repeated evaluation.

903–907

Nit: Use braces all blocks since you need it for the else block below.

mlir/test/lib/Dialect/Test/TestOps.td

1625

Nit: the i-th element -> a specified element

mlir/test/lib/Transforms/TestBufferPlacement.cpp

184–186

Aren't you already in the mlir namespace? Drop mlir:: ?

185

Avoid repeated evaluation of .size() - use i = 0, e = ... form.

[dfki-ehna]

Resolve comments, change setResultConversionKind to a mapping (i.e. setResultConversionKind<RankedTensorType, MemRefType>(AppendToArguementList)), and introduce addDecomposeTypeConversion.

[bondhugula]

This looks really good - thanks!

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

122–133

Could you add a line or two here and further below to summarize what this is testing? (although obvious to you at this point).

163

Please name these functions a bit more descriptively. Here and above.

[bondhugula]

mlir/include/mlir/Transforms/BufferPlacement.h

107–108

You could use llvm::is_one_of instead.

[dfki-ehna]

Resolving comments

[mehdi_amini]

mlir/include/mlir/Transforms/BufferPlacement.h

254

Seems like the error message isn't terminated?

[dfki-ehna]

Resolve the minor comment before landing

[antiagainst]

FYI: this change fails the following tests:

Failed Tests (3):
  MLIR :: Dialect/Linalg/tensors-to-buffers.mlir
  MLIR :: Transforms/buffer-placement-preparation-allowed-memref-results.mlir
  MLIR :: Transforms/buffer-placement-preparation.mlir

See https://buildkite.com/mlir/mlir-core/builds/7580#ad1609b5-eedd-4d04-9fd0-361402d40f36.

I'll revert this. Please resubmit after fixing the tests. :)

[antiagainst]

Reverted as https://github.com/llvm/llvm-project/commit/1b88bbf5eb80b38a4dee129df969d5632993fdd1.

[dfki-ehna]

Reopen to resolve the failure.

[dfki-ehna]

Resolve the failure

[rupprecht]

When trying to integrate this downstream, I'm having trouble figuring out what to do with this block: https://github.com/tensorflow/mlir-hlo/blob/master/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc#L441

BufferAssignmentTypeConverter converter;
...
if (results_escape_function) {
  populateWithBufferAssignmentOpConversionPatterns<
      mlir::ReturnOp, mlir::ReturnOp, lmhlo::CopyOp,
      /*allowMemrefFunctionResults=*/true>(&context, &bufferAssignment,
                                           &converter, &patterns);
} else {
  populateWithBufferAssignmentOpConversionPatterns<
      mlir::ReturnOp, mlir::ReturnOp, lmhlo::CopyOp,
      /*allowMemrefFunctionResults=*/false>(&context, &bufferAssignment,
                                            &converter, &patterns);
}

As the allowMemrefFunctionResults argument went away, I tried using setResultConversionKind as suggested to have the equivalent functionality, however I'm still seeing test failures:

BufferAssignmentTypeConverter converter;
auto kind = results_escape_function
                ? BufferAssignmentTypeConverter::KeepAsFunctionResult
                : BufferAssignmentTypeConverter::AppendToArgumentsList;
converter.setResultConversionKind<RankedTensorType, MemRefType>(kind);
converter.setResultConversionKind<UnrankedTensorType, UnrankedMemRefType>(
    kind);
...
populateWithBufferAssignmentOpConversionPatterns<
    mlir::ReturnOp, mlir::ReturnOp, lmhlo::CopyOp>(
    &context, &bufferAssignment, &converter, &patterns);

Is there anything wrong with the fix applied? (cc @antiagainst)

[dfki-ehna]

@rupprecht I am investigating the downstream issue.

mlir/include/mlir/Transforms/BufferPlacement.h

public:
   /// Computes the actual position to place allocs for the given result.
   OpBuilder::InsertPoint computeAllocPosition(OpResult result);
 
 private:
   /// The operation this analysis was constructed from.
   Operation *operation;
 };
 
+/// 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 BufferAssignmentTypeConverter : 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 };
+
+  BufferAssignmentTypeConverter();
+
+  /// 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> &);
+
+  /// This method tries to decompose a type using provided decompose callback
+  /// functions. If it is unable to do so, the original type is returned.
+  void tryDecomposeType(Type, SmallVectorImpl<Type> &);
+
+  /// This method registers a callback function that will be called to decompose
+  /// a value of a certain type into several values.
+  template <typename FnT,
+            typename T = typename llvm::function_traits<FnT>::template arg_t<2>>
+  void addDecomposeValueConversion(FnT &&callback) {
+    decomposeValueConversions.emplace_back(
+        wrapDecomposeValueConversionCallback<T>(std::forward<FnT>(callback)));
+  }
+
+  /// This method registers a callback function that will be called to decompose
+  /// a type into several types.
+  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(builder, loc, derivedType, value, newValues);
+      return llvm::None;
+    };
+  }
+
+  /// Generate a wrapper for the given decompose type conversion callback.
+  template <typename T, typename FnT>
+  DecomposeTypeConversionCallFn
+  wrapDecomposeTypeConversionCallback(FnT &&callback) {
+    return [callback = std::forward<FnT>(callback)](
+               Type type,
+               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;
+};
+
 /// Helper conversion pattern that encapsulates a BufferAssignmentPlacer
 /// instance. Sample usage:
 /// class CustomConversionPattern : public
 ///     BufferAssignmentOpConversionPattern<MyOpT>
 /// {
 ///   ... matchAndRewrite(...) {
 ///     -> Access stored BufferAssignmentPlacer
 ///     bufferAssignment->computeAllocPosition(resultOp);
 ///   }
 /// };
 template <typename SourceOp>
 class BufferAssignmentOpConversionPattern
     : public OpConversionPattern<SourceOp> {
 public:
   explicit BufferAssignmentOpConversionPattern(
       MLIRContext *context, BufferAssignmentPlacer *bufferAssignment = nullptr,
-      TypeConverter *converter = nullptr, PatternBenefit benefit = 1)
+      BufferAssignmentTypeConverter *converter = nullptr,
+      PatternBenefit benefit = 1)
       : OpConversionPattern<SourceOp>(context, benefit),
-        bufferAssignment(bufferAssignment), converter(converter) {}
+        bufferAssignment(bufferAssignment), converter(converter) {
+    assert(converter && "The type converter has not been defined");
+  }
 
 protected:
   BufferAssignmentPlacer *bufferAssignment;
-  TypeConverter *converter;
-};
-
-/// 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 BufferAssignmentTypeConverter : public TypeConverter {
-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);
+  BufferAssignmentTypeConverter *converter;
 };
 
-namespace detail {
-
-/// 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>
+/// Converts the signature of the function using BufferAssignmentTypeConverter.
+/// Each result type of the function is kept as a function result or appended to
+/// the function arguments list based on ResultConversionKind for the converted
+/// result type.
 class BufferAssignmentFuncOpConverter
     : public BufferAssignmentOpConversionPattern<FuncOp> {
 public:
   using BufferAssignmentOpConversionPattern<
       FuncOp>::BufferAssignmentOpConversionPattern;
 
   /// Performs the actual signature rewriting step.
-  LogicalResult
-  matchAndRewrite(mlir::FuncOp funcOp, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const final {
-    if (!converter)
-      return funcOp.emitError("The type converter has not been defined for "
-                              "BufferAssignmentFuncOpConverter");
-    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);
-    }
-    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();
-  }
+  LogicalResult matchAndRewrite(mlir::FuncOp, ArrayRef<Value>,
+                                ConversionPatternRewriter &) const;
 };
 
 /// Rewrites the `ReturnOp` to conform with the changed function signature.
-/// 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.
+/// Operands that correspond to return values and their types have been set to
+/// AppendToArgumentsList are dropped. In their place, a corresponding copy
+/// operation from the operand to the target function argument is inserted.
 template <typename ReturnOpSourceTy, typename ReturnOpTargetTy,
-          typename CopyOpTy, bool allowMemrefFunctionResults>
+          typename CopyOpTy>
 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();
-    }
+    Location loc = returnOp.getLoc();
 
-    // 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(
-              operand.value().getType(),
-              returnOp.getOperand(operand.index()).getType()))
-        needCopyOperands.push_back(operand.value());
+    // 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 == BufferAssignmentTypeConverter::KeepAsFunctionResult)
+          newOperands.push_back(value.value());
         else
-        newOperands.push_back(operand.value());
+          // kind = BufferAssignmentTypeConverter::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();
-
-    // Find the index of the first destination buffer.
-    assert(needCopyOperands.size() <= numFuncArgs &&
-           "The number of operands of return operation is more than the "
-           "number of function arguments.");
+    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) {
-      // Insert a `CopyOp` for each converted memref-type operand.
-      rewriter.create<CopyOpTy>(returnOp.getLoc(), operand,
+      rewriter.create<CopyOpTy>(loc, operand,
                                 entryBlock.getArgument(destArgNum));
       ++destArgNum;
     }
-
-    // Insert the new target Return operation.
     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
-/// 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 through 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>
+/// BufferAssignmentFuncOpConverter.
 class BufferAssignmentCallOpConverter
     : public BufferAssignmentOpConversionPattern<CallOp> {
 public:
   using BufferAssignmentOpConversionPattern<
       CallOp>::BufferAssignmentOpConversionPattern;
 
-  LogicalResult
-  matchAndRewrite(CallOp callOp, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const final {
-    if (!converter)
-      return callOp.emitError("The type converter has not been defined for "
-                              "BufferAssignmentCallOpConverter");
-    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();
-  }
+  /// Performs the actual rewriting step.
+  LogicalResult matchAndRewrite(CallOp, ArrayRef<Value>,
+                                ConversionPatternRewriter &) const;
 };
-} // end namespace detail
 
 /// Populates `patterns` with the conversion patterns of buffer
 /// assignment.
 template <typename ReturnOpSourceTy, typename ReturnOpTargetTy,
-          typename CopyOpTy, bool allowMemrefFunctionResults>
+          typename CopyOpTy>
 static void populateWithBufferAssignmentOpConversionPatterns(
     MLIRContext *context, BufferAssignmentPlacer *placer,
-    TypeConverter *converter, OwningRewritePatternList *patterns) {
+    BufferAssignmentTypeConverter *converter,
+    OwningRewritePatternList *patterns) {
   // clang-format off
   patterns->insert<
-    detail::BufferAssignmentCallOpConverter<allowMemrefFunctionResults>,
-    detail::BufferAssignmentFuncOpConverter<allowMemrefFunctionResults>,
-    detail::BufferAssignmentReturnOpConverter
-      <ReturnOpSourceTy, ReturnOpTargetTy, CopyOpTy, allowMemrefFunctionResults>
+    BufferAssignmentCallOpConverter,
+    BufferAssignmentFuncOpConverter,
+    BufferAssignmentReturnOpConverter
+      <ReturnOpSourceTy, ReturnOpTargetTy, CopyOpTy>
   >(context, placer, converter);
   // clang-format on
 }
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_BUFFERPLACEMENT_H

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

matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
     return success();
   }
 };
 
 /// 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) {
+    BufferAssignmentTypeConverter *converter,
+    OwningRewritePatternList *patterns) {
   populateWithBufferAssignmentOpConversionPatterns<
-      mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp,
-      /*allowMemrefFunctionResults=*/false>(context, placer, converter,
-                                            patterns);
+      mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>(context, placer,
+                                                      converter, patterns);
   patterns->insert<GenericOpConverter>(context, placer, converter);
 }
 
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
 struct ConvertLinalgOnTensorsToBuffers
     : public LinalgOnTensorsToBuffersBase<ConvertLinalgOnTensorsToBuffers> {
   void runOnOperation() override {
void runOnOperation() override {
     // Mark the function operation illegal as long as an argument is tensor.
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
       return converter.isSignatureLegal(funcOp.getType()) &&
              llvm::none_of(funcOp.getType().getResults(),
                            [&](Type type) { return type.isa<MemRefType>(); }) &&
              converter.isLegal(&funcOp.getBody());
     });
 
+    converter.setResultConversionKind<RankedTensorType, MemRefType>(
+        BufferAssignmentTypeConverter::AppendToArgumentsList);
+
     // Walk over all the functions to apply buffer assignment.
     getOperation().walk([&](FuncOp function) -> WalkResult {
       OwningRewritePatternList patterns;
       BufferAssignmentPlacer placer(function);
       populateConvertLinalgOnTensorsToBuffersPattern(&context, &placer,
                                                      &converter, &patterns);
 
       // Applying full conversion

mlir/lib/Transforms/BufferPlacement.cpp

addConversion([](RankedTensorType type) {
     return (Type)MemRefType::get(type.getShape(), type.getElementType());
   });
   // Convert UnrankedTensorType to UnrankedMemRefType.
   addConversion([](UnrankedTensorType type) {
     return (Type)UnrankedMemRefType::get(type.getElementType(), 0);
   });
 }
 
-/// Checks if `type` has been converted from non-memref type to memref.
-bool BufferAssignmentTypeConverter::isConvertedMemref(Type type, Type before) {
-  return type.isa<BaseMemRefType>() && !before.isa<BaseMemRefType>();
+/// 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 BufferAssignmentTypeConverter::tryDecomposeValue(
+    OpBuilder &builder, Location loc, Type type, Value value,
+    SmallVectorImpl<Value> &results) {
+  for (auto conversion : decomposeValueConversions)
+    if (conversion(builder, loc, type, value, results) != llvm::None)
+      return;
+  results.push_back(value);
+}
+
+/// This method tries to decompose a type using provided decompose callback
+/// functions. If it is unable to do so, the original type is returned.
+void BufferAssignmentTypeConverter::tryDecomposeType(
+    Type type, SmallVectorImpl<Type> &types) {
+  for (auto conversion : decomposeTypeConversions)
+    if (conversion(type, types) != llvm::None)
+      return;
+  types.push_back(type);
+}
+
+/// This method returns ResultConversionKind for the input type.
+BufferAssignmentTypeConverter::ResultConversionKind
+BufferAssignmentTypeConverter::getResultConversionKind(Type origin,
+                                                       Type converted) {
+  for (auto conversion : resultTypeConversions) {
+    auto res = conversion(origin, converted);
+    if (res != llvm::None)
+      return res.getValue();
+  }
+  return KeepAsFunctionResult;
+}
+
+//===----------------------------------------------------------------------===//
+// BufferAssignmentFuncOpConverter
+//===----------------------------------------------------------------------===//
+
+/// Performs the actual function signature rewriting step.
+LogicalResult BufferAssignmentFuncOpConverter::matchAndRewrite(
+    mlir::FuncOp funcOp, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  auto funcType = funcOp.getType();
+
+  // Convert function arguments using the provided TypeConverter.
+  TypeConverter::SignatureConversion conversion(funcType.getNumInputs());
+  for (auto argType : llvm::enumerate(funcType.getInputs())) {
+    SmallVector<Type, 2> decomposedTypes, convertedTypes;
+    converter->tryDecomposeType(argType.value(), decomposedTypes);
+    converter->convertTypes(decomposedTypes, convertedTypes);
+    conversion.addInputs(argType.index(), convertedTypes);
+  }
+
+  // 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 == BufferAssignmentTypeConverter::AppendToArgumentsList)
+        conversion.addInputs(converted);
+      else
+        // kind = BufferAssignmentTypeConverter::KeepAsFunctionResult
+        newResultTypes.push_back(converted);
+    }
+  }
+
+  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();
+}
+
+//===----------------------------------------------------------------------===//
+// BufferAssignmentCallOpConverter
+//===----------------------------------------------------------------------===//
+
+/// Performs the actual rewriting step.
+LogicalResult BufferAssignmentCallOpConverter::matchAndRewrite(
+    CallOp callOp, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+
+  // 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 ResultMapping {
+  public:
+    ResultMapping() { 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 indicies 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.
+      llvm::for_each(
+          toIndicesMapping, [&](const std::pair<unsigned, unsigned> &entry) {
+            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.
+      llvm::for_each(res, [&](const std::pair<unsigned, Value> &entry) {
+        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;
+  };
+
+  Location loc = callOp.getLoc();
+  OpBuilder builder(callOp);
+  SmallVector<Value, 2> newOperands;
+
+  // 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;
+    this->converter->tryDecomposeValue(builder, loc, operand.getType(), operand,
+                                       values);
+    newOperands.append(values.begin(), values.end());
+  }
+
+  // Create the new result types for the new `CallOp` and a mapping from the old
+  // result to new value(s).
+  SmallVector<Type, 2> newResultTypes;
+  SmallVector<ResultMapping, 4> mappings;
+  mappings.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()];
+    for (Type origin : originTypes) {
+      Type converted = converter->convertType(origin);
+      auto kind = converter->getResultConversionKind(origin, converted);
+      if (kind == BufferAssignmentTypeConverter::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 = BufferAssignmentTypeConverter::AppendToArgumentsList
+        OpBuilder::InsertionGuard guard(rewriter);
+        rewriter.restoreInsertionPoint(
+            bufferAssignment->computeAllocPosition(result.value()));
+        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);
+      }
+    }
+  }
+
+  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);
+    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);
+      replacedValues.push_back(packed);
+    }
+  }
+  rewriter.replaceOp(callOp, replacedValues);
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // BufferPlacementPass construction
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> mlir::createBufferPlacementPass() {
   return std::make_unique<BufferPlacementPass>();
 }

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

func @caller(%arg0: tensor<5xf32>) -> tensor<5xf32> {
   %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 BufferAssginmnetCallOpConverter to see if it matches with the
+// signature of the new signature of the callee function when there are tuple typed
+// args and results. BufferAssginmentTypeConverter 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 BufferAssginmnetFuncOpConverter and
+// BufferAssginmentReturnOpConverter to see if the return operation matches with
+// the new function signature when there are tuple typed args and results.
+// BufferAssginmentTypeConverter 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/buffer-placement-preparation.mlir

 // CHECK: (%[[CALLER_ARG:.*]]: memref<5xf32>, %[[CALLER_RESULT:.*]]: memref<5xf32>)
 // CHECK: %[[X0:.*]] = alloc()
 // CHECK: %[[X1:.*]] = call @callee(%[[CALLER_ARG]], %[[X0]])
 // CHECK: %[[Y0:.*]] = alloc()
 // CHECK: %[[Y1:.*]] = call @callee(%[[X0]], %[[Y0]])
 // CHECK: linalg.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>)
+
+// -----
+
+// Test case: Testing BufferAssginmnetCallOpConverter to see if it matches with the
+// signature of the new signature of the callee function when there are tuple typed
+// args and results. BufferAssginmentTypeConverter 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.
+
+// 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>, %[[RESULT0:.*]]: memref<2xf32>, %[[RESULT1:.*]]: memref<5xf32>)
+// CHECK-SAME: i1
+// 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: linalg.copy(%[[FIRST_ELEM]], %[[RESULT0]])
+// CHECK-NEXT: linalg.copy(%[[THIRD_ELEM]], %[[RESULT1]])
+// CHECK-NEXT: return %[[SECOND_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>, %[[RESULT0:.*]]: memref<2xf32>, %[[RESULT1:.*]]: memref<5xf32>)
+// CHECK-SAME: i1
+// 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: %[[FIRST_ALLOC:.*]] = alloc()
+// CHECK-NEXT: %[[SECOND_ALLOC:.*]] = alloc()
+// CHECK-NEXT: %[[CALLEE_RESULT:.*]] = call @callee(%[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]], %[[FIRST_ALLOC]], %[[SECOND_ALLOC]])
+// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>, memref<2xf32>, memref<5xf32>) -> i1
+// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[FIRST_ALLOC]], %[[CALLEE_RESULT]], %[[SECOND_ALLOC]])
+// 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: %[[FIRST_ALLOC:.*]] = alloc()
+// CHECK-NEXT: %[[SECOND_ALLOC:.*]] = alloc()
+// CHECK-NEXT: %[[CALLEE_RESULT:.*]] = call @callee(%[[FIRST_ELEM]], %[[SECOND_ELEM]], %[[THIRD_ELEM]], %[[FIRST_ALLOC]], %[[SECOND_ALLOC]])
+// CHECK-SAME: (memref<2xf32>, i1, memref<5xf32>, memref<2xf32>, memref<5xf32>) -> i1
+// CHECK-NEXT: %[[TUPLE:.*]] = "test.make_tuple"(%[[FIRST_ALLOC]], %[[CALLEE_RESULT]], %[[SECOND_ALLOC]])
+// 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: linalg.copy(%[[FIRST_ELEM]], %[[RESULT0]])
+// CHECK-NEXT: linalg.copy(%[[THIRD_ELEM]], %[[RESULT1]])
+// CHECK-NEXT: return %[[SECOND_ELEM]]
+
+// -----
+
+// Test case: Testing BufferAssginmnetFuncOpConverter and
+// BufferAssginmentReturnOpConverter to see if the return operation matches with
+// the new function signature when there are tuple typed args and results.
+// BufferAssginmentTypeConverter 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.
+
+// 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>, %[[RESULT0:.*]]: memref<5xf32>, %[[RESULT1:.*]]: memref<10xf32>
+// CHECK-SAME: (i1, 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: linalg.copy(%[[SECOND_TUPLE_SECOND_ELEM]], %[[RESULT0]])
+// CHECK-NEXT: linalg.copy(%[[ARG2]], %[[RESULT1]])
+// CHECK-NEXT: return %[[SECOND_TUPLE_FIRST_ELEM]], %[[FIRST_TUPLE_FIRST_ELEM]], %[[FIRST_TUPLE_SECOND_ELEM]]

mlir/test/lib/Dialect/Test/TestOps.td

Block::BlockArgListType getJoinArgs() {
       return getBody(2)->getArguments();
     }
     OperandRange getSuccessorEntryOperands(unsigned index);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Test TableGen generated build() methods
 //===----------------------------------------------------------------------===//
 
 def TableGenConstant : TEST_Op<"tblgen_constant"> {
   let results = (outs AnyType);
 }
 
 // No variadic args or results.
 def TableGenBuildOp0 : TEST_Op<"tblgen_build_0"> {
   let arguments = (ins AnyType:$value);

 // Tests suppression of ambiguous build methods for operations with
 // SameOperandsAndResultType and InferTypeOpInterface.
 def TableGenBuildOp5 : TEST_Op<"tblgen_build_5",
       [SameOperandsAndResultType, InferTypeOpInterface]> {
   let arguments = (ins Variadic<AnyType>:$inputs);
   let results = (outs AnyType:$result);
 
   let extraClassDeclaration = [{
-    static LogicalResult inferReturnTypes(MLIRContext *, 
+    static LogicalResult inferReturnTypes(MLIRContext *,
           Optional<Location> location, ValueRange operands,
           DictionaryAttr attributes, RegionRange regions,
           SmallVectorImpl<Type> &inferredReturnTypes) {
       inferredReturnTypes.assign({operands[0].getType()});
       return success();
     }
    }];
 }
 
+//===----------------------------------------------------------------------===//
+// Test BufferPlacement
+//===----------------------------------------------------------------------===//
+
+def GetTupleElementOp: TEST_Op<"get_tuple_element"> {
+  let description = [{
+    Test op that returns a specified element of the tuple.
+  }];
+
+  let arguments = (ins
+    TupleOf<[AnyType]>,
+    I32Attr:$index
+  );
+  let results = (outs AnyType);
+}
+
+def MakeTupleOp: TEST_Op<"make_tuple"> {
+  let description = [{
+    Test op that creates a tuple value from a list of values.
+  }];
+
+  let arguments = (ins
+    Variadic<AnyType>:$inputs
+  );
+  let results = (outs TupleOf<[AnyType]>);
+}
+
 #endif // TEST_OPS

mlir/test/lib/Transforms/TestBufferPlacement.cpp

 //===- TestBufferPlacement.cpp - Test for buffer placement ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic for testing buffer placement including its
 // utility converters.
 //
 //===----------------------------------------------------------------------===//
 
+#include "TestDialect.h"
+#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/BufferPlacement.h"
 
 using namespace mlir;
matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
       // one.
       rewriter.replaceOp(op, newResults);
       return success();
     }
   };
 
   void populateTensorLinalgToBufferLinalgConversionPattern(
       MLIRContext *context, BufferAssignmentPlacer *placer,
-      TypeConverter *converter, OwningRewritePatternList *patterns) {
+      BufferAssignmentTypeConverter *converter,
+      OwningRewritePatternList *patterns) {
     populateWithBufferAssignmentOpConversionPatterns<
-        mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp,
-        allowMemrefFunctionResults>(context, placer, converter, patterns);
+        mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>(context, placer,
+                                                        converter, patterns);
     patterns->insert<GenericOpConverter>(context, placer, converter);
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<TestDialect>();
     registry.insert<linalg::LinalgDialect>();
   }
 
   void runOnOperation() override {
     MLIRContext &context = this->getContext();
     ConversionTarget target(context);
     BufferAssignmentTypeConverter converter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<StandardOpsDialect>();
+    target.addLegalOp<MakeTupleOp>();
+    target.addLegalOp<GetTupleElementOp>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
     auto isLegalOperation = [&](Operation *op) {
       return converter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(isLegalOperation);
 
     // Mark Standard Return operations illegal as long as one operand is tensor.
     target.addDynamicallyLegalOp<mlir::ReturnOp>([&](mlir::ReturnOp returnOp) {
       return converter.isLegal(returnOp.getOperandTypes());
     });
 
     // Mark Standard Call Operation illegal as long as it operates on tensor.
     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
+                    ? BufferAssignmentTypeConverter::KeepAsFunctionResult
+                    : BufferAssignmentTypeConverter::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,
+           Location loc) -> Optional<Value> {
+          if (inputs.size() == 1)
+            return llvm::None;
+          TypeRange TypeRange = inputs.getTypes();
+          SmallVector<Type, 2> types(TypeRange.begin(), TypeRange.end());
+          TupleType tuple = TupleType::get(types, builder.getContext());
+          mlir::Value value = builder.create<MakeTupleOp>(loc, tuple, inputs);
+          return value;
+        });
+
+    converter.addDecomposeValueConversion([](OpBuilder &builder, Location loc,
+                                             TupleType resultType, Value value,
+                                             SmallVectorImpl<Value> &values) {
+      for (unsigned i = 0, e = resultType.size(); i < e; ++i) {
+        Value res = builder.create<GetTupleElementOp>(
+            loc, resultType.getType(i), value, builder.getI32IntegerAttr(i));
+        values.push_back(res);
+      }
+      return success();
+    });
+
     // Walk over all the functions to apply buffer assignment.
     this->getOperation().walk([&](FuncOp function) -> WalkResult {
       OwningRewritePatternList patterns;
       BufferAssignmentPlacer placer(function);
       populateTensorLinalgToBufferLinalgConversionPattern(
           &context, &placer, &converter, &patterns);
 
       // Applying full conversion