diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-07-18 09:20:07 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-18 09:23:10 -0700 |
| commit | 41d7dc4a02440cef93477445bb58d4ee932a34a4 (patch) | |
| tree | 3d76c84fbe4ae2e0f37c0435086797a765b1debd | |
| parent | 5362d9dbedea6f93b1065e31cd0ee16ab110c7e1 (diff) | |
Change the return type of ProcessState::GetCPUAllocator() from
Allocator* to VisitableAllocator*. All implementations are already
VisitableAllocators, this just makes that fact clear and useful.
PiperOrigin-RevId: 205088826
| -rw-r--r-- | tensorflow/core/common_runtime/gpu/gpu_device.cc | 1 | ||||
| -rw-r--r-- | tensorflow/core/common_runtime/process_state.cc | 2 | ||||
| -rw-r--r-- | tensorflow/core/common_runtime/process_state.h | 4 |
3 files changed, 4 insertions, 3 deletions
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.cc b/tensorflow/core/common_runtime/gpu/gpu_device.cc index 3cb51b0dbc..c87841e2dc 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_device.cc +++ b/tensorflow/core/common_runtime/gpu/gpu_device.cc @@ -41,6 +41,7 @@ limitations under the License. #include "tensorflow/core/common_runtime/gpu/gpu_util.h" #include "tensorflow/core/common_runtime/gpu_device_context.h" #include "tensorflow/core/common_runtime/local_device.h" +#include "tensorflow/core/common_runtime/visitable_allocator.h" #include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/op_kernel.h" diff --git a/tensorflow/core/common_runtime/process_state.cc b/tensorflow/core/common_runtime/process_state.cc index 4d83b25ce6..447338e7bd 100644 --- a/tensorflow/core/common_runtime/process_state.cc +++ b/tensorflow/core/common_runtime/process_state.cc @@ -71,7 +71,7 @@ ProcessState::MemDesc ProcessState::PtrType(const void* ptr) { return MemDesc(); } -Allocator* ProcessState::GetCPUAllocator(int numa_node) { +VisitableAllocator* ProcessState::GetCPUAllocator(int numa_node) { CHECK_GE(numa_node, 0); if (!numa_enabled_) numa_node = 0; mutex_lock lock(mu_); diff --git a/tensorflow/core/common_runtime/process_state.h b/tensorflow/core/common_runtime/process_state.h index 0f4ae230bb..2892677333 100644 --- a/tensorflow/core/common_runtime/process_state.h +++ b/tensorflow/core/common_runtime/process_state.h @@ -65,7 +65,7 @@ class ProcessState { // Returns the one CPUAllocator used for the given numa_node. // TEMPORARY: ignores numa_node. - Allocator* GetCPUAllocator(int numa_node); + VisitableAllocator* GetCPUAllocator(int numa_node); typedef std::unordered_map<const void*, MemDesc> MDMap; @@ -87,7 +87,7 @@ class ProcessState { mutex mu_; - std::vector<Allocator*> cpu_allocators_ GUARDED_BY(mu_); + std::vector<VisitableAllocator*> cpu_allocators_ GUARDED_BY(mu_); virtual ~ProcessState(); |