cubecl_cpp/cuda/
convert.rs

1//! Cuda conversion functions
2#![allow(unused)]
3
4use core::{fmt, ops::Deref};
5
6use cubecl_core::{
7    self as cubecl,
8    ir::types::Fp8Format,
9    ir::{
10        dialect::general::CastOp,
11        interfaces::{ScalarType, TypedExt},
12        match_ty,
13        prelude::*,
14        types::{VectorType, scalar::*},
15    },
16    prelude::*,
17};
18use pliron::{printable::Printable, utils::apfloat::Float8E5M2};
19
20use crate::{
21    cuda::{cuda_op_with_out, ty::*},
22    shared::{
23        CppValue,
24        lowering::LowerOp,
25        ty::{TypeExt, TypeExtCPP, TypedExtCPP},
26    },
27    target::Cuda,
28};
29
30/// special cast function for recursive conversion in the case of minifloat to minifloat conversion
31///
32/// Needs to jump through a lot of hoops to deal with CUDA nonsense.
33/// The overview of available conversions is as follows:
34///
35/// | From                     | To             | Extra args                 |
36/// | ------------------------ | -------------- | -------------------------- |
37/// | f16/bf16/f32/f64         | e4m3/e5m2      | Interpretation, saturation |
38/// | f16/bf16/f32/f64         | e3m2/e2m3/e2m1 | Interpretation, rounding   |
39/// | bf16/f32/f64             | e8m0           | saturation, rounding       |
40/// | e4m3/e5m2/e3m2/e2m3/e2m1 | f16            | Interpretation,            |
41/// | e8m0                     | bf16           |                            |
42///
43/// When the input and output don't match these options, we need to do a two-step conversion.
44/// When the input is a minifloat we always need to cast out to `f16`/`bf16`, and then convert to
45/// the actual out type if it differs. Trying to cast ints also requires an extra conversion, and
46/// so does `f16` to `e8m0` (though it's not recommended to do that anyways, you should be using
47/// `e5m2` for that since you don't have 8 bits of exponent in f16).
48///
49/// See also:
50/// <https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH__FP8__MISC.html>
51/// <https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH__FP6__MISC.html>
52/// <https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH__FP4__MISC.html>
53#[op_interface_impl]
54impl LowerOp<Cuda> for CastOp {
55    fn should_lower(&self, ctx: &Context) -> bool {
56        let input = self.input(ctx);
57        let out = self.get_result(ctx);
58        let should_lower_from = (input.is_fp8_fp6_fp4(ctx) || input.is_float4x2(ctx))
59            && intermediate_for_ty(ctx, input.get_type(ctx)) != out.get_type(ctx);
60        let should_lower_to = (out.is_fp8_fp6_fp4(ctx) || out.is_float4x2(ctx))
61            && !encodes_directly(ctx, input, out.get_type(ctx));
62        should_lower_from || should_lower_to
63    }
64
65    fn lower(&self, scope: &Scope) -> Vec<Value> {
66        let ctx = scope.ctx();
67        let mut current = self.input(ctx);
68        let out_ty = self.get_result(ctx).get_type(ctx);
69        if current.is_fp8_fp6_fp4(ctx) || current.is_float4x2(ctx) {
70            let intermediate = intermediate_for_ty(ctx, current.get_type(ctx));
71            current = cast_value(scope, current, intermediate);
72        }
73        if (out_ty.is_fp8_fp6_fp4(ctx) || out_ty.is_float4x2(ctx))
74            && !encodes_directly(ctx, current, out_ty)
75        {
76            let intermediate = match is_fp8(ctx, out_ty) {
77                true => f32_like(ctx, out_ty),
78                false => intermediate_for_ty(ctx, out_ty),
79            };
80            current = cast_value(scope, current, intermediate);
81        }
82        vec![cast_value(scope, current, out_ty)]
83    }
84}
85
86/// fp8 must convert straight from its source: an f16 detour rounds twice.
87fn encodes_directly(ctx: &Context, input: Value, out_ty: TypeHandle) -> bool {
88    if !is_fp8(ctx, out_ty) {
89        return intermediate_for_ty(ctx, out_ty) == input.get_type(ctx);
90    }
91    let scalar = input.get_type(ctx).scalar_ty(ctx);
92    scalar.is_float16(ctx)
93        || scalar.is_bfloat16(ctx)
94        || scalar.is_float32(ctx)
95        || scalar.is_float64(ctx)
96}
97
98/// `is_float8` also covers e8m0, which only ever converts from bf16.
99fn is_fp8(ctx: &Context, ty: TypeHandle) -> bool {
100    Fp8Format::of_type(ctx, ty.scalar_ty(ctx)).is_some()
101}
102
103fn f32_like(ctx: &Context, ty: TypeHandle) -> TypeHandle {
104    vectorized(ctx, Float32Type::get(ctx).to_handle(), ty.vector_size(ctx))
105}
106
107fn intermediate_for_ty(ctx: &Context, ty: TypeHandle) -> TypeHandle {
108    let vector_size = ty.vector_size(ctx);
109    let intermediate = if ty.scalar_ty(ctx).deref(ctx).is::<Float8E8M0Type>() {
110        BFloat16Type::get(ctx).to_handle()
111    } else if ty.is_float4x2(ctx) {
112        return VectorType::get(ctx, Float16Type::get(ctx).to_handle(), vector_size * 2)
113            .to_handle();
114    } else {
115        Float16Type::get(ctx).to_handle()
116    };
117    vectorized(ctx, intermediate, vector_size)
118}
119
120/// A width-1 `VectorType` is a distinct handle from the bare scalar, which callers compare against.
121fn vectorized(ctx: &Context, scalar: TypeHandle, vector_size: usize) -> TypeHandle {
122    if vector_size > 1 {
123        VectorType::get(ctx, scalar, vector_size).to_handle()
124    } else {
125        scalar
126    }
127}
128
129cuda_op_with_out!(CastOp, |op, ctx| {
130    let input = op.input(ctx);
131    let input_name = input.name(ctx);
132    let out_ty = op.get_result(ctx).get_type(ctx);
133    let out_scalar = out_ty.scalar_ty(ctx).deref(ctx);
134
135    if out_scalar.is::<Complex32Type>() {
136        if input.is_complex(ctx) {
137            format!("make_cuFloatComplex({input_name}.x, {input_name}.y)")
138        } else {
139            format!("make_cuFloatComplex({input_name}, 0.0f)")
140        }
141    } else if out_scalar.is::<Complex64Type>() {
142        if input.is_complex(ctx) {
143            format!("make_cuDoubleComplex({input_name}.x, {input_name}.y)")
144        } else {
145            format!("make_cuDoubleComplex({input_name}, 0.0)")
146        }
147    } else if input.is_complex(ctx) {
148        if out_ty.is_tfloat32(ctx) {
149            format!("nvcuda::wmma::__float_to_tf32({input_name}.x)")
150        } else if out_ty.is_bool(ctx) {
151            format!("({input_name}.x != 0 || {input_name}.y != 0)")
152        } else {
153            format!("{}({input_name}.x)", out_ty.to_cpp(ctx))
154        }
155    } else if input.is_fp8_fp6_fp4(ctx) || input.is_packed_fp6_fp8_fp4(ctx) {
156        cast_minifloat_to_half(ctx, input)
157    } else if out_ty.is_fp8_fp6_fp4(ctx) || out_ty.is_packed_fp6_fp8_fp4(ctx) {
158        cast_half_to_minifloat(ctx, input, out_ty)
159    } else if out_ty.is_tfloat32(ctx) {
160        format!("nvcuda::wmma::__float_to_tf32({input_name})")
161    } else {
162        format!("{}({input_name})", out_ty.to_cpp(ctx))
163    }
164});
165
166// Cast from minifloat to half/bf16. Could be made more generic, but a simple mapping is easier
167// to understand. The naming is very inconsistent (i.e. halfraw2 vs bf162raw)
168fn cast_minifloat_to_half(ctx: &Context, input: Value) -> String {
169    let in_ty = input.get_type(ctx).deref(ctx);
170    let in_val = input.name(ctx);
171    match_ty!((in_ty) {
172        Float8E8M0Type => format!("__nv_bfloat16(__nv_cvt_e8m0_to_bf16raw({in_val}))"),
173        Float8E8M0x2Type => format!("__nv_bfloat162(__nv_cvt_e8m0x2_to_bf162raw({in_val}))"),
174        Float8E4M3Type => format!("__half(__nv_cvt_fp8_to_halfraw({in_val}, __NV_E4M3))"),
175        Float8E4M3x2Type => format!("__half2(__nv_cvt_fp8x2_to_halfraw2({in_val}, __NV_E4M3))"),
176        Float8E5M2Type => format!("__half(__nv_cvt_fp8_to_halfraw({in_val}, __NV_E5M2))"),
177        Float8E5M2x2Type => format!("__half2(__nv_cvt_fp8x2_to_halfraw2({in_val}, __NV_E5M2))"),
178        Float6E2M3Type => format!("__half(__nv_cvt_fp6_to_halfraw({in_val}, __NV_E2M3))"),
179        Float6E2M3x2Type => format!("__half2(__nv_cvt_fp6x2_to_halfraw2({in_val}, __NV_E2M3))"),
180        Float6E3M2Type => format!("__half(__nv_cvt_fp6_to_halfraw({in_val}, __NV_E3M2))"),
181        Float6E3M2x2Type => format!("__half(__nv_cvt_fp6x2_to_halfraw2({in_val}, __NV_E3M2))"),
182        Float4E2M1Type => format!("__half(__nv_cvt_fp4_to_halfraw({in_val}, __NV_E2M1))"),
183        Float4E2M1x2Type => format!("__half2(__nv_cvt_fp4x2_to_halfraw2({in_val}, __NV_E2M1))"),;
184        _ => panic!("Unsupported type {}", in_ty.display(ctx))
185    })
186}
187
188// The naming is very inconsistent (i.e. halfraw2 vs bf162raw). fp8 saturates like the codecs;
189// `__NV_NOSAT` is the header's software path even on sm_89.
190fn cast_half_to_minifloat(ctx: &Context, input: Value, out_ty: TypeHandle) -> String {
191    let in_val = input.name(ctx);
192    let fp8_source = || fp8_source_prefix(ctx, input);
193    match_ty!((out_ty.deref(ctx)) {
194        Float8E8M0Type => format!("__nv_cvt_bfloat16raw_to_e8m0({in_val}, __NV_NOSAT, cudaRoundPosInf)"),
195        Float8E8M0x2Type => format!("__nv_cvt_bfloat162raw_to_e8m0x2({in_val}, __NV_NOSAT, cudaRoundPosInf)"),
196        Float8E4M3Type => format!("__nv_cvt_{}_to_fp8({in_val}, __NV_SATFINITE, __NV_E4M3)", fp8_source()),
197        Float8E4M3x2Type => format!("__nv_cvt_{}_to_fp8x2({in_val}, __NV_SATFINITE, __NV_E4M3)", fp8_source()),
198        Float8E5M2Type => format!("__nv_cvt_{}_to_fp8({in_val}, __NV_SATFINITE, __NV_E5M2)", fp8_source()),
199        Float8E5M2x2Type => format!("__nv_cvt_{}_to_fp8x2({in_val}, __NV_SATFINITE, __NV_E5M2)", fp8_source()),
200        Float6E2M3Type => format!("__nv_cvt_halfraw_to_fp6({in_val}, __NV_E2M3, cudaRoundNearest)"),
201        Float6E2M3x2Type => format!("__nv_cvt_halfraw2_to_fp6x2({in_val}, __NV_E2M3, cudaRoundNearest)"),
202        Float6E3M2Type => format!("__nv_cvt_halfraw_to_fp6({in_val}, __NV_E3M2, cudaRoundNearest)"),
203        Float6E3M2x2Type => format!("__nv_cvt_halfraw2_to_fp6x2({in_val}, __NV_E3M2, cudaRoundNearest)"),
204        Float4E2M1Type => format!("__nv_cvt_halfraw_to_fp4({in_val}, __NV_E2M1, cudaRoundNearest)"),
205        Float4E2M1x2Type => format!("__nv_cvt_halfraw2_to_fp4x2({in_val}, __NV_E2M1, cudaRoundNearest)"),;
206        _ => panic!("Unsupported type {}", out_ty.deref(ctx).display(ctx))
207    })
208}
209
210/// The packed `x2` converters only exist for the 16-bit pairs.
211fn fp8_source_prefix(ctx: &Context, input: Value) -> &'static str {
212    let ty = input.get_type(ctx);
213    let unsupported = || {
214        panic!(
215            "fp8 converts from a float scalar or a packed 16-bit pair, got {}",
216            ty.deref(ctx).display(ctx)
217        )
218    };
219    match_ty!((ty.deref(ctx)) {
220        Float16Type => "halfraw",
221        Float16x2Type => "halfraw2",
222        BFloat16Type => "bfloat16raw",
223        BFloat16x2Type => "bfloat16raw2",
224        Float32Type => "float",
225        Float64Type => "double",;
226        _ => unsupported()
227    })
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233    use crate::{
234        shared::operation::OpToCPP,
235        target::{CtxTarget, Target},
236    };
237    use cubecl_core::ir::{ComplexKind, ConstantValue, ElemType, FloatKind, UIntKind};
238    use pliron::{attribute::boxed_attr_cast, builtin::ops::ConstantOp};
239
240    fn cast(input: ConstantValue, input_ty: ElemType, output_ty: ElemType) -> String {
241        let mut ctx = Context::new();
242        ctx.set_target(Target::Cuda);
243        let input_attr = input.as_attribute(&ctx, input_ty);
244        let input_attr = boxed_attr_cast(input_attr).unwrap();
245        let input = ConstantOp::new(&mut ctx, input_attr).get_result(&ctx);
246        let input_name = input.name(&ctx).to_string();
247        let output_ty = output_ty.to_type(&ctx);
248        let op = CastOp::new(&mut ctx, output_ty, input);
249        let cpp = OpToCPP::<Cuda>::to_cpp(&op, &ctx);
250        cpp.split_once(" = ")
251            .unwrap()
252            .1
253            .replace(&input_name, "input")
254    }
255
256    #[test]
257    fn complex_casts_use_cucomplex_components_and_constructors() {
258        assert_eq!(
259            cast(
260                ConstantValue::UInt(0),
261                UIntKind::U32.into(),
262                ComplexKind::C64.into()
263            ),
264            "make_cuDoubleComplex(input, 0.0);\n"
265        );
266        assert_eq!(
267            cast(
268                ConstantValue::Float(1.0),
269                FloatKind::F64.into(),
270                ComplexKind::C32.into()
271            ),
272            "make_cuFloatComplex(input, 0.0f);\n"
273        );
274        assert_eq!(
275            cast(
276                ConstantValue::Complex(1.0, 2.0),
277                ComplexKind::C64.into(),
278                ComplexKind::C32.into()
279            ),
280            "make_cuFloatComplex(input.x, input.y);\n"
281        );
282        assert_eq!(
283            cast(
284                ConstantValue::Complex(1.0, 2.0),
285                ComplexKind::C32.into(),
286                ComplexKind::C64.into()
287            ),
288            "make_cuDoubleComplex(input.x, input.y);\n"
289        );
290        assert_eq!(
291            cast(
292                ConstantValue::Complex(1.0, 2.0),
293                ComplexKind::C32.into(),
294                FloatKind::F64.into()
295            ),
296            "double(input.x);\n"
297        );
298        assert_eq!(
299            cast(
300                ConstantValue::Complex(0.0, 1.0),
301                ComplexKind::C32.into(),
302                ElemType::Bool
303            ),
304            "(input.x != 0 || input.y != 0);\n"
305        );
306    }
307}
cubecl_cpp/cuda/convert.rs

cubecl_cpp/cuda/
convert.rs
