vortex_compute/lane_kernels/
map_in_place.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! In-place lane kernels: read from an [`IndexedSink`] and write back through
5//! the same sink (no separate output buffer).
6//!
7//! [`IndexedSink`]: crate::lane_kernels::sink::IndexedSink
8
9use vortex_buffer::BitBuffer;
10
11use crate::lane_kernels::CHUNK_LEN;
12use crate::lane_kernels::sink::IndexedSink;
13
14/// Extension trait providing in-place lane-kernel methods on any [`IndexedSink`].
15///
16/// All methods have default implementations and are inherited via the blanket
17/// `impl<S: IndexedSink> IndexedSinkExt for S` below. Bring the trait into scope
18/// (`use vortex_compute::lane_kernels::IndexedSinkExt;`) to call them with
19/// method syntax.
20///
21/// [`IndexedSink`]: crate::lane_kernels::sink::IndexedSink
22pub trait IndexedSinkExt: IndexedSink + Sized {
23    /// In-place counterpart of [`IndexedSourceExt::map_into`]. Each lane
24    /// is replaced with `f(self[i])`.
25    ///
26    /// The closure reads `Self::Item` and returns `Self::Write`. For the common
27    /// case `Self = &mut [T]` both are `T`; for `ReinterpretSink` the read and
28    /// write types can differ (e.g. read `f32`, write `u32`) over the same
29    /// backing memory when sizes and alignments match.
30    ///
31    /// As with [`IndexedSourceExt::map_into`], use this only when the
32    /// input is known non-nullable.
33    ///
34    /// [`IndexedSourceExt::map_into`]: crate::lane_kernels::map_into::IndexedSourceExt::map_into
35    #[inline]
36    fn map_into_in_place<F>(self, mut f: F)
37    where
38        F: FnMut(Self::Item) -> Self::Write,
39    {
40        #[inline(always)]
41        fn chunk<S, F>(values: &mut S, f: &mut F, base: usize, count: usize)
42        where
43            S: IndexedSink,
44            F: FnMut(S::Item) -> S::Write,
45        {
46            for bit_idx in 0..count {
47                let idx = base + bit_idx;
48                // SAFETY: caller guarantees base + count <= len.
49                let val = unsafe { values.get_unchecked(idx) };
50                let result = f(val);
51                // SAFETY: caller guarantees base + count <= len.
52                unsafe { values.set_unchecked(idx, result) };
53            }
54        }
55
56        let mut values = self;
57        let len = values.len();
58        let chunks_count = len / CHUNK_LEN;
59        let remainder = len % CHUNK_LEN;
60
61        for chunk_idx in 0..chunks_count {
62            chunk(&mut values, &mut f, chunk_idx * CHUNK_LEN, CHUNK_LEN);
63        }
64        if remainder != 0 {
65            chunk(&mut values, &mut f, chunks_count * CHUNK_LEN, remainder);
66        }
67    }
68
69    /// In-place counterpart of [`IndexedSourceExt::try_map_into`]. Each
70    /// lane is replaced with `f(self[i])`, or `Self::Write::default()` when `f`
71    /// returns `None`. On failure returns `Err(first_failing_lane)`; the buffer
72    /// state on `Err` is unspecified.
73    ///
74    /// ## Error attribution
75    ///
76    /// Per-lane `is_none()` flags are bit-packed into a `u64` at the lane's
77    /// position — `fail_bits |= (opt.is_none() as u64) << bit_idx`. After the
78    /// 64-lane loop, `trailing_zeros()` of `fail_bits` recovers the first
79    /// failing lane index. `OR + shift` per lane is friendlier to the
80    /// autovectorizer than `min`/`csel` — see [`try_map_masked_in_place`] for
81    /// the same scheme over a masked variant.
82    ///
83    /// [`try_map_masked_in_place`]: IndexedSinkExt::try_map_masked_in_place
84    /// [`IndexedSourceExt::try_map_into`]: crate::lane_kernels::map_into::IndexedSourceExt::try_map_into
85    #[inline]
86    fn try_map_in_place<F>(self, mut f: F) -> Result<(), usize>
87    where
88        Self::Write: Default,
89        F: FnMut(Self::Item) -> Option<Self::Write>,
90    {
91        #[inline(always)]
92        fn chunk<S, F>(values: &mut S, base: usize, count: usize, f: &mut F) -> Option<usize>
93        where
94            S: IndexedSink,
95            S::Write: Default,
96            F: FnMut(S::Item) -> Option<S::Write>,
97        {
98            let mut fail_bits: u64 = 0;
99            for bit_idx in 0..count {
100                let idx = base + bit_idx;
101                // SAFETY: caller guarantees base + count <= len.
102                let val = unsafe { values.get_unchecked(idx) };
103                let opt = f(val);
104                fail_bits |= (opt.is_none() as u64) << bit_idx;
105                let result = opt.unwrap_or_default();
106                // SAFETY: caller guarantees base + count <= len.
107                unsafe { values.set_unchecked(idx, result) };
108            }
109            (fail_bits != 0).then_some(base + fail_bits.trailing_zeros() as usize)
110        }
111
112        let mut values = self;
113        let len = values.len();
114        let chunks_count = len / 64;
115        let remainder = len % 64;
116
117        for chunk_idx in 0..chunks_count {
118            if let Some(failing) = chunk(&mut values, chunk_idx * 64, 64, &mut f) {
119                return Err(failing);
120            }
121        }
122        if remainder != 0
123            && let Some(failing) = chunk(&mut values, chunks_count * 64, remainder, &mut f)
124        {
125            return Err(failing);
126        }
127        Ok(())
128    }
129
130    /// In-place counterpart of [`IndexedSourceExt::try_map_masked_into`]. Each
131    /// lane of `self` is replaced with `f(self[i])`, or `Self::Write::default()`
132    /// if `f` returned `None`. On failure returns `Err(first_failing_lane)`;
133    /// lanes before that point have been written, and lanes within the failing
134    /// chunk hold their unwrapped-or-default result. The buffer state on `Err`
135    /// is intentionally unspecified.
136    ///
137    /// **Null-lane failures are filtered automatically** — same semantics as
138    /// [`IndexedSourceExt::try_map_masked_into`]. The closure has no `valid`
139    /// parameter; the kernel AND-combines `is_none()` with the chunk's validity
140    /// bitmap before folding it into the attribution accumulator.
141    ///
142    /// ## Why in-place is slower at cache-resident sizes
143    ///
144    /// At sizes that fit in L1/L2 the in-place kernel is ~1.5× slower than the
145    /// out-of-place kernel despite having half the memory traffic, because
146    /// input and output share memory and the compiler must be conservative
147    /// reordering loads/stores across iterations. At sizes that exceed L2 the
148    /// in-place kernel wins back the gap by avoiding the second buffer's DRAM
149    /// read+write traffic.
150    ///
151    /// [`IndexedSourceExt::try_map_masked_into`]: crate::lane_kernels::map_into::IndexedSourceExt::try_map_masked_into
152    ///
153    /// # Panics
154    ///
155    /// Panics if `self.len() != mask.len()`.
156    #[inline]
157    fn try_map_masked_in_place<F>(self, mask: &BitBuffer, mut f: F) -> Result<(), usize>
158    where
159        Self::Write: Default,
160        F: FnMut(Self::Item) -> Option<Self::Write>,
161    {
162        /// Bit-pack `is_none()` flags per lane, then AND with `src_chunk` post-loop to
163        /// drop null-lane failures. The per-lane attribution work is `OR + shift`
164        /// (no `min`/`csel`), giving LLVM more freedom to vectorize the value pipeline.
165        #[inline(always)]
166        fn chunk<S, F>(
167            values: &mut S,
168            src_chunk: u64,
169            base: usize,
170            count: usize,
171            f: &mut F,
172        ) -> Option<usize>
173        where
174            S: IndexedSink,
175            S::Write: Default,
176            F: FnMut(S::Item) -> Option<S::Write>,
177        {
178            let mut fail_bits: u64 = 0;
179            for bit_idx in 0..count {
180                let idx = base + bit_idx;
181                // SAFETY: caller guarantees `base + count <= values.len()`.
182                let val = unsafe { values.get_unchecked(idx) };
183                let opt = f(val);
184                fail_bits |= (opt.is_none() as u64) << bit_idx;
185                let result = opt.unwrap_or_default();
186                unsafe { values.set_unchecked(idx, result) };
187            }
188            let valid_failures = fail_bits & src_chunk;
189            (valid_failures != 0).then_some(base + valid_failures.trailing_zeros() as usize)
190        }
191
192        let mut values = self;
193        let len = values.len();
194        assert_eq!(len, mask.len(), "values and mask must have the same length");
195
196        let chunks = mask.chunks();
197        let chunks_count = len / 64;
198        let remainder = len % 64;
199
200        for (chunk_idx, src_chunk) in chunks.iter().enumerate() {
201            if let Some(failing) = chunk(&mut values, src_chunk, chunk_idx * 64, 64, &mut f) {
202                return Err(failing);
203            }
204        }
205        if remainder != 0
206            && let Some(failing) = chunk(
207                &mut values,
208                chunks.remainder_bits(),
209                chunks_count * 64,
210                remainder,
211                &mut f,
212            )
213        {
214            return Err(failing);
215        }
216        Ok(())
217    }
218}
219
220impl<S: IndexedSink> IndexedSinkExt for S {}
221
222#[cfg(test)]
223#[allow(clippy::cast_possible_truncation)]
224mod tests {
225    use vortex_buffer::BitBuffer;
226    use vortex_buffer::BitBufferMut;
227
228    use super::*;
229    use crate::lane_kernels::sink::ReinterpretSink;
230
231    #[test]
232    fn try_map_masked_in_place_all_ok() {
233        let mut values: Vec<u32> = (0..200).collect();
234        let mask = BitBuffer::new_set(200);
235        let res = values
236            .as_mut_slice()
237            .try_map_masked_in_place(&mask, |v| v.checked_mul(2));
238        assert!(res.is_ok());
239        let expected: Vec<u32> = (0..200u32).map(|v| v * 2).collect();
240        assert_eq!(values, expected);
241    }
242
243    #[test]
244    fn try_map_masked_in_place_first_failing_chunk_wins() {
245        let mut values: Vec<u32> = (0..200).collect();
246        values[83] = u32::MAX;
247        values[150] = u32::MAX;
248        let mask = BitBuffer::new_set(200);
249        let res = values
250            .as_mut_slice()
251            .try_map_masked_in_place(&mask, |v| v.checked_mul(2));
252        assert_eq!(res, Err(83));
253    }
254
255    #[test]
256    fn try_map_masked_in_place_within_chunk_reports_lowest() {
257        let mut values: Vec<u32> = (0..200).collect();
258        values[80] = u32::MAX;
259        values[100] = u32::MAX;
260        let mask = BitBuffer::new_set(200);
261        let res = values
262            .as_mut_slice()
263            .try_map_masked_in_place(&mask, |v| v.checked_mul(2));
264        assert_eq!(res, Err(80));
265    }
266
267    #[test]
268    fn try_map_masked_in_place_single_failure_lane_exact() {
269        let mut values: Vec<u32> = (0..200).collect();
270        values[42] = u32::MAX;
271        let mask = BitBuffer::new_set(200);
272        let res = values
273            .as_mut_slice()
274            .try_map_masked_in_place(&mask, |v| v.checked_mul(2));
275        assert_eq!(res, Err(42));
276    }
277
278    #[test]
279    fn try_map_masked_in_place_null_bypass() {
280        let mut values: Vec<u32> = (0..200).collect();
281        values[5] = u32::MAX;
282        let mask = {
283            let mut m = BitBufferMut::with_capacity(200);
284            for i in 0..200 {
285                m.append(i != 5);
286            }
287            m.freeze()
288        };
289        let res = values
290            .as_mut_slice()
291            .try_map_masked_in_place(&mask, |v| v.checked_mul(2));
292        assert!(res.is_ok(), "null-lane overflow should be filtered");
293        assert_eq!(values[5], 0);
294        assert_eq!(values[6], 12);
295    }
296
297    #[test]
298    fn try_map_masked_in_place_remainder_overflow() {
299        let mut values: Vec<u32> = (0..130).collect();
300        values[129] = u32::MAX;
301        let mask = BitBuffer::new_set(130);
302        let res = values
303            .as_mut_slice()
304            .try_map_masked_in_place(&mask, |v| v.checked_mul(2));
305        assert_eq!(res, Err(129));
306    }
307
308    #[test]
309    fn try_map_masked_in_place_sliced_mask() {
310        let big = BitBuffer::new_set(256);
311        let mask = big.slice(13..143);
312        assert_eq!(mask.len(), 130);
313
314        let mut values: Vec<u32> = (0..130).collect();
315        values[77] = u32::MAX;
316        let res = values
317            .as_mut_slice()
318            .try_map_masked_in_place(&mask, |v| v.checked_mul(2));
319        assert_eq!(res, Err(77));
320    }
321
322    #[test]
323    fn reinterpret_sink_same_width_f32_u32() {
324        let mut buf: Vec<f32> = (0..130).map(|i| i as f32).collect();
325        let mask = BitBuffer::new_set(130);
326        ReinterpretSink::<f32, u32>::new(buf.as_mut_slice())
327            .try_map_masked_in_place(&mask, |f| Some(f.to_bits().wrapping_add(1)))
328            .unwrap();
329        // SAFETY: same size + alignment for f32 and u32; every slot now holds a u32 written by
330        // the closure.
331        let as_u32: &[u32] =
332            unsafe { std::slice::from_raw_parts(buf.as_ptr() as *const u32, buf.len()) };
333        for (i, &got) in as_u32.iter().enumerate() {
334            assert_eq!(got, (i as f32).to_bits().wrapping_add(1), "lane {i}");
335        }
336    }
337
338    #[test]
339    fn reinterpret_sink_failure_reports_lane() {
340        let mut buf: Vec<f32> = (0..200).map(|i| i as f32).collect();
341        let mask = BitBuffer::new_set(200);
342        let res = ReinterpretSink::<f32, u32>::new(buf.as_mut_slice()).try_map_masked_in_place(
343            &mask,
344            |f| {
345                if f as u32 == 137 {
346                    None
347                } else {
348                    Some(f as u32)
349                }
350            },
351        );
352        assert_eq!(res, Err(137));
353    }
354
355    #[test]
356    fn try_map_masked_in_place_partial_chunk_success() {
357        let mut values: Vec<u32> = (0..130).collect();
358        let mask = BitBuffer::new_set(130);
359        let res = values
360            .as_mut_slice()
361            .try_map_masked_in_place(&mask, |v| Some(v + 1));
362        assert!(res.is_ok());
363        assert_eq!(values[0], 1);
364        assert_eq!(values[63], 64);
365        assert_eq!(values[64], 65);
366        assert_eq!(values[129], 130);
367    }
368}
vortex_compute/lane_kernels/map_in_place.rs

vortex_compute/lane_kernels/
map_in_place.rs
