Skip to main content

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 into a fixed array without allocation.
8#[must_use]
9pub fn default_priority_offsets_array(total_slots: u32) -> [u32; PRIORITY_OFFSETS_WITH_SENTINEL] {
10    let mut offsets = [0u32; PRIORITY_OFFSETS_WITH_SENTINEL];
11    write_default_priority_offsets_array(total_slots, &mut offsets);
12    offsets
13}
14
15fn write_default_priority_offsets_array(
16    total_slots: u32,
17    offsets: &mut [u32; PRIORITY_OFFSETS_WITH_SENTINEL],
18) {
19    let base_per_pri = total_slots / PRIORITY_LEVELS;
20    let remainder = total_slots % PRIORITY_LEVELS;
21    let mut cursor = 0u32;
22    for pri in 0..PRIORITY_LEVELS_USIZE {
23        offsets[pri] = cursor;
24        let pri_u32 = pri as u32;
25        let size = base_per_pri
26            + if pri_u32 == priority::NORMAL {
27                remainder
28            } else {
29                0
30            };
31        cursor = cursor.saturating_add(size);
32    }
33    offsets[PRIORITY_LEVELS_USIZE] = cursor;
34}
35
36/// Write default priority partition offsets into an encoded control buffer.
37///
38/// # Errors
39///
40/// Returns [`PipelineError::QueueFull`] when the provided control buffer is too
41/// short or not aligned to u32 words.
42pub fn write_default_priority_offsets(
43    control_bytes: &mut [u8],
44    total_slots: u32,
45) -> Result<(), PipelineError> {
46    if control_bytes.len() % 4 != 0 {
47        return Err(PipelineError::QueueFull {
48            queue: "submission",
49            fix: "control buffer byte length is not 4-byte aligned; rebuild it with Megakernel::encode_control",
50        });
51    }
52    let mut offsets = [0u32; PRIORITY_OFFSETS_WITH_SENTINEL];
53    write_default_priority_offsets_array(total_slots, &mut offsets);
54    for (i, value) in offsets.iter().enumerate() {
55        let word_idx = priority_offsets_base_usize()?.checked_add(i).ok_or(
56            PipelineError::QueueFull {
57                queue: "submission",
58                fix: "priority-offset control word index overflowed usize; keep control ABI constants bounded",
59            },
60        )?;
61        let start = word_idx.checked_mul(4).ok_or(PipelineError::QueueFull {
62            queue: "submission",
63            fix: "priority-offset byte index overflowed usize; keep control ABI constants bounded",
64        })?;
65        let end = start.checked_add(4).ok_or(PipelineError::QueueFull {
66            queue: "submission",
67            fix: "priority-offset byte index overflowed usize; keep control ABI constants bounded",
68        })?;
69        let dst = control_bytes.get_mut(start..end).ok_or(PipelineError::QueueFull {
70            queue: "submission",
71            fix: "control buffer is too small for priority partition offsets; rebuild it with Megakernel::encode_control",
72        })?;
73        dst.copy_from_slice(&value.to_le_bytes());
74    }
75    Ok(())
76}
77
78fn priority_offsets_base_usize() -> Result<usize, PipelineError> {
79    usize::try_from(PRIORITY_OFFSETS_BASE).map_err(|_| PipelineError::QueueFull {
80        queue: "submission",
81        fix: "priority-offset base word cannot fit host usize; keep control ABI constants bounded",
82    })
83}