| 35 | namespace paddle::platform { |
| 36 | |
| 37 | void SetNumThreads(int num_threads) { |
| 38 | #ifdef PADDLE_USE_OPENBLAS |
| 39 | // windows has no support for openblas multi-thread |
| 40 | // please refer to: https://github.com/PaddlePaddle/Paddle/issues/7234 |
| 41 | #ifdef _WIN32 |
| 42 | if (num_threads > 1) { |
| 43 | num_threads = 1; |
| 44 | } |
| 45 | #endif |
| 46 | int real_num_threads = num_threads > 1 ? num_threads : 1; |
| 47 | openblas_set_num_threads(real_num_threads); |
| 48 | #elif defined(PADDLE_WITH_MKLML) |
| 49 | int real_num_threads = num_threads > 1 ? num_threads : 1; |
| 50 | phi::dynload::MKL_Set_Num_Threads(real_num_threads); |
| 51 | omp_set_num_threads(real_num_threads); |
| 52 | #elif defined(PADDLE_WITH_HML) |
| 53 | int real_num_threads = num_threads > 1 ? num_threads : 1; |
| 54 | phi::dynload::hml_blas_set_num_threads(real_num_threads); |
| 55 | omp_set_num_threads(real_num_threads); |
| 56 | #elif defined(PADDLE_USE_REFERENCE_CBLAS) |
| 57 | // cblas not support multi-thread |
| 58 | return; |
| 59 | #elif defined(PADDLE_USE_ACCELERATE) |
| 60 | // not sure about apple's blas |
| 61 | return; |
| 62 | #else |
| 63 | PADDLE_THROW(common::errors::Unimplemented( |
| 64 | "This library (except OPENBLAS, MKLML) is not supported yet, so the " |
| 65 | "number of threads cannot be set.")); |
| 66 | #endif |
| 67 | } |
| 68 | |
| 69 | } // namespace paddle::platform |
no outgoing calls