Fix specifying gpu_id, add tests. (#3851)

* Rewrite gpu_id related code.

* Remove normalised/unnormalised operatios.
* Address difference between `Index' and `Device ID'.
* Modify doc for `gpu_id'.
* Better LOG for GPUSet.
* Check specified n_gpus.
* Remove inappropriate `device_idx' term.
* Clarify GpuIdType and size_t.
This commit is contained in:
Jiaming Yuan
2018-11-06 18:17:53 +13:00
committed by GitHub
parent 1698fe64bb
commit f1275f52c1
20 changed files with 341 additions and 203 deletions

View File

@@ -352,8 +352,9 @@ struct GPUSketcher {
dh::ExecuteIndexShards(&shards_, [&](int i, std::unique_ptr<DeviceShard>& shard) {
size_t start = dist_.ShardStart(info.num_row_, i);
size_t size = dist_.ShardSize(info.num_row_, i);
shard = std::unique_ptr<DeviceShard>
(new DeviceShard(dist_.Devices()[i], start, start + size, param_));
shard = std::unique_ptr<DeviceShard>(
new DeviceShard(dist_.Devices().DeviceId(i),
start, start + size, param_));
});
// compute sketches for each shard
@@ -379,8 +380,7 @@ struct GPUSketcher {
}
GPUSketcher(tree::TrainParam param, size_t n_rows) : param_(std::move(param)) {
dist_ = GPUDistribution::Block(GPUSet::All(param_.n_gpus, n_rows).
Normalised(param_.gpu_id));
dist_ = GPUDistribution::Block(GPUSet::All(param_.gpu_id, param_.n_gpus, n_rows));
}
std::vector<std::unique_ptr<DeviceShard>> shards_;