vyre_runtime/megakernel/io/
complete.rs1use 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
15pub 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
30pub 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
124pub 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
136pub 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}