add simd_splat intrinsic

Author: folkertdev

compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs

@@ -348,6 +348,31 @@ pub(super) fn codegen_simd_intrinsic_call<'tcx>(
             ret.write_cvalue(fx, ret_lane);
         }
 
+        sym::simd_splat => {
+            intrinsic_args!(fx, args => (value); intrinsic);
+
+            if !ret.layout().ty.is_simd() {
+                report_simd_type_validation_error(fx, intrinsic, span, ret.layout().ty);
+                return;
+            }
+            let (lane_count, lane_ty) = ret.layout().ty.simd_size_and_type(fx.tcx);
+
+            if value.layout().ty != lane_ty {
+                fx.tcx.dcx().span_fatal(
+                    span,
+                    format!(
+                        "[simd_splat] expected element type {lane_ty:?}, got {got:?}",
+                        got = value.layout().ty
+                    ),
+                );
+            }
+
+            for i in 0..lane_count {
+                let ret_lane = ret.place_lane(fx, i.into());
+                ret_lane.write_cvalue(fx, value);
+            }
+        }
+
         sym::simd_neg
         | sym::simd_bswap
         | sym::simd_bitreverse

compiler/rustc_codegen_gcc/src/intrinsic/simd.rs

@@ -121,6 +121,25 @@ pub fn generic_simd_intrinsic<'a, 'gcc, 'tcx>(
         return Ok(bx.vector_select(vector_mask, arg1, args[2].immediate()));
     }
 
+    if name == sym::simd_splat {
+        require_simd!(ret_ty, InvalidMonomorphization::SimdReturn { span, name, ty: ret_ty });
+        let (out_len, out_ty) = ret_ty.simd_size_and_type(bx.tcx());
+
+        require!(
+            args[0].layout.ty == out_ty,
+            InvalidMonomorphization::ExpectedVectorElementType {
+                span,
+                name,
+                expected_element: out_ty,
+                vector_type: ret_ty,
+            }
+        );
+
+        let elem = args[0].immediate();
+        let elements = vec![elem; out_len as usize];
+        return Ok(bx.context.new_rvalue_from_vector(bx.location, llret_ty, &elements));
+    }
+
     // every intrinsic below takes a SIMD vector as its first argument
     require_simd!(
         args[0].layout.ty,

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1588,6 +1588,32 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
         return Ok(bx.select(m_i1s, args[1].immediate(), args[2].immediate()));
     }
 
+    if name == sym::simd_splat {
+        let (out_len, out_ty) = require_simd!(ret_ty, SimdReturn);
+
+        require!(
+            args[0].layout.ty == out_ty,
+            InvalidMonomorphization::ExpectedVectorElementType {
+                span,
+                name,
+                expected_element: out_ty,
+                vector_type: ret_ty,
+            }
+        );
+
+        // `insertelement <N x elem> poison, elem %x, i32 0`
+        let poison_vec = bx.const_poison(llret_ty);
+        let idx0 = bx.const_i32(0);
+        let v0 = bx.insert_element(poison_vec, args[0].immediate(), idx0);
+
+        // `shufflevector <N x elem> v0, <N x elem> poison, <N x i32> zeroinitializer`
+        // The masks is all zeros, so this splats lane 0 (which has our element in it).
+        let mask = bx.const_vector(&vec![bx.const_i32(0); out_len as usize]);
+        let splat = bx.shuffle_vector(v0, poison_vec, mask);
+
+        return Ok(splat);
+    }
+
     // every intrinsic below takes a SIMD vector as its first argument
     let (in_len, in_elem) = require_simd!(args[0].layout.ty, SimdInput);
     let in_ty = args[0].layout.ty;

compiler/rustc_const_eval/src/interpret/intrinsics/simd.rs

@@ -60,6 +60,15 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 }
                 self.copy_op(&self.project_index(&input, index)?, &dest)?;
             }
+            sym::simd_splat => {
+                let elem = &args[0];
+                let (dest, dest_len) = self.project_to_simd(&dest)?;
+
+                for i in 0..dest_len {
+                    let place = self.project_index(&dest, i)?;
+                    self.copy_op(elem, &place)?;
+                }
+            }
             sym::simd_neg
             | sym::simd_fabs
             | sym::simd_ceil

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -746,6 +746,7 @@ pub(crate) fn check_intrinsic_type(
         sym::simd_extract | sym::simd_extract_dyn => {
             (2, 0, vec![param(0), tcx.types.u32], param(1))
         }
+        sym::simd_splat => (2, 0, vec![param(1)], param(0)),
         sym::simd_cast
         | sym::simd_as
         | sym::simd_cast_ptr

compiler/rustc_span/src/symbol.rs

@@ -2132,6 +2132,7 @@ symbols! {
         simd_shr,
         simd_shuffle,
         simd_shuffle_const_generic,
+        simd_splat,
         simd_sub,
         simd_trunc,
         simd_with_exposed_provenance,

library/core/src/intrinsics/simd.rs

@@ -59,6 +59,10 @@ pub unsafe fn simd_extract_dyn<T, U>(x: T, idx: u32) -> U {
     unsafe { (&raw const x).cast::<U>().add(idx as usize).read() }
 }
 
+#[rustc_nounwind]
+#[rustc_intrinsic]
+pub const unsafe fn simd_splat<T, U>(value: U) -> T;
+
 /// Adds two simd vectors elementwise.
 ///
 /// `T` must be a vector of integers or floats.

tests/codegen-llvm/simd/splat.rs

@@ -0,0 +1,29 @@
+#![crate_type = "lib"]
+#![no_std]
+#![feature(repr_simd, core_intrinsics)]
+use core::intrinsics::simd::simd_splat;
+use core::{mem, ptr};
+
+#[path = "../../auxiliary/minisimd.rs"]
+mod minisimd;
+use minisimd::*;
+
+#[no_mangle]
+unsafe extern "C" fn int(x: u16) -> u16x2 {
+    // CHECK-LABEL: int
+    // CHECK: start:
+    // CHECK-NEXT: %0 = insertelement <2 x i16> poison, i16 %x, i64 0
+    // CHECK-NEXT: %1 = shufflevector <2 x i16> %0, <2 x i16> poison, <2 x i32> zeroinitializer
+    // CHECK-NEXT: ret
+    simd_splat(x)
+}
+
+#[no_mangle]
+unsafe extern "C" fn float(x: f32) -> f32x4 {
+    // CHECK-LABEL: float
+    // CHECK: start:
+    // CHECK-NEXT: %0 = insertelement <4 x float> poison, float %x, i64 0
+    // CHECK-NEXT: %1 = shufflevector <4 x float> %0, <4 x float> poison, <4 x i32> zeroinitializer
+    // CHECK-NEXT: ret
+    simd_splat(x)
+}

tests/ui/simd/intrinsic/splat.rs

@@ -0,0 +1,52 @@
+//@ run-pass
+#![feature(repr_simd, core_intrinsics)]
+
+#[path = "../../../auxiliary/minisimd.rs"]
+mod minisimd;
+use minisimd::*;
+
+use std::intrinsics::simd::simd_splat;
+
+fn main() {
+    unsafe {
+        let x: Simd<u32, 1> = simd_splat(123u32);
+        let y: Simd<u32, 1> = const { simd_splat(123u32) };
+        assert_eq!(x.into_array(), [123; 1]);
+        assert_eq!(x.into_array(), y.into_array());
+
+        let x: u16x2 = simd_splat(42u16);
+        let y: u16x2 = const { simd_splat(42u16) };
+        assert_eq!(x.into_array(), [42; 2]);
+        assert_eq!(x.into_array(), y.into_array());
+
+        let x: u128x4 = simd_splat(42u128);
+        let y: u128x4 = const { simd_splat(42u128) };
+        assert_eq!(x.into_array(), [42; 4]);
+        assert_eq!(x.into_array(), y.into_array());
+
+        let x: i32x4 = simd_splat(-7i32);
+        let y: i32x4 = const { simd_splat(-7i32) };
+        assert_eq!(x.into_array(), [-7; 4]);
+        assert_eq!(x.into_array(), y.into_array());
+
+        let x: f32x4 = simd_splat(42.0f32);
+        let y: f32x4 = const { simd_splat(42.0f32) };
+        assert_eq!(x.into_array(), [42.0; 4]);
+        assert_eq!(x.into_array(), y.into_array());
+
+        let x: f64x2 = simd_splat(42.0f64);
+        let y: f64x2 = const { simd_splat(42.0f64) };
+        assert_eq!(x.into_array(), [42.0; 2]);
+        assert_eq!(x.into_array(), y.into_array());
+
+        static ZERO: u8 = 0u8;
+        let x: Simd<*const u8, 2> = simd_splat(&raw const ZERO);
+        assert_eq!(x.into_array(), [&raw const ZERO; 2]);
+
+        // FIXME: this hits "could not evaluate shuffle_indices at compile time",
+        // emitted in `immediate_const_vector`. const-eval should be able to handle
+        // this though I think? `const { [&raw const ZERO; 2] }` appears to work.
+        // let y: Simd<*const u8, 2> = const { simd_splat(&raw const ZERO) };
+        // assert_eq!(x.into_array(), y.into_array());
+    }
+}