From 194baa820d36926cfa9128211bbd61866d49d501 Mon Sep 17 00:00:00 2001 From: Ralf Jung Date: Tue, 6 Aug 2024 12:25:37 +0200 Subject: [PATCH] simd_shuffle intrinsic: allow argument to be passed as vector (not just as array) --- .../src/intrinsics/simd.rs | 10 +++++++ .../rustc_codegen_gcc/src/intrinsic/simd.rs | 19 ++++++++----- compiler/rustc_codegen_llvm/src/intrinsic.rs | 19 ++++++++----- library/core/src/intrinsics/simd.rs | 2 +- tests/ui/simd/shuffle.rs | 27 +++++++++++++++++-- 5 files changed, 60 insertions(+), 17 deletions(-) diff --git a/compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs b/compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs index ca910dccb0d..604a88393fd 100644 --- a/compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs +++ b/compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs @@ -191,6 +191,14 @@ pub(super) fn codegen_simd_intrinsic_call<'tcx>( }) .try_into() .unwrap(), + _ if idx_ty.is_simd() + && matches!( + idx_ty.simd_size_and_type(fx.tcx).1.kind(), + ty::Uint(ty::UintTy::U32) + ) => + { + idx_ty.simd_size_and_type(fx.tcx).0.try_into().unwrap() + } _ => { fx.tcx.dcx().span_err( span, @@ -213,6 +221,8 @@ pub(super) fn codegen_simd_intrinsic_call<'tcx>( let total_len = lane_count * 2; + // FIXME: this is a terrible abstraction-breaking hack. + // Find a way to reuse `immediate_const_vector` from `codegen_ssa` instead. let indexes = { use rustc_middle::mir::interpret::*; let idx_const = match &idx.node { diff --git a/compiler/rustc_codegen_gcc/src/intrinsic/simd.rs b/compiler/rustc_codegen_gcc/src/intrinsic/simd.rs index 8da1df3be15..96a833ccaf2 100644 --- a/compiler/rustc_codegen_gcc/src/intrinsic/simd.rs +++ b/compiler/rustc_codegen_gcc/src/intrinsic/simd.rs @@ -353,19 +353,24 @@ pub fn generic_simd_intrinsic<'a, 'gcc, 'tcx>( } if name == sym::simd_shuffle { - // Make sure this is actually an array, since typeck only checks the length-suffixed + // Make sure this is actually an array or SIMD vector, since typeck only checks the length-suffixed // version of this intrinsic. - let n: u64 = match *args[2].layout.ty.kind() { + let idx_ty = args[2].layout.ty; + let n: u64 = match idx_ty.kind() { ty::Array(ty, len) if matches!(*ty.kind(), ty::Uint(ty::UintTy::U32)) => { len.try_eval_target_usize(bx.cx.tcx, ty::ParamEnv::reveal_all()).unwrap_or_else( || span_bug!(span, "could not evaluate shuffle index array length"), ) } - _ => return_error!(InvalidMonomorphization::SimdShuffle { - span, - name, - ty: args[2].layout.ty - }), + _ if idx_ty.is_simd() + && matches!( + idx_ty.simd_size_and_type(bx.cx.tcx).1.kind(), + ty::Uint(ty::UintTy::U32) + ) => + { + idx_ty.simd_size_and_type(bx.cx.tcx).0 + } + _ => return_error!(InvalidMonomorphization::SimdShuffle { span, name, ty: idx_ty }), }; require_simd!(ret_ty, InvalidMonomorphization::SimdReturn { span, name, ty: ret_ty }); diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index f5558723d11..5d32ef0d9d6 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -1279,19 +1279,24 @@ fn generic_simd_intrinsic<'ll, 'tcx>( } if name == sym::simd_shuffle { - // Make sure this is actually an array, since typeck only checks the length-suffixed + // Make sure this is actually an array or SIMD vector, since typeck only checks the length-suffixed // version of this intrinsic. - let n: u64 = match args[2].layout.ty.kind() { + let idx_ty = args[2].layout.ty; + let n: u64 = match idx_ty.kind() { ty::Array(ty, len) if matches!(ty.kind(), ty::Uint(ty::UintTy::U32)) => { len.try_eval_target_usize(bx.cx.tcx, ty::ParamEnv::reveal_all()).unwrap_or_else( || span_bug!(span, "could not evaluate shuffle index array length"), ) } - _ => return_error!(InvalidMonomorphization::SimdShuffle { - span, - name, - ty: args[2].layout.ty - }), + _ if idx_ty.is_simd() + && matches!( + idx_ty.simd_size_and_type(bx.cx.tcx).1.kind(), + ty::Uint(ty::UintTy::U32) + ) => + { + idx_ty.simd_size_and_type(bx.cx.tcx).0 + } + _ => return_error!(InvalidMonomorphization::SimdShuffle { span, name, ty: idx_ty }), }; let (out_len, out_ty) = require_simd!(ret_ty, SimdReturn); diff --git a/library/core/src/intrinsics/simd.rs b/library/core/src/intrinsics/simd.rs index 221724d7b4a..59828198099 100644 --- a/library/core/src/intrinsics/simd.rs +++ b/library/core/src/intrinsics/simd.rs @@ -232,7 +232,7 @@ extern "rust-intrinsic" { /// /// `T` must be a vector. /// - /// `U` must be a **const** array of `i32`s. This means it must either refer to a named + /// `U` must be a **const** array or vector of `u32`s. This means it must either refer to a named /// const or be given as an inline const expression (`const { ... }`). /// /// `V` must be a vector with the same element type as `T` and the same length as `U`. diff --git a/tests/ui/simd/shuffle.rs b/tests/ui/simd/shuffle.rs index 09926d95557..dc0d688284e 100644 --- a/tests/ui/simd/shuffle.rs +++ b/tests/ui/simd/shuffle.rs @@ -6,15 +6,20 @@ #![allow(incomplete_features)] #![feature(adt_const_params)] +use std::marker::ConstParamTy; + extern "rust-intrinsic" { fn simd_shuffle(a: T, b: T, i: I) -> U; } -#[derive(Copy, Clone)] +#[derive(Copy, Clone, ConstParamTy, PartialEq, Eq)] #[repr(simd)] struct Simd([T; N]); -pub unsafe fn __shuffle_vector16(x: T, y: T) -> U { +unsafe fn __shuffle_vector16(x: T, y: T) -> U { + simd_shuffle(x, y, IDX) +} +unsafe fn __shuffle_vector16_v2, T, U>(x: T, y: T) -> U { simd_shuffle(x, y, IDX) } @@ -30,6 +35,17 @@ fn main() { let y: Simd = simd_shuffle(a, b, I2); assert_eq!(y.0, [1, 5]); } + // Test that we can also use a SIMD vector instead of a normal array for the shuffle. + const I1_SIMD: Simd = Simd([0, 2, 4, 6]); + const I2_SIMD: Simd = Simd([1, 5]); + unsafe { + let x: Simd = simd_shuffle(a, b, I1_SIMD); + assert_eq!(x.0, [0, 2, 4, 6]); + + let y: Simd = simd_shuffle(a, b, I2_SIMD); + assert_eq!(y.0, [1, 5]); + } + // Test that an indirection (via an unnamed constant) // through a const generic parameter also works. // See https://github.com/rust-lang/rust/issues/113500 for details. @@ -42,4 +58,11 @@ fn main() { Simd, >(a, b); } + unsafe { + __shuffle_vector16_v2::< + { Simd([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]) }, + Simd, + Simd, + >(a, b); + } }