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

[XLA] Stop verifying the HLO inside the pass pipeline on optimized builds.

Browse Source 
Keep performing the verification steps on debug builds. This reduces the overhead of HLO verification on opt builds.

PiperOrigin-RevId: 304403425
Change-Id: Ia36f3b5e9bec6ec3eb2a005b9bb45afadb45acd7
This commit is contained in:
Thomas Joerg 2020-04-02 08:17:07 -07:00 committed by TensorFlower Gardener
parent 4b03ed1391
commit 3a6e2b355e
5 changed files with 30 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
View File

@ -277,8 +277,8 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn(
{ {
auto& pass = auto& pass =
pipeline.AddPass<HloPassFix<HloPassPipeline>>("simplification"); pipeline.AddPass<HloPassFix<HloPassPipeline>>("simplification");
pass.AddInvariantChecker<HloVerifier>(/*layout_sensitive=*/false, pass.AddInvariantCheckerDebug<HloVerifier>(/*layout_sensitive=*/false,
/*allow_mixed_precision=*/false); /*allow_mixed_precision=*/false);
pass.AddPass<TreeReductionRewriter>(); pass.AddPass<TreeReductionRewriter>();
pass.AddPass<ScatterExpander>(); pass.AddPass<ScatterExpander>();
@ -339,18 +339,18 @@ Status CpuCompiler::RunHloPassesAfterLayoutAssn(
LLVMTargetMachineFeatures* target_machine_features) { LLVMTargetMachineFeatures* target_machine_features) {
HloPassPipeline pipeline("HLO passes after layout assignment"); HloPassPipeline pipeline("HLO passes after layout assignment");
// After layout assignment, use a layout-sensitive verifier. // After layout assignment, use a layout-sensitive verifier.
auto& after_layout_assn =
pipeline.AddPass<HloPassPipeline>("after layout assignment"); pipeline.AddPass<HloPassPipeline>("after layout assignment")
after_layout_assn.AddInvariantChecker<HloVerifier>( .AddInvariantCheckerDebug<HloVerifier>(
/*layout_sensitive=*/true, /*layout_sensitive=*/true,
/*allow_mixed_precision=*/false); /*allow_mixed_precision=*/false);
// The LayoutAssignment pass may leave behind kCopy instructions which are // The LayoutAssignment pass may leave behind kCopy instructions which are
// duplicate or NOPs, so remove them with algebraic simplification and CSE. // duplicate or NOPs, so remove them with algebraic simplification and CSE.
{ {
auto& pass = pipeline.AddPass<HloPassFix<HloPassPipeline>>( auto& pass = pipeline.AddPass<HloPassFix<HloPassPipeline>>(
"simplification after layout assignment"); "simplification after layout assignment");
pass.AddInvariantChecker<HloVerifier>( pass.AddInvariantCheckerDebug<HloVerifier>(
/*layout_sensitive=*/true, /*layout_sensitive=*/true,
/*allow_mixed_precision=*/false, /*allow_mixed_precision=*/false,
LayoutAssignment::InstructionCanChangeLayout); LayoutAssignment::InstructionCanChangeLayout);

5
tensorflow/compiler/xla/service/gpu/amdgpu_compiler.cc
View File

@ -76,8 +76,9 @@ Status AMDGPUCompiler::OptimizeHloConvolutionCanonicalization(
// Convert convolutions into CustomCalls to MIOpen, then canonicalize them // Convert convolutions into CustomCalls to MIOpen, then canonicalize them
// (PadInsertion). // (PadInsertion).
HloPassPipeline pipeline("conv_canonicalization"); HloPassPipeline pipeline("conv_canonicalization");
pipeline.AddInvariantChecker<HloVerifier>(/*layout_sensitive=*/false, pipeline.AddInvariantCheckerDebug<HloVerifier>(
/*allow_mixed_precision=*/false); /*layout_sensitive=*/false,
/*allow_mixed_precision=*/false);
pipeline.AddPass<GpuConvRewriter>(); pipeline.AddPass<GpuConvRewriter>();
pipeline.AddPass<GpuConvPaddingLegalization>(); pipeline.AddPass<GpuConvPaddingLegalization>();

11
tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
View File

@ -178,8 +178,9 @@ Status GpuCompiler::OptimizeHloModule(
{ {
auto& pass = auto& pass =
pipeline.AddPass<HloPassFix<HloPassPipeline>>("simplification"); pipeline.AddPass<HloPassFix<HloPassPipeline>>("simplification");
pass.AddInvariantChecker<HloVerifier>(/*layout_sensitive=*/false, pass.AddInvariantCheckerDebug<HloVerifier>(
/*allow_mixed_precision=*/false); /*layout_sensitive=*/false,
/*allow_mixed_precision=*/false);
// If cudnn batchnorms are enabled, rewrite batchnorm HLOs to cudnn calls // If cudnn batchnorms are enabled, rewrite batchnorm HLOs to cudnn calls
// where possible. Not every batchnorm op can be implemented as a call to // where possible. Not every batchnorm op can be implemented as a call to
@ -288,7 +289,7 @@ Status GpuCompiler::OptimizeHloModule(
fusion.AddPass<VariadicOpSplitter>(); fusion.AddPass<VariadicOpSplitter>();
/* TODO(b/117531509): Use LayoutAssignment::InstructionCanChangeLayout after /* TODO(b/117531509): Use LayoutAssignment::InstructionCanChangeLayout after
* fixing the ticket. */ * fixing the ticket. */
fusion.AddInvariantChecker<HloVerifier>( fusion.AddInvariantCheckerDebug<HloVerifier>(
/*layout_sensitive=*/true, /*layout_sensitive=*/true,
/*allow_mixed_precision=*/false, /*allow_mixed_precision=*/false,
LayoutAssignment::InstructionCanChangeLayout); LayoutAssignment::InstructionCanChangeLayout);
@ -336,7 +337,7 @@ Status GpuCompiler::PrepareHloModuleForIrEmitting(HloModule* hlo_module) {
HloPassPipeline pipeline("GPU-ir-emit-prepare"); HloPassPipeline pipeline("GPU-ir-emit-prepare");
/* TODO(b/117531509): Use LayoutAssignment::InstructionCanChangeLayout after /* TODO(b/117531509): Use LayoutAssignment::InstructionCanChangeLayout after
* fixing the ticket. */ * fixing the ticket. */
pipeline.AddInvariantChecker<HloVerifier>( pipeline.AddInvariantCheckerDebug<HloVerifier>(
/*layout_sensitive=*/true, /*layout_sensitive=*/true,
/*allow_mixed_precision=*/false, /*allow_mixed_precision=*/false,
LayoutAssignment::InstructionCanChangeLayout); LayoutAssignment::InstructionCanChangeLayout);
@ -375,7 +376,7 @@ Status GpuCompiler::OptimizeHloPostLayoutAssignment(
HloPassPipeline pipeline("post-layout_assignment"); HloPassPipeline pipeline("post-layout_assignment");
/* TODO(b/117531509): Use LayoutAssignment::InstructionCanChangeLayout after /* TODO(b/117531509): Use LayoutAssignment::InstructionCanChangeLayout after
* fixing the ticket. */ * fixing the ticket. */
pipeline.AddInvariantChecker<HloVerifier>( pipeline.AddInvariantCheckerDebug<HloVerifier>(
/*layout_sensitive=*/true, /*layout_sensitive=*/true,
/*allow_mixed_precision=*/false, /*allow_mixed_precision=*/false,
LayoutAssignment::InstructionCanChangeLayout); LayoutAssignment::InstructionCanChangeLayout);

9
tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc
View File

@ -109,8 +109,9 @@ Status NVPTXCompiler::OptimizeHloConvolutionCanonicalization(
// Convert convolutions into CustomCalls to cudnn, then canonicalize them // Convert convolutions into CustomCalls to cudnn, then canonicalize them
// (GpuConvPaddingLegalization). Also expand cuSolver calls. // (GpuConvPaddingLegalization). Also expand cuSolver calls.
HloPassPipeline pipeline("conv_canonicalization"); HloPassPipeline pipeline("conv_canonicalization");
pipeline.AddInvariantChecker<HloVerifier>(/*layout_sensitive=*/false, pipeline.AddInvariantCheckerDebug<HloVerifier>(
/*allow_mixed_precision=*/false); /*layout_sensitive=*/false,
/*allow_mixed_precision=*/false);
pipeline.AddPass<CusolverRewriter>(); pipeline.AddPass<CusolverRewriter>();
pipeline.AddPass<GpuConvRewriter>(); pipeline.AddPass<GpuConvRewriter>();
pipeline.AddPass<CudnnFusedConvRewriter>(); pipeline.AddPass<CudnnFusedConvRewriter>();
@ -127,8 +128,8 @@ Status NVPTXCompiler::OptimizeHloConvolutionCanonicalization(
{ {
auto& pass = pipeline.AddPass<HloPassFix<HloPassPipeline>>( auto& pass = pipeline.AddPass<HloPassFix<HloPassPipeline>>(
"algebraic_simplification_post_conv_rewriter"); "algebraic_simplification_post_conv_rewriter");
pass.AddInvariantChecker<HloVerifier>(/*layout_sensitive=*/false, pass.AddInvariantCheckerDebug<HloVerifier>(/*layout_sensitive=*/false,
/*allow_mixed_precision=*/false); /*allow_mixed_precision=*/false);
AlgebraicSimplifierOptions options; AlgebraicSimplifierOptions options;
// When transposes appear in a fusion node, we can easily adjust the // When transposes appear in a fusion node, we can easily adjust the

8
tensorflow/compiler/xla/service/hlo_pass_pipeline.h
View File

@ -70,6 +70,14 @@ class HloPassPipeline : public HloPassInterface {
return *pass; return *pass;
} }
// Add an invariant-checking pass to the pipeline on debug builds only.
template <typename T, typename... Args>
void AddInvariantCheckerDebug(Args&&... args) {
#ifndef NDEBUG
AddInvariantChecker<T>(std::forward<Args>(args)...);
#endif // NDEBUG
}
StatusOr<bool> Run(HloModule* module) override; StatusOr<bool> Run(HloModule* module) override;
StatusOr<bool> RunOnModuleGroup(HloModuleGroup* module_group) override; StatusOr<bool> RunOnModuleGroup(HloModuleGroup* module_group) override;