From ea8e87c8e945c5424b130d66a4fbf5fca7e5c9cd Mon Sep 17 00:00:00 2001 From: Alexander Belyaev Date: Thu, 28 May 2020 06:42:39 -0700 Subject: [PATCH] [XLA][MLIR] Register BufferAssignmentTestDialect. It was made harder to shoot yourself in the foot when using operations that are not registered with MLIRContext. In order to use test ops in the buffer_assignment_test.mlir we have to register them first. PiperOrigin-RevId: 313576961 Change-Id: Id3c711a2d1776a6fdee272d31d932ec5010cd0c2 --- tensorflow/compiler/mlir/xla/tests/BUILD | 1 - .../mlir/xla/tests/buffer-assignment.mlir | 18 +-- .../xla/transforms/buffer_assignment_test.cc | 108 ++++++++++-------- 3 files changed, 68 insertions(+), 59 deletions(-) diff --git a/tensorflow/compiler/mlir/xla/tests/BUILD b/tensorflow/compiler/mlir/xla/tests/BUILD index e2f747085c1..ad69383bd98 100644 --- a/tensorflow/compiler/mlir/xla/tests/BUILD +++ b/tensorflow/compiler/mlir/xla/tests/BUILD @@ -6,7 +6,6 @@ package(licenses = ["notice"]) glob_lit_tests( data = [":test_utilities"], driver = "@llvm-project//mlir:run_lit.sh", - exclude = ["buffer-assignment.mlir"], # TODO(b/157616173) test_file_exts = ["mlir"], ) diff --git a/tensorflow/compiler/mlir/xla/tests/buffer-assignment.mlir b/tensorflow/compiler/mlir/xla/tests/buffer-assignment.mlir index ad007d0eb50..d6c164f8160 100644 --- a/tensorflow/compiler/mlir/xla/tests/buffer-assignment.mlir +++ b/tensorflow/compiler/mlir/xla/tests/buffer-assignment.mlir @@ -203,12 +203,12 @@ func @moving_alloc_and_inserting_missing_dealloc(%cond : i1, %arg0 : memref<2xf3 "buffer_assignment_test.unary_lowered"(%arg0, %1) : (memref<2xf32>, memref<2xf32>) -> () br ^exit(%1 : memref<2xf32>) ^exit(%arg2: memref<2xf32>): - "bufer_assignment_test.copy"(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>) -> () + "buffer_assignment_test.copy"(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>) -> () return } // CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc() // CHECK-NEXT: %[[SECOND_ALLOC:.*]] = alloc() -// CHECK: "bufer_assignment_test.copy" +// CHECK: "buffer_assignment_test.copy" // CHECK-NEXT: dealloc // CHECK-NEXT: dealloc // CHECK-NEXT: return @@ -226,11 +226,11 @@ func @moving_invalid_dealloc_op_complex(%cond : i1, %arg0 : memref<2xf32>, %arg1 dealloc %1 : memref<2xf32> br ^exit(%1 : memref<2xf32>) ^exit(%arg2: memref<2xf32>): - "bufer_assignment_test.copy"(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>) -> () + "buffer_assignment_test.copy"(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>) -> () return } // CHECK-NEXT: %[[ALLOC:.*]] = alloc() -// CHECK: bufer_assignment_test.copy +// CHECK: buffer_assignment_test.copy // CHECK-NEXT: dealloc // CHECK-NEXT: return @@ -240,10 +240,10 @@ func @moving_invalid_dealloc_op_complex(%cond : i1, %arg0 : memref<2xf32>, %arg1 func @inserting_missing_dealloc_simple(%arg0 : memref<2xf32>, %arg1: memref<2xf32>){ %0 = alloc() : memref<2xf32> "buffer_assignment_test.unary_lowered"(%arg0, %0) : (memref<2xf32>, memref<2xf32>) -> () - "bufer_assignment_test.copy"(%0, %arg1) : (memref<2xf32>, memref<2xf32>) -> () + "buffer_assignment_test.copy"(%0, %arg1) : (memref<2xf32>, memref<2xf32>) -> () return } -// CHECK: bufer_assignment_test.copy +// CHECK: buffer_assignment_test.copy // CHECK-NEXT: dealloc // ----- @@ -253,8 +253,8 @@ func @moving_invalid_dealloc_op(%arg0 : memref<2xf32>, %arg1: memref<2xf32>){ %0 = alloc() : memref<2xf32> "buffer_assignment_test.unary_lowered"(%arg0, %0) : (memref<2xf32>, memref<2xf32>) -> () dealloc %0 : memref<2xf32> - "bufer_assignment_test.copy"(%0, %arg1) : (memref<2xf32>, memref<2xf32>) -> () + "buffer_assignment_test.copy"(%0, %arg1) : (memref<2xf32>, memref<2xf32>) -> () return } -// CHECK: bufer_assignment_test.copy -// CHECK-NEXT: dealloc \ No newline at end of file +// CHECK: buffer_assignment_test.copy +// CHECK-NEXT: dealloc diff --git a/tensorflow/compiler/mlir/xla/transforms/buffer_assignment_test.cc b/tensorflow/compiler/mlir/xla/transforms/buffer_assignment_test.cc index 5a0d791079c..40c115f4cbc 100644 --- a/tensorflow/compiler/mlir/xla/transforms/buffer_assignment_test.cc +++ b/tensorflow/compiler/mlir/xla/transforms/buffer_assignment_test.cc @@ -29,60 +29,66 @@ limitations under the License. namespace mlir { namespace xla { namespace { + +/// This dialect independent unary operation has been defined only for testing +/// buffer assignment. +class BufferAssignmentTestUnaryOp + : public Op { + public: + using Op::Op; + static StringRef getOperationName() { return "buffer_assignment_test.unary"; } + static void build(OpBuilder& b, OperationState& state, Value source) { + state.addOperands(source); + } +}; + +/// This dialect independent lowered unary operation has been defined only for +/// testing buffer assignment. +class BufferAssignmentTestUnaryLoweredOp + : public Op::Impl> { + public: + using Op::Op; + static StringRef getOperationName() { + return "buffer_assignment_test.unary_lowered"; + } + static void build(OpBuilder& b, OperationState& state, Value source, + Value target) { + state.addOperands(source); + state.addOperands(target); + } +}; + +/// This dialect independent copy operation has been defined only for testing +/// NonVoidToVoidReturnOpConverter +class BufferAssignmentTestCopyOp + : public Op::Impl> { + public: + using Op::Op; + static StringRef getOperationName() { return "buffer_assignment_test.copy"; } + static void build(OpBuilder& b, OperationState& state, Value from, Value to) { + state.addOperands(from); + state.addOperands(to); + } +}; + +class BufferAssignmentTestDialect : public Dialect { + public: + explicit BufferAssignmentTestDialect(MLIRContext* context) + : Dialect(getDialectNamespace(), context) { + addOperations(); + } + static StringRef getDialectNamespace() { return "buffer_assignment_test"; } +}; + /// This pass tests two provided operation converters, /// FunctionAndBlockSignatureConverter and NonVoidToVoidReturnOpConverter, for /// Buffer Assignment. struct BufferAssignmentPreparationTestPass : mlir::PassWrapper { - /// This dialect independent unary operation has been defined only for testing - /// buffer assignment. - class BufferAssignmentTestUnaryOp - : public Op { - public: - using Op::Op; - static StringRef getOperationName() { - return "buffer_assignment_test.unary"; - } - static void build(OpBuilder& b, OperationState& state, Value source) { - state.addOperands(source); - } - }; - - /// This dialect independent lowered unary operation has been defined only for - /// testing buffer assignment. - class BufferAssignmentTestUnaryLoweredOp - : public Op::Impl> { - public: - using Op::Op; - static StringRef getOperationName() { - return "buffer_assignment_test.unary_lowered"; - } - static void build(OpBuilder& b, OperationState& state, Value source, - Value target) { - state.addOperands(source); - state.addOperands(target); - } - }; - - /// This dialect independent copy operation has been defined only for testing - /// NonVoidToVoidReturnOpConverter - class BufferAssignmentTestCopyOp - : public Op::Impl> { - public: - using Op::Op; - static StringRef getOperationName() { - return "buffer_assignment_test.copy"; - } - static void build(OpBuilder& b, OperationState& state, Value from, - Value to) { - state.addOperands(from); - state.addOperands(to); - } - }; - /// A simple converter that legalizes a BufferAssignmentTestUnaryOp to a /// BufferAssignmentTestUnaryLoweredOp and creates buffer allocation for /// the result of the computation. @@ -151,8 +157,12 @@ struct BufferAssignmentPreparationTestPass } }; }; + } // namespace +static mlir::DialectRegistration + buffer_assignment_test_ops; + /// This pass tests helper methods such as computeAllocPosition, /// FunctionAndBlockSignatureConverter, NonVoidToVoidReturnOpConverter /// conversion patterns. Furthermore, it checks buffer-assignment pass that