cubecl_cpp/shared/
value.rs

1use cubecl_core::ir::{
2    attributes::{BoolAttr, ComplexAttr, FloatAttr, IndexAttr, ZeroAttr},
3    types::{barrier::BarrierTokenType, scalar::Complex32Type},
4    verify_attr_succ,
5};
6use pliron::{
7    attribute::{AttrObj, attr_cast},
8    builtin::{attributes::IntegerAttr, ops::ConstantOp},
9    common_traits::Named,
10    context::Context,
11    derive::{attr_interface, attr_interface_impl},
12    identifier::Identifier,
13    r#type::{TypeHandle, Typed},
14    value::Value,
15};
16
17use crate::shared::{
18    shared_op_with_out,
19    ty::{TypeExtCPP, TypedExtCPP},
20};
21
22pub trait CppValue {
23    fn name(&self, ctx: &Context) -> Identifier;
24    fn fmt_left(&self, ctx: &Context) -> String;
25}
26
27impl CppValue for Value {
28    fn name(&self, ctx: &Context) -> Identifier {
29        self.unique_name(ctx)
30    }
31
32    fn fmt_left(&self, ctx: &Context) -> String {
33        let ty = self.get_type(ctx).deref(ctx);
34        let name = self.name(ctx);
35        // C++ has weird semantics so this needs to be mutable for use with `std::move`.
36        // `std::move` preserves constness for the moved value, and the API requires
37        // a non-const `BarrierToken&&`.
38        if ty.is::<BarrierTokenType>() {
39            format!("{} {}", ty.to_cpp(ctx), name)
40        } else {
41            format!("{} const {}", ty.to_cpp(ctx), name)
42        }
43    }
44}
45
46#[attr_interface]
47pub trait CppConstantAttr {
48    verify_attr_succ!();
49    fn as_f64(&self, ctx: &Context) -> f64;
50    fn to_cpp(&self, ctx: &Context) -> String;
51}
52
53#[attr_interface_impl]
54impl CppConstantAttr for IndexAttr {
55    fn as_f64(&self, _ctx: &Context) -> f64 {
56        self.0 as f64
57    }
58    fn to_cpp(&self, _ctx: &Context) -> String {
59        format!("{}", self.0)
60    }
61}
62
63#[attr_interface_impl]
64impl CppConstantAttr for IntegerAttr {
65    fn as_f64(&self, _ctx: &Context) -> f64 {
66        self.value().to_i128() as f64
67    }
68    fn to_cpp(&self, ctx: &Context) -> String {
69        let is_signed = self.get_type().deref(ctx).is_signed();
70        self.value().to_string_decimal(is_signed)
71    }
72}
73
74#[attr_interface_impl]
75impl CppConstantAttr for FloatAttr {
76    fn as_f64(&self, ctx: &Context) -> f64 {
77        self.float_type(ctx).value_to_f64(self.val)
78    }
79    fn to_cpp(&self, ctx: &Context) -> String {
80        // I would prefer to print the bits and use `bit_cast` but that's not well-supported. Keep
81        // an eye on this to make sure it doesn't cause issues.
82        self.float_type(ctx).value_to_string(self.val)
83    }
84}
85
86#[attr_interface_impl]
87impl CppConstantAttr for ComplexAttr {
88    fn as_f64(&self, ctx: &Context) -> f64 {
89        self.float_type(ctx).value_to_f64(self.re)
90    }
91    fn to_cpp(&self, ctx: &Context) -> String {
92        let float_ty = self.float_type(ctx);
93        let re = float_ty.value_to_string(self.re);
94        let im = float_ty.value_to_string(self.im);
95        if self.ty.deref(ctx).is::<Complex32Type>() {
96            format!("make_cuFloatComplex({re}, {im})")
97        } else {
98            format!("make_cuDoubleComplex({re}, {im})")
99        }
100    }
101}
102
103#[attr_interface_impl]
104impl CppConstantAttr for BoolAttr {
105    fn as_f64(&self, _ctx: &Context) -> f64 {
106        self.0 as u8 as f64
107    }
108    fn to_cpp(&self, _ctx: &Context) -> String {
109        self.0.to_string()
110    }
111}
112
113#[attr_interface_impl]
114impl CppConstantAttr for ZeroAttr {
115    fn as_f64(&self, _ctx: &Context) -> f64 {
116        0.0
117    }
118    fn to_cpp(&self, _ctx: &Context) -> String {
119        "{}".to_string()
120    }
121}
122
123shared_op_with_out!(ConstantOp, |op, ctx| {
124    let attr = op.get_attr_builtin_constant_value(ctx).unwrap();
125    format_const(ctx, &attr, op.get_result(ctx).get_type(ctx))
126});
127
128pub(crate) fn format_const(ctx: &Context, value: &AttrObj, ty: TypeHandle) -> String {
129    let const_attr = attr_cast::<dyn CppConstantAttr>(&**value).expect("Should be constant attr");
130    if let Some(attr) = value.downcast_ref::<FloatAttr>() {
131        let val = attr.float_type(ctx).value_to_f64(attr.val);
132        // minifloats are represented as raw bits, so use special handling
133        if ty.is_fp8_fp6_fp4(ctx) {
134            format!("{}", attr.val.to_bits())
135        } else if val.is_nan() {
136            "(0.0f/0.0f)".into()
137        } else if val.is_infinite() && val.is_sign_positive() {
138            "(1.0f/0.0f)".into()
139        } else if val.is_infinite() {
140            "(-1.0f/0.0f)".into()
141        } else {
142            attr.to_cpp(ctx)
143        }
144    } else {
145        const_attr.to_cpp(ctx)
146    }
147}
cubecl_cpp/shared/value.rs

cubecl_cpp/shared/
value.rs
