[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
diff --git a/tensorflow/compiler/mlir/xla/tests/BUILD b/tensorflow/compiler/mlir/xla/tests/BUILD
index e2f7470..ad69383 100644
--- a/tensorflow/compiler/mlir/xla/tests/BUILD
+++ b/tensorflow/compiler/mlir/xla/tests/BUILD
@@ -6,7 +6,6 @@
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 ad007d0..d6c164f 100644
--- a/tensorflow/compiler/mlir/xla/tests/buffer-assignment.mlir
+++ b/tensorflow/compiler/mlir/xla/tests/buffer-assignment.mlir
@@ -203,12 +203,12 @@
"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 @@
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 @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 @@
%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 5a0d791..40c115f 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 @@
namespace mlir {
namespace xla {
namespace {
+
+/// This dialect independent unary operation has been defined only for testing
+/// buffer assignment.
+class BufferAssignmentTestUnaryOp
+ : public Op<BufferAssignmentTestUnaryOp, OpTrait::OneResult,
+ OpTrait::OneOperand> {
+ 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<BufferAssignmentTestUnaryLoweredOp, OpTrait::ZeroResult,
+ OpTrait::NOperands<2>::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<BufferAssignmentTestCopyOp, OpTrait::ZeroResult,
+ OpTrait::NOperands<2>::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<BufferAssignmentTestCopyOp, BufferAssignmentTestUnaryOp,
+ BufferAssignmentTestUnaryLoweredOp>();
+ }
+ static StringRef getDialectNamespace() { return "buffer_assignment_test"; }
+};
+
/// This pass tests two provided operation converters,
/// FunctionAndBlockSignatureConverter and NonVoidToVoidReturnOpConverter, for
/// Buffer Assignment.
struct BufferAssignmentPreparationTestPass
: mlir::PassWrapper<BufferAssignmentPreparationTestPass, FunctionPass> {
- /// This dialect independent unary operation has been defined only for testing
- /// buffer assignment.
- class BufferAssignmentTestUnaryOp
- : public Op<BufferAssignmentTestUnaryOp, OpTrait::OneResult,
- OpTrait::OneOperand> {
- 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<BufferAssignmentTestUnaryLoweredOp, OpTrait::ZeroResult,
- OpTrait::NOperands<2>::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<BufferAssignmentTestCopyOp, OpTrait::ZeroResult,
- OpTrait::NOperands<2>::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 @@
}
};
};
+
} // namespace
+static mlir::DialectRegistration<BufferAssignmentTestDialect>
+ buffer_assignment_test_ops;
+
/// This pass tests helper methods such as computeAllocPosition,
/// FunctionAndBlockSignatureConverter, NonVoidToVoidReturnOpConverter
/// conversion patterns. Furthermore, it checks buffer-assignment pass that
