[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:
parent
84661fa736
commit
2de94bbb89
@ -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);
|
||||
|
@ -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
|
||||
|
@ -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);
|
||||
}
|
||||
|
@ -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 {
|
||||
|
Loading…
Reference in New Issue
Block a user