diff --git a/include/LightGBM/utils/openmp_wrapper.h b/include/LightGBM/utils/openmp_wrapper.h index b48113e10ba4..a337fc353b75 100644 --- a/include/LightGBM/utils/openmp_wrapper.h +++ b/include/LightGBM/utils/openmp_wrapper.h @@ -25,6 +25,15 @@ inline int OMP_NUM_THREADS() { return ret; } +inline void OMP_SET_NUM_THREADS(int num_threads) { + static const int default_omp_num_threads = OMP_NUM_THREADS(); + if (num_threads > 0) { + omp_set_num_threads(num_threads); + } else { + omp_set_num_threads(default_omp_num_threads); + } +} + class ThreadExceptionHelper { public: ThreadExceptionHelper() { @@ -94,6 +103,7 @@ class ThreadExceptionHelper { simulate a single thread running. All #pragma omp should be ignored by the compiler **/ inline void omp_set_num_threads(int) __GOMP_NOTHROW {} // NOLINT (no cast done here) + inline void OMP_SET_NUM_THREADS(int) __GOMP_NOTHROW {} inline int omp_get_num_threads() __GOMP_NOTHROW {return 1;} inline int omp_get_max_threads() __GOMP_NOTHROW {return 1;} inline int omp_get_thread_num() __GOMP_NOTHROW {return 0;} diff --git a/src/application/application.cpp b/src/application/application.cpp index d9a4d7544ebc..b7f55a2ec0e4 100644 --- a/src/application/application.cpp +++ b/src/application/application.cpp @@ -31,9 +31,7 @@ namespace LightGBM { Application::Application(int argc, char** argv) { LoadParameters(argc, argv); // set number of threads for openmp - if (config_.num_threads > 0) { - omp_set_num_threads(config_.num_threads); - } + OMP_SET_NUM_THREADS(config_.num_threads); if (config_.data.size() == 0 && config_.task != TaskType::kConvertModel) { Log::Fatal("No training/prediction data, application quit"); } diff --git a/src/c_api.cpp b/src/c_api.cpp index bc3bfc3b2434..ec3880b60675 100644 --- a/src/c_api.cpp +++ b/src/c_api.cpp @@ -113,9 +113,7 @@ class Booster { const char* parameters) { auto param = Config::Str2Map(parameters); config_.Set(param); - if (config_.num_threads > 0) { - omp_set_num_threads(config_.num_threads); - } + OMP_SET_NUM_THREADS(config_.num_threads); // create boosting if (config_.input_model.size() > 0) { Log::Warning("Continued train from model is not supported for c_api,\n" @@ -314,9 +312,7 @@ class Booster { config_.Set(param); - if (config_.num_threads > 0) { - omp_set_num_threads(config_.num_threads); - } + OMP_SET_NUM_THREADS(config_.num_threads); if (param.count("objective")) { // create objective function @@ -951,9 +947,7 @@ int LGBM_DatasetCreateFromFile(const char* filename, auto param = Config::Str2Map(parameters); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); DatasetLoader loader(config, nullptr, 1, filename); if (reference == nullptr) { if (Network::num_machines() == 1) { @@ -981,9 +975,7 @@ int LGBM_DatasetCreateFromSampledColumn(double** sample_data, auto param = Config::Str2Map(parameters); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); DatasetLoader loader(config, nullptr, 1, nullptr); *out = loader.ConstructFromSampleData(sample_data, sample_indices, ncol, num_per_col, num_sample_row, @@ -1096,9 +1088,7 @@ int LGBM_DatasetCreateFromMats(int32_t nmat, auto param = Config::Str2Map(parameters); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); std::unique_ptr ret; int32_t total_nrow = 0; for (int j = 0; j < nmat; ++j) { @@ -1188,9 +1178,7 @@ int LGBM_DatasetCreateFromCSR(const void* indptr, auto param = Config::Str2Map(parameters); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); std::unique_ptr ret; auto get_row_fun = RowFunctionFromCSR(indptr, indptr_type, indices, data, data_type, nindptr, nelem); int32_t nrow = static_cast(nindptr - 1); @@ -1256,9 +1244,7 @@ int LGBM_DatasetCreateFromCSRFunc(void* get_row_funptr, auto param = Config::Str2Map(parameters); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); std::unique_ptr ret; int32_t nrow = num_rows; if (reference == nullptr) { @@ -1328,9 +1314,7 @@ int LGBM_DatasetCreateFromCSC(const void* col_ptr, auto param = Config::Str2Map(parameters); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); std::unique_ptr ret; int32_t nrow = static_cast(num_row); if (reference == nullptr) { @@ -1409,9 +1393,7 @@ int LGBM_DatasetGetSubset( auto param = Config::Str2Map(parameters); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); auto full_dataset = reinterpret_cast(handle); CHECK_GT(num_used_row_indices, 0); const int32_t lower = 0; @@ -1816,9 +1798,7 @@ int LGBM_BoosterPredictForFile(BoosterHandle handle, auto param = Config::Str2Map(parameter); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); Booster* ref_booster = reinterpret_cast(handle); ref_booster->Predict(start_iteration, num_iteration, predict_type, data_filename, data_has_header, config, result_filename); @@ -1894,9 +1874,7 @@ int LGBM_BoosterPredictForCSR(BoosterHandle handle, auto param = Config::Str2Map(parameter); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); Booster* ref_booster = reinterpret_cast(handle); auto get_row_fun = RowFunctionFromCSR(indptr, indptr_type, indices, data, data_type, nindptr, nelem); int nrow = static_cast(nindptr - 1); @@ -1928,9 +1906,7 @@ int LGBM_BoosterPredictSparseOutput(BoosterHandle handle, auto param = Config::Str2Map(parameter); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); if (matrix_type == C_API_MATRIX_TYPE_CSR) { if (num_col_or_row <= 0) { Log::Fatal("The number of columns should be greater than zero."); @@ -2015,9 +1991,7 @@ int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle, auto param = Config::Str2Map(parameter); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); Booster* ref_booster = reinterpret_cast(handle); auto get_row_fun = RowFunctionFromCSR(indptr, indptr_type, indices, data, data_type, nindptr, nelem); ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config); @@ -2047,9 +2021,7 @@ int LGBM_BoosterPredictForCSRSingleRowFastInit(BoosterHandle handle, data_type, static_cast(num_col))); - if (fastConfig_ptr->config.num_threads > 0) { - omp_set_num_threads(fastConfig_ptr->config.num_threads); - } + OMP_SET_NUM_THREADS(fastConfig_ptr->config.num_threads); fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config); @@ -2095,9 +2067,7 @@ int LGBM_BoosterPredictForCSC(BoosterHandle handle, auto param = Config::Str2Map(parameter); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); int num_threads = OMP_NUM_THREADS(); int ncol = static_cast(ncol_ptr - 1); std::vector> iterators(num_threads, std::vector()); @@ -2140,9 +2110,7 @@ int LGBM_BoosterPredictForMat(BoosterHandle handle, auto param = Config::Str2Map(parameter); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); Booster* ref_booster = reinterpret_cast(handle); auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major); ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, @@ -2165,9 +2133,7 @@ int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle, auto param = Config::Str2Map(parameter); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); Booster* ref_booster = reinterpret_cast(handle); auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, ncol, data_type, is_row_major); ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config); @@ -2191,9 +2157,7 @@ int LGBM_BoosterPredictForMatSingleRowFastInit(BoosterHandle handle, data_type, ncol)); - if (fastConfig_ptr->config.num_threads > 0) { - omp_set_num_threads(fastConfig_ptr->config.num_threads); - } + OMP_SET_NUM_THREADS(fastConfig_ptr->config.num_threads); fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config); @@ -2231,9 +2195,7 @@ int LGBM_BoosterPredictForMats(BoosterHandle handle, auto param = Config::Str2Map(parameter); Config config; config.Set(param); - if (config.num_threads > 0) { - omp_set_num_threads(config.num_threads); - } + OMP_SET_NUM_THREADS(config.num_threads); Booster* ref_booster = reinterpret_cast(handle); auto get_row_fun = RowPairFunctionFromDenseRows(data, ncol, data_type); ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, out_result, out_len);