From e973d82e462038b9ee1e3b31f1987aa3c5d53005 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 16:21:18 -0400 Subject: [PATCH] Use Fn for out-of-place lane kernels Signed-off-by: Connor Tsui --- .../src/arrays/decimal/compute/cast.rs | 6 +- .../src/arrays/primitive/compute/cast.rs | 8 +- .../src/scalar_fn/fns/binary/compare/mod.rs | 4 +- .../scalar_fn/fns/binary/numeric/checked.rs | 14 ++-- .../scalar_fn/fns/binary/numeric/primitive.rs | 2 +- vortex-compute/src/lane_kernels/map_into.rs | 74 +++++++++---------- 6 files changed, 53 insertions(+), 55 deletions(-) diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index 1240010dbf3..8770a9bc8f4 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -235,7 +235,7 @@ where fn cast_decimal_buffer( values: &[F], valid_values: &Mask, - mut cast: impl FnMut(F) -> Option, + cast: impl Fn(F) -> Option, ) -> Result, usize> where F: NativeDecimalType, @@ -244,14 +244,14 @@ where let mut buffer = BufferMut::::with_capacity(values.len()); match valid_values { Mask::AllTrue(_) => { - values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &mut cast)?; + values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &cast)?; } Mask::AllFalse(_) => return Ok(BufferMut::::zeroed(values.len()).freeze()), Mask::Values(mask) => { values.try_map_masked_into( mask.bit_buffer(), &mut buffer.spare_capacity_mut()[..values.len()], - &mut cast, + &cast, )?; } } diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 8f5a0f95fcd..3d183d6be15 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -316,7 +316,7 @@ fn cast_integer_values_to_decimal_buffer( values: &[S], decimal_dtype: DecimalDType, valid_values: &Mask, - cast: impl FnMut(S) -> Option, + cast: impl Fn(S) -> Option, ) -> VortexResult> where S: IntegerPType + ToI256, @@ -402,7 +402,7 @@ fn decimal_value_fits_precision( fn cast_primitive_to_decimal_buffer( values: &[S], valid_values: &Mask, - mut cast: impl FnMut(S) -> Option, + cast: impl Fn(S) -> Option, ) -> Result, usize> where S: NativePType, @@ -415,13 +415,13 @@ where let mut buffer = BufferMut::::with_capacity(values.len()); match valid_values { Mask::AllTrue(_) => { - values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &mut cast)?; + values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &cast)?; } Mask::Values(mask) => { values.try_map_masked_into( mask.bit_buffer(), &mut buffer.spare_capacity_mut()[..values.len()], - &mut cast, + &cast, )?; } Mask::AllFalse(_) => unreachable!("all-null values are handled before allocating"), diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index d25a652ee57..8b5a4522b07 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -293,7 +293,7 @@ pub(super) fn bit_buffer_from_words(words: BufferMut, len: usize) -> BitBuf pub(super) fn collect_zip_bits( lhs: &[T], rhs: &[T], - mut f: impl FnMut(T, T) -> bool, + f: impl Fn(T, T) -> bool, ) -> BitBuffer { let len = lhs.len(); let mut words = BufferMut::::zeroed(len.div_ceil(64)); @@ -302,7 +302,7 @@ pub(super) fn collect_zip_bits( } /// Bit-pack the predicate `f(values[i])` over a slice into a [`BitBuffer`]. -pub(super) fn collect_bits(values: &[T], f: impl FnMut(T) -> bool) -> BitBuffer { +pub(super) fn collect_bits(values: &[T], f: impl Fn(T) -> bool) -> BitBuffer { let len = values.len(); let mut words = BufferMut::::zeroed(len.div_ceil(64)); values.map_bits_into(words.as_mut_slice(), f); diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index afb709abe1c..46ff0f36e15 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -49,7 +49,7 @@ pub(super) fn checked_lanes( where S: IndexedSource, T: Copy + Default, - Apply: FnMut(S::Item) -> Option, + Apply: Fn(S::Item) -> Option, { let len = source.len(); debug_assert_eq!(len, valid_rows.len()); @@ -89,13 +89,13 @@ where pub(super) fn checked_apply_lanes( source: S, valid_rows: &Mask, - mut apply: Apply, + apply: Apply, ) -> Result, usize> where S: IndexedSource + Copy, T: Copy + Default, Fail: Failure, - Apply: FnMut(S::Item) -> (T, Fail), + Apply: Fn(S::Item) -> (T, Fail), { let len = source.len(); debug_assert_eq!(len, valid_rows.len()); @@ -109,14 +109,14 @@ where let mut values = BufferMut::::with_capacity(len); let out = &mut values.spare_capacity_mut()[..len]; - if source.map_checked_into(out, &mut apply) != Fail::default() { - let mut checked = |item: S::Item| { + if source.map_checked_into(out, &apply) != Fail::default() { + let checked = |item: S::Item| { let (value, failure) = apply(item); (failure == Fail::default()).then_some(value) }; match valid_bits { - None => source.try_map_into(out, &mut checked)?, - Some(valid_bits) => source.try_map_masked_into(valid_bits, out, &mut checked)?, + None => source.try_map_into(out, checked)?, + Some(valid_bits) => source.try_map_masked_into(valid_bits, out, checked)?, } } diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 8fd53d15216..3a3c3f1b585 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -211,7 +211,7 @@ where fn checked_op_lanes( source: S, valid_rows: &Mask, - mut to_operands: impl FnMut(S::Item) -> (T, T), + to_operands: impl Fn(S::Item) -> (T, T), ) -> Result, usize> where S: IndexedSource + Copy, diff --git a/vortex-compute/src/lane_kernels/map_into.rs b/vortex-compute/src/lane_kernels/map_into.rs index c1e1107b1b9..5add896c855 100644 --- a/vortex-compute/src/lane_kernels/map_into.rs +++ b/vortex-compute/src/lane_kernels/map_into.rs @@ -18,6 +18,9 @@ use crate::lane_kernels::source::IndexedSource; /// `impl IndexedSourceExt for S` below. Bring the trait into /// scope (`use vortex_compute::lane_kernels::IndexedSourceExt;`) to call /// them with method syntax: `values.try_map_masked_into(&mask, &mut out, f)`. +/// +/// Callbacks implement [`Fn`] because each lane must be independent. A callback that mutates +/// captured state introduces a loop-carried dependency that can prevent vectorization. pub trait IndexedSourceExt: IndexedSource + Sized { /// Fallible map with mask-aware error attribution. `f` returns `Option`; /// `None` indicates a per-lane failure (e.g. range overflow on a narrowing cast). @@ -28,7 +31,7 @@ pub trait IndexedSourceExt: IndexedSource + Sized { /// `is_none()` flags are bit-packed into a `u64` at the lane's position, then /// AND-combined with the chunk's validity bitmap — null-lane bits vanish. /// - /// The closure shape is the same as [`try_map_into`] (`FnMut(Item) -> Option`); + /// The closure shape is the same as [`try_map_into`] (`Fn(Item) -> Option`); /// the mask parameter is what makes this kernel mask-aware. Callers that need to /// distinguish null lanes inside the closure (e.g. to short-circuit an expensive /// computation) should construct their own per-lane validity check externally; for @@ -48,17 +51,17 @@ pub trait IndexedSourceExt: IndexedSource + Sized { self, mask: &BitBuffer, out: &mut [MaybeUninit], - mut f: F, + f: F, ) -> Result<(), usize> where R: Copy + Default, - F: FnMut(Self::Item) -> Option, + F: Fn(Self::Item) -> Option, { #[inline(always)] fn chunk( values: &S, out: &mut [MaybeUninit], - f: &mut F, + f: &F, src_chunk: u64, base: usize, count: usize, @@ -66,7 +69,7 @@ pub trait IndexedSourceExt: IndexedSource + Sized { where S: IndexedSource, R: Copy + Default, - F: FnMut(S::Item) -> Option, + F: Fn(S::Item) -> Option, { let mut fail_bits: u64 = 0; for bit_idx in 0..count { @@ -92,7 +95,7 @@ pub trait IndexedSourceExt: IndexedSource + Sized { let remainder = len % 64; for (chunk_idx, src_chunk) in chunks.iter().enumerate() { - if let Some(idx) = chunk(&values, out, &mut f, src_chunk, chunk_idx * 64, 64) { + if let Some(idx) = chunk(&values, out, &f, src_chunk, chunk_idx * 64, 64) { return Err(idx); } } @@ -100,7 +103,7 @@ pub trait IndexedSourceExt: IndexedSource + Sized { && let Some(idx) = chunk( &values, out, - &mut f, + &f, chunks.remainder_bits(), chunks_count * 64, remainder, @@ -119,20 +122,15 @@ pub trait IndexedSourceExt: IndexedSource + Sized { /// /// Panics if `out.len() != self.len()`. #[inline] - fn map_into(self, out: &mut [MaybeUninit], mut f: F) + fn map_into(self, out: &mut [MaybeUninit], f: F) where - F: FnMut(Self::Item) -> R, + F: Fn(Self::Item) -> R, { #[inline(always)] - fn chunk( - values: &S, - out: &mut [MaybeUninit], - f: &mut F, - base: usize, - count: usize, - ) where + fn chunk(values: &S, out: &mut [MaybeUninit], f: &F, base: usize, count: usize) + where S: IndexedSource, - F: FnMut(S::Item) -> R, + F: Fn(S::Item) -> R, { for bit_idx in 0..count { let idx = base + bit_idx; @@ -150,10 +148,10 @@ pub trait IndexedSourceExt: IndexedSource + Sized { let remainder = len % CHUNK_LEN; for chunk_idx in 0..chunks_count { - chunk(&values, out, &mut f, chunk_idx * CHUNK_LEN, CHUNK_LEN); + chunk(&values, out, &f, chunk_idx * CHUNK_LEN, CHUNK_LEN); } if remainder != 0 { - chunk(&values, out, &mut f, chunks_count * CHUNK_LEN, remainder); + chunk(&values, out, &f, chunks_count * CHUNK_LEN, remainder); } } @@ -179,15 +177,15 @@ pub trait IndexedSourceExt: IndexedSource + Sized { /// /// Panics if `words.len() < self.len().div_ceil(64)`. #[inline] - fn map_bits_into(self, words: &mut [u64], mut f: F) + fn map_bits_into(self, words: &mut [u64], f: F) where - F: FnMut(Self::Item) -> bool, + F: Fn(Self::Item) -> bool, { #[inline(always)] - fn chunk(values: &S, f: &mut F, base: usize, count: usize) -> u64 + fn chunk(values: &S, f: &F, base: usize, count: usize) -> u64 where S: IndexedSource, - F: FnMut(S::Item) -> bool, + F: Fn(S::Item) -> bool, { let mut packed: u64 = 0; for bit_idx in 0..count { @@ -211,10 +209,10 @@ pub trait IndexedSourceExt: IndexedSource + Sized { let remainder = len % 64; for word_idx in 0..full { - words[word_idx] = chunk(&values, &mut f, word_idx * 64, 64); + words[word_idx] = chunk(&values, &f, word_idx * 64, 64); } if remainder != 0 { - words[full] = chunk(&values, &mut f, full * 64, remainder); + words[full] = chunk(&values, &f, full * 64, remainder); } } @@ -240,10 +238,10 @@ pub trait IndexedSourceExt: IndexedSource + Sized { /// /// Panics if `out.len() != self.len()`. #[inline] - fn map_checked_into(self, out: &mut [MaybeUninit], mut apply: Apply) -> Fail + fn map_checked_into(self, out: &mut [MaybeUninit], apply: Apply) -> Fail where Fail: Copy + Default + BitOrAssign, - Apply: FnMut(Self::Item) -> (R, Fail), + Apply: Fn(Self::Item) -> (R, Fail), { const { assert!( @@ -293,10 +291,10 @@ pub trait IndexedSourceExt: IndexedSource + Sized { /// /// Panics if `out.len() != self.len()`. #[inline] - fn try_map_into(self, out: &mut [MaybeUninit], mut f: F) -> Result<(), usize> + fn try_map_into(self, out: &mut [MaybeUninit], f: F) -> Result<(), usize> where R: Copy + Default, - F: FnMut(Self::Item) -> Option, + F: Fn(Self::Item) -> Option, { /// Returns `true` if any lane in `[base, base+count)` failed (OR-reduced); /// the cold attribution path is called at the kernel level so it can be @@ -305,14 +303,14 @@ pub trait IndexedSourceExt: IndexedSource + Sized { fn chunk( values: &S, out: &mut [MaybeUninit], - f: &mut F, + f: &F, base: usize, count: usize, ) -> bool where S: IndexedSource, R: Copy + Default, - F: FnMut(S::Item) -> Option, + F: Fn(S::Item) -> Option, { let mut fail_acc: u64 = 0; for bit_idx in 0..count { @@ -336,14 +334,14 @@ pub trait IndexedSourceExt: IndexedSource + Sized { for chunk_idx in 0..chunks_count { let base = chunk_idx * CHUNK_LEN; - if chunk(&values, out, &mut f, base, CHUNK_LEN) { - return Err(attribute_failure_no_mask(&values, base, CHUNK_LEN, &mut f)); + if chunk(&values, out, &f, base, CHUNK_LEN) { + return Err(attribute_failure_no_mask(&values, base, CHUNK_LEN, &f)); } } if remainder != 0 { let base = chunks_count * CHUNK_LEN; - if chunk(&values, out, &mut f, base, remainder) { - return Err(attribute_failure_no_mask(&values, base, remainder, &mut f)); + if chunk(&values, out, &f, base, remainder) { + return Err(attribute_failure_no_mask(&values, base, remainder, &f)); } } Ok(()) @@ -363,7 +361,7 @@ fn cold_scan( values: &S, base: usize, chunk_len: usize, - mut lane_fails: impl FnMut(usize /* bit_idx */, S::Item) -> bool, + lane_fails: impl Fn(usize /* bit_idx */, S::Item) -> bool, ) -> usize where S: IndexedSource, @@ -382,10 +380,10 @@ where /// Cold attribution for the no-mask variant. Replays `f` over the chunk to find /// the first lane that returns `None`. #[inline] -fn attribute_failure_no_mask(values: &S, base: usize, chunk_len: usize, f: &mut F) -> usize +fn attribute_failure_no_mask(values: &S, base: usize, chunk_len: usize, f: &F) -> usize where S: IndexedSource, - F: FnMut(S::Item) -> Option, + F: Fn(S::Item) -> Option, { cold_scan(values, base, chunk_len, |_bit_idx, val| f(val).is_none()) }