Be more selective when disabling mlir_gpu_lhlo_gen_test.

PiperOrigin-RevId: 294653870
Change-Id: I6e0dcf432f630550e49b776390a738c013a1e891
This commit is contained in:
Stephan Herhut 2020-02-12 06:27:31 -08:00 committed by TensorFlower Gardener
parent dc5443ca68
commit 675799d105
2 changed files with 54 additions and 52 deletions
tensorflow/compiler/xla/service/mlir_gpu/tests

View File

@ -18,14 +18,14 @@ package_group(
],
)
# tf_cc_test(
# name = "mlir_gpu_lhlo_gen_test",
# srcs = ["mlir_gpu_lhlo_gen_test.cc"],
# tags = tf_cuda_tests_tags() + ["no_rocm"],
# deps = [
# "//tensorflow/compiler/xla/service:gpu_plugin_mlir",
# "//tensorflow/compiler/xla/service/mlir_gpu:mlir_irgen_test_base",
# "//tensorflow/core:test_main",
# "//tensorflow/stream_executor/lib",
# ],
# )
tf_cc_test(
name = "mlir_gpu_lhlo_gen_test",
srcs = ["mlir_gpu_lhlo_gen_test.cc"],
tags = tf_cuda_tests_tags() + ["no_rocm"],
deps = [
"//tensorflow/compiler/xla/service:gpu_plugin_mlir",
"//tensorflow/compiler/xla/service/mlir_gpu:mlir_irgen_test_base",
"//tensorflow/core:test_main",
"//tensorflow/stream_executor/lib",
],
)

View File

@ -240,7 +240,8 @@ ENTRY %Add (x: f32[2,2], y: f32[2,2]) -> f32[2,2] {
LoweringStage::KERNEL);
}
TEST_F(LhloGenTest, AddMultiply) {
// TODO(b/149302060) Reenable once fusion is fixed.
TEST_F(LhloGenTest, DISABLED_AddMultiply) {
CompileAndVerifyIr(R"(
HloModule AddMultiply
@ -257,7 +258,7 @@ ENTRY %AddMultiply (x: f32[2,2], y: f32[2,2], z: f32[2,2]) -> f32[2,2] {
;CHECK: %[[REF0:.*]] = tensor_load %[[ARG0]] : [[TYPE]]
;CHECK: %[[REF1:.*]] = tensor_load %[[ARG1]] : [[TYPE]]
;CHECK: %[[REF2:.*]] = tensor_load %[[ARG2]] : [[TYPE]]
;CHECK: %[[ADD:.*]] = xla_hlo.add %[[REF1]], %[[REF2]]
;CHECK: %[[ADD:.*]] = xla_hlo.add %[[REF1]], %[[REF2]]
;CHECK: %[[MUL:.*]] = xla_hlo.mul %[[ADD]], %[[REF0]]
;CHECK: tensor_store %[[MUL]], %[[RESULT]]
;CHECK: "xla_lhlo.terminator"()
@ -265,7 +266,8 @@ ENTRY %AddMultiply (x: f32[2,2], y: f32[2,2], z: f32[2,2]) -> f32[2,2] {
)");
}
TEST_F(LhloGenTest, IotaAddMultiply) {
// TODO(b/149302060) Reenable once fusion is fixed.
TEST_F(LhloGenTest, DISABLED_IotaAddMultiply) {
CompileAndVerifyIr(R"(
HloModule AddMultiply
@ -315,44 +317,44 @@ ENTRY %AddMultiply (x: f32[2,2], y: f32[2,2], z: f32[2,2]) -> f32[2,2] {
}
// TODO(b/137624192): Reenable once we can fuse reductions.
// TEST_F(LhloGenTest, FusedReduce) {
// CompileAndVerifyIr(R"(
// HloModule FusedReduce
//
// %add (x: f32[], y: f32[]) -> f32[] {
// %x = f32[] parameter(0)
// %y = f32[] parameter(1)
// ROOT %add = f32[] add(f32[] %x, f32[] %y)
// }
//
// %fused_computation (param: f32[100,10]) -> f32[10] {
// %param = f32[100,10] parameter(0)
// %constant = f32[] constant(0)
// ROOT %reduce = f32[10]{0} reduce(f32[100,10]{1,0} %param, f32[] %constant),
// dimensions={0}, to_apply=%add
// }
//
// ENTRY %FusedReduce (x: f32[100,10]) -> f32[10] {
// %x = f32[100,10] parameter(0)
// ROOT %fusion = f32[10]{0} fusion(f32[100,10]{1,0} %x), kind=kInput,
// calls=%fused_computation
// }
// )",
// R"(
// ;CHECK: func @fusion(%[[ARG0:.*]]: [[TYPE:.*]], %[[RESULT:.*]]: [[RTYPE:.*]])
// ;CHECK: "xla_lhlo.fusion"() ( {
// ;CHECK: %[[REF0:.*]] = tensor_load %arg0 : [[TYPE]]
// ;CHECK: %[[CT0:.*]] = xla_hlo.constant dense<0.000000e+00>
// ;CHECK: %[[RED:.*]] = "xla_hlo.reduce"(%0, %1) ( {
// ;CHECK: ^bb0(%[[BARG0:.*]]: [[ETYPE:.*]], %[[BARG1:.*]]: [[ETYPE]])
// ;CHECK: %[[ADD:.*]] = xla_hlo.add %[[BARG0]], %[[BARG1]] : [[ETYPE]]
// ;CHECK: "xla_hlo.return"(%[[ADD]])
// ;CHECK: })
// ;CHECK: tensor_store %[[RED]], %[[RESULT]] : [[RTYPE]]
// ;CHECK: "xla_lhlo.terminator"()
// ;CHECK-NEXT: })
// )");
// }
TEST_F(LhloGenTest, DISABLED_FusedReduce) {
CompileAndVerifyIr(R"(
HloModule FusedReduce
%add (x: f32[], y: f32[]) -> f32[] {
%x = f32[] parameter(0)
%y = f32[] parameter(1)
ROOT %add = f32[] add(f32[] %x, f32[] %y)
}
%fused_computation (param: f32[100,10]) -> f32[10] {
%param = f32[100,10] parameter(0)
%constant = f32[] constant(0)
ROOT %reduce = f32[10]{0} reduce(f32[100,10]{1,0} %param, f32[] %constant),
dimensions={0}, to_apply=%add
}
ENTRY %FusedReduce (x: f32[100,10]) -> f32[10] {
%x = f32[100,10] parameter(0)
ROOT %fusion = f32[10]{0} fusion(f32[100,10]{1,0} %x), kind=kInput,
calls=%fused_computation
}
)",
R"(
;CHECK: func @fusion(%[[ARG0:.*]]: [[TYPE:.*]], %[[RESULT:.*]]: [[RTYPE:.*]])
;CHECK: "xla_lhlo.fusion"() ( {
;CHECK: %[[REF0:.*]] = tensor_load %arg0 : [[TYPE]]
;CHECK: %[[CT0:.*]] = xla_hlo.constant dense<0.000000e+00>
;CHECK: %[[RED:.*]] = "xla_hlo.reduce"(%0, %1) ( {
;CHECK: ^bb0(%[[BARG0:.*]]: [[ETYPE:.*]], %[[BARG1:.*]]: [[ETYPE]])
;CHECK: %[[ADD:.*]] = xla_hlo.add %[[BARG0]], %[[BARG1]] : [[ETYPE]]
;CHECK: "xla_hlo.return"(%[[ADD]])
;CHECK: })
;CHECK: tensor_store %[[RED]], %[[RESULT]] : [[RTYPE]]
;CHECK: "xla_lhlo.terminator"()
;CHECK-NEXT: })
)");
}
TEST_F(LhloGenTest, Broadcast) {
CompileAndVerifyIr(R"(