Skip to main content

ferrum_interfaces/vnext/completion/
readback_collection.rs

1use std::cmp::Ordering;
2
3use super::{
4    invalid_completion, BatchOperationIdentity, CompletionReadbackBatchObservation,
5    CompletionReadbackBatchReceipt, CompletionReadbackBatchRequest, CompletionReadbackRequest,
6    VNextError,
7};
8
9/// Canonical terminal readbacks for multiple node/resource groups. Every
10/// group remains a complete participant batch; this type does not weaken the
11/// single-node invariant of [`CompletionReadbackBatchRequest`].
12#[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
122/// Collection receipts use the same ordered, fingerprinted disposition
123/// evidence as a single readback batch.
124pub 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}