Skip to main content

cubecl_ir/
apfloat.rs

1use alloc::string::String;
2use core::{any::TypeId, fmt};
3use cubecl_macros_internal::TypeHash;
4
5use pliron::{
6    builtin::type_interfaces::FloatTypeInterface,
7    context::Context,
8    derive::type_interface,
9    parsable::{ParseResult, StateStream},
10    printable,
11};
12use rustc_apfloat::{
13    Float,
14    ieee::{IeeeFloat, Semantics},
15};
16
17use crate::verify_ty_succ;
18
19/// Type erased floating point
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash)]
21pub struct APFloat {
22    /// Raw bits returned from `IeeeFloat<semantics>::to_bits`
23    bits: u128,
24    /// Marker for the semantics that can be used to guard against incorrect conversions
25    semantics: TypeId,
26}
27
28impl APFloat {
29    pub fn to_bits(self) -> u128 {
30        self.bits
31    }
32
33    pub fn from_bits<S: 'static>(bits: u128) -> Self {
34        APFloat {
35            bits,
36            semantics: TypeId::of::<S>(),
37        }
38    }
39
40    pub fn has_semantics<S: 'static>(&self) -> bool {
41        self.semantics == TypeId::of::<S>()
42    }
43
44    pub fn to_ieee<S: Semantics + 'static>(self) -> IeeeFloat<S> {
45        assert_eq!(TypeId::of::<S>(), self.semantics, "Mismatched semantics");
46        IeeeFloat::from_bits(self.to_bits())
47    }
48
49    pub fn from_ieee<S: Semantics + 'static>(float: IeeeFloat<S>) -> APFloat {
50        Self::from_bits::<S>(float.to_bits())
51    }
52}
53
54#[type_interface]
55pub trait APFloatType: FloatTypeInterface {
56    verify_ty_succ!();
57    fn value_to_f64(&self, val: APFloat) -> f64;
58    fn value_from_f64(&self, val: f64) -> APFloat;
59    fn value_to_string(&self, val: APFloat) -> String;
60
61    fn disp_value(
62        &self,
63        val: APFloat,
64        ctx: &Context,
65        state: &printable::State,
66        f: &mut fmt::Formatter<'_>,
67    ) -> fmt::Result;
68    fn parse_value<'a>(&self, state_stream: &mut StateStream<'a>) -> ParseResult<'a, APFloat>;
69}
70
71macro_rules! apfloat_type {
72    ($ty: ty, $rust_num: ty, $sem: ty) => {
73        #[type_interface_impl]
74        impl APFloatType for $ty {
75            #[allow(unused_imports)]
76            fn value_to_f64(&self, val: APFloat) -> f64 {
77                num_traits::ToPrimitive::to_f64(&<$rust_num>::from_bits(
78                    val.to_bits().try_into().unwrap(),
79                ))
80                .unwrap()
81            }
82            fn value_from_f64(&self, val: f64) -> APFloat {
83                use rustc_apfloat::Float;
84                let val: $rust_num = num_traits::NumCast::from(val).unwrap();
85                APFloat::from_ieee(IeeeFloat::<$sem>::from_bits(val.to_bits().into()))
86            }
87            fn value_to_string(&self, val: APFloat) -> alloc::string::String {
88                alloc::format!("{:#}", val.to_ieee::<$sem>())
89            }
90            fn disp_value(
91                &self,
92                val: APFloat,
93                ctx: &Context,
94                state: &pliron::printable::State,
95                f: &mut fmt::Formatter<'_>,
96            ) -> fmt::Result {
97                let val = val.to_ieee::<$sem>();
98                let val = &val as &dyn pliron::utils::apfloat::DynFloat;
99                pliron::printable::Printable::fmt(val, ctx, state, f)
100            }
101            fn parse_value<'a>(
102                &self,
103                state_stream: &mut StateStream<'a>,
104            ) -> ParseResult<'a, APFloat> {
105                use pliron::parsable::IntoParseResult;
106                Ok(APFloat::from_ieee(
107                    pliron::utils::apfloat::float_parse::<IeeeFloat<$sem>>(state_stream, ())?.0,
108                ))
109                .into_parse_result()
110            }
111        }
112    };
113}
114pub(crate) use apfloat_type;