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 #145574

Author: Zalathar

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -29,8 +29,18 @@ pub(crate) fn apply_to_callsite(callsite: &Value, idx: AttributePlace, attrs: &[
 }
 
 /// Get LLVM attribute for the provided inline heuristic.
-#[inline]
-fn inline_attr<'ll>(cx: &CodegenCx<'ll, '_>, inline: InlineAttr) -> Option<&'ll Attribute> {
+pub(crate) fn inline_attr<'ll, 'tcx>(
+    cx: &CodegenCx<'ll, 'tcx>,
+    instance: ty::Instance<'tcx>,
+) -> Option<&'ll Attribute> {
+    // `optnone` requires `noinline`
+    let codegen_fn_attrs = cx.tcx.codegen_fn_attrs(instance.def_id());
+    let inline = match (codegen_fn_attrs.inline, &codegen_fn_attrs.optimize) {
+        (_, OptimizeAttr::DoNotOptimize) => InlineAttr::Never,
+        (InlineAttr::None, _) if instance.def.requires_inline(cx.tcx) => InlineAttr::Hint,
+        (inline, _) => inline,
+    };
+
     if !cx.tcx.sess.opts.unstable_opts.inline_llvm {
         // disable LLVM inlining
         return Some(AttributeKind::NoInline.create_attr(cx.llcx));
@@ -346,14 +356,6 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
         OptimizeAttr::Speed => {}
     }
 
-    // `optnone` requires `noinline`
-    let inline = match (codegen_fn_attrs.inline, &codegen_fn_attrs.optimize) {
-        (_, OptimizeAttr::DoNotOptimize) => InlineAttr::Never,
-        (InlineAttr::None, _) if instance.def.requires_inline(cx.tcx) => InlineAttr::Hint,
-        (inline, _) => inline,
-    };
-    to_add.extend(inline_attr(cx, inline));
-
     if cx.sess().must_emit_unwind_tables() {
         to_add.push(uwtable_attr(cx.llcx, cx.sess().opts.unstable_opts.use_sync_unwind));
     }
@@ -488,6 +490,14 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
     let function_features =
         codegen_fn_attrs.target_features.iter().map(|f| f.name.as_str()).collect::<Vec<&str>>();
 
+    // Apply function attributes as per usual if there are no user defined
+    // target features otherwise this will get applied at the callsite.
+    if function_features.is_empty() {
+        if let Some(inline_attr) = inline_attr(cx, instance) {
+            to_add.push(inline_attr);
+        }
+    }
+
     let function_features = function_features
         .iter()
         // Convert to LLVMFeatures and filter out unavailable ones
@@ -517,6 +527,7 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
     let function_features = function_features.iter().map(|s| s.as_str());
     let target_features: String =
         global_features.chain(function_features).intersperse(",").collect();
+
     if !target_features.is_empty() {
         to_add.push(llvm::CreateAttrStringValue(cx.llcx, "target-features", &target_features));
     }

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1392,7 +1392,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
     fn call(
         &mut self,
         llty: &'ll Type,
-        fn_attrs: Option<&CodegenFnAttrs>,
+        fn_call_attrs: Option<&CodegenFnAttrs>,
         fn_abi: Option<&FnAbi<'tcx, Ty<'tcx>>>,
         llfn: &'ll Value,
         args: &[&'ll Value],
@@ -1409,10 +1409,10 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         }
 
         // Emit CFI pointer type membership test
-        self.cfi_type_test(fn_attrs, fn_abi, instance, llfn);
+        self.cfi_type_test(fn_call_attrs, fn_abi, instance, llfn);
 
         // Emit KCFI operand bundle
-        let kcfi_bundle = self.kcfi_operand_bundle(fn_attrs, fn_abi, instance, llfn);
+        let kcfi_bundle = self.kcfi_operand_bundle(fn_call_attrs, fn_abi, instance, llfn);
         if let Some(kcfi_bundle) = kcfi_bundle.as_ref().map(|b| b.as_ref()) {
             bundles.push(kcfi_bundle);
         }
@@ -1429,6 +1429,29 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
                 c"".as_ptr(),
             )
         };
+
+        if let Some(instance) = instance {
+            // Attributes on the function definition being called
+            let fn_defn_attrs = self.cx.tcx.codegen_fn_attrs(instance.def_id());
+            if let Some(fn_call_attrs) = fn_call_attrs
+                && !fn_call_attrs.target_features.is_empty()
+                // If there is an inline attribute and a target feature that matches
+                // we will add the attribute to the callsite otherwise we'll omit
+                // this and not add the attribute to prevent soundness issues.
+                && let Some(inlining_rule) = attributes::inline_attr(&self.cx, instance)
+                && self.cx.tcx.is_target_feature_call_safe(
+                    &fn_call_attrs.target_features,
+                    &fn_defn_attrs.target_features,
+                )
+            {
+                attributes::apply_to_callsite(
+                    call,
+                    llvm::AttributePlace::Function,
+                    &[inlining_rule],
+                );
+            }
+        }
+
         if let Some(fn_abi) = fn_abi {
             fn_abi.apply_attrs_callsite(self, call);
         }

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -428,9 +428,16 @@ fn check_result(
     // llvm/llvm-project#70563).
     if !codegen_fn_attrs.target_features.is_empty()
         && matches!(codegen_fn_attrs.inline, InlineAttr::Always)
+        && !tcx.features().target_feature_inline_always()
         && let Some(span) = interesting_spans.inline
     {
-        tcx.dcx().span_err(span, "cannot use `#[inline(always)]` with `#[target_feature]`");
+        feature_err(
+            tcx.sess,
+            sym::target_feature_inline_always,
+            span,
+            "cannot use `#[inline(always)]` with `#[target_feature]`",
+        )
+        .emit();
     }
 
     // warn that inline has no effect when no_sanitize is present

compiler/rustc_feature/src/unstable.rs

@@ -642,6 +642,8 @@ declare_features! (
     (unstable, super_let, "1.88.0", Some(139076)),
     /// Allows subtrait items to shadow supertrait items.
     (unstable, supertrait_item_shadowing, "1.86.0", Some(89151)),
+    /// Allows the use of target_feature when a function is marked inline(always).
+    (unstable, target_feature_inline_always, "CURRENT_RUSTC_VERSION", Some(145574)),
     /// Allows using `#[thread_local]` on `static` items.
     (unstable, thread_local, "1.0.0", Some(29594)),
     /// Allows defining `trait X = A + B;` alias items.

compiler/rustc_lint_defs/src/builtin.rs

@@ -5141,3 +5141,57 @@ declare_lint! {
     "detects tail calls of functions marked with `#[track_caller]`",
     @feature_gate = explicit_tail_calls;
 }
+declare_lint! {
+    /// The `inline_always_mismatching_target_features` lint will trigger when a
+    /// function with the `#[inline(always)]` and `#[target_feature(enable = "...")]`
+    /// attributes is called and cannot be inlined due to missing target features in the caller.
+    ///
+    /// ### Example
+    ///
+    /// ```rust,ignore (fails on x86_64)
+    /// #[inline(always)]
+    /// #[target_feature(enable = "fp16")]
+    /// unsafe fn callee() {
+    ///     // operations using fp16 types
+    /// }
+    ///
+    /// // Caller does not enable the required target feature
+    /// fn caller() {
+    ///     unsafe { callee(); }
+    /// }
+    ///
+    /// fn main() {
+    ///     caller();
+    /// }
+    /// ```
+    ///
+    /// This will produce:
+    ///
+    /// ```text
+    /// warning: call to `#[inline(always)]`-annotated `callee` requires the same target features. Function will not have `alwaysinline` attribute applied
+    ///   --> $DIR/builtin.rs:5192:14
+    ///    |
+    /// 10 |     unsafe { callee(); }
+    ///    |              ^^^^^^^^
+    ///    |
+    /// note: `fp16` target feature enabled in `callee` here but missing from `caller`
+    ///   --> $DIR/builtin.rs:5185:1
+    ///    |
+    /// 3  | #[target_feature(enable = "fp16")]
+    ///    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    /// 4  | unsafe fn callee() {
+    ///    | ------------------
+    ///    = note: `#[warn(inline_always_mismatching_target_features)]` on by default
+    /// warning: 1 warning emitted
+    /// ```
+    ///
+    /// ### Explanation
+    ///
+    /// Inlining a function with a target feature attribute into a caller that
+    /// lacks the corresponding target feature can lead to unsound behavior.
+    /// LLVM may select the wrong instructions or registers, or reorder
+    /// operations, potentially resulting in runtime errors.
+    pub INLINE_ALWAYS_MISMATCHING_TARGET_FEATURES,
+    Warn,
+    r#"detects when a function annotated with `#[inline(always)]` and `#[target_feature(enable = "..")]` is inlined into a caller without the required target feature"#,
+}

compiler/rustc_middle/src/middle/codegen_fn_attrs.rs

@@ -82,7 +82,7 @@ pub enum TargetFeatureKind {
     Forced,
 }
 
-#[derive(Copy, Clone, Debug, TyEncodable, TyDecodable, HashStable)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, TyEncodable, TyDecodable, HashStable)]
 pub struct TargetFeature {
     /// The name of the target feature (e.g. "avx")
     pub name: Symbol,

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,
+                );
+            }
+            _ => (),
+        }
+    }
+}

compiler/rustc_mir_transform/src/errors.rs

@@ -2,13 +2,54 @@ 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;
 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]

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),

compiler/rustc_span/src/symbol.rs

@@ -2162,6 +2162,7 @@ symbols! {
         target_family,
         target_feature,
         target_feature_11,
+        target_feature_inline_always,
         target_has_atomic,
         target_has_atomic_equal_alignment,
         target_has_atomic_load_store,