From fb9f77525b7b00dde321e78fe8fa5ff5c49e526f Mon Sep 17 00:00:00 2001
From: Yunxing Dai <yunxing@google.com>
Date: Thu, 27 Feb 2020 14:51:11 -0800
Subject: [PATCH] [XLA] Add option to print out hlo pass hash.

PiperOrigin-RevId: 297693781
Change-Id: I5a804056643f7c97bd51da6a7e14fe2e8bc77153
---
 tensorflow/compiler/xla/service/BUILD                |  1 +
 tensorflow/compiler/xla/service/hlo_module.cc        | 10 +++++++---
 tensorflow/compiler/xla/service/hlo_module_group.cc  | 10 ++++++++++
 tensorflow/compiler/xla/service/hlo_module_group.h   |  2 ++
 tensorflow/compiler/xla/service/hlo_pass_pipeline.cc |  1 +
 5 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index e7f945bfa99..ad8cfd2cb21 100755
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -1371,6 +1371,7 @@ cc_library(
     deps = [
         ":hlo",
         ":hlo_proto_cc",
+        "//tensorflow/core:lib",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
     ],
diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc
index 8d5af356275..f31c8ad39c3 100644
--- a/tensorflow/compiler/xla/service/hlo_module.cc
+++ b/tensorflow/compiler/xla/service/hlo_module.cc
@@ -716,9 +716,13 @@ HloComputation* HloModule::GetComputationWithName(absl::string_view name) {
 }
 
 uint64 HloModule::Hash() const {
-  return tensorflow::Hash64Combine(
-      entry_computation_layout().Hash(),
-      entry_computation()->root_instruction()->Hash());
+  uint64 result = entry_computation_layout().Hash();
+  for (auto* computation : MakeComputationPostOrder()) {
+    for (auto* instruction : computation->MakeInstructionPostOrder()) {
+      result = tensorflow::Hash64Combine(result, instruction->Hash());
+    }
+  }
+  return result;
 }
 
 /* static */ std::atomic<int> HloModule::next_unique_module_id_(0);
diff --git a/tensorflow/compiler/xla/service/hlo_module_group.cc b/tensorflow/compiler/xla/service/hlo_module_group.cc
index 69d57c3f146..8c49fd2d4f4 100644
--- a/tensorflow/compiler/xla/service/hlo_module_group.cc
+++ b/tensorflow/compiler/xla/service/hlo_module_group.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/service/hlo_module_group.h"
 
+#include "tensorflow/core/lib/hash/hash.h"
+
 namespace xla {
 
 HloModuleGroup::HloModuleGroup(std::unique_ptr<HloModule> module)
@@ -65,6 +67,14 @@ HloModuleGroupProto HloModuleGroup::ToProto() const {
   return proto;
 }
 
+uint64 HloModuleGroup::Hash() const {
+  uint64 result = 0;
+  for (auto& module : modules_) {
+    result = tensorflow::Hash64Combine(result, module->Hash());
+  }
+  return result;
+}
+
 /* static */ StatusOr<HloModuleGroup> HloModuleGroup::CreateFromProto(
     const HloModuleGroupProto& proto,
     absl::Span<const HloModuleConfig> module_configs) {
diff --git a/tensorflow/compiler/xla/service/hlo_module_group.h b/tensorflow/compiler/xla/service/hlo_module_group.h
index 217f65b4a75..2cfb677eaa2 100644
--- a/tensorflow/compiler/xla/service/hlo_module_group.h
+++ b/tensorflow/compiler/xla/service/hlo_module_group.h
@@ -71,6 +71,8 @@ class HloModuleGroup {
     }
   }
 
+  uint64 Hash() const;
+
   // Serialize the module group to/from a proto.
   HloModuleGroupProto ToProto() const;
   static StatusOr<HloModuleGroup> CreateFromProto(
diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
index 62c40546f0a..b07ab10827a 100644
--- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
+++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
@@ -60,6 +60,7 @@ StatusOr<bool> HloPassPipeline::RunPassesInternal(
   for (HloPassInterface* pass : passes) {
     absl::string_view pass_name = pass->name();
     VLOG(1) << "  HLO pass " << pass_name;
+    VLOG(2) << "  Module hash " << hlo->Hash();
     MaybeDumpHlo(*hlo,
                  /*after_pass_name=*/last_pass_name,
                  /*before_pass_name=*/pass_name);