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 HloProgram::Fingerprint() const { tsl::StatusScopedDiagnosticHandler diag_handler(mlir_module_->getContext()); mlir::BytecodeWriterConfig config; - config.attachAttributeCallback( - [](mlir::Attribute attr, - std::optional& group_name_override, - mlir::DialectBytecodeWriter& writer) -> mlir::LogicalResult { - if (llvm::isa_and_nonnull(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 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 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} {