vortex_compute/lane_kernels/
sink.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Writable lane sink — the [`IndexedSink`] trait and the [`ReinterpretSink`] adapter
5//! for in-place type-punning kernels.
6
7use std::marker::PhantomData;
8use std::mem::align_of;
9use std::mem::size_of;
10
11use crate::lane_kernels::source::IndexedSource;
12
13/// An [`IndexedSource`] that also supports unchecked indexed writes — the binding
14/// for in-place kernels.
15///
16/// `Write` is the type written by `set_unchecked` and may differ from
17/// `IndexedSource::Item` (the read type). For the canonical `&mut [T]` impl
18/// both are `T`. The decoupling is what makes [`ReinterpretSink`] possible —
19/// a wrapper that reads `F` and writes `T` over the same backing memory when
20/// the two have identical size and alignment.
21///
22/// Implemented for `&mut [T]`; not implemented for `LaneZip` (you can't write a
23/// `(A, B)` pair back to two separate sources via a single index).
24pub trait IndexedSink: IndexedSource {
25    /// The per-lane write type. Equal to `<Self as IndexedSource>::Item` for
26    /// `&mut [T]`; different for [`ReinterpretSink`].
27    type Write: Copy;
28
29    /// Write `value` into lane `i` without bounds checking.
30    ///
31    /// # Safety
32    ///
33    /// `i` must be strictly less than `self.len()`.
34    unsafe fn set_unchecked(&mut self, i: usize, value: Self::Write);
35}
36
37impl<T: Copy> IndexedSink for &mut [T] {
38    type Write = T;
39    #[inline]
40    unsafe fn set_unchecked(&mut self, i: usize, value: T) {
41        // SAFETY: caller guarantees i < self.len().
42        unsafe { *<[T]>::get_unchecked_mut(self, i) = value };
43    }
44}
45
46/// A sink that reads `F`-values and writes `T`-values over the same backing
47/// slice of `F`, reinterpreting each `T` as `F`-bits on write.
48///
49/// Requires `size_of::<F>() == size_of::<T>()` and `align_of::<F>() == align_of::<T>()`.
50/// Both hold for any pair of `NativePType` primitives with equal byte width
51/// (e.g. `u32` ↔ `f32`, `u64` ↔ `i64`, `f64` ↔ `u64`).
52///
53/// Use this when an in-place kernel needs to convert lanes between two
54/// types of identical width without allocating a second buffer. After the
55/// kernel completes every slot holds a valid `T`-bit pattern; the caller
56/// can recover a typed view via `BufferMut::transmute::<T>()`.
57pub struct ReinterpretSink<'a, F, T> {
58    slice: &'a mut [F],
59    _phantom: PhantomData<T>,
60}
61
62impl<'a, F, T> ReinterpretSink<'a, F, T> {
63    /// Construct a `ReinterpretSink` from `&mut [F]`.
64    ///
65    /// # Panics
66    ///
67    /// Panics if `size_of::<F>() != size_of::<T>()` or
68    /// `align_of::<F>() != align_of::<T>()`.
69    pub fn new(slice: &'a mut [F]) -> Self {
70        assert_eq!(
71            size_of::<F>(),
72            size_of::<T>(),
73            "ReinterpretSink requires F and T to have the same size",
74        );
75        assert_eq!(
76            align_of::<F>(),
77            align_of::<T>(),
78            "ReinterpretSink requires F and T to have the same alignment",
79        );
80        Self {
81            slice,
82            _phantom: PhantomData,
83        }
84    }
85}
86
87impl<F: Copy, T: Copy> IndexedSource for ReinterpretSink<'_, F, T> {
88    type Item = F;
89    #[inline]
90    fn len(&self) -> usize {
91        self.slice.len()
92    }
93    #[inline]
94    unsafe fn get_unchecked(&self, i: usize) -> F {
95        // SAFETY: caller guarantees i < self.slice.len(). Pointer arithmetic
96        // avoids method-resolution ambiguity between `<[F]>::get_unchecked` and
97        // `IndexedSource::get_unchecked`.
98        unsafe { *self.slice.as_ptr().add(i) }
99    }
100}
101
102impl<F: Copy, T: Copy> IndexedSink for ReinterpretSink<'_, F, T> {
103    type Write = T;
104    #[inline]
105    unsafe fn set_unchecked(&mut self, i: usize, value: T) {
106        // SAFETY: caller guarantees i < self.slice.len(); `new` enforces
107        // size_of::<F>() == size_of::<T>() and align_of::<F>() == align_of::<T>(),
108        // so the F-slot can hold a `T` without overflow or misalignment.
109        unsafe {
110            let ptr = self.slice.as_mut_ptr().add(i) as *mut T;
111            ptr.write(value);
112        }
113    }
114}
vortex_compute/lane_kernels/sink.rs

vortex_compute/lane_kernels/
sink.rs
