diff --git a/src/common/common.h b/src/common/common.h index 665870e1a8aa..b7c58e22bf84 100644 --- a/src/common/common.h +++ b/src/common/common.h @@ -155,24 +155,31 @@ class GPUSet { static GPUSet Empty() { return GPUSet(); } - static GPUSet Range(GpuIndex start, GpuIndex n_gpu) { - return n_gpu <=0 ? Empty() : GPUSet{start, n_gpu}; + static GPUSet Range(GpuIndex start, GpuIndex n_gpus) { + return n_gpus <=0 ? Empty() : GPUSet{start, n_gpus}; } /*! \brief ndevices and num_rows both are upper bounds. */ - static GPUSet All(GpuIndex gpu_id, GpuIndex n_gpu, + static GPUSet All(GpuIndex gpu_id, GpuIndex n_gpus, GpuIndex num_rows = std::numeric_limits::max()) { + CHECK_GE(gpu_id, 0) << "gpu_id must be >= 0."; + CHECK_GE(n_gpus, -1) << "n_gpus must be >= -1."; GpuIndex n_devices_visible = AllVisible().Size(); - if (n_gpu == kAll) { // Use all devices starting from `gpu_id'. + if (n_gpus == kAll) { // Use all devices starting from `gpu_id'. CHECK(gpu_id < n_devices_visible || gpu_id == 0) << "\ngpu_id should be less than available devices.\ngpu_id: " << gpu_id << ", number of available devices: " - << n_devices_visible << std::endl; + << n_devices_visible; + GpuIndex n_devices = n_devices_visible - gpu_id; + n_devices = n_devices < num_rows ? n_devices : num_rows; return Range(gpu_id, n_devices_visible - gpu_id); - } else { + } else { // Use devices in ( gpu_id, gpu_id + n_gpus ). GpuIndex n_available_devices = n_devices_visible - gpu_id; - GpuIndex n_devices = - n_available_devices < n_gpu ? n_devices_visible : n_gpu; + CHECK_LE(n_gpus, n_available_devices) + << "Starting from gpu id: " << gpu_id << ", there are only " + << n_available_devices << " available devices, while n_gpus is set to: " + << n_gpus; + GpuIndex n_devices = n_gpus < num_rows ? n_gpus : num_rows; return Range(gpu_id, n_devices); } } diff --git a/src/predictor/gpu_predictor.cu b/src/predictor/gpu_predictor.cu index cf0c7e9ddcb7..d3930f7d316b 100644 --- a/src/predictor/gpu_predictor.cu +++ b/src/predictor/gpu_predictor.cu @@ -1,5 +1,5 @@ /*! - * Copyright by Contributors 2017 + * Copyright 2017-2018 by Contributors */ #include #include