• 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_HLO_ALIAS_ANALYSIS_H_
17 #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_ALIAS_ANALYSIS_H_
18 
19 #include <memory>
20 #include <string>
21 #include <vector>
22 
23 #include "absl/container/flat_hash_map.h"
24 #include "absl/types/span.h"
25 #include "tensorflow/compiler/xla/service/hlo_buffer.h"
26 #include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h"
27 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
28 #include "tensorflow/compiler/xla/service/hlo_module.h"
29 #include "tensorflow/compiler/xla/service/hlo_ordering.h"
30 #include "tensorflow/compiler/xla/status.h"
31 #include "tensorflow/compiler/xla/statusor.h"
32 #include "tensorflow/compiler/xla/types.h"
33 #include "tensorflow/compiler/xla/xla_data.pb.h"
34 #include "tensorflow/core/platform/macros.h"
35 
36 namespace xla {
37 
38 // Analysis which allocates HloBuffers to HloValues.
39 class HloAliasAnalysis {
40  public:
41   // The callgraph of the given HloModule must be flattened
42   // (xla::FlattenCallGraph) prior to running the analysis.
43   static StatusOr<std::unique_ptr<HloAliasAnalysis>> Run(
44       const HloModule* module,
45       const HloDataflowAnalysis::CanShareBuffer& can_share_buffer = nullptr);
46 
47   string ToString() const;
48 
49   // Return the buffer containing the given value.
GetBufferContainingValue(const HloValue & value)50   const HloBuffer& GetBufferContainingValue(const HloValue& value) const {
51     return *value_to_buffer_.at(&value);
52   }
GetBufferContainingValue(const HloValue & value)53   HloBuffer& GetBufferContainingValue(const HloValue& value) {
54     return *value_to_buffer_.at(&value);
55   }
56 
57   // Return the HloBuffer with the given ID.
GetBuffer(HloBuffer::Id buffer_id)58   const HloBuffer& GetBuffer(HloBuffer::Id buffer_id) const {
59     return buffers_.at(buffer_id);
60   }
GetBuffer(HloBuffer::Id buffer_id)61   HloBuffer& GetBuffer(HloBuffer::Id buffer_id) {
62     return buffers_.at(buffer_id);
63   }
64 
65   // Returns the unique buffer at the given position. CHECK fails if the buffer
66   // set at that position does not contain exactly one buffer.
67   const HloBuffer& GetUniqueBufferAt(const HloInstruction* instruction,
68                                      const ShapeIndex& index = {}) const;
69   HloBuffer& GetUniqueBufferAt(const HloInstruction* instruction,
70                                const ShapeIndex& index = {});
71 
72   // Compute the set of buffers at the given instruction and index and return as
73   // a vector. This set is exactly the union of the buffers containing the
74   // HloValues at this position.
75   std::vector<const HloBuffer*> ComputeBuffersAt(
76       const HloInstruction* instruction, const ShapeIndex& index = {}) const;
77 
78   // Return a vector of all HloBuffers stabily sorted by HloBuffer::Id. This
79   // vector is lazily computed. Mutating operations on HloAliasAnalysis may
80   // invalidate the underlying vector requiring recomputation.
buffers()81   const std::vector<HloBuffer>& buffers() const { return buffers_; }
82 
83   // Returns the underlying dataflow analysis used by this alias analysis.
dataflow_analysis()84   HloDataflowAnalysis& dataflow_analysis() const { return *dataflow_analysis_; }
85 
86   // Returns true if any index in the output of the given instruction has more
87   // than one buffer. That is, ComputeBuffersAt returns a vector with more than
88   // one element.
89   bool InstructionBuffersAreAmbiguous(const HloInstruction* instruction) const;
90 
91   // Returns true if no HloBuffer appears in more than one shape index in the
92   // output of the given instruction.
93   bool InstructionBuffersAreDistinct(const HloInstruction* instruction) const;
94 
95   // Merge buffer `from` into buffer `to`. Caller has to make sure no
96   // interference will be introduced after merging. This rebuilds internal data
97   // structure, and invalidates references to all existing buffers.
98   void MergeBuffers(const HloBuffer& to, const HloBuffer& from);
99 
100   // Returns true if any HLO values in the module have interfering live ranges
101   // assuming the given ordering.
102   bool HasLiveRangeInterference(const HloOrdering& ordering) const;
103 
104   // Returns true if a buffer lives out of the module.
BufferLivesOut(const HloBuffer & buffer)105   bool BufferLivesOut(const HloBuffer& buffer) const {
106     return live_out_buffers_.count(&buffer);
107   }
108 
109   // Returns true if a hlo value lives out of the module.
ValueLivesOut(const HloValue & value)110   bool ValueLivesOut(const HloValue& value) const {
111     return live_out_buffers_.count(&GetBufferContainingValue(value));
112   }
113 
LiveOutBuffers()114   std::vector<const HloBuffer*> LiveOutBuffers() const {
115     std::vector<const HloBuffer*> results(live_out_buffers_.begin(),
116                                           live_out_buffers_.end());
117     absl::c_sort(results, [](const HloBuffer* a, const HloBuffer* b) {
118       return a->id() < b->id();
119     });
120     return results;
121   }
122 
123  protected:
124   explicit HloAliasAnalysis(const HloModule* module);
125 
126   // Verify various invariants of the alias analysis.
127   Status Verify() const;
128 
129   const HloModule* module_;
130 
131   // A set of buffers that live out the module.
132   absl::flat_hash_set<const HloBuffer*> live_out_buffers_;
133 
134   // The underlying dataflow analysis used by this alias analysis.
135   std::unique_ptr<HloDataflowAnalysis> dataflow_analysis_;
136 
137   // A map indicating which buffer a value is contained in.
138   absl::flat_hash_map<const HloValue*, HloBuffer*> value_to_buffer_;
139 
140   // A lazily constructed vector containing all HloBuffers sorted by
141   // HloBuffer::Id.
142   std::vector<HloBuffer> buffers_;
143 };
144 
145 }  // namespace xla
146 
147 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_ALIAS_ANALYSIS_H_
148