cubecl_cpp/cuda/
convert.rs

1//! Cuda conversion functions
2#![allow(unused)]
3
4use core::fmt;
5
6use crate::{
7    Dialect,
8    shared::{Component, Elem, FP8Kind, FmtLeft, Instruction, Item, UnaryInstruction, Value},
9};
10
11/// special cast function for recursive conversion in the case of minifloat to minifloat conversion
12///
13/// Needs to jump through a lot of hoops to deal with CUDA nonsense.
14/// The overview of available conversions is as follows:
15///
16/// | From                     | To             | Extra args                 |
17/// | ------------------------ | -------------- | -------------------------- |
18/// | f16/bf16/f32/f64         | e4m3/e5m2      | Interpretation, saturation |
19/// | f16/bf16/f32/f64         | e3m2/e2m3/e2m1 | Interpretation, rounding   |
20/// | bf16/f32/f64             | e8m0           | saturation, rounding       |
21/// | e4m3/e5m2/e3m2/e2m3/e2m1 | f16            | Interpretation,            |
22/// | e8m0                     | bf16           |                            |
23///
24/// When the input and output don't match these options, we need to do a two-step conversion.
25/// When the input is a minifloat we always need to cast out to `f16`/`bf16`, and then convert to
26/// the actual out type if it differs. Trying to cast ints also requires an extra conversion, and
27/// so does `f16` to `e8m0` (though it's not recommended to do that anyways, you should be using
28/// `e5m2` for that since you don't have 8 bits of exponent in f16).
29///
30/// See also:
31/// <https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH__FP8__MISC.html>
32/// <https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH__FP6__MISC.html>
33/// <https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH__FP4__MISC.html>
34pub(crate) fn special_cast<D: Dialect>(
35    f: &mut std::fmt::Formatter,
36    input: &Value<D>,
37    out: &Value<D>,
38) -> fmt::Result {
39    let mut current_in = *input;
40
41    if matches!(
42        input.elem().unpacked(),
43        Elem::FP4(_) | Elem::FP6(_) | Elem::FP8(_)
44    ) {
45        let item = out.item().with_elem(match input.elem().unpacked() {
46            Elem::FP8(FP8Kind::UE8M0) => Elem::BF16,
47            _ => Elem::F16,
48        });
49        let out_var = if item == out.item() {
50            *out
51        } else {
52            Value::tmp(item)
53        };
54        if *item.elem() == Elem::F16 {
55            cast_minifloat_to_half(f, current_in, out_var)?;
56        } else {
57            cast_scale_to_bfloat(f, current_in, out_var)?;
58        }
59        current_in = out_var;
60    }
61
62    let in_vec = match current_in.item() {
63        Item::Scalar(_) => 1,
64        Item::Vector(_, vectorization) | Item::NativeVector(_, vectorization) => vectorization,
65        _ => panic!("Invalid input item for special cast"),
66    };
67
68    // Broadcast scalars to packing factor
69    if out.item().packing_factor() > 1 && in_vec == 1 {
70        let tmp = Value::tmp(Item::new(input.elem(), out.item().packing_factor()));
71        let assign = Instruction::Assign(UnaryInstruction {
72            input: current_in,
73            out: tmp,
74        });
75        writeln!(f, "{assign}")?;
76        current_in = tmp;
77    }
78
79    let in_vec = match current_in.item() {
80        Item::Scalar(_) => 1,
81        Item::Vector(_, vectorization) | Item::NativeVector(_, vectorization) => vectorization,
82        _ => panic!("Invalid input item for special cast"),
83    };
84
85    if matches!(
86        current_in.elem(),
87        Elem::U8
88            | Elem::U16
89            | Elem::U32
90            | Elem::U64
91            | Elem::I8
92            | Elem::I16
93            | Elem::I32
94            | Elem::I64
95            | Elem::Bool
96    ) {
97        // Precision is irrelevant for int, so use bf16 for the range
98        let tmp = Value::tmp(Item::new(Elem::BF16, in_vec));
99        let assign = Instruction::Assign(UnaryInstruction {
100            input: current_in,
101            out: tmp,
102        });
103        writeln!(f, "{assign}")?;
104        current_in = tmp;
105    }
106
107    if matches!(out.elem().unpacked(), Elem::FP4(_) | Elem::FP6(_)) {
108        return cast_to_fp4_fp6(f, current_in, *out);
109    }
110
111    if matches!(out.elem().unpacked(), Elem::FP8(FP8Kind::UE8M0)) {
112        // Scale can't be converted from half...
113        if matches!(current_in.elem(), Elem::F16) {
114            let item = current_in.item().with_elem(Elem::BF16);
115            let tmp = Value::tmp(item);
116            let assign = Instruction::Assign(UnaryInstruction {
117                input: current_in,
118                out: tmp,
119            });
120            writeln!(f, "{assign}")?;
121            current_in = tmp;
122        }
123        return cast_to_scale(f, current_in, *out);
124    }
125
126    if matches!(out.elem().unpacked(), Elem::FP8(_)) {
127        return cast_to_fp8(f, current_in, *out);
128    }
129
130    if current_in.item() != out.item() {
131        let assign = Instruction::Assign(UnaryInstruction {
132            input: current_in,
133            out: *out,
134        });
135        writeln!(f, "{assign}")?;
136    }
137
138    Ok(())
139}
140
141/// Convert any float to fp4/fp6, with round to nearest
142fn cast_to_fp4_fp6<D: Dialect>(
143    f: &mut fmt::Formatter,
144    input: Value<D>,
145    out: Value<D>,
146) -> fmt::Result {
147    let out_opt = out.optimized();
148    let packing = out_opt.item().packing_factor();
149    let packed = packing == 2;
150    let pack_suffix = if packed { "2" } else { "" };
151
152    let (out_ty, interpretation) = match out_opt.elem() {
153        Elem::FP4(kind) => ("fp4", format!("{kind:?}")),
154        Elem::FP4x2(kind) => ("fp4x2", format!("{kind:?}")),
155        Elem::FP6(kind) => ("fp6", format!("{kind:?}")),
156        Elem::FP6x2(kind) => ("fp6x2", format!("{kind:?}")),
157        _ => unreachable!("Must be fp4 or fp6"),
158    };
159
160    let in_ty = match input.elem().unpacked() {
161        Elem::F64 => format!("double{pack_suffix}"),
162        Elem::TF32 | Elem::F32 => format!("float{pack_suffix}"),
163        Elem::F16 => format!("halfraw{pack_suffix}"),
164        Elem::BF16 => format!("bfloat16raw{pack_suffix}"),
165        _ => unreachable!(),
166    };
167
168    let input = input.optimized();
169
170    handle_unroll(f, out, |f, i| {
171        let in_value = float_to_packed(input, i, packing);
172
173        write!(
174            f,
175            "__nv_cvt_{in_ty}_to_{out_ty}({in_value}, __NV_{interpretation}, cudaRoundNearest)",
176        )
177    })
178}
179
180/// Convert any float except f16 to e8m0
181fn cast_to_scale<D: Dialect>(
182    f: &mut fmt::Formatter,
183    input: Value<D>,
184    out: Value<D>,
185) -> fmt::Result {
186    let out_opt = out.optimized();
187    let packing = out_opt.item().packing_factor();
188    let packed = packing > 1;
189    let pack_suffix = if packed { "2" } else { "" };
190
191    let out_ty = match out_opt.elem() {
192        Elem::FP8(_) => "e8m0",
193        Elem::FP8x2(_) => "e8m0x2",
194        _ => unreachable!("Must be scale factor"),
195    };
196
197    let in_ty = match input.elem() {
198        Elem::F64 => format!("double{pack_suffix}"),
199        Elem::TF32 | Elem::F32 => format!("float{pack_suffix}"),
200        Elem::BF16 => format!("bfloat16{pack_suffix}raw"),
201        _ => unreachable!(),
202    };
203
204    let input = input.optimized();
205
206    handle_unroll(f, out, |f, i| {
207        let in_value = float_to_packed(input, i, packing);
208
209        write!(
210            f,
211            "__nv_cvt_{in_ty}_to_{out_ty}({in_value}, __NV_NOSAT, cudaRoundPosInf)",
212        )
213    })
214}
215
216/// Convert any float to fp8 (except e8m0)
217fn cast_to_fp8<D: Dialect>(f: &mut fmt::Formatter, input: Value<D>, out: Value<D>) -> fmt::Result {
218    let out_opt = out.optimized();
219    let packing = out_opt.item().packing_factor();
220    let packed = packing > 1;
221    let pack_suffix = if packed { "2" } else { "" };
222
223    let (out_ty, interpretation) = match out_opt.elem() {
224        Elem::FP8(kind) => ("fp8", format!("{kind:?}")),
225        Elem::FP8x2(kind) => ("fp8x2", format!("{kind:?}")),
226        _ => unreachable!("Must be fp8"),
227    };
228
229    let in_ty = match input.elem() {
230        Elem::F64 => format!("double{pack_suffix}"),
231        Elem::TF32 | Elem::F32 => format!("float{pack_suffix}"),
232        Elem::BF16 => format!("bfloat16raw{pack_suffix}"),
233        Elem::F16 => format!("halfraw{pack_suffix}"),
234        _ => unreachable!(),
235    };
236
237    let input = input.optimized();
238
239    handle_unroll(f, out, |f, i| {
240        let in_value = float_to_packed(input, i, packing);
241
242        write!(
243            f,
244            "__nv_cvt_{in_ty}_to_{out_ty}({in_value}, __NV_NOSAT, __NV_{interpretation})",
245        )
246    })
247}
248
249/// Pack types that normally wouldn't be optimized into a `vec2` for conversion
250fn float_to_packed<D: Dialect>(input: Value<D>, i: usize, packing: usize) -> String {
251    match input.elem() {
252        Elem::TF32 | Elem::F32 => {
253            let i = i * packing;
254            if packing > 1 {
255                format!("float2 {{ {}, {} }}", input.index(i), input.index(i + 1))
256            } else {
257                format!("{}", input.index(i))
258            }
259        }
260        Elem::F64 => {
261            let i = i * packing;
262            if packing > 1 {
263                format!("double2 {{ {}, {} }}", input.index(i), input.index(i + 1))
264            } else {
265                format!("{}", input.index(i))
266            }
267        }
268        Elem::F16 | Elem::F16x2 | Elem::BF16 | Elem::BF16x2 => format!("{}", input.index(i)),
269        _ => unreachable!(),
270    }
271}
272
273/// Convert any FP8/6/4 except e8m0 to half
274fn cast_minifloat_to_half<D: Dialect>(
275    f: &mut fmt::Formatter,
276    input: Value<D>,
277    out: Value<D>,
278) -> fmt::Result {
279    let in_opt = input.optimized();
280    let out_opt = out.optimized().item();
281
282    let (in_ty, interpretation) = match in_opt.elem() {
283        Elem::FP4(kind) => ("fp4", format!("{kind:?}")),
284        Elem::FP4x2(kind) => ("fp4x2", format!("{kind:?}")),
285        Elem::FP6(kind) => ("fp6", format!("{kind:?}")),
286        Elem::FP6x2(kind) => ("fp6x2", format!("{kind:?}")),
287        Elem::FP8(kind) => ("fp8", format!("{kind:?}")),
288        Elem::FP8x2(kind) => ("fp8x2", format!("{kind:?}")),
289        _ => unreachable!("can only cast minifloat"),
290    };
291
292    let out_ty = match out_opt.elem() {
293        Elem::F16 => "halfraw",
294        Elem::F16x2 => "halfraw2",
295        _ => unreachable!("out type must be half"),
296    };
297
298    handle_unroll(f, out, |f, i| {
299        let input = in_opt.index(i);
300        write!(
301            f,
302            "{}(__nv_cvt_{in_ty}_to_{out_ty}({input}, __NV_{interpretation}))",
303            out_opt.elem()
304        )
305    })
306}
307
308/// Convert an e8m0 scaling factor to bf16
309fn cast_scale_to_bfloat<D: Dialect>(
310    f: &mut fmt::Formatter,
311    input: Value<D>,
312    out: Value<D>,
313) -> fmt::Result {
314    let in_opt = input.optimized();
315    let out_opt = out.optimized().item();
316
317    let in_ty = match in_opt.elem() {
318        Elem::FP8(_) => "e8m0",
319        Elem::FP8x2(_) => "e8m0x2",
320        _ => unreachable!("must be scaling factor in e8m0 format"),
321    };
322
323    let out_ty = match out_opt.elem() {
324        Elem::BF16 => "bf16raw",
325        Elem::BF16x2 => "bf162raw",
326        _ => unreachable!("out type must be half"),
327    };
328
329    handle_unroll(f, out, |f, i| {
330        let input = in_opt.index(i);
331        write!(
332            f,
333            "{}(__nv_cvt_{in_ty}_to_{out_ty}({input}))",
334            out_opt.elem()
335        )
336    })
337}
338
339fn handle_unroll<D: Dialect>(
340    f: &mut fmt::Formatter,
341    out: Value<D>,
342    mut op: impl FnMut(&mut fmt::Formatter, usize) -> fmt::Result,
343) -> fmt::Result {
344    let out_opt = out.item().optimized();
345    let vec = match out_opt {
346        Item::Scalar(_) => 1,
347        Item::Vector(_, vectorization) | Item::NativeVector(_, vectorization) => vectorization,
348        _ => panic!("Invalid input item for special cast"),
349    };
350    let out_var = if out.item() != out_opt {
351        Value::tmp(out_opt)
352    } else {
353        out
354    };
355    write!(f, "{} = ", out_var.fmt_left())?;
356    if vec > 1 {
357        writeln!(f, "{out_opt} {{")?;
358    }
359    for i in 0..vec {
360        op(f, i)?;
361        if i + 1 < vec {
362            f.write_str(",\n")?;
363        }
364    }
365    if vec > 1 {
366        write!(f, "\n}}")?;
367    }
368    f.write_str(";\n")?;
369
370    if out.item() != out_opt {
371        writeln!(
372            f,
373            "{} = reinterpret_cast<{}&>({out_var});",
374            out.fmt_left(),
375            out.item()
376        )?;
377    }
378    Ok(())
379}
cubecl_cpp/cuda/convert.rs

cubecl_cpp/cuda/
convert.rs
