[IFRT] Fix incorrect debug info stripping in HloProgram::Fingerprint().

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.

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} {