cubecl_core/
io.rs

1use alloc::string::{String, ToString};
2
3use crate as cubecl;
4use cubecl::prelude::*;
5use cubecl_ir::{Instruction, Operation, Value};
6
7define_scalar!(ElemA);
8define_size!(SizeA);
9
10/// Returns the value at `index` in `list` if `condition` is `true`, otherwise returns `value`.
11#[cube]
12pub fn read_masked<C: CubePrimitive>(mask: bool, list: &[C], index: usize, value: C) -> C {
13    let index = index * usize::cast_from(mask);
14    let input = unsafe { *list.get_unchecked(index) };
15
16    select(mask, input, value)
17}
18
19/// Returns the value at `index` in `list` if the index is in bounds, otherwise returns `value`.
20#[cube]
21pub fn read_checked<C: CubePrimitive>(list: &[C], index: usize) -> C {
22    let fallback = comptime![C::Scalar::default()].runtime();
23    let clamped = index.min(list.len() - 1);
24    let input = unsafe { *list.get_unchecked(clamped) };
25
26    select(index == clamped, input, C::cast_from(fallback))
27}
28
29/// Writes the value only if it is in bounds of the buffer
30#[cube]
31pub fn write_checked<C: CubePrimitive>(list: &mut [C], index: usize, value: C) {
32    if index < list.len() {
33        unsafe { *list.get_unchecked_mut(index) = value };
34    }
35}
36
37/// Returns the value at `index` in tensor within bounds.
38#[cube]
39pub fn checked_index<E: Scalar, N: Size>(
40    index: usize,
41    buffer_len: usize,
42    #[comptime] unroll_factor: usize,
43) -> usize {
44    let len = buffer_len * unroll_factor;
45    index.min(len - 1)
46}
47
48/// Returns the value at `index` in tensor within bounds.
49#[cube]
50pub fn validate_index<E: Scalar, N: Size>(
51    tensor: &[Vector<E, N>],
52    index: usize,
53    buffer_len: usize,
54    #[comptime] unroll_factor: usize,
55    #[comptime] kernel_name: String,
56) -> usize {
57    let len = buffer_len * unroll_factor;
58    let in_bounds = index < len;
59    if !in_bounds {
60        print_oob::<[Vector<E, N>]>(kernel_name, index, len, tensor);
61    }
62
63    index.min(len)
64}
65
66#[cube]
67#[allow(unused)]
68fn print_oob<Out: CubeType<ExpandType: Into<Value>> + ?Sized>(
69    #[comptime] kernel_name: String,
70    index: usize,
71    len: usize,
72    buffer: &Out,
73) {
74    intrinsic!(|scope| {
75        let value: Value = buffer.clone_unchecked().into();
76        let name = value.address_space();
77        __expand_debug_print!(
78            scope,
79            alloc::format!(
80                "[VALIDATION {kernel_name}]: Encountered OOB index in {name} at %u, length is %u\n"
81            ),
82            index,
83            len
84        );
85    })
86}
87
88#[allow(missing_docs)]
89pub fn expand_checked_index(
90    scope: &Scope,
91    list: Value,
92    index: Value,
93    out: Value,
94    unroll_factor: usize,
95) {
96    scope.register_type::<ElemA>(list.ty.storage_type());
97    scope.register_size::<SizeA>(list.ty.vector_size());
98    let len = expand_buffer_length_native(scope, list);
99    let index =
100        checked_index::expand::<ElemA, SizeA>(scope, index.into(), len.into(), unroll_factor);
101    let ptr = index_expand(scope, list, index.expand, false);
102    scope.register(Instruction::new(Operation::Copy(ptr), out));
103}
104
105#[allow(missing_docs)]
106pub fn expand_validate_index(
107    scope: &Scope,
108    list: Value,
109    index: Value,
110    out: Value,
111    unroll_factor: usize,
112    kernel_name: &str,
113) {
114    scope.register_type::<ElemA>(list.ty.storage_type());
115    scope.register_size::<SizeA>(list.ty.vector_size());
116    let len = expand_buffer_length_native(scope, list);
117    let tensor = list.into();
118    let index = validate_index::expand::<ElemA, SizeA>(
119        scope,
120        &tensor,
121        index.into(),
122        len.into(),
123        unroll_factor,
124        kernel_name.to_string(),
125    );
126    let ptr = index_expand(scope, list, index.expand, false);
127    scope.register(Instruction::new(Operation::Copy(ptr), out));
128}
cubecl_core/io.rs

cubecl_core/
io.rs
