cubecl-core 0.9.0

CubeCL core create
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
37

use crate as cubecl;
use cubecl::prelude::*;

/// Returns the value at `index` in `list` if `condition` is `true`, otherwise returns `value`.
#[cube]
pub fn read_masked<C: CubePrimitive>(mask: bool, list: Slice<C>, index: usize, value: C) -> C {
    let index = index * usize::cast_from(mask);
    let input = list.read_unchecked(index);

    select(mask, input, value)
}

/// Returns the value at `index` in tensor within bounds.
#[cube]
pub fn read_tensor_checked<C: CubePrimitive>(
    tensor: Tensor<C>,
    index: usize,
    #[comptime] unroll_factor: usize,
) -> C {
    let len = tensor.buffer_len() * unroll_factor;
    let in_bounds = index < len;
    let index = index.min(len);

    select(in_bounds, tensor.read_unchecked(index), C::cast_from(0u32))
}

/// Returns the value at `index` in tensor within bounds.
#[cube]
pub fn read_tensor_atomic_checked<C: Numeric>(
    tensor: Tensor<Atomic<Line<C>>>,
    index: usize,
    #[comptime] unroll_factor: usize,
) -> Atomic<Line<C>> {
    let index = index.min(tensor.buffer_len() * unroll_factor);

    tensor.read_unchecked(index)
}