From ee129e8f54dbc61d22910becf06ff668b546a0ed Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 30 Oct 2026 12:52:49 -0400 Subject: [PATCH] Keep ARM integer widening loops compact Signed-off-by: Connor Tsui --- .../src/arrays/primitive/compute/cast.rs | 49 +++++- vortex-compute/benches/lane_kernels.rs | 10 +- vortex-compute/src/lane_kernels/map_into.rs | 155 +++++++++++++++--- 3 files changed, 181 insertions(+), 33 deletions(-) diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index af2ad86cf2f..f5f51a00ef2 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -569,8 +569,19 @@ where } None => { let mut buffer = BufferMut::::with_capacity(len); - values.map_into(&mut buffer.spare_capacity_mut()[..len], |v| v.as_()); - // SAFETY: map_into initializes every lane. + let out = &mut buffer.spare_capacity_mut()[..len]; + if const { + cfg!(target_arch = "aarch64") + && matches!( + (F::PTYPE, T::PTYPE), + (PType::I32, PType::I64) | (PType::U16, PType::U32) + ) + } { + values.map_into_widening(out, |v| v.as_()); + } else { + values.map_into(out, |v| v.as_()); + } + // SAFETY: both map methods initialize every lane. unsafe { buffer.set_len(len) }; Ok(PrimitiveArray::new(buffer.freeze(), new_validity).into_array()) } @@ -693,6 +704,7 @@ mod test { use rstest::rstest; use vortex_buffer::BitBuffer; + use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -1296,6 +1308,39 @@ mod test { Ok(()) } + #[rstest] + fn widening_cast_preserves_values_and_validity( + #[values(false, true)] unsigned: bool, + #[values(0, 1, 63, 64, 65, 127, 128, 129, 8192, 16384, 16385)] len: usize, + #[values(0, 1, 2)] validity_kind: u8, + ) -> VortexResult<()> { + let validity = match validity_kind { + 0 => Validity::NonNullable, + 1 => Validity::AllInvalid, + _ => Validity::from_iter((0..len).map(|index| index % 3 != 0)), + }; + let (source, expected) = if unsigned { + let values = [0_u16, 1, u16::MAX].into_iter().cycle().take(len); + ( + PrimitiveArray::new(values.clone().collect::>(), validity.clone()) + .into_array(), + PrimitiveArray::new(values.map(u32::from).collect::>(), validity) + .into_array(), + ) + } else { + let values = [i32::MIN, -1, 0, 1, i32::MAX].into_iter().cycle().take(len); + ( + PrimitiveArray::new(values.clone().collect::>(), validity.clone()) + .into_array(), + PrimitiveArray::new(values.map(i64::from).collect::>(), validity) + .into_array(), + ) + }; + let actual = source.cast(expected.dtype().clone())?; + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); + Ok(()) + } + #[rstest] #[case(buffer![0u8, 1, 2, 3, 255].into_array())] #[case(buffer![0u16, 100, 1000, 65535].into_array())] diff --git a/vortex-compute/benches/lane_kernels.rs b/vortex-compute/benches/lane_kernels.rs index 93c374ecfef..398723d6f1f 100644 --- a/vortex-compute/benches/lane_kernels.rs +++ b/vortex-compute/benches/lane_kernels.rs @@ -222,8 +222,6 @@ fn map_with_mask_widen_u16_u32(bencher: Bencher, n: usize) { }); } -// Both widening variants use the existing map to establish baseline results for the stacked -// optimization before it selects a different loop. Keep their benchmark names stable. #[vortex_bench_support::cpu_features] #[divan::bench(args = WIDEN_SIZES)] fn map_into_i32_i64(bencher: Bencher, n: usize) { @@ -245,7 +243,9 @@ fn map_into_widening_i32_i64(bencher: Bencher, n: usize) { bencher .with_inputs(|| uninit_out::(n)) .bench_local_values(|mut out| { - f.values_i32.as_slice().map_into(&mut out, Into::into); + f.values_i32 + .as_slice() + .map_into_widening(&mut out, Into::into); out }); } @@ -271,7 +271,9 @@ fn map_into_widening_u16_u32(bencher: Bencher, n: usize) { bencher .with_inputs(|| uninit_out::(n)) .bench_local_values(|mut out| { - f.values_u16.as_slice().map_into(&mut out, Into::into); + f.values_u16 + .as_slice() + .map_into_widening(&mut out, Into::into); out }); } diff --git a/vortex-compute/src/lane_kernels/map_into.rs b/vortex-compute/src/lane_kernels/map_into.rs index 04955a7aa20..8f27d9b99ae 100644 --- a/vortex-compute/src/lane_kernels/map_into.rs +++ b/vortex-compute/src/lane_kernels/map_into.rs @@ -211,34 +211,25 @@ pub trait IndexedSourceExt: IndexedSource + Sized { where F: Fn(Self::Item) -> R, { - #[allow(clippy::inline_always)] - #[inline(always)] - fn chunk(values: &S, out: &mut [MaybeUninit], f: &F, base: usize, count: usize) - where - S: IndexedSource, - F: Fn(S::Item) -> R, - { - for bit_idx in 0..count { - let idx = base + bit_idx; - // SAFETY: caller guarantees base + count <= len. - let val = unsafe { values.get_unchecked(idx) }; - unsafe { out.get_unchecked_mut(idx).write(f(val)) }; - } - } - - let values = self; - let len = values.len(); - assert_eq!(out.len(), len, "out must have the same length as values"); - - let chunks_count = len / CHUNK_LEN; - let remainder = len % CHUNK_LEN; + map_into_impl::(self, out, f); + } - for chunk_idx in 0..chunks_count { - chunk(&values, out, &f, chunk_idx * CHUNK_LEN, CHUNK_LEN); - } - if remainder != 0 { - chunk(&values, out, &f, chunks_count * CHUNK_LEN, remainder); - } + /// Map with a loop tuned for numeric widening from contiguous slices. + /// + /// Like [`Self::map_into`], this maps every stored value without consulting validity, + /// initializes every output lane, and performs no allocation. Empty inputs write nothing. + /// Select it only after benchmarking the caller: other source shapes can be slower, even + /// when the output type is wider. + /// + /// # Panics + /// + /// Panics if `out.len() != self.len()`. + #[inline] + fn map_into_widening(self, out: &mut [MaybeUninit], f: F) + where + F: Fn(Self::Item) -> R, + { + map_into_impl::(self, out, f); } /// Apply the predicate `f(value)` lane-by-lane and bit-pack the results into @@ -476,6 +467,56 @@ where cold_scan(values, base, chunk_len, |_bit_idx, val| f(val).is_none()) } +// Keep loop selection at monomorphization time so existing callers retain their loop. +#[inline] +fn map_into_impl(source: S, out: &mut [MaybeUninit], f: F) +where + S: IndexedSource, + F: Fn(S::Item) -> R, +{ + #[allow(clippy::inline_always)] + #[inline(always)] + fn chunk(values: &S, out: &mut [MaybeUninit], f: &F, base: usize, count: usize) + where + S: IndexedSource, + F: Fn(S::Item) -> R, + { + for bit_idx in 0..count { + let idx = base + bit_idx; + // SAFETY: caller guarantees base + count <= len. + let val = unsafe { values.get_unchecked(idx) }; + // SAFETY: the checked output length equals the source length, and the caller + // guarantees that base + count does not exceed it. + unsafe { out.get_unchecked_mut(idx).write(f(val)) }; + } + } + + let values = source; + let len = values.len(); + assert_eq!(out.len(), len, "out must have the same length as values"); + + // NB: On rustc 1.99.1 / LLVM 23.1.1, with 16 CGUs and no LTO, repeated 64-value + // tiles fully unroll i32 widening on M4 (default) and Graviton3 (+neon). A fixed + // prefix also loses the i128-to-i256 gain with CI's AVX-512 features. Keep one + // runtime-bound loop for large maps and retain tiling for the short maps that + // regressed when untiled. Recheck lane_kernels and copy_to_canonical_width, and + // their final IR and assembly, on NEON and AVX-512 before changing this split. + if WIDENING && len > 2 * CHUNK_LEN { + chunk(&values, out, &f, 0, len); + return; + } + + let chunks_count = len / CHUNK_LEN; + let remainder = len % CHUNK_LEN; + + for chunk_idx in 0..chunks_count { + chunk(&values, out, &f, chunk_idx * CHUNK_LEN, CHUNK_LEN); + } + if remainder != 0 { + chunk(&values, out, &f, chunks_count * CHUNK_LEN, remainder); + } +} + #[cfg(test)] #[allow(clippy::cast_possible_truncation)] mod tests { @@ -489,6 +530,66 @@ mod tests { unsafe { std::mem::transmute(out) } } + #[test] + fn map_into_widening_initializes_every_lane() { + for len in [0, 1, 63, 64, 65, 127, 128, 129, 8192, 16384, 16385] { + let signed = [i32::MIN, -1, 0, 1, i32::MAX] + .into_iter() + .cycle() + .take(len) + .collect::>(); + let mut out = vec![MaybeUninit::::uninit(); len]; + signed.as_slice().map_into_widening(&mut out, i64::from); + assert_eq!( + write_t(out), + signed.into_iter().map(i64::from).collect::>() + ); + + let unsigned = [0_u16, 1, u16::MAX] + .into_iter() + .cycle() + .take(len) + .collect::>(); + let mut out = vec![MaybeUninit::::uninit(); len]; + unsigned.as_slice().map_into_widening(&mut out, u32::from); + assert_eq!( + write_t(out), + unsigned.into_iter().map(u32::from).collect::>() + ); + } + } + + #[test] + fn map_into_widening_supports_owned_results() { + let values = [0_u8; 129]; + let mut out = std::iter::repeat_with(MaybeUninit::::uninit) + .take(values.len()) + .collect::>(); + values + .as_slice() + .map_into_widening(&mut out, |value| value.to_string()); + for value in out { + // SAFETY: the map initializes every lane, including the untiled remainder. + assert_eq!(unsafe { value.assume_init() }, "0"); + } + } + + #[test] + #[should_panic(expected = "out must have the same length as values")] + fn map_into_widening_rejects_short_output() { + [1_i32, 2] + .as_slice() + .map_into_widening(&mut [MaybeUninit::::uninit()], i64::from); + } + + #[test] + #[should_panic(expected = "out must have the same length as values")] + fn map_into_widening_rejects_long_output() { + [1_i32] + .as_slice() + .map_into_widening(&mut [MaybeUninit::::uninit(); 2], i64::from); + } + #[test] fn try_map_masked_into_all_ok() { let values: Vec = (0..200).collect();