From 38ad67cea7b27ef35ee30743758cfc4f51f8437c Mon Sep 17 00:00:00 2001 From: Carsten Bauer Date: Sat, 25 May 2024 16:16:20 +0200 Subject: [PATCH] Add `jl_getaffinity` and `jl_setaffinity` (#53402) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR adds two functions `jl_getaffinity` and `jl_setaffinity` to the runtime, which are slim wrappers around `uv_thread_getaffinity` and `uv_thread_setaffinity` and can be used to set the affinity of Julia threads. This will * simplify thread pinning (ThreadPinning.jl currently pins threads by spawning tasks that run the necessary ccalls) and * enable users to also pin GC threads (or, more generally, all Julia threads). **Example:** ```julia bauerc@n2lcn0146 julia git:(cb/affinity) ➜ ./julia -q --startup-file=no --threads 2,3 --gcthreads 4,1 julia> cpumasksize = @ccall uv_cpumask_size()::Cint 1024 julia> mask = zeros(Cchar, cpumasksize); julia> jl_getaffinity(tid, mask, cpumasksize) = ccall(:jl_getaffinity, Int32, (Int16, Ptr{Cchar}, Int32), tid, mask, cpumasksize) jl_getaffinity (generic function with 1 method) julia> jl_getaffinity(1, mask, cpumasksize) 0 julia> print(mask[1:Sys.CPU_THREADS]) Int8[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] julia> mask[1] = 0; julia> jl_setaffinity(tid, mask, cpumasksize) = ccall(:jl_setaffinity, Int32, (Int16, Ptr{Cchar}, Int32), tid, mask, cpumasksize) jl_setaffinity (generic function with 1 method) julia> jl_setaffinity(1, mask, cpumasksize) 0 julia> fill!(mask, 0); julia> jl_getaffinity(1, mask, cpumasksize) 0 julia> print(mask[1:Sys.CPU_THREADS]) Int8[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] ``` (cc @vchuravy, @gbaraldi) Would be great to get this into 1.11 (despite feature freeze) because otherwise we won't be able to pin GC threads until 1.12 (likely not until the end of the year). Closes #53073 --------- Co-authored-by: Valentin Churavy Co-authored-by: Dilum Aluthge --- src/jl_exported_funcs.inc | 2 ++ src/julia_threads.h | 3 +++ src/threading.c | 46 +++++++++++++++++++++++++++++++++++++++ test/threads.jl | 12 ++++++++++ 4 files changed, 63 insertions(+) diff --git a/src/jl_exported_funcs.inc b/src/jl_exported_funcs.inc index e2527e3d7aeab8..5e70566ab310e3 100644 --- a/src/jl_exported_funcs.inc +++ b/src/jl_exported_funcs.inc @@ -194,6 +194,7 @@ XX(jl_generating_output) \ XX(jl_generic_function_def) \ XX(jl_gensym) \ + XX(jl_getaffinity) \ XX(jl_getallocationgranularity) \ XX(jl_getnameinfo) \ XX(jl_getpagesize) \ @@ -406,6 +407,7 @@ XX(jl_safepoint_suspend_thread) \ XX(jl_safepoint_resume_thread) \ XX(jl_SC_CLK_TCK) \ + XX(jl_setaffinity) \ XX(jl_set_ARGS) \ XX(jl_set_const) \ XX(jl_set_errno) \ diff --git a/src/julia_threads.h b/src/julia_threads.h index 3a4e9a66cf5a79..60f156b8fab149 100644 --- a/src/julia_threads.h +++ b/src/julia_threads.h @@ -373,6 +373,9 @@ JL_DLLEXPORT int8_t jl_gc_is_in_finalizer(void) JL_NOTSAFEPOINT; JL_DLLEXPORT void jl_wakeup_thread(int16_t tid); +JL_DLLEXPORT int jl_getaffinity(int16_t tid, char *mask, int cpumasksize); +JL_DLLEXPORT int jl_setaffinity(int16_t tid, char *mask, int cpumasksize); + #ifdef __cplusplus } #endif diff --git a/src/threading.c b/src/threading.c index eb76ac579b5389..d6e28e3d808282 100644 --- a/src/threading.c +++ b/src/threading.c @@ -984,6 +984,52 @@ JL_DLLEXPORT int jl_alignment(size_t sz) return jl_gc_alignment(sz); } +// Return values: +// 0 == success +// 1 == invalid thread id provided +// 2 == ptls2 was NULL +// <0 == uv_thread_getaffinity exit code +JL_DLLEXPORT int jl_getaffinity(int16_t tid, char *mask, int cpumasksize) { + int nthreads = jl_atomic_load_acquire(&jl_n_threads); + if (tid < 0 || tid >= nthreads) + return 1; + + // TODO: use correct lock. system_id is only legal if the thread is alive. + jl_ptls_t ptls2 = jl_atomic_load_relaxed(&jl_all_tls_states)[tid]; + if (ptls2 == NULL) + return 2; + uv_thread_t uvtid = ptls2->system_id; + + int ret_uv = uv_thread_getaffinity(&uvtid, mask, cpumasksize); + if (ret_uv != 0) + return ret_uv; + + return 0; // success +} + +// Return values: +// 0 == success +// 1 == invalid thread id provided +// 2 == ptls2 was NULL +// <0 == uv_thread_getaffinity exit code +JL_DLLEXPORT int jl_setaffinity(int16_t tid, char *mask, int cpumasksize) { + int nthreads = jl_atomic_load_acquire(&jl_n_threads); + if (tid < 0 || tid >= nthreads) + return 1; + + // TODO: use correct lock. system_id is only legal if the thread is alive. + jl_ptls_t ptls2 = jl_atomic_load_relaxed(&jl_all_tls_states)[tid]; + if (ptls2 == NULL) + return 2; + uv_thread_t uvtid = ptls2->system_id; + + int ret_uv = uv_thread_setaffinity(&uvtid, mask, NULL, cpumasksize); + if (ret_uv != 0) + return ret_uv; + + return 0; // success +} + #ifdef __cplusplus } #endif diff --git a/test/threads.jl b/test/threads.jl index 172385d1c130ed..24da56e2b89fd9 100644 --- a/test/threads.jl +++ b/test/threads.jl @@ -354,3 +354,15 @@ end @test istaskfailed(t) end end + +@testset "jl_*affinity" begin + cpumasksize = @ccall uv_cpumask_size()::Cint + if !Sys.iswindows() && cpumasksize > 0 # otherwise affinities are not supported on the platform (UV_ENOTSUP) + mask = zeros(Cchar, cpumasksize); + jl_getaffinity = (tid, mask, cpumasksize) -> ccall(:jl_getaffinity, Int32, (Int16, Ptr{Cchar}, Int32), tid, mask, cpumasksize) + jl_setaffinity = (tid, mask, cpumasksize) -> ccall(:jl_setaffinity, Int32, (Int16, Ptr{Cchar}, Int32), tid, mask, cpumasksize) + @test jl_getaffinity(1, mask, cpumasksize) == 0 + fill!(mask, 1) + @test jl_setaffinity(1, mask, cpumasksize) == 0 + end +end