Add #[rustc_intrinsic_const_vector_arg] to allow vectors to be passed as constants

This allows constant vectors using a repr(simd) type to be propagated
through to the backend by reusing the functionality used to do a similar
thing for the simd_shuffle intrinsic.

fix #118209

Author: GeorgeWort

Cargo.lock

@@ -3496,6 +3496,7 @@ dependencies = [
  "either",
  "itertools",
  "polonius-engine",
+ "rustc_ast",
  "rustc_data_structures",
  "rustc_errors",
  "rustc_fluent_macro",

compiler/rustc_borrowck/Cargo.toml

@@ -8,6 +8,7 @@ edition = "2021"
 either = "1.5.0"
 itertools = "0.11"
 polonius-engine = "0.13.0"
+rustc_ast = { path = "../rustc_ast" }
 rustc_data_structures = { path = "../rustc_data_structures" }
 rustc_errors = { path = "../rustc_errors" }
 rustc_fluent_macro = { path = "../rustc_fluent_macro" }

compiler/rustc_borrowck/messages.ftl

@@ -71,6 +71,8 @@ borrowck_higher_ranked_lifetime_error =
 borrowck_higher_ranked_subtype_error =
     higher-ranked subtype error
 
+borrowck_intrinsic_const_vector_arg_non_const = argument at index {$index} must be a constant
+
 borrowck_lifetime_constraints_error =
     lifetime may not live long enough
 

compiler/rustc_borrowck/src/session_diagnostics.rs

@@ -459,3 +459,11 @@ pub(crate) struct SimdShuffleLastConst {
     #[primary_span]
     pub span: Span,
 }
+
+#[derive(Diagnostic)]
+#[diag(borrowck_intrinsic_const_vector_arg_non_const)]
+pub(crate) struct IntrinsicConstVectorArgNonConst {
+    #[primary_span]
+    pub span: Span,
+    pub index: u128,
+}

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -49,7 +49,9 @@ use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
 use rustc_mir_dataflow::move_paths::MoveData;
 use rustc_mir_dataflow::ResultsCursor;
 
-use crate::session_diagnostics::{MoveUnsized, SimdShuffleLastConst};
+use crate::session_diagnostics::{
+    IntrinsicConstVectorArgNonConst, MoveUnsized, SimdShuffleLastConst,
+};
 use crate::{
     borrow_set::BorrowSet,
     constraints::{OutlivesConstraint, OutlivesConstraintSet},
@@ -1579,6 +1581,35 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                     }
                     _ => {}
                 }
+            } else if let Some(attr) =
+                self.tcx().get_attr(def_id, sym::rustc_intrinsic_const_vector_arg)
+            {
+                match attr.meta_item_list() {
+                    Some(items) => {
+                        items.into_iter().for_each(|item: rustc_ast::NestedMetaItem| match item {
+                            rustc_ast::NestedMetaItem::Lit(rustc_ast::MetaItemLit {
+                                kind: rustc_ast::LitKind::Int(index, _),
+                                ..
+                            }) => {
+                                if index >= args.len() as u128 {
+                                    span_mirbug!(self, term, "index out of bounds");
+                                } else {
+                                    if !matches!(args[index as usize], Operand::Constant(_)) {
+                                        self.tcx().sess.emit_err(IntrinsicConstVectorArgNonConst {
+                                            span: term.source_info.span,
+                                            index,
+                                        });
+                                    }
+                                }
+                            }
+                            _ => {
+                                span_mirbug!(self, term, "invalid index");
+                            }
+                        });
+                    }
+                    // Error is reported by `rustc_attr!`
+                    None => (),
+                }
             }
         }
         debug!(?func_ty);

compiler/rustc_codegen_gcc/src/common.rs

@@ -159,6 +159,11 @@ impl<'gcc, 'tcx> ConstMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
         self.context.new_struct_constructor(None, struct_type.as_type(), None, values)
     }
 
+    fn const_vector(&self, values: &[RValue<'gcc>]) -> RValue<'gcc> {
+        let typ = self.type_vector(values[0].get_type(), values.len() as u64);
+        self.context.new_rvalue_from_vector(None, typ, values)
+    }
+
     fn const_to_opt_uint(&self, _v: RValue<'gcc>) -> Option<u64> {
         // TODO(antoyo)
         None

compiler/rustc_codegen_llvm/src/common.rs

@@ -98,10 +98,6 @@ impl<'ll> CodegenCx<'ll, '_> {
         unsafe { llvm::LLVMConstArray(ty, elts.as_ptr(), elts.len() as c_uint) }
     }
 
-    pub fn const_vector(&self, elts: &[&'ll Value]) -> &'ll Value {
-        unsafe { llvm::LLVMConstVector(elts.as_ptr(), elts.len() as c_uint) }
-    }
-
     pub fn const_bytes(&self, bytes: &[u8]) -> &'ll Value {
         bytes_in_context(self.llcx, bytes)
     }
@@ -217,6 +213,10 @@ impl<'ll, 'tcx> ConstMethods<'tcx> for CodegenCx<'ll, 'tcx> {
         struct_in_context(self.llcx, elts, packed)
     }
 
+    fn const_vector(&self, elts: &[&'ll Value]) -> &'ll Value {
+        unsafe { llvm::LLVMConstVector(elts.as_ptr(), elts.len() as c_uint) }
+    }
+
     fn const_to_opt_uint(&self, v: &'ll Value) -> Option<u64> {
         try_as_const_integral(v).and_then(|v| unsafe {
             let mut i = 0u64;

compiler/rustc_codegen_ssa/messages.ftl

@@ -16,6 +16,8 @@ codegen_ssa_cgu_not_recorded =
 
 codegen_ssa_check_installed_visual_studio = please ensure that Visual Studio 2017 or later, or Build Tools for Visual Studio were installed with the Visual C++ option.
 
+codegen_ssa_const_vector_evaluation = could not evaluate constant vector at compile time
+
 codegen_ssa_copy_path = could not copy {$from} to {$to}: {$error}
 
 codegen_ssa_copy_path_buf = unable to copy {$source_file} to {$output_path}: {$error}

compiler/rustc_codegen_ssa/src/errors.rs

@@ -663,6 +663,13 @@ pub struct ShuffleIndicesEvaluation {
     pub span: Span,
 }
 
+#[derive(Diagnostic)]
+#[diag(codegen_ssa_const_vector_evaluation)]
+pub struct ConstVectorEvaluation {
+    #[primary_span]
+    pub span: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag(codegen_ssa_missing_memory_ordering)]
 pub struct MissingMemoryOrdering;

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -5,12 +5,13 @@ use super::{CachedLlbb, FunctionCx, LocalRef};
 
 use crate::base;
 use crate::common::{self, IntPredicate};
+use crate::errors;
 use crate::meth;
 use crate::traits::*;
 use crate::MemFlags;
 
 use rustc_ast as ast;
-use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
+use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece, LitKind, MetaItemLit, NestedMetaItem};
 use rustc_hir::lang_items::LangItem;
 use rustc_middle::mir::{self, AssertKind, SwitchTargets, UnwindTerminateReason};
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, ValidityRequirement};
@@ -864,7 +865,15 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         // promotes any complex rvalues to constants.
                         if i == 2 && intrinsic == sym::simd_shuffle {
                             if let mir::Operand::Constant(constant) = arg {
-                                let (llval, ty) = self.simd_shuffle_indices(bx, constant);
+                                let (llval, ty) = self.early_evaluate_const_vector(bx, constant);
+                                let llval = llval.unwrap_or_else(|| {
+                                    bx.tcx().sess.emit_err(errors::ShuffleIndicesEvaluation {
+                                        span: constant.span,
+                                    });
+                                    // We've errored, so we don't have to produce working code.
+                                    let llty = bx.backend_type(bx.layout_of(ty));
+                                    bx.const_undef(llty)
+                                });
                                 return OperandRef {
                                     val: Immediate(llval),
                                     layout: bx.layout_of(ty),
@@ -908,9 +917,49 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             (args, None)
         };
 
+        let const_vec_arg_indexes = (|| {
+            if let Some(def) = def
+                && let Some(attr) =
+                    bx.tcx().get_attr(def.def_id(), sym::rustc_intrinsic_const_vector_arg)
+            {
+                attr.meta_item_list()
+                    .iter()
+                    .flatten()
+                    .map(|item: &NestedMetaItem| match item {
+                        NestedMetaItem::Lit(MetaItemLit {
+                            kind: LitKind::Int(index, _), ..
+                        }) => *index as usize,
+                        _ => span_bug!(item.span(), "attribute argument must be an integer"),
+                    })
+                    .collect()
+            } else {
+                Vec::<usize>::new()
+            }
+        })();
+
         let mut copied_constant_arguments = vec![];
         'make_args: for (i, arg) in first_args.iter().enumerate() {
-            let mut op = self.codegen_operand(bx, arg);
+            let mut op = if const_vec_arg_indexes.contains(&i) {
+                // Force the specified argument to be constant by using const-qualification to promote any complex rvalues to constant.
+                if let mir::Operand::Constant(constant) = arg
+                    && constant.ty().is_simd()
+                {
+                    let (llval, ty) = self.early_evaluate_const_vector(bx, &constant);
+                    let llval = llval.unwrap_or_else(|| {
+                        bx.tcx()
+                            .sess
+                            .emit_err(errors::ConstVectorEvaluation { span: constant.span });
+                        // We've errored, so we don't have to produce working code.
+                        let llty = bx.backend_type(bx.layout_of(ty));
+                        bx.const_undef(llty)
+                    });
+                    OperandRef { val: Immediate(llval), layout: bx.layout_of(ty) }
+                } else {
+                    span_bug!(span, "argument at {i} must be a constant vector");
+                }
+            } else {
+                self.codegen_operand(bx, arg)
+            };
 
             if let (0, Some(ty::InstanceDef::Virtual(_, idx))) = (i, def) {
                 match op.val {