Feat/multi streams (#3775)

Author: nathanielsimard

Cargo.lock

@@ -169,9 +169,9 @@ portable-atomic = { version = "1.11.1" }
 portable-atomic-util = { version = "0.2.4", features = ["alloc"] }
 
 ### For the main burn branch. ###
-cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "98eb9d27624375fb11d4a97febd364b049933fdc" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "98eb9d27624375fb11d4a97febd364b049933fdc" }
-cubecl-quant = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "98eb9d27624375fb11d4a97febd364b049933fdc" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "c038fb30634ab4b27bc55f49729eed91c38ecf55" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "c038fb30634ab4b27bc55f49729eed91c38ecf55" }
+cubecl-quant = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "c038fb30634ab4b27bc55f49729eed91c38ecf55" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }

Cargo.toml

@@ -195,6 +195,7 @@ impl<R: FusionRuntime> MultiStream<R> {
     /// Drain a stream
     pub fn drain(&mut self, handles: &mut HandleContainer<R::FusionHandle>, id: StreamId) {
         if let Some(stream) = self.streams.get_mut(&id) {
+            let old = unsafe { StreamId::swap(id) };
             let num_executed = stream.queue.global.len();
             stream.processor.process(
                 Segment::new(&mut stream.queue, handles),
@@ -208,6 +209,9 @@ impl<R: FusionRuntime> MultiStream<R> {
             let to_drop = self.shared_tensors.clear_tensors(cleared);
 
             self.drop_shared_tensors(to_drop, handles, id);
+            unsafe {
+                StreamId::swap(old);
+            };
         }
     }
 
@@ -311,12 +315,8 @@ impl<R: FusionRuntime> MultiStream<R> {
         }
 
         for id in streams_to_sync.drain() {
-            let old = unsafe { StreamId::swap(id) };
             log::info!("Drain stream {id} for use in current {current}");
             self.resolve_stream(handles, id, nodes);
-            unsafe {
-                StreamId::swap(old);
-            };
         }
     }