cubecl-core 0.11.0-pre.3

CubeCL core create
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

use cubecl_ir::{OpInserter, dialect::branch::RangeLoopOp};

use crate::{self as cubecl, ir::Scope, prelude::*};

#[cube]
pub trait SizedContainer<I: CubePrimitive> {
    /// Return the length of the container.
    fn len(&self) -> I;
}

impl<T: CubeIndex<usize> + SizedContainer<usize> + CubeType<ExpandType = NativeExpand<T>>> Iterable
    for NativeExpand<T>
where
    <T::Output as CubeType>::ExpandType: DerefExpand<Target = <T::Output as CubeType>::ExpandType>,
{
    type Item = <T::Output as CubeType>::ExpandType;

    fn expand(self, scope: &Scope, mut body: impl FnMut(&Scope, Self::Item)) {
        let len = self.__expand_len_method(scope);

        let start = scope.const_usize(0);
        let end = len.read_value(scope);
        let step = scope.const_usize(1);

        let range_loop = RangeLoopOp::new(scope.ctx_mut(), start, end, step);
        let i = range_loop.iter_var(scope.ctx());
        let loop_body = range_loop.loop_body(scope.ctx());

        let mut child = scope.loop_child(OpInserter::new_at_block_end(loop_body));

        let index = NativeExpand::new(i.into());
        let item = self
            .__expand_index_method(&child, index)
            .__expand_deref_method(&child);
        body(&mut child, item);
        child.terminate_yield();

        register_range_loop::<usize>(scope, &range_loop, &child);
        scope.set_may_return(&[child]);
    }

    fn expand_unroll(self, _scope: &Scope, _body: impl FnMut(&Scope, Self::Item)) {
        unimplemented!("Can't unroll array iterator")
    }
}