[Mlir-commits] [mlir] 3f6a35e - [mlir] Introduce CallOp converter for buffer placement
Ehsan Toosi
llvmlistbot at llvm.org
Tue Jun 2 02:38:12 PDT 2020
Author: Ehsan Toosi
Date: 2020-06-02T11:35:24+02:00
New Revision: 3f6a35e3ffd49e063bb12464513b3847c535c242
URL: https://github.com/llvm/llvm-project/commit/3f6a35e3ffd49e063bb12464513b3847c535c242
DIFF: https://github.com/llvm/llvm-project/commit/3f6a35e3ffd49e063bb12464513b3847c535c242.diff
LOG: [mlir] Introduce CallOp converter for buffer placement
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.
Differential Revision: https://reviews.llvm.org/D80785
Added:
Modified:
mlir/include/mlir/Transforms/BufferPlacement.h
mlir/lib/Transforms/BufferPlacement.cpp
mlir/test/Transforms/buffer-placement-preparation.mlir
mlir/test/lib/Transforms/TestBufferPlacement.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Transforms/BufferPlacement.h b/mlir/include/mlir/Transforms/BufferPlacement.h
index 10949160fcbd..89cb4b043a19 100644
--- a/mlir/include/mlir/Transforms/BufferPlacement.h
+++ b/mlir/include/mlir/Transforms/BufferPlacement.h
@@ -157,6 +157,21 @@ class BufferAssignmentReturnOpConverter
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
diff --git a/mlir/lib/Transforms/BufferPlacement.cpp b/mlir/lib/Transforms/BufferPlacement.cpp
index 337e26e2a65b..edbaf8e39c97 100644
--- a/mlir/lib/Transforms/BufferPlacement.cpp
+++ b/mlir/lib/Transforms/BufferPlacement.cpp
@@ -468,6 +468,57 @@ LogicalResult FunctionAndBlockSignatureConverter::matchAndRewrite(
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
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Transforms/buffer-placement-preparation.mlir b/mlir/test/Transforms/buffer-placement-preparation.mlir
index 8458154e4985..5cde928e6da2 100644
--- a/mlir/test/Transforms/buffer-placement-preparation.mlir
+++ b/mlir/test/Transforms/buffer-placement-preparation.mlir
@@ -195,3 +195,92 @@ func @compute_allocs_position(%cond: i1, %arg0: tensor<2xf32>) -> tensor<2xf32>{
// 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
diff --git a/mlir/test/lib/Transforms/TestBufferPlacement.cpp b/mlir/test/lib/Transforms/TestBufferPlacement.cpp
index 6152a9b85435..aee12b37a687 100644
--- a/mlir/test/lib/Transforms/TestBufferPlacement.cpp
+++ b/mlir/test/lib/Transforms/TestBufferPlacement.cpp
@@ -106,6 +106,7 @@ struct TestBufferPlacementPreparationPass
TypeConverter *converter, OwningRewritePatternList *patterns) {
// clang-format off
patterns->insert<
+ BufferAssignmentCallOpConverter,
FunctionAndBlockSignatureConverter,
GenericOpConverter,
BufferAssignmentReturnOpConverter<
@@ -137,6 +138,12 @@ struct TestBufferPlacementPreparationPass
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());
More information about the Mlir-commits
mailing list