Repository: coreylowman/dfdx (language: rust)

Issue:
Reduce cuda htod memcpy for binary operations
Currently the forward pass for binary operations does 3 memcpys, and the backward pass does 6. With small data, these copies can be even longer than the actual kernel itself. E.g. the below snapshot was taken from the backward pass while running 06-mnist:

![image](https://user-images.githubusercontent.com/7787278/227634145-5efd23b3-d457-4aff-ab07-743f1afa83f5.png)

The memcpys take ~25us, while the kernel takes ~7us.

One idea is to concatenate all the dims/strides together into a single vec, and then in the kernel just offset the pointer for whatever you need. The backward pass should only need 1 memcpy in this case.


Discussion / hints:
This should probably be done for anything with multiple memcpy, e.g. sum_to

Fix patch:
diff --git a/src/tensor_ops/max_to/cuda_kernel.rs b/src/tensor_ops/max_to/cuda_kernel.rs
index 48341c9b5..f7e29f5af 100644
--- a/src/tensor_ops/max_to/cuda_kernel.rs
+++ b/src/tensor_ops/max_to/cuda_kernel.rs
@@ -4,7 +4,7 @@ use crate::{
     tensor_ops::reduction_utils::*,
 };
 
-use cudarc::driver::{CudaSlice, DeviceSlice, LaunchAsync};
+use cudarc::driver::{DeviceSlice, LaunchAsync};
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/max_to.ptx"));
 
@@ -43,6 +43,7 @@ where
         }
 
         let fill_fn = self.dev.get_func(Self::MOD, Self::FNS[2]).unwrap();
+        let fwd_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
         let mut storage = unsafe {
             let mut storage = self.dev.alloc::<E>(dst.num_elements())?;
             fill_fn.launch(
@@ -52,11 +53,13 @@ where
             storage
         };
 
-        let fwd_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
-
         let (dims, strides) = permute_for_reductions::<_, Ax>(inp.shape.concrete(), inp.strides);
-        let dims: CudaSlice<usize> = self.dev.htod_copy(dims)?;
-        let strides: CudaSlice<usize> = self.dev.htod_copy(strides)?;
+        let num_dims = dims.len();
+
+        let mut info = Vec::with_capacity(num_dims * 2);
+        info.extend(dims);
+        info.extend(strides);
+        let info = self.dev.htod_copy(info)?;
 
         let physical_numel = inp.data.len();
         let (dst_physical_numel, dst_strides) =
@@ -64,13 +67,13 @@ where
         let chunk_len = physical_numel / dst_physical_numel;
 
         let cfg = launch_cfg(physical_numel as u32);
+
         let params = (
             physical_numel,    // const size_t numel,
-            dims.len(),        // const size_t num_dims,
+            num_dims,          // const size_t num_dims,
             chunk_len,         // const size_t chunk_len,
+            &info,             // const size_t *info,
             inp.data.as_ref(), // const float *inp,
-            &dims,             // const size_t *dims,
-            &strides,          // const size_t *strides,
             &mut storage,      // float *out
         );
         unsafe { fwd_fn.launch(cfg, params) }?;
@@ -89,12 +92,8 @@ where
     {
         let bwd_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
 
-        let dims: CudaSlice<usize> = self.dev.htod_copy(inp.shape.concrete().into())?;
-        let inp_strides: CudaSlice<usize> = self.dev.htod_copy(inp.strides.into())?;
         let out_strides: Src::Concrete =
             BroadcastStridesTo::<Src, Ax>::broadcast_strides(&out.shape, out.strides);
-        let out_strides: CudaSlice<usize> = self.dev.htod_copy(out_strides.into())?;
-
         let physical_numel = grad_inp.len();
         let elems_per_thread = E::from_usize(reduction_elems_per_thread::<_, Src>(
             inp.shape.concrete(),
@@ -104,17 +103,22 @@ where
         .unwrap();
 
         let cfg = launch_cfg(physical_numel as u32);
+
+        let mut info = Vec::with_capacity(Src::NUM_DIMS * 3);
+        info.extend(inp.shape.concrete());
+        info.extend(inp.strides);
+        info.extend(out_strides);
+        let info = self.dev.htod_copy(info)?;
+
         let params = (
             physical_numel,    // const size_t numel,
             Src::NUM_DIMS,     // const size_t num_dims,
             elems_per_thread,  // const float elems_per_thread,
-            &dims,             // const size_t *dims,
+            &info,             // const size_t *info,
             inp.data.as_ref(), // const float *inp,
             grad_inp,          // float *grad_inp,
-            &inp_strides,      // const size_t *inp_strides,
             out.data.as_ref(), // const float *out,
             grad_out,          // const float *grad_out,
-            &out_strides,      // const size_t *out_strides
         );
         unsafe { bwd_fn.launch(cfg, params) }?;
         Ok(())
diff --git a/src/tensor_ops/max_to/max_to.cu b/src/tensor_ops/max_to/max_to.cu
index b2c1578ba..846245422 100644
--- a/src/tensor_ops/max_to/max_to.cu
+++ b/src/tensor_ops/max_to/max_to.cu
@@ -70,9 +70,8 @@ __device__ void max_to_fwd(
     const size_t numel,
     const size_t num_dims,
     const size_t chunk_len,
+    const size_t *info,
     const T *inp,
-    const size_t *dims,
-    const size_t *strides,
     T *out
 ) {
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -81,6 +80,9 @@ __device__ void max_to_fwd(
         return;
     }
 
+    const size_t *dims = info;
+    const size_t *strides = info + num_dims;
+
     unsigned int inp_i = get_strided_index(i, num_dims, dims, strides);
     chunk_max(numel, chunk_len, inp[inp_i], out);
 }
@@ -92,13 +94,11 @@ __device__ void max_to_bwd(
     const size_t numel,
     const size_t num_dims,
     const T elems_per_thread,
-    const size_t *dims,
+    const size_t *info,
     const T *inp,
     T *grad_inp,
-    const size_t *inp_strides,
     const T *out,
-    const T *grad_out,
-    const size_t *out_strides
+    const T *grad_out
 ) {
     unsigned int inp_i = blockIdx.x * blockDim.x + threadIdx.x;
 
@@ -106,6 +106,10 @@ __device__ void max_to_bwd(
         return;
     }
 
+    const size_t *dims = info;
+    const size_t *inp_strides = info + num_dims;
+    const size_t *out_strides = info + 2 * num_dims;
+
     unsigned int out_i = restrided(inp_i, num_dims, dims, inp_strides, out_strides);
 
     const T mask = static_cast<T>(inp[inp_i] == out[out_i]);
@@ -117,26 +121,23 @@ extern "C" __global__ void FWD( \
     const size_t numel, \
     const size_t num_dims, \
     const size_t chunk_len, \
+    const size_t *info, \
     const TYPENAME *inp, \
-    const size_t *dims, \
-    const size_t *strides, \
     TYPENAME *out \
 ) { \
-    max_to_fwd(numel, num_dims, chunk_len, inp, dims, strides, out); \
+    max_to_fwd(numel, num_dims, chunk_len, info, inp, out); \
 } \
 extern "C" __global__ void BWD( \
     const size_t numel, \
     const size_t num_dims, \
     const TYPENAME elems_per_thread, \
-    const size_t *dims, \
+    const size_t *info, \
     const TYPENAME *inp, \
     TYPENAME *grad_inp, \
-    const size_t *inp_strides, \
     const TYPENAME *out, \
-    const TYPENAME *grad_out, \
-    const size_t *out_strides \
+    const TYPENAME *grad_out \
 ) { \
-    max_to_bwd(numel, num_dims, elems_per_thread, dims, inp, grad_inp, inp_strides, out, grad_out, out_strides); \
+    max_to_bwd(numel, num_dims, elems_per_thread, info, inp, grad_inp, out, grad_out); \
 }
 
 MAX(float, max_to_fwd_f32, max_to_bwd_f32);
diff --git a/src/tensor_ops/min_to/cuda_kernel.rs b/src/tensor_ops/min_to/cuda_kernel.rs
index 7e67366bb..7927056d7 100644
--- a/src/tensor_ops/min_to/cuda_kernel.rs
+++ b/src/tensor_ops/min_to/cuda_kernel.rs
@@ -4,7 +4,7 @@ use crate::{
     tensor_ops::reduction_utils::*,
 };
 
-use cudarc::driver::{CudaSlice, DeviceSlice, LaunchAsync};
+use cudarc::driver::{DeviceSlice, LaunchAsync};
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/min_to.ptx"));
 
@@ -43,6 +43,8 @@ where
         }
 
         let fill_fn = self.dev.get_func(Self::MOD, Self::FNS[2]).unwrap();
+        let fwd_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
+
         let mut storage = unsafe {
             let mut storage = self.dev.alloc::<E>(dst.num_elements())?;
             fill_fn.launch(
@@ -52,11 +54,13 @@ where
             storage
         };
 
-        let fwd_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
-
         let (dims, strides) = permute_for_reductions::<_, Ax>(inp.shape.concrete(), inp.strides);
-        let dims: CudaSlice<usize> = self.dev.htod_copy(dims)?;
-        let strides: CudaSlice<usize> = self.dev.htod_copy(strides)?;
+        let num_dims = dims.len();
+
+        let mut info = Vec::with_capacity(num_dims * 2);
+        info.extend(dims);
+        info.extend(strides);
+        let info = self.dev.htod_copy(info)?;
 
         let physical_numel = inp.data.len();
         let (dst_physical_numel, dst_strides) =
@@ -66,11 +70,10 @@ where
         let cfg = launch_cfg(physical_numel as u32);
         let params = (
             physical_numel,    // const size_t numel,
-            dims.len(),        // const size_t num_dims,
+            num_dims,          // const size_t num_dims,
             chunk_len,         // const size_t chunk_len,
+            &info,             // const size_t *info,
             inp.data.as_ref(), // const float *inp,
-            &dims,             // const size_t *dims,
-            &strides,          // const size_t *strides,
             &mut storage,      // float *out
         );
         unsafe { fwd_fn.launch(cfg, params) }?;
@@ -89,11 +92,8 @@ where
     {
         let bwd_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
 
-        let dims: CudaSlice<usize> = self.dev.htod_copy(inp.shape.concrete().into())?;
-        let inp_strides: CudaSlice<usize> = self.dev.htod_copy(inp.strides.into())?;
         let out_strides: Src::Concrete =
             BroadcastStridesTo::<Src, Ax>::broadcast_strides(&out.shape, out.strides);
-        let out_strides: CudaSlice<usize> = self.dev.htod_copy(out_strides.into())?;
 
         let physical_numel = grad_inp.len();
         let elems_per_thread = E::from_usize(reduction_elems_per_thread::<_, Src>(
@@ -104,17 +104,22 @@ where
         .unwrap();
 
         let cfg = launch_cfg(physical_numel as u32);
+
+        let mut info = Vec::with_capacity(Src::NUM_DIMS * 3);
+        info.extend(inp.shape.concrete());
+        info.extend(inp.strides);
+        info.extend(out_strides);
+        let info = self.dev.htod_copy(info)?;
+
         let params = (
             physical_numel,    // const size_t numel,
             Src::NUM_DIMS,     // const size_t num_dims,
             elems_per_thread,  // const float elems_per_thread,
-            &dims,             // const size_t *dims,
+            &info,             // const size_t *info,
             inp.data.as_ref(), // const float *inp,
             grad_inp,          // float *grad_inp,
-            &inp_strides,      // const size_t *inp_strides,
             out.data.as_ref(), // const float *out,
             grad_out,          // const float *grad_out,
-            &out_strides,      // const size_t *out_strides
         );
         unsafe { bwd_fn.launch(cfg, params) }?;
         Ok(())
diff --git a/src/tensor_ops/min_to/min_to.cu b/src/tensor_ops/min_to/min_to.cu
index 9e80a3d57..ffd860729 100644
--- a/src/tensor_ops/min_to/min_to.cu
+++ b/src/tensor_ops/min_to/min_to.cu
@@ -70,9 +70,8 @@ __device__ void min_to_fwd(
     const size_t numel,
     const size_t num_dims,
     const size_t chunk_len,
+    const size_t *info,
     const T *inp,
-    const size_t *dims,
-    const size_t *strides,
     T *out
 ) {
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -81,6 +80,9 @@ __device__ void min_to_fwd(
         return;
     }
 
+    const size_t *dims = info;
+    const size_t *strides = info + num_dims;
+
     unsigned int inp_i = get_strided_index(i, num_dims, dims, strides);
     chunk_min(numel, chunk_len, inp[inp_i], out);
 }
@@ -92,13 +94,11 @@ __device__ void min_to_bwd(
     const size_t numel,
     const size_t num_dims,
     const T elems_per_thread,
-    const size_t *dims,
+    const size_t *info,
     const T *inp,
     T *grad_inp,
-    const size_t *inp_strides,
     const T *out,
-    const T *grad_out,
-    const size_t *out_strides
+    const T *grad_out
 ) {
     unsigned int inp_i = blockIdx.x * blockDim.x + threadIdx.x;
 
@@ -106,6 +106,10 @@ __device__ void min_to_bwd(
         return;
     }
 
+    const size_t *dims = info;
+    const size_t *inp_strides = info + num_dims;
+    const size_t *out_strides = info + 2 * num_dims;
+
     unsigned int out_i = restrided(inp_i, num_dims, dims, inp_strides, out_strides);
 
     const T mask = static_cast<T>(inp[inp_i] == out[out_i]);
@@ -117,26 +121,23 @@ extern "C" __global__ void FWD( \
     const size_t numel, \
     const size_t num_dims, \
     const size_t chunk_len, \
+    const size_t *info, \
     const TYPENAME *inp, \
-    const size_t *dims, \
-    const size_t *strides, \
     TYPENAME *out \
 ) { \
-    min_to_fwd(numel, num_dims, chunk_len, inp, dims, strides, out); \
+    min_to_fwd(numel, num_dims, chunk_len, info, inp, out); \
 } \
 extern "C" __global__ void BWD( \
     const size_t numel, \
     const size_t num_dims, \
     const TYPENAME elems_per_thread, \
-    const size_t *dims, \
+    const size_t *info, \
     const TYPENAME *inp, \
     TYPENAME *grad_inp, \
-    const size_t *inp_strides, \
     const TYPENAME *out, \
-    const TYPENAME *grad_out, \
-    const size_t *out_strides \
+    const TYPENAME *grad_out \
 ) { \
-    min_to_bwd(numel, num_dims, elems_per_thread, dims, inp, grad_inp, inp_strides, out, grad_out, out_strides); \
+    min_to_bwd(numel, num_dims, elems_per_thread, info, inp, grad_inp, out, grad_out); \
 }
 
 MIN(float, min_to_fwd_f32, min_to_bwd_f32);
diff --git a/src/tensor_ops/reshape_to/cuda_kernel.rs b/src/tensor_ops/reshape_to/cuda_kernel.rs
index e14b68c53..e6f094eb0 100644
--- a/src/tensor_ops/reshape_to/cuda_kernel.rs
+++ b/src/tensor_ops/reshape_to/cuda_kernel.rs
@@ -50,23 +50,22 @@ where
         let numel = inp.shape.num_elements();
         let mut storage = unsafe { self.dev.alloc::<E>(numel) }?;
 
-        let inp_dims = self.dev.htod_copy(inp.shape.concrete().into())?;
-        let dst_dims = self.dev.htod_copy(dst.concrete().into())?;
-        let inp_strides = self.dev.htod_copy(inp.strides.into())?;
-        let dst_strides = self.dev.htod_copy(dst.strides().into())?;
+        let mut info = Vec::with_capacity(Src::NUM_DIMS * 2 + Dst::NUM_DIMS * 2);
+        info.extend(inp.shape.concrete());
+        info.extend(inp.strides);
+        info.extend(dst.concrete());
+        info.extend(dst.strides());
+        let info = self.dev.htod_copy(info)?;
 
         let fwd_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
         let cfg = launch_cfg(numel as u32);
         let params = (
             numel,             // const size_t numel,
-            inp.data.as_ref(), // const float *inp,
             Src::NUM_DIMS,     // const size_t inp_num_dims,
-            &inp_dims,         // const size_t *inp_dims,
-            &inp_strides,      // const size_t *inp_strides,
-            &mut storage,      // float *out
             Dst::NUM_DIMS,     // const size_t out_num_dims,
-            &dst_dims,         // const size_t *out_dims,
-            &dst_strides,      // const size_t *out_strides,
+            &info,             // const size_t *info,
+            inp.data.as_ref(), // const float *inp,
+            &mut storage,      // float *out
         );
         unsafe { fwd_fn.launch(cfg, params) }?;
 
@@ -83,22 +82,21 @@ where
         let bwd_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
         let numel = grad_inp.len();
 
-        let inp_dims = self.dev.htod_copy(inp.shape.concrete().into())?;
-        let out_dims = self.dev.htod_copy(out.shape.concrete().into())?;
-        let inp_strides = self.dev.htod_copy(inp.strides.into())?;
-        let out_strides = self.dev.htod_copy(out.strides.into())?;
+        let mut info = Vec::with_capacity(Src::NUM_DIMS * 2 + Dst::NUM_DIMS * 2);
+        info.extend(inp.shape.concrete());
+        info.extend(inp.strides);
+        info.extend(out.shape.concrete());
+        info.extend(out.strides);
+        let info = self.dev.htod_copy(info)?;
 
         let cfg = launch_cfg(numel as u32);
         let params = (
             numel,         // const size_t numel,
-            grad_inp,      // float *grad_inp,
             Src::NUM_DIMS, // const size_t inp_num_dims,
-            &inp_dims,     // const size_t *inp_dims,
-            &inp_strides,  // const size_t *inp_strides,
-            grad_out,      // const float *grad_out,
             Dst::NUM_DIMS, // const size_t out_num_dims,
-            &out_dims,     // const size_t *out_dims,
-            &out_strides,  // const size_t *out_strides
+            &info,         // const size_t *info,
+            grad_inp,      // float *grad_inp,
+            grad_out,      // const float *grad_out,
         );
         unsafe { bwd_fn.launch(cfg, params) }?;
         Ok(())
diff --git a/src/tensor_ops/reshape_to/reshape.cu b/src/tensor_ops/reshape_to/reshape.cu
index 9a4380d6d..df85b2e1a 100644
--- a/src/tensor_ops/reshape_to/reshape.cu
+++ b/src/tensor_ops/reshape_to/reshape.cu
@@ -4,20 +4,22 @@
 template<typename T>
 __device__ void reshape_fwd(
     const size_t numel,
-    const T *inp,
     const size_t inp_num_dims,
-    const size_t *inp_dims,
-    const size_t *inp_strides,
-    T *out,
     const size_t out_num_dims,
-    const size_t *out_dims,
-    const size_t *out_strides
+    const size_t *info,
+    const T *inp,
+    T *out
 ) {
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i >= numel) {
         return;
     }
 
+    const size_t *inp_dims = info;
+    const size_t *inp_strides = info + inp_num_dims;
+    const size_t *out_dims = info + 2 * inp_num_dims;
+    const size_t *out_strides = info + 2 * inp_num_dims + out_num_dims;
+
     unsigned int inp_i = get_strided_index(i, inp_num_dims, inp_dims, inp_strides);
     unsigned int out_i = get_strided_index(i, out_num_dims, out_dims, out_strides);
 
@@ -27,20 +29,22 @@ __device__ void reshape_fwd(
 template<typename T>
 __device__ void reshape_bwd(
     const size_t numel,
-    T *grad_inp,
     const size_t inp_num_dims,
-    const size_t *inp_dims,
-    const size_t *inp_strides,
-    const T *grad_out,
     const size_t out_num_dims,
-    const size_t *out_dims,
-    const size_t *out_strides
+    const size_t *info,
+    T *grad_inp,
+    const T *grad_out
 ) {
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i >= numel) {
         return;
     }
 
+    const size_t *inp_dims = info;
+    const size_t *inp_strides = info + inp_num_dims;
+    const size_t *out_dims = info + 2 * inp_num_dims;
+    const size_t *out_strides = info + 2 * inp_num_dims + out_num_dims;
+
     unsigned int inp_i = get_strided_index(i, inp_num_dims, inp_dims, inp_strides);
     unsigned int out_i = get_strided_index(i, out_num_dims, out_dims, out_strides);
 
@@ -50,16 +54,13 @@ __device__ void reshape_bwd(
 #define RESHAPE_FWD(TYPENAME, FN) \
 extern "C" __global__ void FN( \
     const size_t numel, \
-    const TYPENAME *inp, \
     const size_t inp_num_dims, \
-    const size_t *inp_dims, \
-    const size_t *inp_strides, \
-    TYPENAME *out, \
     const size_t out_num_dims, \
-    const size_t *out_dims, \
-    const size_t *out_strides \
+    const size_t *info, \
+    const TYPENAME *inp, \
+    TYPENAME *out \
 ) { \
-    reshape_fwd(numel, inp, inp_num_dims, inp_dims, inp_strides, out, out_num_dims, out_dims, out_strides); \
+    reshape_fwd(numel, inp_num_dims, out_num_dims, info, inp, out); \
 }
 
 #define RESHAPE(TYPENAME, FWD, BWD) \
@@ -67,16 +68,13 @@ RESHAPE_FWD(TYPENAME, FWD) \
 \
 extern "C" __global__ void BWD( \
     const size_t numel, \
-    TYPENAME *grad_inp, \
     const size_t inp_num_dims, \
-    const size_t *inp_dims, \
-    const size_t *inp_strides, \
-    const TYPENAME *grad_out, \
     const size_t out_num_dims, \
-    const size_t *out_dims, \
-    const size_t *out_strides \
+    const size_t *info, \
+    TYPENAME *grad_inp, \
+    const TYPENAME *grad_out \
 ) { \
-    reshape_bwd(numel, grad_inp, inp_num_dims, inp_dims, inp_strides, grad_out, out_num_dims, out_dims, out_strides); \
+    reshape_bwd(numel, inp_num_dims, out_num_dims, info, grad_inp, grad_out); \
 }
 
 RESHAPE(float, reshape_fwd_f32, reshape_bwd_f32);
diff --git a/src/tensor_ops/sum_to/cuda_kernel.rs b/src/tensor_ops/sum_to/cuda_kernel.rs
index c9ec4f36b..abaf5f466 100644
--- a/src/tensor_ops/sum_to/cuda_kernel.rs
+++ b/src/tensor_ops/sum_to/cuda_kernel.rs
@@ -4,7 +4,7 @@ use crate::{
     tensor_ops::reduction_utils::*,
 };
 
-use cudarc::driver::{CudaSlice, DeviceRepr, DeviceSlice, LaunchAsync, ValidAsZeroBits};
+use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync, ValidAsZeroBits};
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/sum_to.ptx"));
 
@@ -43,10 +43,11 @@ where
 
         let (dims, strides) = permute_for_reductions::<_, Ax>(inp.shape.concrete(), inp.strides);
         let num_dims = dims.len();
-        let dims: CudaSlice<usize> = self.dev.htod_copy(dims)?;
-        let strides: CudaSlice<usize> = self.dev.htod_copy(strides)?;
 
-        let mut storage = self.dev.alloc_zeros::<E>(dst.num_elements())?;
+        let mut info = Vec::with_capacity(num_dims * 2);
+        info.extend(dims);
+        info.extend(strides);
+        let info = self.dev.htod_copy(info)?;
 
         let elems_per_thread = E::from_usize(reduction_elems_per_thread::<_, Src>(
             inp.shape.concrete(),
@@ -61,14 +62,15 @@ where
         let chunk_len = physical_numel / dst_physical_numel;
 
         let cfg = launch_cfg(physical_numel as u32);
+
+        let mut storage = self.dev.alloc_zeros::<E>(dst.num_elements())?;
         let params = (
             physical_numel,    // const size_t numel,
             num_dims,          // const size_t num_dims,
             elems_per_thread,  // const float elems_per_thread,
             chunk_len,         // const size_t chunk_len,
+            &info,             // const size_t *info,
             inp.data.as_ref(), // const float *inp,
-            &dims,             // const size_t *dims,
-            &strides,          // const size_t *strides,
             &mut storage,      // float *out
         );
         unsafe { fwd_fn.launch(cfg, params) }?;
@@ -87,12 +89,8 @@ where
     {
         let bwd_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
 
-        let dims: CudaSlice<usize> = self.dev.htod_copy(inp.shape.concrete().into())?;
-        let inp_strides: CudaSlice<usize> = self.dev.htod_copy(inp.strides.into())?;
         let out_strides: Src::Concrete =
             BroadcastStridesTo::<Src, Ax>::broadcast_strides(&out.shape, out.strides);
-        let out_strides: CudaSlice<usize> = self.dev.htod_copy(out_strides.into())?;
-
         let physical_numel = inp.data.len();
         let elems_per_thread = E::from_usize(reduction_elems_per_thread::<_, Src>(
             inp.shape.concrete(),
@@ -102,15 +100,20 @@ where
         .unwrap();
 
         let cfg = launch_cfg(physical_numel as u32);
+
+        let mut info: Vec<usize> = Vec::with_capacity(3 * Src::NUM_DIMS);
+        info.extend(inp.shape.concrete());
+        info.extend(inp.strides);
+        info.extend(out_strides);
+        let info = self.dev.htod_copy(info)?;
+
         let params = (
             physical_numel,   // const size_t numel,
             Src::NUM_DIMS,    // const size_t num_dims,
             elems_per_thread, // const float elems_per_thread,
-            &dims,            // const size_t *dims,
+            &info,            // const size_t *info,
             grad_inp,         // float *grad_inp,
-            &inp_strides,     // const size_t *inp_strides,
             grad_out,         // const float *grad_out,
-            &out_strides,     // const size_t *out_strides
         );
         unsafe { bwd_fn.launch(cfg, params) }?;
         Ok(())
diff --git a/src/tensor_ops/sum_to/sum_to.cu b/src/tensor_ops/sum_to/sum_to.cu
index 3562e7707..776afa37d 100644
--- a/src/tensor_ops/sum_to/sum_to.cu
+++ b/src/tensor_ops/sum_to/sum_to.cu
@@ -8,9 +8,8 @@ __device__ void sum_to_fwd(
     const size_t num_dims,
     const T elems_per_thread,
     const size_t chunk_len,
+    const size_t *info,
     const T *inp,
-    const size_t *dims,
-    const size_t *strides,
     T *out
 ) {
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -19,6 +18,9 @@ __device__ void sum_to_fwd(
         return;
     }
 
+    const size_t *dims = info;
+    const size_t *strides = info + num_dims;
+
     unsigned int inp_i = get_strided_index(i, num_dims, dims, strides);
     chunk_sum(chunk_len, inp[inp_i] * elems_per_thread, out);
 }
@@ -30,11 +32,9 @@ __device__ void sum_to_bwd(
     const size_t numel,
     const size_t num_dims,
     const T elems_per_thread,
-    const size_t *dims,
+    const size_t *info,
     T *grad_inp,
-    const size_t *inp_strides,
-    const T *grad_out,
-    const size_t *out_strides
+    const T *grad_out
 ) {
     unsigned int inp_i = blockIdx.x * blockDim.x + threadIdx.x;
 
@@ -42,6 +42,10 @@ __device__ void sum_to_bwd(
         return;
     }
 
+    const size_t *dims = info;
+    const size_t *inp_strides = info + num_dims;
+    const size_t *out_strides = info + 2 * num_dims;
+
     unsigned int out_i = restrided(inp_i, num_dims, dims, inp_strides, out_strides);
 
     // NOTE: since size of output is less than input, only 1 thread will be writing to inp
@@ -56,24 +60,21 @@ extern "C" __global__ void FWD( \
     const size_t num_dims, \
     const TYPENAME elems_per_thread, \
     const size_t chunk_len, \
+    const size_t *info, \
     const TYPENAME *inp, \
-    const size_t *dims, \
-    const size_t *strides, \
     TYPENAME *out \
 ) { \
-    sum_to_fwd(numel, num_dims, elems_per_thread, chunk_len, inp, dims, strides, out); \
+    sum_to_fwd(numel, num_dims, elems_per_thread, chunk_len, info, inp, out); \
 } \
 extern "C" __global__ void BWD( \
     const size_t numel, \
     const size_t num_dims, \
     const TYPENAME elems_per_thread, \
-    const size_t *dims, \
+    const size_t *info, \
     TYPENAME *grad_inp, \
-    const size_t *inp_strides, \
-    const TYPENAME *grad_out, \
-    const size_t *out_strides \
+    const TYPENAME *grad_out \
 ) { \
-    sum_to_bwd(numel, num_dims, elems_per_thread, dims, grad_inp, inp_strides, grad_out, out_strides); \
+    sum_to_bwd(numel, num_dims, elems_per_thread, info, grad_inp, grad_out); \
 }
 
 SUM(float, sum_to_fwd_f32, sum_to_bwd_f32);
diff --git a/src/tensor_ops/utilities/binary_op_macros.cuh b/src/tensor_ops/utilities/binary_op_macros.cuh
index 132e0e280..dc32dd3eb 100644
--- a/src/tensor_ops/utilities/binary_op_macros.cuh
+++ b/src/tensor_ops/utilities/binary_op_macros.cuh
@@ -5,17 +5,19 @@ extern "C" __global__ void FORWARD( \
     const OP_STRUCT op, \
     const size_t numel, \
     const size_t num_dims, \
-    const size_t *dims, \
+    const size_t *info, \
     const TYPENAME *lhs, \
-    const size_t *lhs_strides, \
     const TYPENAME *rhs, \
-    const size_t *rhs_strides, \
     TYPENAME *out \
 ) { \
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; \
     if (i >= numel) { \
         return; \
     } \
+\
+    const size_t *dims = info; \
+    const size_t *lhs_strides = info + num_dims; \
+    const size_t *rhs_strides = info + 2 * num_dims; \
 \
     unsigned int tmp_i = i; \
     unsigned int lhs_i = 0; \
@@ -40,19 +42,21 @@ extern "C" __global__ void BACKWARD_LHS( \
     const OP_STRUCT op, \
     const size_t numel, \
     const size_t num_dims, \
-    const size_t *dims, \
-    const size_t *out_strides, \
+    const size_t *info, \
     const TYPENAME *lhs, \
     TYPENAME *grad_lhs, \
     const size_t chunk_len, \
     const TYPENAME *rhs, \
-    const size_t *rhs_strides, \
     const TYPENAME *grad_out \
 ) { \
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; \
     if (i >= numel) { \
         return; \
     } \
+\
+    const size_t *dims = info + 0 * num_dims; \
+    const size_t *out_strides = info + 1 * num_dims; \
+    const size_t *rhs_strides = info + 2 * num_dims; \
 \
     unsigned int tmp_i = i; \
     unsigned int out_i = 0; \
@@ -77,10 +81,8 @@ extern "C" __global__ void BACKWARD_RHS( \
     const OP_STRUCT op, \
     const size_t numel, \
     const size_t num_dims, \
-    const size_t *dims, \
-    const size_t *out_strides, \
+    const size_t *info, \
     const TYPENAME *lhs, \
-    const size_t *lhs_strides, \
     const TYPENAME *rhs, \
     TYPENAME *grad_rhs, \
     const size_t chunk_len, \
@@ -90,6 +92,9 @@ extern "C" __global__ void BACKWARD_RHS( \
     if (i >= numel) { \
         return; \
     } \
+    const size_t *dims = info + 3 * num_dims; \
+    const size_t *out_strides = info + 4 * num_dims; \
+    const size_t *lhs_strides = info + 5 * num_dims; \
 \
     unsigned int tmp_i = i; \
     unsigned int lhs_i = 0; \
diff --git a/src/tensor_ops/utilities/cuda_kernels.rs b/src/tensor_ops/utilities/cuda_kernels.rs
index 16f702397..df2968115 100644
--- a/src/tensor_ops/utilities/cuda_kernels.rs
+++ b/src/tensor_ops/utilities/cuda_kernels.rs
@@ -3,7 +3,7 @@ use crate::{
     tensor::{launch_cfg, unique_id, Cuda, Tensor},
     tensor_ops::ops::{BinaryKernel, UnaryKernel},
 };
-use cudarc::driver::{CudaSlice, DeviceRepr, DeviceSlice, LaunchAsync};
+use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync};
 use std::{sync::Arc, vec::Vec};
 
 pub trait UnaryOpCudaKernel<E> {
@@ -148,13 +148,6 @@ where
     tmp.into_iter().map(|(_ord, x)| x).unzip()
 }
 
-fn physical_numel<I: IntoIterator<Item = usize>>(dims: I, strides: I) -> usize {
-    dims.into_iter()
-        .zip(strides.into_iter())
-        .map(|(dim, stride)| if stride == 0 { 1 } else { dim })
-        .product()
-}
-
 pub(crate) use cuda_binary;
 
 impl<E: Dtype, K: BinaryOpCudaKernel<E> + DeviceRepr + Clone> BinaryKernel<K, E> for Cuda {
@@ -173,23 +166,23 @@ impl<E: Dtype, K: BinaryOpCudaKernel<E> + DeviceRepr + Clone> BinaryKernel<K, E>
         let strides = lhs.shape.strides();
         let numel = shape.num_elements();
 
-        let mut storage = unsafe { self.dev.alloc::<E>(numel) }?;
-
-        let dims: CudaSlice<usize> = self.dev.htod_copy(shape.concrete().into())?;
-        let lhs_strides: CudaSlice<usize> = self.dev.htod_copy(lhs.strides.into())?;
-        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 = launch_cfg(numel as u32);
+
+        let mut info: Vec<usize> = Vec::with_capacity(3 * S::NUM_DIMS);
+        info.extend(shape.concrete());
+        info.extend(lhs.strides);
+        info.extend(rhs.strides);
+        let info = self.dev.htod_copy(info)?;
+
+        let mut storage = unsafe { self.dev.alloc::<E>(numel) }?;
         let params = (
             op,
             numel,             // const size_t numel,
             S::NUM_DIMS,       // const size_t num_dims,
-            &dims,             // const size_t *dims,
+            &info,             // const size_t *info,
             lhs.data.as_ref(), // const float *lhs,
-            &lhs_strides,      // const size_t *lhs_strides,
             rhs.data.as_ref(), // const float *rhs,
-            &rhs_strides,      // const size_t *rhs_strides,
             &mut storage,      // float *out,
         );
         unsafe { fwd_fn.launch(cfg, params) }?;
@@ -241,45 +234,41 @@ impl<E: Dtype, K: BinaryOpCudaKernel<E> + DeviceRepr + Clone> BinaryKernel<K, E>
             rhs.strides,
         );
 
-        let out_dims1 = self.dev.htod_copy(out_dims1)?;
-        let out_strides1 = self.dev.htod_copy(out_strides1)?;
-        let rhs_strides1 = self.dev.htod_copy(rhs_strides1)?;
-        let out_dims2 = self.dev.htod_copy(out_dims2)?;
-        let out_strides2 = self.dev.htod_copy(out_strides2)?;
-        let lhs_strides2 = self.dev.htod_copy(lhs_strides2)?;
+        let mut info: Vec<usize> = Vec::with_capacity(6 * S::NUM_DIMS);
+        info.extend(out_dims1);
+        info.extend(out_strides1);
+        info.extend(rhs_strides1);
+        info.extend(out_dims2);
+        info.extend(out_strides2);
+        info.extend(lhs_strides2);
+        let info = self.dev.htod_copy(info)?;
 
-        let chunk_len1 = numel / physical_numel(lhs.shape.concrete(), lhs.strides);
-        let chunk_len2 = numel / physical_numel(rhs.shape.concrete(), rhs.strides);
+        self.par_stream.wait_for_default()?;
 
         let params_lhs = (
-            op.clone(),        // const OP_STRUCT op,
-            numel,             // const size_t numel,
-            S::NUM_DIMS,       // const size_t num_dims,
-            &out_dims1,        // const size_t *dims,
-            &out_strides1,     // const size_t *out_strides,
-            lhs.data.as_ref(), // const TYPENAME *lhs,
-            grad_lhs,          // TYPENAME *grad_lhs,
-            chunk_len1,        // const size_t chunk_len,
-            rhs.data.as_ref(), // const TYPENAME *rhs,
-            &rhs_strides1,     // const size_t *rhs_strides,
-            grad_out,          // const TYPENAME *grad_out
+            op.clone(),             // const OP_STRUCT op,
+            numel,                  // const size_t numel,
+            S::NUM_DIMS,            // const size_t num_dims,
+            &info,                  // const size_t *info,
+            lhs.data.as_ref(),      // const TYPENAME *lhs,
+            grad_lhs,               // TYPENAME *grad_lhs,
+            numel / lhs.data.len(), // const size_t chunk_len,
+            rhs.data.as_ref(),      // const TYPENAME *rhs,
+            grad_out,               // const TYPENAME *grad_out
         );
 
-        self.par_stream.wait_for_default()?;
         unsafe { bwd_lhs_fn.launch_on_stream(&self.par_stream, cfg, params_lhs) }?;
 
         let params_rhs = (
-            op,                // const OP_STRUCT op,
-            numel,             // const size_t numel,
-            S::NUM_DIMS,       // const size_t num_dims,
-            &out_dims2,        // const size_t *dims,
-            &out_strides2,     // const size_t *out_strides,
-            lhs.data.as_ref(), // const TYPENAME *lhs,
-            &lhs_strides2,     // const size_t *lhs_strides,
-            rhs.data.as_ref(), // const TYPENAME *rhs,
-            grad_rhs,          // TYPENAME *grad_rhs,
-            chunk_len2,        // const size_t chunk_len,
-            grad_out,          // const TYPENAME *grad_out
+            op,                     // const OP_STRUCT op,
+            numel,                  // const size_t numel,
+            S::NUM_DIMS,            // const size_t num_dims,
+            &info,                  // const size_T * info,
+            lhs.data.as_ref(),      // const TYPENAME *lhs,
+            rhs.data.as_ref(),      // const TYPENAME *rhs,
+            grad_rhs,               // TYPENAME *grad_rhs,
+            numel / rhs.data.len(), // const size_t chunk_len,
+            grad_out,               // const TYPENAME *grad_out
         );
 
         unsafe { bwd_rhs_fn.launch(cfg, params_rhs) }?;
diff --git a/src/tensor_ops/utilities/reduction_utils.rs b/src/tensor_ops/utilities/reduction_utils.rs
index 8cdc94e9b..fefb9811e 100644
--- a/src/tensor_ops/utilities/reduction_utils.rs
+++ b/src/tensor_ops/utilities/reduction_utils.rs
@@ -75,6 +75,7 @@ where
 /// Returns the physical number of elements and strides of dst so that broadcasted dimensions in
 /// src are also broadcasted in dst
 #[cfg(feature = "cuda")]
+#[inline(always)]
 pub(crate) fn reduction_output_strides<Ax: Axes, Src: Shape, Dst: Shape>(
     src_strides: Src::Concrete,
     dst: Dst,
@@ -101,6 +102,7 @@ pub(crate) fn reduction_output_strides<Ax: Axes, Src: Shape, Dst: Shape>(
 
 /// Gives the product of all dimensions that are being reduced and are broadcasted.
 #[cfg(feature = "cuda")]
+#[inline(always)]
 pub(crate) fn reduction_elems_per_thread<Ax: IntoIterator<Item = isize>, S: Shape>(
     dims: S::Concrete,
     strides: S::Concrete,


Test patch: