Spaces:
Sleeping
Sleeping
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. | |
Licensed under the Apache License, Version 2.0 (the "License"); | |
you may not use this file except in compliance with the License. | |
You may obtain a copy of the License at | |
http://www.apache.org/licenses/LICENSE-2.0 | |
Unless required by applicable law or agreed to in writing, software | |
distributed under the License is distributed on an "AS IS" BASIS, | |
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
See the License for the specific language governing permissions and | |
limitations under the License. | |
==============================================================================*/ | |
// Classes to maintain a static registry of memory allocators | |
namespace tensorflow { | |
// A global AllocatorRegistry is used to hold allocators for CPU backends | |
class AllocatorRegistry { | |
public: | |
// Add an allocator to the registry. Caller releases ownership of | |
// 'allocator'. | |
void Register(const string& name, int priority, Allocator* allocator); | |
// Return allocator with highest priority | |
// If multiple allocators have the same high priority, return one of them | |
Allocator* GetAllocator(); | |
// Returns the global registry of allocators. | |
static AllocatorRegistry* Global(); | |
private: | |
typedef struct { | |
string name; | |
int priority; | |
Allocator* allocator; // not owned | |
} AllocatorRegistryEntry; | |
// Returns the Allocator registered for 'name' and 'priority', | |
// or 'nullptr' if not found. | |
Allocator* GetRegisteredAllocator(const string& name, int priority); | |
std::vector<AllocatorRegistryEntry> allocators_; | |
Allocator* m_curr_allocator_; // not owned | |
}; | |
namespace allocator_registration { | |
class AllocatorRegistration { | |
public: | |
AllocatorRegistration(const string& name, int priority, | |
Allocator* allocator) { | |
AllocatorRegistry::Global()->Register(name, priority, allocator); | |
} | |
}; | |
} // namespace allocator_registration | |
} // namespace tensorflow | |