1use crate::{
2 config::streaming::StreamingLogLevel,
3 logging::ServerLogger,
4 memory_management::{ErrorGraph, FailureId, ManagedMemoryId, SharedMemoryBindings},
5 server::{BufferBinding, ServerError},
6 stream::{FailureStore, Failures, StreamFactory, StreamMemory, StreamPool, base},
7};
8use core::any::Any;
9use cubecl_environment::collections::HashMap;
10use cubecl_environment::stream::StreamId;
11use std::{
12 boxed::Box,
13 format,
14 sync::{Arc, mpsc::SyncSender},
15 vec::Vec,
16};
17
18pub trait EventStreamBackend: 'static {
23 type Stream: core::fmt::Debug + StreamMemory;
26 type Event: Send + 'static;
28
29 fn create_stream(&self) -> Self::Stream;
31 fn handle_cursor(stream: &Self::Stream, handle: &BufferBinding) -> u64;
33 fn flush(stream: &mut Self::Stream, failures: &mut ErrorGraph) -> Self::Event;
40 fn wait_event(stream: &mut Self::Stream, event: Self::Event);
42 fn wait_event_sync(event: Self::Event) -> Result<(), ServerError>;
44}
45
46#[derive(Debug)]
51pub struct MultiStream<B: EventStreamBackend> {
52 streams: StreamPool<EventStreamBackendWrapper<B>>,
54 failures: Failures,
57 pub logger: Arc<ServerLogger>,
59 max_streams: usize,
60 gc: GcThread<B>,
61 shared_bindings_pool: Vec<(ManagedMemoryId, StreamId, u64)>,
62}
63
64pub struct StreamWrapper<B: EventStreamBackend> {
72 stream: B::Stream,
74 cursor: u64,
76 last_synced: HashMap<usize, u64>,
78}
79
80pub struct ResolvedStreams<'a, B: EventStreamBackend> {
82 pub cursor: u64,
86 streams: &'a mut StreamPool<EventStreamBackendWrapper<B>>,
87 failures: &'a mut ErrorGraph,
88 analysis: SharedBindingAnalysis,
89 gc: &'a GcThread<B>,
90 pub current: StreamId,
92}
93
94#[derive(Debug)]
95pub struct GcTask<B: EventStreamBackend> {
97 to_drop: Box<dyn Any + Send + 'static>,
98 event: B::Event,
100}
101
102impl<B: EventStreamBackend> GcTask<B> {
103 pub fn new<T: Send + 'static>(to_drop: T, event: B::Event) -> Self {
105 Self {
106 to_drop: Box::new(to_drop),
107 event,
108 }
109 }
110}
111
112#[derive(Debug)]
113pub struct EventStreamBackendWrapper<B: EventStreamBackend> {
116 backend: B,
117}
118
119impl<B: EventStreamBackend> StreamMemory for StreamWrapper<B> {
120 fn failure(&self, binding: &BufferBinding) -> Option<FailureId> {
121 self.stream.failure(binding)
122 }
123
124 fn taint(&mut self, binding: &BufferBinding, failure: FailureId, failures: &mut ErrorGraph) {
125 self.stream.taint(binding, failure, failures)
126 }
127
128 fn written(&mut self, binding: &BufferBinding, failures: &mut ErrorGraph) {
129 self.stream.written(binding, failures)
130 }
131}
132
133impl<B: EventStreamBackend> StreamFactory for EventStreamBackendWrapper<B> {
134 type Stream = StreamWrapper<B>;
135
136 fn create(&mut self) -> Self::Stream {
137 StreamWrapper {
138 stream: self.backend.create_stream(),
139 cursor: 0,
140 last_synced: Default::default(),
141 }
142 }
143}
144
145#[derive(Debug)]
146struct GcThread<B: EventStreamBackend> {
147 sender: SyncSender<GcTask<B>>,
148}
149
150impl<B: EventStreamBackend> GcThread<B> {
151 fn new() -> GcThread<B> {
152 let (sender, recv) = std::sync::mpsc::sync_channel::<GcTask<B>>(8);
153
154 cubecl_environment::thread::spawn(move || {
155 while let Ok(event) = recv.recv() {
156 B::wait_event_sync(event.event).unwrap();
157 core::mem::drop(event.to_drop);
158 }
159 });
160
161 GcThread { sender }
162 }
163 fn register(&self, task: GcTask<B>) {
164 self.sender.send(task).unwrap()
165 }
166}
167
168fn stream_index(stream_id: &StreamId, max_streams: usize) -> usize {
169 stream_id.value as usize % max_streams
170}
171
172impl<'a, B: EventStreamBackend> ResolvedStreams<'a, B> {
173 pub fn get(&mut self, stream_id: &StreamId) -> &mut B::Stream {
175 let stream = self.streams.get_mut(stream_id);
176 &mut stream.stream
177 }
178
179 pub fn current(&mut self) -> &mut B::Stream {
181 let stream = self.streams.get_mut(&self.current);
182 &mut stream.stream
183 }
184
185 pub fn current_and_failures(&mut self) -> (&mut B::Stream, &mut ErrorGraph) {
190 let stream = self.streams.get_mut(&self.current);
191 (&mut stream.stream, self.failures)
192 }
193
194 pub fn get_and_failures(&mut self, stream_id: &StreamId) -> (&mut B::Stream, &mut ErrorGraph) {
197 let stream = self.streams.get_mut(stream_id);
198 (&mut stream.stream, self.failures)
199 }
200
201 pub fn taint<'b>(
204 &mut self,
205 error: ServerError,
206 written: impl Iterator<Item = &'b BufferBinding>,
207 ) {
208 base::taint(self.streams, error, written, self.failures);
209 }
210
211 pub fn written<'b>(&mut self, written: impl Iterator<Item = &'b BufferBinding>) {
214 base::written(self.streams, written, self.failures);
215 }
216
217 pub fn gc(&mut self, gc: GcTask<B>) {
219 self.gc.sender.send(gc).unwrap();
220 }
221}
222
223impl<'a, B: EventStreamBackend> Drop for ResolvedStreams<'a, B> {
224 fn drop(&mut self) {
225 if self.analysis.pinned.is_empty() {
226 return;
227 }
228
229 let stream = self.streams.get_mut(&self.current);
230 let event_origin = B::flush(&mut stream.stream, self.failures);
231
232 let stream_gc = &mut unsafe { self.streams.get_special(0) }.stream;
233 B::wait_event(stream_gc, event_origin);
234 let event = B::flush(stream_gc, self.failures);
235
236 let pinned = core::mem::take(&mut self.analysis.pinned);
237 self.gc.register(GcTask::new(pinned, event));
238 }
239}
240
241impl<B: EventStreamBackend> MultiStream<B> {
242 pub fn backend_mut(&mut self) -> &mut B {
246 &mut self.streams.factory_mut().backend
247 }
248
249 pub fn new(logger: Arc<ServerLogger>, backend: B, max_streams: u8) -> Self {
251 let wrapper = EventStreamBackendWrapper { backend };
252 Self {
253 streams: StreamPool::new(wrapper, max_streams, 1),
254 failures: Failures::new(logger.clone()),
255 logger,
256 max_streams: max_streams as usize,
257 gc: GcThread::new(),
258 shared_bindings_pool: Vec::new(),
259 }
260 }
261
262 pub fn stream_ids(&self) -> impl Iterator<Item = StreamId> + '_ {
264 self.streams.stream_ids()
265 }
266
267 pub fn gc(&mut self, gc: GcTask<B>) {
269 self.gc.sender.send(gc).unwrap();
270 }
271
272 pub fn try_stream_mut(&mut self, stream_id: &StreamId) -> Option<&mut B::Stream> {
276 self.streams
277 .try_get_mut(stream_id)
278 .map(|wrapper| &mut wrapper.stream)
279 }
280
281 pub fn resolve<'a>(
287 &mut self,
288 stream_id: StreamId,
289 handles: impl Iterator<Item = &'a BufferBinding>,
290 ) -> ResolvedStreams<'_, B> {
291 let analysis = self.align_streams(stream_id, handles);
292
293 let stream = self.streams.get_mut(&stream_id);
294 stream.cursor += 1;
295
296 ResolvedStreams {
297 cursor: stream.cursor,
298 streams: &mut self.streams,
299 failures: self.failures.graph_mut(),
300 current: stream_id,
301 analysis,
302 gc: &self.gc,
303 }
304 }
305
306 fn align_streams<'a>(
311 &mut self,
312 stream_id: StreamId,
313 handles: impl Iterator<Item = &'a BufferBinding>,
314 ) -> SharedBindingAnalysis {
315 let analysis = self.update_shared_bindings(stream_id, handles);
316
317 self.apply_analysis(stream_id, analysis)
318 }
319
320 pub(crate) fn update_shared_bindings<'a>(
325 &mut self,
326 stream_id: StreamId,
327 handles: impl Iterator<Item = &'a BufferBinding>,
328 ) -> SharedBindingAnalysis {
329 self.shared_bindings_pool.clear();
331
332 let mut analysis = SharedBindingAnalysis::default();
333
334 for handle in handles.filter(|handle| handle.stream != stream_id) {
336 let index = stream_index(&handle.stream, self.max_streams);
337 let stream = unsafe { self.streams.get_mut_index(index) };
338 let cursor_handle = B::handle_cursor(&stream.stream, handle);
339
340 self.shared_bindings_pool.push((
341 handle.memory.descriptor().id,
342 handle.stream,
343 cursor_handle,
344 ));
345 analysis.pinned.push(handle.memory.clone());
350 }
351
352 let current = self.streams.get_mut(&stream_id);
353
354 for (handle_id, stream, cursor) in self.shared_bindings_pool.iter() {
355 let index = stream_index(stream, self.max_streams);
356
357 if let Some(last_synced) = current.last_synced.get(&index) {
358 if last_synced < cursor {
359 self.logger.log_streaming(
360 |level| matches!(level, StreamingLogLevel::Full),
361 || {
362 format!(
363 "Binding on {} is shared on {} since it's not sync {} < {}",
364 stream, stream_id, last_synced, cursor
365 )
366 },
367 );
368 analysis.shared(*handle_id, index);
369 }
370 } else {
371 self.logger.log_streaming(
372 |level| matches!(level, StreamingLogLevel::Full),
373 || {
374 format!(
375 "Binding on {} is shared on {} since it was never synced.",
376 stream, stream_id,
377 )
378 },
379 );
380 analysis.shared(*handle_id, index);
381 }
382 }
383
384 analysis
385 }
386
387 pub(crate) fn apply_analysis(
388 &mut self,
389 stream_id: StreamId,
390 analysis: SharedBindingAnalysis,
391 ) -> SharedBindingAnalysis {
392 if analysis.slices.is_empty() {
393 return analysis;
394 }
395
396 let mut events = Vec::with_capacity(analysis.slices.len());
397
398 unsafe {
399 for origin in analysis.slices.keys() {
400 let stream = self.streams.get_mut_index(*origin);
401 let event = B::flush(&mut stream.stream, self.failures.graph_mut());
402
403 events.push(((origin, stream.cursor), event));
404 }
405 }
406
407 let stream = self.streams.get_mut(&stream_id);
408
409 for ((stream_origin, cursor_origin), event) in events {
410 stream.last_synced.insert(*stream_origin, cursor_origin);
411
412 self.logger.log_streaming(
413 |level| !matches!(level, StreamingLogLevel::Disabled),
414 || format!("Waiting on {stream_origin} from {stream_id}",),
415 );
416
417 B::wait_event(&mut stream.stream, event);
418 }
419
420 analysis
421 }
422}
423
424impl<B: EventStreamBackend> FailureStore for MultiStream<B> {
425 type Factory = EventStreamBackendWrapper<B>;
426
427 fn split(&mut self) -> (&mut StreamPool<Self::Factory>, &mut Failures) {
428 (&mut self.streams, &mut self.failures)
429 }
430
431 fn parts(&self) -> (&StreamPool<Self::Factory>, &Failures) {
432 (&self.streams, &self.failures)
433 }
434}
435
436impl<B: EventStreamBackend> core::fmt::Debug for StreamWrapper<B> {
437 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
438 f.debug_struct("StreamWrapper")
439 .field("stream", &self.stream)
440 .field("cursor", &self.cursor)
441 .field("last_synced", &self.last_synced)
442 .finish()
443 }
444}
445
446#[derive(Default, Debug)]
447pub(crate) struct SharedBindingAnalysis {
448 slices: HashMap<usize, Vec<ManagedMemoryId>>,
449 pinned: SharedMemoryBindings,
459}
460
461impl PartialEq for SharedBindingAnalysis {
464 fn eq(&self, other: &Self) -> bool {
465 self.slices == other.slices
466 }
467}
468
469impl Eq for SharedBindingAnalysis {}
470
471impl SharedBindingAnalysis {
472 fn shared(&mut self, id: ManagedMemoryId, index: usize) {
473 match self.slices.get_mut(&index) {
474 Some(bindings) => bindings.push(id),
475 None => {
476 self.slices.insert(index, alloc::vec![id]);
477 }
478 }
479 }
480}
481
482#[cfg(test)]
483mod tests {
484 use crate::server::Handle;
485 use core::sync::atomic::{AtomicBool, Ordering};
486
487 use super::*;
488
489 fn service() -> cubecl_common::device::ServiceId {
491 cubecl_common::device::ServiceId::of::<()>(cubecl_common::device::DeviceId::new(0, 0))
492 }
493
494 const MAX_STREAMS: u8 = 4;
495
496 #[test_log::test]
497 fn test_analysis_shared_bindings_1() {
498 let logger = Arc::new(ServerLogger::default());
499 let stream_1 = StreamId { value: 1 };
500 let stream_2 = StreamId { value: 2 };
501
502 let binding_1 = handle(stream_1);
503 let binding_2 = handle(stream_2);
504
505 let mut ms = MultiStream::new(logger, TestBackend, MAX_STREAMS);
506 ms.resolve(stream_1, [].into_iter());
507 ms.resolve(stream_2, [].into_iter());
508
509 let analysis = ms.update_shared_bindings(stream_1, [&binding_1, &binding_2].into_iter());
510
511 let mut expected = SharedBindingAnalysis::default();
512 expected.shared(
513 binding_2.memory.descriptor().id,
514 ms.streams.stream_index(&binding_2.stream),
515 );
516
517 assert_eq!(analysis, expected);
518 }
519
520 #[test_log::test]
521 fn test_analysis_shared_bindings_2() {
522 let logger = Arc::new(ServerLogger::default());
523 let stream_1 = StreamId { value: 1 };
524 let stream_2 = StreamId { value: 2 };
525
526 let binding_1 = handle(stream_1);
527 let binding_2 = handle(stream_2);
528 let binding_3 = handle(stream_1);
529
530 let mut ms = MultiStream::new(logger, TestBackend, 4);
531 ms.resolve(stream_1, [].into_iter());
532 ms.resolve(stream_2, [].into_iter());
533
534 let analysis =
535 ms.update_shared_bindings(stream_1, [&binding_1, &binding_2, &binding_3].into_iter());
536
537 let mut expected = SharedBindingAnalysis::default();
538 expected.shared(
539 binding_2.memory.descriptor().id,
540 ms.streams.stream_index(&binding_2.stream),
541 );
542
543 assert_eq!(analysis, expected);
544 }
545
546 #[test_log::test]
547 fn test_analysis_no_shared() {
548 let logger = Arc::new(ServerLogger::default());
549 let stream_1 = StreamId { value: 1 };
550 let stream_2 = StreamId { value: 2 };
551
552 let binding_1 = handle(stream_1);
553 let binding_2 = handle(stream_1);
554 let binding_3 = handle(stream_1);
555
556 let mut ms = MultiStream::new(logger, TestBackend, MAX_STREAMS);
557 ms.resolve(stream_1, [].into_iter());
558 ms.resolve(stream_2, [].into_iter());
559
560 let analysis =
561 ms.update_shared_bindings(stream_1, [&binding_1, &binding_2, &binding_3].into_iter());
562
563 let expected = SharedBindingAnalysis::default();
564
565 assert_eq!(analysis, expected);
566 }
567
568 #[test_log::test]
569 fn test_state() {
570 let logger = Arc::new(ServerLogger::default());
571 let stream_1 = StreamId { value: 1 };
572 let stream_2 = StreamId { value: 2 };
573
574 let binding_1 = handle(stream_1);
575 let binding_2 = handle(stream_2);
576 let binding_3 = handle(stream_1);
577
578 let mut ms = MultiStream::new(logger, TestBackend, MAX_STREAMS);
579 ms.resolve(stream_1, [].into_iter());
580 ms.resolve(stream_2, [].into_iter());
581
582 ms.resolve(stream_1, [&binding_1, &binding_2, &binding_3].into_iter());
583
584 let stream1 = ms.streams.get_mut(&stream_1);
585 let index_2 = stream_index(&stream_2, MAX_STREAMS as usize);
586 assert_eq!(stream1.last_synced.get(&index_2), Some(&1));
587 assert_eq!(stream1.cursor, 2);
588
589 let stream2 = ms.streams.get_mut(&stream_2);
590 assert!(stream2.last_synced.is_empty());
591 assert_eq!(stream2.cursor, 1);
592 }
593
594 #[test_log::test]
595 fn test_cross_stream_binding_pinned_until_gc_event() {
596 let logger = Arc::new(ServerLogger::default());
597 let stream_1 = StreamId { value: 1 };
598 let stream_2 = StreamId { value: 2 };
599
600 let gate = Arc::new(AtomicBool::new(false));
601 let mut ms = MultiStream::new(logger, GatedBackend { gate: gate.clone() }, MAX_STREAMS);
602 ms.resolve(stream_1, [].into_iter());
603 ms.resolve(stream_2, [].into_iter());
604
605 let handle = Handle::new(service(), stream_1, 10);
606 let observer = handle.memory.clone();
607 let binding = handle.binding();
608
609 drop(ms.resolve(stream_2, [&binding].into_iter()));
610 drop(binding);
611
612 assert!(
616 !observer.is_free(),
617 "cross-stream binding must stay pinned while the consumer event is pending"
618 );
619
620 gate.store(true, Ordering::Release);
621 wait_until_free(&observer);
622 }
623
624 #[test_log::test]
625 fn test_already_synced_cross_stream_binding_still_pinned() {
626 let logger = Arc::new(ServerLogger::default());
627 let stream_1 = StreamId { value: 1 };
628 let stream_2 = StreamId { value: 2 };
629
630 let gate = Arc::new(AtomicBool::new(true));
631 let mut ms = MultiStream::new(logger, GatedBackend { gate: gate.clone() }, MAX_STREAMS);
632 ms.resolve(stream_1, [].into_iter());
633 ms.resolve(stream_2, [].into_iter());
634
635 let handle_1 = Handle::new(service(), stream_1, 10);
637 let binding_1 = handle_1.binding();
638 drop(ms.resolve(stream_2, [&binding_1].into_iter()));
639 drop(binding_1);
640
641 gate.store(false, Ordering::Release);
643
644 let handle_2 = Handle::new(service(), stream_1, 10);
645 let observer = handle_2.memory.clone();
646 let binding_2 = handle_2.binding();
647
648 let resolved = ms.resolve(stream_2, [&binding_2].into_iter());
653 assert!(resolved.analysis.slices.is_empty());
654 drop(resolved);
655 drop(binding_2);
656
657 assert!(
658 !observer.is_free(),
659 "already-synced cross-stream binding must still be pinned"
660 );
661
662 gate.store(true, Ordering::Release);
663 wait_until_free(&observer);
664 }
665
666 fn wait_until_free(observer: &crate::memory_management::ManagedMemoryHandle) {
667 let start = std::time::Instant::now();
668 while !observer.is_free() {
669 assert!(
670 start.elapsed() < std::time::Duration::from_secs(10),
671 "pinned binding was never released"
672 );
673 std::thread::yield_now();
674 }
675 }
676
677 fn handle(stream: StreamId) -> BufferBinding {
678 Handle::new(service(), stream, 10).binding()
679 }
680
681 struct TestBackend;
682
683 #[derive(Debug, Default)]
684 struct TestStream;
685
686 #[derive(Debug)]
687 struct TestEvent {}
688
689 struct GatedBackend {
692 gate: Arc<AtomicBool>,
693 }
694
695 #[derive(Debug, Default)]
696 struct GatedStream {
697 gate: Arc<AtomicBool>,
698 }
699
700 #[derive(Debug)]
701 struct GatedEvent {
702 gate: Arc<AtomicBool>,
703 }
704
705 impl StreamMemory for GatedStream {
707 fn failure(&self, _binding: &BufferBinding) -> Option<FailureId> {
708 None
709 }
710
711 fn taint(
712 &mut self,
713 _binding: &BufferBinding,
714 _failure: FailureId,
715 _failures: &mut ErrorGraph,
716 ) {
717 }
718
719 fn written(&mut self, _binding: &BufferBinding, _failures: &mut ErrorGraph) {}
720 }
721
722 impl StreamMemory for TestStream {
723 fn failure(&self, _binding: &BufferBinding) -> Option<FailureId> {
724 None
725 }
726
727 fn taint(
728 &mut self,
729 _binding: &BufferBinding,
730 _failure: FailureId,
731 _failures: &mut ErrorGraph,
732 ) {
733 }
734
735 fn written(&mut self, _binding: &BufferBinding, _failures: &mut ErrorGraph) {}
736 }
737
738 impl EventStreamBackend for GatedBackend {
739 type Stream = GatedStream;
740 type Event = GatedEvent;
741
742 fn create_stream(&self) -> Self::Stream {
743 GatedStream {
744 gate: self.gate.clone(),
745 }
746 }
747
748 fn flush(stream: &mut Self::Stream, _failures: &mut ErrorGraph) -> Self::Event {
749 GatedEvent {
750 gate: stream.gate.clone(),
751 }
752 }
753
754 fn wait_event(_stream: &mut Self::Stream, _event: Self::Event) {}
755
756 fn wait_event_sync(event: Self::Event) -> Result<(), ServerError> {
757 while !event.gate.load(Ordering::Acquire) {
758 std::thread::yield_now();
759 }
760 Ok(())
761 }
762
763 fn handle_cursor(_stream: &Self::Stream, _handle: &BufferBinding) -> u64 {
764 0
765 }
766 }
767
768 impl EventStreamBackend for TestBackend {
769 type Stream = TestStream;
770 type Event = TestEvent;
771
772 fn create_stream(&self) -> Self::Stream {
773 TestStream
774 }
775
776 fn flush(_stream: &mut Self::Stream, _failures: &mut ErrorGraph) -> Self::Event {
777 TestEvent {}
778 }
779
780 fn wait_event(_stream: &mut Self::Stream, _event: Self::Event) {}
781
782 fn wait_event_sync(_event: Self::Event) -> Result<(), ServerError> {
783 Ok(())
784 }
785
786 fn handle_cursor(_stream: &Self::Stream, _handle: &BufferBinding) -> u64 {
787 0
788 }
789 }
790}