cubecl_core/post_processing/
predicate.rs

1use alloc::vec::Vec;
2use core::{f32, f64};
3
4use crate as cubecl;
5use cubecl_ir::{
6    Comparison, ElemType, FloatKind, Instruction, Operation, Processor, Scope, ScopeProcessing,
7    UIntKind, Value,
8};
9use half::{bf16, f16};
10
11use crate::prelude::*;
12
13define_scalar!(ElemA);
14define_scalar!(IntB);
15define_size!(SizeA);
16
17#[derive(Debug, Default)]
18pub struct PredicateProcessor;
19
20impl Processor for PredicateProcessor {
21    fn transform(&self, mut processing: cubecl_ir::ScopeProcessing) -> cubecl_ir::ScopeProcessing {
22        let mut instructions = Vec::new();
23        core::mem::swap(&mut processing.instructions, &mut instructions);
24
25        for instruction in instructions {
26            if let Operation::Comparison(comparison) = &instruction.operation {
27                match comparison {
28                    Comparison::IsNan(op) => {
29                        run_polyfill(
30                            &mut processing,
31                            op.input,
32                            instruction.out(),
33                            is_nan::expand::<ElemA, IntB, SizeA>,
34                        );
35                        continue;
36                    }
37                    Comparison::IsInf(op) => {
38                        run_polyfill(
39                            &mut processing,
40                            op.input,
41                            instruction.out(),
42                            is_inf::expand::<ElemA, IntB, SizeA>,
43                        );
44                        continue;
45                    }
46                    _ => {}
47                }
48            }
49            processing.instructions.push(instruction);
50        }
51        processing
52    }
53}
54
55fn run_polyfill<T: CubePrimitive, O: CubePrimitive>(
56    processing: &mut ScopeProcessing,
57    input: Value,
58    out: Value,
59    mut polyfill: impl FnMut(&Scope, NativeExpand<T>, u32, u32) -> NativeExpand<O>,
60) {
61    let scope = Scope::root(false).with_global_state(processing.global_state.clone());
62    scope.register_type::<ElemA>(input.storage_type());
63    scope.register_size::<SizeA>(input.vector_size());
64
65    let out_poly = if let ElemType::Float(kind) = input.elem_type() {
66        let (unsigned_ty, bit_width, mantissa_bits) = match kind {
67            FloatKind::F64 => (
68                UIntKind::U64,
69                f64::size_bits().unwrap(),
70                f64::MANTISSA_DIGITS - 1,
71            ),
72            FloatKind::F32 => (
73                UIntKind::U32,
74                f32::size_bits().unwrap(),
75                f32::MANTISSA_DIGITS - 1,
76            ),
77            FloatKind::F16 => (
78                UIntKind::U16,
79                f16::size_bits().unwrap(),
80                f16::MANTISSA_DIGITS - 1,
81            ),
82            FloatKind::BF16 => (
83                UIntKind::U16,
84                bf16::size_bits().unwrap(),
85                bf16::MANTISSA_DIGITS - 1,
86            ),
87            _ => unreachable!(),
88        };
89        scope.register_type::<IntB>(ElemType::UInt(unsigned_ty).into());
90
91        let exp_bits = bit_width as u32 - mantissa_bits - 1;
92
93        polyfill(&scope, input.into(), mantissa_bits, exp_bits).expand
94    } else {
95        panic!("Should be float")
96    };
97
98    let tmp_processing = scope.process([]);
99
100    processing.instructions.extend(tmp_processing.instructions);
101
102    processing
103        .instructions
104        .push(Instruction::new(Operation::Copy(out_poly), out));
105}
106
107#[cube]
108fn is_nan<F: Float, U: Int, N: Size>(
109    x: Vector<F, N>,
110    #[comptime] mantissa_bits: u32,
111    #[comptime] exp_bits: u32,
112) -> Vector<bool, N> {
113    // Need to mark as u64 otherwise it is coerced into i32 which does not fit the values for f64
114    let inf_bits = comptime![((1u64 << exp_bits as u64) - 1u64) << mantissa_bits as u64];
115    let abs_mask = comptime![(1u64 << (exp_bits as u64 + mantissa_bits as u64)) - 1u64];
116
117    let bits: Vector<U, N> = Vector::<U, N>::reinterpret(x);
118
119    let abs_bits = bits & Vector::new(U::cast_from(abs_mask));
120    let inf_bits = Vector::new(U::cast_from(inf_bits));
121
122    abs_bits.greater_than(&inf_bits)
123}
124
125// Same trick as NaN detection following IEEE 754, but check for all 0 bits equality
126#[cube]
127fn is_inf<F: Float, U: Int, N: Size>(
128    x: Vector<F, N>,
129    #[comptime] mantissa_bits: u32,
130    #[comptime] exp_bits: u32,
131) -> Vector<bool, N> {
132    // Need to mark as u64 otherwise it is coerced into i32 which does not fit the values for f64
133    let inf_bits = comptime![((1u64 << exp_bits as u64) - 1u64) << mantissa_bits as u64];
134    let abs_mask = comptime![(1u64 << (exp_bits as u64 + mantissa_bits as u64)) - 1u64];
135
136    let bits: Vector<U, N> = Vector::<U, N>::reinterpret(x);
137
138    let abs_bits = bits & Vector::new(U::cast_from(abs_mask));
139    let inf_bits = Vector::new(U::cast_from(inf_bits));
140
141    abs_bits.equal(&inf_bits)
142}
cubecl_core/post_processing/predicate.rs

cubecl_core/post_processing/
predicate.rs
