library/alloc/src/vec/sve_retain.rs - third_party/rust - Git at Google

 //! SVE-accelerated `Vec::retain_mut` implementation.
 //!
 //! Two-phase algorithm per 64-element chunk:
 //! - Phase A (scalar): evaluates predicate exactly once, in order, into a bool mask.
 //! - Phase B (SVE): uses `COMPACT` instruction to compress retained elements.

 use core::{cmp, mem, ptr};

 use super::Vec;
 use crate::alloc::Allocator;

 const CHUNK_SIZE: usize = 64;

 pub(super) const MIN_SVE_SIZE_1: usize = 32;
 pub(super) const MIN_SVE_SIZE_2: usize = 32;
 pub(super) const MIN_SVE_SIZE_4: usize = 64;
 pub(super) const MIN_SVE_SIZE_8: usize = 64;

 /// Guard for the SVE retain path. On panic, this guard:
 /// 1. Scalar-compresses the already-decided prefix of the current chunk.
 /// 2. Copies the untouched tail forward.
 /// 3. Sets the Vec length.
 struct PanicGuard<'a, T, A: Allocator> {
     v: &'a mut Vec<T, A>,
     /// Start index of the current chunk within the Vec.
     read: usize,
     /// Write cursor (accumulated from previous chunks).
     write: usize,
     /// How many elements in the current chunk have been decided by the predicate.
     decided: usize,

     mask: &'a mut [bool; CHUNK_SIZE],
     /// Original length of the Vec.
     original_len: usize,
 }

 impl<T, A: Allocator> Drop for PanicGuard<'_, T, A> {
     #[cold]
     fn drop(&mut self) {
         // Scalar-compress the decided prefix: move kept elements to write position.
         let mut dst = self.write;
         for i in 0..self.decided {
             if self.mask[i] {
                 // SAFETY: read + i < original_len (in-bounds).
                 let src = unsafe { self.v.as_ptr().add(self.read + i) };
                 // SAFETY: Continued from above.
                 let dst_ptr = unsafe { self.v.as_mut_ptr().add(dst) };
                 // SAFETY: src and dst_ptr < original_len
                 unsafe { ptr::copy(src, dst_ptr, 1) };
                 dst += 1;
             }
         }

         // Copy the untouched tail.
         let untouched_start = self.read + self.decided;
         let untouched_len = self.original_len - untouched_start;
         if untouched_len > 0 {
             // SAFETY: untouched_start..original_len are valid; dst + untouched_len <= original_len.
             unsafe {
                 ptr::copy(
                     self.v.as_ptr().add(untouched_start),
                     self.v.as_mut_ptr().add(dst),
                     untouched_len,
                 );
             }
             dst += untouched_len;
         }

         // SAFETY: After filling holes, all items are in contiguous memory.
         unsafe { self.v.set_len(dst) };
     }
 }

 /// # Safety
 ///
 /// - `size_of::<T>() == 1/2/4/8`
 /// - Called on aarch64 + SVE target.
 pub(crate) unsafe fn chunked_retain<T, F, A: Allocator>(v: &mut Vec<T, A>, mut f: F)
 where
     F: FnMut(&mut T) -> bool,
 {
     let original_len = v.len();
     let mut guard = PanicGuard {
         v,
         read: 0,
         write: 0,
         decided: 0,
         mask: &mut [false; CHUNK_SIZE],
         original_len,
     };

     while guard.read < guard.original_len {
         let chunk_len = cmp::min(CHUNK_SIZE, guard.original_len - guard.read);

         guard.decided = 0;

         // Phase A: scalar predicate evaluation (exactly once, in order).
         for i in 0..chunk_len {
             // SAFETY: read + i < original_len
             let cur = unsafe { &mut *guard.v.as_mut_ptr().add(guard.read + i) };
             guard.mask[i] = f(cur);
             guard.decided = i + 1;
             if !guard.mask[i] {
                 // SAFETY: read + i < original_len, and after marking `mask` and `decided`,
                 // the guard can properly handles the case where drop_in_place panics.
                 unsafe { ptr::drop_in_place(cur) };
             }
         }

         // Phase B: SVE compress.
         let kept = match mem::size_of::<T>() {
             // SAFETY: write <= read and the dispatch guarantees size_of::<T>() matches
             // the kernel lane width.
             1 => unsafe {
                 compact8_kernel(
                     guard.v.as_mut_ptr().add(guard.read),
                     guard.v.as_mut_ptr().add(guard.write),
                     guard.mask.as_ptr(),
                     chunk_len,
                 )
             },
             // SAFETY: Same as above.
             2 => unsafe {
                 compact16_kernel(
                     guard.v.as_mut_ptr().add(guard.read),
                     guard.v.as_mut_ptr().add(guard.write),
                     guard.mask.as_ptr(),
                     chunk_len,
                 )
             },
             // SAFETY: Same as above.
             4 => unsafe {
                 compact32_kernel(
                     guard.v.as_mut_ptr().add(guard.read),
                     guard.v.as_mut_ptr().add(guard.write),
                     guard.mask.as_ptr(),
                     chunk_len,
                 )
             },
             // SAFETY: Same as above.
             8 => unsafe {
                 compact64_kernel(
                     guard.v.as_mut_ptr().add(guard.read),
                     guard.v.as_mut_ptr().add(guard.write),
                     guard.mask.as_ptr(),
                     chunk_len,
                 )
             },
             _ => unreachable!(),
         };

         guard.write += kept;
         guard.read += chunk_len;
     }

     // SAFETY: write <= original_len, all retained elements are packed at the front.
     unsafe { guard.v.set_len(guard.write) };
     mem::forget(guard);
 }

 macro_rules! sve_compact_kernel {
     (
         $name:ident,
         size = $size:literal,
         lane = $lane:literal,
         mem_lane = $mem_lane:literal,
         shift = $shift:literal,
         inc = $inc:literal
     ) => {
         /// SVE compress kernel: pack retained elements from `src` to `dst`
         /// according to `mask`. Returns the number of retained elements.
         ///
         /// # Safety
         ///
         /// - `src` is valid for `chunk_len` reads of `T`, `dst` for `chunk_len`
         ///   writes, `mask` for `chunk_len` bool reads.
         /// - `size_of::<T>()` equals this kernel's lane width in bytes.
         #[target_feature(enable = "sve")]
         #[inline]
         unsafe fn $name<T>(
             src: *const T,
             dst: *mut T,
             mask: *const bool,
             chunk_len: usize,
         ) -> usize {
             debug_assert_eq!(mem::size_of::<T>(), $size);

             let idx_in = 0usize;
             let mut idx_out = 0usize;
             // SVE intrinsics require treating the data as integers, which doesn't
             // correctly handle provenance and uninitialized padding.
             //
             // SAFETY: whilelo predicates every load/store to the remaining elements.
             unsafe {
                 core::arch::asm!(
                     concat!("whilelo p0.", $lane, ", xzr, {len}"),
                     "2:",
                     concat!("ld1b {{ z0.", $lane, " }}, p0/z, [{mask}, {idx_in}]"),
                     concat!("cmpne p1.", $lane, ", p0/z, z0.", $lane, ", #0"),
                     concat!("ld1", $mem_lane, " {{ z1.", $lane, " }}, p1/z, [{src}, {idx_in}", $shift, "]"),
                     concat!("compact z1.", $lane, ", p1, z1.", $lane),
                     concat!("cntp {kept}, p0, p1.", $lane),
                     concat!("whilelo p2.", $lane, ", xzr, {kept}"),
                     concat!("st1", $mem_lane, " {{ z1.", $lane, " }}, p2, [{dst}, {idx_out}", $shift, "]"),
                     concat!("add {idx_out}, {idx_out}, {kept}"),
                     concat!("inc", $inc, " {idx_in}"),
                     concat!("whilelo p0.", $lane, ", {idx_in}, {len}"),
                     "b.first 2b",
                     src = in(reg) src,
                     dst = in(reg) dst,
                     mask = in(reg) mask,
                     len = in(reg) chunk_len,
                     idx_in = inout(reg) idx_in => _,
                     idx_out = inout(reg) idx_out,
                     kept = out(reg) _,
                     out("p0") _,
                     out("p1") _,
                     out("p2") _,
                     out("z0") _,
                     out("z1") _,
                     options(nostack),
                 );
             }
             idx_out
         }
     };
 }

 sve_compact_kernel!(compact8_kernel, size = 1, lane = "s", mem_lane = "b", shift = "", inc = "w");

 sve_compact_kernel!(
     compact16_kernel,
     size = 2,
     lane = "s",
     mem_lane = "h",
     shift = ", lsl #1",
     inc = "w"
 );

 sve_compact_kernel!(
     compact32_kernel,
     size = 4,
     lane = "s",
     mem_lane = "w",
     shift = ", lsl #2",
     inc = "w"
 );

 sve_compact_kernel!(
     compact64_kernel,
     size = 8,
     lane = "d",
     mem_lane = "d",
     shift = ", lsl #3",
     inc = "d"
 );
	//! SVE-accelerated `Vec::retain_mut` implementation.
	//!
	//! Two-phase algorithm per 64-element chunk:
	//! - Phase A (scalar): evaluates predicate exactly once, in order, into a bool mask.
	//! - Phase B (SVE): uses `COMPACT` instruction to compress retained elements.

	use core::{cmp, mem, ptr};

	use super::Vec;
	use crate::alloc::Allocator;

	const CHUNK_SIZE: usize = 64;

	pub(super) const MIN_SVE_SIZE_1: usize = 32;
	pub(super) const MIN_SVE_SIZE_2: usize = 32;
	pub(super) const MIN_SVE_SIZE_4: usize = 64;
	pub(super) const MIN_SVE_SIZE_8: usize = 64;

	/// Guard for the SVE retain path. On panic, this guard:
	/// 1. Scalar-compresses the already-decided prefix of the current chunk.
	/// 2. Copies the untouched tail forward.
	/// 3. Sets the Vec length.
	struct PanicGuard<'a, T, A: Allocator> {
	v: &'a mut Vec<T, A>,
	/// Start index of the current chunk within the Vec.
	read: usize,
	/// Write cursor (accumulated from previous chunks).
	write: usize,
	/// How many elements in the current chunk have been decided by the predicate.
	decided: usize,

	mask: &'a mut [bool; CHUNK_SIZE],
	/// Original length of the Vec.
	original_len: usize,
	}

	impl<T, A: Allocator> Drop for PanicGuard<'_, T, A> {
	#[cold]
	fn drop(&mut self) {
	// Scalar-compress the decided prefix: move kept elements to write position.
	let mut dst = self.write;
	for i in 0..self.decided {
	if self.mask[i] {
	// SAFETY: read + i < original_len (in-bounds).
	let src = unsafe { self.v.as_ptr().add(self.read + i) };
	// SAFETY: Continued from above.
	let dst_ptr = unsafe { self.v.as_mut_ptr().add(dst) };
	// SAFETY: src and dst_ptr < original_len
	unsafe { ptr::copy(src, dst_ptr, 1) };
	dst += 1;
	}
	}

	// Copy the untouched tail.
	let untouched_start = self.read + self.decided;
	let untouched_len = self.original_len - untouched_start;
	if untouched_len > 0 {
	// SAFETY: untouched_start..original_len are valid; dst + untouched_len <= original_len.
	unsafe {
	ptr::copy(
	self.v.as_ptr().add(untouched_start),
	self.v.as_mut_ptr().add(dst),
	untouched_len,
	);
	}
	dst += untouched_len;
	}

	// SAFETY: After filling holes, all items are in contiguous memory.
	unsafe { self.v.set_len(dst) };
	}
	}

	/// # Safety
	///
	/// - `size_of::<T>() == 1/2/4/8`
	/// - Called on aarch64 + SVE target.
	pub(crate) unsafe fn chunked_retain<T, F, A: Allocator>(v: &mut Vec<T, A>, mut f: F)
	where
	F: FnMut(&mut T) -> bool,
	{
	let original_len = v.len();
	let mut guard = PanicGuard {
	v,
	read: 0,
	write: 0,
	decided: 0,
	mask: &mut [false; CHUNK_SIZE],
	original_len,
	};

	while guard.read < guard.original_len {
	let chunk_len = cmp::min(CHUNK_SIZE, guard.original_len - guard.read);

	guard.decided = 0;

	// Phase A: scalar predicate evaluation (exactly once, in order).
	for i in 0..chunk_len {
	// SAFETY: read + i < original_len
	let cur = unsafe { &mut *guard.v.as_mut_ptr().add(guard.read + i) };
	guard.mask[i] = f(cur);
	guard.decided = i + 1;
	if !guard.mask[i] {
	// SAFETY: read + i < original_len, and after marking `mask` and `decided`,
	// the guard can properly handles the case where drop_in_place panics.
	unsafe { ptr::drop_in_place(cur) };
	}
	}

	// Phase B: SVE compress.
	let kept = match mem::size_of::<T>() {
	// SAFETY: write <= read and the dispatch guarantees size_of::<T>() matches
	// the kernel lane width.
	1 => unsafe {
	compact8_kernel(
	guard.v.as_mut_ptr().add(guard.read),
	guard.v.as_mut_ptr().add(guard.write),
	guard.mask.as_ptr(),
	chunk_len,
	)
	},
	// SAFETY: Same as above.
	2 => unsafe {
	compact16_kernel(
	guard.v.as_mut_ptr().add(guard.read),
	guard.v.as_mut_ptr().add(guard.write),
	guard.mask.as_ptr(),
	chunk_len,
	)
	},
	// SAFETY: Same as above.
	4 => unsafe {
	compact32_kernel(
	guard.v.as_mut_ptr().add(guard.read),
	guard.v.as_mut_ptr().add(guard.write),
	guard.mask.as_ptr(),
	chunk_len,
	)
	},
	// SAFETY: Same as above.
	8 => unsafe {
	compact64_kernel(
	guard.v.as_mut_ptr().add(guard.read),
	guard.v.as_mut_ptr().add(guard.write),
	guard.mask.as_ptr(),
	chunk_len,
	)
	},
	_ => unreachable!(),
	};

	guard.write += kept;
	guard.read += chunk_len;
	}

	// SAFETY: write <= original_len, all retained elements are packed at the front.
	unsafe { guard.v.set_len(guard.write) };
	mem::forget(guard);
	}

	macro_rules! sve_compact_kernel {
	(
	$name:ident,
	size = $size:literal,
	lane = $lane:literal,
	mem_lane = $mem_lane:literal,
	shift = $shift:literal,
	inc = $inc:literal
	) => {
	/// SVE compress kernel: pack retained elements from `src` to `dst`
	/// according to `mask`. Returns the number of retained elements.
	///
	/// # Safety
	///
	/// - `src` is valid for `chunk_len` reads of `T`, `dst` for `chunk_len`
	/// writes, `mask` for `chunk_len` bool reads.
	/// - `size_of::<T>()` equals this kernel's lane width in bytes.
	#[target_feature(enable = "sve")]
	#[inline]
	unsafe fn $name<T>(
	src: *const T,
	dst: *mut T,
	mask: *const bool,
	chunk_len: usize,
	) -> usize {
	debug_assert_eq!(mem::size_of::<T>(), $size);

	let idx_in = 0usize;
	let mut idx_out = 0usize;
	// SVE intrinsics require treating the data as integers, which doesn't
	// correctly handle provenance and uninitialized padding.
	//
	// SAFETY: whilelo predicates every load/store to the remaining elements.
	unsafe {
	core::arch::asm!(
	concat!("whilelo p0.", $lane, ", xzr, {len}"),
	"2:",
	concat!("ld1b {{ z0.", $lane, " }}, p0/z, [{mask}, {idx_in}]"),
	concat!("cmpne p1.", $lane, ", p0/z, z0.", $lane, ", #0"),
	concat!("ld1", $mem_lane, " {{ z1.", $lane, " }}, p1/z, [{src}, {idx_in}", $shift, "]"),
	concat!("compact z1.", $lane, ", p1, z1.", $lane),
	concat!("cntp {kept}, p0, p1.", $lane),
	concat!("whilelo p2.", $lane, ", xzr, {kept}"),
	concat!("st1", $mem_lane, " {{ z1.", $lane, " }}, p2, [{dst}, {idx_out}", $shift, "]"),
	concat!("add {idx_out}, {idx_out}, {kept}"),
	concat!("inc", $inc, " {idx_in}"),
	concat!("whilelo p0.", $lane, ", {idx_in}, {len}"),
	"b.first 2b",
	src = in(reg) src,
	dst = in(reg) dst,
	mask = in(reg) mask,
	len = in(reg) chunk_len,
	idx_in = inout(reg) idx_in => _,
	idx_out = inout(reg) idx_out,
	kept = out(reg) _,
	out("p0") _,
	out("p1") _,
	out("p2") _,
	out("z0") _,
	out("z1") _,
	options(nostack),
	);
	}
	idx_out
	}
	};
	}

	sve_compact_kernel!(compact8_kernel, size = 1, lane = "s", mem_lane = "b", shift = "", inc = "w");

	sve_compact_kernel!(
	compact16_kernel,
	size = 2,
	lane = "s",
	mem_lane = "h",
	shift = ", lsl #1",
	inc = "w"
	);

	sve_compact_kernel!(
	compact32_kernel,
	size = 4,
	lane = "s",
	mem_lane = "w",
	shift = ", lsl #2",
	inc = "w"
	);

	sve_compact_kernel!(
	compact64_kernel,
	size = 8,
	lane = "d",
	mem_lane = "d",
	shift = ", lsl #3",
	inc = "d"
	);