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 // Classes to maintain a static registry of memory allocator factories. 17 #ifndef TENSORFLOW_CORE_FRAMEWORK_ALLOCATOR_REGISTRY_H_ 18 #define TENSORFLOW_CORE_FRAMEWORK_ALLOCATOR_REGISTRY_H_ 19 20 #include <string> 21 #include <vector> 22 23 #include "tensorflow/core/framework/allocator.h" 24 #include "tensorflow/core/platform/macros.h" 25 #include "tensorflow/core/platform/mutex.h" 26 #include "tensorflow/core/platform/numa.h" 27 28 namespace tensorflow { 29 30 class AllocatorFactory { 31 public: ~AllocatorFactory()32 virtual ~AllocatorFactory() {} 33 34 // Returns true if the factory will create a functionally different 35 // SubAllocator for different (legal) values of numa_node. NumaEnabled()36 virtual bool NumaEnabled() { return false; } 37 38 // Create an Allocator. 39 virtual Allocator* CreateAllocator() = 0; 40 41 // Create a SubAllocator. If NumaEnabled() is true, then returned SubAllocator 42 // will allocate memory local to numa_node. If numa_node == kNUMANoAffinity 43 // then allocated memory is not specific to any NUMA node. 44 virtual SubAllocator* CreateSubAllocator(int numa_node) = 0; 45 }; 46 47 // ProcessState is defined in a package that cannot be a dependency of 48 // framework. This definition allows us to access the one method we need. 49 class ProcessStateInterface { 50 public: 51 virtual Allocator* GetCPUAllocator(int numa_node) = 0; 52 }; 53 54 // A singleton registry of AllocatorFactories. 55 // 56 // Allocators should be obtained through ProcessState or cpu_allocator() 57 // (deprecated), not directly through this interface. The purpose of this 58 // registry is to allow link-time discovery of multiple AllocatorFactories among 59 // which ProcessState will obtain the best fit at startup. 60 class AllocatorFactoryRegistry { 61 public: AllocatorFactoryRegistry()62 AllocatorFactoryRegistry() {} ~AllocatorFactoryRegistry()63 ~AllocatorFactoryRegistry() {} 64 65 void Register(const char* source_file, int source_line, const string& name, 66 int priority, AllocatorFactory* factory); 67 68 // Returns 'best fit' Allocator. Find the factory with the highest priority 69 // and return an allocator constructed by it. If multiple factories have 70 // been registered with the same priority, picks one by unspecified criteria. 71 Allocator* GetAllocator(); 72 73 // Returns 'best fit' SubAllocator. First look for the highest priority 74 // factory that is NUMA-enabled. If none is registered, fall back to the 75 // highest priority non-NUMA-enabled factory. If NUMA-enabled, return a 76 // SubAllocator specific to numa_node, otherwise return a NUMA-insensitive 77 // SubAllocator. 78 SubAllocator* GetSubAllocator(int numa_node); 79 80 // Returns the singleton value. 81 static AllocatorFactoryRegistry* singleton(); 82 process_state()83 ProcessStateInterface* process_state() const { return process_state_; } 84 85 protected: 86 friend class ProcessState; 87 ProcessStateInterface* process_state_ = nullptr; 88 89 private: 90 mutex mu_; 91 bool first_alloc_made_ = false; 92 struct FactoryEntry { 93 const char* source_file; 94 int source_line; 95 string name; 96 int priority; 97 std::unique_ptr<AllocatorFactory> factory; 98 std::unique_ptr<Allocator> allocator; 99 // Index 0 corresponds to kNUMANoAffinity, other indices are (numa_node + 100 // 1). 101 std::vector<std::unique_ptr<SubAllocator>> sub_allocators; 102 }; 103 std::vector<FactoryEntry> factories_ GUARDED_BY(mu_); 104 105 // Returns any FactoryEntry registered under 'name' and 'priority', 106 // or 'nullptr' if none found. 107 const FactoryEntry* FindEntry(const string& name, int priority) const 108 EXCLUSIVE_LOCKS_REQUIRED(mu_); 109 110 TF_DISALLOW_COPY_AND_ASSIGN(AllocatorFactoryRegistry); 111 }; 112 113 class AllocatorFactoryRegistration { 114 public: AllocatorFactoryRegistration(const char * file,int line,const string & name,int priority,AllocatorFactory * factory)115 AllocatorFactoryRegistration(const char* file, int line, const string& name, 116 int priority, AllocatorFactory* factory) { 117 AllocatorFactoryRegistry::singleton()->Register(file, line, name, priority, 118 factory); 119 } 120 }; 121 122 #define REGISTER_MEM_ALLOCATOR(name, priority, factory) \ 123 REGISTER_MEM_ALLOCATOR_UNIQ_HELPER(__COUNTER__, __FILE__, __LINE__, name, \ 124 priority, factory) 125 126 #define REGISTER_MEM_ALLOCATOR_UNIQ_HELPER(ctr, file, line, name, priority, \ 127 factory) \ 128 REGISTER_MEM_ALLOCATOR_UNIQ(ctr, file, line, name, priority, factory) 129 130 #define REGISTER_MEM_ALLOCATOR_UNIQ(ctr, file, line, name, priority, factory) \ 131 static AllocatorFactoryRegistration allocator_factory_reg_##ctr( \ 132 file, line, name, priority, new factory) 133 134 } // namespace tensorflow 135 136 #endif // TENSORFLOW_CORE_FRAMEWORK_ALLOCATOR_REGISTRY_H_ 137