Repository: coreylowman/dfdx (language: rust)

Issue:
Add ToDtype trait for converting between dtypes
This should be similar to ToDevice where you can call `tensor.to_dtype::<f32>()`.

Similarly, we can implement for things in nn `module.to_dtype::<f32>()`.

This will require implementing some cuda kernels.

Ideally we should also look some ways to simplify implementation at the nn level, since at the moment we still have to specify ToDtype/ToDevice explicitly for each struct


Discussion / hints:

Fix patch:
diff --git a/src/nn/mod.rs b/src/nn/mod.rs
index 9ee772961..ae483ae28 100644
--- a/src/nn/mod.rs
+++ b/src/nn/mod.rs
@@ -178,6 +178,7 @@ mod num_params;
 mod reset_params;
 pub mod tensor_collection;
 mod to_device;
+mod to_dtype;
 mod zero_grads;
 
 mod module;
@@ -223,6 +224,7 @@ pub use npz::{LoadFromNpz, SaveToNpz};
 pub use num_params::NumParams;
 pub use reset_params::ResetParams;
 pub use to_device::ToDevice;
+pub use to_dtype::ToDtype;
 pub use zero_grads::ZeroGrads;
 
 pub mod modules {
diff --git a/src/nn/to_dtype.rs b/src/nn/to_dtype.rs
new file mode 100644
index 000000000..c01cca0ea
--- /dev/null
+++ b/src/nn/to_dtype.rs
@@ -0,0 +1,54 @@
+use super::tensor_collection::*;
+use crate::{
+    shapes::{Dtype, Shape},
+    tensor::Tensor,
+    tensor_ops::{Device, ToDtypeKernel},
+};
+
+struct Converter<E> {
+    e: core::marker::PhantomData<E>,
+}
+impl<E1: Dtype, E2: Dtype, D: Device<E1> + Device<E2> + ToDtypeKernel<E1, E2>> TensorVisitor<E1, D>
+    for Converter<E2>
+{
+    type Viewer = ViewTensorRef;
+    type Err = D::Err;
+    type E2 = E2;
+    type D2 = D;
+
+    fn visit<S: Shape>(
+        &mut self,
+        _opts: TensorOptions<S, E1, D>,
+        t: &Tensor<S, E1, D>,
+    ) -> Result<Option<Tensor<S, E2, D>>, Self::Err> {
+        Ok(Some(t.clone().try_to_dtype()?))
+    }
+}
+
+/// Something that can be copied to have a different dtype
+pub trait ToDtype<E1: Dtype, E2: Dtype, D: Device<E1> + Device<E2> + ToDtypeKernel<E1, E2>>:
+    TensorCollection<E1, D>
+{
+    /// Fallible version of [ToDevice::to_dtype]
+    fn try_to_dtype(&self) -> Result<Self::To<E2, D>, D::Err> {
+        let out = Self::iter_tensors(&mut RecursiveWalker {
+            m: self,
+            f: &mut Converter {
+                e: Default::default(),
+            },
+        })?;
+        Ok(out.unwrap())
+    }
+
+    /// Create a copy of `self` with dtype E2
+    fn to_dtype(&self) -> Self::To<E2, D> {
+        self.try_to_dtype().unwrap()
+    }
+}
+
+impl<E1: Dtype, E2: Dtype, D: Device<E1> + Device<E2> + ToDtypeKernel<E1, E2>, T> ToDtype<E1, E2, D>
+    for T
+where
+    T: TensorCollection<E1, D>,
+{
+}
diff --git a/src/tensor_ops/mod.rs b/src/tensor_ops/mod.rs
index a11f4ec34..581f80f90 100644
--- a/src/tensor_ops/mod.rs
+++ b/src/tensor_ops/mod.rs
@@ -195,6 +195,7 @@ mod stddev_to;
 mod sub;
 mod sum_to;
 mod tanh;
+mod to_dtype;
 mod var_to;
 
 pub use abs::abs;
@@ -243,8 +244,11 @@ pub use stddev_to::StddevTo;
 pub use sub::{sub, TrySub};
 pub use sum_to::SumTo;
 pub use tanh::tanh;
+pub use to_dtype::to_dtype;
 pub use var_to::VarTo;
 
+pub(crate) use to_dtype::ToDtypeKernel;
+
 #[cfg(feature = "nightly")]
 mod conv2d;
 #[cfg(feature = "nightly")]
diff --git a/src/tensor_ops/to_dtype/cpu_kernel.rs b/src/tensor_ops/to_dtype/cpu_kernel.rs
new file mode 100644
index 000000000..03f5ac467
--- /dev/null
+++ b/src/tensor_ops/to_dtype/cpu_kernel.rs
@@ -0,0 +1,17 @@
+use num_traits::AsPrimitive;
+use std::sync::Arc;
+
+use crate::prelude::{Cpu, Shape, Tensor, Unit};
+
+impl<E1: Unit + AsPrimitive<E2>, E2: Unit> super::ToDtypeKernel<E1, E2> for Cpu {
+    fn forward<S: Shape>(inp: Tensor<S, E1, Self>) -> Result<Tensor<S, E2, Self>, Self::Err> {
+        Ok(Tensor {
+            id: crate::prelude::unique_id(),
+            data: Arc::new(inp.data.iter().map(|x| (*x).as_()).collect()),
+            shape: inp.shape,
+            strides: inp.strides,
+            device: inp.device.clone(),
+            tape: inp.tape,
+        })
+    }
+}
diff --git a/src/tensor_ops/to_dtype/cuda_kernel.rs b/src/tensor_ops/to_dtype/cuda_kernel.rs
new file mode 100644
index 000000000..d91cd8c1c
--- /dev/null
+++ b/src/tensor_ops/to_dtype/cuda_kernel.rs
@@ -0,0 +1,88 @@
+use crate::{
+    shapes::{Shape, Unit},
+    tensor::{cuda::Cuda, unique_id, Tensor},
+};
+use cudarc::driver::{DeviceSlice, LaunchAsync, LaunchConfig};
+use std::{sync::Arc, vec::Vec};
+
+const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/to_dtype.ptx"));
+const MODULE_NAME: &str = "to_dtype";
+
+// Minimize the amount of code in the macro to speed up compilation
+trait DtypeConversion {
+    const FUNC_NAME: &'static str;
+}
+
+macro_rules! conversion {
+    ($e1:ty, $e2:ty) => {
+        impl DtypeConversion for ($e1, $e2) {
+            const FUNC_NAME: &'static str = concat!(stringify!($e1), "_to_", stringify!($e2));
+        }
+    };
+}
+
+impl<E1: Unit, E2: Unit> super::ToDtypeKernel<E1, E2> for Cuda
+where
+    (E1, E2): DtypeConversion,
+{
+    fn forward<S: Shape>(inp: Tensor<S, E1, Self>) -> Result<Tensor<S, E2, Self>, Self::Err> {
+        let fn_name = <(E1, E2)>::FUNC_NAME;
+        let dev = &inp.device.dev;
+
+        if !dev.has_func(MODULE_NAME, fn_name) {
+            dev.load_ptx(PTX.into(), MODULE_NAME, &all_fn_names())?;
+        }
+        let numel = inp.data.len();
+        let mut out = unsafe { dev.alloc::<E2>(numel) }?;
+
+        let fwd_fn = dev.get_func(MODULE_NAME, fn_name).unwrap();
+        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let params = (numel, inp.data.as_ref(), &mut out);
+
+        unsafe { fwd_fn.launch(cfg, params) }?;
+
+        Ok(Tensor {
+            id: unique_id(),
+            data: Arc::new(out),
+            shape: inp.shape,
+            strides: inp.strides,
+            device: inp.device.clone(),
+            tape: Default::default(),
+        })
+    }
+}
+
+// Recursive macro that computes the names of all kernel functions and produces a function that
+// outputs them
+macro_rules! all_fn_names {
+    // perform "nested for loop" with the elements of from_types and to_types to produce the final
+    // list
+    (@ $([$from_types:ident, [$($to_types:ident),+]])+) => {
+        [$($(concat!(stringify!($from_types), "_to_", stringify!($to_types))),+),+]
+    };
+    // repeat to_types for each from_type
+    (@ [$($from_types:ident),+] $to_types:tt) => {
+        all_fn_names!(@ $([$from_types, $to_types])+)
+    };
+    ($from_types:tt $to_types:tt) => {
+        fn all_fn_names() -> Vec<&'static str> {
+            all_fn_names!(@ $from_types $to_types).to_vec()
+        }
+    }
+}
+
+// recursive macro to call "conversions" with all pairs of elements in from_types and to_types
+macro_rules! all_conversions {
+    (@ $from_type:ident [$($to_types:ident),*]) => {
+        $(conversion!($from_type, $to_types);)*
+    };
+    ([$($from_types:ident),+] $to_types:tt) => {
+        all_fn_names!([$($from_types),+] $to_types);
+        $(all_conversions!(@ $from_types $to_types);)+
+    }
+}
+
+all_conversions!(
+    [f32, f64, u8, u16, u32, u64, usize, i8, i16, i32, i64, isize, bool]
+    [f32, f64, u8, u16, u32, u64, usize, i8, i16, i32, i64, isize]
+);
diff --git a/src/tensor_ops/to_dtype/mod.rs b/src/tensor_ops/to_dtype/mod.rs
new file mode 100644
index 000000000..d6a6be466
--- /dev/null
+++ b/src/tensor_ops/to_dtype/mod.rs
@@ -0,0 +1,111 @@
+mod cpu_kernel;
+#[cfg(feature = "cuda")]
+mod cuda_kernel;
+
+use crate::prelude::{DeviceStorage, Shape, Tensor, Unit};
+
+pub trait ToDtypeKernel<E1: Unit, E2: Unit>: DeviceStorage {
+    fn forward<S: Shape>(inp: Tensor<S, E1, Self>) -> Result<Tensor<S, E2, Self>, Self::Err>;
+}
+
+/// Copies the elements of a tensor, converting its data to a different dtype.
+///
+/// Example usage:
+/// ```rust
+/// # use dfdx::prelude::*;
+/// # let dev: Cpu = Default::default();
+/// let a = dev.tensor([1usize, 2, 3, 4, 5]);
+///
+/// let b = a.clone().to_dtype::<i8>();
+/// assert_eq!(b.array(), [1, 2, 3, 4, 5]);
+///
+/// let c = a.clone().to_dtype::<f32>();
+/// assert_eq!(c.array(), [1.0, 2.0, 3.0, 4.0, 5.0]);
+/// ```
+pub fn to_dtype<E2: Unit, S: Shape, E1: Unit, D: ToDtypeKernel<E1, E2>>(
+    tensor: Tensor<S, E1, D>,
+) -> Tensor<S, E2, D> {
+    tensor.to_dtype()
+}
+
+impl<S: Shape, E: Unit, D: DeviceStorage> Tensor<S, E, D> {
+    pub fn try_to_dtype<E2: Unit>(self) -> Result<Tensor<S, E2, D>, D::Err>
+    where
+        D: ToDtypeKernel<E, E2>,
+    {
+        D::forward(self)
+    }
+
+    pub fn to_dtype<E2: Unit>(self) -> Tensor<S, E2, D>
+    where
+        D: ToDtypeKernel<E, E2>,
+    {
+        self.try_to_dtype().unwrap()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{prelude::Rank1, tensor::*, tests::*};
+
+    #[test]
+    fn test_to_dtype_unsigned() {
+        let dev: TestDevice = Default::default();
+        let a = dev.tensor_from_vec(
+            (0..128).map(|x| x as f32).collect(),
+            Rank1::<128>::default(),
+        );
+        let b = a.clone().to_dtype::<u8>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+        let b = a.clone().to_dtype::<u16>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+        let b = a.clone().to_dtype::<u32>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+        let b = a.clone().to_dtype::<u64>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+        let b = a.clone().to_dtype::<usize>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+
+        let a: Tensor<_, bool, _> = dev.tensor([true, true, false, true, false]);
+        let b = a.to_dtype::<usize>();
+        assert_eq!(b.array(), [1, 1, 0, 1, 0]);
+    }
+
+    #[test]
+    fn test_to_dtype_signed() {
+        let dev: TestDevice = Default::default();
+        let a = dev.tensor_from_vec(
+            (0..128).map(|x| x as f32).collect(),
+            Rank1::<128>::default(),
+        );
+        let b = a.clone().to_dtype::<i8>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+        let b = a.clone().to_dtype::<i16>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+        let b = a.clone().to_dtype::<i32>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+        let b = a.clone().to_dtype::<i64>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+        let b = a.clone().to_dtype::<isize>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+
+        let a: Tensor<_, bool, _> = dev.tensor([true, true, false, true, false]);
+        let b = a.to_dtype::<usize>();
+        assert_eq!(b.array(), [1, 1, 0, 1, 0]);
+    }
+
+    #[test]
+    fn test_to_dtype_other() {
+        let dev: TestDevice = Default::default();
+        let a = dev.tensor_from_vec(
+            (0..128).map(|x| x as f32).collect(),
+            Rank1::<128>::default(),
+        );
+        let b = a.clone().to_dtype::<f64>().to_dtype::<f32>();
+        assert_eq!(a.array(), b.array());
+
+        let a: Tensor<_, bool, _> = dev.tensor([true, true, false, true, false]);
+        let b = a.to_dtype::<usize>();
+        assert_eq!(b.array(), [1, 1, 0, 1, 0]);
+    }
+}
diff --git a/src/tensor_ops/to_dtype/to_dtype.cu b/src/tensor_ops/to_dtype/to_dtype.cu
new file mode 100644
index 000000000..86ee9d955
--- /dev/null
+++ b/src/tensor_ops/to_dtype/to_dtype.cu
@@ -0,0 +1,47 @@
+#include <cstdint>
+
+#define CONVERSION_OP_INTO(RTYPE1, RTYPE2, CTYPE1, CTYPE2) \
+extern "C" __global__ void RTYPE1##_to_##RTYPE2( \
+    const size_t numel, \
+    const CTYPE1 *inp, \
+    CTYPE2 *out \
+) { \
+    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; \
+    if (i >= numel) { \
+        return; \
+    } \
+    out[i] = inp[i]; \
+}
+
+#define CONVERSION_OP_FROM(RTYPE1, CTYPE1) \
+    CONVERSION_OP_INTO(RTYPE1, f32, CTYPE1, float); \
+    CONVERSION_OP_INTO(RTYPE1, f64, CTYPE1, double); \
+    \
+    CONVERSION_OP_INTO(RTYPE1, u8, CTYPE1, uint8_t); \
+    CONVERSION_OP_INTO(RTYPE1, u16, CTYPE1, uint16_t); \
+    CONVERSION_OP_INTO(RTYPE1, u32, CTYPE1, uint32_t); \
+    CONVERSION_OP_INTO(RTYPE1, u64, CTYPE1, uint64_t); \
+    CONVERSION_OP_INTO(RTYPE1, usize, CTYPE1, uintptr_t); \
+    \
+    CONVERSION_OP_INTO(RTYPE1, i8, CTYPE1, int8_t); \
+    CONVERSION_OP_INTO(RTYPE1, i16, CTYPE1, int16_t); \
+    CONVERSION_OP_INTO(RTYPE1, i32, CTYPE1, int32_t); \
+    CONVERSION_OP_INTO(RTYPE1, i64, CTYPE1, int64_t); \
+    CONVERSION_OP_INTO(RTYPE1, isize, CTYPE1, intptr_t); \
+
+CONVERSION_OP_FROM(f32, float);
+CONVERSION_OP_FROM(f64, double);
+
+CONVERSION_OP_FROM(u8, uint8_t);
+CONVERSION_OP_FROM(u16, uint16_t);
+CONVERSION_OP_FROM(u32, uint32_t);
+CONVERSION_OP_FROM(u64, uint64_t);
+CONVERSION_OP_FROM(usize, uintptr_t);
+
+CONVERSION_OP_FROM(i8, int8_t);
+CONVERSION_OP_FROM(i16, int16_t);
+CONVERSION_OP_FROM(i32, int32_t);
+CONVERSION_OP_FROM(i64, int64_t);
+CONVERSION_OP_FROM(isize, intptr_t);
+
+CONVERSION_OP_FROM(bool, bool);
diff --git a/src/tensor_ops/utilities/device.rs b/src/tensor_ops/utilities/device.rs
index 8d3b23dd3..d623159c2 100644
--- a/src/tensor_ops/utilities/device.rs
+++ b/src/tensor_ops/utilities/device.rs
@@ -78,6 +78,12 @@ pub trait Device<E: Dtype>:
     + UnaryKernel<super::super::pow::PowfKernelOp<E>, E>
     + UnaryKernel<super::super::pow::PowiKernelOp, E>
 
+    // to_dtype
+    + super::super::to_dtype::ToDtypeKernel<f32, E>
+    + super::super::to_dtype::ToDtypeKernel<f64, E>
+    + super::super::to_dtype::ToDtypeKernel<E, f32>
+    + super::super::to_dtype::ToDtypeKernel<E, f64>
+
     // binary
     + BinaryKernel<super::super::bce::BCEKernelOp, E>
     + BinaryKernel<super::super::huber_error::HuberErrorKernelOp<E>, E>


Test patch: