ruda_kernel/dsl/
io.rs

1use alloc::{
2    borrow::Cow,
3    string::{String, ToString},
4};
5use derive_more::Display;
6
7use crate::dsl::prelude::*;
8use ruda_core::ir::{ManagedVariable, Variable};
9
10define_scalar!(ElemA);
11define_size!(SizeA);
12
13/// Returns the value at `index` in `list` if `condition` is `true`, otherwise returns `value`.
14#[ruda]
15pub fn read_masked<C: RudaPrimitive>(mask: bool, list: Slice<C>, index: usize, value: C) -> C {
16    let index = index * usize::cast_from(mask);
17    let input = list.read_unchecked(index);
18
19    select(mask, input, value)
20}
21
22/// Returns the value at `index` in tensor within bounds.
23#[ruda]
24pub fn read_tensor_checked<C: RudaPrimitive + Default + IntoRuntime>(
25    tensor: Tensor<C>,
26    index: usize,
27    #[comptime] unroll_factor: usize,
28) -> C {
29    let len = tensor.buffer_len() * unroll_factor;
30    let in_bounds = index < len;
31    let index = index.min(len - 1);
32
33    select(in_bounds, tensor.read_unchecked(index), C::default())
34}
35
36/// Returns the value at `index` in tensor within bounds.
37#[ruda]
38pub fn read_tensor_atomic_checked<C: Scalar>(
39    tensor: Tensor<Atomic<C>>,
40    index: usize,
41    #[comptime] unroll_factor: usize,
42) -> Atomic<C> {
43    let index = index.min(tensor.buffer_len() * unroll_factor - 1);
44
45    tensor.read_unchecked(index)
46}
47
48/// Returns the value at `index` in tensor within bounds.
49#[ruda]
50pub fn read_tensor_validate<C: RudaPrimitive + Default + IntoRuntime>(
51    tensor: Tensor<C>,
52    index: usize,
53    #[comptime] unroll_factor: usize,
54    #[comptime] kernel_name: String,
55) -> C {
56    let len = tensor.buffer_len() * unroll_factor;
57    let in_bounds = index < len;
58    if !in_bounds {
59        print_oob::<Tensor<C>>(kernel_name, OobKind::Read, index, len, &tensor);
60    }
61
62    let index = index.min(len - 1);
63
64    select(in_bounds, tensor.read_unchecked(index), C::default())
65}
66
67/// Returns the value at `index` in tensor within bounds.
68#[ruda]
69pub fn read_tensor_atomic_validate<C: Scalar>(
70    tensor: Tensor<Atomic<C>>,
71    index: usize,
72    #[comptime] unroll_factor: usize,
73    #[comptime] kernel_name: String,
74) -> Atomic<C> {
75    let len = tensor.buffer_len() * unroll_factor;
76    if index >= len {
77        print_oob::<Tensor<Atomic<C>>>(kernel_name, OobKind::Read, index, len, &tensor);
78    }
79    let index = index.min(tensor.buffer_len() * unroll_factor - 1);
80
81    tensor.read_unchecked(index)
82}
83
84#[ruda]
85fn checked_index_assign<E: Scalar, N: Size>(
86    index: usize,
87    value: Vector<E, N>,
88    out: &mut Array<Vector<E, N>>,
89    #[comptime] has_buffer_len: bool,
90    #[comptime] unroll_factor: usize,
91) {
92    let array_len = if has_buffer_len {
93        out.buffer_len()
94    } else {
95        out.len()
96    };
97
98    if index < array_len * unroll_factor {
99        unsafe { out.index_assign_unchecked(index, value) };
100    }
101}
102
103#[ruda]
104fn validate_index_assign<E: Scalar, N: Size>(
105    index: usize,
106    value: Vector<E, N>,
107    out: &mut Array<Vector<E, N>>,
108    #[comptime] has_buffer_len: bool,
109    #[comptime] unroll_factor: usize,
110    #[comptime] kernel_name: String,
111) {
112    let array_len = if has_buffer_len {
113        out.buffer_len()
114    } else {
115        out.len()
116    };
117    let len = array_len * unroll_factor;
118
119    if index < len {
120        unsafe { out.index_assign_unchecked(index, value) };
121    } else {
122        print_oob::<Array<Vector<E, N>>>(kernel_name, OobKind::Write, index, len, out);
123    }
124}
125
126#[derive(Display)]
127enum OobKind {
128    #[display("read")]
129    Read,
130    #[display("write")]
131    Write,
132}
133
134#[ruda]
135#[allow(unused)]
136fn print_oob<Out: RudaType<ExpandType: Into<Variable>>>(
137    #[comptime] kernel_name: String,
138    #[comptime] kind: OobKind,
139    index: usize,
140    len: usize,
141    buffer: &Out,
142) {
143    intrinsic!(|scope| {
144        let name = name_of_var(scope, buffer.into());
145        debug_print_expand!(
146            scope,
147            alloc::format!(
148                "[VALIDATION {kernel_name}]: Encountered OOB {kind} in {name} at %u, length is %u\n"
149            ),
150            index,
151            len
152        );
153    })
154}
155
156fn name_of_var(scope: &Scope, var: Variable) -> Cow<'static, str> {
157    let debug_name = scope.debug.variable_names.borrow().get(&var).cloned();
158    debug_name.unwrap_or_else(|| var.to_string().into())
159}
160
161#[allow(missing_docs)]
162pub fn expand_checked_index_assign(
163    scope: &mut Scope,
164    lhs: Variable,
165    rhs: Variable,
166    out: Variable,
167    unroll_factor: usize,
168) {
169    scope.register_type::<ElemA>(rhs.ty.storage_type());
170    scope.register_size::<SizeA>(rhs.ty.vector_size());
171    checked_index_assign::expand::<ElemA, SizeA>(
172        scope,
173        ManagedVariable::Plain(lhs).into(),
174        ManagedVariable::Plain(rhs).into(),
175        ManagedVariable::Plain(out).into(),
176        out.has_buffer_length(),
177        unroll_factor,
178    );
179}
180
181#[allow(missing_docs)]
182pub fn expand_validate_index_assign(
183    scope: &mut Scope,
184    lhs: Variable,
185    rhs: Variable,
186    out: Variable,
187    unroll_factor: usize,
188    kernel_name: &str,
189) {
190    scope.register_type::<ElemA>(rhs.ty.storage_type());
191    scope.register_size::<SizeA>(rhs.ty.vector_size());
192    validate_index_assign::expand::<ElemA, SizeA>(
193        scope,
194        ManagedVariable::Plain(lhs).into(),
195        ManagedVariable::Plain(rhs).into(),
196        ManagedVariable::Plain(out).into(),
197        out.has_buffer_length(),
198        unroll_factor,
199        kernel_name.to_string(),
200    );
201}
ruda_kernel/dsl/io.rs

ruda_kernel/dsl/
io.rs
