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:
~ProcessStateInterface()51   virtual ~ProcessStateInterface() {}
52   virtual Allocator* GetCPUAllocator(int numa_node) = 0;
53 };
54 
55 // A singleton registry of AllocatorFactories.
56 //
57 // Allocators should be obtained through ProcessState or cpu_allocator()
58 // (deprecated), not directly through this interface.  The purpose of this
59 // registry is to allow link-time discovery of multiple AllocatorFactories among
60 // which ProcessState will obtain the best fit at startup.
61 class AllocatorFactoryRegistry {
62  public:
AllocatorFactoryRegistry()63   AllocatorFactoryRegistry() {}
~AllocatorFactoryRegistry()64   ~AllocatorFactoryRegistry() {}
65 
66   void Register(const char* source_file, int source_line, const string& name,
67                 int priority, AllocatorFactory* factory);
68 
69   // Returns 'best fit' Allocator.  Find the factory with the highest priority
70   // and return an allocator constructed by it.  If multiple factories have
71   // been registered with the same priority, picks one by unspecified criteria.
72   Allocator* GetAllocator();
73 
74   // Returns 'best fit' SubAllocator.  First look for the highest priority
75   // factory that is NUMA-enabled.  If none is registered, fall back to the
76   // highest priority non-NUMA-enabled factory.  If NUMA-enabled, return a
77   // SubAllocator specific to numa_node, otherwise return a NUMA-insensitive
78   // SubAllocator.
79   SubAllocator* GetSubAllocator(int numa_node);
80 
81   // Returns the singleton value.
82   static AllocatorFactoryRegistry* singleton();
83 
process_state()84   ProcessStateInterface* process_state() const { return process_state_; }
85 
86  protected:
87   friend class ProcessState;
88   ProcessStateInterface* process_state_ = nullptr;
89 
90  private:
91   mutex mu_;
92   bool first_alloc_made_ = false;
93   struct FactoryEntry {
94     const char* source_file;
95     int source_line;
96     string name;
97     int priority;
98     std::unique_ptr<AllocatorFactory> factory;
99     std::unique_ptr<Allocator> allocator;
100     // Index 0 corresponds to kNUMANoAffinity, other indices are (numa_node +
101     // 1).
102     std::vector<std::unique_ptr<SubAllocator>> sub_allocators;
103   };
104   std::vector<FactoryEntry> factories_ TF_GUARDED_BY(mu_);
105 
106   // Returns any FactoryEntry registered under 'name' and 'priority',
107   // or 'nullptr' if none found.
108   const FactoryEntry* FindEntry(const string& name, int priority) const
109       TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
110 
111   TF_DISALLOW_COPY_AND_ASSIGN(AllocatorFactoryRegistry);
112 };
113 
114 class AllocatorFactoryRegistration {
115  public:
AllocatorFactoryRegistration(const char * file,int line,const string & name,int priority,AllocatorFactory * factory)116   AllocatorFactoryRegistration(const char* file, int line, const string& name,
117                                int priority, AllocatorFactory* factory) {
118     AllocatorFactoryRegistry::singleton()->Register(file, line, name, priority,
119                                                     factory);
120   }
121 };
122 
123 #define REGISTER_MEM_ALLOCATOR(name, priority, factory)                     \
124   REGISTER_MEM_ALLOCATOR_UNIQ_HELPER(__COUNTER__, __FILE__, __LINE__, name, \
125                                      priority, factory)
126 
127 #define REGISTER_MEM_ALLOCATOR_UNIQ_HELPER(ctr, file, line, name, priority, \
128                                            factory)                         \
129   REGISTER_MEM_ALLOCATOR_UNIQ(ctr, file, line, name, priority, factory)
130 
131 #define REGISTER_MEM_ALLOCATOR_UNIQ(ctr, file, line, name, priority, factory) \
132   static AllocatorFactoryRegistration allocator_factory_reg_##ctr(            \
133       file, line, name, priority, new factory)
134 
135 }  // namespace tensorflow
136 
137 #endif  // TENSORFLOW_CORE_FRAMEWORK_ALLOCATOR_REGISTRY_H_
138