Warn when using persistent cache with debug env vars (#7175)

pytorch · Jun 4, 2024 · bcf46c7 · bcf46c7
1 parent 94fa999
commit bcf46c7
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/torch_xla/csrc/runtime/env_hash.cc b/torch_xla/csrc/runtime/env_hash.cc
@@ -93,7 +93,8 @@ torch::lazy::hash_t HashXlaEnvVars() {
   // Both XLA_FLAGS and LIBTPU_INIT_ARGS contain XLA flags which impact
   // the compilation result.
   static std::vector<std::string> flag_vars = {"XLA_FLAGS", "LIBTPU_INIT_ARGS"};
-  static std::vector<std::string> raw_vars = {"TPU_MEGACORE"};
+  static std::vector<std::string> raw_vars = {"TPU_MEGACORE", "XLA_HLO_DEBUG",
+                                              "XLA_IR_DEBUG"};
   return hash_xla_env_vars(flag_vars, raw_vars);
 }
 

diff --git a/torch_xla/csrc/xla_graph_executor.cpp b/torch_xla/csrc/xla_graph_executor.cpp
@@ -100,6 +100,13 @@ XLAGraphExecutor::ComputationCache* CreateComputationCache() {
       return std::make_shared<XLAGraphExecutor::CachedComputation>(
           computation, /*is_sharded=*/UseVirtualDevice());
     };
+    if (runtime::sys_util::GetEnvBool("XLA_HLO_DEBUG", false) ||
+        runtime::sys_util::GetEnvBool("XLA_IR_DEBUG", false)) {
+      TF_LOG(WARNING)
+          << "Using persistent compilation cache with XLA_HLO_DEBUG=1 "
+             "or XLA_IR_DEBUG=1 is not recommended. Changes to the HLO "
+             "metadata will not be reflected in loaded executables.";
+    }
     return new XLAGraphExecutor::PersistentCache(
         kMaxCacheSize, persistentCacheDir, readonlyPersistentCache,
         serialize_fn, deserialize_fn);