1 /* Copyright 2018 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_TF2XLA_CPU_FUNCTION_RUNTIME_H_ 17 #define TENSORFLOW_COMPILER_TF2XLA_CPU_FUNCTION_RUNTIME_H_ 18 19 #include "tensorflow/core/platform/types.h" 20 21 #include <cassert> 22 23 namespace tensorflow { 24 namespace cpu_function_runtime { 25 // Stores information about one buffer used by an XLA:CPU compiled function. 26 // These buffers are used for holding inputs to the computation, outputs from 27 // the computation and as temporary scratch space. 28 class BufferInfo { 29 public: 30 // Creates a BufferInfo from a serialized encoding generated by `Encode`. BufferInfo(std::pair<uint64,uint64> encoding)31 explicit BufferInfo(std::pair<uint64, uint64> encoding) 32 : entry_param_number_(encoding.second) { 33 Kind kind; 34 uint64 size; 35 Unpack(encoding.first, &kind, &size); 36 kind_ = kind; 37 size_ = size; 38 } 39 40 // Returns true if this buffer stores a constant. These never need to be 41 // allocated by the runtime. is_constant()42 bool is_constant() const { return kind() == Kind::kConstant; } 43 44 // Returns true if this buffer stores an entry parameter. These may or may 45 // not need to be allocated by the runtime, depending on 46 // XlaCompiledCpuFunction::AllocMode. is_entry_parameter()47 bool is_entry_parameter() const { return kind() == Kind::kEntryParameter; } 48 49 // Returns the entry parameter number of this buffer. entry_parameter_number()50 uint64 entry_parameter_number() const { 51 assert(is_entry_parameter()); 52 return entry_param_number_; 53 } 54 55 // Returns true if this buffer is temporary scratch space required by the XLA 56 // computations. These are always allocated by the runtime. is_temp_buffer()57 bool is_temp_buffer() const { return kind() == Kind::kTempBuffer; } 58 59 // Returns true if this buffer is allocated on the C stack or into registers. 60 // These buffers are never allocated by the runtime. is_on_stack_buffer()61 bool is_on_stack_buffer() const { return kind() == Kind::kOnStackBuffer; } 62 63 // Returns the size for this buffer. size()64 uint64 size() const { return size_; } 65 66 // Encodes this BufferInfo into two 64 bit integers that can be used to 67 // reconstruct the BufferInfo later using the constructor. We need this 68 // because we use BufferInfo in places where using protocol buffers would 69 // negatively impact binary size. Encode()70 std::pair<uint64, uint64> Encode() const { 71 static_assert(sizeof(*this) == 16, ""); 72 uint64 upper = Pack(kind(), size_); 73 uint64 lower = entry_param_number_; 74 return {upper, lower}; 75 } 76 77 bool operator==(const BufferInfo& buffer_info) const { 78 if (kind() != buffer_info.kind() || size() != buffer_info.size()) { 79 return false; 80 } 81 return !is_entry_parameter() || 82 entry_parameter_number() == buffer_info.entry_parameter_number(); 83 } 84 85 // Factory methods: 86 MakeTempBuffer(uint64 size)87 static BufferInfo MakeTempBuffer(uint64 size) { 88 return BufferInfo(Kind::kTempBuffer, /*size=*/size, 89 /*entry_param_number=*/-1); 90 } MakeConstant(uint64 size)91 static BufferInfo MakeConstant(uint64 size) { 92 return BufferInfo(Kind::kConstant, /*size=*/size, 93 /*entry_param_number=*/-1); 94 } MakeEntryParameter(uint64 size,uint64 param_number)95 static BufferInfo MakeEntryParameter(uint64 size, uint64 param_number) { 96 return BufferInfo(Kind::kEntryParameter, /*size=*/size, 97 /*entry_param_number=*/param_number); 98 } MakeOnStackBuffer(uint64 size)99 static BufferInfo MakeOnStackBuffer(uint64 size) { 100 return BufferInfo(Kind::kOnStackBuffer, /*size=*/size, 101 /*entry_param_number=*/-1); 102 } 103 104 private: 105 BufferInfo() = default; 106 107 enum class Kind : uint64 { 108 kConstant, 109 kTempBuffer, 110 kEntryParameter, 111 kOnStackBuffer 112 }; 113 kind()114 Kind kind() const { return static_cast<Kind>(kind_); } 115 BufferInfo(Kind kind,uint64 size,uint64 entry_param_number)116 explicit BufferInfo(Kind kind, uint64 size, uint64 entry_param_number) 117 : kind_(kind), size_(size), entry_param_number_(entry_param_number) {} 118 Pack(Kind kind,uint64 size)119 static uint64 Pack(Kind kind, uint64 size) { 120 return (static_cast<uint64>(size) << 2) | static_cast<uint64>(kind); 121 } 122 Unpack(uint64 packed,Kind * kind,uint64 * size)123 static void Unpack(uint64 packed, Kind* kind, uint64* size) { 124 *size = packed >> 2; 125 *kind = static_cast<Kind>((packed << 62) >> 62); 126 } 127 128 Kind kind_ : 2; 129 uint64 size_ : 62; 130 int64 entry_param_number_; 131 }; 132 133 // Align to 64-bytes, to mimic tensorflow::Allocator::kAllocatorAlignment. 134 constexpr size_t kAlign = 64; 135 136 // AlignedBufferBytes returns the sum of the size of each buffer in 137 // `buffer_infos`, skipping constants, on-stack buffers and, if 138 // allocate_entry_params is false, entry parameters. There are `n` entries in 139 // `buffer_infos`. Each buffer is aligned to kAlign byte boundaries. 140 size_t AlignedBufferBytes(const BufferInfo* buffer_infos, size_t n, 141 bool allocate_entry_params); 142 143 // MallocContiguousBuffers allocates buffers for use by the entry point 144 // generated by tfcompile. There are `n` entries in `buffer_infos`. If 145 // `annotate_initialized` is set, the allocated memory will be annotated as 146 // having been initialized - this is useful when allocating temporary buffers. 147 // If allocate_entry_params is true then allocates temp buffers and entry 148 // parameters, otherwise allocated only temp buffers. Slots in `bufs` 149 // corresponding to unallocated buffers are set to nullptr. 150 // 151 // A single contiguous block of memory is allocated, and portions of it are 152 // parceled out into `bufs`, which must have space for `n` entries. Returns 153 // the head of the allocated contiguous block, which should be passed to 154 // FreeContiguous when the buffers are no longer in use. 155 void* MallocContiguousBuffers(const BufferInfo* buffer_infos, size_t n, 156 bool allocate_entry_params, void** bufs, 157 bool annotate_initialized); 158 159 // FreeContiguous frees the contiguous block of memory allocated by 160 // MallocContiguousBuffers. 161 void FreeContiguous(void* contiguous); 162 } // namespace cpu_function_runtime 163 } // namespace tensorflow 164 165 #endif // TENSORFLOW_COMPILER_TF2XLA_CPU_FUNCTION_RUNTIME_H_ 166