experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA][MLIR] Register BufferAssignmentTestDialect.

Browse Source 
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
This commit is contained in:
Alexander Belyaev 2020-05-28 06:42:39 -07:00 committed by TensorFlower Gardener
parent 6688e1c23b
commit ea8e87c8e9
3 changed files with 68 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/compiler/mlir/xla/tests/BUILD
View File

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

16
tensorflow/compiler/mlir/xla/tests/buffer-assignment.mlir
View File

@ -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

36
tensorflow/compiler/mlir/xla/transforms/buffer_assignment_test.cc
View File

@ -29,11 +29,7 @@ limitations under the License.
namespace mlir {
namespace xla {
namespace {
/// 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
@ -41,9 +37,7 @@ struct BufferAssignmentPreparationTestPass
OpTrait::OneOperand> {
public:
using Op::Op;
static StringRef getOperationName() {
return "buffer_assignment_test.unary";
}
static StringRef getOperationName() { return "buffer_assignment_test.unary"; }
static void build(OpBuilder& b, OperationState& state, Value source) {
state.addOperands(source);
}
@ -73,16 +67,28 @@ struct BufferAssignmentPreparationTestPass
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) {
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> {
/// 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