Auto merge of #128004 - folkertdev:naked-fn-asm, r=Amanieu

codegen `#[naked]` functions using global asm

tracking issue: https://github.com/rust-lang/rust/issues/90957

Fixes #124375

This implements the approach suggested in the tracking issue: use the existing global assembly infrastructure to emit the body of `#[naked]` functions. The main advantage is that we now have full control over what gets generated, and are no longer dependent on LLVM not sneakily messing with our output (inlining, adding extra instructions, etc).

I discussed this approach with `@Amanieu` and while I think the general direction is correct, there is probably a bunch of stuff that needs to change or move around here. I'll leave some inline comments on things that I'm not sure about.

Combined with https://github.com/rust-lang/rust/pull/127853, if both accepted, I think that resolves all steps from the tracking issue.

r? `@Amanieu`

18 files changed, 635 insertions, 106 deletions

diff --git a/compiler/rustc_codegen_gcc/src/asm.rs b/compiler/rustc_codegen_gcc/src/asm.rs
index a1f9eab10e7..415f8affab9 100644
--- a/compiler/rustc_codegen_gcc/src/asm.rs
+++ b/compiler/rustc_codegen_gcc/src/asm.rs
@@ -867,6 +867,13 @@ impl<'gcc, 'tcx> AsmCodegenMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
         template_str.push_str("\n.popsection");
         self.context.add_top_level_asm(None, &template_str);
     }
+
+    fn mangled_name(&self, instance: Instance<'tcx>) -> String {
+        // TODO(@Amanieu): Additional mangling is needed on
+        // some targets to add a leading underscore (Mach-O)
+        // or byte count suffixes (x86 Windows).
+        self.tcx.symbol_name(instance).name.to_string()
+    }
 }
 
 fn modifier_to_gcc(
diff --git a/compiler/rustc_codegen_llvm/src/asm.rs b/compiler/rustc_codegen_llvm/src/asm.rs
index d1804cb49ad..f257abe2a71 100644
--- a/compiler/rustc_codegen_llvm/src/asm.rs
+++ b/compiler/rustc_codegen_llvm/src/asm.rs
@@ -442,6 +442,14 @@ impl<'tcx> AsmCodegenMethods<'tcx> for CodegenCx<'_, 'tcx> {
             );
         }
     }
+
+    fn mangled_name(&self, instance: Instance<'tcx>) -> String {
+        let llval = self.get_fn(instance);
+        llvm::build_string(|s| unsafe {
+            llvm::LLVMRustGetMangledName(llval, s);
+        })
+        .expect("symbol is not valid UTF-8")
+    }
 }
 
 pub(crate) fn inline_asm_call<'ll>(
diff --git a/compiler/rustc_codegen_llvm/src/attributes.rs b/compiler/rustc_codegen_llvm/src/attributes.rs
index cb958c1d4d7..5552a241060 100644
--- a/compiler/rustc_codegen_llvm/src/attributes.rs
+++ b/compiler/rustc_codegen_llvm/src/attributes.rs
@@ -395,17 +395,9 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
         to_add.push(MemoryEffects::None.create_attr(cx.llcx));
     }
     if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
-        to_add.push(AttributeKind::Naked.create_attr(cx.llcx));
-        // HACK(jubilee): "indirect branch tracking" works by attaching prologues to functions.
-        // And it is a module-level attribute, so the alternative is pulling naked functions into
-        // new LLVM modules. Otherwise LLVM's "naked" functions come with endbr prefixes per
-        // https://github.com/rust-lang/rust/issues/98768
-        to_add.push(AttributeKind::NoCfCheck.create_attr(cx.llcx));
-        if llvm_util::get_version() < (19, 0, 0) {
-            // Prior to LLVM 19, branch-target-enforcement was disabled by setting the attribute to
-            // the string "false". Now it is disabled by absence of the attribute.
-            to_add.push(llvm::CreateAttrStringValue(cx.llcx, "branch-target-enforcement", "false"));
-        }
+        // do nothing; a naked function is converted into an extern function
+        // and a global assembly block. LLVM's support for naked functions is
+        // not used.
     } else {
         // Do not set sanitizer attributes for naked functions.
         to_add.extend(sanitize_attrs(cx, codegen_fn_attrs.no_sanitize));
diff --git a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
index 11744eabab0..ab2d24e8d2d 100644
--- a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
+++ b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
@@ -550,6 +550,13 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
         }
     });
 
+    // naked function MUST NOT be inlined! This attribute is required for the rust compiler itself,
+    // but not for the code generation backend because at that point the naked function will just be
+    // a declaration, with a definition provided in global assembly.
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
+        codegen_fn_attrs.inline = InlineAttr::Never;
+    }
+
     codegen_fn_attrs.optimize = attrs.iter().fold(OptimizeAttr::None, |ia, attr| {
         if !attr.has_name(sym::optimize) {
             return ia;
@@ -634,10 +641,6 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
         }
     }
 
-    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
-        codegen_fn_attrs.inline = InlineAttr::Never;
-    }
-
     // Weak lang items have the same semantics as "std internal" symbols in the
     // sense that they're preserved through all our LTO passes and only
     // strippable by the linker.
diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs
index 0cbc5c45736..62f69af3f2f 100644
--- a/compiler/rustc_codegen_ssa/src/mir/mod.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/mod.rs
@@ -20,6 +20,7 @@ mod coverageinfo;
 pub mod debuginfo;
 mod intrinsic;
 mod locals;
+mod naked_asm;
 pub mod operand;
 pub mod place;
 mod rvalue;
@@ -176,6 +177,11 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     let fn_abi = cx.fn_abi_of_instance(instance, ty::List::empty());
     debug!("fn_abi: {:?}", fn_abi);
 
+    if cx.tcx().codegen_fn_attrs(instance.def_id()).flags.contains(CodegenFnAttrFlags::NAKED) {
+        crate::mir::naked_asm::codegen_naked_asm::<Bx>(cx, &mir, instance);
+        return;
+    }
+
     let debug_context = cx.create_function_debug_context(instance, fn_abi, llfn, mir);
 
     let start_llbb = Bx::append_block(cx, llfn, "start");
diff --git a/compiler/rustc_codegen_ssa/src/mir/naked_asm.rs b/compiler/rustc_codegen_ssa/src/mir/naked_asm.rs
new file mode 100644
index 00000000000..c08758a9796
--- /dev/null
+++ b/compiler/rustc_codegen_ssa/src/mir/naked_asm.rs
@@ -0,0 +1,266 @@
+use rustc_attr::InstructionSetAttr;
+use rustc_middle::mir::mono::{Linkage, MonoItem, MonoItemData, Visibility};
+use rustc_middle::mir::{Body, InlineAsmOperand};
+use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutOf};
+use rustc_middle::ty::{Instance, TyCtxt};
+use rustc_middle::{bug, ty};
+use rustc_span::sym;
+
+use crate::common;
+use crate::traits::{AsmCodegenMethods, BuilderMethods, GlobalAsmOperandRef, MiscCodegenMethods};
+
+pub(crate) fn codegen_naked_asm<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
+    cx: &'a Bx::CodegenCx,
+    mir: &Body<'tcx>,
+    instance: Instance<'tcx>,
+) {
+    let rustc_middle::mir::TerminatorKind::InlineAsm {
+        asm_macro: _,
+        template,
+        ref operands,
+        options,
+        line_spans,
+        targets: _,
+        unwind: _,
+    } = mir.basic_blocks.iter().next().unwrap().terminator().kind
+    else {
+        bug!("#[naked] functions should always terminate with an asm! block")
+    };
+
+    let operands: Vec<_> =
+        operands.iter().map(|op| inline_to_global_operand::<Bx>(cx, instance, op)).collect();
+
+    let item_data = cx.codegen_unit().items().get(&MonoItem::Fn(instance)).unwrap();
+    let name = cx.mangled_name(instance);
+    let (begin, end) = prefix_and_suffix(cx.tcx(), instance, &name, item_data);
+
+    let mut template_vec = Vec::new();
+    template_vec.push(rustc_ast::ast::InlineAsmTemplatePiece::String(begin.into()));
+    template_vec.extend(template.iter().cloned());
+    template_vec.push(rustc_ast::ast::InlineAsmTemplatePiece::String(end.into()));
+
+    cx.codegen_global_asm(&template_vec, &operands, options, line_spans);
+}
+
+fn inline_to_global_operand<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
+    cx: &'a Bx::CodegenCx,
+    instance: Instance<'tcx>,
+    op: &InlineAsmOperand<'tcx>,
+) -> GlobalAsmOperandRef<'tcx> {
+    match op {
+        InlineAsmOperand::Const { value } => {
+            let const_value = instance
+                .instantiate_mir_and_normalize_erasing_regions(
+                    cx.tcx(),
+                    cx.typing_env(),
+                    ty::EarlyBinder::bind(value.const_),
+                )
+                .eval(cx.tcx(), cx.typing_env(), value.span)
+                .expect("erroneous constant missed by mono item collection");
+
+            let mono_type = instance.instantiate_mir_and_normalize_erasing_regions(
+                cx.tcx(),
+                cx.typing_env(),
+                ty::EarlyBinder::bind(value.ty()),
+            );
+
+            let string = common::asm_const_to_str(
+                cx.tcx(),
+                value.span,
+                const_value,
+                cx.layout_of(mono_type),
+            );
+
+            GlobalAsmOperandRef::Const { string }
+        }
+        InlineAsmOperand::SymFn { value } => {
+            let mono_type = instance.instantiate_mir_and_normalize_erasing_regions(
+                cx.tcx(),
+                cx.typing_env(),
+                ty::EarlyBinder::bind(value.ty()),
+            );
+
+            let instance = match mono_type.kind() {
+                &ty::FnDef(def_id, args) => Instance::new(def_id, args),
+                _ => bug!("asm sym is not a function"),
+            };
+
+            GlobalAsmOperandRef::SymFn { instance }
+        }
+        InlineAsmOperand::SymStatic { def_id } => {
+            GlobalAsmOperandRef::SymStatic { def_id: *def_id }
+        }
+        InlineAsmOperand::In { .. }
+        | InlineAsmOperand::Out { .. }
+        | InlineAsmOperand::InOut { .. }
+        | InlineAsmOperand::Label { .. } => {
+            bug!("invalid operand type for naked_asm!")
+        }
+    }
+}
+
+enum AsmBinaryFormat {
+    Elf,
+    Macho,
+    Coff,
+}
+
+impl AsmBinaryFormat {
+    fn from_target(target: &rustc_target::spec::Target) -> Self {
+        if target.is_like_windows {
+            Self::Coff
+        } else if target.is_like_osx {
+            Self::Macho
+        } else {
+            Self::Elf
+        }
+    }
+}
+
+fn prefix_and_suffix<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+    asm_name: &str,
+    item_data: &MonoItemData,
+) -> (String, String) {
+    use std::fmt::Write;
+
+    let asm_binary_format = AsmBinaryFormat::from_target(&tcx.sess.target);
+
+    let is_arm = tcx.sess.target.arch == "arm";
+    let is_thumb = tcx.sess.unstable_target_features.contains(&sym::thumb_mode);
+
+    let attrs = tcx.codegen_fn_attrs(instance.def_id());
+    let link_section = attrs.link_section.map(|symbol| symbol.as_str().to_string());
+    let align = attrs.alignment.map(|a| a.bytes()).unwrap_or(4);
+
+    // See https://sourceware.org/binutils/docs/as/ARM-Directives.html for info on these directives.
+    // In particular, `.arm` can also be written `.code 32` and `.thumb` as `.code 16`.
+    let (arch_prefix, arch_suffix) = if is_arm {
+        (
+            match attrs.instruction_set {
+                None => match is_thumb {
+                    true => ".thumb\n.thumb_func",
+                    false => ".arm",
+                },
+                Some(InstructionSetAttr::ArmT32) => ".thumb\n.thumb_func",
+                Some(InstructionSetAttr::ArmA32) => ".arm",
+            },
+            match is_thumb {
+                true => ".thumb",
+                false => ".arm",
+            },
+        )
+    } else {
+        ("", "")
+    };
+
+    let emit_fatal = |msg| tcx.dcx().span_fatal(tcx.def_span(instance.def_id()), msg);
+
+    // see https://godbolt.org/z/cPK4sxKor.
+    let write_linkage = |w: &mut String| -> std::fmt::Result {
+        match item_data.linkage {
+            Linkage::External => {
+                writeln!(w, ".globl {asm_name}")?;
+            }
+            Linkage::LinkOnceAny | Linkage::LinkOnceODR | Linkage::WeakAny | Linkage::WeakODR => {
+                match asm_binary_format {
+                    AsmBinaryFormat::Elf | AsmBinaryFormat::Coff => {
+                        writeln!(w, ".weak {asm_name}")?;
+                    }
+                    AsmBinaryFormat::Macho => {
+                        writeln!(w, ".globl {asm_name}")?;
+                        writeln!(w, ".weak_definition {asm_name}")?;
+                    }
+                }
+            }
+            Linkage::Internal | Linkage::Private => {
+                // write nothing
+            }
+            Linkage::Appending => emit_fatal("Only global variables can have appending linkage!"),
+            Linkage::Common => emit_fatal("Functions may not have common linkage"),
+            Linkage::AvailableExternally => {
+                // this would make the function equal an extern definition
+                emit_fatal("Functions may not have available_externally linkage")
+            }
+            Linkage::ExternalWeak => {
+                // FIXME: actually this causes a SIGILL in LLVM
+                emit_fatal("Functions may not have external weak linkage")
+            }
+        }
+
+        Ok(())
+    };
+
+    let mut begin = String::new();
+    let mut end = String::new();
+    match asm_binary_format {
+        AsmBinaryFormat::Elf => {
+            let section = link_section.unwrap_or(format!(".text.{asm_name}"));
+
+            let progbits = match is_arm {
+                true => "%progbits",
+                false => "@progbits",
+            };
+
+            let function = match is_arm {
+                true => "%function",
+                false => "@function",
+            };
+
+            writeln!(begin, ".pushsection {section},\"ax\", {progbits}").unwrap();
+            writeln!(begin, ".balign {align}").unwrap();
+            write_linkage(&mut begin).unwrap();
+            if let Visibility::Hidden = item_data.visibility {
+                writeln!(begin, ".hidden {asm_name}").unwrap();
+            }
+            writeln!(begin, ".type {asm_name}, {function}").unwrap();
+            if !arch_prefix.is_empty() {
+                writeln!(begin, "{}", arch_prefix).unwrap();
+            }
+            writeln!(begin, "{asm_name}:").unwrap();
+
+            writeln!(end).unwrap();
+            writeln!(end, ".size {asm_name}, . - {asm_name}").unwrap();
+            writeln!(end, ".popsection").unwrap();
+            if !arch_suffix.is_empty() {
+                writeln!(end, "{}", arch_suffix).unwrap();
+            }
+        }
+        AsmBinaryFormat::Macho => {
+            let section = link_section.unwrap_or("__TEXT,__text".to_string());
+            writeln!(begin, ".pushsection {},regular,pure_instructions", section).unwrap();
+            writeln!(begin, ".balign {align}").unwrap();
+            write_linkage(&mut begin).unwrap();
+            if let Visibility::Hidden = item_data.visibility {
+                writeln!(begin, ".private_extern {asm_name}").unwrap();
+            }
+            writeln!(begin, "{asm_name}:").unwrap();
+
+            writeln!(end).unwrap();
+            writeln!(end, ".popsection").unwrap();
+            if !arch_suffix.is_empty() {
+                writeln!(end, "{}", arch_suffix).unwrap();
+            }
+        }
+        AsmBinaryFormat::Coff => {
+            let section = link_section.unwrap_or(format!(".text.{asm_name}"));
+            writeln!(begin, ".pushsection {},\"xr\"", section).unwrap();
+            writeln!(begin, ".balign {align}").unwrap();
+            write_linkage(&mut begin).unwrap();
+            writeln!(begin, ".def {asm_name}").unwrap();
+            writeln!(begin, ".scl 2").unwrap();
+            writeln!(begin, ".type 32").unwrap();
+            writeln!(begin, ".endef {asm_name}").unwrap();
+            writeln!(begin, "{asm_name}:").unwrap();
+
+            writeln!(end).unwrap();
+            writeln!(end, ".popsection").unwrap();
+            if !arch_suffix.is_empty() {
+                writeln!(end, "{}", arch_suffix).unwrap();
+            }
+        }
+    }
+
+    (begin, end)
+}
diff --git a/compiler/rustc_codegen_ssa/src/mono_item.rs b/compiler/rustc_codegen_ssa/src/mono_item.rs
index 038c5857fac..6749bc63327 100644
--- a/compiler/rustc_codegen_ssa/src/mono_item.rs
+++ b/compiler/rustc_codegen_ssa/src/mono_item.rs
@@ -1,4 +1,5 @@
 use rustc_hir as hir;
+use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::interpret::ErrorHandled;
 use rustc_middle::mir::mono::{Linkage, MonoItem, Visibility};
 use rustc_middle::ty::Instance;
@@ -135,7 +136,13 @@ impl<'a, 'tcx: 'a> MonoItemExt<'a, 'tcx> for MonoItem<'tcx> {
                 cx.predefine_static(def_id, linkage, visibility, symbol_name);
             }
             MonoItem::Fn(instance) => {
-                cx.predefine_fn(instance, linkage, visibility, symbol_name);
+                let attrs = cx.tcx().codegen_fn_attrs(instance.def_id());
+
+                if attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
+                    // do not define this function; it will become a global assembly block
+                } else {
+                    cx.predefine_fn(instance, linkage, visibility, symbol_name);
+                };
             }
             MonoItem::GlobalAsm(..) => {}
         }
diff --git a/compiler/rustc_codegen_ssa/src/traits/asm.rs b/compiler/rustc_codegen_ssa/src/traits/asm.rs
index f4853da1156..7767bffbfbf 100644
--- a/compiler/rustc_codegen_ssa/src/traits/asm.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/asm.rs
@@ -68,4 +68,11 @@ pub trait AsmCodegenMethods<'tcx> {
         options: InlineAsmOptions,
         line_spans: &[Span],
     );
+
+    /// The mangled name of this instance
+    ///
+    /// Additional mangling is used on
+    /// some targets to add a leading underscore (Mach-O)
+    /// or byte count suffixes (x86 Windows).
+    fn mangled_name(&self, instance: Instance<'tcx>) -> String;
 }
diff --git a/compiler/rustc_middle/src/mir/mono.rs b/compiler/rustc_middle/src/mir/mono.rs
index 161716610fe..1f50b67cb50 100644
--- a/compiler/rustc_middle/src/mir/mono.rs
+++ b/compiler/rustc_middle/src/mir/mono.rs
@@ -19,6 +19,7 @@ use rustc_target::spec::SymbolVisibility;
 use tracing::debug;
 
 use crate::dep_graph::{DepNode, WorkProduct, WorkProductId};
+use crate::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use crate::ty::{GenericArgs, Instance, InstanceKind, SymbolName, TyCtxt};
 
 /// Describes how a monomorphization will be instantiated in object files.
@@ -104,7 +105,9 @@ impl<'tcx> MonoItem<'tcx> {
                 let entry_def_id = tcx.entry_fn(()).map(|(id, _)| id);
                 // If this function isn't inlined or otherwise has an extern
                 // indicator, then we'll be creating a globally shared version.
-                if tcx.codegen_fn_attrs(instance.def_id()).contains_extern_indicator()
+                let codegen_fn_attrs = tcx.codegen_fn_attrs(instance.def_id());
+                if codegen_fn_attrs.contains_extern_indicator()
+                    || codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED)
                     || !instance.def.generates_cgu_internal_copy(tcx)
                     || Some(instance.def_id()) == entry_def_id
                 {
diff --git a/tests/codegen/naked-asan.rs b/tests/codegen/naked-asan.rs
index bcaa60baeff..8efedab6ea5 100644
--- a/tests/codegen/naked-asan.rs
+++ b/tests/codegen/naked-asan.rs
@@ -8,14 +8,15 @@
 #![no_std]
 #![feature(abi_x86_interrupt, naked_functions)]
 
-// CHECK: define x86_intrcc void @page_fault_handler(ptr {{.*}}%0, i64 {{.*}}%1){{.*}}#[[ATTRS:[0-9]+]] {
-// CHECK-NOT: memcpy
+pub fn caller() {
+    page_fault_handler(1, 2);
+}
+
+// CHECK: declare x86_intrcc void @page_fault_handler(ptr {{.*}}, i64{{.*}}){{.*}}#[[ATTRS:[0-9]+]]
 #[naked]
 #[no_mangle]
 pub extern "x86-interrupt" fn page_fault_handler(_: u64, _: u64) {
-    unsafe {
-        core::arch::naked_asm!("ud2");
-    }
+    unsafe { core::arch::naked_asm!("ud2") };
 }
 
 // CHECK: #[[ATTRS]] =
diff --git a/tests/codegen/naked-fn/aligned.rs b/tests/codegen/naked-fn/aligned.rs
index 3bbd67981e5..d9dcd7f6c3e 100644
--- a/tests/codegen/naked-fn/aligned.rs
+++ b/tests/codegen/naked-fn/aligned.rs
@@ -6,15 +6,12 @@
 #![feature(naked_functions, fn_align)]
 use std::arch::naked_asm;
 
-// CHECK: Function Attrs: naked
-// CHECK-NEXT: define{{.*}}void @naked_empty()
-// CHECK: align 16
+// CHECK: .balign 16
+// CHECK-LABEL: naked_empty:
 #[repr(align(16))]
 #[no_mangle]
 #[naked]
 pub unsafe extern "C" fn naked_empty() {
-    // CHECK-NEXT: start:
-    // CHECK-NEXT: call void asm
-    // CHECK-NEXT: unreachable
+    // CHECK: ret
     naked_asm!("ret");
 }
diff --git a/tests/codegen/naked-fn/generics.rs b/tests/codegen/naked-fn/generics.rs
new file mode 100644
index 00000000000..a33d213617a
--- /dev/null
+++ b/tests/codegen/naked-fn/generics.rs
@@ -0,0 +1,116 @@
+//@ compile-flags: -O
+//@ only-x86_64
+
+#![crate_type = "lib"]
+#![feature(naked_functions, asm_const)]
+
+use std::arch::naked_asm;
+
+#[no_mangle]
+fn test(x: u64) {
+    // just making sure these symbols get used
+    using_const_generics::<1>(x);
+    using_const_generics::<2>(x);
+
+    generic_function::<i64>(x as i64);
+
+    let foo = Foo(x);
+
+    foo.method();
+    foo.trait_method();
+}
+
+// CHECK: .balign 4
+// CHECK: add rax, 2
+// CHECK: add rax, 42
+
+// CHECK: .balign 4
+// CHECK: add rax, 1
+// CHECK: add rax, 42
+
+#[naked]
+pub extern "C" fn using_const_generics<const N: u64>(x: u64) -> u64 {
+    const M: u64 = 42;
+
+    unsafe {
+        naked_asm!(
+            "xor rax, rax",
+            "add rax, rdi",
+            "add rax, {}",
+            "add rax, {}",
+            "ret",
+            const N,
+            const M,
+        )
+    }
+}
+
+trait Invert {
+    fn invert(self) -> Self;
+}
+
+impl Invert for i64 {
+    fn invert(self) -> Self {
+        -1 * self
+    }
+}
+
+// CHECK: .balign 4
+// CHECK-LABEL: generic_function:
+// CHECK: call
+// CHECK: ret
+
+#[naked]
+#[no_mangle]
+pub extern "C" fn generic_function<T: Invert>(x: i64) -> i64 {
+    unsafe {
+        naked_asm!(
+            "call {}",
+            "ret",
+            sym <T as Invert>::invert,
+        )
+    }
+}
+
+#[derive(Copy, Clone)]
+#[repr(transparent)]
+struct Foo(u64);
+
+// CHECK: .balign 4
+// CHECK-LABEL: method:
+// CHECK: mov rax, rdi
+
+impl Foo {
+    #[naked]
+    #[no_mangle]
+    extern "C" fn method(self) -> u64 {
+        unsafe { naked_asm!("mov rax, rdi", "ret") }
+    }
+}
+
+// CHECK: .balign 4
+// CHECK-LABEL: trait_method:
+// CHECK: mov rax, rdi
+
+trait Bar {
+    extern "C" fn trait_method(self) -> u64;
+}
+
+impl Bar for Foo {
+    #[naked]
+    #[no_mangle]
+    extern "C" fn trait_method(self) -> u64 {
+        unsafe { naked_asm!("mov rax, rdi", "ret") }
+    }
+}
+
+// CHECK: .balign 4
+// CHECK-LABEL: naked_with_args_and_return:
+// CHECK: lea rax, [rdi + rsi]
+
+// this previously ICE'd, see https://github.com/rust-lang/rust/issues/124375
+#[naked]
+#[no_mangle]
+pub unsafe extern "C" fn naked_with_args_and_return(a: isize, b: isize) -> isize {
+    naked_asm!("lea rax, [rdi + rsi]", "ret");
+}
diff --git a/tests/codegen/naked-fn/instruction-set.rs b/tests/codegen/naked-fn/instruction-set.rs
new file mode 100644
index 00000000000..5b790b2034c
--- /dev/null
+++ b/tests/codegen/naked-fn/instruction-set.rs
@@ -0,0 +1,59 @@
+//@ revisions: arm-mode thumb-mode
+//@ [arm-mode] compile-flags: --target armv5te-none-eabi
+//@ [thumb-mode] compile-flags: --target thumbv5te-none-eabi
+//@ [arm-mode] needs-llvm-components: arm
+//@ [thumb-mode] needs-llvm-components: arm
+
+#![crate_type = "lib"]
+#![feature(no_core, lang_items, rustc_attrs, naked_functions)]
+#![no_core]
+
+#[rustc_builtin_macro]
+macro_rules! naked_asm {
+    () => {};
+}
+
+#[lang = "sized"]
+trait Sized {}
+#[lang = "copy"]
+trait Copy {}
+
+// arm-mode: .arm
+// thumb-mode: .thumb
+// CHECK-LABEL: test_unspecified:
+// CHECK: bx lr
+// CHECK: .popsection
+// arm-mode: .arm
+// thumb-mode: .thumb
+#[no_mangle]
+#[naked]
+unsafe extern "C" fn test_unspecified() {
+    naked_asm!("bx lr");
+}
+
+// CHECK: .thumb
+// CHECK: .thumb_func
+// CHECK-LABEL: test_thumb:
+// CHECK: bx lr
+// CHECK: .popsection
+// arm-mode: .arm
+// thumb-mode: .thumb
+#[no_mangle]
+#[naked]
+#[instruction_set(arm::t32)]
+unsafe extern "C" fn test_thumb() {
+    naked_asm!("bx lr");
+}
+
+// CHECK: .arm
+// CHECK-LABEL: test_arm:
+// CHECK: bx lr
+// CHECK: .popsection
+// arm-mode: .arm
+// thumb-mode: .thumb
+#[no_mangle]
+#[naked]
+#[instruction_set(arm::a32)]
+unsafe extern "C" fn test_arm() {
+    naked_asm!("bx lr");
+}
diff --git a/tests/codegen/naked-fn/naked-functions.rs b/tests/codegen/naked-fn/naked-functions.rs
index 3f7447af599..f505d27d48c 100644
--- a/tests/codegen/naked-fn/naked-functions.rs
+++ b/tests/codegen/naked-fn/naked-functions.rs
@@ -1,29 +1,147 @@
-//@ compile-flags: -C no-prepopulate-passes -Copt-level=0
-//@ needs-asm-support
-//@ only-x86_64
+//@ revisions: linux win macos thumb
+//
+//@[linux] compile-flags: --target x86_64-unknown-linux-gnu
+//@[linux] needs-llvm-components: x86
+//@[win] compile-flags: --target x86_64-pc-windows-gnu
+//@[win] needs-llvm-components: x86
+//@[macos] compile-flags: --target aarch64-apple-darwin
+//@[macos] needs-llvm-components: arm
+//@[thumb] compile-flags: --target thumbv7em-none-eabi
+//@[thumb] needs-llvm-components: arm
 
 #![crate_type = "lib"]
-#![feature(naked_functions)]
-use std::arch::naked_asm;
+#![feature(no_core, lang_items, rustc_attrs, naked_functions)]
+#![no_core]
+
+#[rustc_builtin_macro]
+macro_rules! naked_asm {
+    () => {};
+}
+
+#[lang = "sized"]
+trait Sized {}
+#[lang = "copy"]
+trait Copy {}
+
+// linux,win: .intel_syntax
+//
+// linux:   .pushsection .text.naked_empty,\22ax\22, @progbits
+// macos:   .pushsection __TEXT,__text,regular,pure_instructions
+// win: .pushsection .text.naked_empty,\22xr\22
+// thumb:   .pushsection .text.naked_empty,\22ax\22, %progbits
+//
+// CHECK: .balign 4
+//
+// linux,win,thumb: .globl naked_empty
+// macos: .globl _naked_empty
+//
+// CHECK-NOT: .private_extern
+// CHECK-NOT: .hidden
+//
+// linux: .type naked_empty, @function
+//
+// win: .def naked_empty
+// win: .scl 2
+// win: .type 32
+// win: .endef naked_empty
+//
+// thumb: .type naked_empty, %function
+// thumb: .thumb
+// thumb: .thumb_func
+//
+// CHECK-LABEL: naked_empty:
+//
+// linux,macos,win: ret
+// thumb: bx lr
+//
+// CHECK: .popsection
+//
+// thumb: .thumb
+//
+// linux,win: .att_syntax
 
-// CHECK: Function Attrs: naked
-// CHECK-NEXT: define{{.*}}void @naked_empty()
 #[no_mangle]
 #[naked]
 pub unsafe extern "C" fn naked_empty() {
-    // CHECK-NEXT: {{.+}}:
-    // CHECK-NEXT: call void asm
-    // CHECK-NEXT: unreachable
+    #[cfg(not(all(target_arch = "arm", target_feature = "thumb-mode")))]
     naked_asm!("ret");
+
+    #[cfg(all(target_arch = "arm", target_feature = "thumb-mode"))]
+    naked_asm!("bx lr");
 }
 
-// CHECK: Function Attrs: naked
-// CHECK-NEXT: define{{.*}}i{{[0-9]+}} @naked_with_args_and_return(i64 %0, i64 %1)
+// linux,win: .intel_syntax
+//
+// linux:   .pushsection .text.naked_with_args_and_return,\22ax\22, @progbits
+// macos:   .pushsection __TEXT,__text,regular,pure_instructions
+// win: .pushsection .text.naked_with_args_and_return,\22xr\22
+// thumb:   .pushsection .text.naked_with_args_and_return,\22ax\22, %progbits
+//
+// CHECK: .balign 4
+//
+// linux,win,thumb: .globl naked_with_args_and_return
+// macos: .globl _naked_with_args_and_return
+//
+// CHECK-NOT: .private_extern
+// CHECK-NOT: .hidden
+//
+// linux: .type naked_with_args_and_return, @function
+//
+// win: .def naked_with_args_and_return
+// win: .scl 2
+// win: .type 32
+// win: .endef naked_with_args_and_return
+//
+// thumb: .type naked_with_args_and_return, %function
+// thumb: .thumb
+// thumb: .thumb_func
+//
+// CHECK-LABEL: naked_with_args_and_return:
+//
+// linux, win: lea rax, [rdi + rsi]
+// macos: add x0, x0, x1
+// thumb: adds r0, r0, r1
+//
+// linux,macos,win: ret
+// thumb: bx lr
+//
+// CHECK: .popsection
+//
+// thumb: .thumb
+//
+// linux,win: .att_syntax
+
 #[no_mangle]
 #[naked]
 pub unsafe extern "C" fn naked_with_args_and_return(a: isize, b: isize) -> isize {
-    // CHECK-NEXT: {{.+}}:
-    // CHECK-NEXT: call void asm
-    // CHECK-NEXT: unreachable
-    naked_asm!("lea rax, [rdi + rsi]", "ret");
+    #[cfg(any(target_os = "windows", target_os = "linux"))]
+    {
+        naked_asm!("lea rax, [rdi + rsi]", "ret")
+    }
+
+    #[cfg(target_os = "macos")]
+    {
+        naked_asm!("add x0, x0, x1", "ret")
+    }
+
+    #[cfg(all(target_arch = "arm", target_feature = "thumb-mode"))]
+    {
+        naked_asm!("adds r0, r0, r1", "bx lr")
+    }
+}
+
+// linux:   .pushsection .text.some_different_name,\22ax\22, @progbits
+// macos:   .pushsection .text.some_different_name,regular,pure_instructions
+// win: .pushsection .text.some_different_name,\22xr\22
+// thumb:   .pushsection .text.some_different_name,\22ax\22, %progbits
+// CHECK-LABEL: test_link_section:
+#[no_mangle]
+#[naked]
+#[link_section = ".text.some_different_name"]
+pub unsafe extern "C" fn test_link_section() {
+    #[cfg(not(all(target_arch = "arm", target_feature = "thumb-mode")))]
+    naked_asm!("ret");
+
+    #[cfg(all(target_arch = "arm", target_feature = "thumb-mode"))]
+    naked_asm!("bx lr");
 }
diff --git a/tests/codegen/naked-fn/naked-nocoverage.rs b/tests/codegen/naked-fn/naked-nocoverage.rs
deleted file mode 100644
index f63661bcd3a..00000000000
--- a/tests/codegen/naked-fn/naked-nocoverage.rs
+++ /dev/null
@@ -1,19 +0,0 @@
-// Checks that naked functions are not instrumented by -Cinstrument-coverage.
-// Regression test for issue #105170.
-//
-//@ needs-asm-support
-//@ compile-flags: -Zno-profiler-runtime
-//@ compile-flags: -Cinstrument-coverage
-#![crate_type = "lib"]
-#![feature(naked_functions)]
-use std::arch::naked_asm;
-
-#[naked]
-#[no_mangle]
-pub unsafe extern "C" fn f() {
-    // CHECK:       define {{(dso_local )?}}void @f()
-    // CHECK-NEXT:  start:
-    // CHECK-NEXT:    call void asm
-    // CHECK-NEXT:    unreachable
-    naked_asm!("");
-}
diff --git a/tests/codegen/naked-fn/naked-noinline.rs b/tests/codegen/naked-fn/naked-noinline.rs
deleted file mode 100644
index 6ea36d96783..00000000000
--- a/tests/codegen/naked-fn/naked-noinline.rs
+++ /dev/null
@@ -1,31 +0,0 @@
-// Checks that naked functions are never inlined.
-//@ compile-flags: -O -Zmir-opt-level=3
-//@ needs-asm-support
-//@ ignore-wasm32
-#![crate_type = "lib"]
-#![feature(naked_functions)]
-
-use std::arch::naked_asm;
-
-#[naked]
-#[no_mangle]
-pub unsafe extern "C" fn f() {
-    // Check that f has naked and noinline attributes.
-    //
-    // CHECK:       define {{(dso_local )?}}void @f() unnamed_addr [[ATTR:#[0-9]+]]
-    // CHECK-NEXT:  start:
-    // CHECK-NEXT:    call void asm
-    naked_asm!("");
-}
-
-#[no_mangle]
-pub unsafe fn g() {
-    // Check that call to f is not inlined.
-    //
-    // CHECK-LABEL: define {{(dso_local )?}}void @g()
-    // CHECK-NEXT:  start:
-    // CHECK-NEXT:    call void @f()
-    f();
-}
-
-// CHECK: attributes [[ATTR]] = { naked{{.*}}noinline{{.*}} }
diff --git a/tests/crashes/124375.rs b/tests/crashes/124375.rs
deleted file mode 100644
index 1d877caeb8b..00000000000
--- a/tests/crashes/124375.rs
+++ /dev/null
@@ -1,11 +0,0 @@
-//@ known-bug: #124375
-//@ compile-flags: -Zmir-opt-level=0
-//@ only-x86_64
-#![crate_type = "lib"]
-#![feature(naked_functions)]
-use std::arch::naked_asm;
-
-#[naked]
-pub unsafe extern "C" fn naked_with_args_and_return(a: isize, b: isize) -> isize {
-    naked_asm!("lea rax, [rdi + rsi]", "ret");
-}
diff --git a/tests/ui/asm/naked-functions-instruction-set.rs b/tests/ui/asm/naked-functions-instruction-set.rs
index 37c7b52c191..3a6e7a46ce5 100644
--- a/tests/ui/asm/naked-functions-instruction-set.rs
+++ b/tests/ui/asm/naked-functions-instruction-set.rs
@@ -24,7 +24,7 @@ unsafe extern "C" fn test_thumb() {
 
 #[no_mangle]
 #[naked]
-#[instruction_set(arm::t32)]
+#[instruction_set(arm::a32)]
 unsafe extern "C" fn test_arm() {
     naked_asm!("bx lr");
 }