| 60 | } |
| 61 | |
| 62 | Allocator* ProcessState::GetCPUAllocator(int numa_node) { |
| 63 | if (!numa_enabled_ || numa_node == port::kNUMANoAffinity) numa_node = 0; |
| 64 | mutex_lock lock(mu_); |
| 65 | while (cpu_allocators_.size() <= static_cast<size_t>(numa_node)) { |
| 66 | // If visitors have been defined we need an Allocator built from |
| 67 | // a SubAllocator. Prefer BFCAllocator, but fall back to PoolAllocator |
| 68 | // depending on env var setting. |
| 69 | const bool alloc_visitors_defined = |
| 70 | (!cpu_alloc_visitors_.empty() || !cpu_free_visitors_.empty()); |
| 71 | bool use_bfc_allocator = false; |
| 72 | Status status = ReadBoolFromEnvVar( |
| 73 | "TF_CPU_ALLOCATOR_USE_BFC", alloc_visitors_defined, &use_bfc_allocator); |
| 74 | if (!status.ok()) { |
| 75 | LOG(ERROR) << "GetCPUAllocator: " << status.error_message(); |
| 76 | } |
| 77 | Allocator* allocator = nullptr; |
| 78 | SubAllocator* sub_allocator = |
| 79 | (numa_enabled_ || alloc_visitors_defined || use_bfc_allocator) |
| 80 | ? new BasicCPUAllocator( |
| 81 | numa_enabled_ ? numa_node : port::kNUMANoAffinity, |
| 82 | cpu_alloc_visitors_, cpu_free_visitors_) |
| 83 | : nullptr; |
| 84 | if (use_bfc_allocator) { |
| 85 | // TODO(reedwm): evaluate whether 64GB by default is the best choice. |
| 86 | int64 cpu_mem_limit_in_mb = -1; |
| 87 | Status status = ReadInt64FromEnvVar("TF_CPU_BFC_MEM_LIMIT_IN_MB", |
| 88 | 1LL << 16 /*64GB max by default*/, |
| 89 | &cpu_mem_limit_in_mb); |
| 90 | if (!status.ok()) { |
| 91 | LOG(ERROR) << "GetCPUAllocator: " << status.error_message(); |
| 92 | } |
| 93 | int64 cpu_mem_limit = cpu_mem_limit_in_mb * (1LL << 20); |
| 94 | DCHECK(sub_allocator); |
| 95 | allocator = |
| 96 | new BFCAllocator(sub_allocator, cpu_mem_limit, true /*allow_growth*/, |
| 97 | "bfc_cpu_allocator_for_gpu" /*name*/); |
| 98 | VLOG(2) << "Using BFCAllocator with memory limit of " |
| 99 | << cpu_mem_limit_in_mb << " MB for ProcessState CPU allocator"; |
| 100 | } else if (sub_allocator) { |
| 101 | DCHECK(sub_allocator); |
| 102 | allocator = |
| 103 | new PoolAllocator(100 /*pool_size_limit*/, true /*auto_resize*/, |
| 104 | sub_allocator, new NoopRounder, "cpu_pool"); |
| 105 | VLOG(2) << "Using PoolAllocator for ProcessState CPU allocator " |
| 106 | << "numa_enabled_=" << numa_enabled_ |
| 107 | << " numa_node=" << numa_node; |
| 108 | } else { |
| 109 | DCHECK(!sub_allocator); |
| 110 | allocator = cpu_allocator_base(); |
| 111 | } |
| 112 | if (LogMemory::IsEnabled() && !allocator->TracksAllocationSizes()) { |
| 113 | // Wrap the allocator to track allocation ids for better logging |
| 114 | // at the cost of performance. |
| 115 | allocator = new TrackingAllocator(allocator, true); |
| 116 | } |
| 117 | cpu_allocators_.push_back(allocator); |
| 118 | if (!sub_allocator) { |
| 119 | DCHECK(cpu_alloc_visitors_.empty() && cpu_free_visitors_.empty()); |
no test coverage detected