add extract_bits and deposit_bits intrinsics

compiler/rustc_codegen_llvm/src/back/llvm_backend.rs

@@ -216,7 +216,7 @@ impl CodegenBackend for LlvmCodegenBackend {
         // Intrinsics whose fallback body will not be used by the LLVM backend.
         let replaced_intrinsics = {
             #[rustfmt::skip]
-            let will_not_use_fallback = vec![
+            let mut will_not_use_fallback = vec![
                 // These are mapped to LLVM intrinsics instead.
                 sym::unchecked_funnel_shl,
                 sym::unchecked_funnel_shr,
@@ -243,6 +243,11 @@ impl CodegenBackend for LlvmCodegenBackend {
                 sym::copysignf16, sym::copysignf32, sym::copysignf64, sym::copysignf128,
             ];
 
+            if llvm_util::get_version() >= (23, 0, 0) {
+                will_not_use_fallback.push(sym::extract_bits);
+                will_not_use_fallback.push(sym::deposit_bits);
+            }
+
             will_not_use_fallback
         };
 

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -404,7 +404,26 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 let pair = self.insert_value(pair, high, 1);
                 pair
             }
+            // Checking LLVM version as these intrinsics were introduced in LLVM 23.
+            sym::deposit_bits | sym::extract_bits if llvm_version >= (23, 0, 0) => {
+                let ty = args[0].layout.ty;
+                if !ty.is_integral() {
+                    let err = tcx.dcx().emit_err(InvalidMonomorphization::BasicIntegerType {
+                        span,
+                        name,
+                        ty,
+                    });
+                    return IntrinsicResult::Err(err);
+                }
+                let (size, _) = ty.int_size_and_signed(self.tcx);
+                let llty = self.type_ix(size.bits());
 
+                self.call_intrinsic(
+                    if name == sym::extract_bits { "llvm.pext" } else { "llvm.pdep" },
+                    &[llty],
+                    &[args[0].immediate(), args[1].immediate()],
+                )
+            }
             sym::ctlz
             | sym::ctlz_nonzero
             | sym::cttz

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -98,9 +98,11 @@ fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -> hi
         | sym::ctlz
         | sym::ctpop
         | sym::cttz
+        | sym::deposit_bits
         | sym::discriminant_value
         | sym::exp
         | sym::exp2
+        | sym::extract_bits
         | sym::fabs
         | sym::fadd_algebraic
         | sym::fdiv_algebraic
@@ -747,6 +749,8 @@ pub(crate) fn check_intrinsic_type(
 
         sym::return_address => (0, 0, vec![], Ty::new_imm_ptr(tcx, tcx.types.unit)),
 
+        sym::extract_bits | sym::deposit_bits => (1, 0, vec![param(0), param(0)], param(0)),
+
         other => {
             tcx.dcx().emit_err(UnrecognizedIntrinsicFunction { span, name: other });
             return;

compiler/rustc_span/src/symbol.rs

@@ -831,6 +831,7 @@ symbols! {
         default_type_params,
         define_opaque,
         deny,
+        deposit_bits,
         deprecated,
         deprecated_safe,
         deprecated_suggestion,
@@ -967,6 +968,7 @@ symbols! {
         extern_weak,
         external,
         external_doc,
+        extract_bits,
         f16,
         f16_nan,
         f16b,

library/core/src/intrinsics/fallback.rs

@@ -258,3 +258,224 @@ macro_rules! impl_carryless_mul{
 impl_carryless_mul! {
     u8, u16, u32, u64, u128, usize
 }
+
+/// Implementations for `uN::extract_bits` and `uN::deposit_bits`
+///
+/// For the purposes of this implementation, the operations can be thought
+/// of as operating on the input bits as a list, starting from the least
+/// significant bit. Extraction is like `Vec::retain` that deletes bits
+/// where the mask has a zero. Deposition is like doing the inverse by
+/// inserting the zeros that extraction would delete.
+///
+/// Key observation: Each extracted or deposited bit needs to be
+/// shifted by the count of zeros up to the corresponding mask bit.
+///
+/// With that in mind, the general idea is to decompose the operation into
+/// a sequence of stages in `0..log2(BITS)`, where each stage shifts some
+/// of the bits by `n = 1 << stage`. The masks for each stage are computed
+/// via prefix counts of zeros in the mask.
+///
+/// # Extraction
+///
+/// Consider the input as a sequence of runs of data (bitstrings A,B,C,...),
+/// split by fixed-width groups of zeros ('.'), initially at width `n = 1`.
+/// Counting the groups of zeros, each stage shifts the odd-indexed runs of
+/// data right by `n`, effectively swapping them with the preceding zeros.
+/// For the next stage, `n` is doubled as all the zeros are now paired.
+/// ```text
+/// .A.B.C.D.E.F.G.H
+/// ..AB..CD..EF..GH
+/// ....ABCD....EFGH
+/// ........ABCDEFGH
+/// ```
+/// What makes this nontrivial is that the lengths of the bitstrings are not
+/// the same. Using lowercase for individual bits, the above might look like
+/// ```text
+/// .a.bbb.ccccc.dd.e..g.hh
+/// ..abbb..cccccdd..e..ghh
+/// ....abbbcccccdd....eghh
+/// ........abbbcccccddeghh
+/// ```
+///
+/// # Deposition
+///
+/// For `deposit_bits`, the stages are reversed. We start with a single run of
+/// data in the low bits. Each stage then splits each run of data in two by
+/// shifting part of it left by `n`, which is halved each stage.
+/// ```text
+/// ........ABCDEFGH
+/// ....ABCD....EFGH
+/// ..AB..CD..EF..GH
+/// .A.B.C.D.E.F.G.H
+/// ```
+///
+/// # Stage masks
+///
+/// To facilitate the shifts at each stage, we compute a mask that covers both
+/// the bitstrings to shift, and the zeros they shift into.
+/// ```text
+/// .A.B.C.D.E.F.G.H
+///  ##  ##  ##  ##
+/// ..AB..CD..EF..GH
+///   ####    ####
+/// ....ABCD....EFGH
+///     ########
+/// ........ABCDEFGH
+/// ```
+#[rustc_const_unstable(feature = "core_intrinsics_fallbacks", issue = "none")]
+pub const trait ExtractDepositBits: Copy + 'static {
+    fn extract_bits(self, mask: Self) -> Self;
+    fn deposit_bits(self, mask: Self) -> Self;
+}
+
+macro_rules! impl_extract_deposit_bits {
+    ($($U:ident),*) => {$(
+        mod $U {
+            use super::ExtractDepositBits;
+
+            const STAGES: usize = $U::BITS.ilog2() as usize;
+
+            #[inline]
+            const fn prepare(sparse: $U) -> [$U; STAGES] {
+                // We'll start with `zeros` as a mask of the bits to be removed,
+                // and compute into `masks` the parts that shift at each stage.
+                let mut zeros = !sparse;
+                let mut masks = [0; STAGES];
+                let mut stage = 0;
+                while stage < STAGES {
+                    let n = 1 << stage;
+                    // Suppose `zeros` has bits set at ranges `{ a..a+n, b..b+n, ... }`.
+                    // Then `parity` will be computed as `{ a.. } XOR { b.. } XOR ...`,
+                    // which will be the ranges `{ a..b, c..d, e.. }`.
+                    let mut parity = zeros;
+                    let mut len = n;
+                    while len < $U::BITS {
+                        parity ^= parity << len;
+                        len <<= 1;
+                    }
+                    masks[stage] = parity;
+
+                    // Toggle off the bits that are shifted into:
+                    // { a..a+n, b..b+n, ... } & !{ a..b, c..d, e.. }
+                    // == { b..b+n, d..d+n, ... }
+                    zeros &= !parity;
+                    // Expand the remaining ranges down to the bits that were
+                    // shifted from: { b-n..b+n, d-n..d+n, ... }
+                    zeros ^= zeros >> n;
+
+                    stage += 1;
+                }
+                masks
+            }
+
+            #[inline(always)]
+            const fn extract_impl(mut x: $U, sparse: $U) -> $U {
+                let masks = prepare(sparse);
+                x &= sparse;
+                let mut stage = 0;
+                while stage < STAGES {
+                    let n = 1 << stage;
+                    // Consider each two runs of data with their leading
+                    // groups of `n` 0-bits. Suppose that the run that is
+                    // shifted right has length `a`, and the other one has
+                    // length `b`. Assume that only zeros are shifted in.
+                    // ```text
+                    // [0; n], [X; a], [0; n], [Y; b] // x
+                    // [0; n], [X; a], [0; n], [0; b] // q
+                    // [0; n], [0; a   +   n], [Y; b] // x ^= q
+                    // [0; n   +   n], [X; a], [0; b] // q >> n
+                    // [0; n], [0; n], [X; a], [Y; b] // x ^= q >> n
+                    // ```
+                    // Only zeros are shifted out, satisfying the assumption
+                    // for the next group.
+
+                    // In effect, the upper run of data is swapped with the
+                    // group of `n` zeros below it.
+                    let q = x & masks[stage];
+
+                    if cfg!(all(target_arch = "aarch64", target_feature = "neon")) && $U::BITS <= 64
+                    {
+                        // `q >> n` has no set bits in common with `x ^ q`, so addition
+                        // is equivalent to XOR here. For widths up to 64 bits this
+                        // allows AArch64 to use `usra`. For `u128`, this add may introduce
+                        // a slower cross-limb carry chain.
+                        x = (x & !masks[stage]).wrapping_add(q >> n);
+                    } else {
+                        x ^= q;
+                        x ^= q >> n;
+                    }
+                    stage += 1;
+                }
+                x
+            }
+
+            #[inline(always)]
+            const fn deposit_impl(mut x: $U, sparse: $U) -> $U {
+                let masks = prepare(sparse);
+                let mut stage = STAGES;
+                while stage > 0 {
+                    stage -= 1;
+                    let n = 1 << stage;
+                    // Consider each run of data with the `2 * n` arbitrary bits
+                    // above it. Suppose that the run has length `a + b`, with
+                    // `a` being the length of the part that needs to be
+                    // shifted. Assume that only zeros are shifted in.
+                    // ```text
+                    // [_; n], [_; n], [X; a], [Y; b] // x
+                    // [0; n], [_; n], [X; a], [0; b] // q
+                    // [_; n], [0; n   +   a], [Y; b] // x ^= q
+                    // [_; n], [X; a], [0; b   +   n] // q << n
+                    // [_; n], [X; a], [0; n], [Y; b] // x ^= q << n
+                    // ```
+                    // Only zeros are shifted out, satisfying the assumption
+                    // for the next group.
+
+                    // In effect, `n` 0-bits are inserted somewhere in each run
+                    // of data to spread it, and the two groups of `n` bits
+                    // above are XOR'd together.
+                    let q = x & masks[stage];
+                    x ^= q;
+                    x ^= q << n;
+                }
+                x & sparse
+            }
+
+            #[rustc_const_unstable(feature = "core_intrinsics_fallbacks", issue = "none")]
+            const impl ExtractDepositBits for $U {
+                #[inline]
+                fn extract_bits(self, mask: Self) -> Self {
+                    extract_impl(self, mask)
+                }
+
+                #[inline]
+                fn deposit_bits(self, mask: Self) -> Self {
+                    deposit_impl(self, mask)
+                }
+            }
+        }
+    )*};
+}
+
+impl_extract_deposit_bits!(u8, u16, u32, u64, u128);
+
+#[cfg(target_pointer_width = "16")]
+type UsizeActual = u16;
+#[cfg(target_pointer_width = "32")]
+type UsizeActual = u32;
+#[cfg(target_pointer_width = "64")]
+type UsizeActual = u64;
+
+#[rustc_const_unstable(feature = "core_intrinsics_fallbacks", issue = "none")]
+const impl ExtractDepositBits for usize {
+    #[inline]
+    fn extract_bits(self, mask: Self) -> Self {
+        <UsizeActual as ExtractDepositBits>::extract_bits(self as UsizeActual, mask as UsizeActual)
+            as usize
+    }
+
+    #[inline]
+    fn deposit_bits(self, mask: Self) -> Self {
+        <UsizeActual as ExtractDepositBits>::deposit_bits(self as UsizeActual, mask as UsizeActual)
+            as usize
+    }
+}

library/core/src/intrinsics/mod.rs

@@ -3718,3 +3718,27 @@ pub const unsafe fn va_end(ap: &mut VaList<'_>) {
 pub fn return_address() -> *const () {
     core::ptr::null()
 }
+
+/// Returns an integer with the bit locations specified by `mask` packed
+/// contiguously into the least significant bits of the result.
+///
+/// This corresponds to `@llvm.pext.*`.
+#[rustc_intrinsic]
+#[rustc_nounwind]
+#[rustc_const_unstable(feature = "core_intrinsics", issue = "none")]
+#[miri::intrinsic_fallback_is_spec]
+pub const fn extract_bits<T: [const] fallback::ExtractDepositBits>(x: T, mask: T) -> T {
+    x.extract_bits(mask)
+}
+
+/// Returns an integer with the least significant bits of `x`
+/// distributed to the bit locations specified by `mask`.
+///
+/// This corresponds to `@llvm.pdep.*`.
+#[rustc_intrinsic]
+#[rustc_nounwind]
+#[rustc_const_unstable(feature = "core_intrinsics", issue = "none")]
+#[miri::intrinsic_fallback_is_spec]
+pub const fn deposit_bits<T: [const] fallback::ExtractDepositBits>(x: T, mask: T) -> T {
+    x.deposit_bits(mask)
+}

library/core/src/lib.rs

@@ -162,6 +162,7 @@
 #![feature(transparent_unions)]
 #![feature(try_blocks)]
 #![feature(uint_carryless_mul)]
+#![feature(uint_gather_scatter_bits)]
 #![feature(unboxed_closures)]
 #![feature(unsized_fn_params)]
 #![feature(with_negative_coherence)]