| 130 | } |
| 131 | |
| 132 | void LocalDevice::Init(const SessionOptions& options, |
| 133 | const DeviceAttributes& attributes, |
| 134 | const DeviceGlobalThreadPoolOptions& opt) { |
| 135 | // Log info messages if TensorFlow is not compiled with instructions that |
| 136 | // could speed up performance and are available on the current CPU. |
| 137 | port::InfoAboutUnusedCPUFeatures(); |
| 138 | LocalDevice::EigenThreadPoolInfo* tp_info; |
| 139 | |
| 140 | if (OverrideGlobalThreadPoolFromEnvironment()) { |
| 141 | set_use_global_threadpool(false); |
| 142 | } |
| 143 | |
| 144 | if (use_global_threadpool_) { |
| 145 | mutex_lock l(global_tp_mu_); |
| 146 | if (options.config.experimental().use_numa_affinity()) { |
| 147 | int numa_node = attributes.locality().numa_node(); |
| 148 | int num_numa_nodes = port::NUMANumNodes(); |
| 149 | DCHECK_LT(numa_node, num_numa_nodes); |
| 150 | Allocator* numa_allocator = |
| 151 | ProcessState::singleton()->GetCPUAllocator(numa_node); |
| 152 | while (numa_node >= global_tp_info_.size()) { |
| 153 | global_tp_info_.push_back(nullptr); |
| 154 | } |
| 155 | if (!global_tp_info_[numa_node]) { |
| 156 | global_tp_info_[numa_node] = new LocalDevice::EigenThreadPoolInfo( |
| 157 | options, numa_node, numa_allocator); |
| 158 | } |
| 159 | tp_info = global_tp_info_[numa_node]; |
| 160 | } else { |
| 161 | if (opt.global_threadpool_num > 1) { |
| 162 | while (opt.global_threadpool_num > global_tp_info_.size()) { |
| 163 | global_tp_info_.push_back(nullptr); |
| 164 | } |
| 165 | if (!global_tp_info_[opt.device_threadpool_index]) { |
| 166 | global_tp_info_[opt.device_threadpool_index] = |
| 167 | new LocalDevice::EigenThreadPoolInfo( |
| 168 | options, port::kNUMANoAffinity, nullptr); |
| 169 | // set threadpool affinity |
| 170 | if (opt.cpuset.size() > 0) { |
| 171 | std::string msg; |
| 172 | cpu_set_t cpuset; |
| 173 | CPU_ZERO(&cpuset); |
| 174 | for (auto c : opt.cpuset) { |
| 175 | CPU_SET(c, &cpuset); |
| 176 | msg = msg + std::to_string(c) + ", "; |
| 177 | } |
| 178 | global_tp_info_[opt.device_threadpool_index]->SetThreadPoolAffinity(cpuset); |
| 179 | LOG(INFO) << "Intra thread pool #" << opt.device_threadpool_index |
| 180 | << " will be pinned to cpus: " << msg; |
| 181 | } |
| 182 | } |
| 183 | tp_info = global_tp_info_[opt.device_threadpool_index]; |
| 184 | } else { |
| 185 | if (global_tp_info_.empty()) { |
| 186 | global_tp_info_.push_back(new LocalDevice::EigenThreadPoolInfo( |
| 187 | options, port::kNUMANoAffinity, nullptr)); |
| 188 | } |
| 189 | tp_info = global_tp_info_[0]; |
no test coverage detected