Special-case MKL threading setters/getters

This uses the `MKL_Domain_*` functions to get/set the thread count for the BLAS domain only, so FFT, PARDISO, etc... domains in MKL are not affected bt `lbt_set_num_threads()`. It also adds a test to show that this behavior is reasonable when MKL is loaded.
JuliaLinearAlgebra · Jul 31, 2024 · d95a4c1 · d95a4c1
1 parent 5a08240
commit d95a4c1
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 4 deletions.
diff --git a/src/threading.c b/src/threading.c
@@ -1,6 +1,11 @@
 #include "libblastrampoline_internal.h"
 
 #define MAX_THREADING_NAMES     32
+#define max(x, y)  ((x) > (y) ? (x) : (y))
+
+// We need to ask MKL to get/set threads for only the BLAS domain, we therefore pass in
+// this constant to the relevant threading functions to limit our thread setting.
+#define MKL_DOMAIN_BLAS 1 /* From mkl_types.h */
 
 /*
  * We provide a flexible thread getter/setter interface here; by calling `lbt_set_num_threads()`
@@ -10,15 +15,17 @@
  */
 static char * getter_names[MAX_THREADING_NAMES] = {
     "openblas_get_num_threads",
-    "MKL_Get_Max_Threads",
     "bli_thread_get_num_threads",
+    // We special-case MKL in the lookup loop below
+    //"MKL_Domain_Get_Max_Threads",
     NULL
 };
 
 static char * setter_names[MAX_THREADING_NAMES] = {
     "openblas_set_num_threads",
-    "MKL_Set_Num_Threads",
     "bli_thread_set_num_threads",
+    // We special-case MKL in the lookup loop below
+    //"MKL_Domain_Set_Num_Threads",
     NULL
 };
 
@@ -62,9 +69,16 @@ LBT_DLLEXPORT int32_t lbt_get_num_threads() {
             int (*fptr)() = lookup_symbol(lib->handle, symbol_name);
             if (fptr != NULL) {
                 int new_threads = fptr();
-                max_threads = max_threads > new_threads ? max_threads : new_threads;
+                max_threads = max(max_threads, new_threads);
             }
         }
+
+        // Special-case MKL, as we need to specifically ask for the "BLAS" domain
+        int (*fptr)(int) = lookup_symbol(lib->handle, "MKL_Domain_Get_Max_Threads");
+        if (fptr != NULL) {
+            int new_threads = fptr(MKL_DOMAIN_BLAS);
+            max_threads = max(max_threads, new_threads);
+        }
     }
     return max_threads;
 }
@@ -76,15 +90,21 @@ LBT_DLLEXPORT int32_t lbt_get_num_threads() {
  */
 LBT_DLLEXPORT void lbt_set_num_threads(int32_t nthreads) {
     const lbt_config_t * config = lbt_get_config();
+    char symbol_name[MAX_SYMBOL_LEN];
     for (int lib_idx=0; config->loaded_libs[lib_idx] != NULL; ++lib_idx) {
         lbt_library_info_t * lib = config->loaded_libs[lib_idx];
         for (int symbol_idx=0; setter_names[symbol_idx] != NULL; ++symbol_idx) {
-            char symbol_name[MAX_SYMBOL_LEN];
             build_symbol_name(symbol_name, setter_names[symbol_idx], lib->suffix);
             void (*fptr)(int) = lookup_symbol(lib->handle, symbol_name);
             if (fptr != NULL) {
                 fptr(nthreads);
             }
         }
+
+        // Special-case MKL, as we need to specifically ask for the "BLAS" domain
+        int (*fptr)(int, int) = lookup_symbol(lib->handle, "MKL_Domain_Set_Num_Threads");
+        if (fptr != NULL) {
+            fptr(nthreads, MKL_DOMAIN_BLAS);
+        }
     }
 }
diff --git a/test/direct.jl b/test/direct.jl
@@ -331,4 +331,16 @@ if MKL_jll.is_available() && Sys.ARCH == :x86_64
         @test result[1] ≈ ComplexF32(1.47 + 3.83im)
         @test isempty(stacktraces)
     end
+
+    @testset "MKL threading domains" begin
+        nthreads = lbt_get_num_threads(lbt_handle)
+        if nthreads <= 1
+            nthreads = 2
+        else
+            nthreads = div(nthreads, 2)
+        end
+        lbt_set_num_threads(lbt_handle, nthreads)
+        @test ccall((:MKL_Domain_Get_Max_Threads, libmkl_rt), Cint, (Cint,), 1) == nthreads
+        @test ccall((:MKL_Domain_Get_Max_Threads, libmkl_rt), Cint, (Cint,), 2) != nthreads
+    end
 end