ruda_kernel/dsl/post_processing/
predicate.rs

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

ruda_kernel/dsl/post_processing/
predicate.rs
