Add setaffinities and setaffinities_cpuids

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).

Project.toml

@@ -1,7 +1,7 @@
 name = "ThreadPinning"
 uuid = "811555cd-349b-4f26-b7bc-1f208b848042"
 authors = ["Carsten Bauer <crstnbr@gmail.com> and contributors"]
-version = "1.0.2"
+version = "1.1.0-DEV"
 
 [deps]
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"

docs/src/refs/api_pinning.md

@@ -10,6 +10,8 @@ unpinthreads
 unpinthread
 setaffinity
 setaffinity_cpuids
+setaffinities
+setaffinities_cpuids
 ```
 
 ## Pinning - OpenBLAS

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,

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,15 +62,15 @@ 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)
 
 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

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)

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()