From 7fa7aaebd2bd4e9caff2cedcb87ab80144b39e7e Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Sat, 13 Jun 2026 07:46:23 -0700
Subject: [PATCH] [IFRT] Fix incorrect debug info stripping in
 HloProgram::Fingerprint().

It's not enough to strip attributes during the attribute printer because attributes are numbered before they are printed, leading to otherwise equivalent modules having different fingerprints. Instead, we must also elide them during attribute numbering.

PiperOrigin-RevId: 931652173
---
 xla/python/ifrt/hlo/hlo_program.cc      | 13 +------------
 xla/python/ifrt/hlo/hlo_program_test.cc | 23 +++++++++++++++++++++++
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/xla/python/ifrt/hlo/hlo_program.cc b/xla/python/ifrt/hlo/hlo_program.cc
index 49b131c16cf6f..530c0f24369b7 100644
--- a/xla/python/ifrt/hlo/hlo_program.cc
+++ b/xla/python/ifrt/hlo/hlo_program.cc
@@ -168,18 +168,7 @@ absl::StatusOr<uint64_t> HloProgram::Fingerprint() const {
   tsl::StatusScopedDiagnosticHandler diag_handler(mlir_module_->getContext());
 
   mlir::BytecodeWriterConfig config;
-  config.attachAttributeCallback(
-      [](mlir::Attribute attr,
-         std::optional<llvm::StringRef>& group_name_override,
-         mlir::DialectBytecodeWriter& writer) -> mlir::LogicalResult {
-        if (llvm::isa_and_nonnull<mlir::LocationAttr>(attr)) {
-          // Ignore location attributes since they are for debugging only and
-          // do not affect the semantics of the program.
-          return mlir::success();
-        }
-        // Fall back to the default implementation.
-        return mlir::failure();
-      });
+  config.setElideLocations(true);
 
   // Use a version before `kUseListOrdering` due to an MLIR bug where use list
   // ordering is not stable.
diff --git a/xla/python/ifrt/hlo/hlo_program_test.cc b/xla/python/ifrt/hlo/hlo_program_test.cc
index a30e38d1418af..4eecef926d8ae 100644
--- a/xla/python/ifrt/hlo/hlo_program_test.cc
+++ b/xla/python/ifrt/hlo/hlo_program_test.cc
@@ -102,6 +102,29 @@ module @foo {
   EXPECT_EQ(hlo_program1->Fingerprint(), hlo_program2->Fingerprint());
 }
 
+TEST(HloProgramTest, FingerprintIgnoresDebugInfoStructure) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      const std::unique_ptr<xla::ifrt::HloProgram> hlo_program1,
+      ParseHloProgramString(R"(
+module @foo {
+  func.func @main(%arg0: tensor<2x3xi32> loc("foo")) -> tensor<2x3xi32> {
+    return %arg0 : tensor<2x3xi32> loc("foo")
+  } loc("foo")
+} loc("foo")
+)"));
+  TF_ASSERT_OK_AND_ASSIGN(
+      const std::unique_ptr<xla::ifrt::HloProgram> hlo_program2,
+      ParseHloProgramString(R"(
+module @foo {
+  func.func @main(%arg0: tensor<2x3xi32> loc("bar")) -> tensor<2x3xi32> {
+    return %arg0 : tensor<2x3xi32> loc("baz")
+  } loc("qux")
+} loc("quux")
+)"));
+
+  EXPECT_EQ(hlo_program1->Fingerprint(), hlo_program2->Fingerprint());
+}
+
 TEST(HloProgramTest, BytesRoundTrip) {
   static constexpr absl::string_view kModule = R"(
 module @hlo_module attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {