Skip to content

Commit 9d49b8b

Browse files
giordanokshyatt
authored andcommitted
[LinearAlgebra] Initialise number of BLAS threads with jl_effective_threads (#55574)
This is a safer estimate than `Sys.CPU_THREADS` to avoid oversubscribing the machine when running distributed applications, or when the Julia process is constrained by external controls (`taskset`, `cgroups`, etc.). Fix #55572
1 parent ef39d40 commit 9d49b8b

File tree

3 files changed

+19
-2
lines changed

3 files changed

+19
-2
lines changed

NEWS.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,8 @@ Standard library changes
128128
between different eigendecomposition algorithms ([#49355]).
129129
* Added a generic version of the (unblocked) pivoted Cholesky decomposition
130130
(callable via `cholesky[!](A, RowMaximum())`) ([#54619]).
131+
* The number of default BLAS threads now respects process affinity, instead of
132+
using total number of logical threads available on the system ([#55574]).
131133

132134
#### Logging
133135

stdlib/LinearAlgebra/src/LinearAlgebra.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -843,9 +843,9 @@ function __init__()
843843
# https://github.com/xianyi/OpenBLAS/blob/c43ec53bdd00d9423fc609d7b7ecb35e7bf41b85/README.md#setting-the-number-of-threads-using-environment-variables
844844
if !haskey(ENV, "OPENBLAS_NUM_THREADS") && !haskey(ENV, "GOTO_NUM_THREADS") && !haskey(ENV, "OMP_NUM_THREADS")
845845
@static if Sys.isapple() && Base.BinaryPlatforms.arch(Base.BinaryPlatforms.HostPlatform()) == "aarch64"
846-
BLAS.set_num_threads(max(1, Sys.CPU_THREADS))
846+
BLAS.set_num_threads(max(1, @ccall(jl_effective_threads()::Cint)))
847847
else
848-
BLAS.set_num_threads(max(1, Sys.CPU_THREADS ÷ 2))
848+
BLAS.set_num_threads(max(1, @ccall(jl_effective_threads()::Cint) ÷ 2))
849849
end
850850
end
851851
end

test/threads.jl

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -359,3 +359,18 @@ end
359359
@test jl_setaffinity(0, mask, cpumasksize) == 0
360360
end
361361
end
362+
363+
# Make sure default number of BLAS threads respects CPU affinity: issue #55572.
364+
@testset "LinearAlgebra number of default threads" begin
365+
if AFFINITY_SUPPORTED
366+
allowed_cpus = findall(uv_thread_getaffinity())
367+
cmd = addenv(`$(Base.julia_cmd()) --startup-file=no -E 'using LinearAlgebra; BLAS.get_num_threads()'`,
368+
# Remove all variables which could affect the default number of threads
369+
"OPENBLAS_NUM_THREADS"=>nothing,
370+
"GOTO_NUM_THREADS"=>nothing,
371+
"OMP_NUM_THREADS"=>nothing)
372+
for n in 1:min(length(allowed_cpus), 8) # Cap to 8 to avoid too many tests on large systems
373+
@test readchomp(setcpuaffinity(cmd, allowed_cpus[1:n])) == string(max(1, n ÷ 2))
374+
end
375+
end
376+
end

0 commit comments

Comments
 (0)