From 875b3e69655c832c9f0f6733efd27b00e9f2cf56 Mon Sep 17 00:00:00 2001 From: Bixia Zheng Date: Wed, 3 Apr 2019 15:09:26 -0700 Subject: [PATCH] [XLA:CPU/GPU] Run ZeroSizedHloElimination pass to remove zero-sized HLOs from the input module. This is to avoid the need of making other passes, such as algebraic simplification to handle zero-size HLOs. PiperOrigin-RevId: 241814646 --- tensorflow/compiler/xla/service/cpu/cpu_compiler.cc | 5 +++++ tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc | 4 ++++ 2 files changed, 9 insertions(+) diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc index 0f5ac794226..cae1096f8dc 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc @@ -252,6 +252,11 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn( HloPassPipeline pipeline("HLO passes through layout assignment"); pipeline.AddInvariantChecker(/*layout_sensitive=*/false, /*allow_mixed_precision=*/false); + + // Remove zero-sized HLO from the input so that other passes don't have to + // handle it. + pipeline.AddPass(); + pipeline.AddPass(); pipeline.AddPass(); diff --git a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc index 89ef085ae62..31d76906802 100644 --- a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc @@ -183,6 +183,10 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, HloPassPipeline pipeline("optimization"); pipeline.AddInvariantChecker(/*layout_sensitive=*/false, /*allow_mixed_precision=*/false); + // Remove zero-sized HLO from the input so that other passes don't have to + // handle it. + pipeline.AddPass(); + pipeline.AddPass(); pipeline.AddPass(); ReducePrecisionInsertion::AddPasses(