From 3c045ab95c77439981fce1c1090c192032961006 Mon Sep 17 00:00:00 2001
From: Toli Yevtushenko <anatoliyy@google.com>
Date: Tue, 15 Oct 2024 15:22:11 -0700
Subject: [PATCH] More graceful handling AbslHashValue when hlo module doesn't
 have entry computation layout. Without the fix hash breaks on empty module
 such as absl:Hash(HloTestBase::CreateNewVerifiedModule()).

PiperOrigin-RevId: 686257548
---
 xla/hlo/ir/BUILD              | 15 ++++++++++
 xla/hlo/ir/hlo_module.h       |  3 +-
 xla/hlo/ir/hlo_module_test.cc | 52 +++++++++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 xla/hlo/ir/hlo_module_test.cc

diff --git a/xla/hlo/ir/BUILD b/xla/hlo/ir/BUILD
index b13714fbee45e..9ab7a551662a4 100644
--- a/xla/hlo/ir/BUILD
+++ b/xla/hlo/ir/BUILD
@@ -114,6 +114,21 @@ cc_library(
     ],
 )
 
+xla_cc_test(
+    name = "hlo_module_test",
+    srcs = ["hlo_module_test.cc"],
+    deps = [
+        ":hlo",
+        "//xla/hlo/testlib:verified_hlo_module",
+        "//xla/tests:hlo_test_base",
+        "@com_google_absl//absl/hash",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:statusor",
+        "@tsl//tsl/platform:test",
+        "@tsl//tsl/platform:test_main",
+    ],
+)
+
 cc_library(
     name = "backend_config",
     srcs = ["backend_config.cc"],
diff --git a/xla/hlo/ir/hlo_module.h b/xla/hlo/ir/hlo_module.h
index 65004c386e8bf..7bfceca5b1c1f 100644
--- a/xla/hlo/ir/hlo_module.h
+++ b/xla/hlo/ir/hlo_module.h
@@ -287,7 +287,8 @@ class HloModule {
   // with respect to HloInstruction::Identical() method.
   template <typename H>
   friend H AbslHashValue(H h, const HloModule& module) {
-    h = H::combine(std::move(h), module.entry_computation_layout());
+    if (module.config().has_entry_computation_layout())
+      h = H::combine(std::move(h), module.entry_computation_layout());
     // Use MakeComputationSorted() instead of MakeComputationPostOrder()
     // because naming may affect the order of MakeComputationPostOrder() but not
     // MakeComputationSorted().
diff --git a/xla/hlo/ir/hlo_module_test.cc b/xla/hlo/ir/hlo_module_test.cc
new file mode 100644
index 0000000000000..f5bbaf031bbbf
--- /dev/null
+++ b/xla/hlo/ir/hlo_module_test.cc
@@ -0,0 +1,52 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <string_view>
+
+#include <gtest/gtest.h>
+#include "absl/hash/hash.h"
+#include "xla/hlo/testlib/verified_hlo_module.h"
+#include "xla/tests/hlo_test_base.h"
+#include "tsl/platform/statusor.h"
+#include "tsl/platform/test.h"
+
+namespace xla {
+namespace {
+
+using HloModuleTest = HloTestBase;
+
+TEST_F(HloModuleTest, AbslHashValue) {
+  std::unique_ptr<VerifiedHloModule> module1 = CreateNewVerifiedModule();
+  std::unique_ptr<VerifiedHloModule> module2 = CreateNewVerifiedModule();
+  EXPECT_EQ(absl::HashOf(*module1), absl::HashOf(*module2));
+
+  std::string_view hlo = R"(
+      HloModule m1
+        ENTRY main {
+          a = f32[] parameter(0)
+          b = f32[] parameter(1)
+        ROOT res = f32[] multiply(a, b)
+      })";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module3,
+                          ParseAndReturnVerifiedModule(hlo));
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module4,
+                          ParseAndReturnVerifiedModule(hlo));
+  EXPECT_EQ(absl::HashOf(*module3), absl::HashOf(*module4));
+  EXPECT_NE(absl::HashOf(*module1), absl::HashOf(*module4));
+}
+
+}  // namespace
+}  // namespace xla