diff options
author | A. Unique TensorFlower <gardener@tensorflow.org> | 2018-10-04 02:41:25 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-10-04 02:46:00 -0700 |
commit | 98ea840dabc0c4e9417ebe9a0fd10c9d471cda51 (patch) | |
tree | fa3d57dee244ebe86fbb2b095c19d9ccbef26eea /tensorflow/compiler/xla/service | |
parent | 67e0ccb3e5c1a48d62bcc45201fd70d2420dc4eb (diff) |
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
Diffstat (limited to 'tensorflow/compiler/xla/service')
-rw-r--r-- | tensorflow/compiler/xla/service/hlo_memory_scheduler.cc | 10 |
1 files changed, 4 insertions, 6 deletions
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<const LogicalBuffer*, int64> unscheduled_use_count_; + // LogicalBuffer. + absl::flat_hash_map<const LogicalBuffer*, int64> unscheduled_use_count_; // Set of instructions which have been scheduled. absl::flat_hash_set<const HloInstruction*> scheduled_instructions_; |