cubecl_core/frontend/container/
iter.rs

1use cubecl_ir::{OpInserter, dialect::branch::RangeLoopOp};
2
3use crate::{self as cubecl, ir::Scope, prelude::*};
4
5#[cube]
6pub trait SizedContainer<I: CubePrimitive> {
7    /// Return the length of the container.
8    fn len(&self) -> I;
9}
10
11impl<T: CubeIndex<usize> + SizedContainer<usize> + CubeType<ExpandType = NativeExpand<T>>> Iterable
12    for NativeExpand<T>
13where
14    <T::Output as CubeType>::ExpandType: DerefExpand<Target = <T::Output as CubeType>::ExpandType>,
15{
16    type Item = <T::Output as CubeType>::ExpandType;
17
18    fn expand(self, scope: &Scope, mut body: impl FnMut(&Scope, Self::Item)) {
19        let len = self.__expand_len_method(scope);
20
21        let start = scope.const_usize(0);
22        let end = len.read_value(scope);
23        let step = scope.const_usize(1);
24
25        let range_loop = RangeLoopOp::new(scope.ctx_mut(), start, end, step);
26        let i = range_loop.iter_var(scope.ctx());
27        let loop_body = range_loop.loop_body(scope.ctx());
28
29        let mut child = scope.loop_child(OpInserter::new_at_block_end(loop_body));
30
31        let index = NativeExpand::new(i.into());
32        let item = self
33            .__expand_index_method(&child, index)
34            .__expand_deref_method(&child);
35        body(&mut child, item);
36        child.terminate_yield();
37
38        register_range_loop::<usize>(scope, &range_loop, &child);
39        scope.set_may_return(&[child]);
40    }
41
42    fn expand_unroll(self, _scope: &Scope, _body: impl FnMut(&Scope, Self::Item)) {
43        unimplemented!("Can't unroll array iterator")
44    }
45}
cubecl_core/frontend/container/iter.rs

cubecl_core/frontend/container/
iter.rs
