[XLA] Add an option to set the "generate HLO graph" regex without a flag.

Pipes the option through xla.proto ExecutionOptions, to HloModuleConfig, which
can then be accessed throughout the compiler.

PiperOrigin-RevId: 157615458
This commit is contained in:
Eli Bendersky 2017-05-31 11:59:07 -07:00 committed by TensorFlower Gardener
parent 84661fa736
commit 2de94bbb89
4 changed files with 29 additions and 2 deletions

View File

@ -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);

View File

@ -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

View File

@ -323,6 +323,7 @@ StatusOr<std::unique_ptr<HloModuleConfig>> 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);
}

View File

@ -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 {