feat/backend: add ability to synchronize a backend

cheme · Feb 2, 2016 · e1bf014 · e1bf014
1 parent f42db23
commit e1bf014
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 1 deletion.
diff --git a/src/backend.rs b/src/backend.rs
@@ -100,6 +100,9 @@ pub trait IBackend {
 
     /// Returns the backend device.
     fn device(&self) -> &DeviceType;
+
+    /// Synchronize backend.
+    fn synchronize(&self) -> Result<(), ::framework::Error> { Ok(()) }
 }
 
 #[derive(Debug, Clone)]

diff --git a/src/frameworks/cuda/api/driver/context.rs b/src/frameworks/cuda/api/driver/context.rs
@@ -25,11 +25,19 @@ impl API {
         unsafe {API::ffi_destroy_context(context)}
     }
 
+    /// Synchronize the CUDA context associated with the current CPU thread.
+    ///
+    /// Should be called when you want to make sure that previous asynchronous operations
+    /// have been executed.
+    pub fn synchronize_context() -> Result<(), Error> {
+        unsafe {API::ffi_synchronize_context()}
+    }
+
     unsafe fn ffi_create_context(
         dev: CUdevice,
     ) -> Result<CUcontext, Error> {
         let mut context: CUcontext = ptr::null_mut();
-        match cuCtxCreate_v2(&mut context, CU_CTX_SCHED_AUTO, dev) {
+        match cuCtxCreate_v2(&mut context, CU_CTX_SCHED_BLOCKING_SYNC, dev) {
             CUresult::CUDA_SUCCESS => Ok(context),
             CUresult::CUDA_ERROR_DEINITIALIZED => Err(Error::Deinitialized("CUDA got deinitialized.")),
             CUresult::CUDA_ERROR_NOT_INITIALIZED => Err(Error::NotInitialized("CUDA is not initialized.")),
@@ -54,4 +62,15 @@ impl API {
             _ => Err(Error::Unknown("Unable to destroy Cuda context.")),
         }
     }
+
+    unsafe fn ffi_synchronize_context () -> Result<(), Error> {
+        match cuCtxSynchronize() {
+            CUresult::CUDA_SUCCESS => Ok(()),
+            CUresult::CUDA_ERROR_DEINITIALIZED => Err(Error::Deinitialized("CUDA got deinitialized.")),
+            CUresult::CUDA_ERROR_NOT_INITIALIZED => Err(Error::NotInitialized("CUDA is not initialized.")),
+            CUresult::CUDA_ERROR_INVALID_CONTEXT => Err(Error::InvalidContext("No valid context available.")),
+            CUresult::CUDA_ERROR_INVALID_VALUE => Err(Error::InvalidValue("Invalid value provided.")),
+            _ => Err(Error::Unknown("Unable to synchronize CUDA context.")),
+        }
+    }
 }
diff --git a/src/frameworks/cuda/context.rs b/src/frameworks/cuda/context.rs
@@ -56,6 +56,11 @@ impl Context {
     pub fn id_c(&self) -> DriverFFI::CUcontext {
         *self.id as DriverFFI::CUcontext
     }
+
+    /// Synchronize this Context.
+    pub fn synchronize(&self) -> Result<(), DriverError> {
+        Driver::synchronize_context()
+    }
 }
 
 #[cfg(feature = "native")]

diff --git a/src/frameworks/cuda/mod.rs b/src/frameworks/cuda/mod.rs
@@ -89,4 +89,12 @@ impl IBackend for Backend<Cuda> {
     fn device(&self) -> &DeviceType {
         &self.device()
     }
+
+    fn synchronize(&self) -> Result<(), ::framework::Error> {
+        if let &DeviceType::Cuda(ref ctx) = self.device() {
+            Ok(try!(ctx.synchronize()))
+        } else {
+            Err(::framework::Error::Implementation(format!("CUDA backend does not have a CUDA context.")))
+        }
+    }
 }
diff --git a/tests/framework_cuda_specs.rs b/tests/framework_cuda_specs.rs
@@ -4,6 +4,7 @@ extern crate libc;
 #[cfg(test)]
 #[cfg(feature = "cuda")]
 mod framework_cuda_spec {
+    use co::backend::{IBackend, Backend, BackendConfig};
     use co::framework::IFramework;
     use co::frameworks::Cuda;
     use co::device::DeviceType;
@@ -54,4 +55,13 @@ mod framework_cuda_spec {
             let _ = &mut SharedTensor::<f32>::new(&device, &vec![256, 1024, 128]).unwrap();
         }
     }
+
+    #[test]
+    fn it_can_synchronize_context() {
+        let framework = Cuda::new();
+        let hardwares = &framework.hardwares().to_vec();
+        let backend_config = BackendConfig::new(framework, hardwares);
+        let backend = Backend::new(backend_config).unwrap();
+        backend.synchronize().unwrap();
+    }
 }