[XLA] Do not reserve large hash maps when there are many small comput…

…ations.

PiperOrigin-RevId: 380641027
Change-Id: I51ae2adcaefbc94a139eac4a5cafb5df724cb280

tensorflow/compiler/xla/service/hlo_instruction.cc

@@ -3470,12 +3470,7 @@ template <typename Visitor>
 static Status PostOrderDFS(HloInstruction* root, Visitor* visitor,
                            const InternalCompareFunction* operand_order,
                            bool ignore_control_predecessors) {
-  // Calculating the instruction count within a module can be expensive on large
-  // models so only do it if the visit state is empty. This will help when the
-  // same visitor is reused across many computations of a single module.
-  if (visitor->VisitStateCapacity() == 0) {
-    visitor->ReserveVisitStates(root->GetModule()->instruction_count());
-  }
+  visitor->ReserveVisitStates(root->parent()->instruction_count());
 
   // dfs_stack holds pairs of <HloInstruction*->unique_id(), HloInstruction*>.
   //