Correctly pre-reserve visit state in HloInstruction::PostOrderDFS

Previously we pre-reserverd the visit state based on the number of
instructions but then started to index it with the instruction unique ID
what can be larger then the instruction count. This resulted in some
very expensive re-allocations what can be eliminated by reserving the
correctly sized buffer.

PiperOrigin-RevId: 216369849
This commit is contained in:
A. Unique TensorFlower 2018-10-09 09:30:32 -07:00 committed by TensorFlower Gardener
parent 92d533d19c
commit 87d8055c74

View File

@ -2474,7 +2474,7 @@ template <typename Visitor>
static Status PostOrderDFS(HloInstruction* root, Visitor* visitor,
const InternalCompareFunction* operand_order,
bool ignore_control_predecessors) {
visitor->ReserveVisitStates(root->GetModule()->instruction_count());
visitor->ReserveVisitStates(root->GetModule()->NumUniqueInstructionIds());
// dfs_stack holds pairs of <HloInstruction*->unique_id(), HloInstruction*>.
//