[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

@@ -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"],
 )
 

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: buffer_assignment_test.copy
 // CHECK-NEXT: dealloc

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<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 @@ struct BufferAssignmentPreparationTestPass
     }
   };
 };
+
 }  // 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