gpu/
iter.rs

1/// A GPU iterator over a slice of elements.
2/// The Iterator implementation in core crate is not optimized for GPU.
3/// It is recommended to use the provided GPU iterator abstractions for better performance.
4pub struct GpuIter<'a, T: 'a> {
5    /// The pointer to the next element to return, or the past-the-end location
6    /// if the iterator is empty.
7    ///
8    /// This address will be used for all ZST elements, never changed.
9    slice: &'a [T],
10    /// For non-ZSTs, the non-null pointer to the past-the-end element.
11    ///
12    /// For ZSTs, this is `ptr::without_provenance_mut(len)`.
13    idx: usize,
14}
15
16impl<'a, T: 'a> GpuIter<'a, T> {
17    #[inline(always)]
18    #[gpu_codegen::device]
19    pub fn new(s: &'a [T]) -> Self {
20        Self { slice: s, idx: 0 }
21    }
22}
23
24impl<'a, T: 'a> Iterator for GpuIter<'a, T> {
25    type Item = &'a T;
26
27    #[inline(always)]
28    #[gpu_codegen::device]
29    fn next(&mut self) -> Option<Self::Item> {
30        let ret = if self.idx < self.slice.len() {
31            let item = unsafe { self.slice.get_unchecked(self.idx) };
32            Some(item)
33        } else {
34            None
35        };
36        self.idx += 1;
37        ret
38    }
39}