From d2ce93179f6747380488db2a56102dab6fde18ca Mon Sep 17 00:00:00 2001 From: Elliot Saba Date: Thu, 22 Sep 2022 10:38:36 -0700 Subject: [PATCH] Add `OPENBLAS_DEFAULT_NUM_THREADS` This allows Julia to set a default number of threads (usually `1`) to be used when no other thread counts are specified [0], to short-circuit the default OpenBLAS thread initialization routine that spins up a different number of threads than Julia would otherwise choose. The reason to add a new environment variable is that we want to be able to configure OpenBLAS to avoid performing its initial memory allocation/thread startup, as that can consume significant amounts of memory, but we still want to be sensitive to legacy codebases that set things like `OMP_NUM_THREADS` or `GOTOBLAS_NUM_THREADS`. Creating a new environment variable that is openblas-specific and is not already publicly used to control the overall number of threads of programs like Julia seems to be the best way forward. [0] https://github.com/JuliaLang/julia/pull/46844 --- driver/others/init.c | 2 ++ driver/others/openblas_env.c | 8 +++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/driver/others/init.c b/driver/others/init.c index cc3145a62..cd10e8d36 100644 --- a/driver/others/init.c +++ b/driver/others/init.c @@ -823,6 +823,8 @@ void gotoblas_affinity_init(void) { if (numprocs == 0) numprocs = readenv_atoi("OMP_NUM_THREADS"); + if (numprocs == 0) numprocs = readenv_atoi("OPENBLAS_DEFAULT_NUM_THREADS"); + numnodes = 1; if (numprocs == 1) { diff --git a/driver/others/openblas_env.c b/driver/others/openblas_env.c index ef91a08e6..35b2270d4 100644 --- a/driver/others/openblas_env.c +++ b/driver/others/openblas_env.c @@ -67,10 +67,16 @@ void openblas_read_env() { openblas_env_thread_timeout=(unsigned int)ret; ret=0; - if (readenv(p,"OPENBLAS_NUM_THREADS")) ret = atoi(p); + if (readenv(p,"OPENBLAS_DEFAULT_NUM_THREADS")) ret = atoi(p); if(ret<0) ret=0; openblas_env_openblas_num_threads=ret; + ret=0; + if (readenv(p,"OPENBLAS_NUM_THREADS")) ret = atoi(p); + if(ret<0) ret=0; + if(ret != 0 || openblas_env_openblas_num_threads == 0) + openblas_env_openblas_num_threads=ret; + ret=0; if (readenv(p,"GOTO_NUM_THREADS")) ret = atoi(p); if(ret<0) ret=0;