[MLIR] Update the FunctionAndBlockSignatureConverter and NonVoidToVoidReturnOpConverter of Buffer Assignment

[dfki-ehna]

Summary

Making these two converters more generic. FunctionAndBlockSignatureConverter now moves only memref results (after type conversion) to the function argument and keeps other legal function results unchanged. NonVoidToVoidReturnOpConverter is renamed to NoBufferOperandsReturnOpConverter`. It removes only the buffer operands from the operands of the converted ReturnOp and inserts CopyOps to copy each buffer to the target function argument.

Input:

func @example(%arg0: tensor<4x8xf32>, %arg1: i1) -> (tensor<4x8xf32>, i1) {
    return %arg0, %arg1 : i1, tensor<4x8xf32>
}

Output:

func @example(%arg0: memref<4x8xf32>, %arg1: i1, %arg2: memref<4x8xf32>) -> (i1) {
    copy(%arg0, %arg2) : memref<4x8xf32>, memref<4x8xf32>
    return %arg1 : i1
}

[pifon2a]

nice! thanks!

mlir/include/mlir/Transforms/BufferPlacement.h

84

/// Converts the signature of the function using the type converter.
/// It adds an extra argument for each illegally-typed function
/// result to the function arguments. `BufferAssignmentTypeConverter`
/// is a helper `TypeConverter for this purpose. All the non-shaped types
/// of the input function will be converted to memref.

98

/// Converts the source `ReturnOp` to target `ReturnOp`, removes all
/// the buffer operands from the operands list, and inserts `CopyOp`s
/// for all buffer operands instead.

121

these 5 lines are just llvm::erase_if, right?

125

// For each memref type operand of the source `ReturnOp`, a new `CopyOp` is
// inserted that copies the buffer content from the operand to the target.

131

can this be moved outside of the loop?

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

15

please, reformat it so that it fits

[herhut]

Please add a test for the case that is currently asserted.

mlir/include/mlir/Transforms/BufferPlacement.h

114–115

In the case where the return has non-memref operands, this can fail even for correct code as the non-memref operands do not require a function argument. Maybe add a test for such a case?

[dfki-ehna]

Addressing the comments.

[dfki-ehna]

mlir/include/mlir/Transforms/BufferPlacement.h

114–115

Thanks, you are right. I changed the assert and updated the target.addDynamicallyLegalOp<FuncOp>. A few new test cases were also added.

[herhut]

mlir/lib/Transforms/BufferPlacement.cpp

454

it is remained -> it remains unchanged as a function result?

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

34

If this returned a memref before, why would it not return a memref now? We should only rewrite the return values that got rewritten from tensor to memref. Not the ones that already were a memref. I'd assume those would be left alone.

[dfki-ehna]

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

34

The idea of removing all Memref types from function results is proposed to avoid placing a DeallocOp after the return operation using buffer-placement pass. This case would happen if there is a defining operation for a return operand.

func @memref_in_function_results(%arg0: i1, %arg1: f16, %arg2:memref<2xf32>) -> (i1, memref<2xf32>, f16){
  %0 = alloc() : memref<2xf32>,
  lhlo.exp(%arg2, %0)
  return %arg0, %0, %arg1  : i1, memref<2xf32>, f16
}

[herhut]

Making it fail would be good enough for landing.

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

34

I think it is better to just not allow this case. If one has a memref that escapes in the original program, it should still escape in the rewritten program. We cannot support this use case at the moment it seems, so it would be better to fail.

I wonder what it would take, though, to handle this case. Would it suffice to know that the last use is a return and then treat the return like a dealloc from the local scope? This could be an alternative version of this rewriting that turns tensor typed returns into memref typed returns (and not arguments) and lets allocations escape. Would the main logic of placement still work in this setting? What would it take to support both modes?

I see a use for the latter mode in scenarios where the buffers for results are not provided from the context, which is particularly useful in the presence of dynamic shapes where the context might not know sizes.

Such a rewrite would not be globally correct (the caller would need to free the buffer after all and that cannot be decided statically) but at least locally per function it would be. For global correctness we could think in the direction of reference counting but that would be a next step.

[dfki-ehna]

BufferAssignmentPlacer now fails in cases that function returns buffers already. A note has also been added to the description of BufferPlacement.cpp about not supporting these cases. The related TODO in TensorsToBuffers.cpp has also been removed.

[mehdi_amini]

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

34

+1: I would really prefer that we keep passes semantically preserving and be conservative with respect to not changing the original behavior.

[herhut]

mlir/lib/Transforms/BufferPlacement.cpp

435

Instead of an assert, can you fail the rewrite and report an error? That way, the program does not just crash and this also fails outside of debug builds.

Also, please add a test that the error is generated,

[dfki-ehna]

Change Assert to emitError and add a test case for the generated error.

[pifon2a]

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

1

typo in the filename /buffer-placement-prepration.mlir

[herhut]

Thanks. Please fix the filename and land.

mlir/include/mlir/Transforms/BufferPlacement.h

explicit BufferAssignmentOpConversionPattern(
       : OpConversionPattern<SourceOp>(context, benefit),
         bufferAssignment(bufferAssignment), converter(converter) {}
 
 protected:
   BufferAssignmentPlacer *bufferAssignment;
   TypeConverter *converter;
 };
 
-/// This conversion adds an extra argument for each function result which makes
-/// the converted function a void function. A type converter must be provided
-/// for this conversion to convert a non-shaped type to memref.
-/// BufferAssignmentTypeConverter is an helper TypeConverter for this
-/// purpose. All the non-shaped type of the input function will be converted to
-/// memref.
+/// Converts the signature of the function using the type converter.
+/// It adds an extra argument for each illegally-typed function
+/// result to the function arguments. `BufferAssignmentTypeConverter`
+/// is a helper `TypeConverter` for this purpose. All the non-shaped types
+/// of the input function will be converted to memref.
 class FunctionAndBlockSignatureConverter
     : public BufferAssignmentOpConversionPattern<FuncOp> {
 public:
   using BufferAssignmentOpConversionPattern<
       FuncOp>::BufferAssignmentOpConversionPattern;
 
   /// Performs the actual signature rewriting step.
   LogicalResult
   matchAndRewrite(FuncOp funcOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final;
 };
 
-/// This pattern converter transforms a non-void ReturnOpSourceTy into a void
-/// return of type ReturnOpTargetTy. It uses a copy operation of type CopyOpTy
-/// to copy the results to the output buffer.
+/// Converts the source `ReturnOp` to target `ReturnOp`, removes all
+/// the buffer operands from the operands list, and inserts `CopyOp`s
+/// for all buffer operands instead.
 template <typename ReturnOpSourceTy, typename ReturnOpTargetTy,
           typename CopyOpTy>
-class NonVoidToVoidReturnOpConverter
+class NoBufferOperandsReturnOpConverter
     : 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 {
-    unsigned numReturnValues = returnOp.getNumOperands();
     Block &entryBlock = returnOp.getParentRegion()->front();
     unsigned numFuncArgs = entryBlock.getNumArguments();
     Location loc = returnOp.getLoc();
 
-    // Find the corresponding output buffer for each operand.
-    assert(numReturnValues <= numFuncArgs &&
+    // The target `ReturnOp` should not contain any memref operands.
+    SmallVector<Value, 2> newOperands(operands.begin(), operands.end());
+    llvm::erase_if(newOperands, [](Value operand) {
+      return operand.getType().isa<MemRefType>();
+    });
+
+    // Find the index of the first destination buffer.
+    unsigned numBufferOperands = operands.size() - newOperands.size();
+    unsigned destArgNum = numFuncArgs - numBufferOperands;
+
+    rewriter.setInsertionPoint(returnOp);
+    // Find the corresponding destination buffer for each memref operand.
+    for (Value operand : operands)
+      if (operand.getType().isa<MemRefType>()) {
+        assert(destArgNum < numFuncArgs &&
                "The number of operands of return operation is more than the "
                "number of function argument.");
-    unsigned firstReturnParameter = numFuncArgs - numReturnValues;
-    for (auto operand : llvm::enumerate(operands)) {
-      unsigned returnArgNumber = firstReturnParameter + operand.index();
-      BlockArgument dstBuffer = entryBlock.getArgument(returnArgNumber);
-      if (dstBuffer == operand.value())
-        continue;
 
-      // Insert the copy operation to copy before the return.
-      rewriter.setInsertionPoint(returnOp);
-      rewriter.create<CopyOpTy>(loc, operand.value(),
-                                entryBlock.getArgument(returnArgNumber));
+        // For each memref type operand of the source `ReturnOp`, a new `CopyOp`
+        // is inserted that copies the buffer content from the operand to the
+        // target.
+        rewriter.create<CopyOpTy>(loc, operand,
+                                  entryBlock.getArgument(destArgNum));
+        ++destArgNum;
       }
-    // Insert the new target return operation.
-    rewriter.replaceOpWithNewOp<ReturnOpTargetTy>(returnOp);
+
+    // Insert the new target Return operation.
+    rewriter.replaceOpWithNewOp<ReturnOpTargetTy>(returnOp, newOperands);
     return success();
   }
 };
 
 /// 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();
 };
 
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_BUFFERPLACEMENT_H

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

 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/BufferPlacement.h"
 
 using namespace mlir;
 using ReturnOpConverter =
-    NonVoidToVoidReturnOpConverter<mlir::ReturnOp, mlir::ReturnOp,
+    NoBufferOperandsReturnOpConverter<mlir::ReturnOp, mlir::ReturnOp,
                                       linalg::CopyOp>;
 
 namespace {
 /// A pattern to convert Generic Linalg operations which work on tensors to
 /// use buffers. A buffer is allocated using BufferAssignmentPlacer for
 /// each operation result. BufferPlacement pass should be later used to move
 /// Alloc operations to the correct positions and insert the missing Dealloc
 /// operations in the correct places.
 class GenericOpConverter
void runOnOperation() override {
     auto isLegalOperation = [&](Operation *op) {
       return llvm::none_of(op->getOperandTypes(), isIllegalType) &&
              llvm::none_of(op->getResultTypes(), isIllegalType);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(
         Optional<ConversionTarget::DynamicLegalityCallbackFn>(
             isLegalOperation));
 
-    // TODO: Considering the following dynamic legality checks, the current
-    // implementation of FunctionAndBlockSignatureConverter of Buffer Assignment
-    // will convert the function signature incorrectly. This converter moves
-    // all the return values of the function to the input argument list without
-    // considering the return value types and creates a void function. However,
-    // the NonVoidToVoidReturnOpConverter doesn't change the return operation if
-    // its operands are not tensors. The following example leaves the IR in a
-    // broken state.
-    //
-    // @function(%arg0: f32, %arg1: tensor<4xf32>) -> (f32, f32) {
-    //    %0 = mulf %arg0, %arg0 : f32
-    //    return %0, %0 : f32, f32
-    // }
-    //
-    // broken IR after conversion:
-    //
-    // func @function(%arg0: f32, %arg1: memref<4xf32>, f32, f32) {
-    //    %0 = mulf %arg0, %arg0 : f32
-    //    return %0, %0 : f32, f32
-    // }
-    //
-    // This issue must be fixed in FunctionAndBlockSignatureConverter and
-    // NonVoidToVoidReturnOpConverter.
-
     // Mark Standard Return operations illegal as long as one operand is tensor.
     target.addDynamicallyLegalOp<mlir::ReturnOp>([&](mlir::ReturnOp returnOp) {
       return llvm::none_of(returnOp.getOperandTypes(), isIllegalType);
     });
 
     // Mark the function operation illegal as long as an argument is tensor.
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
       return converter.isSignatureLegal(funcOp.getType()) &&

mlir/lib/Transforms/BufferPlacement.cpp

 // after which we have to insert the dealloc. Finding the alloc position is
 // highly similar and non- obvious. Again, we have to consider all potential
 // aliases and find the common dominator block to place the alloc.
 //
 // TODO:
 // The current implementation does not support loops and the resulting code will
 // be invalid with respect to program semantics. The only thing that is
 // currently missing is a high-level loop analysis that allows us to move allocs
-// and deallocs outside of the loop blocks.
+// and deallocs outside of the loop blocks. Furthermore, it doesn't also accept
+// functions which return buffers already.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/BufferPlacement.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"

 LogicalResult FunctionAndBlockSignatureConverter::matchAndRewrite(
     FuncOp funcOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   if (!converter) {
     funcOp.emitError("The type converter has not been defined for "
                      "FunctionAndBlockSignatureConverter");
     return failure();
   }
-  // Converting shaped type arguments to memref type.
   auto funcType = funcOp.getType();
+  TypeRange resultTypes = funcType.getResults();
+  if (llvm::any_of(resultTypes,
+                   [](Type type) { return type.isa<MemRefType>(); }))
+    return funcOp.emitError("BufferAssignmentPlacer doesn't currently support "
+                            "functions which return memref typed values");
+
+  // 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()));
-  // Adding function results to the arguments of the converted function as
-  // memref type. The converted function will be a void function.
-  for (Type resType : funcType.getResults())
-    conversion.addInputs(converter->convertType((resType)));
+
+  // Adding a function argument for each function result which is going to be a
+  // memref type after type conversion.
+  SmallVector<Type, 2> newResultTypes;
+  newResultTypes.reserve(funcOp.getNumResults());
+  for (Type resType : resultTypes) {
+    Type convertedType = converter->convertType(resType);
+
+    // If the result type is memref after the type conversion, a new argument
+    // should be appended to the function arguments list for this result.
+    // Otherwise, it remains unchanged as a function result.
+    if (convertedType.isa<MemRefType>())
+      conversion.addInputs(convertedType);
+    else
+      newResultTypes.push_back(convertedType);
+  }
+
+  // Update the signature of the function.
   rewriter.updateRootInPlace(funcOp, [&] {
-    funcOp.setType(
-        rewriter.getFunctionType(conversion.getConvertedTypes(), llvm::None));
+    funcOp.setType(rewriter.getFunctionType(conversion.getConvertedTypes(),
+                                            newResultTypes));
     rewriter.applySignatureConversion(&funcOp.getBody(), conversion);
   });
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // BufferAssignmentTypeConverter
 //===----------------------------------------------------------------------===//

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

-// RUN: mlir-opt -test-buffer-placement-preparation -split-input-file %s | FileCheck %s -dump-input-on-failure
+// RUN: mlir-opt -test-buffer-placement-preparation -split-input-file -verify-diagnostics %s | FileCheck %s -dump-input-on-failure
 
 // CHECK-LABEL: func @func_signature_conversion
 func @func_signature_conversion(%arg0: tensor<4x8xf32>) {
     return
 }
 // CHECK: ({{.*}}: memref<4x8xf32>) {
 
 // -----
 
+// expected-error @below {{BufferAssignmentPlacer doesn't currently support functions which return memref typed values}}
+// expected-error @below {{failed to legalize operation 'func'}}
+func @memref_in_function_results(%arg0: tensor<4x8xf32>) -> (tensor<4x8xf32>, memref<5xf32>) {
+  %0 = alloc() : memref<5xf32>
+  return %arg0, %0 : tensor<4x8xf32>, memref<5xf32>
+}
+
+// -----
+
+// 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 @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, %[[ARG2:.*]]: memref<5x5xf64>, %[[ARG3:.*]]: f16,
+// CHECK-SAME: %[[RESULT1:.*]]: memref<5x5xf64>, %[[RESULT2:.*]]: memref<4x8xf32>) -> (i1, f16) {
+// CHECK-NEXT: linalg.copy(%[[ARG2]], %[[RESULT1]])
+// CHECK-NEXT: linalg.copy(%[[ARG0]], %[[RESULT2]])
+// CHECK-NEXT: return %[[ARG1]], %[[ARG3]]
+
+// -----
+
 // 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: (%[[ARG0:.*]]: [[TYPE:.*]]<[[RANK:.*]]>, %[[RESULT:.*]]: [[TYPE]]<[[RANK]]>) {
 // CHECK-NEXT: linalg.copy(%[[ARG0]], %[[RESULT]])
 // CHECK-NEXT: return
 

mlir/test/lib/Transforms/TestBufferPlacement.cpp

 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/BufferPlacement.h"
 
 using namespace mlir;
 
 namespace {
 /// This pass tests the computeAllocPosition helper method and two provided
 /// operation converters, FunctionAndBlockSignatureConverter and
-/// NonVoidToVoidReturnOpConverter. Furthermore, this pass converts linalg
+/// NoBufferOperandsReturnOpConverter. Furthermore, this pass converts linalg
 /// operations on tensors to linalg operations on buffers to prepare them for
 /// the BufferPlacement pass that can be applied afterwards.
 struct TestBufferPlacementPreparationPass
     : mlir::PassWrapper<TestBufferPlacementPreparationPass,
                         OperationPass<ModuleOp>> {
 
   /// Converts tensor-type generic linalg operations to memref ones using buffer
   /// assignment.
matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
       rewriter.inlineRegionBefore(op.region(), region, region.end());
 
       // TODO: verify the internal memref-based linalg functionality.
       auto &entryBlock = region.front();
       for (auto result : results) {
         auto type = result.getType().cast<ShapedType>();
         entryBlock.addArgument(type.getElementType());
       }
-
       rewriter.eraseOp(op);
       return success();
     }
   };
 
   void populateTensorLinalgToBufferLinalgConversionPattern(
       MLIRContext *context, BufferAssignmentPlacer *placer,
       TypeConverter *converter, OwningRewritePatternList *patterns) {
     // clang-format off
     patterns->insert<
                    FunctionAndBlockSignatureConverter,
                    GenericOpConverter,
-                   NonVoidToVoidReturnOpConverter<
+                   NoBufferOperandsReturnOpConverter<
                       ReturnOp, ReturnOp, linalg::CopyOp>
     >(context, placer, converter);
     // clang-format on
   }
 
   void runOnOperation() override {
     auto &context = getContext();
     ConversionTarget target(context);
     BufferAssignmentTypeConverter converter;
-    // Make all linalg operations illegal as long as they work on tensors.
     target.addLegalDialect<StandardOpsDialect>();
+
+    // Make all linalg operations illegal as long as they work on tensors.
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(
         Optional<ConversionTarget::DynamicLegalityCallbackFn>(
             [&](Operation *op) {
               auto isIllegalType = [&](Type type) {
                 return !converter.isLegal(type);
               };
               return llvm::none_of(op->getOperandTypes(), isIllegalType) &&
                      llvm::none_of(op->getResultTypes(), isIllegalType);
             }));
 
-    // Mark return operations illegal as long as they return values.
-    target.addDynamicallyLegalOp<mlir::ReturnOp>(
-        [](mlir::ReturnOp returnOp) { return returnOp.getNumOperands() == 0; });
+    // Mark std.ReturnOp illegal as long as an operand is tensor or buffer.
+    target.addDynamicallyLegalOp<mlir::ReturnOp>([&](mlir::ReturnOp returnOp) {
+      return llvm::none_of(returnOp.getOperandTypes(), [&](Type type) {
+        return type.isa<MemRefType>() || !converter.isLegal(type);
+      });
+    });
 
     // Mark the function whose arguments are in tensor-type illegal.
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
       return converter.isSignatureLegal(funcOp.getType());
     });
 
     // Walk over all the functions to apply buffer assignment.
     getOperation().walk([&](FuncOp function) {

 namespace mlir {
 void registerTestBufferPlacementPreparationPass() {
   PassRegistration<TestBufferPlacementPreparationPass>(
       "test-buffer-placement-preparation",
       "Tests buffer placement helper methods including its "
       "operation-conversion patterns");
 }
 } // end namespace mlir
 No newline at end of file