[XLA] fail if the change bit reported by a pass does not agree with HLO hash comparison

Usually we trust that pass truthfully reports if it did update the HLO.
Shockingly there are cases, e.g. algsimp, that wrongly report "unchanged" pass result.
Having a flag makes it easier to detect existing and new cases of such bugs.

PiperOrigin-RevId: 738733528

Author: metaflow

xla/debug_options_flags.cc

@@ -337,6 +337,8 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_hlo_pass_fix_detect_cycles(false);
   opts.set_xla_gpu_experimental_enable_sync_collective_combining(false);
   opts.set_xla_allow_get_default_platform(true);
+  opts.set_xla_unsupported_crash_on_hlo_pass_silent_hlo_change(false);
+  opts.set_xla_unsupported_crash_on_hlo_pass_noop_change(false);
   return opts;
 }
 
@@ -2311,7 +2313,22 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       "If non empty will interpret this variable as a path for performance "
       "tables for collectives. Expects `xla.gpu.DeviceHloInstructionProfiles` "
       "proto."));
-}  // NOLINT(readability/fn_size)1
+  flag_list->push_back(tsl::Flag(
+      "xla_unsupported_crash_on_hlo_pass_silent_hlo_change",
+      bool_setter_for(
+          &DebugOptions::
+              set_xla_unsupported_crash_on_hlo_pass_silent_hlo_change),
+      debug_options->xla_unsupported_crash_on_hlo_pass_silent_hlo_change(),
+      "Crash if a pass reports that it did not change the HLO but in fact it "
+      "did."));
+  flag_list->push_back(tsl::Flag(
+      "xla_unsupported_crash_on_hlo_pass_noop_change",
+      bool_setter_for(
+          &DebugOptions::set_xla_unsupported_crash_on_hlo_pass_noop_change),
+      debug_options->xla_unsupported_crash_on_hlo_pass_noop_change(),
+      "Crash if a pass reports that it did change the HLO but in fact it "
+      "did not."));
+}  // NOLINT(readability/fn_size)
 
 // Allocates flag_values and flag_objects; this function must not be called more
 // than once - its call done via call_once.

xla/hlo/pass/BUILD

@@ -89,6 +89,7 @@ cc_library(
         "//xla:util",
         "//xla:xla_proto_cc",
         "//xla/hlo/ir:hlo",
+        "//xla/hlo/ir:hlo_module_group",
         "//xla/service:compilation_stats",
         "//xla/service:dump",
         "//xla/service:hlo_graph_dumper",
@@ -102,6 +103,7 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/types:optional",
         "@tsl//tsl/profiler/lib:scoped_annotation",
         "@tsl//tsl/profiler/lib:traceme",
     ],

xla/hlo/pass/hlo_pass_pipeline.cc

@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "xla/hlo/pass/hlo_pass_pipeline.h"
 
+#include <cstddef>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -24,6 +26,8 @@ limitations under the License.
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_module_group.h"
+#include "xla/hlo/pass/hlo_pass_interface.h"
 #include "xla/service/dump.h"
 #include "xla/service/hlo_graph_dumper.h"
 #include "xla/service/hlo_proto_util.h"
@@ -169,6 +173,9 @@ absl::StatusOr<bool> HloPassPipeline::RunPassesInternal(
                         /*module_changed=*/false);
 
   bool changed = false;
+  bool verify_pass_changed_report =
+      debug_options.xla_unsupported_crash_on_hlo_pass_silent_hlo_change() ||
+      debug_options.xla_unsupported_crash_on_hlo_pass_noop_change();
   for (int i = 0; i < passes.size(); i++) {
     HloPassInterface* pass = passes[i];
     std::string pass_name = std::string(pass->name());
@@ -178,7 +185,11 @@ absl::StatusOr<bool> HloPassPipeline::RunPassesInternal(
                              pass_name, hlo->name(), UniqueId(*hlo));
     }};
     VLOG(1) << "  HLO pass " << pass_name;
-    VLOG(2) << "  Module hash " << absl::HashOf(*hlo);
+    std::optional<size_t> hash_before = std::nullopt;
+    if (verify_pass_changed_report || VLOG_IS_ON(2)) {
+      hash_before = absl::HashOf(*hlo);
+      VLOG(2) << "  Module hash " << hash_before.value();
+    }
     tsl::profiler::TraceMe traceme(pass->name());
     if (!pass->IsPassPipeline()) {
       compilation_stats_->StartPass(pass_name);
@@ -190,6 +201,27 @@ absl::StatusOr<bool> HloPassPipeline::RunPassesInternal(
           pass_name, absl::StatusCodeToString(status.code()));
     }
     TF_ASSIGN_OR_RETURN(bool pass_changed, status_or_changed);
+    if (verify_pass_changed_report) {
+      size_t hash_after = absl::HashOf(*hlo);
+      // Fail if pass changed HLO but has reported that it didn't.
+      if (!pass_changed && hash_after != hash_before &&
+          debug_options.xla_unsupported_crash_on_hlo_pass_silent_hlo_change()) {
+        LOG(FATAL) << absl::StrFormat(
+            "Pass '%s' in pipeline '%s' reported that it did not change the "
+            "HLO but the hash of HLO was changed from %d to %d. HLO text "
+            "after:\n%s",
+            pass_name, pipeline_name, hash_before.value(), hash_after,
+            hlo->ToString());
+      }
+      // Fail if pass did not change HLO but has reported that it did.
+      if (pass_changed && hash_after == hash_before &&
+          debug_options.xla_unsupported_crash_on_hlo_pass_noop_change()) {
+        LOG(FATAL) << absl::StrFormat(
+            "Pass '%s' in pipeline '%s' reported that it changed the HLO but "
+            "the hash of HLO was not updated. HLO text after:\n%s",
+            pass_name, pipeline_name, hlo->ToString());
+      }
+    }
     if (!dump_regex.empty() && (pass_changed || dump_regex != ".*")) {
       MaybeDumpHloAndSaveFilenames(*hlo,
                                    /*after_pass_name=*/pass_name,

xla/xla.proto

@@ -125,6 +125,10 @@ message DebugOptions {
   bool xla_hlo_pass_fix_detect_cycles = 370;
   // Crash if HloPassFix can not converge after a fixed number of iterations.
   bool xla_unsupported_crash_on_hlo_pass_fix_max_iterations = 363;
+  // Crash if a pass reports that it changes the HLO but in fact it did not.
+  bool xla_unsupported_crash_on_hlo_pass_noop_change = 379;
+  // Crash if a pass reports that it did not change the HLO but in fact it did.
+  bool xla_unsupported_crash_on_hlo_pass_silent_hlo_change = 380;
   // go/keep-sorted end
 
   reserved 346;  // xla_experimental_exec_time_optimization_effort
@@ -1189,7 +1193,7 @@ message DebugOptions {
   // Note: when adding a new flag, please add it to one of the hardware-specific
   // or hardware-agnostic sections at the top of this proto message.
 
-  // Next id: 379
+  // Next id: 381
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.