Repository: coreylowman/dfdx (language: rust)

Issue:
JIT compile concat & stack cuda kernels
Related to #630 

See PR #634 for how this is possible - basically use `cudarc::nvrtc::compile_ptx`, and `cudarc::types::CudaTypeName`


Discussion / hints:

Fix patch:
diff --git a/src/tensor_ops/concat/concat.cu b/src/tensor_ops/concat/concat.cu
deleted file mode 100644
index 8ab3ef2d7..000000000
--- a/src/tensor_ops/concat/concat.cu
+++ /dev/null
@@ -1,14 +0,0 @@
-extern "C" __global__ void concat_bwd_f32(const size_t numel, const float *inp, float *out) {
-    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    if (i >= numel) {
-        return;
-    }
-    out[i] += inp[i];
-}
-extern "C" __global__ void concat_bwd_f64(const size_t numel, const double *inp, double *out) {
-    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    if (i >= numel) {
-        return;
-    }
-    out[i] += inp[i];
-}
diff --git a/src/tensor_ops/concat/cuda_kernel.rs b/src/tensor_ops/concat/cuda_kernel.rs
index 87e76c183..d5a5c8208 100644
--- a/src/tensor_ops/concat/cuda_kernel.rs
+++ b/src/tensor_ops/concat/cuda_kernel.rs
@@ -2,24 +2,13 @@ use crate::{
     shapes::*,
     tensor::{launch_cfg, unique_id, Cuda, Tensor},
 };
-use cudarc::driver::{DeviceSlice, LaunchAsync};
-
-const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/concat.ptx"));
-
-trait HasCudaKernel<E> {
-    const BWD_FN: &'static str;
-}
-impl HasCudaKernel<f32> for Cuda {
-    const BWD_FN: &'static str = "concat_bwd_f32";
-}
-impl HasCudaKernel<f64> for Cuda {
-    const BWD_FN: &'static str = "concat_bwd_f64";
-}
+use cudarc::{
+    driver::{DeviceSlice, LaunchAsync},
+    nvrtc::{compile_ptx_with_opts, CompileOptions},
+    types::CudaTypeName,
+};
 
-impl<E: Dtype> super::ConcatKernel<E> for Cuda
-where
-    Self: HasCudaKernel<E>,
-{
+impl<E: Dtype + CudaTypeName> super::ConcatKernel<E> for Cuda {
     fn forward<A: Shape, B: Shape>(
         &self,
         a: &Tensor<A, E, Self>,
@@ -52,21 +41,27 @@ where
         grad_b: &mut Self::Vec<E>,
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err> {
-        if !self.dev.has_func(Self::BWD_FN, Self::BWD_FN) {
-            self.dev
-                .load_ptx(PTX.into(), Self::BWD_FN, &[Self::BWD_FN])?;
+        let module_name = std::format!("concat_bwd_{}", E::NAME);
+        if !self.dev.has_func(&module_name, "concat_bwd") {
+            let src = BWD_KERNEL.replace("$Ty", E::NAME);
+            let opts = CompileOptions {
+                arch: Some(env!("CUDA_COMPUTE_CAP")),
+                ..Default::default()
+            };
+            let ptx = compile_ptx_with_opts(src, opts).unwrap();
+            self.dev.load_ptx(ptx, &module_name, &["concat_bwd"])?;
         }
 
         let mut offset = 0;
         {
-            let f = self.dev.get_func(Self::BWD_FN, Self::BWD_FN).unwrap();
+            let f = self.dev.get_func(&module_name, "concat_bwd").unwrap();
             let numel = grad_a.len();
             let cfg = launch_cfg(numel as u32);
             unsafe { f.launch(cfg, (numel, &grad_out.slice(0..numel), grad_a)) }?;
             offset += numel;
         }
         {
-            let f = self.dev.get_func(Self::BWD_FN, Self::BWD_FN).unwrap();
+            let f = self.dev.get_func(&module_name, "concat_bwd").unwrap();
             let numel = grad_b.len();
             let cfg = launch_cfg(numel as u32);
             unsafe { f.launch(cfg, (numel, &grad_out.slice(offset..), grad_b)) }?;
@@ -74,3 +69,10 @@ where
         Ok(())
     }
 }
+
+const BWD_KERNEL: &str = "
+extern \"C\" __global__ void concat_bwd(const size_t numel, const $Ty *inp, $Ty *out) {
+    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < numel) { out[i] += inp[i]; }
+}
+";
diff --git a/src/tensor_ops/slice/cuda_kernel.rs b/src/tensor_ops/slice/cuda_kernel.rs
index eec114f15..5c328b98d 100644
--- a/src/tensor_ops/slice/cuda_kernel.rs
+++ b/src/tensor_ops/slice/cuda_kernel.rs
@@ -1,9 +1,9 @@
 use crate::{
     prelude::cpu::NdIndex,
     shapes::*,
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
 };
-use cudarc::driver::{CudaSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{CudaSlice, LaunchAsync};
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/slice.ptx"));
 
@@ -61,7 +61,7 @@ where
         let strides: CudaSlice<usize> = self.dev.htod_copy(strides.into())?;
 
         let fwd_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (
             numel,             // const size_t numel,
             Src::NUM_DIMS,     // const size_t num_dims,
@@ -97,7 +97,7 @@ where
         let strides: CudaSlice<usize> = self.dev.htod_copy(strides.into())?;
 
         let bwd_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (
             numel,         // const size_t numel,
             Src::NUM_DIMS, // const size_t num_dims,
diff --git a/src/tensor_ops/stack/cuda_kernel.rs b/src/tensor_ops/stack/cuda_kernel.rs
index dde7624db..a28833aa6 100644
--- a/src/tensor_ops/stack/cuda_kernel.rs
+++ b/src/tensor_ops/stack/cuda_kernel.rs
@@ -2,35 +2,14 @@ use crate::{
     shapes::*,
     tensor::{launch_cfg, Cuda, Tensor},
 };
-use cudarc::driver::{DeviceSlice, LaunchAsync};
+use cudarc::{
+    driver::{DeviceSlice, LaunchAsync},
+    nvrtc::{compile_ptx_with_opts, CompileOptions},
+    types::CudaTypeName,
+};
 use std::vec::Vec;
 
-const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/stack.ptx"));
-
-trait HasCudaKernel<E> {
-    const MOD: &'static str;
-    const FNS: &'static [&'static str];
-}
-
-impl HasCudaKernel<f32> for Cuda {
-    const MOD: &'static str = "stack_f32";
-    const FNS: &'static [&'static str] = &["sum_f32"];
-}
-
-impl HasCudaKernel<f64> for Cuda {
-    const MOD: &'static str = "stack_f64";
-    const FNS: &'static [&'static str] = &["sum_f64"];
-}
-
-impl HasCudaKernel<usize> for Cuda {
-    const MOD: &'static str = "stack_usize";
-    const FNS: &'static [&'static str] = &["sum_usize"];
-}
-
-impl<E: Dtype> super::StackKernel<E> for Cuda
-where
-    Self: HasCudaKernel<E>,
-{
+impl<E: Dtype + CudaTypeName> super::StackKernel<E> for Cuda {
     fn forward<S: Shape, Num: Dim>(
         &self,
         num: Num,
@@ -76,12 +55,20 @@ where
         mut grad_inp: Vec<&mut Self::Vec<E>>,
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err> {
-        if !self.dev.has_func(Self::MOD, Self::FNS[0]) {
-            self.dev.load_ptx(PTX.into(), Self::MOD, Self::FNS)?;
+        let module_name = std::format!("stack_bwd_{}", E::NAME);
+        if !self.dev.has_func(&module_name, "stack_bwd") {
+            let src = BWD_KERNEL.replace("$Ty", E::NAME);
+            let opts = CompileOptions {
+                arch: Some(env!("CUDA_COMPUTE_CAP")),
+                ..Default::default()
+            };
+            let ptx = compile_ptx_with_opts(src, opts).unwrap();
+            self.dev.load_ptx(ptx, &module_name, &["stack_bwd"])?;
         }
+
         let mut offset = 0;
         for item in grad_inp.drain(..) {
-            let f = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
+            let f = self.dev.get_func(&module_name, "stack_bwd").unwrap();
             let numel: usize = item.len();
             let cfg = launch_cfg(numel as u32);
             let sub = grad_out.slice(offset..offset + numel);
@@ -92,3 +79,10 @@ where
         Ok(())
     }
 }
+
+const BWD_KERNEL: &str = "
+extern \"C\" __global__ void stack_bwd(const size_t numel, const $Ty *inp, $Ty *out) {
+    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < numel) { out[i] += inp[i]; }
+}
+";
diff --git a/src/tensor_ops/stack/stack.cu b/src/tensor_ops/stack/stack.cu
deleted file mode 100644
index e9436ae92..000000000
--- a/src/tensor_ops/stack/stack.cu
+++ /dev/null
@@ -1,35 +0,0 @@
-extern "C" __global__ void sum_f32(
-    const size_t numel,
-    const float *inp,
-    float *out
-) {
-    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    if (i >= numel) {
-        return;
-    }
-    out[i] += inp[i];
-}
-
-extern "C" __global__ void sum_f64(
-    const size_t numel,
-    const double *inp,
-    double *out
-) {
-    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    if (i >= numel) {
-        return;
-    }
-    out[i] += inp[i];
-}
-
-extern "C" __global__ void sum_usize(
-    const size_t numel,
-    const size_t *inp,
-    size_t *out
-) {
-    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    if (i >= numel) {
-        return;
-    }
-    out[i] += inp[i];
-}


Test patch: