Auto merge of rust-lang#134954 - scottmcm:more-swap-tuning, r=<try>

Redo the swap code for better tail & padding handling

A couple of parts here

## Fixes rust-lang#134713

Actually using the type passed to `ptr::swap_nonoverlapping` for anything other than its size + align turns out to not work, so this redo goes back to *always* swapping via not-`!noundef` integers.

(Except in `const`, which keeps doing the same thing as before to preserve `@RalfJung's` fix from rust-lang#134689)

## Fixes rust-lang#134946

I'd previously moved the swapping to use auto-vectorization, but someone pointed out on Discord that the tail loop handling from that left a whole bunch of byte-by-byte swapping around.  This PR goes back to a manual chunk, with at most logarithmic more instructions for the tail.

(There are other ways that could potentially handle the tail even better, but this seems to be pretty good, since it's how LLVM ends up lowering operations on types like `i56`.)

## Polymorphization

Since it's silly to have separate copies of swapping -- especially *untyped* swapping! -- for `u32`, `i32`, `f32`, `[u16; 2]`, etc, this sends everything to byte versions, but still mono'd by alignment.  That should make it more ok that the code is somewhat more complex, since we only get 7 monomorphizations of the complicated bit.

(One day we'll be able to remove some of the hacks by being able to just call `foo::<{align_of::<T>()}>`, but since alignments are only powers of two, the manual dispatch out isn't a big deal.)

Author: bors

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -498,6 +498,23 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 }
             }
 
+            sym::untyped_swap_nonoverlapping => {
+                // The fallback impl uses memcpy, which leaves around allocas
+                // that don't optimize out for certain widths, so force it to
+                // use SSA registers instead.
+
+                let chunk_ty = fn_args.type_at(0);
+                let layout = self.layout_of(chunk_ty).layout;
+                let integer_ty = self.type_ix(layout.size().bits());
+                let a = args[0].immediate();
+                let b = args[1].immediate();
+                let a_val = self.load(integer_ty, a, layout.align().abi);
+                let b_val = self.load(integer_ty, b, layout.align().abi);
+                self.store(b_val, a, layout.align().abi);
+                self.store(a_val, b, layout.align().abi);
+                return Ok(());
+            }
+
             sym::compare_bytes => {
                 // Here we assume that the `memcmp` provided by the target is a NOP for size 0.
                 let cmp = self.call_intrinsic("memcmp", &[

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -504,6 +504,12 @@ pub fn check_intrinsic_type(
             sym::typed_swap_nonoverlapping => {
                 (1, 0, vec![Ty::new_mut_ptr(tcx, param(0)); 2], tcx.types.unit)
             }
+            sym::untyped_swap_nonoverlapping => (
+                1,
+                0,
+                vec![Ty::new_mut_ptr(tcx, Ty::new_maybe_uninit(tcx, param(0))); 2],
+                tcx.types.unit,
+            ),
 
             sym::discriminant_value => {
                 let assoc_items = tcx.associated_item_def_ids(

compiler/rustc_span/src/symbol.rs

@@ -2142,6 +2142,7 @@ symbols! {
                           unstable location; did you mean to load this crate \
                           from crates.io via `Cargo.toml` instead?",
         untagged_unions,
+        untyped_swap_nonoverlapping,
         unused_imports,
         unwind,
         unwind_attributes,

library/core/src/intrinsics/mod.rs

@@ -66,7 +66,7 @@
 
 use crate::marker::{DiscriminantKind, Tuple};
 use crate::mem::SizedTypeProperties;
-use crate::{ptr, ub_checks};
+use crate::{mem, ptr, ub_checks};
 
 pub mod fallback;
 pub mod mir;
@@ -4003,7 +4003,37 @@ pub use typed_swap as typed_swap_nonoverlapping;
 pub const unsafe fn typed_swap_nonoverlapping<T>(x: *mut T, y: *mut T) {
     // SAFETY: The caller provided single non-overlapping items behind
     // pointers, so swapping them with `count: 1` is fine.
-    unsafe { ptr::swap_nonoverlapping(x, y, 1) };
+    unsafe { crate::swapping::swap_nonoverlapping(x, y, 1) };
+}
+
+/// Swaps the `N` untyped & non-overlapping bytes behind the two pointers.
+///
+/// Split out from `typed_swap` for the internal swaps in `swap_nonoverlapping`
+/// which would otherwise cause cycles between the fallback implementations on
+/// backends where neither is overridden.
+///
+/// # Safety
+///
+/// `x` and `y` are readable and writable as `MaybeUninit<C>` and non-overlapping.
+#[inline]
+#[rustc_nounwind]
+#[cfg_attr(not(bootstrap), rustc_intrinsic)]
+#[miri::intrinsic_fallback_is_spec]
+#[rustc_const_stable_indirect]
+pub const unsafe fn untyped_swap_nonoverlapping<C>(
+    x: *mut mem::MaybeUninit<C>,
+    y: *mut mem::MaybeUninit<C>,
+) {
+    // This intentionally uses untyped memory copies, not reads/writes,
+    // to avoid any risk of losing padding in things like (u16, u8).
+    let mut temp = mem::MaybeUninit::<C>::uninit();
+    // SAFETY: Caller promised that x and y are non-overlapping & read/writeable,
+    // and our fresh local is always disjoint from anything otherwise readable.
+    unsafe {
+        (&raw mut temp).copy_from_nonoverlapping(x, 1);
+        x.copy_from_nonoverlapping(y, 1);
+        y.copy_from_nonoverlapping(&raw const temp, 1);
+    }
 }
 
 /// Returns whether we should perform some UB-checking at runtime. This eventually evaluates to

library/core/src/lib.rs

@@ -376,6 +376,7 @@ pub mod alloc;
 // note: does not need to be public
 mod bool;
 mod escape;
+pub(crate) mod swapping;
 mod tuple;
 mod unit;
 

library/core/src/ptr/mod.rs

@@ -395,7 +395,6 @@
 #![allow(clippy::not_unsafe_ptr_arg_deref)]
 
 use crate::cmp::Ordering;
-use crate::intrinsics::const_eval_select;
 use crate::marker::FnPtr;
 use crate::mem::{self, MaybeUninit, SizedTypeProperties};
 use crate::{fmt, hash, intrinsics, ub_checks};
@@ -1092,84 +1091,8 @@ pub const unsafe fn swap_nonoverlapping<T>(x: *mut T, y: *mut T, count: usize) {
         }
     );
 
-    const_eval_select!(
-        @capture[T] { x: *mut T, y: *mut T, count: usize }:
-        if const {
-            // At compile-time we want to always copy this in chunks of `T`, to ensure that if there
-            // are pointers inside `T` we will copy them in one go rather than trying to copy a part
-            // of a pointer (which would not work).
-            // SAFETY: Same preconditions as this function
-            unsafe { swap_nonoverlapping_simple_untyped(x, y, count) }
-        } else {
-            macro_rules! attempt_swap_as_chunks {
-                ($ChunkTy:ty) => {
-                    if mem::align_of::<T>() >= mem::align_of::<$ChunkTy>()
-                        && mem::size_of::<T>() % mem::size_of::<$ChunkTy>() == 0
-                    {
-                        let x: *mut $ChunkTy = x.cast();
-                        let y: *mut $ChunkTy = y.cast();
-                        let count = count * (mem::size_of::<T>() / mem::size_of::<$ChunkTy>());
-                        // SAFETY: these are the same bytes that the caller promised were
-                        // ok, just typed as `MaybeUninit<ChunkTy>`s instead of as `T`s.
-                        // The `if` condition above ensures that we're not violating
-                        // alignment requirements, and that the division is exact so
-                        // that we don't lose any bytes off the end.
-                        return unsafe { swap_nonoverlapping_simple_untyped(x, y, count) };
-                    }
-                };
-            }
-
-            // Split up the slice into small power-of-two-sized chunks that LLVM is able
-            // to vectorize (unless it's a special type with more-than-pointer alignment,
-            // because we don't want to pessimize things like slices of SIMD vectors.)
-            if mem::align_of::<T>() <= mem::size_of::<usize>()
-            && (!mem::size_of::<T>().is_power_of_two()
-                || mem::size_of::<T>() > mem::size_of::<usize>() * 2)
-            {
-                attempt_swap_as_chunks!(usize);
-                attempt_swap_as_chunks!(u8);
-            }
-
-            // SAFETY: Same preconditions as this function
-            unsafe { swap_nonoverlapping_simple_untyped(x, y, count) }
-        }
-    )
-}
-
-/// Same behavior and safety conditions as [`swap_nonoverlapping`]
-///
-/// LLVM can vectorize this (at least it can for the power-of-two-sized types
-/// `swap_nonoverlapping` tries to use) so no need to manually SIMD it.
-#[inline]
-const unsafe fn swap_nonoverlapping_simple_untyped<T>(x: *mut T, y: *mut T, count: usize) {
-    let x = x.cast::<MaybeUninit<T>>();
-    let y = y.cast::<MaybeUninit<T>>();
-    let mut i = 0;
-    while i < count {
-        // SAFETY: By precondition, `i` is in-bounds because it's below `n`
-        let x = unsafe { x.add(i) };
-        // SAFETY: By precondition, `i` is in-bounds because it's below `n`
-        // and it's distinct from `x` since the ranges are non-overlapping
-        let y = unsafe { y.add(i) };
-
-        // If we end up here, it's because we're using a simple type -- like
-        // a small power-of-two-sized thing -- or a special type with particularly
-        // large alignment, particularly SIMD types.
-        // Thus, we're fine just reading-and-writing it, as either it's small
-        // and that works well anyway or it's special and the type's author
-        // presumably wanted things to be done in the larger chunk.
-
-        // SAFETY: we're only ever given pointers that are valid to read/write,
-        // including being aligned, and nothing here panics so it's drop-safe.
-        unsafe {
-            let a: MaybeUninit<T> = read(x);
-            let b: MaybeUninit<T> = read(y);
-            write(x, b);
-            write(y, a);
-        }
-
-        i += 1;
-    }
+    // SAFETY: Same preconditions as this function
+    unsafe { crate::swapping::swap_nonoverlapping(x, y, count) }
 }
 
 /// Moves `src` into the pointed `dst`, returning the previous `dst` value.

library/core/src/swapping.rs

@@ -0,0 +1,182 @@
+use crate::{hint, intrinsics, mem, ptr};
+
+//#[rustc_const_stable_indirect]
+//#[rustc_allow_const_fn_unstable(const_eval_select)]
+#[rustc_const_unstable(feature = "const_swap_nonoverlapping", issue = "133668")]
+#[inline]
+pub(crate) const unsafe fn swap_nonoverlapping<T>(x: *mut T, y: *mut T, count: usize) {
+    intrinsics::const_eval_select!(
+        @capture[T] { x: *mut T, y: *mut T, count: usize }:
+        if const {
+            // At compile-time we want to always copy this in chunks of `T`, to ensure that if there
+            // are pointers inside `T` we will copy them in one go rather than trying to copy a part
+            // of a pointer (which would not work).
+            // SAFETY: Same preconditions as this function
+            unsafe { swap_nonoverlapping_const(x, y, count) }
+        } else {
+            // At runtime we want to make sure not to swap byte-for-byte for types like [u8; 15],
+            // and swapping as `MaybeUninit<T>` doesn't actually work as untyped for things like
+            // T = (u16, u8), so we type-erase to raw bytes and swap that way.
+            // SAFETY: Same preconditions as this function
+            unsafe { swap_nonoverlapping_runtime(x, y, count) }
+        }
+    )
+}
+
+/// Same behavior and safety conditions as [`swap_nonoverlapping`]
+#[rustc_const_stable_indirect]
+#[inline]
+const unsafe fn swap_nonoverlapping_const<T>(x: *mut T, y: *mut T, count: usize) {
+    let x = x.cast::<mem::MaybeUninit<T>>();
+    let y = y.cast::<mem::MaybeUninit<T>>();
+    let mut i = 0;
+    while i < count {
+        // SAFETY: By precondition, `i` is in-bounds because it's below `n`
+        // and because the two input ranges are non-overlapping and read/writeable,
+        // these individual items inside them are too.
+        unsafe {
+            intrinsics::untyped_swap_nonoverlapping::<T>(x.add(i), y.add(i));
+        }
+
+        i += 1;
+    }
+}
+
+// Scale the monomorphizations with the size of the machine, roughly.
+const MAX_ALIGN: usize = align_of::<usize>().pow(2);
+
+/// Same behavior and safety conditions as [`swap_nonoverlapping`]
+#[inline]
+unsafe fn swap_nonoverlapping_runtime<T>(x: *mut T, y: *mut T, count: usize) {
+    let bytes = {
+        let slice = ptr::slice_from_raw_parts(x, count);
+        // SAFETY: Because they both exist in memory and don't overlap, they
+        // must be legal slice sizes (below `isize::MAX` bytes).
+        unsafe { mem::size_of_val_raw(slice) }
+    };
+
+    // Generating *untyped* loops for every type is silly, so we polymorphize away
+    // the actual type, but we want to take advantage of alignment if possible,
+    // so monomorphize for a restricted set of possible alignments.
+    macro_rules! delegate_by_alignment {
+        ($($p:pat => $align:expr,)+) => {{
+            #![allow(unreachable_patterns)]
+            match const { align_of::<T>() } {
+                $(
+                    $p => {
+                        swap_nonoverlapping_bytes::<$align>(x.cast(), y.cast(), bytes);
+                    }
+                )+
+            }
+        }};
+    }
+
+    // SAFETY:
+    unsafe {
+        delegate_by_alignment! {
+            MAX_ALIGN.. => MAX_ALIGN,
+            64.. => 64,
+            32.. => 32,
+            16.. => 16,
+            8.. => 8,
+            4.. => 4,
+            2.. => 2,
+            _ => 1,
+        }
+    }
+}
+
+/// # Safety:
+/// - `x` and `y` must be aligned to `ALIGN`
+/// - `bytes` must be a multiple of `ALIGN`
+/// - They must be readable, writable, and non-overlapping for `bytes` bytes
+#[inline]
+unsafe fn swap_nonoverlapping_bytes<const ALIGN: usize>(
+    x: *mut mem::MaybeUninit<u8>,
+    y: *mut mem::MaybeUninit<u8>,
+    bytes: usize,
+) {
+    // SAFETY: Two legal non-overlapping regions can't be bigger than this.
+    // (And they couldn't have made allocations any bigger either anyway.)
+    // FIXME: Would be nice to have a type for this instead of the assume.
+    unsafe { hint::assert_unchecked(bytes < isize::MAX as usize) };
+
+    let mut i = 0;
+    macro_rules! swap_next_n {
+        ($n:expr) => {{
+            let x: *mut mem::MaybeUninit<[u8; $n]> = x.add(i).cast();
+            let y: *mut mem::MaybeUninit<[u8; $n]> = y.add(i).cast();
+            swap_nonoverlapping_aligned_chunk::<ALIGN, [u8; $n]>(
+                x.as_mut_unchecked(),
+                y.as_mut_unchecked(),
+            );
+            i += $n;
+        }};
+    }
+
+    while bytes - i >= MAX_ALIGN {
+        const { assert!(MAX_ALIGN >= ALIGN) };
+        // SAFETY: the const-assert above confirms we're only ever called with
+        // an alignment equal to or smaller than max align, so this is necessarily
+        // aligned, and the while loop ensures there's enough read/write memory.
+        unsafe {
+            swap_next_n!(MAX_ALIGN);
+        }
+    }
+
+    macro_rules! handle_tail {
+        ($($n:literal)+) => {$(
+            if const { $n % ALIGN == 0 } {
+                // Checking this way simplifies the block end to just add+test,
+                // rather than needing extra math before the check.
+                if (bytes & $n) != 0 {
+                    // SAFETY: The above swaps were bigger, so could not have
+                    // impacted the `$n`-relevant bit, so checking `bytes & $n`
+                    // was equivalent to `bytes - i >= $n`, and thus we have
+                    // enough space left to swap another `$n` bytes.
+                    unsafe {
+                        swap_next_n!($n);
+                    }
+                }
+            }
+        )+};
+    }
+    const { assert!(MAX_ALIGN <= 64) };
+    handle_tail!(32 16 8 4 2 1);
+
+    debug_assert_eq!(i, bytes);
+}
+
+/// Swaps the `C` behind `x` and `y` as untyped memory
+///
+/// # Safety
+///
+/// Both `x` and `y` must be aligned to `ALIGN`, in addition to their normal alignment.
+/// They must be readable and writeable for `sizeof(C)` bytes, as usual for `&mut`s.
+///
+/// (The actual instantiations are usually `C = [u8; _]`, so we get the alignment
+/// information from the loads by `assume`ing the passed-in alignment.)
+// Don't let MIR inline this, because we really want it to keep its noalias metadata
+#[rustc_no_mir_inline]
+#[inline]
+unsafe fn swap_nonoverlapping_aligned_chunk<const ALIGN: usize, C>(
+    x: &mut mem::MaybeUninit<C>,
+    y: &mut mem::MaybeUninit<C>,
+) {
+    assert!(size_of::<C>() % ALIGN == 0);
+
+    let x = ptr::from_mut(x);
+    let y = ptr::from_mut(y);
+
+    // SAFETY: One of our preconditions.
+    unsafe {
+        hint::assert_unchecked(x.is_aligned_to(ALIGN));
+        hint::assert_unchecked(y.is_aligned_to(ALIGN));
+    }
+
+    // SAFETY: The memory is readable and writable because these were passed to
+    // us as mutable references, and the untyped swap doesn't need validity.
+    unsafe {
+        intrinsics::untyped_swap_nonoverlapping::<C>(x, y);
+    }
+}