add simd_extract_dyn and simd_insert_dyn

Author: folkertdev

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1451,6 +1451,23 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
             bx.const_i32(idx as i32),
         ));
     }
+    if name == sym::simd_insert_dyn {
+        require!(
+            in_elem == arg_tys[2],
+            InvalidMonomorphization::InsertedType {
+                span,
+                name,
+                in_elem,
+                in_ty,
+                out_ty: arg_tys[2]
+            }
+        );
+        return Ok(bx.insert_element(
+            args[0].immediate(),
+            args[2].immediate(),
+            args[1].immediate(),
+        ));
+    }
     if name == sym::simd_extract {
         require!(
             ret_ty == in_elem,
@@ -1469,6 +1486,13 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
         }
         return Ok(bx.extract_element(args[0].immediate(), bx.const_i32(idx as i32)));
     }
+    if name == sym::simd_extract_dyn {
+        require!(
+            ret_ty == in_elem,
+            InvalidMonomorphization::ReturnType { span, name, in_elem, in_ty, ret_ty }
+        );
+        return Ok(bx.extract_element(args[0].immediate(), args[1].immediate()));
+    }
 
     if name == sym::simd_select {
         let m_elem_ty = in_elem;

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -682,8 +682,12 @@ pub 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

@@ -1874,6 +1874,7 @@ symbols! {
         simd_eq,
         simd_expose_provenance,
         simd_extract,
+        simd_extract_dyn,
         simd_fabs,
         simd_fcos,
         simd_fexp,
@@ -1894,6 +1895,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
 ///
@@ -16,20 +16,57 @@ pub unsafe fn simd_insert<T, U>(_x: T, _idx: u32, _val: U) -> T {
     unreachable!()
 }
 
-/// Extracts an element from a vector.
+/// Inserts an element into a vector, returning the updated vector.
 ///
 /// `T` must be a vector with element type `U`.
 ///
+/// If the index is `const`, using [`simd_insert`] instead may emit better assembly.
+///
 /// # Safety
 ///
 /// `idx` must be in-bounds of the vector.
+#[rustc_nounwind]
+#[inline]
+#[cfg_attr(not(bootstrap), rustc_intrinsic)]
+pub const unsafe fn simd_insert_dyn<T, U>(mut x: T, idx: u32, val: U) -> T {
+    // SAFETY: `idx` must be in-bounds
+    unsafe { (&mut x as *mut T as *mut U).add(idx as usize).write(val) }
+    x
+}
+
+/// Extracts an element from a vector.
+///
+/// `T` must be a vector with element type `U`, and `idx` must be const.
+///
+/// # Safety
+///
+/// `idx` must be const and in-bounds of the vector.
 #[rustc_intrinsic]
 #[rustc_intrinsic_must_be_overridden]
 #[rustc_nounwind]
 pub unsafe fn simd_extract<T, U>(_x: T, _idx: u32) -> U {
     unreachable!()
 }
 
+/// Extracts an element from a vector.
+///
+/// `T` must be a vector with element type `U`.
+///
+/// If the index is `const`, using [`simd_extract`] instead may emit better assembly.
+///
+/// # Safety
+///
+/// `idx` must be in-bounds of the vector.
+#[rustc_nounwind]
+#[inline]
+#[cfg_attr(not(bootstrap), rustc_intrinsic)]
+pub const unsafe fn simd_extract_dyn<T, U>(x: T, idx: u32) -> U {
+    // SAFETY: `idx` must be in-bounds
+    let e = unsafe { (&x as *const T as *const U).add(idx as usize).read() };
+    core::mem::forget(x);
+    e
+}
+
 /// Adds two simd vectors elementwise.
 ///
 /// `T` must be a vector of integer or floating point primitive types.

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

@@ -0,0 +1,88 @@
+//@compile-flags: -Copt-level=3 -Z merge-functions=disabled
+
+#![feature(core_intrinsics, s390x_target_feature, repr_simd)]
+#![no_std]
+#![crate_type = "lib"]
+#![allow(non_camel_case_types)]
+
+// Test that
+//
+// - `core::intrinsics::simd::simd_extract_dyn` is lowered to an `extractelement`
+// - `core::intrinsics::simd::simd_insert_dyn` is lowered to an `insertelement`
+//
+// On many platforms (e.g. x86_64) this will ultimately lower to a pointer load, but some platforms
+// (e.g. s390x, powerpc) have dedicated instructions for these operations.
+
+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
+// CHECK-NEXT: ret
+#[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, i64 7
+// CHECK-NEXT: ret
+#[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, i64 7
+// CHECK-NEXT: ret
+#[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, i64 7
+// CHECK-NEXT: ret
+#[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
+// CHECK-NEXT: ret
+#[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, i64 7
+// CHECK-NEXT: ret
+#[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, i64 7
+// CHECK-NEXT: ret
+#[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, i64 7
+// CHECK-NEXT: ret
+#[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

@@ -19,9 +19,14 @@ struct i32x8([i32; 8]);
 #[rustc_intrinsic]
 unsafe fn simd_insert<T, E>(x: T, idx: u32, y: E) -> T;
 
+#[rustc_intrinsic]
+unsafe fn simd_insert_dyn<T, E>(x: T, idx: u32, y: E) -> T;
+
 #[rustc_intrinsic]
 unsafe fn simd_extract<T, E>(x: T, idx: u32) -> E;
 
+#[rustc_intrinsic]
+unsafe fn simd_extract_dyn<T, E>(x: T, idx: u32) -> E;
 
 #[rustc_intrinsic]
 unsafe fn simd_shuffle<T, I, U>(x: T, y: T, idx: I) -> U;
@@ -79,6 +84,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]);