simd_shuffle intrinsic: allow argument to be passed as vector (not just as array)

This commit is contained in:
Ralf Jung 2024-08-06 12:25:37 +02:00
parent e9e27ab0cf
commit 194baa820d
5 changed files with 60 additions and 17 deletions

View File

@ -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 {

View File

@ -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 });

View File

@ -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);

View File

@ -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`.

View File

@ -6,15 +6,20 @@
#![allow(incomplete_features)]
#![feature(adt_const_params)]
use std::marker::ConstParamTy;
extern "rust-intrinsic" {
fn simd_shuffle<T, I, U>(a: T, b: T, i: I) -> U;
}
#[derive(Copy, Clone)]
#[derive(Copy, Clone, ConstParamTy, PartialEq, Eq)]
#[repr(simd)]
struct Simd<T, const N: usize>([T; N]);
pub unsafe fn __shuffle_vector16<const IDX: [u32; 16], T, U>(x: T, y: T) -> U {
unsafe fn __shuffle_vector16<const IDX: [u32; 16], T, U>(x: T, y: T) -> U {
simd_shuffle(x, y, IDX)
}
unsafe fn __shuffle_vector16_v2<const IDX: Simd<u32, 16>, T, U>(x: T, y: T) -> U {
simd_shuffle(x, y, IDX)
}
@ -30,6 +35,17 @@ fn main() {
let y: Simd<u8, 2> = 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<u32, 4> = Simd([0, 2, 4, 6]);
const I2_SIMD: Simd<u32, 2> = Simd([1, 5]);
unsafe {
let x: Simd<u8, 4> = simd_shuffle(a, b, I1_SIMD);
assert_eq!(x.0, [0, 2, 4, 6]);
let y: Simd<u8, 2> = 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<u8, 16>,
>(a, b);
}
unsafe {
__shuffle_vector16_v2::<
{ Simd([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]) },
Simd<u8, 16>,
Simd<u8, 16>,
>(a, b);
}
}