ferrum_interfaces/vnext/completion/
readback_collection.rs1use std::cmp::Ordering;
2
3use super::{
4 invalid_completion, BatchOperationIdentity, CompletionReadbackBatchObservation,
5 CompletionReadbackBatchReceipt, CompletionReadbackBatchRequest, CompletionReadbackRequest,
6 VNextError,
7};
8
9#[derive(Debug, Clone, PartialEq, Eq)]
13#[must_use = "a completion readback collection must be consumed by one terminal wait"]
14pub struct CompletionReadbackCollectionRequest {
15 batches: Vec<CompletionReadbackBatchRequest>,
16}
17
18impl CompletionReadbackCollectionRequest {
19 pub fn new(mut batches: Vec<CompletionReadbackBatchRequest>) -> Result<Self, VNextError> {
20 if batches.is_empty() || u32::try_from(batches.len()).is_err() {
21 return Err(invalid_completion(
22 "completion readback collection is empty or its group count exceeds u32",
23 ));
24 }
25 let participant_count = batches[0].len();
26 if batches.iter().any(|batch| batch.len() != participant_count) {
27 return Err(invalid_completion(
28 "completion readback collection groups must cover the same participant count",
29 ));
30 }
31 batches.sort_by(compare_readback_batches);
32 if batches
33 .windows(2)
34 .any(|pair| compare_readback_batches(&pair[0], &pair[1]) == Ordering::Equal)
35 {
36 return Err(invalid_completion(
37 "completion readback collection contains a duplicate typed physical range",
38 ));
39 }
40 Ok(Self { batches })
41 }
42
43 pub fn batches(&self) -> &[CompletionReadbackBatchRequest] {
44 &self.batches
45 }
46
47 pub fn len(&self) -> usize {
48 self.batches.len()
49 }
50
51 pub fn is_empty(&self) -> bool {
52 self.batches.is_empty()
53 }
54
55 pub fn request_count(&self) -> usize {
56 self.batches
57 .iter()
58 .map(CompletionReadbackBatchRequest::len)
59 .sum()
60 }
61
62 pub(super) fn validate_for(
63 &self,
64 batch_identity: &BatchOperationIdentity,
65 ) -> Result<(), VNextError> {
66 for batch in &self.batches {
67 batch.validate_for(batch_identity)?;
68 }
69 Ok(())
70 }
71
72 pub(super) fn into_requests(self) -> Vec<CompletionReadbackRequest> {
73 self.batches
74 .into_iter()
75 .flat_map(CompletionReadbackBatchRequest::into_requests)
76 .collect()
77 }
78}
79
80fn compare_readback_batches(
81 left: &CompletionReadbackBatchRequest,
82 right: &CompletionReadbackBatchRequest,
83) -> Ordering {
84 let left_first = &left.requests()[0];
85 let right_first = &right.requests()[0];
86 let group_order = left_first
87 .node_id()
88 .cmp(right_first.node_id())
89 .then_with(|| left_first.resource_id().cmp(right_first.resource_id()))
90 .then_with(|| {
91 left_first
92 .expected_usage()
93 .cmp(&right_first.expected_usage())
94 })
95 .then_with(|| {
96 left_first
97 .logical_offset_bytes()
98 .cmp(&right_first.logical_offset_bytes())
99 })
100 .then_with(|| left.len().cmp(&right.len()));
101 if group_order != Ordering::Equal {
102 return group_order;
103 }
104 left.requests()
105 .iter()
106 .zip(right.requests())
107 .find_map(|(left, right)| {
108 let layout_order = left
109 .output_layout()
110 .element_type()
111 .cmp(&right.output_layout().element_type())
112 .then_with(|| {
113 left.output_layout()
114 .element_count()
115 .cmp(&right.output_layout().element_count())
116 });
117 (layout_order != Ordering::Equal).then_some(layout_order)
118 })
119 .unwrap_or(Ordering::Equal)
120}
121
122pub type CompletionReadbackCollectionReceipt = CompletionReadbackBatchReceipt;
125pub type CompletionReadbackCollectionObservation = CompletionReadbackBatchObservation;
126
127#[cfg(test)]
128mod tests {
129 use super::{CompletionReadbackBatchRequest, CompletionReadbackCollectionRequest};
130 use crate::vnext::{
131 BufferUsage, CompletionReadbackRequest, ElementType, HostTransferLayout, NodeId, ResourceId,
132 };
133
134 fn request(
135 participant_index: u32,
136 logical_offset_bytes: u64,
137 element_count: u64,
138 ) -> CompletionReadbackRequest {
139 CompletionReadbackRequest::new_typed(
140 NodeId::new("node/readback").unwrap(),
141 participant_index,
142 ResourceId::new("resource/readback").unwrap(),
143 BufferUsage::State,
144 logical_offset_bytes,
145 HostTransferLayout::new(ElementType::U8, element_count).unwrap(),
146 )
147 .unwrap()
148 }
149
150 fn batch(logical_offset_bytes: u64, element_count: u64) -> CompletionReadbackBatchRequest {
151 CompletionReadbackBatchRequest::new(vec![request(0, logical_offset_bytes, element_count)])
152 .unwrap()
153 }
154
155 #[test]
156 fn participant_layouts_may_follow_distinct_active_token_extents() {
157 let batch =
158 CompletionReadbackBatchRequest::new(vec![request(0, 0, 4), request(1, 0, 7)]).unwrap();
159 assert_eq!(batch.requests()[0].output_layout().element_count(), 4);
160 assert_eq!(batch.requests()[1].output_layout().element_count(), 7);
161 let wrong_element = CompletionReadbackRequest::new_typed(
162 NodeId::new("node/readback").unwrap(),
163 1,
164 ResourceId::new("resource/readback").unwrap(),
165 BufferUsage::State,
166 0,
167 HostTransferLayout::new(ElementType::F16, 2).unwrap(),
168 )
169 .unwrap();
170 assert!(
171 CompletionReadbackBatchRequest::new(vec![request(0, 0, 4), wrong_element]).is_err()
172 );
173 }
174
175 #[test]
176 fn collection_keys_the_complete_typed_physical_range() {
177 let first = batch(0, 4);
178 assert!(
179 CompletionReadbackCollectionRequest::new(vec![first.clone(), first.clone(),]).is_err()
180 );
181 let collection =
182 CompletionReadbackCollectionRequest::new(vec![first, batch(4, 4), batch(0, 8)])
183 .unwrap();
184 assert_eq!(collection.len(), 3);
185 }
186
187 #[test]
188 fn collection_is_not_limited_to_the_legacy_sixty_four_groups() {
189 let batches = (0..65).map(|offset| batch(offset, 1)).collect::<Vec<_>>();
190 assert_eq!(
191 CompletionReadbackCollectionRequest::new(batches)
192 .unwrap()
193 .len(),
194 65
195 );
196 }
197}