[EXPERIMENT] Keep ARM integer widening loops compact

vortex-array/src/arrays/primitive/compute/cast.rs

@@ -569,8 +569,19 @@ where
             }
             None => {
                 let mut buffer = BufferMut::<T>::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::<Buffer<_>>(), validity.clone())
+                    .into_array(),
+                PrimitiveArray::new(values.map(u32::from).collect::<Buffer<_>>(), validity)
+                    .into_array(),
+            )
+        } else {
+            let values = [i32::MIN, -1, 0, 1, i32::MAX].into_iter().cycle().take(len);
+            (
+                PrimitiveArray::new(values.clone().collect::<Buffer<_>>(), validity.clone())
+                    .into_array(),
+                PrimitiveArray::new(values.map(i64::from).collect::<Buffer<_>>(), 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())]

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::<i64>(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::<u32>(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
         });
 }

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<S, R, F>(values: &S, out: &mut [MaybeUninit<R>], 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::<false, _, _, _>(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<R, F>(self, out: &mut [MaybeUninit<R>], f: F)
+    where
+        F: Fn(Self::Item) -> R,
+    {
+        map_into_impl::<true, _, _, _>(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<const WIDENING: bool, S, R, F>(source: S, out: &mut [MaybeUninit<R>], f: F)
+where
+    S: IndexedSource,
+    F: Fn(S::Item) -> R,
+{
+    #[allow(clippy::inline_always)]
+    #[inline(always)]
+    fn chunk<S, R, F>(values: &S, out: &mut [MaybeUninit<R>], 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::<Vec<_>>();
+            let mut out = vec![MaybeUninit::<i64>::uninit(); len];
+            signed.as_slice().map_into_widening(&mut out, i64::from);
+            assert_eq!(
+                write_t(out),
+                signed.into_iter().map(i64::from).collect::<Vec<_>>()
+            );
+
+            let unsigned = [0_u16, 1, u16::MAX]
+                .into_iter()
+                .cycle()
+                .take(len)
+                .collect::<Vec<_>>();
+            let mut out = vec![MaybeUninit::<u32>::uninit(); len];
+            unsigned.as_slice().map_into_widening(&mut out, u32::from);
+            assert_eq!(
+                write_t(out),
+                unsigned.into_iter().map(u32::from).collect::<Vec<_>>()
+            );
+        }
+    }
+
+    #[test]
+    fn map_into_widening_supports_owned_results() {
+        let values = [0_u8; 129];
+        let mut out = std::iter::repeat_with(MaybeUninit::<String>::uninit)
+            .take(values.len())
+            .collect::<Vec<_>>();
+        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::<i64>::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::<i64>::uninit(); 2], i64::from);
+    }
+
     #[test]
     fn try_map_masked_into_all_ok() {
         let values: Vec<u64> = (0..200).collect();