diff --git a/src/objective/rank_objective.hpp b/src/objective/rank_objective.hpp index 6bd5324812f8..b9c45ae62410 100644 --- a/src/objective/rank_objective.hpp +++ b/src/objective/rank_objective.hpp @@ -289,17 +289,12 @@ class LambdarankNDCG : public RankingObjective { void UpdatePositionBiasFactors(const score_t* lambdas, const score_t* hessians) const override { /// get number of threads - int num_threads = 1; - #pragma omp parallel - #pragma omp master - { - num_threads = omp_get_num_threads(); - } + int num_threads = OMP_NUM_THREADS(); // create per-thread buffers for first and second derivatives of utility w.r.t. position bias factors std::vector bias_first_derivatives(num_position_ids_ * num_threads, 0.0); std::vector bias_second_derivatives(num_position_ids_ * num_threads, 0.0); std::vector instance_counts(num_position_ids_ * num_threads, 0); - #pragma omp parallel for schedule(guided) + #pragma omp parallel for schedule(guided) num_threads(num_threads) for (data_size_t i = 0; i < num_data_; i++) { // get thread ID const int tid = omp_get_thread_num(); @@ -310,7 +305,7 @@ class LambdarankNDCG : public RankingObjective { bias_second_derivatives[offset] -= hessians[i]; instance_counts[offset]++; } - #pragma omp parallel for schedule(guided) + #pragma omp parallel for schedule(guided) num_threads(num_threads) for (data_size_t i = 0; i < num_position_ids_; i++) { double bias_first_derivative = 0.0; double bias_second_derivative = 0.0; diff --git a/src/treelearner/gpu_tree_learner.cpp b/src/treelearner/gpu_tree_learner.cpp index 294be28b6f86..4cd2c0b2731e 100644 --- a/src/treelearner/gpu_tree_learner.cpp +++ b/src/treelearner/gpu_tree_learner.cpp @@ -344,7 +344,7 @@ void GPUTreeLearner::AllocateGPUMemory() { // for data transfer time auto start_time = std::chrono::steady_clock::now(); // Now generate new data structure feature4, and copy data to the device - int nthreads = std::min(omp_get_max_threads(), static_cast(dense_feature_group_map_.size()) / dword_features_); + int nthreads = std::min(OMP_NUM_THREADS(), static_cast(dense_feature_group_map_.size()) / dword_features_); nthreads = std::max(nthreads, 1); std::vector host4_vecs(nthreads); std::vector host4_bufs(nthreads); diff --git a/src/treelearner/linear_tree_learner.cpp b/src/treelearner/linear_tree_learner.cpp index 071a6405709c..0850abe8dcf1 100644 --- a/src/treelearner/linear_tree_learner.cpp +++ b/src/treelearner/linear_tree_learner.cpp @@ -52,7 +52,7 @@ void LinearTreeLearner::InitLinear(const Dataset* train_data, const int max_leav } XTHX_by_thread_.clear(); XTg_by_thread_.clear(); - int max_threads = omp_get_max_threads(); + int max_threads = OMP_NUM_THREADS(); for (int i = 0; i < max_threads; ++i) { XTHX_by_thread_.push_back(XTHX_); XTg_by_thread_.push_back(XTg_);