cubecl_runtime/stream/
base.rs

1use alloc::vec::Vec;
2use cubecl_common::stream_id::StreamId;
3
4/// Trait for creating streams, used by the stream pool to generate streams as needed.
5pub trait StreamFactory {
6    /// The type of stream produced by this factory.
7    type Stream;
8    /// Creates a new stream instance.
9    fn create(&mut self) -> Self::Stream;
10}
11
12/// Represents a pool of streams, managing a collection of streams created by a factory.
13#[derive(Debug)]
14pub struct StreamPool<F: StreamFactory> {
15    /// Vector storing optional streams, where None indicates an uninitialized stream.
16    streams: Vec<Option<F::Stream>>,
17    /// The factory used to create new streams when needed.
18    factory: F,
19    /// Maximum number of regular streams (excludes special streams).
20    max_streams: usize,
21}
22
23impl<F: StreamFactory> StreamPool<F> {
24    /// Creates a new stream pool with the given backend factory and capacity constraints.
25    pub fn new(backend: F, max_streams: u8, num_special: u8) -> Self {
26        // Initialize a vector with capacity for regular and special streams.
27        let mut streams = Vec::with_capacity(max_streams as usize);
28        // Pre-populate the vector with None to reserve space for all streams.
29        for _ in 0..(max_streams.saturating_add(num_special)) {
30            streams.push(None);
31        }
32
33        Self {
34            streams,
35            factory: backend,
36            max_streams: max_streams as usize,
37        }
38    }
39
40    /// Read-only iterator over initialized streams (unlike [`Self::get_mut`], never creates one).
41    pub fn streams(&self) -> impl Iterator<Item = &F::Stream> {
42        self.streams.iter().flatten()
43    }
44
45    /// Synthetic [`StreamId`]s, one per initialized regular pool slot.
46    ///
47    /// Each id round-trips through [`Self::get_mut`] to the same slot it
48    /// came from (slot `i` is reachable via `StreamId { value: i }` since
49    /// indexing is `value % max_streams`), so it's safe to feed these
50    /// ids back into per-stream APIs.
51    pub fn stream_ids(&self) -> impl Iterator<Item = StreamId> + '_ {
52        self.streams[..self.max_streams]
53            .iter()
54            .enumerate()
55            .filter_map(|(i, s)| s.as_ref().map(|_| StreamId { value: i as u64 }))
56    }
57
58    /// Retrieves a mutable reference to a stream for a given stream ID.
59    pub fn get_mut(&mut self, stream_id: &StreamId) -> &mut F::Stream {
60        // Calculate the index for the stream ID.
61        let index = self.stream_index(stream_id);
62
63        // Use unsafe method to retrieve the stream, assuming the index is valid.
64        //
65        // # Safety
66        //
67        // * The `stream_index` function ensures the index is within bounds.
68        unsafe { self.get_mut_index(index) }
69    }
70
71    /// Retrieves a mutable reference to a stream at the specified index, initializing it if needed.
72    ///
73    /// # Safety
74    ///
75    /// * Caller must ensure the index is valid (less than `max_streams + num_special`).
76    /// * Lifetimes still follow the Rust rules.
77    pub unsafe fn get_mut_index(&mut self, index: usize) -> &mut F::Stream {
78        unsafe {
79            // Access the stream entry without bounds checking for performance.
80            let entry = self.streams.get_unchecked_mut(index);
81            match entry {
82                // If the stream exists, return it.
83                Some(val) => val,
84                // If the stream is None, create a new one using the factory.
85                None => {
86                    let stream = self.factory.create();
87                    // Store the new stream in the vector.
88                    *entry = Some(stream);
89
90                    // Re-access the entry, which is now guaranteed to be Some.
91                    match entry {
92                        Some(val) => val,
93                        // Unreachable because we just set it to Some.
94                        None => unreachable!(),
95                    }
96                }
97            }
98        }
99    }
100
101    /// Retrieves a mutable reference to a special stream at the given index.
102    ///
103    /// # Safety
104    ///
105    /// * Caller must ensure the index corresponds to a valid special stream.
106    /// * Lifetimes still follow the Rust rules.
107    pub unsafe fn get_special(&mut self, index: u8) -> &mut F::Stream {
108        // Calculate the index for the special stream (offset by max_streams).
109        unsafe { self.get_mut_index(self.max_streams + index as usize) }
110    }
111
112    /// Calculates the index for a given stream ID, mapping it to the pool's capacity.
113    pub fn stream_index(&mut self, id: &StreamId) -> usize {
114        stream_index(id, self.max_streams)
115    }
116
117    /// Mutable access to the factory, e.g. to change the configuration new
118    /// streams are created with. Already-created streams are unaffected.
119    pub fn factory_mut(&mut self) -> &mut F {
120        &mut self.factory
121    }
122}
123
124/// Maps a stream ID to an index within the pool's capacity using modulo arithmetic.
125pub fn stream_index(stream_id: &StreamId, max_streams: usize) -> usize {
126    stream_id.value as usize % max_streams
127}
cubecl_runtime/stream/base.rs

cubecl_runtime/stream/
base.rs
