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 rust-lang#118209

Author: GeorgeWort

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, SimdShuffleLastConst,
+};
 use crate::{
     borrow_set::BorrowSet,
     constraints::{OutlivesConstraint, OutlivesConstraintSet},
@@ -1630,6 +1632,35 @@ 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: 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

@@ -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}

compiler/rustc_codegen_ssa/src/errors.rs

@@ -598,6 +598,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,13 +5,13 @@ 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;
 
 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, BasicBlock, SwitchTargets, UnwindTerminateReason};
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, ValidityRequirement};
@@ -925,7 +925,15 @@ 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);
+                                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),
@@ -1003,9 +1011,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.node);
+            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.node
+                    && 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.node)
+            };
 
             if let (0, Some(ty::InstanceKind::Virtual(_, idx))) = (i, def) {
                 match op.val {

compiler/rustc_codegen_ssa/src/mir/constant.rs

@@ -1,4 +1,3 @@
-use crate::errors;
 use crate::mir::operand::OperandRef;
 use crate::traits::*;
 use rustc_middle::mir;
@@ -29,7 +28,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,12 +59,12 @@ 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>) {
+    ) -> (Option<Bx::Value>, Ty<'tcx>) {
         let ty = self.monomorphize(constant.ty());
         let val = self
             .eval_unevaluated_mir_constant_to_valtree(constant)
@@ -85,17 +84,15 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                             };
                             bx.scalar_to_backend(prim, scalar, bx.immediate_backend_type(layout))
                         } else {
-                            bug!("simd shuffle field {:?}", field)
+                            bug!("field is not a scalar {:?}", field)
                         }
                     })
                     .collect();
-                bx.const_struct(&values, false)
-            })
-            .unwrap_or_else(|| {
-                bx.tcx().dcx().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)
+                if ty.is_simd() {
+                    bx.const_vector(&values)
+                } else {
+                    bx.const_struct(&values, false)
+                }
             });
         (val, ty)
     }

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, INTERNAL_UNSTABLE
+    ),
 
     // ==========================================================================
     // Internal attributes, Layout related:

compiler/rustc_passes/messages.ftl

@@ -601,6 +601,18 @@ passes_rustc_allow_const_fn_unstable =
 passes_rustc_dirty_clean =
     attribute requires -Z query-dep-graph to be enabled
 
+passes_rustc_intrinsic_const_vector_arg =
+    attribute should be applied to functions in `extern "unadjusted"` modules
+    .label = not a function in an `extern "unadjusted"` module
+
+passes_rustc_intrinsic_const_vector_arg_invalid = attribute requires a parameter index
+
+passes_rustc_intrinsic_const_vector_arg_non_vector = parameter at index {$index} must be a simd type
+    .label = parameter is a non-simd type
+
+passes_rustc_intrinsic_const_vector_arg_out_of_bounds = function does not have a parameter at index {$index}
+    .label = function has {$arg_count} arguments
+
 passes_rustc_layout_scalar_valid_range_arg =
     expected exactly one integer literal argument