vyre_driver_wgpu/engine/multi_gpu/
partition.rs

1use std::cmp::{Ordering, Reverse};
2use std::collections::BinaryHeap;
3
4/// One pending unit of GPU work.
5#[derive(Clone, Debug, Eq, PartialEq)]
6pub struct WeightedWorkItem {
7    /// Stable work identifier used by callers to map results back.
8    pub id: usize,
9    /// Relative cost estimate. Zero-cost work is rejected because it cannot
10    /// contribute to a meaningful load balance.
11    pub cost: u64,
12}
13
14/// Current device load snapshot.
15#[derive(Clone, Debug, Eq, PartialEq)]
16pub struct DeviceLoad {
17    /// Device ordinal in the caller's adapter list.
18    pub device_index: usize,
19    /// Cost already queued on the device before this partitioning pass.
20    pub queued_cost: u64,
21}
22
23/// Work assigned to one device.
24#[derive(Clone, Debug, Eq, PartialEq)]
25pub struct Partition {
26    /// Device ordinal receiving this partition.
27    pub device_index: usize,
28    /// Work item identifiers assigned to the device.
29    pub item_ids: Vec<usize>,
30    /// Total assigned cost including pre-existing queued cost.
31    pub total_cost: u64,
32}
33
34#[derive(Clone, Copy, Debug, Eq, PartialEq)]
35struct PartitionHeapEntry {
36    total_cost: u64,
37    device_index: usize,
38    partition_index: usize,
39}
40
41impl Ord for PartitionHeapEntry {
42    fn cmp(&self, other: &Self) -> Ordering {
43        self.total_cost
44            .cmp(&other.total_cost)
45            .then_with(|| self.device_index.cmp(&other.device_index))
46            .then_with(|| self.partition_index.cmp(&other.partition_index))
47    }
48}
49
50impl PartialOrd for PartitionHeapEntry {
51    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
52        Some(self.cmp(other))
53    }
54}
55
56/// Partition work by repeatedly assigning the largest remaining item to the
57/// least-loaded device.
58///
59/// # Errors
60///
61/// Returns an actionable error when no devices are available, duplicate device
62/// ordinals are supplied, or a work item has zero cost.
63pub fn partition_work_stealing(
64    devices: &[DeviceLoad],
65    items: &[WeightedWorkItem],
66) -> Result<Vec<Partition>, String> {
67    validate_inputs(devices, items)?;
68    let mut partitions = Vec::new();
69    vyre_driver::allocation::try_reserve_vec_to_capacity(&mut partitions, devices.len()).map_err(|error| {
70        format!(
71            "partition table allocation failed for {} GPU devices: {error}. Fix: lower the adapter fanout or memory pressure before scheduling.",
72            devices.len()
73        )
74    })?;
75    let mut least_loaded = BinaryHeap::new();
76    vyre_foundation::allocation::try_reserve_binary_heap_to_capacity(
77        &mut least_loaded,
78        devices.len(),
79    )
80    .map_err(|error| {
81        format!(
82            "partition heap allocation failed for {} GPU devices: {error}. Fix: lower the adapter fanout or memory pressure before scheduling.",
83            devices.len()
84        )
85    })?;
86    let target_item_capacity = items.len().div_ceil(devices.len());
87    for device in devices {
88        let partition_index = partitions.len();
89        let mut item_ids = Vec::new();
90        vyre_driver::allocation::try_reserve_vec_to_capacity(&mut item_ids, target_item_capacity).map_err(|error| {
91            format!(
92                "partition assignment allocation failed for GPU device {} and {} target work slots: {error}. Fix: split the multi-GPU batch.",
93                device.device_index, target_item_capacity
94            )
95        })?;
96        partitions.push(Partition {
97            device_index: device.device_index,
98            item_ids,
99            total_cost: device.queued_cost,
100        });
101        least_loaded.push(Reverse(PartitionHeapEntry {
102            total_cost: device.queued_cost,
103            device_index: device.device_index,
104            partition_index,
105        }));
106    }
107
108    let mut ordered = Vec::new();
109    vyre_driver::allocation::try_reserve_vec_to_capacity(&mut ordered, items.len()).map_err(|error| {
110        format!(
111            "partition work-order allocation failed for {} items: {error}. Fix: split the multi-GPU batch.",
112            items.len()
113        )
114    })?;
115    ordered.extend(items.iter());
116    ordered.sort_by(|left, right| {
117        right
118            .cost
119            .cmp(&left.cost)
120            .then_with(|| left.id.cmp(&right.id))
121    });
122
123    for item in ordered {
124        let Some(mut target) = least_loaded.pop().map(|entry| entry.0) else {
125            return Err(
126                "partition target not found. Fix: validate non-empty device list before partitioning."
127                    .to_string()
128            );
129        };
130        let partition = &mut partitions[target.partition_index];
131        ensure_vec_spare(&mut partition.item_ids, 1, "partition assignment list")?;
132        partition.item_ids.push(item.id);
133        partition.total_cost = partition.total_cost.checked_add(item.cost).ok_or_else(|| {
134            "partition cost overflow. Fix: split the batch before multi-GPU scheduling.".to_string()
135        })?;
136        target.total_cost = partition.total_cost;
137        least_loaded.push(Reverse(target));
138    }
139    Ok(partitions)
140}
141
142fn validate_inputs(devices: &[DeviceLoad], items: &[WeightedWorkItem]) -> Result<(), String> {
143    if devices.is_empty() {
144        return Err(
145            "no GPU devices supplied. Fix: probe adapters before partitioning.".to_string(),
146        );
147    }
148    let mut seen = rustc_hash::FxHashSet::default();
149    vyre_foundation::allocation::try_reserve_hash_set_to_capacity(&mut seen, devices.len()).map_err(|error| {
150        format!(
151            "GPU device validation allocation failed for {} devices: {error}. Fix: lower adapter fanout or memory pressure before scheduling.",
152            devices.len()
153        )
154    })?;
155    for device in devices {
156        if !seen.insert(device.device_index) {
157            return Err(format!(
158                "duplicate GPU device index {}. Fix: pass each adapter exactly once.",
159                device.device_index
160            ));
161        }
162    }
163    for item in items {
164        if item.cost == 0 {
165            return Err(format!(
166                "work item {} has zero cost. Fix: assign at least one cost unit or remove it.",
167                item.id
168            ));
169        }
170    }
171    let mut seen_items = rustc_hash::FxHashSet::default();
172    vyre_foundation::allocation::try_reserve_hash_set_to_capacity(&mut seen_items, items.len()).map_err(|error| {
173        format!(
174            "work item validation allocation failed for {} items: {error}. Fix: split the multi-GPU batch.",
175            items.len()
176        )
177    })?;
178    for item in items {
179        if !seen_items.insert(item.id) {
180            return Err(format!(
181                "duplicate work item id {}. Fix: assign a unique stable id to every multi-GPU work item.",
182                item.id
183            ));
184        }
185    }
186    Ok(())
187}
188
189fn ensure_vec_spare<T>(vec: &mut Vec<T>, additional: usize, label: &str) -> Result<(), String> {
190    let spare = vec.capacity().saturating_sub(vec.len());
191    if spare >= additional {
192        return Ok(());
193    }
194    vec.try_reserve_exact(additional - spare).map_err(|error| {
195        format!(
196            "{label} allocation failed while extending by {additional} entries: {error}. Fix: split the multi-GPU batch."
197        )
198    })
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    #[test]
206    fn multi_gpu_partition_unit() {
207        let devices = [
208            DeviceLoad {
209                device_index: 0,
210                queued_cost: 0,
211            },
212            DeviceLoad {
213                device_index: 1,
214                queued_cost: 4,
215            },
216        ];
217        let items = [
218            WeightedWorkItem { id: 10, cost: 9 },
219            WeightedWorkItem { id: 11, cost: 4 },
220            WeightedWorkItem { id: 12, cost: 4 },
221            WeightedWorkItem { id: 13, cost: 1 },
222        ];
223
224        let partitions = partition_work_stealing(&devices, &items)
225            .expect("Fix: valid synthetic device loads must partition");
226        let mut assigned = partitions
227            .iter()
228            .flat_map(|partition| partition.item_ids.iter().copied())
229            .collect::<Vec<_>>();
230        assigned.sort_unstable();
231
232        assert_eq!(assigned, vec![10, 11, 12, 13]);
233        let spread = partitions
234            .iter()
235            .map(|partition| partition.total_cost)
236            .max()
237            .zip(
238                partitions
239                    .iter()
240                    .map(|partition| partition.total_cost)
241                    .min(),
242            )
243            .map(|(max, min)| max - min)
244            .expect("Fix: partitions must be non-empty");
245        assert!(
246            spread <= 5,
247            "synthetic work stealing left an avoidable load spread: {partitions:?}"
248        );
249    }
250
251    #[test]
252    fn rejects_duplicate_device_ordinals() {
253        let devices = [
254            DeviceLoad {
255                device_index: 0,
256                queued_cost: 0,
257            },
258            DeviceLoad {
259                device_index: 0,
260                queued_cost: 1,
261            },
262        ];
263
264        let error = partition_work_stealing(&devices, &[WeightedWorkItem { id: 1, cost: 1 }])
265            .expect_err("Fix: duplicate synthetic device indices must be rejected");
266        assert!(error.contains("duplicate GPU device index"));
267    }
268
269    #[test]
270    fn partition_source_has_no_release_path_infallible_allocation() {
271        let source = include_str!("partition.rs");
272        let production = source
273            .split("#[cfg(test)]")
274            .next()
275            .expect("Fix: partition production source must precede tests");
276        assert!(
277            !production.contains("Vec::with_capacity")
278                && !production.contains("BinaryHeap::with_capacity")
279                && !production.contains("with_capacity_and_hasher"),
280            "Fix: WGPU multi-GPU partitioning must report allocation pressure instead of aborting on infallible capacity constructors."
281        );
282        assert!(
283            production.contains("try_reserve_vec_to_capacity")
284                && production.contains("try_reserve_binary_heap_to_capacity")
285                && production.contains("try_reserve_hash_set_to_capacity")
286                && production.contains("ensure_vec_spare"),
287            "Fix: WGPU multi-GPU partitioning must use fallible allocation for schedule tables and per-device assignments."
288        );
289    }
290}
vyre_driver_wgpu/engine/multi_gpu/partition.rs

vyre_driver_wgpu/engine/multi_gpu/
partition.rs
