Update to rust-cuda with async kernel launch async return

necsim/core/Cargo.toml

@@ -20,7 +20,7 @@ contracts = "0.6.3"
 serde = { version = "1.0", default-features = false, features = ["derive"] }
 
 [target.'cfg(target_os = "cuda")'.dependencies]
-rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "c74b542", features = ["derive"], optional = true }
+rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "4148959", features = ["derive"], optional = true }
 
 [target.'cfg(not(target_os = "cuda"))'.dependencies]
-rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "c74b542", features = ["derive", "host"], optional = true }
+rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "4148959", features = ["derive", "host"], optional = true }

necsim/impls/cuda/Cargo.toml

@@ -15,7 +15,7 @@ contracts = "0.6.3"
 serde = { version = "1.0", default-features = false, features = ["derive"] }
 
 [target.'cfg(target_os = "cuda")'.dependencies]
-rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "c74b542", features = ["derive"] }
+rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "4148959", features = ["derive"] }
 
 [target.'cfg(not(target_os = "cuda"))'.dependencies]
-rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "c74b542", features = ["derive", "host"] }
+rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "4148959", features = ["derive", "host"] }

necsim/impls/no-std/Cargo.toml

@@ -30,7 +30,7 @@ fnv = { version = "1.0", default-features = false, features = [] }
 rand_core = "0.6"
 
 [target.'cfg(target_os = "cuda")'.dependencies]
-rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "c74b542", features = ["derive", "final"], optional = true }
+rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "4148959", features = ["derive", "final"], optional = true }
 
 [target.'cfg(not(target_os = "cuda"))'.dependencies]
-rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "c74b542", features = ["derive", "final", "host"], optional = true }
+rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "4148959", features = ["derive", "final", "host"], optional = true }

rustcoalescence/algorithms/cuda/Cargo.toml

@@ -23,4 +23,4 @@ thiserror = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_state = "0.4"
 serde_derive_state = "0.4"
-rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "c74b542", features = ["host"] }
+rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "4148959", features = ["host"] }

rustcoalescence/algorithms/cuda/cpu-kernel/Cargo.toml

@@ -14,4 +14,4 @@ necsim-impls-no-std = { path = "../../../../necsim/impls/no-std", features = ["c
 necsim-impls-cuda = { path = "../../../../necsim/impls/cuda" }
 rustcoalescence-algorithms-cuda-gpu-kernel = { path = "../gpu-kernel" }
 
-rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "c74b542", features = ["host"] }
+rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "4148959", features = ["host"] }

rustcoalescence/algorithms/cuda/gpu-kernel/Cargo.toml

@@ -17,7 +17,7 @@ necsim-impls-no-std = { path = "../../../../necsim/impls/no-std", features = ["c
 necsim-impls-cuda = { path = "../../../../necsim/impls/cuda" }
 
 [target.'cfg(target_os = "cuda")'.dependencies]
-rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "c74b542", features = ["derive", "device", "kernel"] }
+rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "4148959", features = ["derive", "device", "kernel"] }
 
 [target.'cfg(not(target_os = "cuda"))'.dependencies]
-rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "c74b542", features = ["derive", "kernel"] }
+rust-cuda = { git = "https://github.com/juntyr/rust-cuda", rev = "4148959", features = ["derive", "kernel"] }

rustcoalescence/algorithms/cuda/src/launch.rs

@@ -177,7 +177,7 @@ where
     };
 
     let (mut status, time, steps, lineages) = with_initialised_cuda(args.device, || {
-        let stream = CudaDropWrapper::from(Stream::new(StreamFlags::NON_BLOCKING, None)?);
+        let mut stream = CudaDropWrapper::from(Stream::new(StreamFlags::NON_BLOCKING, None)?);
 
         let mut kernel = TypedPtxKernel::new::<Ptx>(Some(Box::new(|kernel| {
             crate::info::print_kernel_function_attributes("simulate", kernel);
@@ -190,21 +190,23 @@ where
             ptx_jit: args.ptx_jit,
         };
 
-        let launcher = Launcher {
-            stream: &stream,
-            kernel: &mut kernel,
-            config,
-        };
-
-        parallelisation::monolithic::simulate(
-            &mut simulation,
-            launcher,
-            (args.dedup_cache, args.step_slice),
-            lineages,
-            event_slice,
-            pause_before,
-            local_partition,
-        )
+        rust_cuda::host::Stream::with(&mut stream, |stream| {
+            let launcher = Launcher {
+                stream,
+                kernel: &mut kernel,
+                config,
+            };
+
+            parallelisation::monolithic::simulate(
+                &mut simulation,
+                launcher,
+                (args.dedup_cache, args.step_slice),
+                lineages,
+                event_slice,
+                pause_before,
+                local_partition,
+            )
+        })
     })
     .map_err(CudaError::from)?;
 

rustcoalescence/algorithms/cuda/src/parallelisation/monolithic.rs

@@ -249,14 +249,14 @@ pub fn simulate<
                         let mut task_list_cuda_async =
                             task_list.move_to_device_async(launcher.stream)?;
 
-                        launcher.launch9_async(
-                            simulation_cuda_repr.as_async(launcher.stream).as_ref(),
-                            task_list_cuda_async.as_mut_async().proj_mut(),
-                            event_buffer_cuda_async.as_mut_async().proj_mut(),
-                            min_spec_sample_buffer_cuda_async.as_mut_async().proj_mut(),
-                            next_event_time_buffer_cuda_async.as_mut_async().proj_mut(),
-                            total_time_max.as_ref().as_async(launcher.stream).as_ref(),
-                            total_steps_sum.as_ref().as_async(launcher.stream).as_ref(),
+                        let launch = launcher.launch9_async(
+                            simulation_cuda_repr.as_async(launcher.stream).extract_ref(),
+                            task_list_cuda_async.as_mut_async(),
+                            event_buffer_cuda_async.as_mut_async(),
+                            min_spec_sample_buffer_cuda_async.as_mut_async(),
+                            next_event_time_buffer_cuda_async.as_mut_async(),
+                            total_time_max.as_ref().as_async(launcher.stream).extract_ref(),
+                            total_steps_sum.as_ref().as_async(launcher.stream).extract_ref(),
                             step_slice.get(),
                             level_time,
                         )?;
@@ -276,6 +276,8 @@ pub fn simulate<
                         next_event_time_buffer = next_event_time_buffer_host_async.synchronize()?;
                         min_spec_sample_buffer = min_spec_sample_buffer_host_async.synchronize()?;
 
+                        launch.synchronize()?;
+
                         // Fetch the completion of the tasks
                         for ((mut spec_sample, mut next_event_time), mut task) in
                             min_spec_sample_buffer