add BLAS.get_num_threads

stdlib/LinearAlgebra/src/blas.jl

@@ -107,26 +107,65 @@ end
 openblas_get_config() = strip(unsafe_string(ccall((@blasfunc(openblas_get_config), libblas), Ptr{UInt8}, () )))
 
 """
-    set_num_threads(n)
+    set_num_threads(n::Integer)
+    set_num_threads(::Nothing)
+
+Set the number of threads the BLAS library should use equal to `n::Integer`.
+If `nothing` is passed to this function, julia tries to figure out the optimial number of threads.
+The exact heuristic is an implementation detail.
 
-Set the number of threads the BLAS library should use.
+On exotic variants of `BLAS` this function can fail.
 """
-function set_num_threads(n::Integer)
+function set_num_threads(n::Integer)::Nothing
     blas = vendor()
-    if blas === :openblas
-        return ccall((:openblas_set_num_threads, libblas), Cvoid, (Int32,), n)
-    elseif blas === :openblas64
-        return ccall((:openblas_set_num_threads64_, libblas), Cvoid, (Int32,), n)
+    if blas === :openblas || blas == :openblas64
+        return ccall((@blasfunc(openblas_set_num_threads), libblas), Cvoid, (Cint,), n)
     elseif blas === :mkl
         # MKL may let us set the number of threads in several ways
         return ccall((:MKL_Set_Num_Threads, libblas), Cvoid, (Cint,), n)
-    end
-
-    # OSX BLAS looks at an environment variable
-    @static if Sys.isapple()
+    elseif Sys.isapple()
+        # OSX BLAS looks at an environment variable
         ENV["VECLIB_MAXIMUM_THREADS"] = n
+    else
+        @warn "Failed to set number of BLAS threads." maxlog=1
     end
+    return nothing
+end
+
+_tryparse_env_int(key) = tryparse(Int, get(ENV, key, ""))
+
+function set_num_threads(::Nothing)
+    n = something(
+        _tryparse_env_int("OPENBLAS_NUM_THREADS"),
+        _tryparse_env_int("OMP_NUM_THREADS"),
+        max(1, Sys.CPU_THREADS ÷ 2),
+    )
+    set_num_threads(n)
+end
+
+"""
+    get_num_threads()
 
+Get the number of threads the BLAS library is using.
+
+On exotic variants of `BLAS` this function can fail, which is indicated by returning `nothing`.
+"""
+function get_num_threads()::Union{Int, Nothing}
+    blas = LinearAlgebra.BLAS.vendor()
+    if blas === :openblas || blas === :openblas64
+        return Int(ccall((@blasfunc(openblas_get_num_threads), libblas), Cint, ()))
+    elseif blas == :mkl
+        return Int(ccall((:mkl_get_max_threads, libblas), Cint, ()))
+    elseif Sys.isapple()
+        key = "VECLIB_MAXIMUM_THREADS"
+        nt = _tryparse_env_int(key)
+        if nt === nothing
+            @warn "Failed to read environment variable $key" maxlog=1
+        else
+            return nt
+        end
+    end
+    @warn "Could not get number of BLAS threads. Returning `nothing` instead." maxlog=1
     return nothing
 end
 

stdlib/LinearAlgebra/test/blas.jl

@@ -553,4 +553,14 @@ Base.stride(A::WrappedArray, i::Int) = stride(A.A, i)
     end
 end
 
+@testset "get_set_num_threads" begin
+    default = BLAS.get_num_threads()
+    @test default isa Integer
+    @test default > 0
+    BLAS.set_num_threads(1)
+    @test BLAS.get_num_threads() == 1
+    BLAS.set_num_threads(default)
+    @test BLAS.get_num_threads() == default
+end
+
 end # module TestBLAS