Skip to content

Commit 4f08893

Browse files
veluca93veluca93
committed
Implement struct_target_features for generic functions.
1 parent 5ae1957 commit 4f08893

File tree

12 files changed

+128
-45
lines changed

12 files changed

+128
-45
lines changed

‎compiler/rustc_codegen_llvm/src/attributes.rs

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use rustc_attr::{InlineAttr, InstructionSetAttr, OptimizeAttr};
44
use rustc_codegen_ssa::traits::*;
55
use rustc_hir::def_id::DefId;
66
use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, PatchableFunctionEntry};
7-
use rustc_middle::ty::{self, TyCtxt};
7+
use rustc_middle::ty::{self, ParamEnv, TyCtxt};
88
use rustc_session::config::{BranchProtection, FunctionReturn, OptLevel, PAuthKey, PacRet};
99
use rustc_span::symbol::sym;
1010
use rustc_target::spec::{FramePointer, SanitizerSet, StackProbeType, StackProtector};
@@ -500,7 +500,9 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
500500
to_add.extend(tune_cpu_attr(cx));
501501

502502
let function_features =
503-
codegen_fn_attrs.target_features.iter().map(|f| f.name.as_str()).collect::<Vec<&str>>();
503+
codegen_fn_attrs.target_features_for_instance(cx.tcx, ParamEnv::reveal_all(), instance);
504+
let function_features =
505+
function_features.iter().map(|f| f.name.as_str()).collect::<Vec<&str>>();
504506

505507
if let Some(f) = llvm_util::check_tied_features(
506508
cx.tcx.sess,

‎compiler/rustc_codegen_ssa/src/codegen_attrs.rs

Lines changed: 8 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,11 @@ use rustc_hir::weak_lang_items::WEAK_LANG_ITEMS;
99
use rustc_hir::{LangItem, lang_items};
1010
use rustc_middle::middle::codegen_fn_attrs::{
1111
CodegenFnAttrFlags, CodegenFnAttrs, PatchableFunctionEntry, TargetFeature,
12+
extend_with_struct_target_features,
1213
};
1314
use rustc_middle::mir::mono::Linkage;
1415
use rustc_middle::query::Providers;
15-
use rustc_middle::ty::{self as ty, Ty, TyCtxt};
16+
use rustc_middle::ty::{self as ty, TyCtxt};
1617
use rustc_session::lint;
1718
use rustc_session::parse::feature_err;
1819
use rustc_span::symbol::Ident;
@@ -292,7 +293,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
292293
tcx,
293294
attr,
294295
supported_target_features,
295-
&mut codegen_fn_attrs.target_features,
296+
&mut codegen_fn_attrs.def_target_features,
296297
Some(&mut codegen_fn_attrs.target_features_from_args),
297298
);
298299
}
@@ -599,8 +600,8 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
599600
let owner_id = tcx.parent(did.to_def_id());
600601
if tcx.def_kind(owner_id).has_codegen_attrs() {
601602
codegen_fn_attrs
602-
.target_features
603-
.extend(tcx.codegen_fn_attrs(owner_id).target_features.iter().copied());
603+
.def_target_features
604+
.extend(tcx.codegen_fn_attrs(owner_id).def_target_features.iter().copied());
604605
}
605606
}
606607

@@ -622,15 +623,16 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
622623
);
623624
}
624625
codegen_fn_attrs
625-
.target_features
626+
.def_target_features
626627
.extend(additional_tf.iter().map(|tf| TargetFeature { implied: true, ..*tf }));
627628
}
628629

629630
// If a function uses non-default target_features it can't be inlined into general
630631
// purpose functions as they wouldn't have the right target features
631632
// enabled. For that reason we also forbid #[inline(always)] as it can't be
632633
// respected.
633-
if !codegen_fn_attrs.target_features.is_empty() && codegen_fn_attrs.inline == InlineAttr::Always
634+
if !codegen_fn_attrs.def_target_features.is_empty()
635+
&& codegen_fn_attrs.inline == InlineAttr::Always
634636
{
635637
if let Some(span) = inline_span {
636638
tcx.dcx().span_err(
@@ -778,30 +780,6 @@ fn struct_target_features(tcx: TyCtxt<'_>, def_id: LocalDefId) -> &[TargetFeatur
778780
tcx.arena.alloc_slice(&features)
779781
}
780782

781-
fn extend_with_struct_target_features<'tcx>(
782-
tcx: TyCtxt<'tcx>,
783-
env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>,
784-
target_features: &mut Vec<TargetFeature>,
785-
) {
786-
// Collect target features from types reachable from `env.value` by dereferencing a certain
787-
// number of references and resolving aliases.
788-
789-
let mut ty = env.value;
790-
if matches!(ty.kind(), ty::Alias(..)) {
791-
ty = match tcx.try_normalize_erasing_regions(env.param_env, ty) {
792-
Ok(ty) => ty,
793-
Err(_) => return,
794-
};
795-
}
796-
while let ty::Ref(_, inner, _) = ty.kind() {
797-
ty = *inner;
798-
}
799-
800-
if let ty::Adt(adt_def, ..) = ty.kind() {
801-
target_features.extend_from_slice(&tcx.struct_target_features(adt_def.did()));
802-
}
803-
}
804-
805783
pub(crate) fn provide(providers: &mut Providers) {
806784
*providers = Providers {
807785
codegen_fn_attrs,

‎compiler/rustc_codegen_ssa/src/target_features.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ fn asm_target_features(tcx: TyCtxt<'_>, did: DefId) -> &FxIndexSet<Symbol> {
137137
let mut target_features = tcx.sess.unstable_target_features.clone();
138138
if tcx.def_kind(did).has_codegen_attrs() {
139139
let attrs = tcx.codegen_fn_attrs(did);
140-
target_features.extend(attrs.target_features.iter().map(|feature| feature.name));
140+
target_features.extend(attrs.def_target_features.iter().map(|feature| feature.name));
141141
match attrs.instruction_set {
142142
None => {}
143143
Some(InstructionSetAttr::ArmA32) => {

‎compiler/rustc_hir_analysis/src/check/entry.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ fn check_main_fn_ty(tcx: TyCtxt<'_>, main_def_id: DefId) {
105105
error = true;
106106
}
107107

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()
109109
// Calling functions with `#[target_feature]` is not unsafe on WASM, see #84988
110110
&& !tcx.sess.target.is_like_wasm
111111
&& !tcx.sess.opts.actually_rustdoc

‎compiler/rustc_hir_typeck/src/coercion.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -915,7 +915,7 @@ impl<'f, 'tcx> Coerce<'f, 'tcx> {
915915
&& self
916916
.tcx
917917
.codegen_fn_attrs(def_id)
918-
.target_features
918+
.def_target_features
919919
.iter()
920920
.any(|x| !x.implied)
921921
{

‎compiler/rustc_middle/src/middle/codegen_fn_attrs.rs

Lines changed: 82 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use rustc_target::abi::Align;
55
use rustc_target::spec::SanitizerSet;
66

77
use crate::mir::mono::Linkage;
8+
use crate::ty::{self, Instance, ParamEnv, Ty, TyCtxt};
89

910
#[derive(Clone, TyEncodable, TyDecodable, HashStable, Debug)]
1011
pub struct CodegenFnAttrs {
@@ -28,7 +29,7 @@ pub struct CodegenFnAttrs {
2829
pub link_ordinal: Option<u16>,
2930
/// All the target features that are enabled for this function. Some features might be enabled
3031
/// implicitly.
31-
pub target_features: Vec<TargetFeature>,
32+
pub def_target_features: Vec<TargetFeature>,
3233
/// The `#[linkage = "..."]` attribute on Rust-defined items and the value we found.
3334
pub linkage: Option<Linkage>,
3435
/// The `#[linkage = "..."]` attribute on foreign items and the value we found.
@@ -139,6 +140,30 @@ bitflags::bitflags! {
139140
}
140141
rustc_data_structures::external_bitflags_debug! { CodegenFnAttrFlags }
141142

143+
pub fn extend_with_struct_target_features<'tcx>(
144+
tcx: TyCtxt<'tcx>,
145+
env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>,
146+
target_features: &mut Vec<TargetFeature>,
147+
) {
148+
// Collect target features from types reachable from `env.value` by dereferencing a certain
149+
// number of references and resolving aliases.
150+
151+
let mut ty = env.value;
152+
if matches!(ty.kind(), ty::Alias(..)) {
153+
ty = match tcx.try_normalize_erasing_regions(env.param_env, ty) {
154+
Ok(ty) => ty,
155+
Err(_) => return,
156+
};
157+
}
158+
while let ty::Ref(_, inner, _) = ty.kind() {
159+
ty = *inner;
160+
}
161+
162+
if let ty::Adt(adt_def, ..) = ty.kind() {
163+
target_features.extend_from_slice(&tcx.struct_target_features(adt_def.did()));
164+
}
165+
}
166+
142167
impl CodegenFnAttrs {
143168
pub const EMPTY: &'static Self = &Self::new();
144169

@@ -150,7 +175,7 @@ impl CodegenFnAttrs {
150175
export_name: None,
151176
link_name: None,
152177
link_ordinal: None,
153-
target_features: vec![],
178+
def_target_features: vec![],
154179
linkage: None,
155180
import_linkage: None,
156181
link_section: None,
@@ -177,4 +202,59 @@ impl CodegenFnAttrs {
177202
Some(_) => true,
178203
}
179204
}
205+
206+
pub fn target_features_for_instance<'tcx>(
207+
&self,
208+
tcx: TyCtxt<'tcx>,
209+
param_env: ParamEnv<'tcx>,
210+
instance: Instance<'tcx>,
211+
) -> Vec<TargetFeature> {
212+
if !self.target_features_from_args {
213+
return self.def_target_features.clone();
214+
}
215+
let inputs = match tcx.type_of(instance.def_id()).skip_binder().kind() {
216+
ty::Closure(..) => {
217+
let closure = instance.args.as_closure();
218+
let mut inputs =
219+
tcx.instantiate_bound_regions_with_erased(closure.sig()).inputs().to_vec();
220+
inputs.extend(closure.upvar_tys());
221+
inputs
222+
}
223+
ty::CoroutineClosure(..) => {
224+
let closure = instance.args.as_coroutine_closure();
225+
// FIXME: might be missing inputs to the closure
226+
closure.upvar_tys().to_vec()
227+
}
228+
ty::Coroutine(..) => {
229+
let coro = instance.args.as_coroutine();
230+
coro.upvar_tys().to_vec()
231+
}
232+
_ => {
233+
let ty = match tcx.try_instantiate_and_normalize_erasing_regions(
234+
instance.args,
235+
param_env,
236+
tcx.type_of(instance.def_id()),
237+
) {
238+
Ok(ty) => ty,
239+
Err(_) => {
240+
return self.def_target_features.clone();
241+
}
242+
};
243+
let sig = tcx.instantiate_bound_regions_with_erased(ty.fn_sig(tcx));
244+
sig.inputs().to_vec()
245+
}
246+
};
247+
let mut additional_features = vec![];
248+
for input in inputs {
249+
extend_with_struct_target_features(tcx, param_env.and(input), &mut additional_features);
250+
}
251+
if additional_features.is_empty() {
252+
self.def_target_features.clone()
253+
} else {
254+
additional_features.extend_from_slice(&self.def_target_features);
255+
additional_features.sort_by_key(|a| (a.name, a.implied));
256+
additional_features.dedup_by_key(|a| a.name);
257+
additional_features
258+
}
259+
}
180260
}

‎compiler/rustc_middle/src/ty/context.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -364,7 +364,7 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
364364
}
365365

366366
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()
368368
}
369369

370370
fn require_lang_item(self, lang_item: TraitSolverLangItem) -> DefId {

‎compiler/rustc_mir_build/src/check_unsafety.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -476,7 +476,7 @@ impl<'a, 'tcx> Visitor<'a, 'tcx> for UnsafetyVisitor<'a, 'tcx> {
476476
// Implicit target features are OK because they are either a consequence of some
477477
// explicit target feature (which is checked to be present in the caller) or
478478
// come from a witness argument.
479-
let callee_features = &self.tcx.codegen_fn_attrs(func_did).target_features;
479+
let callee_features = &self.tcx.codegen_fn_attrs(func_did).def_target_features;
480480
if !self.tcx.sess.target.options.is_like_wasm
481481
&& !callee_features.iter().all(|feature| {
482482
feature.implied
@@ -1143,7 +1143,7 @@ pub(crate) fn check_unsafety(tcx: TyCtxt<'_>, def: LocalDefId) {
11431143
SafetyContext::Safe
11441144
}
11451145
});
1146-
let body_target_features = &tcx.body_codegen_attrs(def.to_def_id()).target_features;
1146+
let body_target_features = &tcx.body_codegen_attrs(def.to_def_id()).def_target_features;
11471147
let mut warnings = Vec::new();
11481148
let mut visitor = UnsafetyVisitor {
11491149
tcx,

‎compiler/rustc_mir_transform/src/inline.rs

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -470,8 +470,19 @@ impl<'tcx> Inliner<'tcx> {
470470
return Err("incompatible instruction set");
471471
}
472472

473-
let callee_feature_names = callee_attrs.target_features.iter().map(|f| f.name);
474-
let this_feature_names = self.codegen_fn_attrs.target_features.iter().map(|f| f.name);
473+
if callee_attrs.target_features_from_args || self.codegen_fn_attrs.target_features_from_args
474+
{
475+
// Since these functions inherit features from their arguments and might be
476+
// non-fully-instantiated generics, we give up MIR inlining.
477+
// FIXME: check if these are indeed non-fully-instantiated generics.
478+
// FIXME: we actually don't need to check target_features_from_args in the *caller*
479+
// once #127731 lands and is completed for all targets. Relatedly, we also won't need
480+
// to check equality below.
481+
return Err("using #[target_feature(from_args)]");
482+
}
483+
484+
let callee_feature_names = callee_attrs.def_target_features.iter().map(|f| f.name);
485+
let this_feature_names = self.codegen_fn_attrs.def_target_features.iter().map(|f| f.name);
475486
if callee_feature_names.ne(this_feature_names) {
476487
// In general it is not correct to inline a callee with target features that are a
477488
// subset of the caller. This is because the callee might contain calls, and the ABI of

‎compiler/rustc_trait_selection/src/error_reporting/traits/fulfillment_errors.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -464,7 +464,7 @@ impl<'a, 'tcx> TypeErrCtxt<'a, 'tcx> {
464464
let is_target_feature_fn = if let ty::FnDef(def_id, _) =
465465
*leaf_trait_ref.skip_binder().self_ty().kind()
466466
{
467-
self.tcx.codegen_fn_attrs(def_id).target_features.iter().any(|x| !x.implied)
467+
self.tcx.codegen_fn_attrs(def_id).def_target_features.iter().any(|x| !x.implied)
468468
} else {
469469
false
470470
};

‎compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -482,7 +482,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
482482
ty::FnDef(def_id, _) => {
483483
let tcx = self.tcx();
484484
if tcx.fn_sig(def_id).skip_binder().is_fn_trait_compatible()
485-
&& tcx.codegen_fn_attrs(def_id).target_features.is_empty()
485+
&& tcx.codegen_fn_attrs(def_id).def_target_features.is_empty()
486486
{
487487
candidates.vec.push(AsyncClosureCandidate);
488488
}
@@ -551,7 +551,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
551551
ty::FnDef(def_id, args) => {
552552
let tcx = self.tcx();
553553
if tcx.fn_sig(def_id).skip_binder().is_fn_trait_compatible()
554-
&& !tcx.codegen_fn_attrs(def_id).target_features.iter().any(|x| !x.implied)
554+
&& !tcx.codegen_fn_attrs(def_id).def_target_features.iter().any(|x| !x.implied)
555555
{
556556
candidates.vec.push(FnPointerCandidate {
557557
fn_host_effect: tcx

‎tests/assembly/struct-target-features.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,3 +37,15 @@ pub fn add_fma_combined(_: &Avx, _: &Fma, v: __m256) -> (__m256, __m256) {
3737
let r2 = unsafe { _mm256_fmadd_ps(v, v, v) };
3838
(r1, r2)
3939
}
40+
41+
#[target_feature(from_args)]
42+
fn add_generic<S>(_: S, v: __m256) -> __m256 {
43+
// CHECK-NOT: call
44+
// CHECK: vaddps
45+
unsafe { _mm256_add_ps(v, v) }
46+
}
47+
48+
pub fn add_using_generic(v: __m256) -> __m256 {
49+
assert!(is_x86_feature_detected!("avx"));
50+
add_generic(unsafe { Avx {} }, v)
51+
}

0 commit comments

Comments
 (0)