cubecl_core/
io.rs

1use crate::{self as cubecl, frontend::buffer_len::expand_buffer_length_native};
2use cubecl::prelude::*;
3use cubecl_ir::pliron::{common_traits::Named, value::Value};
4
5/// Returns the value at `index` in `list` if `condition` is `true`, otherwise returns `value`.
6#[cube]
7pub fn read_masked<C: CubePrimitive>(mask: bool, list: &[C], index: usize, value: C) -> C {
8    let index = index * usize::cast_from(mask);
9    let input = unsafe { *list.get_unchecked(index) };
10
11    select(mask, input, value)
12}
13
14/// Returns the value at `index` in `list` if the index is in bounds, otherwise returns `value`.
15#[cube]
16pub fn read_checked<C: CubePrimitive>(list: &[C], index: usize) -> C {
17    let fallback = comptime![C::Scalar::default()].runtime();
18    let clamped = index.min(list.len() - 1);
19    let input = unsafe { *list.get_unchecked(clamped) };
20
21    select(index == clamped, input, C::cast_from(fallback))
22}
23
24/// Writes the value only if it is in bounds of the buffer
25#[cube]
26pub fn write_checked<C: CubePrimitive>(list: &mut [C], index: usize, value: C) {
27    if index < list.len() {
28        unsafe { *list.get_unchecked_mut(index) = value };
29    }
30}
31
32/// Returns the value at `index` in tensor within bounds.
33#[cube]
34pub fn checked_index(index: usize, buffer_len: usize) -> usize {
35    index.min(buffer_len - 1)
36}
37
38/// Returns the value at `index` in tensor within bounds.
39#[cube]
40pub fn validate_index(
41    #[comptime] buffer_name: &str,
42    index: usize,
43    len: usize,
44    #[comptime] kernel_name: &str,
45) -> usize {
46    let in_bounds = index < len;
47    if !in_bounds {
48        print_oob(kernel_name, index, len, buffer_name);
49    }
50
51    index.min(len)
52}
53
54#[cube]
55#[allow(unused)]
56fn print_oob(
57    #[comptime] kernel_name: &str,
58    index: usize,
59    len: usize,
60    #[comptime] buffer_name: &str,
61) {
62    intrinsic!(|scope| {
63        __expand_debug_print!(
64            scope,
65            alloc::format!(
66                "[VALIDATION {kernel_name}]: Encountered OOB index in {buffer_name} at %u, length is %u\n"
67            ),
68            index,
69            len
70        );
71    })
72}
73
74#[allow(missing_docs)]
75pub fn expand_checked_index(scope: &Scope, list: Value, index: Value) -> Value {
76    let len = expand_buffer_length_native(scope, list);
77    let index = checked_index::expand(scope, index.into(), len.into());
78    index_expand(scope, list, index.value(scope), false)
79}
80
81#[allow(missing_docs)]
82pub fn expand_validate_index(scope: &Scope, list: Value, index: Value, kernel_name: &str) -> Value {
83    let len = expand_buffer_length_native(scope, list);
84    let buffer_name = list.given_name(scope.ctx());
85    let buffer_name = buffer_name
86        .as_ref()
87        .map(|it| it.as_ref())
88        .unwrap_or("buffer");
89    let index = validate_index::expand(scope, buffer_name, index.into(), len.into(), kernel_name);
90    index_expand(scope, list, index.value(scope), false)
91}
cubecl_core/io.rs

cubecl_core/
io.rs
