cubecl-runtime 0.6.0

Crate that helps creating high performance async runtimes for CubeCL.
Documentation 
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

use crate::memory_id_type;
use crate::memory_management::MemoryHandle;
use crate::{id::HandleRef, server::Handle};
use alloc::vec::Vec;

// The SliceId allows to keep track of how many references there are to a specific slice.
memory_id_type!(SliceId, SliceHandle, SliceBinding);

impl MemoryHandle<SliceBinding> for SliceHandle {
    fn can_mut(&self) -> bool {
        HandleRef::can_mut(self)
    }

    fn binding(self) -> SliceBinding {
        self.binding()
    }
}

/// Take a list of sub-slices of a buffer and create a list of offset handles.
/// Sizes must be in bytes and aligned to the memory alignment.
pub fn offset_handles(base_handle: Handle, sizes_bytes: &[usize]) -> Vec<Handle> {
    let total_size: usize = sizes_bytes.iter().sum();
    let mut offset = 0;
    let mut out = Vec::new();

    for size in sizes_bytes {
        let handle = base_handle
            .clone()
            .offset_start(offset as u64)
            .offset_end((total_size - offset - size) as u64);
        out.push(handle);
        offset += size;
    }

    out
}