Add BLAS.get_num_threads

Project.toml

@@ -1,6 +1,6 @@
 name = "Compat"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
-version = "3.12.0"
+version = "3.13.0"
 
 [deps]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

README.md

@@ -52,6 +52,8 @@ changes in `julia`.
 
 ## Supported features
 
+* `Compat.get_num_threads()` adds the functionality of `LinearAlgebra.BLAS.get_num_threads()`, and has matching `Compat.set_num_threads(n)` ([#36360]). (since Compat 3.13.0)
+
 * `@inferred [AllowedType] f(x)` is defined ([#27516]). (since Compat 3.12.0)
 
 * Search a character in a string with `findfirst`, `findnext`, `findlast` and `findprev` ([#31664]). (since Compat 3.12.0)
@@ -171,3 +173,4 @@ Note that you should specify the correct minimum version for `Compat` in the
 [#34251]: https://github.com/JuliaLang/julia/pull/34251
 [#35577]: https://github.com/JuliaLang/julia/pull/35577
 [#27516]: https://github.com/JuliaLang/julia/pull/27516
+[#36360]: https://github.com/JuliaLang/julia/pull/36360

src/Compat.jl

@@ -1,7 +1,7 @@
 module Compat
 
 import LinearAlgebra
-using LinearAlgebra: Adjoint, Diagonal, Transpose, UniformScaling, RealHermSymComplexHerm
+using LinearAlgebra: Adjoint, Diagonal, Transpose, UniformScaling, RealHermSymComplexHerm, BLAS
 
 include("compatmacro.jl")
 
@@ -464,6 +464,102 @@ if VERSION < v"1.2.0-DEV.77"
     #export @inferred
 end
 
+# https://github.com/JuliaLang/julia/pull/36360
+if VERSION < v"1.6.0-DEV.322" # b8110f8d1ec6349bee77efb5022621fdf50bd4a5
+
+    function guess_vendor()
+        # like determine_vendor, but guesses blas in some cases
+        # where determine_vendor returns :unknown
+        ret = BLAS.vendor()
+        if Base.Sys.isapple() && (ret == :unknown)
+            ret = :osxblas
+        end
+        ret
+    end
+
+    """
+        Compat.set_num_threads(n)
+
+    Set the number of threads the BLAS library should use.
+
+    Also accepts `nothing`, in which case julia tries to guess the default number of threads.
+    Passing `nothing` is discouraged and mainly exists because,
+    on exotic variants of BLAS, `nothing` may be returned by `get_num_threads()`.
+    Thus the following pattern may fail to set the number of threads, but will not error:
+    ```julia
+    old = get_num_threads()
+    set_num_threads(1)
+    @threads for i in 1:10
+        # single-threaded BLAS calls
+    end
+    set_num_threads(old)
+    ```
+    """
+    set_num_threads(n)::Nothing = _set_num_threads(n)
+
+    function _set_num_threads(n::Integer; _blas = guess_vendor())
+        if _blas === :openblas || _blas == :openblas64
+            return ccall((BLAS.@blasfunc(openblas_set_num_threads), BLAS.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, BLAS.libblas), Cvoid, (Cint,), n)
+            elseif _blas === :osxblas
+            # OSX BLAS looks at an environment variable
+            ENV["VECLIB_MAXIMUM_THREADS"] = n
+        else
+            @assert _blas === :unknown
+            @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; _blas = guess_vendor())
+        n = something(
+            _tryparse_env_int("OPENBLAS_NUM_THREADS"),
+            _tryparse_env_int("OMP_NUM_THREADS"),
+            max(1, Base.Sys.CPU_THREADS ÷ 2),
+        )
+        _set_num_threads(n; _blas = _blas)
+    end
+
+    """
+        Compat.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`.
+
+    In Julia 1.6 this is `LinearAlgebra.BLAS.get_num_threads()`
+    """
+    get_num_threads(;_blas=guess_vendor())::Union{Int, Nothing} = _get_num_threads()
+
+    function _get_num_threads(; _blas = guess_vendor())::Union{Int, Nothing}
+        if _blas === :openblas || _blas === :openblas64
+            return Int(ccall((BLAS.@blasfunc(openblas_get_num_threads), BLAS.libblas), Cint, ()))
+        elseif _blas === :mkl
+            return Int(ccall((:mkl_get_max_threads, BLAS.libblas), Cint, ()))
+        elseif _blas === :osxblas
+            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
+        else
+            @assert _blas === :unknown
+        end
+        @warn "Could not get number of BLAS threads. Returning `nothing` instead." maxlog=1
+        return nothing
+    end
+
+else
+    # Ensure that these can still be accessed as Compat.get_num_threads() etc:
+    import LinearAlgebra.BLAS: set_num_threads, get_num_threads
+end
+
 include("deprecated.jl")
 
 end # module Compat

test/runtests.jl

@@ -464,4 +464,39 @@ end
     @inferred Missing g(10)
 end
 
+# https://github.com/JuliaLang/julia/pull/36360
+@testset "get_set_num_threads" begin
+    default = Compat.get_num_threads()
+    @test default isa Int # seems dodgy, could be nothing!
+    @test default > 0
+    Compat.set_num_threads(1)
+    @test Compat.get_num_threads() === 1
+    Compat.set_num_threads(default)
+    @test Compat.get_num_threads() === default
+
+    # Run the ::Nothing method, to check no error:
+    Compat.set_num_threads(nothing)
+    Compat.set_num_threads(default)
+
+    if VERSION < v"1.6.0-DEV.322"
+        # These tests from PR rely on internal functions which would be BLAS. not Compat.
+        @test_logs (:warn,) match_mode=:any Compat._set_num_threads(1, _blas=:unknown)
+        if Compat.guess_vendor() !== :osxblas
+            # test osxblas which is not covered by CI
+            withenv("VECLIB_MAXIMUM_THREADS" => nothing) do
+                @test @test_logs(
+                    (:warn,),
+                    (:warn,),
+                    match_mode=:any,
+                    Compat._get_num_threads(_blas=:osxblas),
+                ) === nothing
+                @test_logs Compat._set_num_threads(1, _blas=:osxblas)
+                @test @test_logs(Compat._get_num_threads(_blas=:osxblas)) === 1
+                @test_logs Compat._set_num_threads(2, _blas=:osxblas)
+                @test @test_logs(Compat._get_num_threads(_blas=:osxblas)) === 2
+            end
+        end
+    end
+end
+
 nothing