vyre_runtime/megakernel/io/
queue.rs

1//! [`MegakernelIoQueue`]  -  high-level wrapper around the raw queue bytes.
2
3use std::sync::atomic::{fence, Ordering};
4
5use crate::PipelineError;
6
7use super::super::protocol::slot;
8use super::helpers::try_queue_word_index;
9use super::{io_op, io_status, io_word, IoCompletion, IO_SLOT_COUNT, IO_SLOT_WORDS};
10
11/// Host-side handle to the megakernel IO queue. Wraps a `Vec<u32>` slot ring
12/// and exposes typed poll/publish/complete entry points.
13#[derive(Debug, Clone)]
14pub struct MegakernelIoQueue {
15    words: Vec<u32>,
16    slot_count: u32,
17}
18
19impl MegakernelIoQueue {
20    /// Allocate an empty queue with `slot_count` entries.
21    ///
22    /// # Errors
23    ///
24    /// Returns [`PipelineError::QueueFull`] when `slot_count` is zero or
25    /// exceeds the IR/program's fixed poll window of [`IO_SLOT_COUNT`].
26    pub fn new(slot_count: u32) -> Result<Self, PipelineError> {
27        if slot_count == 0 {
28            return Err(PipelineError::QueueFull {
29                queue: "submission",
30                fix: "MegakernelIoQueue requires at least one slot",
31            });
32        }
33        if slot_count > IO_SLOT_COUNT {
34            return Err(PipelineError::QueueFull {
35                queue: "submission",
36                fix: "MegakernelIoQueue exceeds the compiled IO poll window of 64 slots; enlarge IO_SLOT_COUNT and rebuild the megakernel before publishing more than 64 completions",
37            });
38        }
39        let word_count = slot_count
40            .checked_mul(IO_SLOT_WORDS)
41            .ok_or(PipelineError::QueueFull {
42                queue: "submission",
43                fix: "io_queue word count overflows u32; shard the queue before allocating",
44            })?;
45        let word_count = usize::try_from(word_count).map_err(|error| {
46            PipelineError::Backend(format!(
47                "io_queue word count cannot fit host usize: {error}. Fix: shard the queue before allocating."
48            ))
49        })?;
50        Ok(Self {
51            words: vec![0; word_count],
52            slot_count,
53        })
54    }
55
56    /// Borrow the raw bytes for backend upload / readback.
57    #[must_use]
58    pub fn as_bytes(&self) -> &[u8] {
59        bytemuck::cast_slice(&self.words)
60    }
61
62    /// Mutably borrow the raw bytes for backend upload / host updates.
63    #[must_use]
64    pub fn as_mut_bytes(&mut self) -> &mut [u8] {
65        bytemuck::cast_slice_mut(&mut self.words)
66    }
67
68    /// Queue capacity in slots.
69    #[must_use]
70    pub fn slot_count(&self) -> u32 {
71        self.slot_count
72    }
73
74    /// Publish a completed DMA slot so the megakernel can consume it.
75    ///
76    /// The host writes the metadata first, then flips `STATUS` to
77    /// `slot::PUBLISHED` as the publication barrier.
78    ///
79    /// # Errors
80    ///
81    /// Returns [`PipelineError::QueueFull`] when the slot is out of bounds or
82    /// still owned by the GPU/host from a prior ingest.
83    pub fn publish_slot(
84        &mut self,
85        queue_slot: u32,
86        mapped_slot: u32,
87        byte_count: u32,
88        tag: u32,
89    ) -> Result<(), PipelineError> {
90        self.publish_dma_read(
91            queue_slot,
92            0,
93            mapped_slot,
94            byte_count,
95            tag,
96            "io_queue slot exceeds MegakernelIoQueue::slot_count; enlarge the queue or publish into a valid slot id",
97            "io_queue slot still in flight; wait for the GPU to recycle it before publishing again",
98        )
99    }
100
101    /// Submit a DMA-read request to the IO queue.
102    ///
103    /// This is the GPU-initiated path: the caller writes the request metadata,
104    /// then flips `STATUS` to `slot::PUBLISHED` so the host/runtime can claim
105    /// and service it.
106    ///
107    /// # Errors
108    ///
109    /// Returns [`PipelineError::QueueFull`] when the slot is out of bounds or
110    /// not empty.
111    pub fn submit_dma_read(
112        &mut self,
113        queue_slot: u32,
114        src_handle: u32,
115        dst_handle: u32,
116        byte_count: u32,
117        tag: u32,
118    ) -> Result<(), PipelineError> {
119        self.publish_dma_read(
120            queue_slot,
121            src_handle,
122            dst_handle,
123            byte_count,
124            tag,
125            "io_queue slot exceeds MegakernelIoQueue::slot_count; enlarge the queue or submit into a valid slot id",
126            "io_queue slot still in flight; wait for completion before submitting a new request",
127        )
128    }
129
130    #[allow(clippy::too_many_arguments)]
131    fn publish_dma_read(
132        &mut self,
133        queue_slot: u32,
134        src_handle: u32,
135        dst_handle: u32,
136        byte_count: u32,
137        tag: u32,
138        out_of_bounds_fix: &'static str,
139        in_flight_fix: &'static str,
140    ) -> Result<(), PipelineError> {
141        if queue_slot >= self.slot_count {
142            return Err(PipelineError::QueueFull {
143                queue: "submission",
144                fix: out_of_bounds_fix,
145            });
146        }
147        if self.read_word(queue_slot, io_word::STATUS)? != slot::EMPTY {
148            return Err(PipelineError::QueueFull {
149                queue: "submission",
150                fix: in_flight_fix,
151            });
152        }
153        self.write_word_unfenced(queue_slot, io_word::OP_TYPE, io_op::READ)?;
154        self.write_word_unfenced(queue_slot, io_word::SRC_HANDLE, src_handle)?;
155        self.write_word_unfenced(queue_slot, io_word::DST_HANDLE, dst_handle)?;
156        self.write_word_unfenced(queue_slot, io_word::OFFSET_LO, 0)?;
157        self.write_word_unfenced(queue_slot, io_word::OFFSET_HI, 0)?;
158        self.write_word_unfenced(queue_slot, io_word::BYTE_COUNT, byte_count)?;
159        self.write_word_unfenced(queue_slot, io_word::TAG, tag)?;
160        fence(Ordering::Release);
161        self.write_word_unfenced(queue_slot, io_word::STATUS, slot::PUBLISHED)?;
162        fence(Ordering::Release);
163        Ok(())
164    }
165
166    /// Read the queue slot back as a completion record.
167    ///
168    /// # Panics
169    /// Panics when a queue word index overflows, which the slot bounds check above already
170    /// rules out. Keep `slot_count <= IO_SLOT_COUNT`.
171    #[must_use]
172    pub fn completion(&self, queue_slot: u32) -> Option<IoCompletion> {
173        if queue_slot >= self.slot_count {
174            return None;
175        }
176        // INVARIANT: queue_slot < slot_count <= IO_SLOT_COUNT (64) and io_word
177        // constants are all <= 7, so slot_idx * IO_SLOT_WORDS + word <= 511 which
178        // cannot overflow usize on any supported platform. The expect here documents
179        // that invariant loudly rather than silently reading the wrong word.
180        let status = self
181            .read_word(queue_slot, io_word::STATUS)
182            .expect("IO queue word index overflow is impossible after slot bounds check; Fix: ensure slot_count <= IO_SLOT_COUNT before calling completion");
183        if status == slot::EMPTY {
184            return None;
185        }
186        Some(IoCompletion {
187            slot_idx: queue_slot,
188            mapped_slot: self
189                .read_word_unfenced(queue_slot, io_word::DST_HANDLE)
190                .expect("IO queue word index overflow is impossible after slot bounds check; Fix: ensure slot_count <= IO_SLOT_COUNT before calling completion"),
191            byte_count: self
192                .read_word_unfenced(queue_slot, io_word::BYTE_COUNT)
193                .expect("IO queue word index overflow is impossible after slot bounds check; Fix: ensure slot_count <= IO_SLOT_COUNT before calling completion"),
194            tag: self
195                .read_word_unfenced(queue_slot, io_word::TAG)
196                .expect("IO queue word index overflow is impossible after slot bounds check; Fix: ensure slot_count <= IO_SLOT_COUNT before calling completion"),
197        })
198    }
199
200    /// Return true when the GPU has recycled the slot to `EMPTY`.
201    ///
202    /// # Panics
203    /// Panics when a queue word index overflows, which the slot bounds check above already
204    /// rules out. Keep `slot_count <= IO_SLOT_COUNT`.
205    #[must_use]
206    pub fn is_recycled(&self, queue_slot: u32) -> bool {
207        if queue_slot >= self.slot_count {
208            return false;
209        }
210        // INVARIANT: same as completion, slot_idx * IO_SLOT_WORDS + word fits usize
211        // after the queue_slot < slot_count guard above.
212        let status = self
213            .read_word(queue_slot, io_word::STATUS)
214            .expect("IO queue word index overflow is impossible after slot bounds check; Fix: ensure slot_count <= IO_SLOT_COUNT before calling is_recycled");
215        match status {
216            slot::EMPTY => true,
217            slot::PUBLISHED | slot::CLAIMED | io_status::OK | io_status::ERROR | slot::DONE => {
218                false
219            }
220            _ => false,
221        }
222    }
223
224    fn read_word(&self, slot_idx: u32, word: u32) -> Result<u32, PipelineError> {
225        let idx = try_queue_word_index(slot_idx, word)?;
226        fence(Ordering::Acquire);
227        Ok(self.words[idx])
228    }
229
230    fn read_word_unfenced(&self, slot_idx: u32, word: u32) -> Result<u32, PipelineError> {
231        let idx = try_queue_word_index(slot_idx, word)?;
232        Ok(self.words[idx])
233    }
234
235    fn write_word_unfenced(
236        &mut self,
237        slot_idx: u32,
238        word: u32,
239        value: u32,
240    ) -> Result<(), PipelineError> {
241        let idx = try_queue_word_index(slot_idx, word)?;
242        self.words[idx] = value;
243        Ok(())
244    }
245}
vyre_runtime/megakernel/io/queue.rs

vyre_runtime/megakernel/io/
queue.rs
