diff --git a/CHANGELOG.md b/CHANGELOG.md index 8cc6e74..8f8883c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ ThreadPinning.jl Changelog ========================= +Version 1.1 +------------- +- ![Feature][badge-feature] Added `setaffinities` and `setaffinities_cpuids` as multi-thread analogs of `setaffinity` and `setaffinity_cpuids`, respectively. They allow setting the affinity of all Julia threads at once, similar to how `pinthreads` relates to `pinthread`. A single mask/cpuid vector can be passed to broadcast the same affinity to all threads (e.g. `setaffinities_cpuids(numa(1))`). + Version 1.0 ------------- - ![Feature][badge-feature] OpenBLAS: Almost all pinning and querying functions now have `openblas_*` analogs that provide (almost) all of the same features as for regular Julia threads. Example: `openblas_pinthreads(:cores)` now works. You can also visualize the placement of OpenBLAS threads via `threadinfo(; blas=true)`. These functions are now also part of the official API (and SemVer). diff --git a/Project.toml b/Project.toml index 47bcf51..4941106 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "ThreadPinning" uuid = "811555cd-349b-4f26-b7bc-1f208b848042" authors = ["Carsten Bauer and contributors"] -version = "1.0.2" +version = "1.1.0-DEV" [deps] DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab" diff --git a/docs/src/refs/api_pinning.md b/docs/src/refs/api_pinning.md index ba66716..92ffd23 100644 --- a/docs/src/refs/api_pinning.md +++ b/docs/src/refs/api_pinning.md @@ -10,6 +10,8 @@ unpinthreads unpinthread setaffinity setaffinity_cpuids +setaffinities +setaffinities_cpuids ``` ## Pinning - OpenBLAS diff --git a/src/ThreadPinning.jl b/src/ThreadPinning.jl index b4d17c7..ce3e225 100644 --- a/src/ThreadPinning.jl +++ b/src/ThreadPinning.jl @@ -24,6 +24,8 @@ else unpinthread(args...; kwargs...) = nothing setaffinity(args...; kwargs...) = nothing setaffinity_cpuids(args...; kwargs...) = nothing + setaffinities(args...; kwargs...) = nothing + setaffinities_cpuids(args...; kwargs...) = nothing with_pinthreads(f, args...; kwargs...) = f() pinthreads_likwidpin(args...; kwargs...) = nothing mpi_pinthreads(args...; kwargs...) = nothing @@ -52,7 +54,7 @@ export openblas_getaffinity, openblas_getcpuid, openblas_getcpuids, ## pinning export pinthread, pinthreads, with_pinthreads, unpinthread, unpinthreads -export setaffinity, setaffinity_cpuids +export setaffinity, setaffinity_cpuids, setaffinities, setaffinities_cpuids export pinthreads_likwidpin, likwidpin_domains, likwidpin_to_cpuids export mpi_pinthreads, mpi_getcpuids, mpi_gethostnames, mpi_getlocalrank export distributed_pinthreads, distributed_getcpuids, distributed_gethostnames, diff --git a/src/mkl.jl b/src/mkl.jl index 4aa76d2..4547a52 100644 --- a/src/mkl.jl +++ b/src/mkl.jl @@ -1,8 +1,10 @@ module MKL using LinearAlgebra: LinearAlgebra, BLAS +using Libdl: Libdl const MKL_PATH = Ref{Union{Nothing, String}}(nothing) +const MKL_HANDLE = Ref{Ptr{Cvoid}}(C_NULL) """ Returns the full path to the `libmkl_rt` library if the latter is loaded. Will try to @@ -11,13 +13,32 @@ locate the library and, if successfull, will cache the result. Throws an error o To force an update of the cache, provide `force_update=true`. """ function mkl_fullpath(; force_update = false) - if isnothing(MKL_PATH[]) || force_update + if isnothing(MKL_PATH[]) + # First-time lookup: just cache the path mklpath = _find_mkl() MKL_PATH[] = mklpath + elseif force_update + # Force a re-discovery of the path + mklpath = _find_mkl() + if MKL_PATH[] != mklpath + # Path changed: close any existing handle and invalidate cache + if MKL_HANDLE[] != C_NULL + Libdl.dlclose(MKL_HANDLE[]) + end + MKL_HANDLE[] = C_NULL + MKL_PATH[] = mklpath + end end return MKL_PATH[] end +function _mkl_handle() + if MKL_HANDLE[] == C_NULL + MKL_HANDLE[] = Libdl.dlopen(mkl_fullpath()) + end + return MKL_HANDLE[] +end + """ Check whether Intel MKL is currently loaded via libblastrampoline """ @@ -41,7 +62,7 @@ end mkl_get_dynamic() Wrapper around the MKL function [`mkl_get_dynamic`](https://www.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-fortran/top/support-functions/threading-control/mkl-get-dynamic.html). """ -mkl_get_dynamic() = @ccall mkl_fullpath().mkl_get_dynamic()::Cint +mkl_get_dynamic() = ccall(Libdl.dlsym(_mkl_handle(), :mkl_get_dynamic), Cint, ()) """ mkl_set_dynamic(flag::Integer) @@ -49,7 +70,7 @@ mkl_get_dynamic() = @ccall mkl_fullpath().mkl_get_dynamic()::Cint Wrapper around the MKL function [`mkl_set_dynamic`](https://www.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-c/top/support-functions/threading-control/mkl-set-dynamic.html). """ function mkl_set_dynamic(flag::Integer) - @ccall mkl_fullpath().MKL_Set_Dynamic(flag::Cint)::Cvoid + ccall(Libdl.dlsym(_mkl_handle(), :MKL_Set_Dynamic), Cvoid, (Cint,), flag) end end # module diff --git a/src/pinning.jl b/src/pinning.jl index ca118da..04db561 100644 --- a/src/pinning.jl +++ b/src/pinning.jl @@ -152,10 +152,10 @@ function unpinthreads end Set the affinity of a Julia thread to the given CPU-threads. *Examples:* -* `setaffinity(socket(1))` # set the affinity to the first socket -* `setaffinity(numa(2))` # set the affinity to the second NUMA domain -* `setaffinity(socket(1, 1:3))` # set the affinity to the first three cores in the first NUMA domain -* `setaffinity([1,3,5])` # set the affinity to the CPU-threads with the IDs 1, 3, and 5. +* `setaffinity_cpuids(socket(1))` # set the affinity to the first socket +* `setaffinity_cpuids(numa(2))` # set the affinity to the second NUMA domain +* `setaffinity_cpuids(socket(1, 1:3))` # set the affinity to the first three cores in the first NUMA domain +* `setaffinity_cpuids([1,3,5])` # set the affinity to the CPU-threads with the IDs 1, 3, and 5. """ function setaffinity_cpuids end @@ -166,6 +166,29 @@ Set the affinity of a Julia thread based on the given mask (a vector of ones and """ function setaffinity end +""" +Set the affinity of multiple Julia threads to the given CPU-threads. + +`cpuids_vec` should be a vector of CPU ID vectors (one per Julia thread), or a single +CPU ID vector to set all threads to the same affinity. + +*Examples:* +* `setaffinities_cpuids(numa(1))` # set all threads' affinity to the first NUMA domain +* `setaffinities_cpuids([socket(1), socket(2)])` # set first thread's affinity to socket 1, second to socket 2 +* `setaffinities_cpuids([numa(1), numa(2)])` # set affinities across NUMA domains +""" +function setaffinities_cpuids end + +""" + setaffinities(masks; threadpool = :default) + +Set the affinity of multiple Julia threads based on the given masks. + +`masks` should be a vector of affinity masks (one per Julia thread), or a single mask to +set all threads to the same affinity. Each mask is a vector of ones and zeros. +""" +function setaffinities end + # OpenBLAS """ openblas_setaffinity(mask; threadid) @@ -230,7 +253,7 @@ function openblas_unpinthreads end module Pinning import ThreadPinning: pinthread, pinthreads, with_pinthreads, unpinthread, unpinthreads -import ThreadPinning: setaffinity, setaffinity_cpuids +import ThreadPinning: setaffinity, setaffinity_cpuids, setaffinities, setaffinities_cpuids import ThreadPinning: openblas_setaffinity, openblas_setaffinity_cpuids, openblas_pinthread, openblas_pinthreads, openblas_unpinthread, openblas_unpinthreads @@ -282,6 +305,41 @@ function setaffinity_cpuids(cpuids::AbstractVector{<:Integer}; kwargs...) return end +function setaffinities(masks::AbstractVector{<:AbstractVector{<:Integer}}; + threadpool::Symbol = :default) + tids = ThreadPinningCore.threadids(; threadpool) + length(masks) == length(tids) || + throw(ArgumentError("Number of masks ($(length(masks))) must match the number of " * + "threads ($(length(tids))) in the `$(threadpool)` threadpool.")) + for (i, threadid) in pairs(tids) + setaffinity(copy(masks[i]); threadid) + end + return +end + +function setaffinities(mask::AbstractVector{<:Integer}; threadpool::Symbol = :default, kwargs...) + tids = ThreadPinningCore.threadids(; threadpool) + setaffinities([mask for _ in tids]; threadpool, kwargs...) +end + +function setaffinities_cpuids(cpuids_vec::AbstractVector{<:AbstractVector{<:Integer}}; + threadpool::Symbol = :default) + tids = ThreadPinningCore.threadids(; threadpool) + length(cpuids_vec) == length(tids) || + throw(ArgumentError("Number of CPU ID vectors ($(length(cpuids_vec))) must match " * + "the number of threads ($(length(tids))) in the " * + "`$(threadpool)` threadpool.")) + for (i, threadid) in pairs(tids) + setaffinity_cpuids(cpuids_vec[i]; threadid) + end + return +end + +function setaffinities_cpuids(cpuids::AbstractVector{<:Integer}; threadpool::Symbol = :default, kwargs...) + tids = ThreadPinningCore.threadids(; threadpool) + setaffinities_cpuids([cpuids for _ in tids]; threadpool, kwargs...) +end + function openblas_setaffinity_cpuids(cpuids::AbstractVector{<:Integer}; kwargs...) _check_cpuids(cpuids) mask = Utility.cpuids2affinitymask(cpuids) diff --git a/test/tests_pinning.jl b/test/tests_pinning.jl index 8c6a839..03bc8a2 100644 --- a/test/tests_pinning.jl +++ b/test/tests_pinning.jl @@ -181,6 +181,70 @@ function pinning_tests() @test isnothing(setaffinity_cpuids([cpuid2, cpuid1]; threadid = randtid)) end + @testset "setaffinities" begin + tids = ThreadPinning.threadids(; threadpool = :default) + nt = length(tids) + + @testset "setaffinities (vector of masks)" begin + cpuid1, cpuid2 = get_two_cpuids() + pinthreads([isodd(i) ? cpuid1 : cpuid2 for i in 1:nt]; warn = false) + masks = [getaffinity(; threadid = tid) for tid in tids] + expected_cpuids = getcpuids() + pinthreads(:random; warn = false) + @test isnothing(setaffinities(masks)) + for (j, tid) in pairs(tids) + @test getaffinity(; threadid = tid) == masks[j] + end + @test getcpuids() == expected_cpuids + end + + @testset "setaffinities (single mask broadcast)" begin + cpuid1, _ = get_two_cpuids() + pinthreads(fill(cpuid1, nt); warn = false) + mask = getaffinity(; threadid = firsttid) + pinthreads(:random; warn = false) + @test isnothing(setaffinities(mask)) + @test all(tid -> getaffinity(; threadid = tid) == mask, tids) + @test all(i -> getcpuid(; threadid = i) == cpuid1, tids) + end + + @testset "setaffinities_cpuids (vector of cpuid vectors)" begin + cpuid1, cpuid2 = get_two_cpuids() + cpuids_vec = fill([cpuid1], nt) + cpuids_vec[2] = [cpuid2] + @test isnothing(setaffinities_cpuids(cpuids_vec)) + @test all(tid -> ispinned(; threadid = tid), tids) + # Under faking, cpuids2affinitymask maps by list position, so + # getcpuid may not match the requested cpuid on systems with + # non-contiguous CPU IDs (e.g. FUGAKU). Verify cpuids on host only. + if !ThreadPinning.Faking.isfaking() + @test getcpuid(; threadid = tids[1]) == cpuid1 + @test getcpuid(; threadid = tids[2]) == cpuid2 + end + end + + @testset "setaffinities_cpuids (single cpuid vector broadcast)" begin + cpuid1, _ = get_two_cpuids() + @test isnothing(setaffinities_cpuids([cpuid1])) + @test all(tid -> ispinned(; threadid = tid), tids) + # Same caveat as above: cpuids2affinitymask position mismatch under faking. + if !ThreadPinning.Faking.isfaking() + @test all(i -> getcpuid(; threadid = i) == cpuid1, tids) + end + end + + @testset "setaffinities error: wrong number of masks" begin + masks = [getaffinity(; threadid = i) for i in tids] + @test_throws ArgumentError setaffinities(masks[1:end-1]) + end + + @testset "setaffinities_cpuids error: wrong number of cpuid vectors" begin + cpuid1, _ = get_two_cpuids() + cpuids_vec = fill([cpuid1], nt) + @test_throws ArgumentError setaffinities_cpuids(cpuids_vec[1:end-1]) + end + end + @testset "force = false" begin cpuid1, cpuid2 = get_two_cpuids() ThreadPinningCore.Internals.forget_pin_attempts()