rust-lang/rust · #162810

vectorized vec::dedup pre-scan

ywxt · merged Oct 12, 20263 files · 390 + / 108 −
library/alloc/src/vec/mod.rs159 + / 4 −
@@ -2717,7 +2717,24 @@ impl<T, A: Allocator> Vec<T, A> {         if first_duplicate_idx == len {             return;         }+        // SAFETY: first_duplicate_idx always in range [1..len)+        unsafe {+            self.dedup_tail(same_bucket, first_duplicate_idx);+        }+    } +    /// Continues deduplication after the first duplicate is found.+    ///+    /// The function treats `self[first_duplicate]` as the first duplicated+    /// element and drops it.+    ///+    /// # Safety+    ///+    /// `1 <= first_duplicate < self.len()`+    unsafe fn dedup_tail<F>(&mut self, mut same_bucket: F, first_duplicate: usize)+    where+        F: FnMut(&mut T, &mut T) -> bool,+    {         /* INVARIANT: vec.len() > read > write > write-1 >= 0 */         struct FillGapOnDrop<'a, T, A: core::alloc::Allocator> {             /* Offset of the element we want to check if it is duplicate */@@ -2764,16 +2781,20 @@ impl<T, A: Allocator> Vec<T, A> {             }         } +        let start = self.as_mut_ptr();+        let len = self.len();+         /* Drop items while going through Vec, it should be more efficient than          * doing slice partition_dedup + truncate */          // Construct gap first and then drop item to avoid memory corruption if `T::drop` panics.         let mut gap =-            FillGapOnDrop { read: first_duplicate_idx + 1, write: first_duplicate_idx, vec: self };-        // SAFETY: we checked that first_duplicate_idx in bounds before.+            FillGapOnDrop { read: first_duplicate + 1, write: first_duplicate, vec: self };++        // SAFETY: By this function's safety contract, `first_duplicate < self.len()`.         // If drop panics, `gap` would remove this item without drop.         unsafe {-            ptr::drop_in_place(start.add(first_duplicate_idx));+            ptr::drop_in_place(start.add(first_duplicate));         }          // SAFETY: Because of the invariant, read_ptr, prev_ptr and write_ptr@@ -3792,8 +3813,142 @@ impl<T: PartialEq, A: Allocator> Vec<T, A> {     #[stable(feature = "rust1", since = "1.0.0")]     #[inline]     pub fn dedup(&mut self) {-        self.dedup_by(|a, b| a == b)+        // Vectorize the process of finding first duplicate if possible+        if let Some(chunk_size) = dedup_simd_chunk_size::<T>() {+            // Unlike dedup_by, we don't guarantee the order+            // and time in which we call `PartialEq::eq`.+            let Some(first_duplicate) = find_first_duplicate(self.as_slice(), chunk_size) else {+                return;+            };++            // SAFETY: `1 <= first_duplicate < self.len()`+            unsafe {+                self.dedup_tail(|a, b| a == b, first_duplicate);+            }+        } else {+            self.dedup_by(|a, b| a == b);+        }+    }+}++const fn dedup_simd_chunk_size<T>() -> Option<usize> {+    // The minimal SIMD width on supported platforms+    #[allow(unused)]+    const SIMD_WIDTH: usize = 16;+    let size = size_of::<T>();++    // We only have SIMD implementations for these element sizes.+    if !matches!(size, 1 | 2 | 4 | 8) {+        return None;+    }++    let simd_width: Option<usize> = cfg_select! {+        all(target_arch = "arm", target_feature = "neon") => {+            // Packed 64-bit integer lane equality is unavailable on arm neon+            if size == 8 { None } else { Some(SIMD_WIDTH) }+        }++        all(target_arch = "aarch64", target_feature = "neon") => {+            Some(SIMD_WIDTH)+        }++        all(+            any(target_arch = "x86", target_arch = "x86_64"),+            target_feature = "sse4.1"+        ) => Some(SIMD_WIDTH),++        all(+            any(target_arch = "x86", target_arch = "x86_64"),+            target_feature = "sse2"+        ) => {+            // Packed 64-bit integer lane equality is only available on sse4.1+            if size == 8 { None } else { Some(SIMD_WIDTH) }+        }++        all(+            any(target_arch = "riscv32", target_arch = "riscv64"),+            target_feature = "v"+        ) => Some(SIMD_WIDTH),++        all(+            any(target_arch = "loongarch32", target_arch = "loongarch64"),+            target_feature = "lsx"+        ) => Some(SIMD_WIDTH),++        _ => None,+    };++    if let Some(simd_width) = simd_width { Some(simd_width * 2 / size) } else { None }+}++fn find_first_duplicate<T: PartialEq>(v: &[T], chunk_size: usize) -> Option<usize> {+    let len = v.len();+    if len <= 1 {+        return None;+    }++    let mut i = 1;++    // The scalar fallback path if the length is not long enough+    if len < chunk_size + 1 {+        while i < len {+            // It keeps the same comparison order as `dedup_tail`+            //+            // SAFETY: i always in range [1..len)+            if unsafe { v.get_unchecked(i) == v.get_unchecked(i - 1) } {+                return Some(i);+            }+            i += 1;+        }+        return None;+    }++    /// # Safety+    /// - `start > 0` and `start + chunk_size <= v.len()`+    /// - `0 < chunk_size <= 32`+    #[inline(always)]+    unsafe fn find_duplicate<T: PartialEq>(+        v: &[T],+        start: usize,+        chunk_size: usize,+    ) -> Option<usize> {+        let mut mask: u32 = 0;+        debug_assert!(+            u32::BITS as usize >= chunk_size,+            "the bit width of mask must be greater than or equal to chunk_size"+        );+        for j in 0..chunk_size {+            // SAFETY: start + j always in range [1..v.len())+            let eq = unsafe { v.get_unchecked(start + j) == v.get_unchecked(start + j - 1) };+            mask |= (eq as u32) << j;+        }+        if mask != 0 {+            return Some(start + mask.trailing_zeros() as usize);+        }+        None     }++    // This loop performs chunk_size comparisons in parallel without early exit,+    // which is what allows for SIMD optimization.+    while i + chunk_size <= len {+        // SAFETY: `i >= 1` and `i + chunk_size <= len`+        if let Some(duplicate) = unsafe { find_duplicate(v, i, chunk_size) } {+            return Some(duplicate);+        }+        i += chunk_size;+    }++    // This may check the same elements multiple times.+    // However we don't guarantee `PartialEq::eq` is only called once per pair.+    if i < len {+        let start = len - chunk_size;+        // SAFETY: `len - chunk_size >= 1` and `start + chunk_size <= len`+        if let Some(duplicate) = unsafe { find_duplicate(v, start, chunk_size) } {+            return Some(duplicate);+        }+    }++    None }  ////////////////////////////////////////////////////////////////////////////////
library/alloctests/benches/vec.rs143 + / 104 −
@@ -688,137 +688,176 @@ fn bench_dedup_slice_truncate(b: &mut Bencher, sz: usize) {     }); } -// Measures performance of Vec::dedup on random data.-fn bench_vec_dedup_random(b: &mut Bencher, sz: usize) {-    let mut template = vec![0u32; sz];-    b.bytes = size_of_val(template.as_slice()) as u64;-    random_sorted_fill(0x43, &mut template);--    let mut vec = template.clone();-    b.iter(|| {-        let vec = black_box(&mut vec);-        vec.dedup();-        black_box(vec.first());-        let vec = black_box(vec);-        vec.clear();-        vec.extend_from_slice(&template);-    });-}--// Measures performance of Vec::dedup when there is no items removed-fn bench_vec_dedup_none(b: &mut Bencher, sz: usize) {-    let mut template = vec![0u32; sz];-    b.bytes = size_of_val(template.as_slice()) as u64;-    template.chunks_exact_mut(2).for_each(|w| {-        w[0] = black_box(0);-        w[1] = black_box(5);-    });--    let mut vec = template.clone();-    b.iter(|| {-        let vec = black_box(&mut vec);-        vec.dedup();-        black_box(vec.first());-        // Unlike other benches of `dedup`-        // this doesn't reinitialize vec-        // because we measure how efficient dedup is-        // when no memory written-    });-}--// Measures performance of Vec::dedup when there is all items removed-fn bench_vec_dedup_all(b: &mut Bencher, sz: usize) {-    let mut template = vec![0u32; sz];-    b.bytes = size_of_val(template.as_slice()) as u64;-    template.iter_mut().for_each(|w| {-        *w = black_box(0);-    });--    let mut vec = template.clone();-    b.iter(|| {-        let vec = black_box(&mut vec);-        vec.dedup();-        black_box(vec.first());-        let vec = black_box(vec);-        vec.clear();-        vec.extend_from_slice(&template);-    });-}- #[bench] fn bench_dedup_slice_truncate_100(b: &mut Bencher) {     bench_dedup_slice_truncate(b, 100); }-#[bench]-fn bench_dedup_random_100(b: &mut Bencher) {-    bench_vec_dedup_random(b, 100);-}--#[bench]-fn bench_dedup_none_100(b: &mut Bencher) {-    bench_vec_dedup_none(b, 100);-}--#[bench]-fn bench_dedup_all_100(b: &mut Bencher) {-    bench_vec_dedup_all(b, 100);-}- #[bench] fn bench_dedup_slice_truncate_1000(b: &mut Bencher) {     bench_dedup_slice_truncate(b, 1000); } #[bench]-fn bench_dedup_random_1000(b: &mut Bencher) {-    bench_vec_dedup_random(b, 1000);+fn bench_dedup_slice_truncate_10000(b: &mut Bencher) {+    bench_dedup_slice_truncate(b, 10000); }- #[bench]-fn bench_dedup_none_1000(b: &mut Bencher) {-    bench_vec_dedup_none(b, 1000);+fn bench_dedup_slice_truncate_100000(b: &mut Bencher) {+    bench_dedup_slice_truncate(b, 100000); } -#[bench]-fn bench_dedup_all_1000(b: &mut Bencher) {-    bench_vec_dedup_all(b, 1000);+#[derive(Clone, Copy)]+enum DedupDist {+    Random,+    None,+    All, } -#[bench]-fn bench_dedup_slice_truncate_10000(b: &mut Bencher) {-    bench_dedup_slice_truncate(b, 10000);+trait DedupBenchValue {+    fn bench_bytes(&self) -> usize; }-#[bench]-fn bench_dedup_random_10000(b: &mut Bencher) {-    bench_vec_dedup_random(b, 10000);++macro_rules! dedup_bench_value_by_size {+    ($($ty:ty),+ $(,)?) => {+        $(+            impl DedupBenchValue for $ty {+                fn bench_bytes(&self) -> usize {+                    size_of_val(self)+                }+            }+        )++    }; } -#[bench]-fn bench_dedup_none_10000(b: &mut Bencher) {-    bench_vec_dedup_none(b, 10000);+dedup_bench_value_by_size!(u8, u16, u32, u64, [u64; 4]);++impl DedupBenchValue for String {+    fn bench_bytes(&self) -> usize {+        self.len()+    } } -#[bench]-fn bench_dedup_all_10000(b: &mut Bencher) {-    bench_vec_dedup_all(b, 10000);+fn bench_vec_dedup<T: DedupBenchValue + Ord + Clone>(+    b: &mut Bencher,+    dist: DedupDist,+    sz: usize,+    make: impl Fn(u32) -> T,+) {+    let template: Vec<T> = match dist {+        DedupDist::Random => {+            let mut seeds = vec![0u32; sz];+            random_sorted_fill(0x43, &mut seeds);+            let mut v: Vec<T> = seeds.into_iter().map(make).collect();+            v.sort();+            v+        }+        DedupDist::None => {+            (0..sz).map(|i| make(black_box(if i % 2 == 0 { 0 } else { 5 }))).collect()+        }+        DedupDist::All => (0..sz).map(|_| make(black_box(0))).collect(),+    };+    b.bytes = template.iter().map(DedupBenchValue::bench_bytes).sum::<usize>() as u64;+    let mut vec = template.clone();+    b.iter(|| {+        let vec = black_box(&mut vec);+        vec.dedup();+        black_box(vec.first());+        // Unlike other benches of `dedup`+        // `DedupDist::None` doesn't reinitialize vec+        // because we measure how efficient dedup is+        // when no memory written+        if !matches!(dist, DedupDist::None) {+            black_box(vec).clone_from(&template);+        }+    }); } -#[bench]-fn bench_dedup_slice_truncate_100000(b: &mut Bencher) {-    bench_dedup_slice_truncate(b, 100000);+fn dedup_gen_string<const LEN: usize>(x: u32) -> String {+    const DIGITS: usize = 10;+    let mut s = String::with_capacity(LEN);+    s.extend(std::iter::repeat_n('p', LEN - DIGITS));+    s.push_str(&format!("{x:0DIGITS$}"));+    s }-#[bench]-fn bench_dedup_random_100000(b: &mut Bencher) {-    bench_vec_dedup_random(b, 100000);++fn dedup_gen_wide(x: u32) -> [u64; 4] {+    [x as u64, (x as u64) << 8, (x as u64) << 16, (x as u64) << 24] } -#[bench]-fn bench_dedup_none_100000(b: &mut Bencher) {-    bench_vec_dedup_none(b, 100000);+macro_rules! dedup_typed {+    ($($name:ident: $gen:expr, $dist:ident, $sz:literal;)*) => {+        $(+            #[bench]+            fn $name(b: &mut Bencher) {+                bench_vec_dedup(b, DedupDist::$dist, $sz, $gen);+            }+        )*+    }; } -#[bench]-fn bench_dedup_all_100000(b: &mut Bencher) {-    bench_vec_dedup_all(b, 100000);+dedup_typed! {+    bench_dedup_random_100: |x| x, Random, 100;+    bench_dedup_none_100: |x| x, None, 100;+    bench_dedup_all_100: |x| x, All, 100;+    bench_dedup_random_1000: |x| x, Random, 1000;+    bench_dedup_none_1000: |x| x, None, 1000;+    bench_dedup_all_1000: |x| x, All, 1000;+    bench_dedup_random_10000: |x| x, Random, 10000;+    bench_dedup_none_10000: |x| x, None, 10000;+    bench_dedup_all_10000: |x| x, All, 10000;+    bench_dedup_random_100000: |x| x, Random, 100000;+    bench_dedup_none_100000: |x| x, None, 100000;+    bench_dedup_all_100000: |x| x, All, 100000;+    bench_dedup_string_short_random_100: dedup_gen_string::<10>, Random, 100;+    bench_dedup_string_short_none_100: dedup_gen_string::<10>, None, 100;+    bench_dedup_string_short_all_100: dedup_gen_string::<10>, All, 100;+    bench_dedup_string_short_random_1000: dedup_gen_string::<10>, Random, 1000;+    bench_dedup_string_short_none_1000: dedup_gen_string::<10>, None, 1000;+    bench_dedup_string_short_all_1000: dedup_gen_string::<10>, All, 1000;+    bench_dedup_string_short_random_10000: dedup_gen_string::<10>, Random, 10000;+    bench_dedup_string_short_none_10000: dedup_gen_string::<10>, None, 10000;+    bench_dedup_string_short_all_10000: dedup_gen_string::<10>, All, 10000;+    bench_dedup_string_medium_random_100: dedup_gen_string::<64>, Random, 100;+    bench_dedup_string_medium_none_100: dedup_gen_string::<64>, None, 100;+    bench_dedup_string_medium_all_100: dedup_gen_string::<64>, All, 100;+    bench_dedup_string_medium_random_1000: dedup_gen_string::<64>, Random, 1000;+    bench_dedup_string_medium_none_1000: dedup_gen_string::<64>, None, 1000;+    bench_dedup_string_medium_all_1000: dedup_gen_string::<64>, All, 1000;+    bench_dedup_string_medium_random_10000: dedup_gen_string::<64>, Random, 10000;+    bench_dedup_string_medium_none_10000: dedup_gen_string::<64>, None, 10000;+    bench_dedup_string_medium_all_10000: dedup_gen_string::<64>, All, 10000;+    bench_dedup_string_long_random_100: dedup_gen_string::<1024>, Random, 100;+    bench_dedup_string_long_none_100: dedup_gen_string::<1024>, None, 100;+    bench_dedup_string_long_all_100: dedup_gen_string::<1024>, All, 100;+    bench_dedup_string_long_random_1000: dedup_gen_string::<1024>, Random, 1000;+    bench_dedup_string_long_none_1000: dedup_gen_string::<1024>, None, 1000;+    bench_dedup_string_long_all_1000: dedup_gen_string::<1024>, All, 1000;+    bench_dedup_string_long_random_10000: dedup_gen_string::<1024>, Random, 10000;+    bench_dedup_string_long_none_10000: dedup_gen_string::<1024>, None, 10000;+    bench_dedup_string_long_all_10000: dedup_gen_string::<1024>, All, 10000;+    bench_dedup_u8_random_1000: |x| x as u8, Random, 1000;+    bench_dedup_u8_none_1000: |x| x as u8, None, 1000;+    bench_dedup_u8_all_1000: |x| x as u8, All, 1000;+    bench_dedup_u8_random_10000: |x| x as u8, Random, 10000;+    bench_dedup_u8_none_10000: |x| x as u8, None, 10000;+    bench_dedup_u8_all_10000: |x| x as u8, All, 10000;+    bench_dedup_u16_random_1000: |x| x as u16, Random, 1000;+    bench_dedup_u16_none_1000: |x| x as u16, None, 1000;+    bench_dedup_u16_all_1000: |x| x as u16, All, 1000;+    bench_dedup_u16_random_10000: |x| x as u16, Random, 10000;+    bench_dedup_u16_none_10000: |x| x as u16, None, 10000;+    bench_dedup_u16_all_10000: |x| x as u16, All, 10000;+    bench_dedup_u64_random_1000: |x| x as u64, Random, 1000;+    bench_dedup_u64_none_1000: |x| x as u64, None, 1000;+    bench_dedup_u64_all_1000: |x| x as u64, All, 1000;+    bench_dedup_u64_random_10000: |x| x as u64, Random, 10000;+    bench_dedup_u64_none_10000: |x| x as u64, None, 10000;+    bench_dedup_u64_all_10000: |x| x as u64, All, 10000;+    bench_dedup_wide_random_1000: dedup_gen_wide, Random, 1000;+    bench_dedup_wide_none_1000: dedup_gen_wide, None, 1000;+    bench_dedup_wide_all_1000: dedup_gen_wide, All, 1000;+    bench_dedup_wide_random_10000: dedup_gen_wide, Random, 10000;+    bench_dedup_wide_none_10000: dedup_gen_wide, None, 10000;+    bench_dedup_wide_all_10000: dedup_gen_wide, All, 10000; }  #[bench]
library/alloctests/tests/vec.rs88 + / 0 −
@@ -2534,6 +2534,94 @@ fn test_vec_dedup_panicking() {     } } +#[test]+fn test_vec_dedup_long() {+    fn check<T: Clone + Debug + PartialEq>(template: &[T]) {+        let mut expected = template.to_vec();+        let (kept, _) = expected.partition_dedup();+        let kept_len = kept.len();+        expected.truncate(kept_len);++        let mut vec = template.to_vec();+        vec.dedup();+        assert_eq!(vec, expected, "input: {template:?}");+    }++    macro_rules! check_type {+        ($ty:ty) => {+            for len in [0, 1, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65] {+                let distinct: Vec<$ty> = (0..len).map(|x| x as $ty).collect();+                check(&distinct);+                check(&vec![7 as $ty; len]);++                for dup_at in 1..len {+                    let mut v = distinct.clone();+                    v[dup_at] = v[dup_at - 1];+                    check(&v);++                    if dup_at + 1 < len {+                        let mut w = v.clone();+                        w[len - 1] = w[len - 2];+                        check(&w);+                    }+                }+            }+        };+    }++    check_type!(u8);+    check_type!(u16);+    check_type!(u32);+    check_type!(u64);+}++#[test]+fn test_vec_dedup_long_boxed() {+    for len in [9, 10, 16, 17, 24, 25, 33] {+        for dup_at in 1..len {+            let mut v: Vec<Box<u32>> = (0..len as u32).map(Box::new).collect();+            v[dup_at] = Box::new(*v[dup_at - 1]);+            let expected: Vec<u32> = (0..len as u32).filter(|&x| x != dup_at as u32).collect();+            v.dedup();+            let got: Vec<u32> = v.iter().map(|b| **b).collect();+            assert_eq!(got, expected, "len={len} dup_at={dup_at}");+        }+    }+}++#[test]+#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]+fn test_vec_dedup_eq_panics_for_simd_sized_type() {+    static EQ_COUNT: AtomicU32 = AtomicU32::new(0);+    static DROP_COUNT: AtomicU32 = AtomicU32::new(0);++    #[derive(Debug)]+    #[repr(transparent)]+    struct Panic(u32);++    impl PartialEq for Panic {+        fn eq(&self, other: &Self) -> bool {+            if EQ_COUNT.fetch_add(1, Ordering::SeqCst) == 2 {+                panic!("comparison panic");+            }+            self.0 == other.0+        }+    }++    impl Drop for Panic {+        fn drop(&mut self) {+            DROP_COUNT.fetch_add(1, Ordering::SeqCst);+        }+    }++    EQ_COUNT.store(0, Ordering::SeqCst);+    DROP_COUNT.store(0, Ordering::SeqCst);+    let mut vec: Vec<_> = [0, 0, 1, 2, 3, 4, 5, 6, 7].into_iter().map(Panic).collect();+    assert!(catch_unwind(AssertUnwindSafe(|| vec.dedup())).is_err());+    drop(vec);+    assert_eq!(DROP_COUNT.load(Ordering::SeqCst), 9);+}+ // Regression test for issue #82533 #[test] #[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]