diff options
| author |  Michael Kuperstein <mkuper@google.com> | 2018-09-25 10:14:53 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-25 10:18:50 -0700 |
| commit | faee2023f9764de44a804c3208be6f68dac04917 (patch) | |
| tree | cde9610e0649cf68f0fc39f9f5104b478b550541 /tensorflow/compiler/xla/service/hlo_computation.cc | |
| parent | 954d6a0ace9b96cdd54659b99e9378a1138a7266 (diff) | |
[XLA] Make HloComputation::instruction_count() constant-time.
* Use a FlatMap for instruction_iterators_, and actually remove elements from it (which is cheap for a FlatMap).
* Use the size of the map (which is O(1)) rather than the size of the list (which is O(n)) for instruction_count().
PiperOrigin-RevId: 214459259
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_computation.cc')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_computation.cc | 9 |
1 files changed, 5 insertions, 4 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc index e9e70b2c57..0e5920af7a 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.cc +++ b/tensorflow/compiler/xla/service/hlo_computation.cc @@ -272,10 +272,11 @@ Status HloComputation::RemoveInstruction(HloInstruction* instruction) { << "instruction " << instruction->name() << " has control successors and cannot be removed"; - TF_RET_CHECK(instruction_iterators_.count(instruction) != 0); - auto inst_it = instruction_iterators_.at(instruction); - (*inst_it)->set_parent(nullptr); - instructions_.erase(inst_it); + auto inst_it = instruction_iterators_.find(instruction); + TF_RET_CHECK(inst_it != instruction_iterators_.end()); + (*inst_it->second)->set_parent(nullptr); + instructions_.erase(inst_it->second); + instruction_iterators_.erase(inst_it); return Status::OK(); } |