cubecl-cpu 0.10.0-pre.3

CPU runtime for CubeCL
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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use cubecl_core::ir::{self, ConstantValue, VariableKind};
use tracel_llvm::mlir_rs::{
    Context,
    dialect::{arith, ods::vector},
    ir::{
        Attribute, Type, Value,
        attribute::{FloatAttribute, IntegerAttribute},
        r#type::IntegerType,
    },
};

use super::prelude::*;

impl IntoType for ir::Type {
    fn to_type<'a>(self, context: &'a Context) -> Type<'a> {
        let inner_type = self.storage_type().to_type(context);
        match self.vector_size() {
            size if size > 1 => Type::vector(&[size as u64], inner_type),
            _ => inner_type,
        }
    }
    fn is_vectorized(&self) -> bool {
        self.vector_size() > 1
    }
}

impl<'a> Visitor<'a> {
    pub fn into_attribute(
        context: &'a Context,
        var: Variable,
        item: ir::Type,
    ) -> Option<Attribute<'a>> {
        let r#type = item.storage_type().to_type(context);
        match var.kind {
            VariableKind::Constant(ConstantValue::Float(float)) => {
                if item.is_float() {
                    Some(FloatAttribute::new(context, r#type, float).into())
                } else {
                    Some(IntegerAttribute::new(r#type, float as i64).into())
                }
            }
            VariableKind::Constant(ConstantValue::Bool(bool)) => {
                if item.is_float() {
                    Some(FloatAttribute::new(context, r#type, bool as i64 as f64).into())
                } else {
                    Some(IntegerAttribute::new(r#type, bool as i64).into())
                }
            }
            VariableKind::Constant(ConstantValue::Int(int)) => {
                if item.is_float() {
                    Some(FloatAttribute::new(context, r#type, int as f64).into())
                } else {
                    Some(IntegerAttribute::new(r#type, int).into())
                }
            }
            VariableKind::Constant(ConstantValue::UInt(u_int)) => {
                if item.is_float() {
                    Some(FloatAttribute::new(context, r#type, u_int as f64).into())
                } else {
                    Some(IntegerAttribute::new(r#type, u_int as i64).into())
                }
            }
            _ => None,
        }
    }

    pub fn create_float_constant_from_item(&self, item: ir::Type, constant: f64) -> Value<'a, 'a> {
        let float = item.storage_type().to_type(self.context);
        let constant = FloatAttribute::new(self.context, float, constant);
        let constant = self.append_operation_with_result(arith::constant(
            self.context,
            constant.into(),
            self.location,
        ));
        let result_type = item.to_type(self.context);
        match item.is_vectorized() {
            true => self.append_operation_with_result(vector::splat(
                self.context,
                result_type,
                constant,
                self.location,
            )),
            false => constant,
        }
    }

    pub fn create_int_constant_from_item(&self, item: ir::Type, constant: i64) -> Value<'a, 'a> {
        let integer = item.storage_type().to_type(self.context);
        let constant = IntegerAttribute::new(integer, constant);
        let constant = self.append_operation_with_result(arith::constant(
            self.context,
            constant.into(),
            self.location,
        ));
        let result_type = item.to_type(self.context);
        match item.is_vectorized() {
            true => self.append_operation_with_result(vector::splat(
                self.context,
                result_type,
                constant,
                self.location,
            )),
            false => constant,
        }
    }

    pub fn cast_to_bool(&self, value: Value<'a, 'a>, item: ir::Type) -> Value<'a, 'a> {
        let mut bool = IntegerType::new(self.context, 1).into();
        if item.is_vectorized() {
            bool = Type::vector(&[item.vector_size() as u64], bool);
        }
        self.append_operation_with_result(arith::trunci(value, bool, self.location))
    }

    pub fn cast_to_u8(&self, value: Value<'a, 'a>, item: ir::Type) -> Value<'a, 'a> {
        let mut byte = IntegerType::new(self.context, 8).into();
        if item.is_vectorized() {
            byte = Type::vector(&[item.vector_size() as u64], byte);
        }
        self.append_operation_with_result(arith::extui(value, byte, self.location))
    }
}