Repository: coreylowman/dfdx (language: rust)

Issue:
Figure out a better number of threads to launch kernels with
Apparently always using 1024 (the max) is not the best, because it could overload threads. This is what cudarc currently does.

For example doing a small change to using 128, I see a slight boost in speedup for a small resnet model.

What are ways we could improve this?

Additionally, could some kernels be improved by using 2d or 3d block/grid dims?


Discussion / hints:

Fix patch:
diff --git a/src/optim/adam/cuda_kernel.rs b/src/optim/adam/cuda_kernel.rs
index 3723fb429..15d239621 100644
--- a/src/optim/adam/cuda_kernel.rs
+++ b/src/optim/adam/cuda_kernel.rs
@@ -1,5 +1,10 @@
-use crate::{optim::optimizer::*, shapes::*, tensor::Cuda};
-use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync, LaunchConfig};
+use crate::{
+    optim::optimizer::*,
+    shapes::*,
+    tensor::{launch_cfg, Cuda},
+};
+
+use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync};
 
 #[repr(C)]
 struct CudaAdamConfig<E> {
@@ -63,7 +68,7 @@ where
         let opt_cfg = adam_config_to_cuda(cfg);
         let numel = param.len();
         let func = self.dev.get_func(Self::MOD, Self::FWD).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let t = <E>::from_i32(t).unwrap();
         let params = (opt_cfg, numel, t, param, moment1, moment2, grad);
         unsafe { func.launch(cfg, params) }?;
diff --git a/src/optim/rmsprop/cuda_kernel.rs b/src/optim/rmsprop/cuda_kernel.rs
index ae16197eb..6d6e18b43 100644
--- a/src/optim/rmsprop/cuda_kernel.rs
+++ b/src/optim/rmsprop/cuda_kernel.rs
@@ -1,6 +1,11 @@
 use super::RMSpropConfig;
-use crate::{optim::optimizer::*, shapes::*, tensor::Cuda};
-use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync, LaunchConfig};
+use crate::{
+    optim::optimizer::*,
+    shapes::*,
+    tensor::{launch_cfg, Cuda},
+};
+
+use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync};
 
 #[repr(C)]
 struct CudaRMSpropConfig<E> {
@@ -73,7 +78,7 @@ where
         let opt_cfg = rmsprop_config_to_cuda(cfg);
         let numel = param.len();
         let func = self.dev.get_func(Self::MOD, Self::FWD).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (opt_cfg, numel, param, momentum, square_avg, grad_avg, grad);
         unsafe { func.launch(cfg, params) }?;
         Ok(())
diff --git a/src/optim/sgd/cuda_kernel.rs b/src/optim/sgd/cuda_kernel.rs
index 81f5d34e8..9216f0870 100644
--- a/src/optim/sgd/cuda_kernel.rs
+++ b/src/optim/sgd/cuda_kernel.rs
@@ -1,6 +1,11 @@
 use super::SgdConfig;
-use crate::{optim::optimizer::*, shapes::*, tensor::Cuda};
-use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync, LaunchConfig};
+use crate::{
+    optim::optimizer::*,
+    shapes::*,
+    tensor::{launch_cfg, Cuda},
+};
+
+use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync};
 
 #[repr(C)]
 struct CudaSgdConfig<E> {
@@ -61,7 +66,7 @@ where
         let opt_cfg = sgd_config_to_cuda(cfg);
         let numel = param.len();
         let func = self.dev.get_func(Self::MOD, Self::FWD).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         unsafe { func.launch(cfg, (opt_cfg, numel, param, velocity, grad)) }?;
         Ok(())
     }
diff --git a/src/tensor/cuda/mod.rs b/src/tensor/cuda/mod.rs
index 9d9d1966e..db328591d 100644
--- a/src/tensor/cuda/mod.rs
+++ b/src/tensor/cuda/mod.rs
@@ -2,3 +2,13 @@ mod allocate;
 mod device;
 
 pub use device::{Cuda, CudaError};
+
+pub(crate) fn launch_cfg(n: u32) -> cudarc::driver::LaunchConfig {
+    const NUM_THREADS: u32 = 128;
+    let num_blocks = (n + NUM_THREADS - 1) / NUM_THREADS;
+    cudarc::driver::LaunchConfig {
+        grid_dim: (num_blocks, 1, 1),
+        block_dim: (NUM_THREADS, 1, 1),
+        shared_mem_bytes: 0,
+    }
+}
diff --git a/src/tensor/mod.rs b/src/tensor/mod.rs
index ff68b0338..5969e2645 100644
--- a/src/tensor/mod.rs
+++ b/src/tensor/mod.rs
@@ -143,6 +143,8 @@ pub use cpu::{Cpu, CpuError};
 #[cfg(not(feature = "cuda"))]
 pub type AutoDevice = Cpu;
 
+#[cfg(feature = "cuda")]
+pub(crate) use cuda::launch_cfg;
 #[cfg(feature = "cuda")]
 pub use cuda::{Cuda, CudaError};
 #[cfg(feature = "cuda")]
diff --git a/src/tensor_ops/attention_reshape/cuda_kernel.rs b/src/tensor_ops/attention_reshape/cuda_kernel.rs
index 63f372eb0..c0642498c 100644
--- a/src/tensor_ops/attention_reshape/cuda_kernel.rs
+++ b/src/tensor_ops/attention_reshape/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use super::*;
 use crate::tensor::cuda::Cuda;
-use cudarc::driver::{DeviceRepr, LaunchAsync, LaunchConfig};
+use cudarc::driver::{DeviceRepr, LaunchAsync};
 
 const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/attention_reshape.ptx"));
 
@@ -72,7 +72,7 @@ where
             sequence_length,
             past_length,
         };
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (
             op,
             qkv.data.as_ref(),
diff --git a/src/tensor_ops/axpy/cuda_kernel.rs b/src/tensor_ops/axpy/cuda_kernel.rs
index 9d97367df..6b2e900df 100644
--- a/src/tensor_ops/axpy/cuda_kernel.rs
+++ b/src/tensor_ops/axpy/cuda_kernel.rs
@@ -1,6 +1,9 @@
-use crate::{shapes::*, tensor::Cuda};
+use crate::{
+    shapes::*,
+    tensor::{launch_cfg, Cuda},
+};
 
-use cudarc::driver::{DeviceSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{DeviceSlice, LaunchAsync};
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/axpy.ptx"));
 
@@ -30,7 +33,7 @@ where
         }
         let numel = a.len();
         let fwd_fn = self.dev.get_func(Self::FN, Self::FN).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         unsafe { fwd_fn.launch(cfg, (numel, a, alpha, b, beta)) }?;
         Ok(())
     }
diff --git a/src/tensor_ops/boolean/cuda_kernels.rs b/src/tensor_ops/boolean/cuda_kernels.rs
index cdf8aa0d4..32da141da 100644
--- a/src/tensor_ops/boolean/cuda_kernels.rs
+++ b/src/tensor_ops/boolean/cuda_kernels.rs
@@ -1,10 +1,7 @@
 use super::BooleanKernel;
 use crate::{
     shapes::Shape,
-    tensor::{
-        cuda::{Cuda, CudaError},
-        Tensor,
-    },
+    tensor::{launch_cfg, Cuda, CudaError, Tensor},
 };
 use cudarc::driver::*;
 
@@ -35,7 +32,7 @@ impl Cuda {
         let rhs_strides: CudaSlice<usize> = self.dev.htod_copy(rhs.strides.into())?;
 
         let fwd_fn = self.dev.get_func(MODULE_NAME, fn_name).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (
             numel,             // const size_t numel,
             S::NUM_DIMS,       // const size_t num_dims,
@@ -65,7 +62,7 @@ impl BooleanKernel for Cuda {
         let mut storage = unsafe { self.dev.alloc(numel) }?;
 
         let fwd_fn = self.dev.get_func(MODULE_NAME, "boolean_not").unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (
             numel,             // const size_t numel,
             inp.data.as_ref(), // const bool *inp,
diff --git a/src/tensor_ops/choose/cuda_kernel.rs b/src/tensor_ops/choose/cuda_kernel.rs
index f7a2bca5e..d8999bdf6 100644
--- a/src/tensor_ops/choose/cuda_kernel.rs
+++ b/src/tensor_ops/choose/cuda_kernel.rs
@@ -1,8 +1,8 @@
 use crate::{
     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"), "/choose.ptx"));
 
@@ -47,7 +47,7 @@ where
         let rhs_strides: CudaSlice<usize> = self.dev.htod_copy(rhs.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,
             S::NUM_DIMS,        // const size_t num_dims,
@@ -81,7 +81,7 @@ where
         let cond_strides: CudaSlice<usize> = self.dev.htod_copy(cond.strides.into())?;
         let rhs_strides: CudaSlice<usize> = self.dev.htod_copy(rhs.strides.into())?;
 
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (
             numel,              // const size_t numel,
             S::NUM_DIMS,        // const size_t num_dims,
diff --git a/src/tensor_ops/cmp/cuda_kernels.rs b/src/tensor_ops/cmp/cuda_kernels.rs
index f09f972a9..becf8e739 100644
--- a/src/tensor_ops/cmp/cuda_kernels.rs
+++ b/src/tensor_ops/cmp/cuda_kernels.rs
@@ -1,8 +1,8 @@
 use crate::{
     shapes::{Shape, Unit},
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
 };
-use cudarc::driver::{CudaSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{CudaSlice, LaunchAsync};
 
 use super::{
     CmpKernel, EqKernelOp, GeKernelOp, GtKernelOp, LeKernelOp, LtKernelOp, NeKernelOp,
@@ -56,7 +56,7 @@ impl<E: Unit, Op: CmpOpCudaKernel<E>> CmpKernel<Op, E> for Cuda {
         let out_strides: CudaSlice<usize> = self.dev.htod_copy(strides.into())?;
 
         let fwd_fn = self.dev.get_func(Op::MODULE_NAME, Op::FWD_FN_NAME).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (
             numel,             // const size_t numel,
             S::NUM_DIMS,       // const size_t num_dims,
@@ -95,7 +95,7 @@ impl<E: Unit, Op: ScalarCmpOpCudaKernel<E>> ScalarCmpKernel<Op, E> for Cuda {
         let out_strides: CudaSlice<usize> = self.dev.htod_copy(strides.into())?;
 
         let fwd_fn = self.dev.get_func(Op::MODULE_NAME, Op::FWD_FN_NAME).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (
             numel,             // const size_t numel,
             S::NUM_DIMS,       // const size_t num_dims,
diff --git a/src/tensor_ops/concat/cuda_kernel.rs b/src/tensor_ops/concat/cuda_kernel.rs
index 0eb416c93..995163522 100644
--- a/src/tensor_ops/concat/cuda_kernel.rs
+++ b/src/tensor_ops/concat/cuda_kernel.rs
@@ -1,8 +1,8 @@
 use crate::{
     shapes::*,
-    tensor::{unique_id, Cuda, Tensor},
+    tensor::{launch_cfg, unique_id, Cuda, Tensor},
 };
-use cudarc::driver::{DeviceSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{DeviceSlice, LaunchAsync};
 
 const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/concat.ptx"));
 
@@ -67,14 +67,14 @@ where
         {
             let f = self.dev.get_func(Self::BWD_FN, Self::BWD_FN).unwrap();
             let numel = grad_a.len();
-            let cfg = LaunchConfig::for_num_elems(numel as u32);
+            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 numel = grad_b.len();
-            let cfg = LaunchConfig::for_num_elems(numel as u32);
+            let cfg = launch_cfg(numel as u32);
             unsafe { f.launch(cfg, (numel, &grad_out.slice(offset..), grad_b)) }?;
         }
         Ok(())
diff --git a/src/tensor_ops/conv2d/cuda_kernel.rs b/src/tensor_ops/conv2d/cuda_kernel.rs
index 83b9f90aa..a1eb4b885 100644
--- a/src/tensor_ops/conv2d/cuda_kernel.rs
+++ b/src/tensor_ops/conv2d/cuda_kernel.rs
@@ -1,10 +1,10 @@
 use cudarc::cublas::{CudaBlas, Gemm};
-use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync, LaunchConfig, ValidAsZeroBits};
+use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync, ValidAsZeroBits};
 
 use crate::tensor_ops::matmul::cuda_kernel::sgemm_batch;
 use crate::{
     shapes::*,
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
 };
 
 use std::sync::Arc;
@@ -66,7 +66,7 @@ where
         let mut patches = self.dev.alloc_zeros::<E>(patches_numel)?;
         let img_strides = self.dev.htod_copy(make_4d::<L>(lhs.strides).into())?;
         let unfold_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
-        let cfg = LaunchConfig::for_num_elems(patches.len() as u32);
+        let cfg = launch_cfg(patches.len() as u32);
         let params = (op, lhs.data.as_ref(), &img_strides, &mut patches);
         unsafe { unfold_fn.launch(cfg, params) }?;
 
@@ -108,7 +108,7 @@ where
         {
             // unfold grad_out into patches
             let unfold_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
-            let cfg = LaunchConfig::for_num_elems(patches_numel as u32);
+            let cfg = launch_cfg(patches_numel as u32);
             unsafe { unfold_fn.launch(cfg, (op, grad_out, &mut patches)) }?;
         }
 
@@ -121,7 +121,7 @@ where
             // prepare filters for backward operations by
             // swapping dims 0 and 1 and adding a batch dimension
             let tr_fn = self.dev.get_func(Self::MOD, Self::FNS[2]).unwrap();
-            let cfg = LaunchConfig::for_num_elems(rhs.shape.num_elements() as u32);
+            let cfg = launch_cfg(rhs.shape.num_elements() as u32);
             unsafe { tr_fn.launch(cfg, (op, rhs.data.as_ref(), &f_strides, &mut f_b1023)) }?;
         }
 
@@ -171,7 +171,7 @@ where
             // sum all the gradients collected in our broadcasted grad_f
             // into grad_rhs
             let sum_fn = self.dev.get_func(Self::MOD, Self::FNS[3]).unwrap();
-            let cfg = LaunchConfig::for_num_elems(rhs.shape.num_elements() as u32);
+            let cfg = launch_cfg(rhs.shape.num_elements() as u32);
             unsafe { sum_fn.launch(cfg, (op, &grad_f_b1023, grad_rhs, &f_strides)) }?;
         }
 
diff --git a/src/tensor_ops/dropout/cuda_kernel.rs b/src/tensor_ops/dropout/cuda_kernel.rs
index dada0bb5b..b99b35ae7 100644
--- a/src/tensor_ops/dropout/cuda_kernel.rs
+++ b/src/tensor_ops/dropout/cuda_kernel.rs
@@ -1,11 +1,11 @@
 use crate::{
     shapes::*,
-    tensor::{cuda::Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
 };
 
 use std::vec::Vec;
 
-use cudarc::driver::{DeviceSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{DeviceSlice, LaunchAsync};
 
 use rand::{rngs::StdRng, Rng, SeedableRng};
 use rand_distr::{Distribution, Standard};
@@ -52,7 +52,7 @@ where
         let mut storage = unsafe { self.dev.alloc::<E>(numel) }?;
 
         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 = (op.prob, numel, inp.data.as_ref(), &noise, &mut storage);
         unsafe { fwd_fn.launch(cfg, params) }?;
         Ok(self.build_tensor(inp.shape, inp.strides, storage))
@@ -72,7 +72,7 @@ where
         }?;
         let bwd_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
         let numel = inp.data.len();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (op.prob, numel, &noise, grad_inp, grad_out);
         unsafe { bwd_fn.launch(cfg, params) }?;
         Ok(())
diff --git a/src/tensor_ops/max_to/cuda_kernel.rs b/src/tensor_ops/max_to/cuda_kernel.rs
index ec63db27f..48341c9b5 100644
--- a/src/tensor_ops/max_to/cuda_kernel.rs
+++ b/src/tensor_ops/max_to/cuda_kernel.rs
@@ -1,10 +1,10 @@
 use crate::{
     shapes::*,
-    tensor::{cuda::Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
     tensor_ops::reduction_utils::*,
 };
 
-use cudarc::driver::{CudaSlice, DeviceSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{CudaSlice, DeviceSlice, LaunchAsync};
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/max_to.ptx"));
 
@@ -46,7 +46,7 @@ where
         let mut storage = unsafe {
             let mut storage = self.dev.alloc::<E>(dst.num_elements())?;
             fill_fn.launch(
-                LaunchConfig::for_num_elems(dst.num_elements() as u32),
+                launch_cfg(dst.num_elements() as u32),
                 (&mut storage, Self::INIT, dst.num_elements()),
             )?;
             storage
@@ -63,7 +63,7 @@ where
             reduction_output_strides::<Ax, Src, Dst>(inp.strides, dst);
         let chunk_len = physical_numel / dst_physical_numel;
 
-        let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
+        let cfg = launch_cfg(physical_numel as u32);
         let params = (
             physical_numel,    // const size_t numel,
             dims.len(),        // const size_t num_dims,
@@ -103,7 +103,7 @@ where
         ))
         .unwrap();
 
-        let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
+        let cfg = launch_cfg(physical_numel as u32);
         let params = (
             physical_numel,    // const size_t numel,
             Src::NUM_DIMS,     // const size_t num_dims,
diff --git a/src/tensor_ops/min_to/cuda_kernel.rs b/src/tensor_ops/min_to/cuda_kernel.rs
index 300b16c8e..7e67366bb 100644
--- a/src/tensor_ops/min_to/cuda_kernel.rs
+++ b/src/tensor_ops/min_to/cuda_kernel.rs
@@ -1,10 +1,10 @@
 use crate::{
     shapes::*,
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
     tensor_ops::reduction_utils::*,
 };
 
-use cudarc::driver::{CudaSlice, DeviceSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{CudaSlice, DeviceSlice, LaunchAsync};
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/min_to.ptx"));
 
@@ -46,7 +46,7 @@ where
         let mut storage = unsafe {
             let mut storage = self.dev.alloc::<E>(dst.num_elements())?;
             fill_fn.launch(
-                LaunchConfig::for_num_elems(dst.num_elements() as u32),
+                launch_cfg(dst.num_elements() as u32),
                 (&mut storage, Self::INIT, dst.num_elements()),
             )?;
             storage
@@ -63,7 +63,7 @@ where
             reduction_output_strides::<Ax, Src, Dst>(inp.strides, dst);
         let chunk_len = physical_numel / dst_physical_numel;
 
-        let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
+        let cfg = launch_cfg(physical_numel as u32);
         let params = (
             physical_numel,    // const size_t numel,
             dims.len(),        // const size_t num_dims,
@@ -103,7 +103,7 @@ where
         ))
         .unwrap();
 
-        let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
+        let cfg = launch_cfg(physical_numel as u32);
         let params = (
             physical_numel,    // const size_t numel,
             Src::NUM_DIMS,     // const size_t num_dims,
diff --git a/src/tensor_ops/pool2d/cuda_kernel.rs b/src/tensor_ops/pool2d/cuda_kernel.rs
index ae22c3611..4ede1488a 100644
--- a/src/tensor_ops/pool2d/cuda_kernel.rs
+++ b/src/tensor_ops/pool2d/cuda_kernel.rs
@@ -1,11 +1,11 @@
 use crate::{
     shapes::*,
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
 };
 
 use std::sync::Arc;
 
-use cudarc::driver::{DeviceRepr, LaunchAsync, LaunchConfig};
+use cudarc::driver::{DeviceRepr, LaunchAsync};
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/pool2d.ptx"));
 
@@ -35,7 +35,7 @@ macro_rules! pool_impl {
                 let inp_strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
                 let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides).into())?;
                 let fwd_fn = self.dev.get_func($Fwd, $Fwd).unwrap();
-                let cfg = LaunchConfig::for_num_elems(out.shape().num_elements() as u32);
+                let cfg = launch_cfg(out.shape().num_elements() as u32);
                 let params = (
                     op,                           // const Pool2dOp op,
                     &inp_strides,                 // const size_t *inp_strides,
@@ -57,7 +57,7 @@ macro_rules! pool_impl {
                 let inp_strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
                 let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides).into())?;
                 let bwd_fn = self.dev.get_func($Fwd, $Bwd).unwrap();
-                let cfg = LaunchConfig::for_num_elems(inp.shape().num_elements() as u32);
+                let cfg = launch_cfg(inp.shape().num_elements() as u32);
                 let params = (
                     op,                // const Pool2dOp op,
                     &inp_strides,      // const size_t *inp_strides,
diff --git a/src/tensor_ops/reshape_to/cuda_kernel.rs b/src/tensor_ops/reshape_to/cuda_kernel.rs
index 14d350e4e..29aa684ea 100644
--- a/src/tensor_ops/reshape_to/cuda_kernel.rs
+++ b/src/tensor_ops/reshape_to/cuda_kernel.rs
@@ -1,8 +1,8 @@
 use crate::{
     shapes::*,
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
 };
-use cudarc::driver::{DeviceSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{DeviceSlice, LaunchAsync};
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/reshape.ptx"));
 
@@ -43,7 +43,7 @@ where
         let dst_strides = self.dev.htod_copy(dst.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,
             inp.data.as_ref(), // const float *inp,
@@ -75,7 +75,7 @@ where
         let inp_strides = self.dev.htod_copy(inp.strides.into())?;
         let out_strides = self.dev.htod_copy(out.strides.into())?;
 
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (
             numel,         // const size_t numel,
             grad_inp,      // float *grad_inp,
diff --git a/src/tensor_ops/select_and_gather/cuda_kernel.rs b/src/tensor_ops/select_and_gather/cuda_kernel.rs
index 2aefe610c..6f95520d1 100644
--- a/src/tensor_ops/select_and_gather/cuda_kernel.rs
+++ b/src/tensor_ops/select_and_gather/cuda_kernel.rs
@@ -1,8 +1,8 @@
 use crate::{
     shapes::{RemoveDimTo, ReplaceDimTo, Shape},
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
 };
-use cudarc::driver::{DeviceSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{DeviceSlice, LaunchAsync};
 
 const GATHER_PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/gather.ptx"));
 const SELECT_PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/select.ptx"));
@@ -36,7 +36,7 @@ macro_rules! impl_cuda_kernels {
                 let idx_strides = self.dev.htod_copy(idx.strides.into())?;
 
                 let fwd_fn = self.dev.get_func($GatherMod, $GatherFwd).unwrap();
-                let cfg = LaunchConfig::for_num_elems(numel as u32);
+                let cfg = launch_cfg(numel as u32);
                 let params = (
                     numel,             // const size_t numel,
                     inp.data.as_ref(), // const float *inp,
@@ -74,7 +74,7 @@ macro_rules! impl_cuda_kernels {
                 let inp_strides = self.dev.htod_copy(inp.strides.into())?;
                 let idx_strides = self.dev.htod_copy(idx.strides.into())?;
 
-                let cfg = LaunchConfig::for_num_elems(numel as u32);
+                let cfg = launch_cfg(numel as u32);
                 let params = (
                     numel,             // const size_t numel,
                     grad_inp,          // float *grad_inp,
@@ -122,7 +122,7 @@ macro_rules! impl_cuda_kernels {
                 let dst_strides = self.dev.htod_copy(dst.strides().into())?;
 
                 let fwd_fn = self.dev.get_func($SelectMod, $SelectFwd).unwrap();
-                let cfg = LaunchConfig::for_num_elems(numel as u32);
+                let cfg = launch_cfg(numel as u32);
                 let params = (
                     numel,             // const size_t numel,
                     inp.data.as_ref(), // const float *inp,
@@ -163,7 +163,7 @@ macro_rules! impl_cuda_kernels {
                 let idx_strides = self.dev.htod_copy(idx.strides.into())?;
                 let out_strides = self.dev.htod_copy(out.strides.into())?;
 
-                let cfg = LaunchConfig::for_num_elems(numel as u32);
+                let cfg = launch_cfg(numel as u32);
                 let params = (
                     numel,             // const size_t numel,
                     grad_inp,          // float *grad_inp,
diff --git a/src/tensor_ops/stack/cuda_kernel.rs b/src/tensor_ops/stack/cuda_kernel.rs
index c3a760a1a..b0bb2d056 100644
--- a/src/tensor_ops/stack/cuda_kernel.rs
+++ b/src/tensor_ops/stack/cuda_kernel.rs
@@ -1,8 +1,8 @@
 use crate::{
     shapes::*,
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
 };
-use cudarc::driver::{DeviceSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{DeviceSlice, LaunchAsync};
 use std::vec::Vec;
 
 const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/stack.ptx"));
@@ -78,7 +78,7 @@ where
         for item in grad_inp.drain(..) {
             let f = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
             let numel: usize = item.len();
-            let cfg = LaunchConfig::for_num_elems(numel as u32);
+            let cfg = launch_cfg(numel as u32);
             let sub = grad_out.slice(offset..offset + numel);
             unsafe { f.launch(cfg, (numel, &sub, item)) }?;
             offset += numel;
diff --git a/src/tensor_ops/sum_to/cuda_kernel.rs b/src/tensor_ops/sum_to/cuda_kernel.rs
index a7e75a5cd..c9ec4f36b 100644
--- a/src/tensor_ops/sum_to/cuda_kernel.rs
+++ b/src/tensor_ops/sum_to/cuda_kernel.rs
@@ -1,12 +1,10 @@
 use crate::{
     shapes::*,
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
     tensor_ops::reduction_utils::*,
 };
 
-use cudarc::driver::{
-    CudaSlice, DeviceRepr, DeviceSlice, LaunchAsync, LaunchConfig, ValidAsZeroBits,
-};
+use cudarc::driver::{CudaSlice, DeviceRepr, DeviceSlice, LaunchAsync, ValidAsZeroBits};
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/sum_to.ptx"));
 
@@ -62,7 +60,7 @@ where
             reduction_output_strides::<Ax, Src, Dst>(inp.strides, dst);
         let chunk_len = physical_numel / dst_physical_numel;
 
-        let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
+        let cfg = launch_cfg(physical_numel as u32);
         let params = (
             physical_numel,    // const size_t numel,
             num_dims,          // const size_t num_dims,
@@ -103,7 +101,7 @@ where
         ))
         .unwrap();
 
-        let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
+        let cfg = launch_cfg(physical_numel as u32);
         let params = (
             physical_numel,   // const size_t numel,
             Src::NUM_DIMS,    // const size_t num_dims,
diff --git a/src/tensor_ops/to_dtype/cuda_kernel.rs b/src/tensor_ops/to_dtype/cuda_kernel.rs
index d91cd8c1c..4e74d0fac 100644
--- a/src/tensor_ops/to_dtype/cuda_kernel.rs
+++ b/src/tensor_ops/to_dtype/cuda_kernel.rs
@@ -1,8 +1,8 @@
 use crate::{
     shapes::{Shape, Unit},
-    tensor::{cuda::Cuda, unique_id, Tensor},
+    tensor::{cuda::Cuda, launch_cfg, unique_id, Tensor},
 };
-use cudarc::driver::{DeviceSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{DeviceSlice, LaunchAsync};
 use std::{sync::Arc, vec::Vec};
 
 const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/to_dtype.ptx"));
@@ -36,7 +36,7 @@ where
         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 cfg = launch_cfg(numel as u32);
         let params = (numel, inp.data.as_ref(), &mut out);
 
         unsafe { fwd_fn.launch(cfg, params) }?;
diff --git a/src/tensor_ops/utilities/cuda_kernels.rs b/src/tensor_ops/utilities/cuda_kernels.rs
index 10dfd233c..f1a180e64 100644
--- a/src/tensor_ops/utilities/cuda_kernels.rs
+++ b/src/tensor_ops/utilities/cuda_kernels.rs
@@ -1,9 +1,9 @@
 use crate::{
     shapes::{Dtype, Shape},
-    tensor::{cuda::Cuda, unique_id, Tensor},
+    tensor::{launch_cfg, unique_id, Cuda, Tensor},
     tensor_ops::ops::{BinaryKernel, UnaryKernel},
 };
-use cudarc::driver::{CudaSlice, DeviceRepr, DeviceSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{CudaSlice, DeviceRepr, DeviceSlice, LaunchAsync};
 use std::{sync::Arc, vec::Vec};
 
 pub trait UnaryOpCudaKernel<E> {
@@ -50,7 +50,7 @@ impl<E: Dtype, K: UnaryOpCudaKernel<E> + DeviceRepr> UnaryKernel<K, E> for Cuda
         let mut storage = unsafe { self.dev.alloc::<E>(numel) }?;
 
         let fwd_fn = self.dev.get_func(K::MODULE_NAME, K::FWD_FN_NAME).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (op, numel, inp.data.as_ref(), &mut storage);
         unsafe { fwd_fn.launch(cfg, params) }?;
 
@@ -73,7 +73,7 @@ impl<E: Dtype, K: UnaryOpCudaKernel<E> + DeviceRepr> UnaryKernel<K, E> for Cuda
     ) -> Result<(), Self::Err> {
         let bwd_fn = self.dev.get_func(K::MODULE_NAME, K::BWD_FN_NAME).unwrap();
         let numel = inp.data.len();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (op, numel, inp.data.as_ref(), grad_inp, grad_out);
         unsafe { bwd_fn.launch(cfg, params) }?;
         Ok(())
@@ -180,7 +180,7 @@ impl<E: Dtype, K: BinaryOpCudaKernel<E> + DeviceRepr + Clone> BinaryKernel<K, E>
         let rhs_strides: CudaSlice<usize> = self.dev.htod_copy(rhs.strides.into())?;
 
         let fwd_fn = self.dev.get_func(K::MODULE_NAME, K::FWD_FN_NAME).unwrap();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
         let params = (
             op,
             numel,             // const size_t numel,
@@ -225,7 +225,7 @@ impl<E: Dtype, K: BinaryOpCudaKernel<E> + DeviceRepr + Clone> BinaryKernel<K, E>
             .unwrap();
 
         let numel = lhs.shape.num_elements();
-        let cfg = LaunchConfig::for_num_elems(numel as u32);
+        let cfg = launch_cfg(numel as u32);
 
         let ((out_dims1, out_strides1), rhs_strides1) = permute_for_binary_backward(
             lhs.shape.concrete(),


Test patch: