[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

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() &&
+  string generate_hlo_graph_regex;
-      RE2::PartialMatch(module.name(), flags->xla_generate_hlo_graph)) {
+  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);

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

tensorflow/compiler/xla/service/service.cc

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

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 {