Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

simd_shuffle: require index argument to be a vector #130268

Merged
merged 1 commit into from
Sep 14, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 13 additions & 27 deletions compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs
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
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
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
80 changes: 28 additions & 52 deletions compiler/rustc_codegen_llvm/src/intrinsic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -573,6 +573,8 @@ impl<'ll, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'_, 'll, 'tcx> {
span,
) {
Ok(llval) => llval,
// If there was an error, just skip this invocation... we'll abort compilation anyway,
// but we can keep codegen'ing to find more errors.
Err(()) => return Ok(()),
}
}
Expand Down Expand Up @@ -1290,24 +1292,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 +1314,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 {
return_error!(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 All @@ -1371,13 +1349,12 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
.const_to_opt_u128(args[1].immediate(), false)
.expect("typeck should have ensure that this is a const");
if idx >= in_len.into() {
bx.sess().dcx().emit_err(InvalidMonomorphization::SimdIndexOutOfBounds {
return_error!(InvalidMonomorphization::SimdIndexOutOfBounds {
span,
name,
arg_idx: 1,
total_len: in_len.into(),
});
return Ok(bx.const_null(llret_ty));
}
return Ok(bx.insert_element(
args[0].immediate(),
Expand All @@ -1394,13 +1371,12 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
.const_to_opt_u128(args[1].immediate(), false)
.expect("typeck should have ensure that this is a const");
if idx >= in_len.into() {
bx.sess().dcx().emit_err(InvalidMonomorphization::SimdIndexOutOfBounds {
return_error!(InvalidMonomorphization::SimdIndexOutOfBounds {
span,
name,
arg_idx: 1,
total_len: in_len.into(),
});
return Ok(bx.const_null(llret_ty));
}
return Ok(bx.extract_element(args[0].immediate(), bx.const_i32(idx as i32)));
}
Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_codegen_ssa/messages.ftl
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
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,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❤️ glad to see this hack go away lol

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah me too :D

// 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
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
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
Loading
Loading