Add rustc_intrinsic_const_vector_arg attribute to allow vectors to be passed as constants

Cargo.lock

@@ -3779,6 +3779,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.12"
 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

@@ -74,6 +74,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

@@ -479,3 +479,11 @@ pub(crate) struct SimdIntrinsicArgConst {
     pub arg: usize,
     pub intrinsic: String,
 }
+
+#[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

@@ -51,7 +51,9 @@ use rustc_mir_dataflow::move_paths::MoveData;
 use rustc_mir_dataflow::ResultsCursor;
 
 use crate::renumber::RegionCtxt;
-use crate::session_diagnostics::{MoveUnsized, SimdIntrinsicArgConst};
+use crate::session_diagnostics::{
+    IntrinsicConstVectorArgNonConst, MoveUnsized, SimdIntrinsicArgConst,
+};
 use crate::{
     borrow_set::BorrowSet,
     constraints::{OutlivesConstraint, OutlivesConstraintSet},
@@ -1630,6 +1632,33 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                         intrinsic: name.to_string(),
                     });
                 }
+            } 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| {
+                            if let Some(rustc_ast::LitKind::Int(index, _)) =
+                                item.lit().map(|lit| &lit.kind)
+                            {
+                                if let Some(arg) = args.get::<usize>(index.get() as usize) {
+                                    if !matches!(arg, Spanned { node: Operand::Constant(_), .. }) {
+                                        self.tcx().dcx().emit_err(
+                                            IntrinsicConstVectorArgNonConst {
+                                                span: term.source_info.span,
+                                                index: index.get(),
+                                            },
+                                        );
+                                    }
+                                }
+                            } else {
+                                span_mirbug!(self, term, "literal kind");
+                            }
+                        });
+                    }
+                    None => (),
+                }
+                // Error is reported by `rustc_attr!`
             }
         }
         debug!(?func_ty);

compiler/rustc_codegen_gcc/src/common.rs

@@ -161,6 +161,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

@@ -224,6 +224,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

@@ -19,6 +19,8 @@ codegen_ssa_check_installed_visual_studio = please ensure that Visual Studio 201
 codegen_ssa_compiler_builtins_cannot_call =
     `compiler_builtins` cannot call functions through upstream monomorphizations; encountered invalid call from `{$caller}` to `{$callee}`
 
+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}
@@ -31,6 +33,7 @@ codegen_ssa_expected_used_symbol = expected `used`, `used(compiler)` or `used(li
 
 codegen_ssa_extern_funcs_not_found = some `extern` functions couldn't be found; some native libraries may need to be installed or have their path specified
 
+
 codegen_ssa_extract_bundled_libs_archive_member = failed to get data from archive member '{$rlib}': {$error}
 codegen_ssa_extract_bundled_libs_convert_name = failed to convert name '{$rlib}': {$error}
 codegen_ssa_extract_bundled_libs_mmap_file = failed to mmap file '{$rlib}': {$error}
@@ -60,6 +63,8 @@ codegen_ssa_incorrect_cgu_reuse_type =
 
 codegen_ssa_insufficient_vs_code_product = VS Code is a different product, and is not sufficient.
 
+codegen_ssa_invalid_intrinsic_const_vector_arg = invalid argument, must be integer literal
+
 codegen_ssa_invalid_link_ordinal_nargs = incorrect number of arguments to `#[link_ordinal]`
     .note = the attribute requires exactly one argument
 

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -1,4 +1,4 @@
-use rustc_ast::{ast, attr, MetaItemKind, NestedMetaItem};
+use rustc_ast::{ast, attr, LitKind, MetaItemKind, MetaItemLit, NestedMetaItem};
 use rustc_attr::{list_contains_name, InlineAttr, InstructionSetAttr, OptimizeAttr};
 use rustc_errors::{codes::*, struct_span_code_err, DiagMessage, SubdiagMessage};
 use rustc_hir as hir;
@@ -523,6 +523,26 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
                     ))
                 })
             }
+            sym::rustc_intrinsic_const_vector_arg => {
+                let mut indicies: Vec<usize> = Vec::new();
+                if let Some(item_list) = attr.meta_item_list() {
+                    for item in item_list {
+                        match item {
+                            NestedMetaItem::Lit(MetaItemLit {
+                                kind: LitKind::Int(index, _),
+                                ..
+                            }) => {
+                                indicies.push(index.get() as usize);
+                            }
+                            _ => {
+                                tcx.dcx()
+                                    .emit_err(errors::IntrinsicConstVectorArg { span: attr.span });
+                            }
+                        }
+                    }
+                }
+                codegen_fn_attrs.const_vector_indices = Some(indicies);
+            }
             _ => {}
         }
     }

compiler/rustc_codegen_ssa/src/errors.rs

@@ -598,6 +598,20 @@ 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_invalid_intrinsic_const_vector_arg)]
+pub struct IntrinsicConstVectorArg {
+    #[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,7 +5,7 @@ use super::{CachedLlbb, FunctionCx, LocalRef};
 
 use crate::base::{self, is_call_from_compiler_builtins_to_upstream_monomorphization};
 use crate::common::{self, IntPredicate};
-use crate::errors::CompilerBuiltinsCannotCall;
+use crate::errors;
 use crate::meth;
 use crate::traits::*;
 use crate::MemFlags;
@@ -169,7 +169,7 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
                 if destination.is_some() {
                     let caller = with_no_trimmed_paths!(tcx.def_path_str(fx.instance.def_id()));
                     let callee = with_no_trimmed_paths!(tcx.def_path_str(instance.def_id()));
-                    tcx.dcx().emit_err(CompilerBuiltinsCannotCall { caller, callee });
+                    tcx.dcx().emit_err(errors::CompilerBuiltinsCannotCall { caller, callee });
                 } else {
                     info!(
                         "compiler_builtins call to diverging function {:?} replaced with abort",
@@ -925,7 +925,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         // checked by the type-checker.
                         if i == 2 && intrinsic.name == sym::simd_shuffle {
                             if let mir::Operand::Constant(constant) = &arg.node {
-                                let (llval, ty) = self.simd_shuffle_indices(bx, constant);
+                                let (llval, ty) = self.early_evaluate_const_vector(bx, constant);
+
                                 return OperandRef {
                                     val: Immediate(llval),
                                     layout: bx.layout_of(ty),
@@ -1003,9 +1004,32 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             (args, None)
         };
 
+        let const_vec_arg_indexes = if let Some(def) = def {
+            let val = bx.tcx().codegen_fn_attrs(def.def_id());
+            let const_vector_indices = &val.const_vector_indices;
+            if let Some(const_vector_i) = const_vector_indices {
+                &const_vector_i.as_ref()
+            } else {
+                &Vec::<usize>::new()
+            }
+        } 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.node);
+            let mut op = if const_vec_arg_indexes.contains(&i) {
+                if let mir::Operand::Constant(constant) = &arg.node
+                    && constant.ty().is_simd()
+                {
+                    let (llval, ty) = self.early_evaluate_const_vector(bx, &constant);
+                    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.node)
+            };
 
             if let (0, Some(ty::InstanceKind::Virtual(_, idx))) = (i, def) {
                 match op.val {

compiler/rustc_codegen_ssa/src/mir/constant.rs

@@ -4,6 +4,7 @@ use crate::traits::*;
 use rustc_middle::mir;
 use rustc_middle::mir::interpret::ErrorHandled;
 use rustc_middle::ty::layout::HasTyCtxt;
+use rustc_middle::ty::ValTree;
 use rustc_middle::ty::{self, Ty};
 use rustc_middle::{bug, span_bug};
 use rustc_target::abi::Abi;
@@ -29,7 +30,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             .expect("erroneous constant missed by mono item collection")
     }
 
-    /// This is a convenience helper for `simd_shuffle_indices`. It has the precondition
+    /// This is a convenience helper for `early_evaluate_const_vector`. It has the precondition
     /// that the given `constant` is an `Const::Unevaluated` and must be convertible to
     /// a `ValTree`. If you want a more general version of this, talk to `wg-const-eval` on zulip.
     ///
@@ -60,36 +61,48 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         self.cx.tcx().const_eval_resolve_for_typeck(ty::ParamEnv::reveal_all(), uv, constant.span)
     }
 
-    /// process constant containing SIMD shuffle indices
-    pub fn simd_shuffle_indices(
+    /// process constant SIMD vector or constant containing SIMD shuffle indices
+    pub fn early_evaluate_const_vector(
         &mut self,
         bx: &Bx,
         constant: &mir::ConstOperand<'tcx>,
     ) -> (Bx::Value, Ty<'tcx>) {
         let ty = self.monomorphize(constant.ty());
+        let ty_is_simd = ty.is_simd();
+        let field_ty = if ty_is_simd {
+            ty.simd_size_and_type(bx.tcx()).1
+        } else {
+            ty.builtin_index().unwrap()
+        };
         let val = self
             .eval_unevaluated_mir_constant_to_valtree(constant)
             .ok()
             .map(|x| x.ok())
             .flatten()
             .map(|val| {
-                let field_ty = ty.builtin_index().unwrap();
-                let values: Vec<_> = val
-                    .unwrap_branch()
-                    .iter()
-                    .map(|field| {
-                        if let Some(prim) = field.try_to_scalar() {
-                            let layout = bx.layout_of(field_ty);
-                            let Abi::Scalar(scalar) = layout.abi else {
-                                bug!("from_const: invalid ByVal layout: {:#?}", layout);
-                            };
-                            bx.scalar_to_backend(prim, scalar, bx.immediate_backend_type(layout))
-                        } else {
-                            bug!("simd shuffle field {:?}", field)
+                let mut values: Vec<Bx::Value> = Vec::new();
+                // For reliably being able to handle either:
+                // pub struct defn(i32, i32)
+                // call(1, 0);
+                // OR
+                // pub struct defn([i32, 2])
+                // call([1, 0]);
+                //
+                // And outputting: @call(<2 x i32> <i32 0, i32 1>)
+                // thus treating them the same if they are a const vector
+                for field in val.unwrap_branch().iter() {
+                    match field {
+                        ValTree::Branch(_) => {
+                            let scalars = self.flatten_branch_to_scalars(bx, &field, field_ty);
+                            values.extend(scalars);
+                        }
+                        ValTree::Leaf(_) => {
+                            let scalar = self.extract_scalar(bx, &field, field_ty);
+                            values.push(scalar);
                         }
-                    })
-                    .collect();
-                bx.const_struct(&values, false)
+                    }
+                }
+                if ty_is_simd { bx.const_vector(&values) } else { bx.const_struct(&values, false) }
             })
             .unwrap_or_else(|| {
                 bx.tcx().dcx().emit_err(errors::ShuffleIndicesEvaluation { span: constant.span });
@@ -99,4 +112,34 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             });
         (val, ty)
     }
+
+    fn flatten_branch_to_scalars(
+        &mut self,
+        bx: &Bx,
+        branch: &ValTree<'tcx>,
+        field_ty: Ty<'tcx>,
+    ) -> Vec<Bx::Value> {
+        branch
+            .unwrap_branch()
+            .iter()
+            .map(|field| match field {
+                ValTree::Branch(_) => {
+                    bug!("Cannot have arbitrarily nested const vectors: {:#?}", field)
+                }
+                ValTree::Leaf(_) => self.extract_scalar(bx, field, field_ty),
+            })
+            .collect()
+    }
+
+    fn extract_scalar(&mut self, bx: &Bx, field: &ValTree<'tcx>, field_ty: Ty<'tcx>) -> Bx::Value {
+        if let Some(prim) = field.try_to_scalar() {
+            let layout = bx.layout_of(field_ty);
+            let Abi::Scalar(scalar) = layout.abi else {
+                bug!("from_const: invalid ByVal layout: {:#?}", layout);
+            };
+            bx.scalar_to_backend(prim, scalar, bx.immediate_backend_type(layout))
+        } else {
+            bug!("field is not a scalar {:?}", field)
+        }
+    }
 }

compiler/rustc_codegen_ssa/src/traits/consts.rs

@@ -29,6 +29,7 @@ pub trait ConstMethods<'tcx>: BackendTypes {
 
     fn const_str(&self, s: &str) -> (Self::Value, Self::Value);
     fn const_struct(&self, elts: &[Self::Value], packed: bool) -> Self::Value;
+    fn const_vector(&self, elts: &[Self::Value]) -> Self::Value;
 
     fn const_to_opt_uint(&self, v: Self::Value) -> Option<u64>;
     fn const_to_opt_u128(&self, v: Self::Value, sign_ext: bool) -> Option<u128>;

compiler/rustc_feature/src/builtin_attrs.rs

@@ -837,6 +837,9 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         rustc_runtime, Normal, template!(Word), WarnFollowing,
         EncodeCrossCrate::No, INTERNAL_UNSTABLE
     ),
+    rustc_attr!(
+        rustc_intrinsic_const_vector_arg, Normal, template!(List: "arg1_index, arg2_index, ..."), ErrorFollowing, EncodeCrossCrate::Yes, INTERNAL_UNSTABLE
+    ),
 
     // ==========================================================================
     // Internal attributes, Layout related: