vyre_runtime/
scheduler.rs

1//! Multi-GPU work stealing scheduler (Innovation I.7).
2//!
3//! Partitions a large Program or batch of Programs across all
4//! registered physical devices.
5
6use std::ops::Range;
7use std::sync::atomic::{AtomicUsize, Ordering};
8use std::sync::Arc;
9use vyre_driver::{BackendError, VyreBackend};
10
11/// A unit of work assigned to one GPU.
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub struct Shard {
14    /// Stable backend identifier for the GPU backend receiving this shard.
15    pub backend_id: &'static str,
16    /// Half-open byte/item range assigned to the backend.
17    pub work_range: Range<usize>,
18}
19
20/// Dynamic work-stealing scheduler.
21pub struct WorkStealingScheduler {
22    backends: Vec<Arc<dyn VyreBackend>>,
23    /// Atomic work index used by dispatch loops to let fast backends
24    /// steal more fine-grained work units. Worker threads call
25    /// [`Self::claim_next_unit`] which atomically increments the index;
26    /// the returned value is the unit index they own. This is the
27    /// work-stealing primitive  -  fast backends pull more units, slow
28    /// backends pull fewer.
29    work_index: AtomicUsize,
30}
31
32impl WorkStealingScheduler {
33    /// Create a scheduler over the live runtime backends available to the process.
34    pub fn new(backends: Vec<Arc<dyn VyreBackend>>) -> Self {
35        Self {
36            backends,
37            work_index: AtomicUsize::new(0),
38        }
39    }
40
41    /// Partition a large haystack across available GPUs.
42    pub fn partition(&self, total_len: usize) -> Vec<Shard> {
43        match self.try_partition(total_len) {
44            Ok(shards) => shards,
45            Err(_error) => Vec::new(),
46        }
47    }
48
49    /// Partition a large haystack across available GPUs with explicit staging
50    /// allocation failure reporting.
51    pub fn try_partition(&self, total_len: usize) -> Result<Vec<Shard>, BackendError> {
52        let mut shards = Vec::new();
53        self.try_partition_into(total_len, &mut shards)?;
54        Ok(shards)
55    }
56
57    /// Atomically claim the next fine-grained work unit. Worker threads
58    /// call this in a loop; the returned value is the unit index they
59    /// own. When the returned index is `>= num_units`, the worker is
60    /// done. This is the work-stealing primitive: fast backends call
61    /// `claim_next_unit` more times in the same wall-clock window.
62    ///
63    /// # Examples
64    ///
65    /// ```
66    /// use vyre_runtime::scheduler::WorkStealingScheduler;
67    /// let scheduler = WorkStealingScheduler::new(Vec::new());
68    /// assert_eq!(scheduler.claim_next_unit(), 0);
69    /// assert_eq!(scheduler.claim_next_unit(), 1);
70    /// scheduler.reset_unit_cursor();
71    /// assert_eq!(scheduler.claim_next_unit(), 0);
72    /// ```
73    #[must_use]
74    pub fn claim_next_unit(&self) -> usize {
75        self.work_index.fetch_add(1, Ordering::AcqRel)
76    }
77
78    /// Reset the work-unit cursor to zero. Call between dispatches that
79    /// reuse the same scheduler.
80    pub fn reset_unit_cursor(&self) {
81        self.work_index.store(0, Ordering::Release);
82    }
83
84    /// Partition a large haystack into many fine-grained work units
85    /// assigned round-robin to backends. A caller-side dispatch loop
86    /// uses [`Self::claim_next_unit`] to let worker threads atomically
87    /// claim units so fast backends steal more work.
88    ///
89    /// # Panics
90    /// Panics when `total_len` cannot be partitioned into work units. Callers that must
91    /// recover use the `try_` twin.
92    pub fn partition_into(&self, total_len: usize, out: &mut Vec<Shard>) {
93        // Clearing to empty on failure leaves NO work units, the dispatch loop
94        // then claims nothing and the scan silently processes nothing (Law 10
95        // recall loss). Fail loud; callers use try_partition_into.
96        if let Err(error) = self.try_partition_into(total_len, out) {
97            panic!("vyre-runtime work-unit partition failed: {error}");
98        }
99    }
100
101    /// Partition into caller-owned storage with explicit staging allocation
102    /// failure reporting.
103    pub fn try_partition_into(
104        &self,
105        total_len: usize,
106        out: &mut Vec<Shard>,
107    ) -> Result<(), BackendError> {
108        let n = self.backends.len();
109        out.clear();
110        if n == 0 || total_len == 0 {
111            return Ok(());
112        }
113        let work_unit_size = partition_work_unit_size(total_len, n);
114        let num_units = total_len.div_ceil(work_unit_size);
115        vyre_foundation::allocation::try_reserve_vec_to_capacity(out, num_units).map_err(
116            |error| BackendError::InvalidProgram {
117                fix: format!(
118                    "Fix: scheduler could not reserve {num_units} GPU work shard(s): {error}. Shard the workload before work-stealing partitioning."
119                ),
120            },
121        )?;
122        let mut start = 0;
123        for i in 0..num_units {
124            let end = (start + work_unit_size).min(total_len);
125            out.push(Shard {
126                backend_id: self.backends[i % n].id(),
127                work_range: start..end,
128            });
129            start = end;
130        }
131        Ok(())
132    }
133}
134
135fn partition_work_unit_size(total_len: usize, backend_count: usize) -> usize {
136    if total_len == 0 || backend_count == 0 {
137        return 1;
138    }
139    let denominator = backend_count.checked_mul(4).unwrap_or(usize::MAX);
140    (total_len / denominator.max(1)).max(1)
141}
142
143#[cfg(test)]
144fn partition_ranges(total_len: usize, backend_count: usize) -> Vec<Range<usize>> {
145    if backend_count == 0 || total_len == 0 {
146        return Vec::new();
147    }
148    let work_unit_size = partition_work_unit_size(total_len, backend_count);
149    let num_units = total_len.div_ceil(work_unit_size);
150    let mut ranges = Vec::with_capacity(num_units);
151    let mut start = 0;
152    for _ in 0..num_units {
153        let end = (start + work_unit_size).min(total_len);
154        ranges.push(start..end);
155        start = end;
156    }
157    ranges
158}
159
160#[cfg(test)]
161mod tests {
162    use super::{partition_ranges, WorkStealingScheduler};
163    use std::sync::Arc;
164    use vyre_driver::backend::{DispatchConfig, VyreBackend};
165    use vyre_foundation::ir::Program;
166
167    struct TestBackend(&'static str);
168
169    impl vyre_driver::backend::private::Sealed for TestBackend {}
170
171    impl VyreBackend for TestBackend {
172        fn id(&self) -> &'static str {
173            self.0
174        }
175
176        fn dispatch(
177            &self,
178            _program: &Program,
179            _inputs: &[Vec<u8>],
180            _config: &DispatchConfig,
181        ) -> Result<Vec<Vec<u8>>, vyre_driver::BackendError> {
182            Ok(Vec::new())
183        }
184    }
185
186    #[test]
187    fn partition_ranges_produces_fine_grained_units() {
188        let ranges = partition_ranges(10, 3);
189        assert_eq!(ranges.len(), 10);
190        assert_eq!(
191            ranges,
192            vec![0..1, 1..2, 2..3, 3..4, 4..5, 5..6, 6..7, 7..8, 8..9, 9..10]
193        );
194    }
195
196    #[test]
197    fn partition_ranges_never_emits_empty_shards() {
198        let ranges = partition_ranges(2, 8);
199        assert_eq!(ranges, vec![0..1, 1..2]);
200    }
201
202    #[test]
203    fn partition_ranges_uses_overflow_safe_work_unit_math() {
204        let ranges = partition_ranges(2, usize::MAX);
205        assert_eq!(ranges[0], 0..1);
206        assert_eq!(ranges[1], 1..2);
207        assert_eq!(
208            super::partition_work_unit_size(2, usize::MAX),
209            1,
210            "backend_count * 4 overflow must not panic or enlarge the work unit"
211        );
212    }
213
214    #[test]
215    fn scheduler_partition_into_reuses_output_storage() {
216        let scheduler = WorkStealingScheduler::new(vec![
217            Arc::new(TestBackend("a")),
218            Arc::new(TestBackend("b")),
219            Arc::new(TestBackend("c")),
220        ]);
221        let mut shards = Vec::with_capacity(10);
222
223        scheduler.partition_into(10, &mut shards);
224        let ptr = shards.as_ptr();
225        scheduler.partition_into(10, &mut shards);
226
227        assert_eq!(shards.as_ptr(), ptr);
228        assert_eq!(shards.len(), 10);
229        assert_eq!(shards[0].backend_id, "a");
230        assert_eq!(shards[0].work_range, 0..1);
231        assert_eq!(shards[1].backend_id, "b");
232        assert_eq!(shards[1].work_range, 1..2);
233        assert_eq!(shards[9].backend_id, "a");
234        assert_eq!(shards[9].work_range, 9..10);
235        assert_eq!(
236            scheduler
237                .work_index
238                .load(std::sync::atomic::Ordering::Relaxed),
239            0
240        );
241    }
242}
vyre_runtime/scheduler.rs

vyre_runtime/
scheduler.rs
