• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_BUFFER_LIVENESS_H_
17 #define TENSORFLOW_COMPILER_XLA_SERVICE_BUFFER_LIVENESS_H_
18 
19 #include <memory>
20 #include <string>
21 #include <utility>
22 
23 #include "absl/container/flat_hash_set.h"
24 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
25 #include "tensorflow/compiler/xla/service/hlo_module.h"
26 #include "tensorflow/compiler/xla/service/hlo_ordering.h"
27 #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h"
28 #include "tensorflow/compiler/xla/statusor.h"
29 #include "tensorflow/compiler/xla/types.h"
30 #include "tensorflow/core/lib/core/status.h"
31 
32 namespace xla {
33 
34 // Class which computes liveness of the output buffers of HLOs and their
35 // interference.
36 class BufferLiveness {
37  public:
38   using Colorer = std::function<Status(const BufferLiveness& buffer_liveness)>;
39 
40   // Constructs a buffer liveness object for the given module assuming the given
41   // HLO instruction ordering.
42   static StatusOr<std::unique_ptr<BufferLiveness>> Run(
43       const HloModule* module, std::unique_ptr<HloOrdering> hlo_ordering);
44 
45   // Returns true if the live range of the buffer containing the output of 'a'
46   // may overlap with the live range of the buffer of 'b'. If instruction 'a'
47   // interferes with instruction 'b' then they cannot share the same buffer.
48   bool MayInterfere(const LogicalBuffer& a, const LogicalBuffer& b) const;
49 
50   // Returns true if the buffer for the given instruction may be live out of the
51   // module. That is, the instruction's buffer may be included in the output of
52   // the entry computation.
53   bool MaybeLiveOut(const LogicalBuffer& buffer) const;
54 
55   // Returns the complete set of buffers that may be live out of the module.
maybe_live_out_buffers()56   const PointsToSet::BufferSet& maybe_live_out_buffers() const {
57     return maybe_live_out_buffers_;
58   }
59 
60   // Returns the underlying points-to analysis used for this liveness analysis.
points_to_analysis()61   const TuplePointsToAnalysis& points_to_analysis() const {
62     return *points_to_analysis_;
63   }
64 
65   // Returns the underlying hlo ordering used for this liveness analysis.
hlo_ordering()66   const HloOrdering& hlo_ordering() const { return *hlo_ordering_; }
67 
module()68   const HloModule& module() const { return *module_; }
69 
70   string ToString() const;
71 
DefaultColorer()72   static Colorer DefaultColorer() {
73     return [](const BufferLiveness& buffer_liveness) {
74       for (LogicalBuffer::Id id = 0;
75            id < buffer_liveness.points_to_analysis().num_logical_buffers();
76            id++) {
77         auto& buffer = buffer_liveness.points_to_analysis().logical_buffer(id);
78         buffer.set_color(LogicalBuffer::Color(0));
79       }
80       return Status::OK();
81     };
82   }
83 
84  private:
BufferLiveness(const HloModule * module,std::unique_ptr<HloOrdering> hlo_ordering)85   explicit BufferLiveness(const HloModule* module,
86                           std::unique_ptr<HloOrdering> hlo_ordering)
87       : module_(module), hlo_ordering_(std::move(hlo_ordering)) {}
88 
89   // Perform buffer liveness analysis. This method must be called prior to
90   // MayInterfere or MaybeLiveOut.
91   Status Analyze();
92 
93   // Returns true if the live range of the buffer of 'a' is strictly before the
94   // live range of the buffer of 'b' (they do not overlap).
95   bool live_range_strictly_before(const LogicalBuffer& a,
96                                   const LogicalBuffer& b) const;
97 
98   const HloModule* module_;
99   std::unique_ptr<HloOrdering> hlo_ordering_;
100 
101   // Set of LogicalBuffers which are aliased in the output of other
102   // instructions. For example, a LogicalBuffer which is inserted into a tuple
103   // is considered to be aliased and will be in this set.
104   absl::flat_hash_set<const LogicalBuffer*> aliased_buffers_;
105 
106   // LogicalBuffers that may be live out of the entry computation.
107   PointsToSet::BufferSet maybe_live_out_buffers_;
108 
109   std::unique_ptr<TuplePointsToAnalysis> points_to_analysis_;
110 };
111 
112 }  // namespace xla
113 
114 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_BUFFER_LIVENESS_H_
115