Repository: coreylowman/dfdx (language: rust)

Issue:
Update conv transpose cuda kernels to match conv
The conv kernels/launching logic has gone through a number of updates lately, want to make sure they are shared by conv transpose


Discussion / hints:

Fix patch:
diff --git a/src/nn/mod.rs b/src/nn/mod.rs
index 95c564764..1b3cd88a8 100644
--- a/src/nn/mod.rs
+++ b/src/nn/mod.rs
@@ -240,6 +240,7 @@ pub mod modules {
     pub use super::bias2d::Bias2D;
     #[cfg(feature = "nightly")]
     pub use super::conv::Conv2D;
+    #[cfg(feature = "nightly")]
     pub use super::convtrans::ConvTrans2D;
     pub use super::dropout::{Dropout, DropoutOneIn};
     pub use super::embedding::Embedding;
@@ -273,6 +274,7 @@ pub mod builders {
     pub use super::bias2d::builder::Bias2D;
     #[cfg(feature = "nightly")]
     pub use super::conv::builder::Conv2D;
+    #[cfg(feature = "nightly")]
     pub use super::convtrans::builder::ConvTrans2D;
     pub use super::dropout::{Dropout, DropoutOneIn};
     pub use super::embedding::builder::Embedding;
@@ -292,5 +294,6 @@ pub mod builders {
         TransformerEncoder, TransformerEncoderBlock,
     };
     pub use super::unbiased_linear::builder::UnbiasedLinear;
+    pub use super::upscale::Upscale2D;
     pub use super::*;
 }
diff --git a/src/nn/upscale.rs b/src/nn/upscale.rs
index 3f85d71e4..9efa80738 100644
--- a/src/nn/upscale.rs
+++ b/src/nn/upscale.rs
@@ -87,7 +87,7 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{prelude::Bilinear, shapes::*, tensor::*, tests::*};
+    use crate::{shapes::*, tensor::*, tests::*};
 
     #[test]
     fn test_upscale2d() {
@@ -96,12 +96,13 @@ mod tests {
         let _: Tensor<Rank3<3, 8, 8>, _, _> = Upscale2D::<8>::default().forward(x.clone());
         let _: Tensor<Rank3<3, 8, 12>, _, _> = Upscale2D::<8, 12>::default().forward(x.clone());
         let _: Tensor<Rank3<3, 9, 9>, _, _> =
-            Upscale2D::<9, 9, NearestNeighbor>::default().forward(x.clone());
+            Upscale2D::<9, 9, NearestNeighbor>::default().forward(x);
     }
 
     #[cfg(feature = "nightly")]
     #[test]
     fn test_upscale2dby() {
+        use crate::prelude::Bilinear;
         let dev: TestDevice = Default::default();
         let x: Tensor<Rank3<3, 4, 4>, TestDtype, _> = dev.zeros();
         let _: Tensor<Rank3<3, 8, 8>, _, _> = Upscale2DBy::<2>::default().forward(x.clone());
diff --git a/src/tensor_ops/convtrans2d/convtrans2d.cu b/src/tensor_ops/convtrans2d/convtrans2d.cu
index d2d98abd1..ff267c416 100644
--- a/src/tensor_ops/convtrans2d/convtrans2d.cu
+++ b/src/tensor_ops/convtrans2d/convtrans2d.cu
@@ -19,54 +19,40 @@ __device__ void unfold_input_into_patches(
     T *patches // 6d (Batch, Channels, KernelSize, KernelSize, HeightOut, WidthOut)
 ) {
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    const auto patches_numel = op.batch * op.chan_in * op.kernel * op.kernel * op.h_out * op.w_out;
-    if (i >= patches_numel) {
+    if (i >= op.batch * op.chan_in * op.h_out * op.w_out) {
         return;
     }
 
-    // patches shape is (B, C, K, K, h_out, w_out)
     unsigned int idx = i;
     const size_t ow = idx % op.w_out;
     idx /= op.w_out;
     const size_t oh = idx % op.h_out;
     idx /= op.h_out;
-    const size_t k2 = idx % op.kernel;
-    idx /= op.kernel;
-    const size_t k1 = idx % op.kernel;
-    idx /= op.kernel;
     const size_t c = idx % op.chan_in;
     idx /= op.chan_in;
     const size_t b = idx % op.batch;
     idx /= op.batch;
 
-    size_t y = oh + op.padding;
-    if (y < k1) {
-        return;
+    image += b * strides[0] + c * strides[1];
+    patches += oh * op.w_out + ow;
+    patches += c * (op.kernel * op.kernel * op.h_out * op.w_out);
+    patches += b * (op.chan_in * op.kernel * op.kernel * op.h_out * op.w_out);
+
+    for (int k1 = 0;k1 < op.kernel;k1++) {
+        const size_t y_ks = oh + op.padding;
+        const size_t y_s = y_ks - k1;
+        const size_t y = y_s / op.stride;
+        const bool k1_invalid = (y_ks < k1 || y_s % op.stride != 0 || y >= op.h_in);
+        for (int k2 = 0;k2 < op.kernel;k2++) {
+            const size_t x_ks = ow + op.padding;
+            const size_t x_s = x_ks - k2;
+            const size_t x = x_s / op.stride;
+        
+            const bool invalid = k1_invalid || (x_ks < k2 || x_s % op.stride != 0 || x >= op.w_in);
+            *patches = invalid ? 0.0 : image[y * strides[2] + x * strides[3]];
+            patches += op.h_out * op.w_out;
+        }
     }
-    y -= k1;
-    if (y % op.stride != 0) {
-        return;
-    }
-    y /= op.stride;
-    if (y >= op.h_in) {
-        return;
-    }
-    
-    size_t x = ow + op.padding;
-    if (x < k2) {
-        return;
-    }
-    x -= k2;
-    if (x % op.stride != 0) {
-        return;
-    }
-    x /= op.stride;
-    if (x >= op.w_in) {
-        return;
-    }
-
-    const size_t i_image = b * strides[0] + c * strides[1] + y * strides[2] + x * strides[3];
-    patches[i] = image[i_image];
 }
 
 template<typename T>
@@ -76,8 +62,7 @@ __device__ void unfold_output_into_patches(
     T *patches // 6d (Batch, ChanOut, KernelSize, KernelSize, HeightIn, WidthIn)
 ) {
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    const auto patches_numel = op.batch * op.chan_out * op.kernel * op.kernel * op.h_in * op.w_in;
-    if (i >= patches_numel) {
+    if (i >= op.batch * op.chan_out * op.h_in * op.w_in) {
         return;
     }
 
@@ -86,47 +71,34 @@ __device__ void unfold_output_into_patches(
     idx /= op.w_in;
     const size_t y = idx % op.h_in;
     idx /= op.h_in;
-    const size_t k2 = idx % op.kernel;
-    idx /= op.kernel;
-    const size_t k1 = idx % op.kernel;
-    idx /= op.kernel;
     const size_t o = idx % op.chan_out;
     idx /= op.chan_out;
     const size_t b = idx % op.batch;
-    idx /= op.batch;
 
-    const size_t y_plus_p = y * op.stride + k1;
-    if (y_plus_p < op.padding) {
-        return;
-    }
-    const size_t oh = y_plus_p - op.padding;
-    if (oh >= op.h_out) {
-        return;
-    }
+    image_out += b * (op.chan_out * op.h_out * op.w_out) + o * (op.h_out * op.w_out);
+    patches += y * op.w_in + x;
+    patches += o * (op.kernel * op.kernel * op.h_in * op.w_in);
+    patches += b * (op.chan_in * op.kernel * op.kernel * op.h_in * op.w_in);
 
-    const size_t x_plus_p = x * op.stride + k2;
-    if (x_plus_p < op.padding) {
-        return;
-    }
-    const size_t ow = x_plus_p - op.padding;
-    if (ow >= op.w_out) {
-        return;
+    for (int k1 = 0;k1 < op.kernel;k1++) {
+        const size_t oh = y * op.stride + k1 - op.padding;
+        for (int k2 = 0;k2 < op.kernel;k2++) {
+            const size_t ow = x * op.stride + k2 - op.padding;
+            *patches = (oh >= op.h_out || ow >= op.w_out) ? 0.0 : image_out[oh * op.w_out + ow];
+            patches += op.h_in * op.w_in;
+        }
     }
-
-    size_t image_i = b * (op.chan_out * op.h_out * op.w_out) + o * (op.h_out * op.w_out) + oh * (op.w_out)  + ow;
-    patches[i] = image_out[image_i];
 }
 
 template<typename T>
-__device__ void transpose_and_broadcast_filters(
+__device__ void transpose_filters(
     const Conv2DOp op,
     const T *filters, // 4d (ChanOut, ChanIn, KernelSize, KernelSize)
     const size_t *strides, // 4d filters strides
-    T *filters_tr // 5d (Batch, ChanIn, ChanOut, KernelSize, KernelSize)
+    T *filters_tr // 4d (ChanIn, ChanOut, KernelSize, KernelSize)
 ) {
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    auto numel = op.chan_in * op.chan_out * op.kernel * op.kernel;
-    if (i >= numel) {
+    if (i >= op.chan_in * op.chan_out * op.kernel * op.kernel) {
         return;
     }
 
@@ -135,18 +107,13 @@ __device__ void transpose_and_broadcast_filters(
     idx /= op.kernel;
     const size_t k1 = idx % op.kernel;
     idx /= op.kernel;
-    const size_t c = idx % op.chan_in;
-    idx /= op.chan_in;
     const size_t o = idx % op.chan_out;
     idx /= op.chan_out;
+    const size_t c = idx % op.chan_in;
 
-    auto i_tr = c * (op.chan_out * op.kernel * op.kernel) + o * (op.kernel * op.kernel) + k1 * (op.kernel) + k2;
     auto i_no = o * strides[0] + c * strides[1] + k1 * strides[2] + k2 * strides[3];
 
-    const T f = filters[i_no];
-    for (auto b = 0; b < op.batch; b++) {
-        filters_tr[b * numel + i_tr] = f;
-    }
+    filters_tr[i] = filters[i_no];
 }
 
 template<typename T>
@@ -175,9 +142,12 @@ __device__ void sum_transposed_filters(
     auto i_tr = c * (op.chan_out * op.kernel * op.kernel) + o * (op.kernel * op.kernel) + k1 * (op.kernel) + k2;
     auto i_no = o * strides[0] + c * strides[1] + k1 * strides[2] + k2 * strides[3];
 
+    const T *ptr = filters_tr + i_tr;
+
     T tmp = 0.0;
     for (auto b = 0; b < op.batch; b++) {
-        tmp += filters_tr[b * numel + i_tr];
+        tmp += *ptr;
+        ptr += numel;
     }
 
     filters[i_no] += tmp;
@@ -205,7 +175,7 @@ extern "C" __global__ void TR_FILTERS( \
     const size_t *strides, \
     TYPENAME *filters_tr \
 ) { \
-    transpose_and_broadcast_filters(op, filters, strides, filters_tr); \
+    transpose_filters(op, filters, strides, filters_tr); \
 } \
 extern "C" __global__ void SUM_TR_FILTERS( \
     const Conv2DOp op, \
@@ -220,13 +190,13 @@ CONV_OP(
     float,
     unfold_input_into_patches_f32,
     unfold_output_into_patches_f32,
-    transpose_and_broadcast_filters_f32,
+    transpose_filters_f32,
     sum_transposed_filters_f32
 );
 CONV_OP(
     double,
     unfold_input_into_patches_f64,
     unfold_output_into_patches_f64,
-    transpose_and_broadcast_filters_f64,
+    transpose_filters_f64,
     sum_transposed_filters_f64
 );
\ No newline at end of file
diff --git a/src/tensor_ops/convtrans2d/cuda_kernel.rs b/src/tensor_ops/convtrans2d/cuda_kernel.rs
index a7bf44a95..c69af99df 100644
--- a/src/tensor_ops/convtrans2d/cuda_kernel.rs
+++ b/src/tensor_ops/convtrans2d/cuda_kernel.rs
@@ -1,9 +1,9 @@
 use cudarc::cublas::{CudaBlas, Gemm};
-use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync, LaunchConfig, ValidAsZeroBits};
+use cudarc::driver::{DeviceRepr, LaunchAsync, ValidAsZeroBits};
 
 use crate::{
     shapes::*,
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
 };
 
 use std::sync::Arc;
@@ -22,7 +22,7 @@ impl HasCudaKernel<f32> for Cuda {
     const FNS: &'static [&'static str] = &[
         "unfold_input_into_patches_f32",
         "unfold_output_into_patches_f32",
-        "transpose_and_broadcast_filters_f32",
+        "transpose_filters_f32",
         "sum_transposed_filters_f32",
     ];
 }
@@ -32,7 +32,7 @@ impl HasCudaKernel<f64> for Cuda {
     const FNS: &'static [&'static str] = &[
         "unfold_input_into_patches_f64",
         "unfold_output_into_patches_f64",
-        "transpose_and_broadcast_filters_f64",
+        "transpose_filters_f64",
         "sum_transposed_filters_f64",
     ];
 }
@@ -62,10 +62,12 @@ where
         }
 
         let patches_numel = op.batch * op.chan_in * op.kernel * op.kernel * op.h_out * op.w_out;
-        let mut patches = self.dev.alloc_zeros::<E>(patches_numel)?;
+        let mut patches = unsafe { self.get_workspace::<E>(patches_numel) }?;
+        let mut patches = unsafe { patches.transmute_mut::<E>(patches_numel).unwrap() };
+
         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((op.batch * op.chan_in * op.h_out * op.w_out) as u32);
         let params = (op, lhs.data.as_ref(), &img_strides, &mut patches);
         unsafe { unfold_fn.launch(cfg, params) }?;
 
@@ -81,7 +83,7 @@ where
                 &patches,
                 [k * n, n, 1],
                 Default::default(),
-                Arc::make_mut(&mut out.data),
+                Arc::get_mut(&mut out.data).unwrap(),
                 [m * n, n, 1],
             )
             .unwrap();
@@ -101,35 +103,46 @@ where
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err> {
         let patches_numel = op.batch * op.chan_out * op.kernel * op.kernel * op.h_in * op.w_in;
-        let mut patches = self.dev.alloc_zeros::<E>(patches_numel)?;
+        let filters_numel = op.chan_in * op.chan_out * op.kernel * op.kernel;
+
+        let mut patches = unsafe { self.get_workspace::<E>(patches_numel) }?;
+        let mut patches = unsafe { patches.transmute_mut::<E>(patches_numel).unwrap() };
+
+        let mut f_b1023 = unsafe { self.dev.alloc::<E>(filters_numel) }?;
+        let mut grad_f_b1023 = unsafe { self.dev.alloc::<E>(op.batch * filters_numel) }?;
+        let f_strides = self.dev.htod_copy(rhs.strides.into())?;
+
+        self.par_stream.wait_for_default()?;
 
         {
             // 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((op.batch * op.chan_out * op.h_in * op.w_in) as u32);
             unsafe { unfold_fn.launch(cfg, (op, grad_out, &mut patches)) }?;
         }
 
-        let filters_numel = op.batch * op.chan_in * op.chan_out * op.kernel * op.kernel;
-        let mut f_b1023 = self.dev.alloc_zeros::<E>(filters_numel)?;
-        let mut grad_f_b1023 = self.dev.alloc_zeros::<E>(filters_numel)?;
-        let f_strides = self.dev.htod_copy(rhs.strides.into())?;
-
         {
             // 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);
-            unsafe { tr_fn.launch(cfg, (op, rhs.data.as_ref(), &f_strides, &mut f_b1023)) }?;
-        }
+            let cfg = launch_cfg(rhs.shape.num_elements() as u32);
+            unsafe {
+                tr_fn.launch_on_stream(
+                    self.par_stream.as_ref(),
+                    cfg,
+                    (op, rhs.data.as_ref(), &f_strides, &mut f_b1023),
+                )
+            }?;
+
+            self.par_stream.wait_for_default()?;
 
-        {
             // img_g += filters * patches
             // (B, C, H * W) += (B, C, O * K * K) * (B, O * K * K, H * W)
             let m = op.chan_in;
             let k = op.chan_out * op.kernel * op.kernel;
             let n = op.h_in * op.w_in;
             unsafe {
+                self.blas.set_stream(Some(self.par_stream.as_ref()))?;
                 self.gemm_batch(
                     (op.batch, m, k, n),
                     &f_b1023,
@@ -141,6 +154,7 @@ where
                     [m * n, n, 1],
                 )
                 .unwrap();
+                self.blas.set_stream(None)?;
             }
         }
 
@@ -157,7 +171,7 @@ where
                     [m * k, k, 1],
                     &patches,
                     [k * n, 1, k],
-                    <E>::ONE,
+                    Default::default(),
                     &mut grad_f_b1023,
                     [m * n, n, 1],
                 )
@@ -167,10 +181,12 @@ 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)) }?;
         }
 
+        self.dev.wait_for(self.par_stream.as_ref())?;
+
         Ok(())
     }
 }


Test patch: