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    ComplexKind, 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
145macro_rules! complex_type {
146    ($name: literal, $ty: ident, $kind: ident, $size: literal) => {
147        #[pliron_type(name = $name, format = "", generate_get = true, verifier = "succ")]
148        #[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
149        pub struct $ty;
150        scalar!($ty);
151        not_packed!($ty);
152        aligned!($ty, $size / 2);
153        sized!($ty, $size);
154
155        #[type_interface_impl]
156        impl ScalarType for $ty {
157            fn elem_type(&self, _ctx: &Context) -> ElemType {
158                ElemType::Complex(ComplexKind::$kind)
159            }
160        }
161    };
162}
163
164complex_type!("cube.c32", Complex32Type, C32, 8);
165complex_type!("cube.c64", Complex64Type, C64, 16);
166
167float_type!("cube.f64", Float64Type, F64, 8);
168float_semantics!(Float64Type, apfloat::Double);
169apfloat_type!(Float64Type, f64, ieee::DoubleS);
170
171float_type!("cube.f32", Float32Type, F32, 4);
172float_semantics!(Float32Type, apfloat::Single);
173apfloat_type!(Float32Type, f32, ieee::SingleS);
174
175float_type!("cube.tf32", TFloat32Type, TF32, 4);
176apfloat_type!(TFloat32Type, tf32, ieee::SingleS);
177
178float_type!("cube.flex32", FloatFlex32Type, Flex32, 4);
179float_semantics!(FloatFlex32Type, apfloat::Single);
180apfloat_type!(FloatFlex32Type, flex32, ieee::SingleS);
181
182float_type!("cube.f16", Float16Type, F16, 2);
183float_semantics!(Float16Type, apfloat::Half);
184apfloat_type!(Float16Type, f16, ieee::HalfS);
185
186float_type!("cube.bf16", BFloat16Type, BF16, 2);
187float_semantics!(BFloat16Type, apfloat::BFloat);
188apfloat_type!(BFloat16Type, bf16, ieee::BFloatS);
189
190float_type!("cube.ue8m0", Float8E8M0Type, UE8M0, 1);
191
192float_type!("cube.e5m2", Float8E5M2Type, E5M2, 1);
193float_semantics!(Float8E5M2Type, apfloat::Float8E5M2);
194apfloat_type!(Float8E5M2Type, e5m2, ieee::Float8E5M2S);
195
196float_type!("cube.e4m3", Float8E4M3Type, E4M3, 1);
197float_semantics!(Float8E4M3Type, apfloat::Float8E4M3FN);
198apfloat_type!(Float8E4M3Type, e4m3, ieee::Float8E4M3FNS);
199
200float_type!("cube.e3m2", Float6E3M2Type, E3M2, 1);
201
202float_type!("cube.e2m3", Float6E2M3Type, E2M3, 1);
203
204float_type!("cube.e2m1", Float4E2M1Type, E2M1, 1, 4);
205apfloat_type!(Float4E2M1Type, e2m1, Float4E2M1S);
206
207#[type_interface_impl]
208impl FloatTypeInterface for TFloat32Type {
209    fn get_semantics(&self) -> Semantics {
210        let precision = 11;
211        Semantics {
212            bits: 19,
213            exp_bits: 8,
214            precision,
215            nonfinite_behavior: NonfiniteBehavior::IEEE754,
216            max_exp: 127,
217            ieee_max_exp: 127,
218            min_exp: -126,
219            ieee_min_exp: -126,
220            nan_significand_base: 0,
221            nan_payload_mask: (1u128 << (precision - 1)) - 1,
222            qnan_significand: 1u128 << (precision - 2),
223        }
224    }
225}
226
227pub struct Float6E3M2S;
228impl ieee::Semantics for Float6E3M2S {
229    const BITS: usize = 6;
230    const EXP_BITS: usize = 3;
231    const NONFINITE_BEHAVIOR: NonfiniteBehavior = NonfiniteBehavior::NanOnly;
232}
233
234pub struct Float6E2M3S;
235impl ieee::Semantics for Float6E2M3S {
236    const BITS: usize = 6;
237    const EXP_BITS: usize = 2;
238    const NONFINITE_BEHAVIOR: NonfiniteBehavior = NonfiniteBehavior::NanOnly;
239}
240
241pub struct Float4E2M1S;
242impl ieee::Semantics for Float4E2M1S {
243    const BITS: usize = 4;
244    const EXP_BITS: usize = 2;
245    const NONFINITE_BEHAVIOR: NonfiniteBehavior = NonfiniteBehavior::NanOnly;
246}
247
248#[type_interface_impl]
249impl FloatTypeInterface for Float8E8M0Type {
250    fn get_semantics(&self) -> Semantics {
251        let precision = 1;
252        Semantics {
253            bits: 8,
254            exp_bits: 8,
255            precision,
256            nonfinite_behavior: NonfiniteBehavior::NanOnly,
257            max_exp: 127,
258            ieee_max_exp: 127,
259            min_exp: -126,
260            ieee_min_exp: -126,
261            nan_significand_base: 0,
262            nan_payload_mask: 0,
263            qnan_significand: 0,
264        }
265    }
266}
267
268#[type_interface_impl]
269impl FloatTypeInterface for Float6E3M2Type {
270    fn get_semantics(&self) -> Semantics {
271        IeeeFloat::<Float6E3M2S>::get_semantics()
272    }
273}
274
275#[type_interface_impl]
276impl FloatTypeInterface for Float6E2M3Type {
277    fn get_semantics(&self) -> Semantics {
278        IeeeFloat::<Float6E2M3S>::get_semantics()
279    }
280}
281
282#[type_interface_impl]
283impl FloatTypeInterface for Float4E2M1Type {
284    fn get_semantics(&self) -> Semantics {
285        IeeeFloat::<Float4E2M1S>::get_semantics()
286    }
287}
288
289/// `IeeeFloat::from_bits` assumes the presence of a sign bit and will overflow when extracting this
290/// non-existent bit. So we need to do custom conversion here.
291#[type_interface_impl]
292impl APFloatType for Float8E8M0Type {
293    fn value_to_f64(&self, val: APFloat) -> f64 {
294        assert!(val.has_semantics::<ue8m0>(), "Should me ue8m0");
295        ue8m0::from_bits(val.to_bits() as u8).to_f64()
296    }
297    fn value_from_f64(&self, val: f64) -> APFloat {
298        let bits = ue8m0::from_f64(val).to_bits() as u128;
299        APFloat::from_bits::<ue8m0>(bits)
300    }
301    fn value_to_string(&self, val: APFloat) -> String {
302        assert!(val.has_semantics::<ue8m0>(), "Should me ue8m0");
303        ue8m0::from_bits(val.to_bits() as u8).to_string()
304    }
305    fn disp_value(
306        &self,
307        val: APFloat,
308        _: &Context,
309        _: &printable::State,
310        f: &mut fmt::Formatter<'_>,
311    ) -> fmt::Result {
312        assert!(val.has_semantics::<ue8m0>(), "Should me ue8m0");
313        write!(f, "{}", ue8m0::from_bits(val.to_bits() as u8))
314    }
315    fn parse_value<'a>(&self, input: &mut StateStream<'a>) -> ParseResult<'a, APFloat> {
316        let val = single_to_f32(float_parse::<apfloat::Single>(input, ())?.0);
317        let val = APFloat::from_bits::<ue8m0>(ue8m0::from_f32(val).to_bits() as u128);
318        Ok(val).into_parse_result()
319    }
320}
321
322#[pliron_type(
323    name = "cube.e2m1x2",
324    format = "",
325    generate_get = true,
326    verifier = "succ"
327)]
328#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
329pub struct Float4E2M1x2Type;
330scalar!(Float4E2M1x2Type);
331aligned!(Float4E2M1x2Type, 1);
332sized!(Float4E2M1x2Type, 1);
333
334#[type_interface_impl]
335impl MaybePackedType for Float4E2M1x2Type {
336    fn packing_factor(&self, _ctx: &Context) -> usize {
337        2
338    }
339}
340
341#[type_interface_impl]
342impl ScalarType for Float4E2M1x2Type {
343    fn elem_type(&self, _ctx: &Context) -> ElemType {
344        FloatKind::E2M1x2.into()
345    }
346}
347
348#[pliron_type(
349    name = "cube.bool",
350    format = "",
351    generate_get = true,
352    verifier = "succ"
353)]
354#[derive(Hash, PartialEq, Eq, Debug, Clone)]
355pub struct BoolType;
356aligned!(BoolType, 1);
357scalar!(BoolType);
358not_packed!(BoolType);
359
360#[type_interface_impl]
361impl ScalarType for BoolType {
362    fn elem_type(&self, _ctx: &Context) -> ElemType {
363        ElemType::Bool
364    }
365}
cubecl_ir/types/scalar.rs

cubecl_ir/types/
scalar.rs
