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        if queue_slot >= self.slot_count {
91            return Err(PipelineError::QueueFull {
92                queue: "submission",
93                fix: "io_queue slot exceeds MegakernelIoQueue::slot_count; enlarge the queue or publish into a valid slot id",
94            });
95        }
96        let current_status = self.read_word(queue_slot, io_word::STATUS)?;
97        if current_status != slot::EMPTY {
98            return Err(PipelineError::QueueFull {
99                queue: "submission",
100                fix: "io_queue slot still in flight; wait for the GPU to recycle it before publishing again",
101            });
102        }
103        self.write_word_unfenced(queue_slot, io_word::OP_TYPE, io_op::READ)?;
104        self.write_word_unfenced(queue_slot, io_word::SRC_HANDLE, 0)?;
105        self.write_word_unfenced(queue_slot, io_word::DST_HANDLE, mapped_slot)?;
106        self.write_word_unfenced(queue_slot, io_word::OFFSET_LO, 0)?;
107        self.write_word_unfenced(queue_slot, io_word::OFFSET_HI, 0)?;
108        self.write_word_unfenced(queue_slot, io_word::BYTE_COUNT, byte_count)?;
109        self.write_word_unfenced(queue_slot, io_word::TAG, tag)?;
110        fence(Ordering::Release);
111        self.write_word_unfenced(queue_slot, io_word::STATUS, slot::PUBLISHED)?;
112        fence(Ordering::Release);
113        Ok(())
114    }
115
116    /// Submit a DMA-read request to the IO queue.
117    ///
118    /// This is the GPU-initiated path: the caller writes the request metadata,
119    /// then flips `STATUS` to `slot::PUBLISHED` so the host/runtime can claim
120    /// and service it.
121    ///
122    /// # Errors
123    ///
124    /// Returns [`PipelineError::QueueFull`] when the slot is out of bounds or
125    /// not empty.
126    pub fn submit_dma_read(
127        &mut self,
128        queue_slot: u32,
129        src_handle: u32,
130        dst_handle: u32,
131        byte_count: u32,
132        tag: u32,
133    ) -> Result<(), PipelineError> {
134        if queue_slot >= self.slot_count {
135            return Err(PipelineError::QueueFull {
136                queue: "submission",
137                fix: "io_queue slot exceeds MegakernelIoQueue::slot_count; enlarge the queue or submit into a valid slot id",
138            });
139        }
140        let current_status = self.read_word(queue_slot, io_word::STATUS)?;
141        if current_status != slot::EMPTY {
142            return Err(PipelineError::QueueFull {
143                queue: "submission",
144                fix: "io_queue slot still in flight; wait for completion before submitting a new request",
145            });
146        }
147        self.write_word_unfenced(queue_slot, io_word::OP_TYPE, io_op::READ)?;
148        self.write_word_unfenced(queue_slot, io_word::SRC_HANDLE, src_handle)?;
149        self.write_word_unfenced(queue_slot, io_word::DST_HANDLE, dst_handle)?;
150        self.write_word_unfenced(queue_slot, io_word::OFFSET_LO, 0)?;
151        self.write_word_unfenced(queue_slot, io_word::OFFSET_HI, 0)?;
152        self.write_word_unfenced(queue_slot, io_word::BYTE_COUNT, byte_count)?;
153        self.write_word_unfenced(queue_slot, io_word::TAG, tag)?;
154        fence(Ordering::Release);
155        self.write_word_unfenced(queue_slot, io_word::STATUS, slot::PUBLISHED)?;
156        fence(Ordering::Release);
157        Ok(())
158    }
159
160    /// Read the queue slot back as a completion record.
161    #[must_use]
162    pub fn completion(&self, queue_slot: u32) -> Option<IoCompletion> {
163        if queue_slot >= self.slot_count {
164            return None;
165        }
166        // INVARIANT: queue_slot < slot_count <= IO_SLOT_COUNT (64) and io_word
167        // constants are all <= 7, so slot_idx * IO_SLOT_WORDS + word <= 511 which
168        // cannot overflow usize on any supported platform. The expect here documents
169        // that invariant loudly rather than silently reading the wrong word.
170        let status = self
171            .read_word(queue_slot, io_word::STATUS)
172            .expect("IO queue word index overflow is impossible after slot bounds check; Fix: ensure slot_count <= IO_SLOT_COUNT before calling completion");
173        if status == slot::EMPTY {
174            return None;
175        }
176        Some(IoCompletion {
177            slot_idx: queue_slot,
178            mapped_slot: self
179                .read_word_unfenced(queue_slot, io_word::DST_HANDLE)
180                .expect("IO queue word index overflow is impossible after slot bounds check; Fix: ensure slot_count <= IO_SLOT_COUNT before calling completion"),
181            byte_count: self
182                .read_word_unfenced(queue_slot, io_word::BYTE_COUNT)
183                .expect("IO queue word index overflow is impossible after slot bounds check; Fix: ensure slot_count <= IO_SLOT_COUNT before calling completion"),
184            tag: self
185                .read_word_unfenced(queue_slot, io_word::TAG)
186                .expect("IO queue word index overflow is impossible after slot bounds check; Fix: ensure slot_count <= IO_SLOT_COUNT before calling completion"),
187        })
188    }
189
190    /// Return true when the GPU has recycled the slot to `EMPTY`.
191    #[must_use]
192    pub fn is_recycled(&self, queue_slot: u32) -> bool {
193        if queue_slot >= self.slot_count {
194            return false;
195        }
196        // INVARIANT: same as completion, slot_idx * IO_SLOT_WORDS + word fits usize
197        // after the queue_slot < slot_count guard above.
198        let status = self
199            .read_word(queue_slot, io_word::STATUS)
200            .expect("IO queue word index overflow is impossible after slot bounds check; Fix: ensure slot_count <= IO_SLOT_COUNT before calling is_recycled");
201        match status {
202            slot::EMPTY => true,
203            slot::PUBLISHED | slot::CLAIMED | io_status::OK | io_status::ERROR | slot::DONE => {
204                false
205            }
206            _ => false,
207        }
208    }
209
210    fn read_word(&self, slot_idx: u32, word: u32) -> Result<u32, PipelineError> {
211        let idx = try_queue_word_index(slot_idx, word)?;
212        fence(Ordering::Acquire);
213        Ok(self.words[idx])
214    }
215
216    fn read_word_unfenced(&self, slot_idx: u32, word: u32) -> Result<u32, PipelineError> {
217        let idx = try_queue_word_index(slot_idx, word)?;
218        Ok(self.words[idx])
219    }
220
221    fn write_word_unfenced(
222        &mut self,
223        slot_idx: u32,
224        word: u32,
225        value: u32,
226    ) -> Result<(), PipelineError> {
227        let idx = try_queue_word_index(slot_idx, word)?;
228        self.words[idx] = value;
229        Ok(())
230    }
231}
vyre_runtime/megakernel/io/queue.rs

vyre_runtime/megakernel/io/
queue.rs
