Auto merge of #130268 - RalfJung:simd-shuffle-idx-vector, 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 #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 #128738, see that issue for more context.
rust-lang · Sep 14, 2024 · 6390091 · 6390091
2 parents f9567d0 + 7a1b5e5
commit 6390091
Show file tree

Hide file tree

Showing 24 changed files with 224 additions and 301 deletions.
diff --git a/compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs b/compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs
@@ -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());

diff --git a/compiler/rustc_codegen_gcc/src/builder.rs b/compiler/rustc_codegen_gcc/src/builder.rs
@@ -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.

diff --git a/compiler/rustc_codegen_gcc/src/intrinsic/simd.rs b/compiler/rustc_codegen_gcc/src/intrinsic/simd.rs
@@ -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};
@@ -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 });
 

diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs
@@ -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);
@@ -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 {

diff --git a/compiler/rustc_codegen_ssa/messages.ftl b/compiler/rustc_codegen_ssa/messages.ftl
@@ -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}`
 

diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -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

diff --git a/compiler/rustc_codegen_ssa/src/mir/constant.rs b/compiler/rustc_codegen_ssa/src/mir/constant.rs
@@ -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;
 
@@ -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);
@@ -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 });

diff --git 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 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`.

diff --git a/library/portable-simd/crates/core_simd/src/swizzle.rs b/library/portable-simd/crates/core_simd/src/swizzle.rs
@@ -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,
@@ -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)
                 },
             )
         }
@@ -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,
@@ -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)
                 },
             )
         }

diff --git a/src/tools/miri/src/intrinsics/simd.rs b/src/tools/miri/src/intrinsics/simd.rs
@@ -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);