add BLAS.get_num_threads

stdlib/LinearAlgebra/src/blas.jl

@@ -113,23 +113,46 @@ Set the number of threads the BLAS library should use.
 """
 function set_num_threads(n::Integer)
     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, (Int32,), 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."
     end
 
     return nothing
 end
 
+"""
+    get_num_threads()
+
+Get the number of threads the BLAS library is using.
+"""
+function get_num_threads()
+    blas = LinearAlgebra.BLAS.vendor()
+    if blas === :openblas || blas === :openblas64
+        return ccall((@blasfunc(openblas_get_num_threads), libblas), Cint, ())
+    elseif blas == :mkl
+        return ccall((:mkl_get_max_threads, libblas), Cint, ())
+    elseif Sys.isapple()
+        key = "VECLIB_MAXIMUM_THREADS"
+        s = get(ENV, key, "")
+        nt = Base.tryparse(Cint, s)
+        if nt === nothing
+            @warn "Failed to read environment variable $key"
+        else
+            return nt
+        end
+    end
+    @warn "Could not get number of BLAS threads. Returning Sys.CPU_THREADS instead."
+    return Sys.CPU_THREADS
+end
+
 const _testmat = [1.0 0.0; 0.0 -1.0]
 function check()
     blas = vendor()

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