[MLIR][BUFFER_PLACEMENT] Introduce CallOp converter for buffer placement

[dfki-ehna]

Summary

Add BufferAssignmentCallOpConverter as a pattern rewriter for Buffer Placement. It matches the signature of the caller operation with the callee after rewriting the callee with FunctionAndBlockSignatureConverter.

[pifon2a]

mlir/lib/Transforms/BufferPlacement.cpp

508

nit: No replacing is required.

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

214

Please, format IR to fit 80 chars, smth like:

%0 = linalg.generic {
       args_in = 1 : i64, 
       args_out = 1 : i64,
       indexing_maps = [#map0, #map0], 
       iterator_types = ["parallel"]
     } %arg1 {

257

that's nice that we can already do that!

[mehdi_amini]

(Please void upper caps in your commit title)
(I'd also avoid [BUFFER_PLACEMENT] as a tag in the first place, the commit title is descriptive enough as Introduce CallOp converter for buffer placement)

[herhut]

Thanks!

[dfki-ehna]

Resolve the comments.

mlir/include/mlir/Transforms/BufferPlacement.h

for (Value operand : needCopyOperands) {
       ++destArgNum;
     }
 
     // Insert the new target Return operation.
     rewriter.replaceOpWithNewOp<ReturnOpTargetTy>(returnOp, newOperands);
     return success();
   }
 };
+
+/// Converts `CallOp` to match its operands and results with the
+/// the callee after rewriting the callee with
+/// FunctionAndBlockSignatureConverter.
+class BufferAssignmentCallOpConverter
+    : public BufferAssignmentOpConversionPattern<CallOp> {
+public:
+  using BufferAssignmentOpConversionPattern<
+      CallOp>::BufferAssignmentOpConversionPattern;
+
+  /// Performs the actual `CallOp` conversion step.
+  LogicalResult
+  matchAndRewrite(CallOp callOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final;
+};
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_BUFFERPLACEMENT_H

mlir/lib/Transforms/BufferPlacement.cpp

rewriter.updateRootInPlace(funcOp, [&] {
     funcOp.setType(rewriter.getFunctionType(conversion.getConvertedTypes(),
                                             newResultTypes));
     rewriter.applySignatureConversion(&funcOp.getBody(), conversion);
   });
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// BufferAssignmentCallOpConverter
+//===----------------------------------------------------------------------===//
+
+// Performs `CallOp` conversion to match its operands and results with the
+// signature of the callee after rewriting the callee with
+// FunctionAndBlockSignatureConverter.
+LogicalResult BufferAssignmentCallOpConverter::matchAndRewrite(
+    CallOp callOp, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+
+  Location loc = callOp.getLoc();
+  SmallVector<Value, 2> newOperands, replacingValues;
+  SmallVector<Type, 2> newResultTypes;
+  unsigned numResults = callOp.getNumResults();
+  newOperands.reserve(numResults + operands.size());
+  newOperands.append(operands.begin(), operands.end());
+  newResultTypes.reserve(numResults);
+  replacingValues.reserve(numResults);
+
+  // For each memref result of `CallOp` which has not been a memref before type
+  // conversion, a new buffer is allocated and passed to the operands list of
+  // the new `CallOp`. Otherwise, it remains as a caller result.
+  for (Value result : callOp.getResults()) {
+    Type currType = result.getType();
+    Type newType = converter->convertType(result.getType());
+    if (BufferAssignmentTypeConverter::isConvertedMemref(newType, currType)) {
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.restoreInsertionPoint(
+          bufferAssignment->computeAllocPosition(result.dyn_cast<OpResult>()));
+      Value alloc =
+          rewriter.create<AllocOp>(loc, newType.dyn_cast<MemRefType>());
+      newOperands.push_back(alloc);
+      replacingValues.push_back(alloc);
+    } else {
+      newResultTypes.push_back(currType);
+
+      // No replacing is required.
+      replacingValues.push_back(nullptr);
+    }
+  }
+
+  // Creating the new `CallOp`.
+  rewriter.create<CallOp>(loc, callOp.getCallee(), newResultTypes, newOperands);
+
+  // Replacing the results of the old `CallOp`.
+  rewriter.replaceOp(callOp, replacingValues);
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // BufferAssignmentTypeConverter
 //===----------------------------------------------------------------------===//
 
 /// Registers conversions into BufferAssignmentTypeConverter
 BufferAssignmentTypeConverter::BufferAssignmentTypeConverter() {
   // Keep all types unchanged.
   addConversion([](Type type) { return type; });
   // A type conversion that converts ranked-tensor type to memref type.

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

 //      CHECK: %[[ALLOC5:.*]] = alloc()
 // CHECK-NEXT: linalg.generic {{.*}} %[[ALLOC4]], %[[ALLOC5]]
 //      CHECK: br ^[[EXIT]]
 // CHECK-NEXT: ^[[EXIT]]
 // CHECK-NEXT: %[[ALLOC6:.*]] = alloc()
 // CHECK-NEXT: linalg.generic {{.*}} %[[ARG0]], %[[ALLOC6]]
 //      CHECK: %[[ALLOC7:.*]] = alloc()
 // CHECK-NEXT: linalg.generic {{.*}} %[[ALLOC6]], %[[ALLOC7]]
+
+// -----
+
+// Test case: Checking BufferAssignmentCallOpConverter and
+// FunctionAndBlockSignatureConverter and BufferAssignmentReturnOpConverter all
+// together. The signature of `callee` after signature conversion would be:
+
+// func @callee(%arg0: memref<5xf32>,%arg1: memref<5xf32>) -> ()
+
+// The operands and results of caller and return operations must be matched
+// respectively.
+
+#map0 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: func @callee
+func @callee(%arg1: tensor<5xf32>) -> tensor<5xf32> {
+  %0 = linalg.generic {
+                        args_in = 1 : i64,
+                        args_out = 1 : i64,
+                        indexing_maps = [#map0, #map0],
+                        iterator_types = ["parallel"]
+                      } %arg1 {
+  ^bb0(%gen1_arg0: f32):
+    %tmp1 = exp %gen1_arg0 : f32
+    linalg.yield %tmp1 : f32
+  }: tensor<5xf32> -> tensor<5xf32>
+  return %0 : tensor<5xf32>
+}
+// CHECK: (%[[CALLEE_ARG:.*]]: memref<5xf32>, %[[CALLEE_RESULT:.*]]: memref<5xf32>)
+// CHECK: %[[ALLOC:.*]] = alloc()
+// CHECK: linalg.generic
+// CHECK: linalg.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>, %[[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: linalg.copy(%[[SECOND_ALLOC]], %[[CALLER_RESULT]])
+// CHECK: return
+
+// -----
+
+// Test case: Checking BufferAssignmentCallOpConverter and
+// FunctionAndBlockSignatureConverter and BufferAssignmentReturnOpConverter all
+// together on functions that also have memref typed results. The signature of
+// `callee` after signature conversion would be:
+
+// func @callee(%arg0: memref<5xf32>,%arg1: memref<5xf32>)-> memref<2xf32>
+
+// where %arg0 is the input and %arg1 is the output buffer and the original memref
+// type result remain as the function result. Then, the rewriter should match the
+// caller's signature with the callee. Thus, two buffers will be allocated instead
+// 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>, %[[CALLEE_RESULT:.*]]: memref<5xf32>)
+// CHECK-SAME: memref<2xf32>
+//      CHECK: %[[ALLOC:.*]] = alloc()
+//      CHECK: linalg.copy(%[[CALLEE_ARG]], %[[CALLEE_RESULT]])
+//      CHECK: return %[[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>, %[[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

mlir/test/lib/Transforms/TestBufferPlacement.cpp

public:
     }
   };
 
   void populateTensorLinalgToBufferLinalgConversionPattern(
       MLIRContext *context, BufferAssignmentPlacer *placer,
       TypeConverter *converter, OwningRewritePatternList *patterns) {
     // clang-format off
     patterns->insert<
+                   BufferAssignmentCallOpConverter,
                    FunctionAndBlockSignatureConverter,
                    GenericOpConverter,
                    BufferAssignmentReturnOpConverter<
                       ReturnOp, ReturnOp, linalg::CopyOp>
     >(context, placer, converter);
     // clang-format on
   }
 
target.addDynamicallyLegalDialect<linalg::LinalgDialect>(
         Optional<ConversionTarget::DynamicLegalityCallbackFn>(
             isLegalOperation));
 
     // 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 Standard Call Operation illegal as long as it operates on tensor.
+    target.addDynamicallyLegalOp<mlir::CallOp>([&](mlir::CallOp callOp) {
+      return llvm::none_of(callOp.getOperandTypes(), isIllegalType) &&
+             llvm::none_of(callOp.getResultTypes(), isIllegalType);
+    });
+
     // Mark the function whose arguments are in tensor-type illegal.
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
       return converter.isSignatureLegal(funcOp.getType());
     });
 
     // Walk over all the functions to apply buffer assignment.
     getOperation().walk([&](FuncOp function) -> WalkResult {
       OwningRewritePatternList patterns;

 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