cubecl-cpu 0.11.0-pre.1

CPU runtime for CubeCL
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

use std::mem;

use super::CoreId;
use libc::{
    CPU_ISSET, CPU_SET, CPU_SETSIZE, SYS_gettid, cpu_set_t, sched_getaffinity, sched_setaffinity,
    syscall,
};

pub fn get_active_cores() -> impl Iterator<Item = CoreId> {
    let affinity_mask = get_affinity_mask();

    (0..CPU_SETSIZE as usize)
        .filter(move |i| unsafe { CPU_ISSET(*i, &affinity_mask) })
        .map(CoreId)
}

pub fn set_for_current(core_id: CoreId) {
    let mut set = new_cpu_set();
    let tid = unsafe { syscall(SYS_gettid) } as libc::id_t;
    unsafe { libc::setpriority(libc::PRIO_PROCESS, tid, 0) };
    unsafe { CPU_SET(core_id.0, &mut set) };
    unsafe { sched_setaffinity(0, mem::size_of::<cpu_set_t>(), &set) };
}

fn get_affinity_mask() -> cpu_set_t {
    let mut set = new_cpu_set();
    unsafe { sched_getaffinity(0, mem::size_of::<cpu_set_t>(), &mut set) };
    set
}

fn new_cpu_set() -> cpu_set_t {
    unsafe { std::mem::zeroed::<cpu_set_t>() }
}