From 4d5197c63c7c625fed31b66ed03f086feee0b430 Mon Sep 17 00:00:00 2001 From: wangsiyu Date: Fri, 15 Jan 2021 17:31:15 +0800 Subject: [PATCH 1/3] [BugFix] Fix bug when reduce dimension becomes emptry in reduction_degenerate_dim_remover pass --- .../xla/service/gpu/reduction_degenerate_dim_remover.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorflow/compiler/xla/service/gpu/reduction_degenerate_dim_remover.cc b/tensorflow/compiler/xla/service/gpu/reduction_degenerate_dim_remover.cc index 2c786b577fc..3defd4ccbff 100644 --- a/tensorflow/compiler/xla/service/gpu/reduction_degenerate_dim_remover.cc +++ b/tensorflow/compiler/xla/service/gpu/reduction_degenerate_dim_remover.cc @@ -65,6 +65,12 @@ class ReductionDegenerateDimRemoverVisitor : public DfsHloRewriteVisitor { } } + if (updated_reduced_dimensions.empty()) { + std::unique_ptr reshape = HloInstruction::CreateBitcast( + reduce_shape, reduced_op); + return ReplaceWithNewInstruction(instr, std::move(reshape)); + } + HloInstruction *input_reshape = instr->parent()->AddInstruction( HloInstruction::CreateBitcast(canonical_input_shape, reduced_op)); From 09df86ebe40eb13704ad9d73039986351fd37a69 Mon Sep 17 00:00:00 2001 From: wangsiyu Date: Sun, 17 Jan 2021 15:14:27 +0800 Subject: [PATCH 2/3] Add unit test for reduction remover when reduction dimension is empty --- .../reduction_degenerate_dim_remover_test.cc | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc b/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc index 92f558ee98d..a84185c8e3c 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc @@ -69,6 +69,36 @@ ENTRY main { )"); } +TEST_F(ReductionDegenerateDimRemoverTest, DegenerateWithEmptyDimension) { + const char* hlo_text = R"( +HloModule ReduceWithDegenerateDimensions + +add { + accum = f32[] parameter(0) + op = f32[] parameter(1) + ROOT out = f32[] add(accum, op) +} + +ENTRY main { + input = f32[1,3,1,4,1,5,1] parameter(0) + zero = f32[] constant(0) + + ROOT out = f32[3,4,5,1] reduce(input, zero), dimensions={0,2,4}, to_apply=add +} + +)"; + + EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{1e-5, 1e-5})); + // Copy insertion + MatchOptimizedHloWithShapes(hlo_text, + R"( +// CHECK: ENTRY %main (input: f32[1,3,1,4,1,5,1]) -> f32[3,4,5,1] { +// CHECK: %input = f32[1,3,1,4,1,5,1]{6,5,4,3,2,1,0} parameter(0) +// CHECK: %bitcast{{.+}} = f32[3,4,5,1]{3,2,1,0} bitcast(f32[1,3,1,4,1,5,1]{6,5,4,3,2,1,0} %input) +// CHECK: ROOT %copy{{.+}} = f32[3,4,5,1]{3,2,1,0} copy(f32[3,4,5,1]{3,2,1,0} %bitcast{{.+}}) + )"); +} + } // namespace } // namespace gpu } // namespace xla From f519f072f45c7871ace1f317756edaf8b7188463 Mon Sep 17 00:00:00 2001 From: wangsiyu Date: Sun, 17 Jan 2021 15:20:21 +0800 Subject: [PATCH 3/3] Add more comments --- .../gpu/tests/reduction_degenerate_dim_remover_test.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc b/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc index a84185c8e3c..f5031817818 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc @@ -89,7 +89,9 @@ ENTRY main { )"; EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{1e-5, 1e-5})); - // Copy insertion + // Copy instruction is added after bitcast because of copy-insertion pass, + // so we check the entire hlo module to verify there is no reduce instruction + // in this case. MatchOptimizedHloWithShapes(hlo_text, R"( // CHECK: ENTRY %main (input: f32[1,3,1,4,1,5,1]) -> f32[3,4,5,1] {