Unify thread configuration. (#6186)

This commit is contained in:
Jiaming Yuan
2020-10-19 16:05:42 +08:00
committed by GitHub
parent 7f6ed5780c
commit ddf37cca30
5 changed files with 39 additions and 12 deletions

View File

@@ -19,6 +19,7 @@
#include "../common/math.h"
#include "../common/version.h"
#include "../common/group_data.h"
#include "../common/threading_utils.h"
#include "../data/adapter.h"
#include "../data/iterative_device_dmatrix.h"
@@ -843,10 +844,7 @@ void SparsePage::Push(const SparsePage &batch) {
template <typename AdapterBatchT>
uint64_t SparsePage::Push(const AdapterBatchT& batch, float missing, int nthread) {
// Set number of threads but keep old value so we can reset it after
const int nthreadmax = omp_get_max_threads();
if (nthread <= 0) nthread = nthreadmax;
const int nthread_original = omp_get_max_threads();
omp_set_num_threads(nthread);
int nthread_original = common::OmpSetNumThreads(&nthread);
auto& offset_vec = offset.HostVector();
auto& data_vec = data.HostVector();
@@ -865,7 +863,7 @@ uint64_t SparsePage::Push(const AdapterBatchT& batch, float missing, int nthread
}
}
size_t batch_size = batch.Size();
const size_t thread_size = batch_size/nthread;
const size_t thread_size = batch_size / nthread;
builder.InitBudget(expected_rows+1, nthread);
uint64_t max_columns = 0;
if (batch_size == 0) {