cubecl_core/
io.rs

1use crate as cubecl;
2use cubecl::prelude::*;
3
4/// Returns the value at `index` in `list` if `condition` is `true`, otherwise returns `value`.
5#[cube]
6pub fn read_masked<C: CubePrimitive>(mask: bool, list: Slice<C>, index: u32, value: C) -> C {
7    let index = index * u32::cast_from(mask);
8    let input = list.read_unchecked(index);
9
10    select(mask, input, value)
11}
12
13/// Returns the value at `index` in tensor within bounds.
14#[cube]
15pub fn read_tensor_checked<C: CubePrimitive>(tensor: Tensor<C>, index: u32) -> C {
16    let mask = index < tensor.buffer_len();
17
18    read_masked::<C>(mask, tensor.to_slice(), index, C::cast_from(0u32))
19}
20
21/// Returns the value at `index` in tensor within bounds.
22#[cube]
23pub fn read_tensor_atomic_checked<C: Numeric>(
24    tensor: Tensor<Atomic<Line<C>>>,
25    index: u32,
26) -> Atomic<Line<C>> {
27    let mask = index < tensor.buffer_len();
28
29    let mask_num = u32::cast_from(mask);
30    let index = index * mask_num;
31
32    tensor.read_unchecked(index)
33}