Skip to content

Commit eba25b4

Browse files
RalfJungRalfJung
committed
simd_shuffle: require index argument to be a vector
1 parent 6c65d4f commit eba25b4

File tree

22 files changed

+184
-275
lines changed

22 files changed

+184
-275
lines changed

Diff for: ‎compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs

+13-27
Original file line numberDiff line numberDiff line change
@@ -180,34 +180,20 @@ pub(super) fn codegen_simd_intrinsic_call<'tcx>(
180180
return;
181181
}
182182

183-
// Make sure this is actually an array, since typeck only checks the length-suffixed
184-
// version of this intrinsic.
183+
// Make sure this is actually a SIMD vector.
185184
let idx_ty = fx.monomorphize(idx.node.ty(fx.mir, fx.tcx));
186-
let n: u16 = match idx_ty.kind() {
187-
ty::Array(ty, len) if matches!(ty.kind(), ty::Uint(ty::UintTy::U32)) => len
188-
.try_eval_target_usize(fx.tcx, ty::ParamEnv::reveal_all())
189-
.unwrap_or_else(|| {
190-
span_bug!(span, "could not evaluate shuffle index array length")
191-
})
192-
.try_into()
193-
.unwrap(),
194-
_ if idx_ty.is_simd()
195-
&& matches!(
196-
idx_ty.simd_size_and_type(fx.tcx).1.kind(),
197-
ty::Uint(ty::UintTy::U32)
198-
) =>
199-
{
200-
idx_ty.simd_size_and_type(fx.tcx).0.try_into().unwrap()
201-
}
202-
_ => {
203-
fx.tcx.dcx().span_err(
204-
span,
205-
format!("simd_shuffle index must be an array of `u32`, got `{}`", idx_ty),
206-
);
207-
// Prevent verifier error
208-
fx.bcx.ins().trap(TrapCode::UnreachableCodeReached);
209-
return;
210-
}
185+
let n: u16 = if idx_ty.is_simd()
186+
&& matches!(idx_ty.simd_size_and_type(fx.tcx).1.kind(), ty::Uint(ty::UintTy::U32))
187+
{
188+
idx_ty.simd_size_and_type(fx.tcx).0.try_into().unwrap()
189+
} else {
190+
fx.tcx.dcx().span_err(
191+
span,
192+
format!("simd_shuffle index must be a SIMD vector of `u32`, got `{}`", idx_ty),
193+
);
194+
// Prevent verifier error
195+
fx.bcx.ins().trap(TrapCode::UnreachableCodeReached);
196+
return;
211197
};
212198

213199
assert_eq!(x.layout(), y.layout());

Diff for: ‎compiler/rustc_codegen_gcc/src/builder.rs

+11-27
Original file line numberDiff line numberDiff line change
@@ -1939,33 +1939,17 @@ impl<'a, 'gcc, 'tcx> Builder<'a, 'gcc, 'tcx> {
19391939
self.int_type
19401940
};
19411941

1942-
let mut mask_elements = if let Some(vector_type) = mask.get_type().dyncast_vector() {
1943-
let mask_num_units = vector_type.get_num_units();
1944-
let mut mask_elements = vec![];
1945-
for i in 0..mask_num_units {
1946-
let index = self.context.new_rvalue_from_long(self.cx.type_u32(), i as _);
1947-
mask_elements.push(self.context.new_cast(
1948-
self.location,
1949-
self.extract_element(mask, index).to_rvalue(),
1950-
mask_element_type,
1951-
));
1952-
}
1953-
mask_elements
1954-
} else {
1955-
let struct_type = mask.get_type().is_struct().expect("mask should be of struct type");
1956-
let mask_num_units = struct_type.get_field_count();
1957-
let mut mask_elements = vec![];
1958-
for i in 0..mask_num_units {
1959-
let field = struct_type.get_field(i as i32);
1960-
mask_elements.push(self.context.new_cast(
1961-
self.location,
1962-
mask.access_field(self.location, field).to_rvalue(),
1963-
mask_element_type,
1964-
));
1965-
}
1966-
mask_elements
1967-
};
1968-
let mask_num_units = mask_elements.len();
1942+
let vector_type = mask.get_type().dyncast_vector().unwrap();
1943+
let mask_num_units = vector_type.get_num_units();
1944+
let mut mask_elements = vec![];
1945+
for i in 0..mask_num_units {
1946+
let index = self.context.new_rvalue_from_long(self.cx.type_u32(), i as _);
1947+
mask_elements.push(self.context.new_cast(
1948+
self.location,
1949+
self.extract_element(mask, index).to_rvalue(),
1950+
mask_element_type,
1951+
));
1952+
}
19691953

19701954
// NOTE: the mask needs to be the same length as the input vectors, so add the missing
19711955
// elements in the mask if needed.

Diff for: ‎compiler/rustc_codegen_gcc/src/intrinsic/simd.rs

+7-18
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ use rustc_codegen_ssa::traits::{BaseTypeMethods, BuilderMethods};
1414
#[cfg(feature = "master")]
1515
use rustc_hir as hir;
1616
use rustc_middle::mir::BinOp;
17-
use rustc_middle::span_bug;
1817
use rustc_middle::ty::layout::HasTyCtxt;
1918
use rustc_middle::ty::{self, Ty};
2019
use rustc_span::{sym, Span, Symbol};
@@ -353,24 +352,14 @@ pub fn generic_simd_intrinsic<'a, 'gcc, 'tcx>(
353352
}
354353

355354
if name == sym::simd_shuffle {
356-
// Make sure this is actually an array or SIMD vector, since typeck only checks the length-suffixed
357-
// version of this intrinsic.
355+
// Make sure this is actually a SIMD vector.
358356
let idx_ty = args[2].layout.ty;
359-
let n: u64 = match idx_ty.kind() {
360-
ty::Array(ty, len) if matches!(*ty.kind(), ty::Uint(ty::UintTy::U32)) => {
361-
len.try_eval_target_usize(bx.cx.tcx, ty::ParamEnv::reveal_all()).unwrap_or_else(
362-
|| span_bug!(span, "could not evaluate shuffle index array length"),
363-
)
364-
}
365-
_ if idx_ty.is_simd()
366-
&& matches!(
367-
idx_ty.simd_size_and_type(bx.cx.tcx).1.kind(),
368-
ty::Uint(ty::UintTy::U32)
369-
) =>
370-
{
371-
idx_ty.simd_size_and_type(bx.cx.tcx).0
372-
}
373-
_ => return_error!(InvalidMonomorphization::SimdShuffle { span, name, ty: idx_ty }),
357+
let n: u64 = if idx_ty.is_simd()
358+
&& matches!(idx_ty.simd_size_and_type(bx.cx.tcx).1.kind(), ty::Uint(ty::UintTy::U32))
359+
{
360+
idx_ty.simd_size_and_type(bx.cx.tcx).0
361+
} else {
362+
return_error!(InvalidMonomorphization::SimdShuffle { span, name, ty: idx_ty })
374363
};
375364
require_simd!(ret_ty, InvalidMonomorphization::SimdReturn { span, name, ty: ret_ty });
376365

Diff for: ‎compiler/rustc_codegen_llvm/src/intrinsic.rs

+24-48
Original file line numberDiff line numberDiff line change
@@ -1290,24 +1290,14 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
12901290
}
12911291

12921292
if name == sym::simd_shuffle {
1293-
// Make sure this is actually an array or SIMD vector, since typeck only checks the length-suffixed
1294-
// version of this intrinsic.
1293+
// Make sure this is actually a SIMD vector.
12951294
let idx_ty = args[2].layout.ty;
1296-
let n: u64 = match idx_ty.kind() {
1297-
ty::Array(ty, len) if matches!(ty.kind(), ty::Uint(ty::UintTy::U32)) => {
1298-
len.try_eval_target_usize(bx.cx.tcx, ty::ParamEnv::reveal_all()).unwrap_or_else(
1299-
|| span_bug!(span, "could not evaluate shuffle index array length"),
1300-
)
1301-
}
1302-
_ if idx_ty.is_simd()
1303-
&& matches!(
1304-
idx_ty.simd_size_and_type(bx.cx.tcx).1.kind(),
1305-
ty::Uint(ty::UintTy::U32)
1306-
) =>
1307-
{
1308-
idx_ty.simd_size_and_type(bx.cx.tcx).0
1309-
}
1310-
_ => return_error!(InvalidMonomorphization::SimdShuffle { span, name, ty: idx_ty }),
1295+
let n: u64 = if idx_ty.is_simd()
1296+
&& matches!(idx_ty.simd_size_and_type(bx.cx.tcx).1.kind(), ty::Uint(ty::UintTy::U32))
1297+
{
1298+
idx_ty.simd_size_and_type(bx.cx.tcx).0
1299+
} else {
1300+
return_error!(InvalidMonomorphization::SimdShuffle { span, name, ty: idx_ty })
13111301
};
13121302

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

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

1325-
let vector = args[2].immediate();
1326-
1327-
let indices: Option<Vec<_>> = (0..n)
1328-
.map(|i| {
1329-
let arg_idx = i;
1330-
let val = bx.const_get_elt(vector, i as u64);
1331-
match bx.const_to_opt_u128(val, true) {
1332-
None => {
1333-
bug!("typeck should have already ensured that these are const")
1334-
}
1335-
Some(idx) if idx >= total_len => {
1336-
bx.sess().dcx().emit_err(InvalidMonomorphization::SimdIndexOutOfBounds {
1337-
span,
1338-
name,
1339-
arg_idx,
1340-
total_len,
1341-
});
1342-
None
1343-
}
1344-
Some(idx) => Some(bx.const_i32(idx as i32)),
1345-
}
1346-
})
1347-
.collect();
1348-
let Some(indices) = indices else {
1349-
return Ok(bx.const_null(llret_ty));
1350-
};
1315+
// Check that the indices are in-bounds.
1316+
let indices = args[2].immediate();
1317+
for i in 0..n {
1318+
let val = bx.const_get_elt(indices, i as u64);
1319+
let idx = bx
1320+
.const_to_opt_u128(val, true)
1321+
.unwrap_or_else(|| bug!("typeck should have already ensured that these are const"));
1322+
if idx >= total_len {
1323+
bx.sess().dcx().emit_err(InvalidMonomorphization::SimdIndexOutOfBounds {
1324+
span,
1325+
name,
1326+
arg_idx: i,
1327+
total_len,
1328+
});
1329+
}
1330+
}
13511331

1352-
return Ok(bx.shuffle_vector(
1353-
args[0].immediate(),
1354-
args[1].immediate(),
1355-
bx.const_vector(&indices),
1356-
));
1332+
return Ok(bx.shuffle_vector(args[0].immediate(), args[1].immediate(), indices));
13571333
}
13581334

13591335
if name == sym::simd_insert {

Diff for: ‎compiler/rustc_codegen_ssa/messages.ftl

+1-1
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ codegen_ssa_invalid_monomorphization_simd_return = invalid monomorphization of `
132132
133133
codegen_ssa_invalid_monomorphization_simd_second = invalid monomorphization of `{$name}` intrinsic: expected SIMD second type, found non-SIMD `{$ty}`
134134
135-
codegen_ssa_invalid_monomorphization_simd_shuffle = invalid monomorphization of `{$name}` intrinsic: simd_shuffle index must be an array of `u32`, got `{$ty}`
135+
codegen_ssa_invalid_monomorphization_simd_shuffle = invalid monomorphization of `{$name}` intrinsic: simd_shuffle index must be a SIMD vector of `u32`, got `{$ty}`
136136
137137
codegen_ssa_invalid_monomorphization_simd_third = invalid monomorphization of `{$name}` intrinsic: expected SIMD third type, found non-SIMD `{$ty}`
138138

Diff for: ‎compiler/rustc_codegen_ssa/src/mir/block.rs

+2-26
Original file line numberDiff line numberDiff line change
@@ -915,32 +915,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
915915
}
916916
};
917917

918-
let args: Vec<_> = args
919-
.iter()
920-
.enumerate()
921-
.map(|(i, arg)| {
922-
// The indices passed to simd_shuffle in the
923-
// third argument must be constant. This is
924-
// checked by the type-checker.
925-
if i == 2 && intrinsic.name == sym::simd_shuffle {
926-
// FIXME: the simd_shuffle argument is actually an array,
927-
// not a vector, so we need this special hack to make sure
928-
// it is passed as an immediate. We should pass the
929-
// shuffle indices as a vector instead to avoid this hack.
930-
if let mir::Operand::Constant(constant) = &arg.node {
931-
let (llval, ty) = self.immediate_const_vector(bx, constant);
932-
return OperandRef {
933-
val: Immediate(llval),
934-
layout: bx.layout_of(ty),
935-
};
936-
} else {
937-
span_bug!(span, "shuffle indices must be constant");
938-
}
939-
}
940-
941-
self.codegen_operand(bx, &arg.node)
942-
})
943-
.collect();
918+
let args: Vec<_> =
919+
args.iter().map(|arg| self.codegen_operand(bx, &arg.node)).collect();
944920

945921
if matches!(intrinsic, ty::IntrinsicDef { name: sym::caller_location, .. }) {
946922
let location = self

Diff for: ‎compiler/rustc_codegen_ssa/src/mir/constant.rs

+11-24
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use rustc_middle::mir::interpret::ErrorHandled;
22
use rustc_middle::ty::layout::HasTyCtxt;
3-
use rustc_middle::ty::{self, Ty, ValTree};
3+
use rustc_middle::ty::{self, Ty};
44
use rustc_middle::{bug, mir, span_bug};
55
use rustc_target::abi::Abi;
66

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

7972
let val = self
8073
.eval_unevaluated_mir_constant_to_valtree(constant)
8174
.ok()
8275
.map(|x| x.ok())
8376
.flatten()
8477
.map(|val| {
85-
// Depending on whether this is a SIMD type with an array field
86-
// or a type with many fields (one for each elements), the valtree
87-
// is either a single branch with N children, or a root node
88-
// with exactly one child which then in turn has many children.
89-
// So we look at the first child to determine whether it is a
90-
// leaf or whether we have to go one more layer down.
91-
let branch_or_leaf = val.unwrap_branch();
92-
let first = branch_or_leaf.get(0).unwrap();
93-
let field_iter = match first {
94-
ValTree::Branch(_) => first.unwrap_branch().iter(),
95-
ValTree::Leaf(_) => branch_or_leaf.iter(),
96-
};
97-
let values: Vec<_> = field_iter
78+
// A SIMD type has a single field, which is an array.
79+
let fields = val.unwrap_branch();
80+
assert_eq!(fields.len(), 1);
81+
let array = fields[0].unwrap_branch();
82+
// Iterate over the array elements to obtain the values in the vector.
83+
let values: Vec<_> = array
84+
.iter()
9885
.map(|field| {
9986
if let Some(prim) = field.try_to_scalar() {
10087
let layout = bx.layout_of(field_ty);
@@ -107,7 +94,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
10794
}
10895
})
10996
.collect();
110-
if ty_is_simd { bx.const_vector(&values) } else { bx.const_struct(&values, false) }
97+
bx.const_vector(&values)
11198
})
11299
.unwrap_or_else(|| {
113100
bx.tcx().dcx().emit_err(errors::ShuffleIndicesEvaluation { span: constant.span });

Diff for: ‎library/core/src/intrinsics/simd.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -232,7 +232,7 @@ extern "rust-intrinsic" {
232232
///
233233
/// `T` must be a vector.
234234
///
235-
/// `U` must be a **const** array or vector of `u32`s. This means it must either refer to a named
235+
/// `U` must be a **const** vector of `u32`s. This means it must either refer to a named
236236
/// const or be given as an inline const expression (`const { ... }`).
237237
///
238238
/// `V` must be a vector with the same element type as `T` and the same length as `U`.

Diff for: ‎library/portable-simd/crates/core_simd/src/swizzle.rs

+12-4
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ pub trait Swizzle<const N: usize> {
8585
LaneCount<N>: SupportedLaneCount,
8686
LaneCount<M>: SupportedLaneCount,
8787
{
88-
// Safety: `vector` is a vector, and the index is a const array of u32.
88+
// Safety: `vector` is a vector, and the index is a const vector of u32.
8989
unsafe {
9090
core::intrinsics::simd::simd_shuffle(
9191
vector,
@@ -103,7 +103,11 @@ pub trait Swizzle<const N: usize> {
103103
output[i] = index as u32;
104104
i += 1;
105105
}
106-
output
106+
107+
// The index list needs to be returned as a vector.
108+
#[repr(simd)]
109+
struct SimdShuffleIdx<const LEN: usize>([u32; LEN]);
110+
SimdShuffleIdx(output)
107111
},
108112
)
109113
}
@@ -121,7 +125,7 @@ pub trait Swizzle<const N: usize> {
121125
LaneCount<N>: SupportedLaneCount,
122126
LaneCount<M>: SupportedLaneCount,
123127
{
124-
// Safety: `first` and `second` are vectors, and the index is a const array of u32.
128+
// Safety: `first` and `second` are vectors, and the index is a const vector of u32.
125129
unsafe {
126130
core::intrinsics::simd::simd_shuffle(
127131
first,
@@ -139,7 +143,11 @@ pub trait Swizzle<const N: usize> {
139143
output[i] = index as u32;
140144
i += 1;
141145
}
142-
output
146+
147+
// The index list needs to be returned as a vector.
148+
#[repr(simd)]
149+
struct SimdShuffleIdx<const LEN: usize>([u32; LEN]);
150+
SimdShuffleIdx(output)
143151
},
144152
)
145153
}

Diff for: ‎src/tools/miri/src/intrinsics/simd.rs

+1-7
Original file line numberDiff line numberDiff line change
@@ -664,15 +664,9 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
664664
let [left, right, index] = check_arg_count(args)?;
665665
let (left, left_len) = this.operand_to_simd(left)?;
666666
let (right, right_len) = this.operand_to_simd(right)?;
667+
let (index, index_len) = this.operand_to_simd(index)?;
667668
let (dest, dest_len) = this.mplace_to_simd(dest)?;
668669

669-
// `index` is an array or a SIMD type
670-
let (index, index_len) = match index.layout.ty.kind() {
671-
// FIXME: remove this once `index` must always be a SIMD vector.
672-
ty::Array(..) => (index.assert_mem_place(), index.len(this)?),
673-
_ => this.operand_to_simd(index)?,
674-
};
675-
676670
assert_eq!(left_len, right_len);
677671
assert_eq!(index_len, dest_len);
678672

0 commit comments

Comments
 (0)