From 24c1a0f6acf422ae3e9128de869811d173539090 Mon Sep 17 00:00:00 2001
From: Sebastien Boisvert <sebhtml@protonmail.com>
Date: Wed, 12 Jun 2024 06:37:40 -0400
Subject: [PATCH] Implement Send + Sync for CudaStream

---
 src/driver/safe/core.rs | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/src/driver/safe/core.rs b/src/driver/safe/core.rs
index 62f1f94..21a70da 100644
--- a/src/driver/safe/core.rs
+++ b/src/driver/safe/core.rs
@@ -484,6 +484,9 @@ pub struct CudaStream {
     device: Arc<CudaDevice>,
 }
 
+unsafe impl Send for CudaStream {}
+unsafe impl Sync for CudaStream {}
+
 impl CudaDevice {
     /// Allocates a new stream that can execute kernels concurrently to the default stream.
     ///
@@ -906,6 +909,8 @@ impl<R: RangeBounds<usize>> RangeHelper for R {
 
 #[cfg(test)]
 mod tests {
+    use std::thread;
+
     use super::*;
 
     #[test]
@@ -930,4 +935,17 @@ mod tests {
         assert!(unsafe { slice.transmute_mut::<f32>(25) }.is_some());
         assert!(unsafe { slice.transmute_mut::<f32>(26) }.is_none());
     }
+
+    #[test]
+    fn test_send_dev() {
+        let dev = CudaDevice::new(0).unwrap();
+        thread::spawn(|| dev);
+    }
+
+    #[test]
+    fn test_send_stream() {
+        let dev = CudaDevice::new(0).unwrap();
+        let stream = dev.fork_default_stream().unwrap();
+        thread::spawn(|| stream);
+    }
 }