cubecl_core/frontend/iter/
range.rs

1use cubecl_ir::ExpandValue;
2
3use crate::prelude::*;
4
5impl<I: Int> IteratorExpand for RangeExpand<I> {}
6
7impl<I: Int> IndexedIteratorExpand for RangeExpand<I> {
8    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
9        range_len_expand(scope, self.start, self.end)
10    }
11
12    fn get(&mut self, scope: &Scope, index: NativeExpand<usize>) -> NativeExpand<I> {
13        range_get_expand(scope, self.start, index)
14    }
15}
16
17impl<I: Int> IteratorExpand for RangeInclusiveExpand<I> {}
18
19impl<I: Int> IndexedIteratorExpand for RangeInclusiveExpand<I> {
20    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
21        // `last + 1` would wrap when `last` is the largest value.
22        let not_empty = self.start.__expand_le_method(scope, &self.last);
23        let len = range_len_expand(scope, self.start, self.last)
24            .__expand_add_method(scope, NativeExpand::from_lit(scope, 1));
25        select::expand(scope, not_empty, len, NativeExpand::from_lit(scope, 0))
26    }
27
28    fn get(&mut self, scope: &Scope, index: NativeExpand<usize>) -> NativeExpand<I> {
29        range_get_expand(scope, self.start, index)
30    }
31}
32
33impl<I: Int + Into<ExpandValue>> IteratorExpand for SteppedRangeExpand<I> {}
34
35/// The number of items in `start..end`, which is empty rather than negative when `start > end`.
36fn range_len_expand<I: Int>(
37    scope: &Scope,
38    start: NativeExpand<I>,
39    end: NativeExpand<I>,
40) -> NativeExpand<usize> {
41    let len = end
42        .__expand_max_method(scope, start)
43        .__expand_sub_method(scope, start);
44    usize::__expand_cast_from(scope, len)
45}
46
47fn range_get_expand<I: Int>(
48    scope: &Scope,
49    start: NativeExpand<I>,
50    index: NativeExpand<usize>,
51) -> NativeExpand<I> {
52    let index = I::__expand_cast_from(scope, index);
53    start.__expand_add_method(scope, index)
54}
cubecl_core/frontend/iter/range.rs

cubecl_core/frontend/iter/
range.rs
