From 435e604e2b59195127c6dbb17dcf7680f22df776 Mon Sep 17 00:00:00 2001
From: Minh Vu <vuhoangminh97@gmail.com>
Date: Fri, 12 Jun 2026 19:11:23 +0200
Subject: [PATCH] Preserve device contexts on duplicate init

---
 crates/cuda-async/src/device_context.rs | 32 ++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/crates/cuda-async/src/device_context.rs b/crates/cuda-async/src/device_context.rs
index e984ee572d..659d3d0103 100644
--- a/crates/cuda-async/src/device_context.rs
+++ b/crates/cuda-async/src/device_context.rs
@@ -105,9 +105,12 @@ pub fn init_device_contexts(
     num_devices: usize,
 ) -> Result<(), DeviceError> {
     DEVICE_CONTEXTS.with(|ctx| {
+        let devices = ctx.devices.take();
+        let is_uninitialized = devices.is_none();
+        ctx.devices.set(devices);
         device_assert(
             default_device_id,
-            ctx.devices.replace(None).is_none(),
+            is_uninitialized,
             "Context already initialized.",
         )
     })?;
@@ -327,3 +330,30 @@ pub fn get_cuda_function(
         Ok(Arc::clone(function))
     })?
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn assert_duplicate_init_error(result: Result<(), DeviceError>) {
+        assert!(matches!(
+            result,
+            Err(DeviceError::Context {
+                device_id: 0,
+                message,
+            }) if message == "Context already initialized."
+        ));
+    }
+
+    #[test]
+    fn duplicate_init_preserves_existing_device_contexts() {
+        std::thread::spawn(|| {
+            init_device_contexts(0, 1).expect("initial context initialization should succeed");
+
+            assert_duplicate_init_error(init_device_contexts(0, 1));
+            assert_duplicate_init_error(init_device_contexts(0, 1));
+        })
+        .join()
+        .expect("test thread should not panic");
+    }
+}