From 59c55339af73b5345ec28d7d830c7bd3feee45b3 Mon Sep 17 00:00:00 2001 From: Folkert de Vries Date: Sat, 22 Feb 2025 22:01:18 +0100 Subject: [PATCH] add `simd_insert_dyn` and `simd_extract_dyn` --- compiler/rustc_codegen_llvm/src/intrinsic.rs | 69 +++++++++-------- .../rustc_hir_analysis/src/check/intrinsic.rs | 8 +- compiler/rustc_span/src/symbol.rs | 2 + library/core/src/intrinsics/simd.rs | 39 +++++++++- tests/codegen/simd/extract-insert-dyn.rs | 75 +++++++++++++++++++ .../simd/intrinsic/generic-elements-pass.rs | 41 +++++++++- 6 files changed, 197 insertions(+), 37 deletions(-) create mode 100644 tests/codegen/simd/extract-insert-dyn.rs diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 660fc7ec4c404..a14600f925fb5 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -1417,7 +1417,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 { @@ -1428,40 +1428,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 { diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs index 42d785c8dd0fe..21d40b01e9a1c 100644 --- a/compiler/rustc_hir_analysis/src/check/intrinsic.rs +++ b/compiler/rustc_hir_analysis/src/check/intrinsic.rs @@ -674,8 +674,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 diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index 8a8bec35d8194..0b6e3f64ad10f 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1885,6 +1885,7 @@ symbols! { simd_eq, simd_expose_provenance, simd_extract, + simd_extract_dyn, simd_fabs, simd_fcos, simd_fexp, @@ -1903,6 +1904,7 @@ symbols! { simd_ge, simd_gt, simd_insert, + simd_insert_dyn, simd_le, simd_lt, simd_masked_load, diff --git a/library/core/src/intrinsics/simd.rs b/library/core/src/intrinsics/simd.rs index ae6e1a779ed58..9ac6ee8553558 100644 --- a/library/core/src/intrinsics/simd.rs +++ b/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(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(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(mut x: T, idx: u32, val: U) -> T { + // SAFETY: `idx` must be in-bounds + unsafe { (&raw mut x).cast::().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(x: T, idx: u32) -> U { + // SAFETY: `idx` must be in-bounds + unsafe { (&raw const x).cast::().add(idx as usize).read() } +} + /// Adds two simd vectors elementwise. /// /// `T` must be a vector of integers or floats. diff --git a/tests/codegen/simd/extract-insert-dyn.rs b/tests/codegen/simd/extract-insert-dyn.rs new file mode 100644 index 0000000000000..584e2c7887adc --- /dev/null +++ b/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) +} diff --git a/tests/ui/simd/intrinsic/generic-elements-pass.rs b/tests/ui/simd/intrinsic/generic-elements-pass.rs index 4dc2e4d5a80a9..e4d47cdb38184 100644 --- a/tests/ui/simd/intrinsic/generic-elements-pass.rs +++ b/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]);