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_CORE_COMMON_RUNTIME_GPU_GPU_PROCESS_STATE_H_ 17 #define TENSORFLOW_CORE_COMMON_RUNTIME_GPU_GPU_PROCESS_STATE_H_ 18 19 #include <functional> 20 #include <map> 21 #include <unordered_map> 22 #include <vector> 23 24 #include "tensorflow/core/common_runtime/gpu/gpu_id.h" 25 #include "tensorflow/core/common_runtime/process_state.h" 26 #include "tensorflow/core/common_runtime/shared_counter.h" 27 #include "tensorflow/core/framework/allocator.h" 28 #include "tensorflow/core/platform/mutex.h" 29 #include "tensorflow/core/platform/thread_annotations.h" 30 #include "tensorflow/core/platform/types.h" 31 #include "tensorflow/core/protobuf/config.pb.h" 32 33 namespace tensorflow { 34 35 class Allocator; 36 class GPUBFCAllocator; 37 class PoolAllocator; 38 class SharedCounter; 39 40 // Singleton that manages per-process state when GPUs are present. 41 class GPUProcessState { 42 public: 43 // If ps == nullptr, returns pointer to the single instance of this class to 44 // be used within this process. 45 // 46 // If ps != nullptrs, accepts a value to be returned by all subsequent calls. 47 // A non-null ps may ONLY be provided during program static storage 48 // initialization. Must not be called more than once with a non-null ps. 49 // 50 // If a derived class of GPUProcessState is ever used in a process, it must 51 // always be used in place of this class. In order to ensure that existing 52 // calls to GPUProcessState::singleton() all resolve to the derived instance 53 // instead, this function must be called once during startup, supplying the 54 // derived instance value, prior to any accessor call to this function. 55 static GPUProcessState* singleton(GPUProcessState* ps = nullptr); 56 57 // Query whether any GPU device has been created so far. 58 // Disable thread safety analysis since a race is benign here. HasGPUDevice()59 bool HasGPUDevice() const TF_NO_THREAD_SAFETY_ANALYSIS { 60 return gpu_device_enabled_; 61 } 62 63 // Set the flag to indicate a GPU device has been created. 64 // Disable thread safety analysis since a race is benign here. EnableGPUDevice()65 void EnableGPUDevice() TF_NO_THREAD_SAFETY_ANALYSIS { 66 gpu_device_enabled_ = true; 67 } 68 69 // Returns the one GPU allocator used for the indexed GPU. 70 // Note that this is a system GPU index, not (necessarily) a brain 71 // device index. 72 // 73 // 'total_bytes' is the total number of bytes that should be made 74 // available to the allocator. The first call to this function for 75 // a given tf_device_id creates the allocator, so only the total_bytes 76 // used on that first call is used. 77 // 78 // "Allocator type" describes the type of algorithm to use for the 79 // underlying allocator. REQUIRES: Must be a valid type (see 80 // config.proto for the list of supported strings.). 81 // 82 // REQUIRES: tf_device_id must be a valid id for a BaseGPUDevice available in 83 // the current system environment. Otherwise returns nullptr. 84 virtual Allocator* GetGPUAllocator( 85 const GPUOptions& options, TfDeviceId tf_device_id, size_t total_bytes, 86 const std::vector<TfDeviceId>& peer_gpu_ids); 87 GetGPUAllocator(TfDeviceId tf_device_id)88 Allocator* GetGPUAllocator(TfDeviceId tf_device_id) { 89 return GetGPUAllocator(/*options=*/{}, tf_device_id, /*total_bytes=*/0, 90 /*peer_gpu_ids=*/{}); 91 } 92 NumGPUAllocators()93 int NumGPUAllocators() { 94 mutex_lock l(mu_); 95 return gpu_allocators_.size(); 96 } 97 98 virtual Allocator* GetGpuHostAllocator(int numa_node); 99 100 // Registers a Visitor to be invoked on new chunks of memory allocated by the 101 // SubAllocator of every GPU proximate to the specified bus. The AllocVisitor 102 // is provided with a memory pointer, a GPU id, and the size of the area it 103 // identifies. The pointer is not guaranteed to be valid after the call 104 // terminates. The intention is for this interface to be used for network 105 // device memory registration. "bus_id" is platform-specific. On many 106 // platforms it should be 0. On machines with multiple PCIe buses, it should 107 // be the index of one of the PCIe buses (maybe the NUMA node at which the 108 // PCIe is rooted). If the bus_id is invalid, results are undefined. 109 virtual void AddGPUAllocVisitor(int bus_id, 110 const SubAllocator::Visitor& visitor); 111 112 // Registers a Visitor to be invoked on new chunks of memory allocated by 113 // the SubAllocator of the GpuHostAllocator for the given numa_node. 114 virtual void AddGpuHostAllocVisitor(int numa_node, 115 const SubAllocator::Visitor& visitor); 116 117 // Registers a Visitor to be invoked on each chunk handed back for freeing to 118 // the SubAllocator of the GpuHostAllocator for the given numa_node. 119 virtual void AddGpuHostFreeVisitor(int numa_node, 120 const SubAllocator::Visitor& visitor); 121 122 // Returns bus_id for the given GPU id. 123 virtual int BusIdForGPU(TfDeviceId tf_device_id); 124 125 SharedCounter* GPUAllocatorCounter(TfDeviceId tf_device_id); 126 127 protected: 128 // GPUProcessState is a singleton that should not normally be deleted except 129 // at process shutdown. 130 GPUProcessState(); ~GPUProcessState()131 virtual ~GPUProcessState() {} 132 friend class GPUDeviceTest; 133 134 // Helper method for unit tests to reset the ProcessState singleton by 135 // cleaning up everything. Never use in production. 136 virtual void TestOnlyReset(); 137 mem_desc_map()138 ProcessState::MDMap* mem_desc_map() { 139 if (process_state_) return &process_state_->mem_desc_map_; 140 return nullptr; 141 } 142 143 static GPUProcessState* instance_; 144 ProcessState* process_state_; // Not owned. 145 bool gpu_device_enabled_; 146 147 mutex mu_; 148 149 struct AllocatorParts { 150 std::unique_ptr<Allocator> allocator; 151 std::unique_ptr<SharedCounter> counter; 152 GPUBFCAllocator* bfc_allocator; 153 SubAllocator* sub_allocator; // owned by allocator 154 std::unique_ptr<Allocator> recording_allocator; 155 }; 156 std::vector<AllocatorParts> gpu_allocators_ TF_GUARDED_BY(mu_); 157 std::vector<std::vector<SubAllocator::Visitor>> gpu_visitors_ 158 TF_GUARDED_BY(mu_); 159 160 std::vector<AllocatorParts> gpu_host_allocators_ TF_GUARDED_BY(mu_); 161 std::vector<std::vector<SubAllocator::Visitor>> gpu_host_alloc_visitors_ 162 TF_GUARDED_BY(mu_); 163 std::vector<std::vector<SubAllocator::Visitor>> gpu_host_free_visitors_ 164 TF_GUARDED_BY(mu_); 165 }; 166 167 } // namespace tensorflow 168 #endif // TENSORFLOW_CORE_COMMON_RUNTIME_GPU_GPU_PROCESS_STATE_H_ 169