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 = unsafe { list.index_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}