Skip to content

Commit

Permalink
Add jl_getaffinity and jl_setaffinity (#53402)
Browse files Browse the repository at this point in the history
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 <vchuravy@users.noreply.github.com>
Co-authored-by: Dilum Aluthge <dilum@aluthge.com>
(cherry picked from commit 065aeb6)
  • Loading branch information
carstenbauer authored and KristofferC committed Jun 7, 2024
1 parent 4bc3d07 commit 6b4b5ef
Show file tree
Hide file tree
Showing 4 changed files with 63 additions and 0 deletions.
2 changes: 2 additions & 0 deletions src/jl_exported_funcs.inc
Original file line number Diff line number Diff line change
Expand Up @@ -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) \
Expand Down Expand Up @@ -407,6 +408,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) \
Expand Down
3 changes: 3 additions & 0 deletions src/julia_threads.h
Original file line number Diff line number Diff line change
Expand Up @@ -363,6 +363,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
Expand Down
46 changes: 46 additions & 0 deletions src/threading.c
Original file line number Diff line number Diff line change
Expand Up @@ -960,6 +960,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
12 changes: 12 additions & 0 deletions test/threads.jl
Original file line number Diff line number Diff line change
Expand Up @@ -341,3 +341,15 @@ end
@testset "Base.Threads docstrings" begin
@test isempty(Docs.undocumented_names(Threads))
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

0 comments on commit 6b4b5ef

Please sign in to comment.