use cudarc::driver::{DeviceRepr, ValidAsZeroBits};
use super::Scalar;
#[cfg(feature = "cuda-oxide")]
pub use cuda_core::DeviceCopy as OxideCopy;
#[cfg(not(feature = "cuda-oxide"))]
pub trait OxideCopy {}
#[cfg(not(feature = "cuda-oxide"))]
impl<T> OxideCopy for T {}
pub enum CudaType {
F32,
F64,
}
pub trait ScalarCuda: Scalar + ValidAsZeroBits + DeviceRepr + OxideCopy {
fn as_enum() -> CudaType;
fn as_f64(self) -> f64 {
panic!("Unsupported type for as_f64");
}
fn as_str() -> &'static str {
match Self::as_enum() {
CudaType::F32 => "f32",
CudaType::F64 => "f64",
}
}
#[cfg(feature = "cuda-oxide")]
fn shuffle_xor(self, lane_mask: u32) -> Self;
#[cfg(feature = "cuda-oxide")]
fn warp_reduce_sum(self) -> Self;
#[cfg(feature = "cuda-oxide")]
fn to_max_bits(self) -> u64;
}
impl ScalarCuda for f64 {
fn as_enum() -> CudaType {
CudaType::F64
}
fn as_f64(self) -> f64 {
self
}
#[cfg(feature = "cuda-oxide")]
#[inline(always)]
fn shuffle_xor(self, lane_mask: u32) -> Self {
cuda_device::warp::shuffle_xor_f64(self, lane_mask)
}
#[cfg(feature = "cuda-oxide")]
#[inline(always)]
fn warp_reduce_sum(self) -> Self {
cuda_device::warp::reduce_sum_f64(self)
}
#[cfg(feature = "cuda-oxide")]
#[inline(always)]
fn to_max_bits(self) -> u64 {
self.to_bits()
}
}
impl ScalarCuda for f32 {
fn as_enum() -> CudaType {
CudaType::F32
}
fn as_f64(self) -> f64 {
self as f64
}
#[cfg(feature = "cuda-oxide")]
#[inline(always)]
fn shuffle_xor(self, lane_mask: u32) -> Self {
cuda_device::warp::shuffle_xor_f32(self, lane_mask)
}
#[cfg(feature = "cuda-oxide")]
#[inline(always)]
fn warp_reduce_sum(self) -> Self {
cuda_device::warp::reduce_sum_f32(self)
}
#[cfg(feature = "cuda-oxide")]
#[inline(always)]
fn to_max_bits(self) -> u64 {
(self as f64).to_bits()
}
}