cubecl_core/post_processing/
bf16.rs

1//! Conversion to `bf16` from a source wider than `f32`.
2//!
3//! A backend converts to `bf16` from `f32` with one rounding, but from an `f64` or a 32- or 64-bit
4//! integer it goes through the nearest `f32` and rounds twice: a value just past a `bf16` tie can
5//! round onto the tie in `f32`, then to even the wrong way. [`LowerWideBf16Cast`] narrows such a
6//! source to `f32` by rounding to odd instead, which `f32`'s width keeps exact for the conversion
7//! to `bf16` that follows.
8
9use alloc::vec;
10use cubecl_ir::{
11    NamedRewrite, Scope, dialect::general::CastOp, interfaces::TypedExt, prelude::*,
12    types::scalar::BFloat16Type,
13};
14use pliron::r#type::TypeHandle;
15
16use crate::{self as cubecl, prelude::*};
17
18define_size!(N);
19
20const F32_SIGNIFICAND_BITS: u32 = f32::MANTISSA_DIGITS;
21const F32_MANTISSA_BITS: u32 = f32::MANTISSA_DIGITS - 1;
22const F32_EXPONENT_BIAS: u32 = (f32::MAX_EXP - 1) as u32;
23
24pub type LowerWideBf16CastPass = MatchRewritePass<LowerWideBf16Cast>;
25
26/// Rewrites a cast to `bf16` from a source wider than `f32` as the cast from an `f32` rounded to
27/// odd, which the backend converts natively with the one rounding the source needs.
28#[derive(new, Clone, Copy, Debug, Default, NamedRewrite)]
29pub struct LowerWideBf16Cast;
30
31impl MatchRewrite for LowerWideBf16Cast {
32    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
33        op.is_op::<CastOp>(ctx)
34            && is_bf16(ctx, op.result(ctx))
35            && needs_rounding_to_odd(ctx, op.operand(ctx, 0).scalar_ty(ctx))
36    }
37
38    fn rewrite(
39        &mut self,
40        ctx: &mut Context,
41        rewriter: &mut MatchRewriter,
42        op: Ptr<Operation>,
43    ) -> Result<()> {
44        let scope = Scope::from_context_and_inserter(ctx, rewriter);
45        let input = op.operand(ctx, 0);
46        let result_ty = op.result(ctx).get_type(ctx);
47        let lanes = input.vector_size(ctx);
48        debug_assert_eq!(
49            lanes,
50            result_ty.vector_size(ctx),
51            "A cast keeps its vectorization, so one `N` describes both sides"
52        );
53        scope.register_size::<N>(lanes);
54
55        let value = cast_value(&scope, to_f32_rounding_once(&scope, input), result_ty);
56        rewriter.replace_operation_with_values(ctx, op, vec![value]);
57        Ok(())
58    }
59}
60
61/// Narrows to `f32` by rounding to odd: an inexact result keeps its low bit set, so rounding it
62/// once more to `bf16` lands where rounding `value` straight to `bf16` would. Rounding to the
63/// nearest `f32` first would round twice.
64#[cube]
65fn f64_to_f32_round_to_odd<N: Size>(value: Vector<f64, N>) -> Vector<f32, N> {
66    let nearest = Vector::<f32, N>::cast_from(value);
67    let widened = Vector::<f64, N>::cast_from(nearest);
68    let bits = Vector::<u32, N>::reinterpret(nearest);
69    // The bits of an `f32` order as its magnitude does, so a step down is one ulp toward zero.
70    let rounded_away = widened.abs().greater_than(&value.abs());
71    let toward_zero = select_many(rounded_away, bits - Vector::new(1u32), bits);
72    let odd = select_many(widened.equal(&value), bits, toward_zero | Vector::new(1u32));
73    Vector::<f32, N>::reinterpret(odd)
74}
75
76/// [`f64_to_f32_round_to_odd`] for an unsigned integer: its top 24 significant bits, with a
77/// sticky low bit for whatever was dropped below them. Generic over the width, so a 32-bit
78/// source does not pay for 64-bit arithmetic.
79#[cube]
80fn unsigned_to_f32_round_to_odd<U: Int, N: Size>(magnitude: Vector<U, N>) -> Vector<f32, N> {
81    let bits = Vector::new(U::size_bits().comptime() as u32);
82    let significant = bits - Vector::leading_zeros(magnitude);
83    let excess = select_many(
84        significant.greater_than(&Vector::new(F32_SIGNIFICAND_BITS)),
85        significant - Vector::new(F32_SIGNIFICAND_BITS),
86        Vector::new(0u32),
87    );
88    let shift = Vector::<U, N>::cast_from(excess);
89    let kept = magnitude >> shift;
90    let dropped = (kept << shift).not_equal(&magnitude);
91    let odd = kept | Vector::<U, N>::cast_from(dropped);
92    // `odd` fits the significand and the scale is a power of two: neither step rounds.
93    let scale = Vector::<f32, N>::reinterpret(
94        (excess + Vector::new(F32_EXPONENT_BIAS)) << Vector::new(F32_MANTISSA_BITS),
95    );
96    Vector::<f32, N>::cast_from(odd) * scale
97}
98
99/// [`unsigned_to_f32_round_to_odd`] on the magnitude, the sign put back after. `U` is the
100/// unsigned type of `S`'s width.
101#[cube]
102fn signed_to_f32_round_to_odd<S: Int, U: Int, N: Size>(value: Vector<S, N>) -> Vector<f32, N> {
103    let negative = value.less_than(&Vector::zero());
104    // The most negative value negates to itself, whose bits are its magnitude.
105    let magnitude =
106        Vector::<U, N>::reinterpret(select_many(negative, Vector::zero() - value, value));
107    let rounded = unsigned_to_f32_round_to_odd::<U, N>(magnitude);
108    select_many(negative, Vector::new(0.0f32) - rounded, rounded)
109}
110
111fn is_bf16(ctx: &Context, value: impl Typed) -> bool {
112    value
113        .try_get_scalar_ty(ctx)
114        .is_some_and(|scalar| scalar.deref(ctx).is::<BFloat16Type>())
115}
116
117/// `N` lanes as an `f32` that rounds to `bf16` as the lanes themselves would. Only a source with
118/// more significant bits than `f32` needs rounding to odd; any other converts exactly.
119fn to_f32_rounding_once(scope: &Scope, value: Value) -> Value {
120    let scalar = value.scalar_ty(scope.ctx());
121    if !needs_rounding_to_odd(scope.ctx(), scalar) {
122        return cast_value(scope, value, Vector::<f32, N>::__expand_as_type(scope));
123    }
124    let signed = scalar.is_signed_int(scope.ctx());
125    if scalar.is_float64(scope.ctx()) {
126        let value = cast_value(scope, value, Vector::<f64, N>::__expand_as_type(scope));
127        f64_to_f32_round_to_odd::expand::<N>(scope, value.into()).read_value(scope)
128    } else if scalar.size_bits(scope.ctx()) <= u32::BITS as usize {
129        int_to_f32_round_to_odd::<i32, u32>(scope, value, signed)
130    } else {
131        int_to_f32_round_to_odd::<i64, u64>(scope, value, signed)
132    }
133}
134
135/// [`signed_to_f32_round_to_odd`] or [`unsigned_to_f32_round_to_odd`] at the width of `S`/`U`.
136fn int_to_f32_round_to_odd<S: Int, U: Int>(scope: &Scope, value: Value, signed: bool) -> Value {
137    match signed {
138        true => {
139            let value = cast_value(scope, value, Vector::<S, N>::__expand_as_type(scope));
140            signed_to_f32_round_to_odd::expand::<S, U, N>(scope, value.into()).read_value(scope)
141        }
142        false => {
143            let value = cast_value(scope, value, Vector::<U, N>::__expand_as_type(scope));
144            unsigned_to_f32_round_to_odd::expand::<U, N>(scope, value.into()).read_value(scope)
145        }
146    }
147}
148
149/// Whether a `scalar` source carries more significant bits than `f32`, which would round twice on
150/// its way to `bf16` through a nearest `f32`.
151fn needs_rounding_to_odd(ctx: &Context, scalar: TypeHandle) -> bool {
152    let integer = scalar.is_int(ctx) || scalar.is_index(ctx);
153    scalar.is_float64(ctx) || integer && scalar.size_bits(ctx) > F32_SIGNIFICAND_BITS as usize
154}
cubecl_core/post_processing/bf16.rs

cubecl_core/post_processing/
bf16.rs
