cubecl_core/frontend/container/
iter.rs

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

cubecl_core/frontend/container/
iter.rs
