Refactor BiDirectionalMap to fully hide internals and expose a symmetric API.

PiperOrigin-RevId: 780793848

Author: dubstack

third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_eval.cc

@@ -67,29 +67,28 @@ int64_t CountSplitAllegianceParental(const HloGumgraph& left,
                                      const HloGumgraph& right,
                                      const HloGumgraphMappings& mappings) {
   int64_t count = 0;
-  for (const auto& it : mappings.left_to_right_instruction_map.left) {
-    if (it.first->children.size() != it.second.node->children.size()) {
+  for (const auto& [left, right] : mappings.left_to_right_instruction_map) {
+    if (left->children.size() != right->children.size()) {
       continue;
     }
     bool children_opcode_mismatch = false;
-    for (int i = 0; i < it.first->children.size(); ++i) {
-      if (it.first->children[i]->instruction->opcode() !=
-          it.second.node->children[i]->instruction->opcode()) {
+    for (int i = 0; i < left->children.size(); ++i) {
+      if (left->children[i]->instruction->opcode() !=
+          right->children[i]->instruction->opcode()) {
         children_opcode_mismatch = true;
         break;
       }
     }
     if (children_opcode_mismatch) {
       continue;
     }
-    for (int i = 0; i < it.first->children.size(); ++i) {
-      if (auto cit = mappings.left_to_right_instruction_map.left.find(
-              it.first->children[i]);
-          cit == mappings.left_to_right_instruction_map.left.end() ||
-          cit->second.node != it.second.node->children[i]) {
+    for (int i = 0; i < left->children.size(); ++i) {
+      const auto* left_child = left->children[i];
+      const auto* expected_right_child = right->children[i];
+      if (auto mapped_right_child =
+              mappings.left_to_right_instruction_map.GetRight(left_child);
+          mapped_right_child != expected_right_child) {
         count++;
-        // LOG(INFO) << it.first->instruction->name() << " has split child: "
-        //           << it.first->children[i]->instruction->name();
       }
     }
   }

third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_result.cc

@@ -71,11 +71,9 @@ std::unique_ptr<const DiffResult> ConstructDiffResult(
 
     // The node is matched
     const HloInstructionNode* right_node =
-        mappings.left_to_right_instruction_map.left.find(left_node)
-            ->second.node;
-    const std::optional<HloInstructionNodeMappingProps>& mapping_props =
-        mappings.left_to_right_instruction_map.left.find(left_node)
-            ->second.props;
+        *mappings.left_to_right_instruction_map.GetRight(left_node);
+    const std::optional<HloInstructionNodeMappingProps> mapping_props =
+        mappings.left_to_right_instruction_map.GetPropsByLeft(left_node);
 
     // Fill in matcher debug info.
     diff_result->map_by[std::make_pair(left_node->instruction,

third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_summary.cc

@@ -20,6 +20,7 @@
 #include <algorithm>
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <ostream>
 #include <string>
 #include <utility>
@@ -60,25 +61,21 @@ InstructionBimap ConstructInstructionBimap(const DiffResult& diff_result) {
 
 // Returns the mapped instruction node of the given instruction in the given
 // direction. Returns nullptr if the instruction is not mapped.
-const HloInstruction* FindMappedInstruction(const InstructionBimap& mapping,
-                                            const HloInstruction* instruction,
-                                            DiffSide side) {
+std::optional<const HloInstruction*> FindMappedInstruction(
+    const InstructionBimap& mapping, const HloInstruction* instruction,
+    DiffSide side) {
   switch (side) {
     case DiffSide::kLeft: {
-      if (auto it = mapping.left.find(instruction); it != mapping.left.end()) {
-        return it->second.node;
-      }
+      return mapping.GetRight(instruction);
       break;
     }
     case DiffSide::kRight: {
-      if (auto it = mapping.right.find(instruction);
-          it != mapping.right.end()) {
-        return it->second;
-      }
+      return mapping.GetLeft(instruction);
       break;
     }
   }
-  return nullptr;
+
+  return std::nullopt;
 }
 
 // Result of finding the main matched computation.
@@ -99,11 +96,11 @@ MainMatchedComputationResult FindMainMatchedComputation(
   int mapped_instruction_count = 0;
   const HloComputation* main_matched_computation = nullptr;
   for (const HloInstruction* instruction : computation->instructions()) {
-    if (const HloInstruction* const mapped_instruction =
+    if (std::optional<const HloInstruction*> mapped_instruction =
             FindMappedInstruction(mapping, instruction, side);
-        mapped_instruction != nullptr) {
+        mapped_instruction.has_value()) {
       ++mapped_instruction_count;
-      const HloComputation* right_computation = mapped_instruction->parent();
+      const HloComputation* right_computation = (*mapped_instruction)->parent();
       const int count = ++matched_instruction_count[right_computation];
       if (count > max_count) {
         max_count = count;

third_party/xla/xla/hlo/tools/hlo_diff/hlo_gumgraph_diff.cc

@@ -56,10 +56,10 @@ absl::StatusOr<std::unique_ptr<const HloGumgraphMappings>> FindMappings(
 
   TF_RETURN_IF_ERROR(left.GetCallGraph().VisitNodes(
       [&](const CallGraphNode& node) {
-        if (auto it = mappings->left_to_right_computation_map.left.find(&node);
-            it != mappings->left_to_right_computation_map.left.end()) {
-          MatchComputationGraphs(left, right, node, *it->second.node,
-                                 *mappings);
+        if (auto right_node =
+                mappings->left_to_right_computation_map.GetRight(&node);
+            right_node.has_value()) {
+          MatchComputationGraphs(left, right, node, **right_node, *mappings);
         }
         return absl::OkStatus();
       },

third_party/xla/xla/hlo/tools/hlo_diff/hlo_gumgraph_mappings.h

@@ -94,13 +94,13 @@ struct HloGumgraphMappings {
   // Returns true if the left node is mapped to a right node.
   inline bool InstructionMapContainsLeft(
       const HloInstructionNode* left_node) const {
-    return left_to_right_instruction_map.left.contains(left_node);
+    return left_to_right_instruction_map.ContainsLeft(left_node);
   }
 
   // Returns true if the right node is mapped to a left node.
   inline bool InstructionMapContainsRight(
       const HloInstructionNode* right_node) const {
-    return left_to_right_instruction_map.right.contains(right_node);
+    return left_to_right_instruction_map.ContainsRight(right_node);
   }
 };
 

third_party/xla/xla/hlo/tools/hlo_diff/matchers/bottom_up_matcher.cc

@@ -86,9 +86,9 @@ double DiceSimLimitedSubgraph(const HloInstructionNode* absl_nonnull left,
 
   int common = 0;
   for (const HloInstructionNode* left_node : left_nodes) {
-    if (auto it = mappings.left_to_right_instruction_map.left.find(left_node);
-        it != mappings.left_to_right_instruction_map.left.end() &&
-        right_nodes_set.contains(it->second.node)) {
+    if (auto right_node =
+            mappings.left_to_right_instruction_map.GetRight(left_node);
+        right_node.has_value() && right_nodes_set.contains(*right_node)) {
       ++common;
     }
   }
@@ -127,10 +127,9 @@ double AllOperandHloValuesMatchedScore(
         left.GetNode(left_hlo_values[i]->defining_instruction());
     HloInstructionNode* right_hlo_value_node =
         right.GetNode(right_hlo_values[i]->defining_instruction());
-    if (auto it = mappings.left_to_right_instruction_map.left.find(
+    if (auto right_node = mappings.left_to_right_instruction_map.GetRight(
             left_hlo_value_node);
-        it == mappings.left_to_right_instruction_map.left.end() ||
-        it->second.node != right_hlo_value_node) {
+        right_node != right_hlo_value_node) {
       mappings_matched = false;
     }
     if (left_hlo_value_node->props.fingerprint !=
@@ -177,9 +176,10 @@ void GreedyLimitedCandidatesBottomUpMatcher::Match(
     HloGumgraphBfs(
         *left_node,
         [&](const HloInstructionNode& node, int distance) {
-          if (auto it = mappings.left_to_right_instruction_map.left.find(&node);
-              it != mappings.left_to_right_instruction_map.left.end()) {
-            right_seeds.push_back(it->second.node);
+          if (auto right_node =
+                  mappings.left_to_right_instruction_map.GetRight(&node);
+              right_node.has_value()) {
+            right_seeds.push_back(*right_node);
           }
           // Don't pursue subgraphs with too many childrens. Allows us to visit
           // deeper subgraphs without getting stuck on a single node with a

third_party/xla/xla/hlo/tools/hlo_diff/matchers/exact_subgraph_matcher.cc

@@ -95,8 +95,11 @@ void MapSubgraph(const HloInstructionNode* absl_nonnull left,
       exact_mapped_subgraph_roots.insert(left_subgraph[i]);
       exact_mapped_subgraph_roots.insert(right_subgraph[i]);
       if (i != 0) {
-        mappings.left_to_right_instruction_map.left.find(left_subgraph[i])
-            ->second.props->unchanged = true;
+        auto props = mappings.left_to_right_instruction_map.GetPropsByLeft(
+            left_subgraph[i]);
+        props->unchanged = true;
+        mappings.left_to_right_instruction_map.SetPropsByLeft(left_subgraph[i],
+                                                              *props);
       }
     }
   }

third_party/xla/xla/hlo/tools/hlo_diff/matchers/hlo_call_graph_matcher.cc

@@ -138,16 +138,20 @@ void MapCalledComputations(const HloInstruction* left_instruction,
 void ProcessCallGraphNode(const CallGraphNode& left_computation,
                           const HloGumgraph& left, const HloGumgraph& right,
                           HloGumgraphMappings& mappings) {
+  if (left_computation.callees().empty()) {
+    return;
+  }
+
   // Only match called computations if current computation is already matched.
-  auto it = mappings.left_to_right_computation_map.left.find(&left_computation);
-  if (it == mappings.left_to_right_computation_map.left.end() ||
-      left_computation.callees().empty()) {
+  auto right_computation =
+      mappings.left_to_right_computation_map.GetRight(&left_computation);
+  if (!right_computation.has_value()) {
     return;
   }
 
-  const CallGraphNode* right_computation = it->second.node;
   HloComputation::CachingPostOrder left_cpo(left_computation.computation());
-  HloComputation::CachingPostOrder right_cpo(right_computation->computation());
+  HloComputation::CachingPostOrder right_cpo(
+      (*right_computation)->computation());
 
   // Phase 1: Match called computations to computations with matching
   // computation fingerprints, i.e. exact matches.
@@ -158,7 +162,7 @@ void ProcessCallGraphNode(const CallGraphNode& left_computation,
     left_callees_by_fingerprint[left_props.fingerprint].insert(
         left_props.call_graph_node);
   }
-  for (const HloComputation* callee : right_computation->callees()) {
+  for (const HloComputation* callee : (*right_computation)->callees()) {
     CallGraphNodeProps right_props = right.AllComputationProps().at(callee);
     right_callees_by_fingerprint[right_props.fingerprint].insert(
         right_props.call_graph_node);
@@ -185,7 +189,7 @@ void ProcessCallGraphNode(const CallGraphNode& left_computation,
   for (const HloInstruction* instruction : left_cpo.PostOrder()) {
     bool all_called_computations_matched = true;
     for (const HloComputation* callee : instruction->called_computations()) {
-      if (!mappings.left_to_right_computation_map.left.contains(
+      if (!mappings.left_to_right_computation_map.ContainsLeft(
               &left.GetCallGraph().GetNode(callee))) {
         all_called_computations_matched = false;
         break;
@@ -202,7 +206,7 @@ void ProcessCallGraphNode(const CallGraphNode& left_computation,
   for (const HloInstruction* instruction : right_cpo.PostOrder()) {
     bool all_called_computations_matched = true;
     for (const HloComputation* callee : instruction->called_computations()) {
-      if (!mappings.left_to_right_computation_map.right.contains(
+      if (!mappings.left_to_right_computation_map.ContainsRight(
               &right.GetCallGraph().GetNode(callee))) {
         all_called_computations_matched = false;
         break;
@@ -256,7 +260,7 @@ void ProcessCallGraphNode(const CallGraphNode& left_computation,
   absl::flat_hash_map<std::string, absl::flat_hash_set<const CallGraphNode*>>
       unmatched_left_callees, unmatched_right_callees;
   for (const HloComputation* callee : left_computation.callees()) {
-    if (!mappings.left_to_right_computation_map.left.contains(
+    if (!mappings.left_to_right_computation_map.ContainsLeft(
             &left.GetCallGraph().GetNode(callee))) {
       const CallGraphNode& callee_node = left.GetCallGraph().GetNode(callee);
       std::string opcode_and_name;
@@ -276,8 +280,8 @@ void ProcessCallGraphNode(const CallGraphNode& left_computation,
       unmatched_left_callees[opcode_name_shape].insert(&callee_node);
     }
   }
-  for (const HloComputation* callee : right_computation->callees()) {
-    if (!mappings.left_to_right_computation_map.right.contains(
+  for (const HloComputation* callee : (*right_computation)->callees()) {
+    if (!mappings.left_to_right_computation_map.ContainsRight(
             &right.GetCallGraph().GetNode(callee))) {
       const CallGraphNode& callee_node = right.GetCallGraph().GetNode(callee);
       std::string opcode_and_name;
@@ -364,10 +368,9 @@ void MatchCallGraphs(const HloGumgraph& left, const HloGumgraph& right,
       });
 
   int signature_match_count = 0, exact_match_count = 0;
-  for (auto it = mappings.left_to_right_computation_map.left.begin();
-       it != mappings.left_to_right_computation_map.left.end(); ++it) {
-    if (it->second.props->computation_match_type ==
-        ComputationMatchType::kSignature) {
+  for (const auto& [left, right] : mappings.left_to_right_computation_map) {
+    auto props = mappings.left_to_right_computation_map.GetPropsByLeft(left);
+    if (props->computation_match_type == ComputationMatchType::kSignature) {
       ++signature_match_count;
     } else {
       ++exact_match_count;

third_party/xla/xla/hlo/tools/hlo_diff/matchers/hlo_computation_graph_matcher.cc

@@ -61,17 +61,16 @@ void MatchComputationGraphs(const HloGumgraph& left, const HloGumgraph& right,
                             const CallGraphNode& left_computation,
                             const CallGraphNode& right_computation,
                             HloGumgraphMappings& mappings) {
-  auto it = mappings.left_to_right_computation_map.left.find(&left_computation);
-  if (it == mappings.left_to_right_computation_map.left.end()) {
+  if (!mappings.left_to_right_computation_map.ContainsLeft(&left_computation)) {
     return;
   }
 
   MatchCallSites(left, right, left_computation, right_computation, mappings);
 
   // If the two computations are exact matches, we can match all
   // instructions in the two computations.
-  if (it->second.props->computation_match_type ==
-      ComputationMatchType::kExact) {
+  if (mappings.left_to_right_computation_map.GetPropsByLeft(&left_computation)
+          ->computation_match_type == ComputationMatchType::kExact) {
     auto left_instructions =
         left_computation.computation()->MakeInstructionPostOrder();
     auto right_instructions =

third_party/xla/xla/hlo/tools/hlo_diff/matchers/top_down_matcher.cc

@@ -67,20 +67,21 @@ void GreedyTopDownMatcher::Match(HloGumgraphMappings& mappings) const {
   HloGumgraphDfs(
       left_.GetRoot(),
       [&](const HloInstructionNode& left_node) {
-        auto it = mappings.left_to_right_instruction_map.left.find(&left_node);
-        if (it == mappings.left_to_right_instruction_map.left.end()) {
+        auto right_node =
+            mappings.left_to_right_instruction_map.GetRight(&left_node);
+        if (!right_node) {
           return;
         }
 
         if (require_same_children_ &&
-            ShouldSkipMatching(left_node, *it->second.node)) {
+            ShouldSkipMatching(left_node, **right_node)) {
           return;
         }
 
         std::vector<HloInstructionNode*> left_children =
             GetUnmatchedChildren(left_node.children, mappings);
         std::vector<HloInstructionNode*> right_children =
-            GetUnmatchedChildren(it->second.node->children, mappings);
+            GetUnmatchedChildren((*right_node)->children, mappings);
 
         MatchInstructions(left_, right_, left_children, right_children,
                           mappings, type_, MapByPositionMode::kAlways);