Implement struct_target_features for generic functions.

veluca93 · veluca93 · commit fba90ce21eb4 · 2024-10-14T11:50:42.000+02:00
diff --git a/compiler/rustc_codegen_gcc/src/attributes.rs b/compiler/rustc_codegen_gcc/src/attributes.rs
@@ -6,7 +6,8 @@ use rustc_attr::InlineAttr;
 use rustc_attr::InstructionSetAttr;
 #[cfg(feature = "master")]
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
-use rustc_middle::ty;
+use rustc_middle::ty::{self, ParamEnv};
+use rustc_span::symbol::sym;
 
 use crate::context::CodegenCx;
 use crate::gcc_util::to_gcc_features;
@@ -70,8 +71,9 @@ pub fn from_fn_attrs<'gcc, 'tcx>(
         }
     }
 
-    let mut function_features = codegen_fn_attrs
-        .target_features
+    let function_features =
+        codegen_fn_attrs.target_features_for_instance(cx.tcx, ParamEnv::reveal_all(), instance);
+    let mut function_features = function_features
         .iter()
         .map(|features| features.name.as_str())
         .flat_map(|feat| to_gcc_features(cx.tcx.sess, feat).into_iter())
diff --git a/compiler/rustc_codegen_llvm/src/attributes.rs b/compiler/rustc_codegen_llvm/src/attributes.rs
@@ -4,7 +4,7 @@ use rustc_attr::{InlineAttr, InstructionSetAttr, OptimizeAttr};
 use rustc_codegen_ssa::traits::*;
 use rustc_hir::def_id::DefId;
 use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, PatchableFunctionEntry};
-use rustc_middle::ty::{self, TyCtxt};
+use rustc_middle::ty::{self, ParamEnv, TyCtxt};
 use rustc_session::config::{BranchProtection, FunctionReturn, OptLevel, PAuthKey, PacRet};
 use rustc_target::spec::{FramePointer, SanitizerSet, StackProbeType, StackProtector};
 use smallvec::SmallVec;
@@ -499,7 +499,9 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
     to_add.extend(tune_cpu_attr(cx));
 
     let function_features =
-        codegen_fn_attrs.target_features.iter().map(|f| f.name.as_str()).collect::<Vec<&str>>();
+        codegen_fn_attrs.target_features_for_instance(cx.tcx, ParamEnv::reveal_all(), instance);
+    let function_features =
+        function_features.iter().map(|f| f.name.as_str()).collect::<Vec<&str>>();
 
     let function_features = function_features
         .iter()
diff --git a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
@@ -10,6 +10,7 @@ use rustc_hir::weak_lang_items::WEAK_LANG_ITEMS;
 use rustc_hir::{LangItem, lang_items};
 use rustc_middle::middle::codegen_fn_attrs::{
     CodegenFnAttrFlags, CodegenFnAttrs, PatchableFunctionEntry, TargetFeature,
+    extend_with_struct_target_features,
 };
 use rustc_middle::mir::mono::Linkage;
 use rustc_middle::query::Providers;
@@ -293,7 +294,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
                     tcx,
                     attr,
                     supported_target_features,
-                    &mut codegen_fn_attrs.target_features,
+                    &mut codegen_fn_attrs.def_target_features,
                     Some(&mut codegen_fn_attrs.target_features_from_args),
                 );
             }
@@ -600,8 +601,8 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
         let owner_id = tcx.parent(did.to_def_id());
         if tcx.def_kind(owner_id).has_codegen_attrs() {
             codegen_fn_attrs
-                .target_features
-                .extend(tcx.codegen_fn_attrs(owner_id).target_features.iter().copied());
+                .def_target_features
+                .extend(tcx.codegen_fn_attrs(owner_id).def_target_features.iter().copied());
         }
     }
 
@@ -623,15 +624,16 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
             );
         }
         codegen_fn_attrs
-            .target_features
+            .def_target_features
             .extend(additional_tf.iter().map(|tf| TargetFeature { implied: true, ..*tf }));
     }
 
     // If a function uses non-default target_features it can't be inlined into general
     // purpose functions as they wouldn't have the right target features
     // enabled. For that reason we also forbid #[inline(always)] as it can't be
     // respected.
-    if !codegen_fn_attrs.target_features.is_empty() && codegen_fn_attrs.inline == InlineAttr::Always
+    if !codegen_fn_attrs.def_target_features.is_empty()
+        && codegen_fn_attrs.inline == InlineAttr::Always
     {
         if let Some(span) = inline_span {
             tcx.dcx().span_err(
@@ -697,7 +699,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
     if let Some(features) = check_tied_features(
         tcx.sess,
         &codegen_fn_attrs
-            .target_features
+            .def_target_features
             .iter()
             .map(|features| (features.name.as_str(), true))
             .collect(),
@@ -819,30 +821,6 @@ fn struct_target_features(tcx: TyCtxt<'_>, def_id: LocalDefId) -> &[TargetFeatur
     tcx.arena.alloc_slice(&features)
 }
 
-fn extend_with_struct_target_features<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>,
-    target_features: &mut Vec<TargetFeature>,
-) {
-    // Collect target features from types reachable from `env.value` by dereferencing a certain
-    // number of references and resolving aliases.
-
-    let mut ty = env.value;
-    if matches!(ty.kind(), ty::Alias(..)) {
-        ty = match tcx.try_normalize_erasing_regions(env.param_env, ty) {
-            Ok(ty) => ty,
-            Err(_) => return,
-        };
-    }
-    while let ty::Ref(_, inner, _) = ty.kind() {
-        ty = *inner;
-    }
-
-    if let ty::Adt(adt_def, ..) = ty.kind() {
-        target_features.extend_from_slice(&tcx.struct_target_features(adt_def.did()));
-    }
-}
-
 pub(crate) fn provide(providers: &mut Providers) {
     *providers = Providers {
         codegen_fn_attrs,
diff --git a/compiler/rustc_codegen_ssa/src/target_features.rs b/compiler/rustc_codegen_ssa/src/target_features.rs
@@ -137,7 +137,7 @@ fn asm_target_features(tcx: TyCtxt<'_>, did: DefId) -> &FxIndexSet<Symbol> {
     let mut target_features = tcx.sess.unstable_target_features.clone();
     if tcx.def_kind(did).has_codegen_attrs() {
         let attrs = tcx.codegen_fn_attrs(did);
-        target_features.extend(attrs.target_features.iter().map(|feature| feature.name));
+        target_features.extend(attrs.def_target_features.iter().map(|feature| feature.name));
         match attrs.instruction_set {
             None => {}
             Some(InstructionSetAttr::ArmA32) => {
diff --git a/compiler/rustc_hir_analysis/src/check/entry.rs b/compiler/rustc_hir_analysis/src/check/entry.rs
@@ -105,7 +105,7 @@ fn check_main_fn_ty(tcx: TyCtxt<'_>, main_def_id: DefId) {
         error = true;
     }
 
-    if !tcx.codegen_fn_attrs(main_def_id).target_features.is_empty()
+    if !tcx.codegen_fn_attrs(main_def_id).def_target_features.is_empty()
         // Calling functions with `#[target_feature]` is not unsafe on WASM, see #84988
         && !tcx.sess.target.is_like_wasm
         && !tcx.sess.opts.actually_rustdoc
diff --git a/compiler/rustc_hir_typeck/src/coercion.rs b/compiler/rustc_hir_typeck/src/coercion.rs
@@ -926,7 +926,7 @@ impl<'f, 'tcx> Coerce<'f, 'tcx> {
                         && self
                             .tcx
                             .codegen_fn_attrs(def_id)
-                            .target_features
+                            .def_target_features
                             .iter()
                             .any(|x| !x.implied)
                     {
diff --git a/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs b/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs
@@ -5,6 +5,7 @@ use rustc_target::abi::Align;
 use rustc_target::spec::SanitizerSet;
 
 use crate::mir::mono::Linkage;
+use crate::ty::{self, Instance, ParamEnv, Ty, TyCtxt};
 
 #[derive(Clone, TyEncodable, TyDecodable, HashStable, Debug)]
 pub struct CodegenFnAttrs {
@@ -28,7 +29,7 @@ pub struct CodegenFnAttrs {
     pub link_ordinal: Option<u16>,
     /// All the target features that are enabled for this function. Some features might be enabled
     /// implicitly.
-    pub target_features: Vec<TargetFeature>,
+    pub def_target_features: Vec<TargetFeature>,
     /// The `#[linkage = "..."]` attribute on Rust-defined items and the value we found.
     pub linkage: Option<Linkage>,
     /// The `#[linkage = "..."]` attribute on foreign items and the value we found.
@@ -139,6 +140,30 @@ bitflags::bitflags! {
 }
 rustc_data_structures::external_bitflags_debug! { CodegenFnAttrFlags }
 
+pub fn extend_with_struct_target_features<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>,
+    target_features: &mut Vec<TargetFeature>,
+) {
+    // Collect target features from types reachable from `env.value` by dereferencing a certain
+    // number of references and resolving aliases.
+
+    let mut ty = env.value;
+    if matches!(ty.kind(), ty::Alias(..)) {
+        ty = match tcx.try_normalize_erasing_regions(env.param_env, ty) {
+            Ok(ty) => ty,
+            Err(_) => return,
+        };
+    }
+    while let ty::Ref(_, inner, _) = ty.kind() {
+        ty = *inner;
+    }
+
+    if let ty::Adt(adt_def, ..) = ty.kind() {
+        target_features.extend_from_slice(&tcx.struct_target_features(adt_def.did()));
+    }
+}
+
 impl CodegenFnAttrs {
     pub const EMPTY: &'static Self = &Self::new();
 
@@ -150,7 +175,7 @@ impl CodegenFnAttrs {
             export_name: None,
             link_name: None,
             link_ordinal: None,
-            target_features: vec![],
+            def_target_features: vec![],
             linkage: None,
             import_linkage: None,
             link_section: None,
@@ -177,4 +202,59 @@ impl CodegenFnAttrs {
                 Some(_) => true,
             }
     }
+
+    pub fn target_features_for_instance<'tcx>(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        param_env: ParamEnv<'tcx>,
+        instance: Instance<'tcx>,
+    ) -> Vec<TargetFeature> {
+        if !self.target_features_from_args {
+            return self.def_target_features.clone();
+        }
+        let inputs = match tcx.type_of(instance.def_id()).skip_binder().kind() {
+            ty::Closure(..) => {
+                let closure = instance.args.as_closure();
+                let mut inputs =
+                    tcx.instantiate_bound_regions_with_erased(closure.sig()).inputs().to_vec();
+                inputs.extend(closure.upvar_tys());
+                inputs
+            }
+            ty::CoroutineClosure(..) => {
+                let closure = instance.args.as_coroutine_closure();
+                // FIXME: might be missing inputs to the closure
+                closure.upvar_tys().to_vec()
+            }
+            ty::Coroutine(..) => {
+                let coro = instance.args.as_coroutine();
+                coro.upvar_tys().to_vec()
+            }
+            _ => {
+                let ty = match tcx.try_instantiate_and_normalize_erasing_regions(
+                    instance.args,
+                    param_env,
+                    tcx.type_of(instance.def_id()),
+                ) {
+                    Ok(ty) => ty,
+                    Err(_) => {
+                        return self.def_target_features.clone();
+                    }
+                };
+                let sig = tcx.instantiate_bound_regions_with_erased(ty.fn_sig(tcx));
+                sig.inputs().to_vec()
+            }
+        };
+        let mut additional_features = vec![];
+        for input in inputs {
+            extend_with_struct_target_features(tcx, param_env.and(input), &mut additional_features);
+        }
+        if additional_features.is_empty() {
+            self.def_target_features.clone()
+        } else {
+            additional_features.extend_from_slice(&self.def_target_features);
+            additional_features.sort_by_key(|a| (a.name, a.implied));
+            additional_features.dedup_by_key(|a| a.name);
+            additional_features
+        }
+    }
 }
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
@@ -364,7 +364,7 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
     }
 
     fn has_target_features(self, def_id: DefId) -> bool {
-        !self.codegen_fn_attrs(def_id).target_features.is_empty()
+        !self.codegen_fn_attrs(def_id).def_target_features.is_empty()
     }
 
     fn require_lang_item(self, lang_item: TraitSolverLangItem) -> DefId {
diff --git a/compiler/rustc_mir_build/src/check_unsafety.rs b/compiler/rustc_mir_build/src/check_unsafety.rs
@@ -476,7 +476,7 @@ impl<'a, 'tcx> Visitor<'a, 'tcx> for UnsafetyVisitor<'a, 'tcx> {
                     // Implicit target features are OK because they are either a consequence of some
                     // explicit target feature (which is checked to be present in the caller) or
                     // come from a witness argument.
-                    let callee_features = &self.tcx.codegen_fn_attrs(func_did).target_features;
+                    let callee_features = &self.tcx.codegen_fn_attrs(func_did).def_target_features;
                     if !self.tcx.sess.target.options.is_like_wasm
                         && !callee_features.iter().all(|feature| {
                             feature.implied
@@ -1143,7 +1143,7 @@ pub(crate) fn check_unsafety(tcx: TyCtxt<'_>, def: LocalDefId) {
             SafetyContext::Safe
         }
     });
-    let body_target_features = &tcx.body_codegen_attrs(def.to_def_id()).target_features;
+    let body_target_features = &tcx.body_codegen_attrs(def.to_def_id()).def_target_features;
     let mut warnings = Vec::new();
     let mut visitor = UnsafetyVisitor {
         tcx,
diff --git a/compiler/rustc_mir_transform/src/inline.rs b/compiler/rustc_mir_transform/src/inline.rs
@@ -470,8 +470,19 @@ impl<'tcx> Inliner<'tcx> {
             return Err("incompatible instruction set");
         }
 
-        let callee_feature_names = callee_attrs.target_features.iter().map(|f| f.name);
-        let this_feature_names = self.codegen_fn_attrs.target_features.iter().map(|f| f.name);
+        if callee_attrs.target_features_from_args || self.codegen_fn_attrs.target_features_from_args
+        {
+            // Since these functions inherit features from their arguments and might be
+            // non-fully-instantiated generics, we give up MIR inlining.
+            // FIXME: check if these are indeed non-fully-instantiated generics.
+            // FIXME: we actually don't need to check target_features_from_args in the *caller*
+            // once #127731 lands and is completed for all targets. Relatedly, we also won't need
+            // to check equality below.
+            return Err("using #[target_feature(from_args)]");
+        }
+
+        let callee_feature_names = callee_attrs.def_target_features.iter().map(|f| f.name);
+        let this_feature_names = self.codegen_fn_attrs.def_target_features.iter().map(|f| f.name);
         if callee_feature_names.ne(this_feature_names) {
             // In general it is not correct to inline a callee with target features that are a
             // subset of the caller. This is because the callee might contain calls, and the ABI of
diff --git a/compiler/rustc_trait_selection/src/error_reporting/traits/fulfillment_errors.rs b/compiler/rustc_trait_selection/src/error_reporting/traits/fulfillment_errors.rs
@@ -464,7 +464,7 @@ impl<'a, 'tcx> TypeErrCtxt<'a, 'tcx> {
                         let is_target_feature_fn = if let ty::FnDef(def_id, _) =
                             *leaf_trait_ref.skip_binder().self_ty().kind()
                         {
-                            self.tcx.codegen_fn_attrs(def_id).target_features.iter().any(|x| !x.implied)
+                            self.tcx.codegen_fn_attrs(def_id).def_target_features.iter().any(|x| !x.implied)
                         } else {
                             false
                         };
diff --git a/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs b/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs
@@ -482,7 +482,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             ty::FnDef(def_id, _) => {
                 let tcx = self.tcx();
                 if tcx.fn_sig(def_id).skip_binder().is_fn_trait_compatible()
-                    && tcx.codegen_fn_attrs(def_id).target_features.is_empty()
+                    && tcx.codegen_fn_attrs(def_id).def_target_features.is_empty()
                 {
                     candidates.vec.push(AsyncClosureCandidate);
                 }
@@ -551,7 +551,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             ty::FnDef(def_id, args) => {
                 let tcx = self.tcx();
                 if tcx.fn_sig(def_id).skip_binder().is_fn_trait_compatible()
-                    && !tcx.codegen_fn_attrs(def_id).target_features.iter().any(|x| !x.implied)
+                    && !tcx.codegen_fn_attrs(def_id).def_target_features.iter().any(|x| !x.implied)
                 {
                     candidates.vec.push(FnPointerCandidate {
                         fn_host_effect: tcx
diff --git a/src/tools/miri/src/machine.rs b/src/tools/miri/src/machine.rs
@@ -15,7 +15,9 @@ use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::static_assert_size;
 use rustc_middle::mir;
 use rustc_middle::query::TyCtxtAt;
-use rustc_middle::ty::layout::{HasTyCtxt, LayoutCx, LayoutError, LayoutOf, TyAndLayout};
+use rustc_middle::ty::layout::{
+    HasParamEnv, HasTyCtxt, LayoutCx, LayoutError, LayoutOf, TyAndLayout,
+};
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_session::config::InliningThreshold;
 use rustc_span::def_id::{CrateNum, DefId};
@@ -964,12 +966,12 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> {
     ) -> InterpResult<'tcx> {
         let attrs = ecx.tcx.codegen_fn_attrs(instance.def_id());
         if attrs
-            .target_features
+            .target_features_for_instance(ecx.tcx.tcx, ecx.param_env(), instance)
             .iter()
             .any(|feature| !ecx.tcx.sess.target_features.contains(&feature.name))
         {
             let unavailable = attrs
-                .target_features
+                .target_features_for_instance(ecx.tcx.tcx, ecx.param_env(), instance)
                 .iter()
                 .filter(|&feature| {
                     !feature.implied && !ecx.tcx.sess.target_features.contains(&feature.name)
diff --git a/tests/assembly/struct-target-features.rs b/tests/assembly/struct-target-features.rs
@@ -37,3 +37,15 @@ pub fn add_fma_combined(_: &Avx, _: &Fma, v: __m256) -> (__m256, __m256) {
     let r2 = unsafe { _mm256_fmadd_ps(v, v, v) };
     (r1, r2)
 }
+
+#[target_feature(from_args)]
+fn add_generic<S>(_: S, v: __m256) -> __m256 {
+    // CHECK-NOT: call
+    // CHECK: vaddps
+    unsafe { _mm256_add_ps(v, v) }
+}
+
+pub fn add_using_generic(v: __m256) -> __m256 {
+    assert!(is_x86_feature_detected!("avx"));
+    add_generic(unsafe { Avx {} }, v)
+}

Original file line number	Diff line number	Diff line change
`@@ -105,7 +105,7 @@ fn check_main_fn_ty(tcx: TyCtxt<'_>, main_def_id: DefId) {`
`105`	`105`	`error = true;`
`106`	`106`	`}`
`107`	`107`
`108`		`- if !tcx.codegen_fn_attrs(main_def_id).target_features.is_empty()`
	`108`	`+ if !tcx.codegen_fn_attrs(main_def_id).def_target_features.is_empty()`
`109`	`109`	// Calling functions with `#[target_feature]` is not unsafe on WASM, see #84988
`110`	`110`	`&& !tcx.sess.target.is_like_wasm`
`111`	`111`	`&& !tcx.sess.opts.actually_rustdoc`
Original file line number	Diff line number	Diff line change
`@@ -926,7 +926,7 @@ impl<'f, 'tcx> Coerce<'f, 'tcx> {`
`926`	`926`	`&& self`
`927`	`927`	`.tcx`
`928`	`928`	`.codegen_fn_attrs(def_id)`
`929`		`- .target_features`
	`929`	`+ .def_target_features`
`930`	`930`	`.iter()`
`931`	`931`	`.any(\|x\| !x.implied)`
`932`	`932`	`{`
Original file line number	Diff line number	Diff line change
`@@ -364,7 +364,7 @@ impl<'tcx> Interner for TyCtxt<'tcx> {`
`364`	`364`	`}`
`365`	`365`
`366`	`366`	`fn has_target_features(self, def_id: DefId) -> bool {`
`367`		`- !self.codegen_fn_attrs(def_id).target_features.is_empty()`
	`367`	`+ !self.codegen_fn_attrs(def_id).def_target_features.is_empty()`
`368`	`368`	`}`
`369`	`369`
`370`	`370`	`fn require_lang_item(self, lang_item: TraitSolverLangItem) -> DefId {`