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_ALLOCATION_TRACKER_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_ALLOCATION_TRACKER_H_ 18 19 #include <map> 20 #include <memory> 21 #include <set> 22 #include <string> 23 #include <vector> 24 25 #include "absl/container/flat_hash_map.h" 26 #include "tensorflow/compiler/xla/service/backend.h" 27 #include "tensorflow/compiler/xla/statusor.h" 28 #include "tensorflow/compiler/xla/types.h" 29 #include "tensorflow/compiler/xla/xla_data.pb.h" 30 #include "tensorflow/core/platform/macros.h" 31 #include "tensorflow/core/platform/mutex.h" 32 #include "tensorflow/core/platform/thread_annotations.h" 33 #include "tensorflow/core/platform/types.h" 34 35 namespace xla { 36 37 // Tracks allocations for the XLA service; allocations can be registered 38 // with shape/device/tag and resolved from a handle for later use. 39 class AllocationTracker { 40 public: 41 // The allocator is used for deallocating memory when allocations are 42 // deregistered. All registered allocations must have the same platform as the 43 // allocator. AllocationTracker(Backend * backend)44 AllocationTracker(Backend* backend) : backend_(backend), next_handle_(1) {} 45 46 // Registers a shaped buffer of device memory, and returns a corresponding 47 // handle that can be used for talking to XLA clients. The given shaped buffer 48 // will be treated as the buffer corresponding to the only replica. 49 StatusOr<GlobalDataHandle> Register(ScopedShapedBuffer shaped_buffer, 50 const string& tag); 51 52 // Registers a vector of shaped buffers of device memory, one per replica, and 53 // returns a corresponding handle that can be used for talking to XLA clients. 54 StatusOr<GlobalDataHandle> RegisterReplicatedBuffers( 55 std::vector<ScopedShapedBuffer> replicated_buffers, const string& tag); 56 57 // Unregister the allocation for the given data handle. 58 Status Unregister(const GlobalDataHandle& data); 59 60 // Returns a vector of global data handles that point to the tuple elements. 61 StatusOr<std::vector<GlobalDataHandle>> DeconstructTuple( 62 const GlobalDataHandle& Data); 63 64 // Resolve a handle from an XLA client to a vector of shaped buffers, one per 65 // replica, or provide an error status to say whether any of those buffers 66 // were not found (or found, but found deallocated). 67 StatusOr<std::vector<const ShapedBuffer*>> Resolve( 68 const GlobalDataHandle& data) const; 69 70 // Resolves a handle from an XLA client and replica id to a shaped buffer, or 71 // provide an error status to say whether it was not found (or found, but 72 // found deallocated). 73 StatusOr<const ShapedBuffer*> ResolveForReplica(const GlobalDataHandle& data, 74 int replica_id) const; 75 76 private: 77 // Data structure encapsulating single memory allocation on the device. 78 struct Allocation { 79 // The pointer to this allocation. 80 OwningDeviceMemory device_memory; 81 82 // This is the number of times this memory allocation is referred to by 83 // registered data handles. 84 int ref_count; 85 }; 86 87 // Internal helper which resolves the given GlobalDataHandle to a 88 // list of ScopedShapedBuffers. 89 StatusOr<std::vector<const ShapedBuffer*>> ResolveInternal( 90 const GlobalDataHandle& data) const EXCLUSIVE_LOCKS_REQUIRED(mutex_); 91 92 // Internal helper which registers a vector of shaped buffers, one per 93 // replica. ShapedBufferTy is either ScopedShapedBuffer or ShapedBuffer. If 94 // it's ShapedBuffer, all of the given buffers must already be tracked by this 95 // object -- presumably this is a call from DeconstructTuple. 96 template <typename ShapedBufferTy> 97 StatusOr<GlobalDataHandle> RegisterInternal( 98 std::vector<ShapedBufferTy> replicated_buffers, const string& tag) 99 EXCLUSIVE_LOCKS_REQUIRED(mutex_); 100 101 // Adds the given device address to the allocation tracker, or if it already 102 // exists, then increment its reference count. 103 void AddAllocationOrIncrementRefCount(se::DeviceMemoryBase device_memory, 104 int device_ordinal) 105 EXCLUSIVE_LOCKS_REQUIRED(mutex_); 106 107 // Decrements the reference count of the given device memory. Then, if it is 108 // zero, deallocate the memory. 109 Status DecrementRefCount(se::DeviceMemoryBase device_memory, 110 int device_ordinal) EXCLUSIVE_LOCKS_REQUIRED(mutex_); 111 112 // A map from device memory opaque value to allocation. One such map is 113 // maintained per device ordinal. 114 using AllocationMap = absl::flat_hash_map<const void*, Allocation>; 115 116 mutable tensorflow::mutex mutex_; 117 118 // Backend to use with this tracker. The backend supplies the memory allocator 119 // to use when deallocating memory. 120 Backend* backend_; 121 122 // The next handle to assign to an allocation, guarded by the same mutex as 123 // the mapping as they'll be mutated at the same time. 124 int64 next_handle_ GUARDED_BY(mutex_); 125 126 // A map from device ordinal to AllocationMap. 127 absl::flat_hash_map<int, AllocationMap> opaque_to_allocation_map_ 128 GUARDED_BY(mutex_); 129 130 // A map from data handle to a vector of shaped buffers that represent the 131 // buffers for different replicas. 132 // 133 // The ShapedBuffers in this map's vectors need to be unique_ptrs, because our 134 // public API returns pointers to them. We expect the concrete class to be 135 // ShapedBuffer and never ScopedShapedBuffer; deallocation of buffers is 136 // handled by opaque_to_allocation_map_. 137 // 138 // The elements of the vectors need to be unique_ptrs because we return 139 // pointers to them. (In theory we could use std::list or something instead, 140 // but we also want to be able to null out these elements.) 141 // 142 // The reason that the elements can't be unique_ptr<ScopedShapedBuffer>s is 143 // the existence of DeconstructTuple(). This function allows us to create a 144 // non-owning "view" into a tuple's sub-buffers. The sub-buffers are then 145 // free'd when both the view *and* the original tuple are Unregistered. This 146 // refcounting is managed in opaque_to_allocation_map_. 147 absl::flat_hash_map<int64, std::vector<std::unique_ptr<ShapedBuffer>>> 148 handle_to_shaped_buffers_ GUARDED_BY(mutex_); 149 150 TF_DISALLOW_COPY_AND_ASSIGN(AllocationTracker); 151 }; 152 153 } // namespace xla 154 155 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_ALLOCATION_TRACKER_H_ 156