Skip to main content

vyre_runtime/megakernel/io/
complete.rs

1//! Completion-write helpers. Update the queue with the result of a
2//! serviced IO operation so the GPU sees COMPLETE for that slot.
3
4use std::sync::atomic::{fence, Ordering};
5
6use crate::PipelineError;
7
8use super::super::protocol::slot;
9use super::queue_words::{
10    read_queue_word, try_queue_word_index, validate_io_queue_view, write_queue_word_unfenced,
11    IoQueueView,
12};
13use super::{io_status, io_word};
14
15/// Strictly write a completion status for a serviced IO request.
16///
17/// # Errors
18///
19/// Returns [`PipelineError`] when the target slot is outside the queue byte
20/// view, the view is not aligned to complete IO slots, or the view exceeds the
21/// compiled poll window.
22pub fn try_complete_io_request(
23    io_queue_bytes: &mut [u8],
24    slot_idx: u32,
25    success: bool,
26) -> Result<(), PipelineError> {
27    try_complete_io_requests_batch(io_queue_bytes, &[(slot_idx, success)])
28}
29
30/// Strictly complete several claimed IO requests after one validation pass.
31///
32/// # Errors
33///
34/// Returns [`PipelineError`] without mutating the queue when any completion
35/// references an invalid slot or a slot not currently owned as `CLAIMED`.
36pub fn try_complete_io_requests_batch(
37    io_queue_bytes: &mut [u8],
38    completions: &[(u32, bool)],
39) -> Result<(), PipelineError> {
40    let view = validate_io_queue_view(io_queue_bytes.len())?;
41    if let Ok(words) = bytemuck::try_cast_slice_mut::<u8, u32>(io_queue_bytes) {
42        return complete_io_requests_words(words, view, completions);
43    }
44    for (slot_idx, _) in completions {
45        let base_word = completion_base_word(*slot_idx, view)?;
46        let current_status = read_queue_word(io_queue_bytes, base_word, io_word::STATUS)?;
47        if current_status != slot::CLAIMED {
48            return Err(PipelineError::QueueFull {
49                queue: "submission",
50                fix: "io_queue completion requires a CLAIMED request; poll with claim_io_requests_into before completing so the same DMA is not completed without ownership",
51            });
52        }
53    }
54    for (slot_idx, success) in completions {
55        let base_word = completion_base_word(*slot_idx, view)?;
56        let status = if *success {
57            io_status::OK
58        } else {
59            io_status::ERROR
60        };
61        write_queue_word_unfenced(io_queue_bytes, base_word, io_word::STATUS, status)?;
62    }
63    fence(Ordering::Release);
64    Ok(())
65}
66
67fn complete_io_requests_words(
68    words: &mut [u32],
69    view: IoQueueView,
70    completions: &[(u32, bool)],
71) -> Result<(), PipelineError> {
72    fence(Ordering::Acquire);
73    for (slot_idx, _) in completions {
74        let status_index = completion_status_word(*slot_idx, view)?;
75        let current_status = u32::from_le(*words.get(status_index).ok_or_else(|| {
76            PipelineError::Backend(format!(
77                "io_queue completion status word index {status_index} is outside the aligned word view. Fix: validate io_queue byte length before completion."
78            ))
79        })?);
80        if current_status != slot::CLAIMED {
81            return Err(PipelineError::QueueFull {
82                queue: "submission",
83                fix: "io_queue completion requires a CLAIMED request; poll with claim_io_requests_into before completing so the same DMA is not completed without ownership",
84            });
85        }
86    }
87    for (slot_idx, success) in completions {
88        let status_index = completion_status_word(*slot_idx, view)?;
89        let status = if *success {
90            io_status::OK
91        } else {
92            io_status::ERROR
93        };
94        *words.get_mut(status_index).ok_or_else(|| {
95            PipelineError::Backend(format!(
96                "io_queue completion status word index {status_index} is outside the aligned word view. Fix: validate io_queue byte length before completion."
97            ))
98        })? = status.to_le();
99    }
100    fence(Ordering::Release);
101    Ok(())
102}
103
104fn completion_base_word(slot_idx: u32, view: IoQueueView) -> Result<usize, PipelineError> {
105    let slot = usize::try_from(slot_idx).map_err(|error| {
106        PipelineError::Backend(format!(
107            "io_queue completion slot {slot_idx} cannot fit usize: {error}. Fix: shard completion batches before host processing."
108        ))
109    })?;
110    if slot >= view.slot_count {
111        return Err(PipelineError::QueueFull {
112            queue: "submission",
113            fix: "io_queue completion slot exceeds queue length; complete a valid slot id",
114        });
115    }
116    try_queue_word_index(slot_idx, 0)
117}
118
119fn completion_status_word(slot_idx: u32, view: IoQueueView) -> Result<usize, PipelineError> {
120    let _ = completion_base_word(slot_idx, view)?;
121    try_queue_word_index(slot_idx, io_word::STATUS)
122}
123
124/// Complete several serviced IO requests.
125///
126/// # Errors
127///
128/// See [`try_complete_io_requests_batch`].
129pub fn complete_io_requests_batch(
130    io_queue_bytes: &mut [u8],
131    completions: &[(u32, bool)],
132) -> Result<(), PipelineError> {
133    try_complete_io_requests_batch(io_queue_bytes, completions)
134}
135
136/// Write a completion status for a serviced IO request.
137///
138/// # Errors
139///
140/// Returns [`PipelineError`] when the target slot is outside the queue byte
141/// view, the view is not aligned to complete IO slots, or the view exceeds the
142/// compiled poll window.
143pub fn complete_io_request(
144    io_queue_bytes: &mut [u8],
145    slot_idx: u32,
146    success: bool,
147) -> Result<(), PipelineError> {
148    try_complete_io_request(io_queue_bytes, slot_idx, success)
149}