Skip to content

Commit 6009ae9

Browse files
authored
[LinearAlgebra] Support more env variables to set OpenBLAS threads (#46118)
1 parent d75843d commit 6009ae9

File tree

1 file changed

+22
-6
lines changed

1 file changed

+22
-6
lines changed

stdlib/LinearAlgebra/src/LinearAlgebra.jl

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -555,14 +555,29 @@ function versioninfo(io::IO=stdout)
555555
"JULIA_NUM_THREADS",
556556
"MKL_DYNAMIC",
557557
"MKL_NUM_THREADS",
558-
"OPENBLAS_NUM_THREADS",
558+
# OpenBLAS has a hierarchy of environment variables for setting the
559+
# number of threads, see
560+
# https://github.com/xianyi/OpenBLAS/blob/c43ec53bdd00d9423fc609d7b7ecb35e7bf41b85/README.md#setting-the-number-of-threads-using-environment-variables
561+
("OPENBLAS_NUM_THREADS", "GOTO_NUM_THREADS", "OMP_NUM_THREADS"),
559562
]
560563
printed_at_least_one_env_var = false
564+
print_var(io, indent, name) = println(io, indent, name, " = ", ENV[name])
561565
for name in env_var_names
562-
if haskey(ENV, name)
563-
value = ENV[name]
564-
println(io, indent, name, " = ", value)
565-
printed_at_least_one_env_var = true
566+
if name isa Tuple
567+
# If `name` is a Tuple, then find the first environment which is
568+
# defined, and disregard the following ones.
569+
for nm in name
570+
if haskey(ENV, nm)
571+
print_var(io, indent, nm)
572+
printed_at_least_one_env_var = true
573+
break
574+
end
575+
end
576+
else
577+
if haskey(ENV, name)
578+
print_var(io, indent, name)
579+
printed_at_least_one_env_var = true
580+
end
566581
end
567582
end
568583
if !printed_at_least_one_env_var
@@ -581,7 +596,8 @@ function __init__()
581596
# register a hook to disable BLAS threading
582597
Base.at_disable_library_threading(() -> BLAS.set_num_threads(1))
583598

584-
if !haskey(ENV, "OPENBLAS_NUM_THREADS")
599+
# https://github.com/xianyi/OpenBLAS/blob/c43ec53bdd00d9423fc609d7b7ecb35e7bf41b85/README.md#setting-the-number-of-threads-using-environment-variables
600+
if !haskey(ENV, "OPENBLAS_NUM_THREADS") && !haskey(ENV, "GOTO_NUM_THREADS") && !haskey(ENV, "OMP_NUM_THREADS")
585601
@static if Sys.isapple() && Base.BinaryPlatforms.arch(Base.BinaryPlatforms.HostPlatform()) == "aarch64"
586602
BLAS.set_num_threads(max(1, Sys.CPU_THREADS))
587603
else

0 commit comments

Comments
 (0)