HloPrintOptions::Fingerprint() canonicalizes computations by sorting them by names.
PiperOrigin-RevId: 274484074
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h
index 6a27092..e6d9aad 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.h
+++ b/tensorflow/compiler/xla/service/hlo_instruction.h
@@ -91,7 +91,8 @@
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 @@
// 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 @@
.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 @@
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 @@
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 @@
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
diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc
index d068b77..51c83a3 100644
--- a/tensorflow/compiler/xla/service/hlo_module.cc
+++ b/tensorflow/compiler/xla/service/hlo_module.cc
@@ -222,7 +222,10 @@
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 @@
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()) {
diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h
index 745ad56..40f3c97 100644
--- a/tensorflow/compiler/xla/service/hlo_module.h
+++ b/tensorflow/compiler/xla/service/hlo_module.h
@@ -189,6 +189,9 @@
// 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
