| 94 | } |
| 95 | |
| 96 | Allocator* GPUProcessState::GetGPUAllocator(const GPUOptions& options, |
| 97 | TfGpuId tf_gpu_id, |
| 98 | size_t total_bytes) { |
| 99 | CHECK(process_state_); |
| 100 | #if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ |
| 101 | (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) |
| 102 | const string& allocator_type = options.allocator_type(); |
| 103 | mutex_lock lock(mu_); |
| 104 | GpuIdUtil::CheckValidTfGpuId(tf_gpu_id); |
| 105 | |
| 106 | if (tf_gpu_id.value() >= static_cast<int64>(gpu_allocators_.size())) { |
| 107 | gpu_allocators_.resize(tf_gpu_id.value() + 1); |
| 108 | } |
| 109 | |
| 110 | AllocatorParts& allocator_parts = gpu_allocators_[tf_gpu_id.value()]; |
| 111 | if (allocator_parts.allocator == nullptr) { |
| 112 | // Validate allocator types. |
| 113 | if (!allocator_type.empty() && allocator_type != "BFC") { |
| 114 | LOG(ERROR) << "Invalid allocator type: " << allocator_type; |
| 115 | return nullptr; |
| 116 | } |
| 117 | |
| 118 | PlatformGpuId platform_gpu_id; |
| 119 | TF_CHECK_OK(GpuIdManager::TfToPlatformGpuId(tf_gpu_id, &platform_gpu_id)); |
| 120 | int bus_id = BusIdForGPU(tf_gpu_id); |
| 121 | DCHECK_GE(bus_id, 0); |
| 122 | while (bus_id >= gpu_visitors_.size()) { |
| 123 | gpu_visitors_.push_back({}); |
| 124 | } |
| 125 | |
| 126 | bool use_mps = GpuIdUtil::EnableMPS(); |
| 127 | se::StreamExecutor* stream_exec = nullptr; |
| 128 | GPUMemAllocator* sub_allocator = nullptr; |
| 129 | if (use_mps) { |
| 130 | stream_exec = GpuIdUtil::ExecutorForTfGpuId(platform_gpu_id, tf_gpu_id).ValueOrDie(); |
| 131 | sub_allocator = new GPUMemAllocator( |
| 132 | stream_exec, |
| 133 | tf_gpu_id, |
| 134 | (options.per_process_gpu_memory_fraction() > 1.0 || |
| 135 | options.experimental().use_unified_memory()), |
| 136 | gpu_visitors_[bus_id], {}); |
| 137 | } else { |
| 138 | stream_exec = GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(); |
| 139 | sub_allocator = new GPUMemAllocator( |
| 140 | stream_exec, |
| 141 | platform_gpu_id, |
| 142 | (options.per_process_gpu_memory_fraction() > 1.0 || |
| 143 | options.experimental().use_unified_memory()), |
| 144 | gpu_visitors_[bus_id], {}); |
| 145 | } |
| 146 | Allocator* gpu_allocator = nullptr; |
| 147 | GPUBFCAllocator* gpu_bfc_allocator = nullptr; |
| 148 | if (useTensorPoolAllocator()) { |
| 149 | gpu_allocator = |
| 150 | new GPUTensorPoolAllocator(sub_allocator, |
| 151 | strings::StrCat("GPU_", tf_gpu_id.value(), "_tensorpool"), |
| 152 | total_bytes); |
| 153 | } else { |