diff --git a/tensorflow/compiler/xla/service/executable.cc b/tensorflow/compiler/xla/service/executable.cc index e0cf4b19f70..3a9f8dc79ee 100644 --- a/tensorflow/compiler/xla/service/executable.cc +++ b/tensorflow/compiler/xla/service/executable.cc @@ -30,8 +30,15 @@ namespace xla { const HloExecutionProfile* profile) { VLOG(2) << "module name = " << module.name(); legacy_flags::ServiceFlags* flags = legacy_flags::GetServiceFlags(); - if (!flags->xla_generate_hlo_graph.empty() && - RE2::PartialMatch(module.name(), flags->xla_generate_hlo_graph)) { + string generate_hlo_graph_regex; + if (!flags->xla_generate_hlo_graph.empty()) { + generate_hlo_graph_regex = flags->xla_generate_hlo_graph; + } else { + generate_hlo_graph_regex = + module.config().debug_options().xla_generate_hlo_graph(); + } + if (!generate_hlo_graph_regex.empty() && + RE2::PartialMatch(module.name(), generate_hlo_graph_regex)) { hlo_graph_dumper::DumpGraph(*module.entry_computation(), label, flags->xla_hlo_graph_addresses, flags->xla_hlo_graph_layout, profile); diff --git a/tensorflow/compiler/xla/service/hlo_module_config.h b/tensorflow/compiler/xla/service/hlo_module_config.h index c750da940c0..023102d8ff7 100644 --- a/tensorflow/compiler/xla/service/hlo_module_config.h +++ b/tensorflow/compiler/xla/service/hlo_module_config.h @@ -94,6 +94,12 @@ class HloModuleConfig { // executable. string compilation_cache_key() const; + const DebugOptions& debug_options() const { return debug_options_; } + + void set_debug_options(const DebugOptions& debug_options) { + debug_options_ = debug_options; + } + private: // If you add new members, be sure to update compilation_cache_key. @@ -119,6 +125,8 @@ class HloModuleConfig { int64 replica_count_ = 1; bool fast_math_disabled_ = false; + + DebugOptions debug_options_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc index 59d828acf23..327ca9cb4f5 100644 --- a/tensorflow/compiler/xla/service/service.cc +++ b/tensorflow/compiler/xla/service/service.cc @@ -323,6 +323,7 @@ StatusOr> Service::CreateModuleConfig( module_config->set_replica_count(backend->Replicas().size()); module_config->set_fast_math_disabled(execution_options.disable_fast_math()); module_config->set_seed(execution_options.seed()); + module_config->set_debug_options(execution_options.debug_options()); return std::move(module_config); } diff --git a/tensorflow/compiler/xla/xla.proto b/tensorflow/compiler/xla/xla.proto index 39d095d9461..716eb424424 100644 --- a/tensorflow/compiler/xla/xla.proto +++ b/tensorflow/compiler/xla/xla.proto @@ -20,6 +20,15 @@ import "tensorflow/compiler/xla/service/session.proto"; package xla; +// Debugging options for XLA. These options may change at any time - there are +// no guarantees about backward or forward compatibility for these fields. +message DebugOptions { + // HLO modules matching this regex will be dumped to a .dot file throughout + // various stages in compilation (file names are LOG(INFO)'d). Set to ".*" to + // dump *all* HLO modules. + string xla_generate_hlo_graph = 1; +} + // These settings control how XLA compiles and/or runs code. Not all settings // will have an effect on every platform. // @@ -46,6 +55,8 @@ message ExecutionOptions { // // TODO(b/32083678): Changing the seed unnecessarily forces a recompilation. uint64 seed = 3; + + DebugOptions debug_options = 4; } message SnapshotComputationRequest {