From 98ea840dabc0c4e9417ebe9a0fd10c9d471cda51 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 4 Oct 2018 02:41:25 -0700 Subject: Improve the performance of the ListMemoryScheduler This CL replaces a std::unordered_map with an absl::flat_hash_map and removes an unnecessary map lookup. This two change can improve the performance of the scheduler on large graphs by up to 2x. PiperOrigin-RevId: 215707921 --- tensorflow/compiler/xla/service/hlo_memory_scheduler.cc | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) (limited to 'tensorflow/compiler') diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc b/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc index 55314d0ae9..bf30764488 100644 --- a/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc +++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc @@ -263,9 +263,8 @@ class ListScheduler { }; for (auto* instruction : computation_.instructions()) { - // Instruction with no operands or control predecessors will - // not be in the map. - if (unscheduled_pred_count.count(instruction) == 0) { + if (instruction->operands().empty() && + instruction->control_predecessors().empty()) { add_to_ready_queue(instruction); } } @@ -356,9 +355,8 @@ class ListScheduler { buffer_uses_; // A map containing the count of unscheduled HLOs which using a particular - // LogicalBuffer. We rely on iterator stability in this map, and that the map - // entries are std::pair's. - std::unordered_map unscheduled_use_count_; + // LogicalBuffer. + absl::flat_hash_map unscheduled_use_count_; // Set of instructions which have been scheduled. absl::flat_hash_set scheduled_instructions_; -- cgit v1.2.3