1 /* Copyright 2020 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_MEMORY_SPACE_ASSIGNMENT_REPACKING_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_REPACKING_H_ 18 19 #include "tensorflow/compiler/xla/statusor.h" 20 #include "tensorflow/compiler/xla/types.h" 21 22 namespace xla { 23 24 // An interface to define allocation repacking algorithms. 25 class MemorySpaceAssignmentRepacker { 26 public: MemorySpaceAssignmentRepacker(int64 max_size,int64 alignment)27 MemorySpaceAssignmentRepacker(int64 max_size, int64 alignment) 28 : max_size_(max_size), alignment_(alignment) {} 29 virtual ~MemorySpaceAssignmentRepacker() = default; 30 31 // A contiguous block of allocation consisting of start and end (logical) 32 // times, size, and the initial offset. After repacking, if the repacking was 33 // successful and the allocations were modified, the offset field holds the 34 // new offset. To support aliased allocations, AllocationBlock also includes a 35 // vector of AllocationBlock pointers, called colocations. All AllocationBlock 36 // objects within the colocations must get the same offset. The id should be 37 // unique and is used to ensure determinism for comparison tie-breaker. 38 struct AllocationBlock { 39 int64 start_time; 40 int64 end_time; 41 int64 size; 42 int64 offset; 43 int64 initial_offset; 44 int64 id; 45 std::vector<AllocationBlock*> colocations; 46 ToStringAllocationBlock47 std::string ToString() const { 48 return absl::StrCat("[", start_time, ", ", end_time, "] : size = ", size, 49 ", offset = ", offset, 50 " initial offset = ", initial_offset); 51 } 52 53 // This is required by BufferIntervalCompare as a tie breaker. Use a unique 54 // and deterministic id. 55 bool operator<(const AllocationBlock& other) const { return id < other.id; } 56 }; 57 58 // Repack the AllocationBlocks provided in the parameter. Returns true if 59 // allocations have been modified and false if not. Returns a non-ok status if 60 // there was an error. 61 virtual StatusOr<bool> Repack(absl::Span<AllocationBlock*> allocations) = 0; 62 63 protected: 64 int64 max_size_; 65 int64 alignment_; 66 }; 67 68 } // namespace xla 69 70 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_REPACKING_H_ 71