cubecl-ir 0.11.0-pre.3

Intermediate representation for CubeCL
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
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

use alloc::string::String;
use core::{any::TypeId, fmt};
use cubecl_macros_internal::TypeHash;

use pliron::{
    builtin::type_interfaces::FloatTypeInterface,
    context::Context,
    derive::type_interface,
    parsable::{ParseResult, StateStream},
    printable,
};
use rustc_apfloat::{
    Float,
    ieee::{IeeeFloat, Semantics},
};

use crate::verify_ty_succ;

/// Type erased floating point
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash)]
pub struct APFloat {
    /// Raw bits returned from `IeeeFloat<semantics>::to_bits`
    bits: u128,
    /// Marker for the semantics that can be used to guard against incorrect conversions
    semantics: TypeId,
}

impl APFloat {
    pub fn to_bits(self) -> u128 {
        self.bits
    }

    pub fn from_bits<S: 'static>(bits: u128) -> Self {
        APFloat {
            bits,
            semantics: TypeId::of::<S>(),
        }
    }

    pub fn has_semantics<S: 'static>(&self) -> bool {
        self.semantics == TypeId::of::<S>()
    }

    pub fn to_ieee<S: Semantics + 'static>(self) -> IeeeFloat<S> {
        assert_eq!(TypeId::of::<S>(), self.semantics, "Mismatched semantics");
        IeeeFloat::from_bits(self.to_bits())
    }

    pub fn from_ieee<S: Semantics + 'static>(float: IeeeFloat<S>) -> APFloat {
        Self::from_bits::<S>(float.to_bits())
    }
}

#[type_interface]
pub trait APFloatType: FloatTypeInterface {
    verify_ty_succ!();
    fn value_to_f64(&self, val: APFloat) -> f64;
    fn value_from_f64(&self, val: f64) -> APFloat;
    fn value_to_string(&self, val: APFloat) -> String;

    fn disp_value(
        &self,
        val: APFloat,
        ctx: &Context,
        state: &printable::State,
        f: &mut fmt::Formatter<'_>,
    ) -> fmt::Result;
    fn parse_value<'a>(&self, state_stream: &mut StateStream<'a>) -> ParseResult<'a, APFloat>;
}

macro_rules! apfloat_type {
    ($ty: ty, $rust_num: ty, $sem: ty) => {
        #[type_interface_impl]
        impl APFloatType for $ty {
            #[allow(unused_imports)]
            fn value_to_f64(&self, val: APFloat) -> f64 {
                num_traits::ToPrimitive::to_f64(&<$rust_num>::from_bits(
                    val.to_bits().try_into().unwrap(),
                ))
                .unwrap()
            }
            fn value_from_f64(&self, val: f64) -> APFloat {
                use rustc_apfloat::Float;
                let val: $rust_num = num_traits::NumCast::from(val).unwrap();
                APFloat::from_ieee(IeeeFloat::<$sem>::from_bits(val.to_bits().into()))
            }
            fn value_to_string(&self, val: APFloat) -> alloc::string::String {
                alloc::format!("{:#}", val.to_ieee::<$sem>())
            }
            fn disp_value(
                &self,
                val: APFloat,
                ctx: &Context,
                state: &pliron::printable::State,
                f: &mut fmt::Formatter<'_>,
            ) -> fmt::Result {
                let val = val.to_ieee::<$sem>();
                let val = &val as &dyn pliron::utils::apfloat::DynFloat;
                pliron::printable::Printable::fmt(val, ctx, state, f)
            }
            fn parse_value<'a>(
                &self,
                state_stream: &mut StateStream<'a>,
            ) -> ParseResult<'a, APFloat> {
                use pliron::parsable::IntoParseResult;
                Ok(APFloat::from_ieee(
                    pliron::utils::apfloat::float_parse::<IeeeFloat<$sem>>(state_stream, ())?.0,
                ))
                .into_parse_result()
            }
        }
    };
}
pub(crate) use apfloat_type;