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>(
16    tensor: Tensor<C>,
17    index: u32,
18    #[comptime] unroll_factor: u32,
19) -> C {
20    let len = tensor.buffer_len() * unroll_factor;
21    let in_bounds = index < len;
22    let index = Min::min(index, len);
23
24    select(in_bounds, tensor.read_unchecked(index), C::cast_from(0u32))
25}
26
27/// Returns the value at `index` in tensor within bounds.
28#[cube]
29pub fn read_tensor_atomic_checked<C: Numeric>(
30    tensor: Tensor<Atomic<Line<C>>>,
31    index: u32,
32    #[comptime] unroll_factor: u32,
33) -> Atomic<Line<C>> {
34    let index = Min::min(index, tensor.buffer_len() * unroll_factor);
35
36    tensor.read_unchecked(index)
37}
cubecl_core/io.rs

cubecl_core/
io.rs
