* @brief Construct a general solver * @param device_ids list of GPU ids, {} for auto * @param _num_sampler_per_worker number of sampler thread per GPU * @param _gpu_memory_limit memory limit for each GPU */
| 182 | * @param _gpu_memory_limit memory limit for each GPU |
| 183 | */ |
| 184 | SolverMixin(std::vector<int> device_ids = {}, int num_sampler_per_worker = kAuto, size_t _gpu_memory_limit = kAuto) : |
| 185 | gpu_memory_limit(_gpu_memory_limit), edge_table(-1), batch_id(0) { |
| 186 | if (device_ids.empty()) { |
| 187 | CUDA_CHECK(cudaGetDeviceCount(&num_worker)); |
| 188 | CHECK(num_worker) << "No GPU devices found"; |
| 189 | for (int i = 0; i < num_worker; i++) |
| 190 | device_ids.push_back(i); |
| 191 | } else |
| 192 | num_worker = device_ids.size(); |
| 193 | if (num_sampler_per_worker == kAuto) |
| 194 | num_sampler_per_worker = std::thread::hardware_concurrency() / num_worker - 1; |
| 195 | num_sampler = num_sampler_per_worker * num_worker; |
| 196 | num_thread = num_worker + num_sampler; |
| 197 | LOG_IF(WARNING, num_sampler > std::thread::hardware_concurrency()) |
| 198 | << "#CPU threads is beyond the hardware concurrency"; |
| 199 | if (gpu_memory_limit == kAuto) { |
| 200 | gpu_memory_limit = GiB(32); |
| 201 | size_t free_memory; |
| 202 | for (auto &&device_id : device_ids) { |
| 203 | CUDA_CHECK(cudaSetDevice(device_id)); |
| 204 | CUDA_CHECK(cudaMemGetInfo(&free_memory, nullptr)); |
| 205 | gpu_memory_limit = std::min(gpu_memory_limit, free_memory); |
| 206 | } |
| 207 | } |
| 208 | for (int i = 0; i < num_sampler_per_worker; i++) |
| 209 | for (auto &&device_id : device_ids) |
| 210 | samplers.push_back(new Sampler(this, device_id)); |
| 211 | for (auto &&device_id : device_ids) |
| 212 | workers.push_back(new Worker(this, device_id)); |
| 213 | } |
| 214 | |
| 215 | SolverMixin(const SolverMixin &) = delete; |
| 216 |