Make inline attributes apply to the generated poll in async fn

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -72,8 +72,15 @@ fn process_builtin_attrs(
             AttributeKind::Cold => codegen_fn_attrs.flags |= CodegenFnAttrFlags::COLD,
             AttributeKind::ExportName { name, .. } => codegen_fn_attrs.symbol_name = Some(*name),
             AttributeKind::Inline(inline, span) => {
-                codegen_fn_attrs.inline = *inline;
-                interesting_spans.inline = Some(*span);
+                // On an `async fn` (or similar), inlining hints are meant for the code the user
+                // wrote, which lives in the body coroutine. The constructor is trivial and left to
+                // the usual heuristics. `#[rustc_force_inline]` is rejected on such items anyway.
+                if matches!(inline, InlineAttr::Force { .. })
+                    || !constructs_desugared_body_coroutine(tcx, did)
+                {
+                    codegen_fn_attrs.inline = *inline;
+                    interesting_spans.inline = Some(*span);
+                }
             }
             AttributeKind::Naked(_) => codegen_fn_attrs.flags |= CodegenFnAttrFlags::NAKED,
             AttributeKind::RustcAlign { align, .. } => codegen_fn_attrs.alignment = Some(*align),
@@ -297,9 +304,52 @@ fn process_builtin_attrs(
         }
     }
 
+    // The body coroutine can't carry attributes of its own, so it takes the ones describing the
+    // code the user wrote from the `async fn` (or similar) it was desugared from. This makes e.g.
+    // `#[inline(never)] async fn` affect `Future::poll`. See #129347.
+    if is_desugared_body_coroutine(tcx, did) {
+        let parent = tcx.local_parent(did);
+        if let Some((inline, span)) = find_attr!(
+            tcx,
+            parent,
+            Inline(inline, span) if !matches!(inline, InlineAttr::Force { .. }) => (*inline, *span)
+        ) {
+            codegen_fn_attrs.inline = inline;
+            interesting_spans.inline = Some(span);
+        }
+        // Unlike inlining hints, `#[cold]` stays on the constructor as well: calling it is just as
+        // unlikely, and its call sites are only recognized as cold through that.
+        if find_attr!(tcx, parent, Cold) {
+            codegen_fn_attrs.flags |= CodegenFnAttrFlags::COLD;
+        }
+    }
+
     interesting_spans
 }
 
+/// Whether `did` is the coroutine holding the body of an `async fn`, `gen fn`, `async gen fn`, or
+/// `async` closure, as opposed to one written as an `async {}` block.
+fn is_desugared_body_coroutine(tcx: TyCtxt<'_>, did: LocalDefId) -> bool {
+    matches!(
+        tcx.coroutine_kind(did),
+        Some(hir::CoroutineKind::Desugared(
+            _,
+            hir::CoroutineSource::Fn | hir::CoroutineSource::Closure
+        ))
+    )
+}
+
+/// Whether `did` is an `async fn`, `gen fn`, `async gen fn`, or `async` closure, i.e. its body
+/// does nothing but construct a [desugared body coroutine](is_desugared_body_coroutine).
+fn constructs_desugared_body_coroutine(tcx: TyCtxt<'_>, did: LocalDefId) -> bool {
+    tcx.hir_maybe_body_owned_by(did).is_some_and(|body| {
+        matches!(
+            body.value.kind,
+            hir::ExprKind::Closure(closure) if is_desugared_body_coroutine(tcx, closure.def_id)
+        )
+    })
+}
+
 /// Applies overrides for codegen fn attrs. These often have a specific reason why they're necessary.
 /// Please comment why when adding a new one!
 fn apply_overrides(tcx: TyCtxt<'_>, did: LocalDefId, codegen_fn_attrs: &mut CodegenFnAttrs) {
@@ -336,7 +386,18 @@ fn apply_overrides(tcx: TyCtxt<'_>, did: LocalDefId, codegen_fn_attrs: &mut Code
     // its parent function, which effectively inherits the features anyway. Boxing this closure
     // would result in this closure being compiled without the inherited target features, but this
     // is probably a poor usage of `#[inline(always)]` and easily avoided by not using the attribute.
-    if tcx.is_closure_like(did.to_def_id()) && codegen_fn_attrs.inline != InlineAttr::Always {
+    //
+    // That reasoning doesn't hold for the body coroutine of an `async fn` or `gen fn`: it is never
+    // inlined into the function constructing it, yet its body was checked as part of that function.
+    // It always inherits, which makes `check_result` reject `#[inline(always)]` there just like it
+    // does on the function itself.
+    if tcx.is_closure_like(did.to_def_id())
+        && (codegen_fn_attrs.inline != InlineAttr::Always
+            || matches!(
+                tcx.coroutine_kind(did),
+                Some(hir::CoroutineKind::Desugared(_, hir::CoroutineSource::Fn))
+            ))
+    {
         let owner_id = tcx.parent(did.to_def_id());
         if tcx.def_kind(owner_id).has_codegen_attrs() {
             codegen_fn_attrs

tests/codegen-llvm/async-fn-cold-attr.rs

@@ -0,0 +1,51 @@
+// Checks that `#[cold]` on `async fn`s and async closures applies to the coroutine holding their
+// body (i.e. the `Future::poll` implementation), in addition to the function or closure that
+// constructs it.
+//
+//@ edition: 2024
+//@ compile-flags: -Cno-prepopulate-passes -Csymbol-mangling-version=v0 -Zinline-mir=no
+#![crate_type = "lib"]
+#![feature(stmt_expr_attributes)]
+
+use std::future::Future;
+use std::pin::pin;
+use std::task::Context;
+
+// CHECK:      ; async_fn_cold_attr::cold{{$}}
+// CHECK-NEXT: ; Function Attrs: cold
+// CHECK:      ; async_fn_cold_attr::cold::{closure#0}
+// CHECK-NEXT: ; Function Attrs: cold
+#[cold]
+pub async fn cold() {}
+
+// CHECK:      ; async_fn_cold_attr::not_cold::{closure#0}
+// CHECK-NEXT: ; Function Attrs:
+// CHECK-NOT:  cold
+// CHECK-SAME: {{$}}
+pub async fn not_cold() {}
+
+pub struct S;
+
+impl S {
+    // CHECK:      ; <async_fn_cold_attr::S>::method{{$}}
+    // CHECK-NEXT: ; Function Attrs: cold
+    // CHECK:      ; <async_fn_cold_attr::S>::method::{closure#0}
+    // CHECK-NEXT: ; Function Attrs: cold
+    #[cold]
+    pub async fn method(&self) {}
+}
+
+// CHECK:      ; async_fn_cold_attr::poll_all{{$}}
+pub fn poll_all(cx: &mut Context<'_>) {
+    let _ = pin!(cold()).poll(cx);
+    let _ = pin!(not_cold()).poll(cx);
+    let _ = pin!(S.method()).poll(cx);
+
+    // CHECK:      ; async_fn_cold_attr::poll_all::{closure#0}{{$}}
+    // CHECK-NEXT: ; Function Attrs: cold
+    // CHECK:      ; async_fn_cold_attr::poll_all::{closure#0}::{closure#0}::<i16>
+    // CHECK-NEXT: ; Function Attrs: cold
+    let closure = #[cold]
+    async || {};
+    let _ = pin!(closure()).poll(cx);
+}

tests/codegen-llvm/async-fn-inline-attr.rs

@@ -0,0 +1,71 @@
+// Checks that inline attributes on `async fn`s and async closures apply to the coroutine holding
+// their body (i.e. the `Future::poll` implementation) rather than to the function or closure that
+// merely constructs it. See #129347.
+//
+//@ edition: 2024
+//@ compile-flags: -Cno-prepopulate-passes -Csymbol-mangling-version=v0 -Zinline-mir=no
+#![crate_type = "lib"]
+#![feature(stmt_expr_attributes)]
+
+use std::future::Future;
+use std::pin::pin;
+use std::task::Context;
+
+// CHECK:      ; async_fn_inline_attr::never{{$}}
+// CHECK-NEXT: ; Function Attrs:
+// CHECK-NOT:  noinline
+// CHECK-SAME: {{$}}
+// CHECK:      ; async_fn_inline_attr::never::{closure#0}
+// CHECK-NEXT: ; Function Attrs: noinline
+#[inline(never)]
+pub async fn never() {}
+
+// CHECK:      ; async_fn_inline_attr::always{{$}}
+// CHECK-NEXT: ; Function Attrs:
+// CHECK-NOT:  alwaysinline
+// CHECK-SAME: {{$}}
+// CHECK:      ; async_fn_inline_attr::always::{closure#0}
+// CHECK-NEXT: ; Function Attrs: alwaysinline
+#[inline(always)]
+pub async fn always() {}
+
+// The body coroutine would get an inline hint regardless, like all closures.
+// CHECK:      ; async_fn_inline_attr::hint{{$}}
+// CHECK-NEXT: ; Function Attrs:
+// CHECK-NOT:  inlinehint
+// CHECK-SAME: {{$}}
+// CHECK:      ; async_fn_inline_attr::hint::{closure#0}
+// CHECK-NEXT: ; Function Attrs: inlinehint
+#[inline]
+pub async fn hint() {}
+
+pub struct S;
+
+impl S {
+    // CHECK:      ; <async_fn_inline_attr::S>::method{{$}}
+    // CHECK-NEXT: ; Function Attrs:
+    // CHECK-NOT:  noinline
+    // CHECK-SAME: {{$}}
+    // CHECK:      ; <async_fn_inline_attr::S>::method::{closure#0}
+    // CHECK-NEXT: ; Function Attrs: noinline
+    #[inline(never)]
+    pub async fn method(&self) {}
+}
+
+// CHECK:      ; async_fn_inline_attr::poll_all{{$}}
+pub fn poll_all(cx: &mut Context<'_>) {
+    let _ = pin!(never()).poll(cx);
+    let _ = pin!(always()).poll(cx);
+    let _ = pin!(hint()).poll(cx);
+    let _ = pin!(S.method()).poll(cx);
+
+    // CHECK:      ; async_fn_inline_attr::poll_all::{closure#0}{{$}}
+    // CHECK-NEXT: ; Function Attrs:
+    // CHECK-NOT:  noinline
+    // CHECK-SAME: {{$}}
+    // CHECK:      ; async_fn_inline_attr::poll_all::{closure#0}::{closure#0}::<i16>
+    // CHECK-NEXT: ; Function Attrs: noinline
+    let closure = #[inline(never)]
+    async || {};
+    let _ = pin!(closure()).poll(cx);
+}

tests/ui/target-feature/inline-always-async-fn.aarch64.stderr

@@ -0,0 +1,10 @@
+error: cannot use `#[inline(always)]` with `#[target_feature]`
+  --> $DIR/inline-always-async-fn.rs:13:1
+   |
+LL | #[inline(always)]
+   | ^^^^^^^^^^^^^^^^^
+   |
+   = note: See this issue for full discussion: https://github.com/rust-lang/rust/issues/145574
+
+error: aborting due to 1 previous error
+

tests/ui/target-feature/inline-always-async-fn.rs

@@ -0,0 +1,29 @@
+// Inlining attributes on an `async fn` apply to its body coroutine instead of the function
+// constructing it (#129347). The body coroutine inherits the function's target features, so
+// `#[inline(always)]` must still be rejected together with `#[target_feature]`.
+//
+//@ revisions: x86_64 aarch64
+//@[x86_64] only-x86_64
+//@[aarch64] only-aarch64
+//@ edition: 2024
+
+#![crate_type = "lib"]
+#![feature(stmt_expr_attributes)]
+
+#[inline(always)]
+//~^ ERROR cannot use `#[inline(always)]` with `#[target_feature]`
+#[cfg_attr(target_arch = "x86_64", target_feature(enable = "sse2"))]
+#[cfg_attr(target_arch = "aarch64", target_feature(enable = "neon"))]
+pub async fn always() {}
+
+#[inline(never)]
+#[cfg_attr(target_arch = "x86_64", target_feature(enable = "sse2"))]
+#[cfg_attr(target_arch = "aarch64", target_feature(enable = "neon"))]
+pub async fn never() {}
+
+#[cfg_attr(target_arch = "x86_64", target_feature(enable = "sse2"))]
+#[cfg_attr(target_arch = "aarch64", target_feature(enable = "neon"))]
+pub fn with_async_closure() {
+    // Just like a regular `#[inline(always)]` closure, this doesn't inherit the target features.
+    let _ = #[inline(always)] async || {};
+}

tests/ui/target-feature/inline-always-async-fn.x86_64.stderr

@@ -0,0 +1,10 @@
+error: cannot use `#[inline(always)]` with `#[target_feature]`
+  --> $DIR/inline-always-async-fn.rs:13:1
+   |
+LL | #[inline(always)]
+   | ^^^^^^^^^^^^^^^^^
+   |
+   = note: See this issue for full discussion: https://github.com/rust-lang/rust/issues/145574
+
+error: aborting due to 1 previous error
+