diff options
| author |  Justin Lebar <jlebar@google.com> | 2017-09-28 23:07:12 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2017-09-28 23:14:07 -0700 |
| commit | 9b1b5d85b9ce3c812dc772da1f3f5d09581e5b49 (patch) | |
| tree | 798207232f7ae8240d5ee2e7f934669ed1318bef /tensorflow/compiler/xla/service/hlo_computation.h | |
| parent | 75e07e01a41434fdf40eea6291fe7bc47ad74312 (diff) | |
[XLA] Make HloComputation::instructions() return a view of HloInstruction*s.
Currently it returns a view of unique_ptr<HloInstruction>s. But the
fact that these are unique_ptrs is an implementation detail, and it's
ugly to leak it everywhere.
PiperOrigin-RevId: 170445375
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_computation.h')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_computation.h | 21 |
1 files changed, 19 insertions, 2 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h index ab902312ad..b929b41bad 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.h +++ b/tensorflow/compiler/xla/service/hlo_computation.h @@ -24,6 +24,7 @@ limitations under the License. #include <utility> #include <vector> +#include "tensorflow/compiler/xla/iterator_util.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" @@ -142,8 +143,24 @@ class HloComputation { // Returns a serialized representation of this computation. HloComputationProto ToProto() const; - const std::list<std::unique_ptr<HloInstruction>>& instructions() const { - return instructions_; + // Gets the instructions in this computation. + // + // The returned type is a range of HloInstruction*s, so you can iterate over + // it using a range-based for loop in the natural way: + // + // for (HloInstruction* instr : computation->instructions()) { ... } + // + tensorflow::gtl::iterator_range<UnwrappingIterator< + std::list<std::unique_ptr<HloInstruction>>::const_iterator>> + instructions() const { + return {MakeUnwrappingIterator(instructions_.begin()), + MakeUnwrappingIterator(instructions_.end())}; + } + tensorflow::gtl::iterator_range< + UnwrappingIterator<std::list<std::unique_ptr<HloInstruction>>::iterator>> + instructions() { + return {MakeUnwrappingIterator(instructions_.begin()), + MakeUnwrappingIterator(instructions_.end())}; } // Compute and return a post-order of the instructions in the computation. In |