Skip to content

Commit

Permalink
Rollup merge of rust-lang#130268 - RalfJung:simd-shuffle-idx-vector, …
Browse files Browse the repository at this point in the history 
…r=compiler-errors

simd_shuffle: require index argument to be a vector

Remove some codegen hacks by forcing the SIMD shuffle `index` argument to be a vector, which means (thanks to rust-lang#128537) that it will automatically be passed as an immediate in LLVM. The only special-casing we still have is for the extra sanity-checks we add that ensure that the indices are all in-bounds. (And the GCC backend needs to do a bunch of work since the Rust intrinsic is modeled after what LLVM expects, which seems to be quite different from what GCC expects.)

Fixes rust-lang#128738, see that issue for more context.
  • Loading branch information
@Zalathar
Zalathar authored Sep 14, 2024
2 parents 2b9a0bd + 7a1b5e5 commit c0aa9f9
Show file tree
Hide file tree
Showing 24 changed files with 224 additions and 301 deletions.
40 changes: 13 additions & 27 deletions compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -180,34 +180,20 @@ pub(super) fn codegen_simd_intrinsic_call<'tcx>(
return;
}

// Make sure this is actually an array, since typeck only checks the length-suffixed
// version of this intrinsic.
// Make sure this is actually a SIMD vector.
let idx_ty = fx.monomorphize(idx.node.ty(fx.mir, fx.tcx));
let n: u16 = match idx_ty.kind() {
ty::Array(ty, len) if matches!(ty.kind(), ty::Uint(ty::UintTy::U32)) => len
.try_eval_target_usize(fx.tcx, ty::ParamEnv::reveal_all())
.unwrap_or_else(|| {
span_bug!(span, "could not evaluate shuffle index array length")
})
.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,
format!("simd_shuffle index must be an array of `u32`, got `{}`", idx_ty),
);
// Prevent verifier error
fx.bcx.ins().trap(TrapCode::UnreachableCodeReached);
return;
}
let n: u16 = 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()
} else {
fx.tcx.dcx().span_err(
span,
format!("simd_shuffle index must be a SIMD vector of `u32`, got `{}`", idx_ty),
);
// Prevent verifier error
fx.bcx.ins().trap(TrapCode::UnreachableCodeReached);
return;
};

assert_eq!(x.layout(), y.layout());
Expand Down
39 changes: 12 additions & 27 deletions compiler/rustc_codegen_gcc/src/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1939,33 +1939,18 @@ impl<'a, 'gcc, 'tcx> Builder<'a, 'gcc, 'tcx> {
self.int_type
};

let mut mask_elements = if let Some(vector_type) = mask.get_type().dyncast_vector() {
let mask_num_units = vector_type.get_num_units();
let mut mask_elements = vec![];
for i in 0..mask_num_units {
let index = self.context.new_rvalue_from_long(self.cx.type_u32(), i as _);
mask_elements.push(self.context.new_cast(
self.location,
self.extract_element(mask, index).to_rvalue(),
mask_element_type,
));
}
mask_elements
} else {
let struct_type = mask.get_type().is_struct().expect("mask should be of struct type");
let mask_num_units = struct_type.get_field_count();
let mut mask_elements = vec![];
for i in 0..mask_num_units {
let field = struct_type.get_field(i as i32);
mask_elements.push(self.context.new_cast(
self.location,
mask.access_field(self.location, field).to_rvalue(),
mask_element_type,
));
}
mask_elements
};
let mask_num_units = mask_elements.len();
let vector_type =
mask.get_type().dyncast_vector().expect("simd_shuffle mask should be of vector type");
let mask_num_units = vector_type.get_num_units();
let mut mask_elements = vec![];
for i in 0..mask_num_units {
let index = self.context.new_rvalue_from_long(self.cx.type_u32(), i as _);
mask_elements.push(self.context.new_cast(
self.location,
self.extract_element(mask, index).to_rvalue(),
mask_element_type,
));
}

// NOTE: the mask needs to be the same length as the input vectors, so add the missing
// elements in the mask if needed.
Expand Down
25 changes: 7 additions & 18 deletions compiler/rustc_codegen_gcc/src/intrinsic/simd.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ use rustc_codegen_ssa::traits::{BaseTypeMethods, BuilderMethods};
#[cfg(feature = "master")]
use rustc_hir as hir;
use rustc_middle::mir::BinOp;
use rustc_middle::span_bug;
use rustc_middle::ty::layout::HasTyCtxt;
use rustc_middle::ty::{self, Ty};
use rustc_span::{sym, Span, Symbol};
Expand Down Expand Up @@ -353,24 +352,14 @@ pub fn generic_simd_intrinsic<'a, 'gcc, 'tcx>(
}

if name == sym::simd_shuffle {
// Make sure this is actually an array or SIMD vector, since typeck only checks the length-suffixed
// version of this intrinsic.
// Make sure this is actually a SIMD vector.
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"),
)
}
_ 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 n: u64 = 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
} else {
return_error!(InvalidMonomorphization::SimdShuffle { span, name, ty: idx_ty })
};
require_simd!(ret_ty, InvalidMonomorphization::SimdReturn { span, name, ty: ret_ty });

Expand Down
72 changes: 24 additions & 48 deletions compiler/rustc_codegen_llvm/src/intrinsic.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1290,24 +1290,14 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
}

if name == sym::simd_shuffle {
// Make sure this is actually an array or SIMD vector, since typeck only checks the length-suffixed
// version of this intrinsic.
// Make sure this is actually a SIMD vector.
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"),
)
}
_ 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 n: u64 = 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
} else {
return_error!(InvalidMonomorphization::SimdShuffle { span, name, ty: idx_ty })
};

let (out_len, out_ty) = require_simd!(ret_ty, SimdReturn);
Expand All @@ -1322,38 +1312,24 @@ fn generic_simd_intrinsic<'ll, 'tcx>(

let total_len = u128::from(in_len) * 2;

let vector = args[2].immediate();

let indices: Option<Vec<_>> = (0..n)
.map(|i| {
let arg_idx = i;
let val = bx.const_get_elt(vector, i as u64);
match bx.const_to_opt_u128(val, true) {
None => {
bug!("typeck should have already ensured that these are const")
}
Some(idx) if idx >= total_len => {
bx.sess().dcx().emit_err(InvalidMonomorphization::SimdIndexOutOfBounds {
span,
name,
arg_idx,
total_len,
});
None
}
Some(idx) => Some(bx.const_i32(idx as i32)),
}
})
.collect();
let Some(indices) = indices else {
return Ok(bx.const_null(llret_ty));
};
// Check that the indices are in-bounds.
let indices = args[2].immediate();
for i in 0..n {
let val = bx.const_get_elt(indices, i as u64);
let idx = bx
.const_to_opt_u128(val, true)
.unwrap_or_else(|| bug!("typeck should have already ensured that these are const"));
if idx >= total_len {
bx.sess().dcx().emit_err(InvalidMonomorphization::SimdIndexOutOfBounds {
span,
name,
arg_idx: i,
total_len,
});
}
}

return Ok(bx.shuffle_vector(
args[0].immediate(),
args[1].immediate(),
bx.const_vector(&indices),
));
return Ok(bx.shuffle_vector(args[0].immediate(), args[1].immediate(), indices));
}

if name == sym::simd_insert {
Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_codegen_ssa/messages.ftl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ codegen_ssa_invalid_monomorphization_simd_return = invalid monomorphization of `
codegen_ssa_invalid_monomorphization_simd_second = invalid monomorphization of `{$name}` intrinsic: expected SIMD second type, found non-SIMD `{$ty}`
codegen_ssa_invalid_monomorphization_simd_shuffle = invalid monomorphization of `{$name}` intrinsic: simd_shuffle index must be an array of `u32`, got `{$ty}`
codegen_ssa_invalid_monomorphization_simd_shuffle = invalid monomorphization of `{$name}` intrinsic: simd_shuffle index must be a SIMD vector of `u32`, got `{$ty}`
codegen_ssa_invalid_monomorphization_simd_third = invalid monomorphization of `{$name}` intrinsic: expected SIMD third type, found non-SIMD `{$ty}`
Expand Down
28 changes: 2 additions & 26 deletions compiler/rustc_codegen_ssa/src/mir/block.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -915,32 +915,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
}
};

let args: Vec<_> = args
.iter()
.enumerate()
.map(|(i, arg)| {
// The indices passed to simd_shuffle in the
// third argument must be constant. This is
// checked by the type-checker.
if i == 2 && intrinsic.name == sym::simd_shuffle {
// FIXME: the simd_shuffle argument is actually an array,
// not a vector, so we need this special hack to make sure
// it is passed as an immediate. We should pass the
// shuffle indices as a vector instead to avoid this hack.
if let mir::Operand::Constant(constant) = &arg.node {
let (llval, ty) = self.immediate_const_vector(bx, constant);
return OperandRef {
val: Immediate(llval),
layout: bx.layout_of(ty),
};
} else {
span_bug!(span, "shuffle indices must be constant");
}
}

self.codegen_operand(bx, &arg.node)
})
.collect();
let args: Vec<_> =
args.iter().map(|arg| self.codegen_operand(bx, &arg.node)).collect();

if matches!(intrinsic, ty::IntrinsicDef { name: sym::caller_location, .. }) {
let location = self
Expand Down
35 changes: 11 additions & 24 deletions compiler/rustc_codegen_ssa/src/mir/constant.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use rustc_middle::mir::interpret::ErrorHandled;
use rustc_middle::ty::layout::HasTyCtxt;
use rustc_middle::ty::{self, Ty, ValTree};
use rustc_middle::ty::{self, Ty};
use rustc_middle::{bug, mir, span_bug};
use rustc_target::abi::Abi;

Expand Down Expand Up @@ -66,35 +66,22 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
constant: &mir::ConstOperand<'tcx>,
) -> (Bx::Value, Ty<'tcx>) {
let ty = self.monomorphize(constant.ty());
let ty_is_simd = ty.is_simd();
// FIXME: ideally we'd assert that this is a SIMD type, but simd_shuffle
// in its current form relies on a regular array being passed as an
// immediate argument. This hack can be removed once that is fixed.
let field_ty = if ty_is_simd {
ty.simd_size_and_type(bx.tcx()).1
} else {
ty.builtin_index().unwrap()
};
assert!(ty.is_simd());
let field_ty = ty.simd_size_and_type(bx.tcx()).1;

let val = self
.eval_unevaluated_mir_constant_to_valtree(constant)
.ok()
.map(|x| x.ok())
.flatten()
.map(|val| {
// Depending on whether this is a SIMD type with an array field
// or a type with many fields (one for each elements), the valtree
// is either a single branch with N children, or a root node
// with exactly one child which then in turn has many children.
// So we look at the first child to determine whether it is a
// leaf or whether we have to go one more layer down.
let branch_or_leaf = val.unwrap_branch();
let first = branch_or_leaf.get(0).unwrap();
let field_iter = match first {
ValTree::Branch(_) => first.unwrap_branch().iter(),
ValTree::Leaf(_) => branch_or_leaf.iter(),
};
let values: Vec<_> = field_iter
// A SIMD type has a single field, which is an array.
let fields = val.unwrap_branch();
assert_eq!(fields.len(), 1);
let array = fields[0].unwrap_branch();
// Iterate over the array elements to obtain the values in the vector.
let values: Vec<_> = array
.iter()
.map(|field| {
if let Some(prim) = field.try_to_scalar() {
let layout = bx.layout_of(field_ty);
Expand All @@ -107,7 +94,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
}
})
.collect();
if ty_is_simd { bx.const_vector(&values) } else { bx.const_struct(&values, false) }
bx.const_vector(&values)
})
.unwrap_or_else(|| {
bx.tcx().dcx().emit_err(errors::ShuffleIndicesEvaluation { span: constant.span });
Expand Down
2 changes: 1 addition & 1 deletion library/core/src/intrinsics/simd.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -232,7 +232,7 @@ extern "rust-intrinsic" {
///
/// `T` must be a vector.
///
/// `U` must be a **const** array or vector of `u32`s. This means it must either refer to a named
/// `U` must be a **const** 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`.
Expand Down
16 changes: 12 additions & 4 deletions library/portable-simd/crates/core_simd/src/swizzle.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ pub trait Swizzle<const N: usize> {
LaneCount<N>: SupportedLaneCount,
LaneCount<M>: SupportedLaneCount,
{
// Safety: `vector` is a vector, and the index is a const array of u32.
// Safety: `vector` is a vector, and the index is a const vector of u32.
unsafe {
core::intrinsics::simd::simd_shuffle(
vector,
Expand All @@ -103,7 +103,11 @@ pub trait Swizzle<const N: usize> {
output[i] = index as u32;
i += 1;
}
output

// The index list needs to be returned as a vector.
#[repr(simd)]
struct SimdShuffleIdx<const LEN: usize>([u32; LEN]);
SimdShuffleIdx(output)
},
)
}
Expand All @@ -121,7 +125,7 @@ pub trait Swizzle<const N: usize> {
LaneCount<N>: SupportedLaneCount,
LaneCount<M>: SupportedLaneCount,
{
// Safety: `first` and `second` are vectors, and the index is a const array of u32.
// Safety: `first` and `second` are vectors, and the index is a const vector of u32.
unsafe {
core::intrinsics::simd::simd_shuffle(
first,
Expand All @@ -139,7 +143,11 @@ pub trait Swizzle<const N: usize> {
output[i] = index as u32;
i += 1;
}
output

// The index list needs to be returned as a vector.
#[repr(simd)]
struct SimdShuffleIdx<const LEN: usize>([u32; LEN]);
SimdShuffleIdx(output)
},
)
}
Expand Down
8 changes: 1 addition & 7 deletions src/tools/miri/src/intrinsics/simd.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -664,15 +664,9 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
let [left, right, index] = check_arg_count(args)?;
let (left, left_len) = this.project_to_simd(left)?;
let (right, right_len) = this.project_to_simd(right)?;
let (index, index_len) = this.project_to_simd(index)?;
let (dest, dest_len) = this.project_to_simd(dest)?;

// `index` is an array or a SIMD type
let (index, index_len) = match index.layout.ty.kind() {
// FIXME: remove this once `index` must always be a SIMD vector.
ty::Array(..) => (index.clone(), index.len(this)?),
_ => this.project_to_simd(index)?,
};

assert_eq!(left_len, right_len);
assert_eq!(index_len, dest_len);

Expand Down
Loading

0 comments on commit c0aa9f9

Please sign in to comment.