Rollup merge of #137447 - folkertdev:simd-extract-insert-dyn, r=scottmcm

add `core::intrinsics::simd::{simd_extract_dyn, simd_insert_dyn}`

fixes #137372

adds `core::intrinsics::simd::{simd_extract_dyn, simd_insert_dyn}`, which contrary to their non-dyn counterparts allow a non-const index. Many platforms (but notably not x86_64 or aarch64) have dedicated instructions for this operation, which stdarch can emit with this change.

Future work is to also make the `Index` operation on the `Simd` type emit this operation, but the intrinsic can't be used directly. We'll need some MIR shenanigans for that.

r? `@ghost`

Author: Zalathar

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1421,7 +1421,7 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
         return Ok(bx.shuffle_vector(args[0].immediate(), args[1].immediate(), indices));
     }
 
-    if name == sym::simd_insert {
+    if name == sym::simd_insert || name == sym::simd_insert_dyn {
         require!(
             in_elem == arg_tys[2],
             InvalidMonomorphization::InsertedType {
@@ -1432,40 +1432,49 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
                 out_ty: arg_tys[2]
             }
         );
-        let idx = bx
-            .const_to_opt_u128(args[1].immediate(), false)
-            .expect("typeck should have ensure that this is a const");
-        if idx >= in_len.into() {
-            return_error!(InvalidMonomorphization::SimdIndexOutOfBounds {
-                span,
-                name,
-                arg_idx: 1,
-                total_len: in_len.into(),
-            });
-        }
-        return Ok(bx.insert_element(
-            args[0].immediate(),
-            args[2].immediate(),
-            bx.const_i32(idx as i32),
-        ));
+
+        let index_imm = if name == sym::simd_insert {
+            let idx = bx
+                .const_to_opt_u128(args[1].immediate(), false)
+                .expect("typeck should have ensure that this is a const");
+            if idx >= in_len.into() {
+                return_error!(InvalidMonomorphization::SimdIndexOutOfBounds {
+                    span,
+                    name,
+                    arg_idx: 1,
+                    total_len: in_len.into(),
+                });
+            }
+            bx.const_i32(idx as i32)
+        } else {
+            args[1].immediate()
+        };
+
+        return Ok(bx.insert_element(args[0].immediate(), args[2].immediate(), index_imm));
     }
-    if name == sym::simd_extract {
+    if name == sym::simd_extract || name == sym::simd_extract_dyn {
         require!(
             ret_ty == in_elem,
             InvalidMonomorphization::ReturnType { span, name, in_elem, in_ty, ret_ty }
         );
-        let idx = bx
-            .const_to_opt_u128(args[1].immediate(), false)
-            .expect("typeck should have ensure that this is a const");
-        if idx >= in_len.into() {
-            return_error!(InvalidMonomorphization::SimdIndexOutOfBounds {
-                span,
-                name,
-                arg_idx: 1,
-                total_len: in_len.into(),
-            });
-        }
-        return Ok(bx.extract_element(args[0].immediate(), bx.const_i32(idx as i32)));
+        let index_imm = if name == sym::simd_extract {
+            let idx = bx
+                .const_to_opt_u128(args[1].immediate(), false)
+                .expect("typeck should have ensure that this is a const");
+            if idx >= in_len.into() {
+                return_error!(InvalidMonomorphization::SimdIndexOutOfBounds {
+                    span,
+                    name,
+                    arg_idx: 1,
+                    total_len: in_len.into(),
+                });
+            }
+            bx.const_i32(idx as i32)
+        } else {
+            args[1].immediate()
+        };
+
+        return Ok(bx.extract_element(args[0].immediate(), index_imm));
     }
 
     if name == sym::simd_select {

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -659,8 +659,12 @@ pub(crate) fn check_intrinsic_type(
             sym::simd_masked_load => (3, 0, vec![param(0), param(1), param(2)], param(2)),
             sym::simd_masked_store => (3, 0, vec![param(0), param(1), param(2)], tcx.types.unit),
             sym::simd_scatter => (3, 0, vec![param(0), param(1), param(2)], tcx.types.unit),
-            sym::simd_insert => (2, 0, vec![param(0), tcx.types.u32, param(1)], param(0)),
-            sym::simd_extract => (2, 0, vec![param(0), tcx.types.u32], param(1)),
+            sym::simd_insert | sym::simd_insert_dyn => {
+                (2, 0, vec![param(0), tcx.types.u32, param(1)], param(0))
+            }
+            sym::simd_extract | sym::simd_extract_dyn => {
+                (2, 0, vec![param(0), tcx.types.u32], param(1))
+            }
             sym::simd_cast
             | sym::simd_as
             | sym::simd_cast_ptr

compiler/rustc_span/src/symbol.rs

@@ -1913,6 +1913,7 @@ symbols! {
         simd_eq,
         simd_expose_provenance,
         simd_extract,
+        simd_extract_dyn,
         simd_fabs,
         simd_fcos,
         simd_fexp,
@@ -1931,6 +1932,7 @@ symbols! {
         simd_ge,
         simd_gt,
         simd_insert,
+        simd_insert_dyn,
         simd_le,
         simd_lt,
         simd_masked_load,

library/core/src/intrinsics/simd.rs

@@ -4,7 +4,7 @@
 
 /// Inserts an element into a vector, returning the updated vector.
 ///
-/// `T` must be a vector with element type `U`.
+/// `T` must be a vector with element type `U`, and `idx` must be `const`.
 ///
 /// # Safety
 ///
@@ -15,15 +15,48 @@ pub const unsafe fn simd_insert<T, U>(x: T, idx: u32, val: U) -> T;
 
 /// Extracts an element from a vector.
 ///
-/// `T` must be a vector with element type `U`.
+/// `T` must be a vector with element type `U`, and `idx` must be `const`.
 ///
 /// # Safety
 ///
-/// `idx` must be in-bounds of the vector.
+/// `idx` must be const and in-bounds of the vector.
 #[rustc_intrinsic]
 #[rustc_nounwind]
 pub const unsafe fn simd_extract<T, U>(x: T, idx: u32) -> U;
 
+/// Inserts an element into a vector, returning the updated vector.
+///
+/// `T` must be a vector with element type `U`.
+///
+/// If the index is `const`, [`simd_insert`] may emit better assembly.
+///
+/// # Safety
+///
+/// `idx` must be in-bounds of the vector.
+#[rustc_nounwind]
+#[cfg_attr(not(bootstrap), rustc_intrinsic)]
+pub unsafe fn simd_insert_dyn<T, U>(mut x: T, idx: u32, val: U) -> T {
+    // SAFETY: `idx` must be in-bounds
+    unsafe { (&raw mut x).cast::<U>().add(idx as usize).write(val) }
+    x
+}
+
+/// Extracts an element from a vector.
+///
+/// `T` must be a vector with element type `U`.
+///
+/// If the index is `const`, [`simd_extract`] may emit better assembly.
+///
+/// # Safety
+///
+/// `idx` must be in-bounds of the vector.
+#[rustc_nounwind]
+#[cfg_attr(not(bootstrap), rustc_intrinsic)]
+pub unsafe fn simd_extract_dyn<T, U>(x: T, idx: u32) -> U {
+    // SAFETY: `idx` must be in-bounds
+    unsafe { (&raw const x).cast::<U>().add(idx as usize).read() }
+}
+
 /// Adds two simd vectors elementwise.
 ///
 /// `T` must be a vector of integers or floats.

tests/codegen/simd/extract-insert-dyn.rs

@@ -0,0 +1,75 @@
+//@compile-flags: -C opt-level=3 -C no-prepopulate-passes
+
+#![feature(core_intrinsics, repr_simd)]
+#![no_std]
+#![crate_type = "lib"]
+#![allow(non_camel_case_types)]
+
+// Test that `core::intrinsics::simd::{simd_extract_dyn, simd_insert_dyn}`
+// lower to an LLVM extractelement or insertelement operation.
+
+use core::intrinsics::simd::{simd_extract, simd_extract_dyn, simd_insert, simd_insert_dyn};
+
+#[repr(simd)]
+#[derive(Clone, Copy)]
+pub struct u32x16([u32; 16]);
+
+#[repr(simd)]
+#[derive(Clone, Copy)]
+pub struct i8x16([i8; 16]);
+
+// CHECK-LABEL: dyn_simd_extract
+// CHECK: extractelement <16 x i8> %x, i32 %idx
+#[no_mangle]
+unsafe extern "C" fn dyn_simd_extract(x: i8x16, idx: u32) -> i8 {
+    simd_extract_dyn(x, idx)
+}
+
+// CHECK-LABEL: literal_dyn_simd_extract
+// CHECK: extractelement <16 x i8> %x, i32 7
+#[no_mangle]
+unsafe extern "C" fn literal_dyn_simd_extract(x: i8x16) -> i8 {
+    simd_extract_dyn(x, 7)
+}
+
+// CHECK-LABEL: const_dyn_simd_extract
+// CHECK: extractelement <16 x i8> %x, i32 7
+#[no_mangle]
+unsafe extern "C" fn const_dyn_simd_extract(x: i8x16) -> i8 {
+    simd_extract_dyn(x, const { 3 + 4 })
+}
+
+// CHECK-LABEL: const_simd_extract
+// CHECK: extractelement <16 x i8> %x, i32 7
+#[no_mangle]
+unsafe extern "C" fn const_simd_extract(x: i8x16) -> i8 {
+    simd_extract(x, const { 3 + 4 })
+}
+
+// CHECK-LABEL: dyn_simd_insert
+// CHECK: insertelement <16 x i8> %x, i8 %e, i32 %idx
+#[no_mangle]
+unsafe extern "C" fn dyn_simd_insert(x: i8x16, e: i8, idx: u32) -> i8x16 {
+    simd_insert_dyn(x, idx, e)
+}
+
+// CHECK-LABEL: literal_dyn_simd_insert
+// CHECK: insertelement <16 x i8> %x, i8 %e, i32 7
+#[no_mangle]
+unsafe extern "C" fn literal_dyn_simd_insert(x: i8x16, e: i8) -> i8x16 {
+    simd_insert_dyn(x, 7, e)
+}
+
+// CHECK-LABEL: const_dyn_simd_insert
+// CHECK: insertelement <16 x i8> %x, i8 %e, i32 7
+#[no_mangle]
+unsafe extern "C" fn const_dyn_simd_insert(x: i8x16, e: i8) -> i8x16 {
+    simd_insert_dyn(x, const { 3 + 4 }, e)
+}
+
+// CHECK-LABEL: const_simd_insert
+// CHECK: insertelement <16 x i8> %x, i8 %e, i32 7
+#[no_mangle]
+unsafe extern "C" fn const_simd_insert(x: i8x16, e: i8) -> i8x16 {
+    simd_insert(x, const { 3 + 4 }, e)
+}

tests/ui/simd/intrinsic/generic-elements-pass.rs

@@ -1,8 +1,10 @@
 //@ run-pass
 
-#![feature(repr_simd, core_intrinsics)]
+#![feature(repr_simd, intrinsics, core_intrinsics)]
 
-use std::intrinsics::simd::{simd_extract, simd_insert, simd_shuffle};
+use std::intrinsics::simd::{
+    simd_extract, simd_extract_dyn, simd_insert, simd_insert_dyn, simd_shuffle,
+};
 
 #[repr(simd)]
 #[derive(Copy, Clone, Debug, PartialEq)]
@@ -70,6 +72,41 @@ fn main() {
         all_eq!(simd_extract(x8, 6), 86);
         all_eq!(simd_extract(x8, 7), 87);
     }
+    unsafe {
+        all_eq!(simd_insert_dyn(x2, 0, 100), i32x2([100, 21]));
+        all_eq!(simd_insert_dyn(x2, 1, 100), i32x2([20, 100]));
+
+        all_eq!(simd_insert_dyn(x4, 0, 100), i32x4([100, 41, 42, 43]));
+        all_eq!(simd_insert_dyn(x4, 1, 100), i32x4([40, 100, 42, 43]));
+        all_eq!(simd_insert_dyn(x4, 2, 100), i32x4([40, 41, 100, 43]));
+        all_eq!(simd_insert_dyn(x4, 3, 100), i32x4([40, 41, 42, 100]));
+
+        all_eq!(simd_insert_dyn(x8, 0, 100), i32x8([100, 81, 82, 83, 84, 85, 86, 87]));
+        all_eq!(simd_insert_dyn(x8, 1, 100), i32x8([80, 100, 82, 83, 84, 85, 86, 87]));
+        all_eq!(simd_insert_dyn(x8, 2, 100), i32x8([80, 81, 100, 83, 84, 85, 86, 87]));
+        all_eq!(simd_insert_dyn(x8, 3, 100), i32x8([80, 81, 82, 100, 84, 85, 86, 87]));
+        all_eq!(simd_insert_dyn(x8, 4, 100), i32x8([80, 81, 82, 83, 100, 85, 86, 87]));
+        all_eq!(simd_insert_dyn(x8, 5, 100), i32x8([80, 81, 82, 83, 84, 100, 86, 87]));
+        all_eq!(simd_insert_dyn(x8, 6, 100), i32x8([80, 81, 82, 83, 84, 85, 100, 87]));
+        all_eq!(simd_insert_dyn(x8, 7, 100), i32x8([80, 81, 82, 83, 84, 85, 86, 100]));
+
+        all_eq!(simd_extract_dyn(x2, 0), 20);
+        all_eq!(simd_extract_dyn(x2, 1), 21);
+
+        all_eq!(simd_extract_dyn(x4, 0), 40);
+        all_eq!(simd_extract_dyn(x4, 1), 41);
+        all_eq!(simd_extract_dyn(x4, 2), 42);
+        all_eq!(simd_extract_dyn(x4, 3), 43);
+
+        all_eq!(simd_extract_dyn(x8, 0), 80);
+        all_eq!(simd_extract_dyn(x8, 1), 81);
+        all_eq!(simd_extract_dyn(x8, 2), 82);
+        all_eq!(simd_extract_dyn(x8, 3), 83);
+        all_eq!(simd_extract_dyn(x8, 4), 84);
+        all_eq!(simd_extract_dyn(x8, 5), 85);
+        all_eq!(simd_extract_dyn(x8, 6), 86);
+        all_eq!(simd_extract_dyn(x8, 7), 87);
+    }
 
     let y2 = i32x2([120, 121]);
     let y4 = i32x4([140, 141, 142, 143]);