Rollup merge of #145932 - JamieCunliffe:target-feature-inlining, r=jackh726

Allow `inline(always)` with a target feature behind a unstable feature `target_feature_inline_always`.

Rather than adding the inline always attribute to the function definition, we add it to the callsite. We can then check that the target features match and that the call would be safe to inline. If the function isn't inlined due to a mismatch, we emit a warning informing the user that the function can't be inlined due to the target feature mismatch.

See tracking issue rust-lang/rust#145574

3 files changed, 133 insertions, 0 deletions

diff --git a/compiler/rustc_mir_transform/src/check_inline_always_target_features.rs b/compiler/rustc_mir_transform/src/check_inline_always_target_features.rs
new file mode 100644
index 00000000000..abad28f0a8f
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/check_inline_always_target_features.rs
@@ -0,0 +1,88 @@
+use rustc_hir::attrs::InlineAttr;
+use rustc_middle::middle::codegen_fn_attrs::{TargetFeature, TargetFeatureKind};
+use rustc_middle::mir::{Body, TerminatorKind};
+use rustc_middle::ty::{self, TyCtxt};
+
+use crate::pass_manager::MirLint;
+
+pub(super) struct CheckInlineAlwaysTargetFeature;
+
+impl<'tcx> MirLint<'tcx> for CheckInlineAlwaysTargetFeature {
+    fn run_lint(&self, tcx: TyCtxt<'tcx>, body: &Body<'tcx>) {
+        check_inline_always_target_features(tcx, body)
+    }
+}
+
+/// `#[target_feature]`-annotated functions can be marked `#[inline]` and will only be inlined if
+/// the target features match (as well as all of the other inlining heuristics). `#[inline(always)]`
+/// will always inline regardless of matching target features, which can result in errors from LLVM.
+/// However, it is desirable to be able to always annotate certain functions (e.g. SIMD intrinsics)
+/// as `#[inline(always)]` but check the target features match in Rust to avoid the LLVM errors.
+///
+/// We check the caller and callee target features to ensure that this can
+/// be done or emit a lint.
+#[inline]
+fn check_inline_always_target_features<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) {
+    let caller_def_id = body.source.def_id().expect_local();
+    if !tcx.def_kind(caller_def_id).has_codegen_attrs() {
+        return;
+    }
+
+    let caller_codegen_fn_attrs = tcx.codegen_fn_attrs(caller_def_id);
+
+    for bb in body.basic_blocks.iter() {
+        let terminator = bb.terminator();
+        match &terminator.kind {
+            TerminatorKind::Call { func, .. } | TerminatorKind::TailCall { func, .. } => {
+                let fn_ty = func.ty(body, tcx);
+                let ty::FnDef(callee_def_id, _) = *fn_ty.kind() else {
+                    continue;
+                };
+
+                if !tcx.def_kind(callee_def_id).has_codegen_attrs() {
+                    continue;
+                }
+                let callee_codegen_fn_attrs = tcx.codegen_fn_attrs(callee_def_id);
+                if callee_codegen_fn_attrs.inline != InlineAttr::Always
+                    || callee_codegen_fn_attrs.target_features.is_empty()
+                {
+                    continue;
+                }
+
+                // Scan the users defined target features and ensure they
+                // match the caller.
+                if tcx.is_target_feature_call_safe(
+                    &callee_codegen_fn_attrs.target_features,
+                    &caller_codegen_fn_attrs
+                        .target_features
+                        .iter()
+                        .cloned()
+                        .chain(tcx.sess.target_features.iter().map(|feat| TargetFeature {
+                            name: *feat,
+                            kind: TargetFeatureKind::Implied,
+                        }))
+                        .collect::<Vec<_>>(),
+                ) {
+                    continue;
+                }
+
+                let callee_only: Vec<_> = callee_codegen_fn_attrs
+                    .target_features
+                    .iter()
+                    .filter(|it| !caller_codegen_fn_attrs.target_features.contains(it))
+                    .filter(|it| !matches!(it.kind, TargetFeatureKind::Implied))
+                    .map(|it| it.name.as_str())
+                    .collect();
+
+                crate::errors::emit_inline_always_target_feature_diagnostic(
+                    tcx,
+                    terminator.source_info.span,
+                    callee_def_id,
+                    caller_def_id.into(),
+                    &callee_only,
+                );
+            }
+            _ => (),
+        }
+    }
+}
diff --git a/compiler/rustc_mir_transform/src/errors.rs b/compiler/rustc_mir_transform/src/errors.rs
index ad9635aae33..775f5f9a7cf 100644
--- a/compiler/rustc_mir_transform/src/errors.rs
+++ b/compiler/rustc_mir_transform/src/errors.rs
@@ -2,6 +2,7 @@ use rustc_errors::codes::*;
 use rustc_errors::{Diag, LintDiagnostic};
 use rustc_macros::{Diagnostic, LintDiagnostic, Subdiagnostic};
 use rustc_middle::mir::AssertKind;
+use rustc_middle::query::Key;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::lint::{self, Lint};
 use rustc_span::def_id::DefId;
@@ -9,6 +10,46 @@ use rustc_span::{Ident, Span, Symbol};
 
 use crate::fluent_generated as fluent;
 
+/// Emit diagnostic for calls to `#[inline(always)]`-annotated functions with a
+/// `#[target_feature]` attribute where the caller enables a different set of target features.
+pub(crate) fn emit_inline_always_target_feature_diagnostic<'a, 'tcx>(
+    tcx: TyCtxt<'tcx>,
+    call_span: Span,
+    callee_def_id: DefId,
+    caller_def_id: DefId,
+    callee_only: &[&'a str],
+) {
+    let callee = tcx.def_path_str(callee_def_id);
+    let caller = tcx.def_path_str(caller_def_id);
+
+    tcx.node_span_lint(
+        lint::builtin::INLINE_ALWAYS_MISMATCHING_TARGET_FEATURES,
+        tcx.local_def_id_to_hir_id(caller_def_id.as_local().unwrap()),
+        call_span,
+        |lint| {
+            lint.primary_message(format!(
+                "call to `#[inline(always)]`-annotated `{callee}` \
+                requires the same target features to be inlined"
+            ));
+            lint.note("function will not be inlined");
+
+            lint.note(format!(
+                "the following target features are on `{callee}` but missing from `{caller}`: {}",
+                callee_only.join(", ")
+            ));
+            lint.span_note(callee_def_id.default_span(tcx), format!("`{callee}` is defined here"));
+
+            let feats = callee_only.join(",");
+            lint.span_suggestion(
+                tcx.def_span(caller_def_id).shrink_to_lo(),
+                format!("add `#[target_feature]` attribute to `{caller}`"),
+                format!("#[target_feature(enable = \"{feats}\")]\n"),
+                lint::Applicability::MaybeIncorrect,
+            );
+        },
+    );
+}
+
 #[derive(LintDiagnostic)]
 #[diag(mir_transform_unconditional_recursion)]
 #[help]
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 8f319e64916..1663dfa744f 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/compiler/rustc_mir_transform/src/lib.rs
@@ -117,6 +117,7 @@ declare_passes! {
     mod add_subtyping_projections : Subtyper;
     mod check_inline : CheckForceInline;
     mod check_call_recursion : CheckCallRecursion, CheckDropRecursion;
+    mod check_inline_always_target_features: CheckInlineAlwaysTargetFeature;
     mod check_alignment : CheckAlignment;
     mod check_enums : CheckEnums;
     mod check_const_item_mutation : CheckConstItemMutation;
@@ -384,6 +385,9 @@ fn mir_built(tcx: TyCtxt<'_>, def: LocalDefId) -> &Steal<Body<'_>> {
             // MIR-level lints.
             &Lint(check_inline::CheckForceInline),
             &Lint(check_call_recursion::CheckCallRecursion),
+            // Check callee's target features match callers target features when
+            // using `#[inline(always)]`
+            &Lint(check_inline_always_target_features::CheckInlineAlwaysTargetFeature),
             &Lint(check_packed_ref::CheckPackedRef),
             &Lint(check_const_item_mutation::CheckConstItemMutation),
             &Lint(function_item_references::FunctionItemReferences),