cubecl_ir/types/
scalar.rs

1use alloc::string::{String, ToString};
2use core::fmt;
3
4use cubecl_common::{e2m1, e4m3, e5m2, flex32, tf32, ue8m0};
5use half::{bf16, f16};
6use pliron::{
7    builtin::{
8        type_interfaces::FloatTypeInterface,
9        types::{IntegerType, Signedness},
10    },
11    context::Context,
12    derive::{pliron_type, type_interface_impl},
13    parsable::{IntoParseResult, ParseResult, StateStream},
14    printable,
15    utils::apfloat::{self, GetSemantics, Semantics, float_parse, single_to_f32},
16};
17use rustc_apfloat::ieee::{self, IeeeFloat, NonfiniteBehavior};
18
19use crate::{
20    ContextExt, ElemType, FloatKind, IntKind, UIntKind, aligned,
21    apfloat::{APFloat, APFloatType, apfloat_type},
22    interfaces::{AlignedType, MaybePackedType, ScalarType, SizedType, not_packed},
23    scalar, sized,
24};
25
26scalar!(IntegerType);
27not_packed!(IntegerType);
28
29#[type_interface_impl]
30impl AlignedType for IntegerType {
31    fn align(&self, _ctx: &Context) -> usize {
32        self.width().div_ceil(8) as usize
33    }
34}
35
36#[type_interface_impl]
37impl SizedType for IntegerType {
38    fn size(&self, _ctx: &Context) -> usize {
39        self.width().div_ceil(8) as usize
40    }
41}
42
43#[type_interface_impl]
44impl ScalarType for IntegerType {
45    fn elem_type(&self, _ctx: &Context) -> ElemType {
46        match (self.width(), self.signedness()) {
47            (8, Signedness::Signed) => IntKind::I8.into(),
48            (16, Signedness::Signed) => IntKind::I16.into(),
49            (32, Signedness::Signed) => IntKind::I32.into(),
50            (64, Signedness::Signed) => IntKind::I64.into(),
51            (8, _) => UIntKind::U8.into(),
52            (16, _) => UIntKind::U16.into(),
53            (32, _) => UIntKind::U32.into(),
54            (64, _) => UIntKind::U64.into(),
55            _ => unreachable!("Unsupported bit width"),
56        }
57    }
58}
59
60#[pliron_type(
61    name = "cube.poison",
62    format = "",
63    generate_get = true,
64    verifier = "succ"
65)]
66#[derive(Hash, PartialEq, Eq, Debug, Clone)]
67pub struct PoisonType;
68
69#[pliron_type(
70    name = "cube.index",
71    format = "",
72    generate_get = true,
73    verifier = "succ"
74)]
75#[derive(Hash, PartialEq, Eq, Debug, Clone)]
76pub struct IndexType;
77scalar!(IndexType);
78not_packed!(IndexType);
79
80#[type_interface_impl]
81impl AlignedType for IndexType {
82    fn align(&self, ctx: &Context) -> usize {
83        self.size(ctx)
84    }
85}
86
87#[type_interface_impl]
88impl SizedType for IndexType {
89    fn size(&self, ctx: &Context) -> usize {
90        ctx.address_type().size()
91    }
92}
93
94#[type_interface_impl]
95impl ScalarType for IndexType {
96    fn elem_type(&self, _ctx: &Context) -> ElemType {
97        ElemType::Index
98    }
99}
100
101macro_rules! float_type {
102    ($name: literal, $ty: ident, $kind: ident, $size: literal, $size_bits: expr) => {
103        #[pliron_type(name = $name, format = "", generate_get = true, verifier = "succ")]
104        #[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
105        pub struct $ty;
106        scalar!($ty);
107        not_packed!($ty);
108        aligned!($ty, $size);
109
110        #[type_interface_impl]
111        impl ScalarType for $ty {
112            fn elem_type(&self, _ctx: &Context) -> ElemType {
113                FloatKind::$kind.into()
114            }
115        }
116
117        #[type_interface_impl]
118        impl SizedType for $ty {
119            fn size(&self, _ctx: &Context) -> usize {
120                $size
121            }
122
123            fn size_bits(&self, _ctx: &Context) -> usize {
124                $size_bits
125            }
126        }
127    };
128    ($name: literal, $ty: ident, $kind: ident, $size: literal) => {
129        float_type!($name, $ty, $kind, $size, $size * 8);
130    };
131}
132
133// Not all floats have semantics that fit the apfloat model, so separate this out
134macro_rules! float_semantics {
135    ($ty: ty, $semantics: ty) => {
136        #[type_interface_impl]
137        impl FloatTypeInterface for $ty {
138            fn get_semantics(&self) -> Semantics {
139                <$semantics>::get_semantics()
140            }
141        }
142    };
143}
144
145float_type!("cube.f64", Float64Type, F64, 8);
146float_semantics!(Float64Type, apfloat::Double);
147apfloat_type!(Float64Type, f64, ieee::DoubleS);
148
149float_type!("cube.f32", Float32Type, F32, 4);
150float_semantics!(Float32Type, apfloat::Single);
151apfloat_type!(Float32Type, f32, ieee::SingleS);
152
153float_type!("cube.tf32", TFloat32Type, TF32, 4);
154apfloat_type!(TFloat32Type, tf32, ieee::SingleS);
155
156float_type!("cube.flex32", FloatFlex32Type, Flex32, 4);
157float_semantics!(FloatFlex32Type, apfloat::Single);
158apfloat_type!(FloatFlex32Type, flex32, ieee::SingleS);
159
160float_type!("cube.f16", Float16Type, F16, 2);
161float_semantics!(Float16Type, apfloat::Half);
162apfloat_type!(Float16Type, f16, ieee::HalfS);
163
164float_type!("cube.bf16", BFloat16Type, BF16, 2);
165float_semantics!(BFloat16Type, apfloat::BFloat);
166apfloat_type!(BFloat16Type, bf16, ieee::BFloatS);
167
168float_type!("cube.ue8m0", Float8E8M0Type, UE8M0, 1);
169
170float_type!("cube.e5m2", Float8E5M2Type, E5M2, 1);
171float_semantics!(Float8E5M2Type, apfloat::Float8E5M2);
172apfloat_type!(Float8E5M2Type, e5m2, ieee::Float8E5M2S);
173
174float_type!("cube.e4m3", Float8E4M3Type, E4M3, 1);
175float_semantics!(Float8E4M3Type, apfloat::Float8E4M3FN);
176apfloat_type!(Float8E4M3Type, e4m3, ieee::Float8E4M3FNS);
177
178float_type!("cube.e3m2", Float6E3M2Type, E3M2, 1);
179
180float_type!("cube.e2m3", Float6E2M3Type, E2M3, 1);
181
182float_type!("cube.e2m1", Float4E2M1Type, E2M1, 1, 4);
183apfloat_type!(Float4E2M1Type, e2m1, Float4E2M1S);
184
185#[type_interface_impl]
186impl FloatTypeInterface for TFloat32Type {
187    fn get_semantics(&self) -> Semantics {
188        let precision = 11;
189        Semantics {
190            bits: 19,
191            exp_bits: 8,
192            precision,
193            nonfinite_behavior: NonfiniteBehavior::IEEE754,
194            max_exp: 127,
195            ieee_max_exp: 127,
196            min_exp: -126,
197            ieee_min_exp: -126,
198            nan_significand_base: 0,
199            nan_payload_mask: (1u128 << (precision - 1)) - 1,
200            qnan_significand: 1u128 << (precision - 2),
201        }
202    }
203}
204
205pub struct Float6E3M2S;
206impl ieee::Semantics for Float6E3M2S {
207    const BITS: usize = 6;
208    const EXP_BITS: usize = 3;
209    const NONFINITE_BEHAVIOR: NonfiniteBehavior = NonfiniteBehavior::NanOnly;
210}
211
212pub struct Float6E2M3S;
213impl ieee::Semantics for Float6E2M3S {
214    const BITS: usize = 6;
215    const EXP_BITS: usize = 2;
216    const NONFINITE_BEHAVIOR: NonfiniteBehavior = NonfiniteBehavior::NanOnly;
217}
218
219pub struct Float4E2M1S;
220impl ieee::Semantics for Float4E2M1S {
221    const BITS: usize = 4;
222    const EXP_BITS: usize = 2;
223    const NONFINITE_BEHAVIOR: NonfiniteBehavior = NonfiniteBehavior::NanOnly;
224}
225
226#[type_interface_impl]
227impl FloatTypeInterface for Float8E8M0Type {
228    fn get_semantics(&self) -> Semantics {
229        let precision = 1;
230        Semantics {
231            bits: 8,
232            exp_bits: 8,
233            precision,
234            nonfinite_behavior: NonfiniteBehavior::NanOnly,
235            max_exp: 127,
236            ieee_max_exp: 127,
237            min_exp: -126,
238            ieee_min_exp: -126,
239            nan_significand_base: 0,
240            nan_payload_mask: 0,
241            qnan_significand: 0,
242        }
243    }
244}
245
246#[type_interface_impl]
247impl FloatTypeInterface for Float6E3M2Type {
248    fn get_semantics(&self) -> Semantics {
249        IeeeFloat::<Float6E3M2S>::get_semantics()
250    }
251}
252
253#[type_interface_impl]
254impl FloatTypeInterface for Float6E2M3Type {
255    fn get_semantics(&self) -> Semantics {
256        IeeeFloat::<Float6E2M3S>::get_semantics()
257    }
258}
259
260#[type_interface_impl]
261impl FloatTypeInterface for Float4E2M1Type {
262    fn get_semantics(&self) -> Semantics {
263        IeeeFloat::<Float4E2M1S>::get_semantics()
264    }
265}
266
267/// `IeeeFloat::from_bits` assumes the presence of a sign bit and will overflow when extracting this
268/// non-existent bit. So we need to do custom conversion here.
269#[type_interface_impl]
270impl APFloatType for Float8E8M0Type {
271    fn value_to_f64(&self, val: APFloat) -> f64 {
272        assert!(val.has_semantics::<ue8m0>(), "Should me ue8m0");
273        ue8m0::from_bits(val.to_bits() as u8).to_f64()
274    }
275    fn value_from_f64(&self, val: f64) -> APFloat {
276        let bits = ue8m0::from_f64(val).to_bits() as u128;
277        APFloat::from_bits::<ue8m0>(bits)
278    }
279    fn value_to_string(&self, val: APFloat) -> String {
280        assert!(val.has_semantics::<ue8m0>(), "Should me ue8m0");
281        ue8m0::from_bits(val.to_bits() as u8).to_string()
282    }
283    fn disp_value(
284        &self,
285        val: APFloat,
286        _: &Context,
287        _: &printable::State,
288        f: &mut fmt::Formatter<'_>,
289    ) -> fmt::Result {
290        assert!(val.has_semantics::<ue8m0>(), "Should me ue8m0");
291        write!(f, "{}", ue8m0::from_bits(val.to_bits() as u8))
292    }
293    fn parse_value<'a>(&self, input: &mut StateStream<'a>) -> ParseResult<'a, APFloat> {
294        let val = single_to_f32(float_parse::<apfloat::Single>(input, ())?.0);
295        let val = APFloat::from_bits::<ue8m0>(ue8m0::from_f32(val).to_bits() as u128);
296        Ok(val).into_parse_result()
297    }
298}
299
300#[pliron_type(
301    name = "cube.e2m1x2",
302    format = "",
303    generate_get = true,
304    verifier = "succ"
305)]
306#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
307pub struct Float4E2M1x2Type;
308scalar!(Float4E2M1x2Type);
309aligned!(Float4E2M1x2Type, 1);
310sized!(Float4E2M1x2Type, 1);
311
312#[type_interface_impl]
313impl MaybePackedType for Float4E2M1x2Type {
314    fn packing_factor(&self, _ctx: &Context) -> usize {
315        2
316    }
317}
318
319#[type_interface_impl]
320impl ScalarType for Float4E2M1x2Type {
321    fn elem_type(&self, _ctx: &Context) -> ElemType {
322        FloatKind::E2M1x2.into()
323    }
324}
325
326#[pliron_type(
327    name = "cube.bool",
328    format = "",
329    generate_get = true,
330    verifier = "succ"
331)]
332#[derive(Hash, PartialEq, Eq, Debug, Clone)]
333pub struct BoolType;
334aligned!(BoolType, 1);
335scalar!(BoolType);
336not_packed!(BoolType);
337
338#[type_interface_impl]
339impl ScalarType for BoolType {
340    fn elem_type(&self, _ctx: &Context) -> ElemType {
341        ElemType::Bool
342    }
343}
cubecl_ir/types/scalar.rs

cubecl_ir/types/
scalar.rs
