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_HEAP_SIMULATOR_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_HEAP_SIMULATOR_H_ 18 19 #include <memory> 20 #include <set> 21 #include <utility> 22 #include <vector> 23 24 #include "absl/container/flat_hash_map.h" 25 #include "absl/container/flat_hash_set.h" 26 #include "tensorflow/compiler/xla/service/buffer_value.h" 27 #include "tensorflow/compiler/xla/service/buffer_value_containers.h" 28 #include "tensorflow/compiler/xla/service/hlo.pb.h" 29 #include "tensorflow/compiler/xla/service/hlo_computation.h" 30 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 31 #include "tensorflow/compiler/xla/service/hlo_ordering.h" 32 #include "tensorflow/compiler/xla/service/hlo_schedule.h" 33 #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" 34 #include "tensorflow/compiler/xla/statusor.h" 35 36 namespace xla { 37 38 // Forward declare classes defined below. 39 class HeapAlgorithm; 40 class NoFragmentationStatsHeap; 41 42 // HeapSimulator assigns buffer offsets by running a simulation of a regular 43 // memory heap with Alloc and Free calls. It only works for completely 44 // sequential instruction sequences. Unlike regular heaps, we have the 45 // advantage that the sequence of Alloc and Free calls is known up-front; we 46 // don't need to return the assignment of buffer offsets until the very end. 47 class HeapSimulator { 48 public: 49 // Chunk represents a contiguous piece of memory. Each BufferValue will be 50 // associated with a chunk in the assignment result. 51 struct Chunk { 52 int64 offset; 53 int64 size; 54 chunk_endChunk55 int64 chunk_end() const { return offset + size; } 56 }; 57 58 // Result represents the result of the heap simulation. 59 struct Result { 60 // The assignment of buffers to chunks. 61 absl::flat_hash_map<const BufferValue*, Chunk> chunk_map; 62 63 // The total size in bytes of the heap, containing all assigned chunks. 64 int64 heap_size = 0; 65 66 // The total size in bytes of heap fragmentation. 67 int64 fragmentation_size = 0; 68 69 // A trace of heap simulation events. 70 HeapSimulatorTrace debug_trace; 71 }; 72 73 // The different options to be passed to the Run() APIs. 74 struct Options { OptionsOptions75 Options() 76 : may_reuse_operand_buffers(true), 77 alloc_constants(false), 78 buffers_to_assign(nullptr) {} 79 80 // Whether a buffer about to be Free()-ed, can be recycled for a new born 81 // one, hence collapsing Free()+Alloc() calls (default true). 82 bool may_reuse_operand_buffers; 83 // Whether to issue Alloc() and Free() calls for constants (default false). 84 bool alloc_constants; 85 // If 'buffers_to_assign' is provided, only those buffers are assigned 86 // offsets, otherwise all buffers defined by the instructions are assigned. 87 const BufferValueFlatSet* buffers_to_assign; 88 }; 89 90 // Returns the minimum memory required to compute an HLO module where all 91 // computations have been scheduled (represented by the given 92 // schedule), assuming no fragmentation. 93 static StatusOr<int64> MinimumMemoryForModule( 94 const HloSchedule& schedule, 95 const LogicalBuffer::SizeFunction& size_function); 96 97 // Returns the minimum memory required to compute the given computation, 98 // assuming no fragmentation. 99 static StatusOr<int64> MinimumMemoryForComputation( 100 const HloComputation& computation, const HloInstructionSequence& sequence, 101 const TuplePointsToAnalysis& points_to_analysis, 102 const LogicalBuffer::SizeFunction& size_function, 103 const absl::flat_hash_map<const HloComputation*, int64>* 104 memory_by_computation = nullptr); 105 106 // Run the heap simulation with the given algorithm, assuming the given 107 // schedule, which must contain a topologically-consistent total 108 // ordering of all instructions within each computation. The result is invalid 109 // if instructions are not run in exactly this sequence. 110 // 111 // Running heap simulation on the whole module tends to save memory, compared 112 // to running on a per-computation basis, since we can re-use buffer space for 113 // called sub-computations. 114 // 115 static StatusOr<Result> Run(std::unique_ptr<HeapAlgorithm> algorithm, 116 const HloModule& module, 117 const HloSchedule& schedule, 118 const TuplePointsToAnalysis& points_to_analysis, 119 const BufferValue::SizeFunction& size_fn, 120 const Options& options = Options()); 121 122 // Same as above, but runs on a single computation. The 'instruction_sequence' 123 // must contain a topologically-consistent total ordering of all instructions 124 // in the computation. The result is invalid if instructions are not run in 125 // exactly this sequence. 126 static StatusOr<Result> Run( 127 std::unique_ptr<HeapAlgorithm> algorithm, 128 const HloComputation& computation, 129 const HloInstructionSequence& instruction_sequence, 130 const TuplePointsToAnalysis& points_to_analysis, 131 const BufferValue::SizeFunction& size_fn, 132 const Options& options = Options(), 133 const absl::flat_hash_map<const HloComputation*, int64>* 134 memory_by_computation = nullptr); 135 136 private: 137 // If 'schedule' is non-null, it is used to find kCall and kWhile 138 // sub-computations, and the heap simulation for those sub-computations will 139 // be run recursively. I.e. the simulation is run over the whole module. 140 HeapSimulator(std::unique_ptr<HeapAlgorithm> algorithm, 141 const BufferValue::SizeFunction& size_fn, 142 const Options& options, const HloSchedule* schedule = nullptr, 143 const absl::flat_hash_map<const HloComputation*, int64>* 144 memory_by_computation = nullptr); 145 ~HeapSimulator(); 146 147 Status RunComputation(const HloComputation& computation, 148 const HloInstructionSequence& instruction_sequence, 149 const TuplePointsToAnalysis& points_to_analysis); 150 151 bool IgnoreBuffer(const BufferValue* buffer) const; 152 void Alloc(const BufferValue* buffer, const HloInstruction* instruction); 153 void Free(const BufferValue* buffer, const HloInstruction* instruction); 154 void ShareBuffer(const BufferValue* buffer, const BufferValue* shared, 155 const HloInstruction* instruction); 156 Result Finish(); 157 158 void FillDebugTrace(HeapSimulatorTrace::Event::Kind kind, 159 const BufferValue* buffer, 160 const HloInstruction* instruction, 161 const BufferValue* share_with_canonical); 162 163 // Counterintuitive: the algorithm_ itself can be a NoFragmentationStatsHeap, 164 // in which case we are calculating the same allocs/frees twice in the 165 // simulation. 166 const std::unique_ptr<NoFragmentationStatsHeap> no_fragmentation_stats_; 167 const std::unique_ptr<HeapAlgorithm> algorithm_; 168 const BufferValue::SizeFunction size_fn_; 169 const Options options_; 170 // schedule_ is set by buffer assignment, and memory_by_computation_ is 171 // set by hlo scheduling. Then, in RunComputation, we check both in order to 172 // handle subcomputations. It would be good to unify the handling of 173 // subcomputations, but it's not clear how. 174 const HloSchedule* schedule_; 175 const absl::flat_hash_map<const HloComputation*, int64>* 176 memory_by_computation_; 177 178 // In addition to Alloc and Free, the heap simulator exposes a concept of 179 // buffer sharing. When ShareBuffer is called, instead of allocating new 180 // space for the buffer, it associates the buffer with a previously allocated 181 // (or shared) buffer. Each group of mutually-shared buffers points to a 182 // single SharedGroup instance, which is a shared control block. 183 // 184 // This forced buffer sharing is hidden from the underlying heap algorithm, 185 // which only sees a regular Alloc call on the canonical buffer. The 186 // corresponding Free call is delayed until the liveness of all shared buffers 187 // in the group has expired, which is tracked via the refcount. The results 188 // are post-processed in Finish to add chunks for shared buffers. 189 // 190 // The shared_buffers_ map associates each shared buffer (including the 191 // canonical) to its SharedGroup control block. 192 struct SharedGroup { 193 const BufferValue* canonical = nullptr; 194 int64 refcount = 0; 195 }; 196 absl::flat_hash_map<const BufferValue*, std::shared_ptr<SharedGroup>> 197 shared_buffers_; 198 199 // Hold some sets for error-checking the sequence of Alloc and Free calls. 200 absl::flat_hash_set<const BufferValue*> allocated_buffers_; 201 absl::flat_hash_set<const BufferValue*> freed_buffers_; 202 203 // Debugging information filled in while the heap simulator runs. 204 HeapSimulatorTrace debug_trace_; 205 }; 206 207 // Abstract base class describing a heap simulation algorithm that assigns 208 // offsets to buffers. A sequence of Alloc / Free calls will be made, with the 209 // same semantics as a regular memory heap. Finish will be called at the end to 210 // collect the simulation results. 211 class HeapAlgorithm { 212 public: 213 using Chunk = HeapSimulator::Chunk; 214 using Result = HeapSimulator::Result; 215 216 virtual ~HeapAlgorithm() = default; 217 218 // Alloc allocates a buffer of 'size' bytes. 219 virtual void Alloc(const BufferValue* buffer, int64 size) = 0; 220 221 // Takes memory usage of subcomputations into account when calculating the 222 // memory usage of a computation. Currently, we don't handle buffer aliasing 223 // between computations entirely correctly. We are careful to not double count 224 // for the output buffers of whiles/conds/calls. But we don't take into 225 // account other aliases, such as for the while init. A more thorough solution 226 // would require something like BufferAssignment::BuildColocatedBufferSets. 227 // TODO(b/65835246): 228 // Since TuplePointsToAnalysis is being replaced with a module-aware alias 229 // analysis, it's not worth making major changes to HeapSimulator now. AccountForSubcomputationMemory(const HloInstruction * instruction,int64 alloc_size_by_instruction,const absl::flat_hash_map<const HloComputation *,int64> & memory_by_computation)230 virtual void AccountForSubcomputationMemory( 231 const HloInstruction* instruction, 232 // The total number of bytes allocated by instruction. 233 int64 alloc_size_by_instruction, 234 const absl::flat_hash_map<const HloComputation*, int64>& 235 memory_by_computation) {} 236 237 // Free de-allocates a previously allocated buffer. 238 virtual void Free(const BufferValue* buffer, int64 size) = 0; 239 240 // Finish collects the buffer offset assignment results. Free may only be 241 // called once, after the Alloc and Free calls. 242 virtual Result Finish() = 0; 243 }; 244 245 // NoFragmentationStatsHeap computes the heap size assuming no fragmentation; 246 // this is the absolute minimum size for a given instruction sequence. The 247 // result.chunk_map returned in Finish is always empty, since we only collect 248 // stats, and don't actually compute chunk assignments. 249 class NoFragmentationStatsHeap : public HeapAlgorithm { 250 public: 251 NoFragmentationStatsHeap() = default; 252 ~NoFragmentationStatsHeap() override = default; 253 254 void Alloc(const BufferValue* buffer, int64 size) override; 255 256 void AccountForSubcomputationMemory( 257 const HloInstruction* instruction, int64 alloc_size_by_instruction, 258 const absl::flat_hash_map<const HloComputation*, int64>& 259 memory_by_computation) override; 260 261 void Free(const BufferValue* buffer, int64 size) override; 262 263 Result Finish() override; 264 265 private: 266 int64 current_heap_size_ = 0; 267 int64 max_heap_size_ = 0; 268 }; 269 270 // DecreasingSizeRunsHeap collects runs of Alloc and Free calls, sorts them by 271 // decreasing size, and delegates the actual calls to another heap algorithm. 272 // This greedy heuristic tends to reduce fragmentation for all algorithms. 273 class DecreasingSizeRunsHeap : public HeapAlgorithm { 274 public: DecreasingSizeRunsHeap(std::unique_ptr<HeapAlgorithm> algorithm)275 DecreasingSizeRunsHeap(std::unique_ptr<HeapAlgorithm> algorithm) 276 : algorithm_(std::move(algorithm)) {} ~DecreasingSizeRunsHeap()277 ~DecreasingSizeRunsHeap() override {} 278 279 void Alloc(const BufferValue* buffer, int64 size) override; 280 void Free(const BufferValue* buffer, int64 size) override; 281 Result Finish() override; 282 283 private: 284 // A single Alloc or Free operation that we've buffered in run_. 285 struct Op { 286 const BufferValue* buffer; 287 int64 size; 288 }; 289 290 // Current collection mode; kInit means no ops have been collected yet. 291 enum Mode { kInit, kAlloc, kFree }; 292 293 void SetMode(Mode mode); 294 void CallAndDrainRun(); 295 296 const std::unique_ptr<HeapAlgorithm> algorithm_; 297 std::vector<Op> run_; 298 Mode mode_ = kInit; 299 }; 300 301 // LazyBestFitHeap is a variant of the traditional best-fit heap. This is a 302 // greedy heuristic, based on the idea that delaying offset assignment helps 303 // reduce fragmentation. Here's an example of a "bad" offset assignment, where 304 // a tiny buffer A prevents adjacent free chunks from being coalesced: 305 // BAD: | free |A| free | 306 // If we could have delayed the assignment of A, we might have ended up with: 307 // GOOD: | free |A| 308 // 309 // In general it's actually hard to say whether GOOD is better than BAD; the 310 // heuristic we use is we try to leave large contiguous chunks free, and we try 311 // to avoid growing the overall heap size unless necessary. 312 // 313 // Just like regular best-fit, in Alloc we look for the smallest free chunk that 314 // fits the requested size. Unlike regular best-fit, we postpone offset 315 // assignment for buffers that cannot re-use existing free chunks (and force us 316 // to grow the heap); these buffers are "lazily" assigned offsets in Free. 317 class LazyBestFitHeap : public HeapAlgorithm { 318 public: LazyBestFitHeap(int64 alignment)319 LazyBestFitHeap(int64 alignment) : alignment_(alignment) {} ~LazyBestFitHeap()320 ~LazyBestFitHeap() override {} 321 322 void Alloc(const BufferValue* buffer, int64 size) override; 323 void Free(const BufferValue* buffer, int64 size) override; 324 Result Finish() override; 325 326 private: 327 // Sentry value used to indicate a chunk that wasn't assigned an offset in 328 // Alloc, and will instead be assigned an offset in Free. 329 enum { kLazyAllocOffset = -1 }; 330 331 struct OrderChunkByIncreasingSize { operatorOrderChunkByIncreasingSize332 bool operator()(const Chunk& a, const Chunk& b) const { 333 if (a.size != b.size) return a.size < b.size; 334 return a.offset < b.offset; 335 } 336 }; 337 338 void AddFreeChunk(int64 offset, int64 size); 339 340 const int64 alignment_; 341 Result result_; 342 343 // Maintain the set of free chunks, ordered by increasing size. 344 std::set<Chunk, OrderChunkByIncreasingSize> free_; 345 }; 346 347 // GlobalDecreasingSizeBestFitHeap collects the live intervals of all buffers, 348 // then allocates them in decreasing sizes regardless of the alloc/free time. It 349 // internally tracks the allocated buffers and their live intervals; when 350 // allocating a buffer, it finds the best-fit free chunk during its live 351 // interval. 352 class GlobalDecreasingSizeBestFitHeap : public HeapAlgorithm { 353 public: GlobalDecreasingSizeBestFitHeap(int64 alignment)354 GlobalDecreasingSizeBestFitHeap(int64 alignment) : alignment_(alignment) {} ~GlobalDecreasingSizeBestFitHeap()355 ~GlobalDecreasingSizeBestFitHeap() override {} 356 357 void Alloc(const BufferValue* buffer, int64 size) override; 358 void Free(const BufferValue* buffer, int64 size) override; 359 Result Finish() override; 360 361 private: 362 int64 alignment_; 363 Result result_; 364 365 // The current time represented as an integer. It increments by 1 at each 366 // Alloc or Free call. 367 int64 current_time_ = 0; 368 369 // BufferInterval stores a buffer's size and time interval. 370 struct BufferInterval { 371 const BufferValue* buffer; 372 int64 size; 373 // Alloc time of the buffer. 374 int64 start; 375 // Free time of the buffer. 376 int64 end; 377 }; 378 absl::flat_hash_map<const BufferValue*, BufferInterval> buffer_intervals_; 379 }; 380 381 // A heap algorithm that chooses the best results from other algorithms added to 382 // it. 383 class ChooseBestHeapAlgorithm : public HeapAlgorithm { 384 public: ChooseBestHeapAlgorithm(std::unique_ptr<std::vector<std::unique_ptr<HeapAlgorithm>>> algorithms)385 ChooseBestHeapAlgorithm( 386 std::unique_ptr<std::vector<std::unique_ptr<HeapAlgorithm>>> algorithms) 387 : algorithms_(std::move(*algorithms)) {} ~ChooseBestHeapAlgorithm()388 ~ChooseBestHeapAlgorithm() override {} 389 Alloc(const BufferValue * buffer,int64 size)390 void Alloc(const BufferValue* buffer, int64 size) override { 391 for (auto& algorithm : algorithms_) { 392 algorithm->Alloc(buffer, size); 393 } 394 } 395 Free(const BufferValue * buffer,int64 size)396 void Free(const BufferValue* buffer, int64 size) override { 397 for (auto& algorithm : algorithms_) { 398 algorithm->Free(buffer, size); 399 } 400 } 401 402 Result Finish() override; 403 404 private: 405 std::vector<std::unique_ptr<HeapAlgorithm>> algorithms_; 406 }; 407 408 } // namespace xla 409 410 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HEAP_SIMULATOR_H_ 411