HloPrintOptions::Fingerprint() canonicalizes computations by sorting them by names.

PiperOrigin-RevId: 274484074

tensorflow/compiler/xla/service/hlo_instruction.h

@@ -91,7 +91,8 @@ class HloPrintOptions {
         canonicalize_instruction_names_(false),
         indent_amount_(0),
         is_in_nested_computation_(false),
-        print_ids_(true) {}
+        print_ids_(true),
+        canonicalize_computations_(false) {}
 
   static HloPrintOptions ShortParsable() {
     return HloPrintOptions()
@@ -124,7 +125,7 @@ class HloPrintOptions {
   // Options to produce a fingerprint of an HLO.
   static HloPrintOptions Fingerprint() {
     return HloPrintOptions()
-        .set_print_subcomputation_mode(PrintSubcomputationMode::kNameOnly)
+        .set_print_subcomputation_mode(PrintSubcomputationMode::kFullBodies)
         .set_print_metadata(false)
         .set_print_backend_config(false)
         .set_compact_operands(true)
@@ -134,7 +135,8 @@ class HloPrintOptions {
         .set_print_percent(false)
         .set_print_control_dependencies(false)
         .set_canonicalize_instruction_names(true)
-        .set_print_ids(false);
+        .set_print_ids(false)
+        .set_canonicalize_computations(true);
   }
 
   // If true, large constants will be printed out.
@@ -218,6 +220,12 @@ class HloPrintOptions {
     return *this;
   }
 
+  // If true, canonicalizes computations, sorting by computations' names.
+  HloPrintOptions& set_canonicalize_computations(bool value) {
+    canonicalize_computations_ = value;
+    return *this;
+  }
+
   // The indent of the hlo text block.
   HloPrintOptions& set_indent_amount(int value) {
     indent_amount_ = value;
@@ -250,6 +258,7 @@ class HloPrintOptions {
   bool canonicalize_instruction_names() const {
     return canonicalize_instruction_names_;
   }
+  bool canonicalize_computations() const { return canonicalize_computations_; }
   int indent_amount() const { return indent_amount_; }
   int is_in_nested_computation() const { return is_in_nested_computation_; }
 
@@ -269,6 +278,7 @@ class HloPrintOptions {
   int indent_amount_;
   bool is_in_nested_computation_;
   bool print_ids_;
+  bool canonicalize_computations_;
 };
 
 // For canonical string output, we need to have a canonical way to rename

tensorflow/compiler/xla/service/hlo_module.cc

@@ -222,7 +222,10 @@ string HloModule::ToString(const HloPrintOptions& options) const {
     s << ", is_scheduled=true";
   }
   s << "\n\n";
-  for (const HloComputation* computation : MakeComputationPostOrder()) {
+  const auto& computations = options.canonicalize_computations()
+                                 ? MakeComputationPostOrderAndSortedByNames()
+                                 : MakeComputationPostOrder();
+  for (const HloComputation* computation : computations) {
     if (computation == entry_computation()) {
       s << "ENTRY ";
     }
@@ -592,6 +595,16 @@ std::vector<HloComputation*> HloModule::MakeComputationPostOrder() const {
   return post_order;
 }
 
+std::vector<HloComputation*>
+HloModule::MakeComputationPostOrderAndSortedByNames() const {
+  auto result = MakeComputationPostOrder();
+  std::sort(result.begin(), result.end(),
+            [](HloComputation* a, HloComputation* b) {
+              return a->name() < b->name();
+            });
+  return result;
+}
+
 std::vector<HloComputation*> HloModule::MakeNonfusionComputations() const {
   std::vector<HloComputation*> result;
   for (auto* c : computations()) {

tensorflow/compiler/xla/service/hlo_module.h

@@ -189,6 +189,9 @@ class HloModule {
   // computation B, then A will appear after B in the sort.
   std::vector<HloComputation*> MakeComputationPostOrder() const;
 
+  // Same as MakeComputationPostOrder() but sorting the computations by names.
+  std::vector<HloComputation*> MakeComputationPostOrderAndSortedByNames() const;
+
   // Gets the computations in this module which aren't for fusion nodes.
   //
   // Postcondition: All computations in the returned list have