From 065aeb6c27ae2a0d5dc616f567053be4b028a565 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 e2527e3d7aeab..5e70566ab310e 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 3a4e9a66cf5a7..60f156b8fab14 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 eb76ac579b538..d6e28e3d80828 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 172385d1c130e..24da56e2b89fd 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