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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

use std::sync::mpsc;

use crate::compute::{
    affinity::get_active_cores,
    threadpool::{ThreadTask, compute_task::ComputeTask, scheduler::Worker},
};

pub struct NaiveScheduler {
    tx: Vec<mpsc::Sender<ComputeTask>>,
}

impl Default for NaiveScheduler {
    fn default() -> Self {
        Self::new()
    }
}

impl NaiveScheduler {
    pub fn new() -> Self {
        let tx = get_active_cores()
            .map(|core_id| {
                let (worker, tx) = NaiveWorker::new();
                worker.spawn_thread(core_id);
                tx
            })
            .collect();

        Self { tx }
    }

    pub fn send(&mut self, index: usize, task: ComputeTask) {
        let _ = self.tx[index].send(task);
    }
}

pub struct NaiveWorker {
    rx: mpsc::Receiver<ComputeTask>,
}

impl NaiveWorker {
    fn new() -> (Self, mpsc::Sender<ComputeTask>) {
        let (tx, rx) = mpsc::channel();
        let rx = Self { rx };
        (rx, tx)
    }
}

impl Worker for NaiveWorker {
    fn work(self) {
        loop {
            let mut task = self.rx.recv().unwrap();
            while !task.is_ready() {
                std::hint::spin_loop();
            }
            task.compute();
        }
    }
}