vyre_runtime/megakernel/scheduler/
offsets.rs

1use super::{priority, PRIORITY_LEVELS, PRIORITY_OFFSETS_BASE};
2use crate::PipelineError;
3
4const PRIORITY_LEVELS_USIZE: usize = 5;
5const PRIORITY_OFFSETS_WITH_SENTINEL: usize = PRIORITY_LEVELS_USIZE + 1;
6
7/// Encode default priority partition offsets for uniform distribution.
8///
9/// Each priority level gets `total_slots / PRIORITY_LEVELS` slots.
10/// Any remainder goes to the NORMAL partition.
11#[must_use]
12pub fn default_priority_offsets(total_slots: u32) -> Vec<u32> {
13    default_priority_offsets_array(total_slots).to_vec()
14}
15
16/// Encode default priority partition offsets with fallible host staging.
17///
18/// # Errors
19///
20/// Returns [`PipelineError::Backend`] when the host cannot reserve the small
21/// compatibility vector used by legacy callers.
22pub fn try_default_priority_offsets(total_slots: u32) -> Result<Vec<u32>, PipelineError> {
23    let mut offsets = Vec::new();
24    vyre_foundation::allocation::try_reserve_vec_to_capacity(
25        &mut offsets,
26        PRIORITY_OFFSETS_WITH_SENTINEL,
27    )
28    .map_err(|source| {
29            PipelineError::Backend(format!(
30                "priority offset vector reservation failed for {PRIORITY_OFFSETS_WITH_SENTINEL} words: {source}. Fix: use default_priority_offsets_array in hot paths or reduce host memory pressure."
31            ))
32        })?;
33    for value in default_priority_offsets_array(total_slots) {
34        offsets.push(value);
35    }
36    Ok(offsets)
37}
38
39/// Encode default priority partition offsets into a fixed array.
40///
41/// Hot callers that immediately write offsets into a control buffer can use
42/// this path to avoid allocating the compatibility `Vec` returned by
43/// [`default_priority_offsets`].
44#[must_use]
45pub fn default_priority_offsets_array(total_slots: u32) -> [u32; PRIORITY_OFFSETS_WITH_SENTINEL] {
46    let mut offsets = [0u32; PRIORITY_OFFSETS_WITH_SENTINEL];
47    write_default_priority_offsets_array(total_slots, &mut offsets);
48    offsets
49}
50
51fn write_default_priority_offsets_array(
52    total_slots: u32,
53    offsets: &mut [u32; PRIORITY_OFFSETS_WITH_SENTINEL],
54) {
55    let base_per_pri = total_slots / PRIORITY_LEVELS;
56    let remainder = total_slots % PRIORITY_LEVELS;
57    let mut cursor = 0u32;
58    for pri in 0..PRIORITY_LEVELS_USIZE {
59        offsets[pri] = cursor;
60        let pri_u32 = pri as u32;
61        let size = base_per_pri
62            + if pri_u32 == priority::NORMAL {
63                remainder
64            } else {
65                0
66            };
67        cursor = cursor.saturating_add(size);
68    }
69    offsets[PRIORITY_LEVELS_USIZE] = cursor;
70}
71
72/// Write default priority partition offsets into an encoded control buffer.
73///
74/// # Errors
75///
76/// Returns [`PipelineError::QueueFull`] when the provided control buffer is too
77/// short or not aligned to u32 words.
78pub fn write_default_priority_offsets(
79    control_bytes: &mut [u8],
80    total_slots: u32,
81) -> Result<(), PipelineError> {
82    if control_bytes.len() % 4 != 0 {
83        return Err(PipelineError::QueueFull {
84            queue: "submission",
85            fix: "control buffer byte length is not 4-byte aligned; rebuild it with Megakernel::encode_control",
86        });
87    }
88    let mut offsets = [0u32; PRIORITY_OFFSETS_WITH_SENTINEL];
89    write_default_priority_offsets_array(total_slots, &mut offsets);
90    for (i, value) in offsets.iter().enumerate() {
91        let word_idx = priority_offsets_base_usize()?.checked_add(i).ok_or(
92            PipelineError::QueueFull {
93                queue: "submission",
94                fix: "priority-offset control word index overflowed usize; keep control ABI constants bounded",
95            },
96        )?;
97        let start = word_idx.checked_mul(4).ok_or(PipelineError::QueueFull {
98            queue: "submission",
99            fix: "priority-offset byte index overflowed usize; keep control ABI constants bounded",
100        })?;
101        let end = start.checked_add(4).ok_or(PipelineError::QueueFull {
102            queue: "submission",
103            fix: "priority-offset byte index overflowed usize; keep control ABI constants bounded",
104        })?;
105        let dst = control_bytes.get_mut(start..end).ok_or(PipelineError::QueueFull {
106            queue: "submission",
107            fix: "control buffer is too small for priority partition offsets; rebuild it with Megakernel::encode_control",
108        })?;
109        dst.copy_from_slice(&value.to_le_bytes());
110    }
111    Ok(())
112}
113
114fn priority_offsets_base_usize() -> Result<usize, PipelineError> {
115    usize::try_from(PRIORITY_OFFSETS_BASE).map_err(|_| PipelineError::QueueFull {
116        queue: "submission",
117        fix: "priority-offset base word cannot fit host usize; keep control ABI constants bounded",
118    })
119}
vyre_runtime/megakernel/scheduler/offsets.rs

vyre_runtime/megakernel/scheduler/
offsets.rs
