Compiler Explorer

Source code

#![feature(cold_path)]

use core::arch::asm;
use core::sync::atomic::{AtomicU32, Ordering};

#[allow(unused_macros)] // only used on some architectures
macro_rules! select_once {
    (
        sig: fn($($arg:ident: $ArgTy:ty),*) -> $RetTy:ty,
        init: $init:expr,
        call: $call:expr,
    ) => {{
        use core::mem;
        use core::sync::atomic::{AtomicPtr, Ordering};

type Func = unsafe fn($($arg: $ArgTy),*) -> $RetTy;

/// Stores a pointer that is immediately jumped to. By default it is an init function
        /// that sets FUNC to something else.
        static FUNC: AtomicPtr<()> = AtomicPtr::new((initializer as Func) as *mut ());

/// Run once to set the function that will be used for all subsequent calls.
        fn initializer($($arg: $ArgTy),*) -> $RetTy {
            // Select an implementation, ensuring a 'static lifetime.
            let fn_ptr: Func = $init();
            FUNC.store(fn_ptr as *mut (), Ordering::Relaxed);

// Forward the call to the selected function.
            $call(fn_ptr)
        }

let raw: *mut () = unsafe { *FUNC.as_ptr() };
        // let raw: *mut () = FUNC.load(Ordering::Relaxed);

// SAFETY: will only ever be `initializer` or another function pointer that has the
        // 'static lifetime.
        let fn_ptr: Func = unsafe { mem::transmute::<*mut (), Func>(raw) };

$call(fn_ptr)
    }}
}

const FMA: u32 = 0;
const FMA4: u32 = 1;
const FALLBACK: u32 = 2;
const UNINIT: u32 = u32::MAX;
pub static FEATURES: AtomicU32 = AtomicU32::new(UNINIT);

#[unsafe(no_mangle)]
pub fn fma_but_not_what_cfi_likes(x: f64, y: f64, z: f64) -> f64 {
    select_once! {
        sig: fn(x: f64, y: f64, z: f64) -> f64,
        init: || {
            let feat = FEATURES.load(Ordering::Relaxed);
            if feat == FMA {
                fma_with_fma
            } else if feat == FMA4 {
                fma_with_fma4
            } else if feat == FALLBACK {
                fma_fallback
            } else {
                core::hint::cold_path();
                todo!("init")
            }
        },
        // SAFETY: `fn_ptr` is the result of `init`, preconditions have been checked.
        call: |fn_ptr: Func| unsafe { fn_ptr(x, y, z) },
    }
}

#[unsafe(no_mangle)]
pub fn fma_that_makes_cfi_happy(x: f64, y: f64, z: f64) -> f64 {
    let feat = FEATURES.load(Ordering::Relaxed);
    if feat == FMA {
        unsafe { fma_with_fma(x, y, z) }
    } else if feat == FMA4 {
        unsafe { fma_with_fma4(x, y, z) }
    } else if feat == FALLBACK {
        unsafe { fma_fallback(x, y, z) }
    } else {
        core::hint::cold_path();
        todo!("init")
    }
}

/// # Safety
///
/// Must have +fma available.
unsafe fn fma_with_fma(mut x: f64, y: f64, z: f64) -> f64 {
    // debug_assert!(get_cpu_features().contains(cpu_flags::FMA));

// SAFETY: fma is asserted available by precondition, which provides the instruction. No
    // memory access or side effects.
    unsafe {
        asm!(
            "vfmadd213sd {x}, {y}, {z}",
            x = inout(xmm_reg) x,
            y = in(xmm_reg) y,
            z = in(xmm_reg) z,
            options(nostack, nomem, pure),
        );
    }
    x
}

unsafe fn fma_with_fma4(_x: f64, _y: f64, _z: f64) -> f64 {
    todo!("fma4")
}

unsafe fn fma_fallback(x: f64, y: f64, z: f64) -> f64 {
    x.mul_add(y, z)
}