cubecl-core 0.2.0

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
46

use cubecl_core as cubecl;
use cubecl_core::prelude::*;

#[cube]
pub fn kernel<T: Numeric>(input: &Tensor<T>) {
    let _shape = input.shape(1);
    let _stride = input.stride(1);
    let _length = input.len();
}

mod tests {
    use super::*;
    use cubecl_core::{
        cpa,
        ir::{Item, Operation, Variable},
    };

    type ElemType = F32;

    #[test]
    fn cube_support_tensor_metadata() {
        let mut context = CubeContext::root();
        let input = context.input(0, Item::new(ElemType::as_elem()));

        kernel::__expand::<ElemType>(&mut context, input.into());
        assert_eq!(context.into_scope().operations, inline_macro_ref());
    }

    fn inline_macro_ref() -> Vec<Operation> {
        let mut context = CubeContext::root();
        let item = Item::new(ElemType::as_elem());
        let input = context.input(0, item);

        let mut scope = context.into_scope();
        let input: Variable = input.into();
        let x = scope.create_local(Item::new(UInt::as_elem()));
        let y = scope.create_local(Item::new(UInt::as_elem()));
        let z = scope.create_local(Item::new(UInt::as_elem()));

        cpa!(&mut scope, x = shape(input, 1u32));
        cpa!(&mut scope, y = stride(input, 1u32));
        cpa!(&mut scope, z = len(input));

        scope.operations
    }
}