Repository: coreylowman/dfdx (language: rust)

Issue:
Add synchronize method to DeviceStorage & use in benchmarks
Related to #383 

Currently, only the Cuda device has a synchronize method. This issue is so that both devices will have synchronize method in a generic context.

The Cpu implementation can just do nothing.

The existing Cuda method can just be moved to the DeviceStorage implementation

All benchmarks (under `benches/`) should have `dev.synchronize().unwrap()` added right before elapsed duration is captured.


Discussion / hints:

Fix patch:
diff --git a/benches/batchnorm2d.rs b/benches/batchnorm2d.rs
index 81e4b77eb..43ba72a17 100644
--- a/benches/batchnorm2d.rs
+++ b/benches/batchnorm2d.rs
@@ -29,10 +29,12 @@ fn main() {
         let start = Instant::now();
         let out = m.forward_mut(img.traced(grads));
         let loss = out.square().mean();
+        dev.synchronize();
         let fwd_dur = start.elapsed();
 
         let start = Instant::now();
         let _ = loss.backward();
+        dev.synchronize();
         let bwd_dur = start.elapsed();
         println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
     }
diff --git a/benches/conv2d.rs b/benches/conv2d.rs
index dc0118983..58eaf686a 100644
--- a/benches/conv2d.rs
+++ b/benches/conv2d.rs
@@ -31,10 +31,12 @@ fn main() {
         let start = Instant::now();
         let out = m.forward_mut(img.leaky_traced());
         let loss = out.square().mean();
+        dev.synchronize();
         let fwd_dur = start.elapsed();
 
         let start = Instant::now();
         let _ = loss.backward();
+        dev.synchronize();
         let bwd_dur = start.elapsed();
         println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
     }
diff --git a/benches/softmax.rs b/benches/softmax.rs
index fc638fb44..517615aa1 100644
--- a/benches/softmax.rs
+++ b/benches/softmax.rs
@@ -27,10 +27,12 @@ fn main() {
 
         let start = Instant::now();
         let y = img.traced(grads).softmax::<Ax>();
+        dev.synchronize();
         let fwd_dur = start.elapsed();
 
         let start = Instant::now();
         let _ = y.sum().backward();
+        dev.synchronize();
         let bwd_dur = start.elapsed();
         println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
     }
diff --git a/benches/sum.rs b/benches/sum.rs
index 6efbb9240..2e490d5d1 100644
--- a/benches/sum.rs
+++ b/benches/sum.rs
@@ -26,10 +26,12 @@ fn main() {
 
         let start = Instant::now();
         let y = img.traced(grads).sum();
+        dev.synchronize();
         let fwd_dur = start.elapsed();
 
         let start = Instant::now();
         let _ = y.backward();
+        dev.synchronize();
         let bwd_dur = start.elapsed();
         println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
     }
diff --git a/src/tensor/cpu/device.rs b/src/tensor/cpu/device.rs
index c06b956e6..b5698269f 100644
--- a/src/tensor/cpu/device.rs
+++ b/src/tensor/cpu/device.rs
@@ -86,4 +86,8 @@ impl DeviceStorage for Cpu {
         }
         buf
     }
+
+    fn try_synchronize(&self) -> Result<(), Self::Err> {
+        Ok(())
+    }
 }
diff --git a/src/tensor/cuda/device.rs b/src/tensor/cuda/device.rs
index 6e791ae0d..0a37acb97 100644
--- a/src/tensor/cuda/device.rs
+++ b/src/tensor/cuda/device.rs
@@ -82,20 +82,6 @@ impl Cuda {
             workspace,
         })
     }
-
-    /// Block until kernels finish processing. Useful for benchmarking.
-    ///
-    /// Examples:
-    /// ```rust
-    /// # use dfdx::prelude::*;
-    /// let dev: Cuda = Default::default();
-    /// let a = dev.tensor([1., 2., 3.]);
-    /// let _b = a.square();
-    /// dev.synchronize().unwrap(); // blocks until square kernel finishes.
-    /// ```
-    pub fn synchronize(&self) -> Result<(), CudaError> {
-        self.dev.synchronize().map_err(CudaError::from)
-    }
 }
 
 impl Cuda {
@@ -149,4 +135,8 @@ impl DeviceStorage for Cuda {
         }
         contiguous
     }
+
+    fn try_synchronize(&self) -> Result<(), CudaError> {
+        self.dev.synchronize().map_err(CudaError::from)
+    }
 }
diff --git a/src/tensor/storage_traits.rs b/src/tensor/storage_traits.rs
index fb5fe9fa2..68927793a 100644
--- a/src/tensor/storage_traits.rs
+++ b/src/tensor/storage_traits.rs
@@ -28,6 +28,14 @@ pub trait DeviceStorage: 'static + std::fmt::Debug + Default + Clone + HasErr {
     fn try_alloc_grad<E: Unit>(&self, storage: &Self::Vec<E>) -> Result<Self::Vec<E>, Self::Err>;
 
     fn tensor_to_vec<S: Shape, E: Unit, T>(&self, tensor: &Tensor<S, E, Self, T>) -> Vec<E>;
+
+    /// Blocks until all work on device to complete. Useful for benchmarking.
+    fn synchronize(&self) {
+        self.try_synchronize().unwrap()
+    }
+
+    /// Blocks until all work on device to complete. Useful for benchmarking.
+    fn try_synchronize(&self) -> Result<(), Self::Err>;
 }
 
 /// Internal trait - Represents something that can allocate its own gradient.


Test patch: