Skip to content

Commit

Permalink
Rollup merge of rust-lang#128004 - folkertdev:naked-fn-asm, r=Amanieu
Browse files Browse the repository at this point in the history
codegen `#[naked]` functions using global asm

tracking issue: rust-lang#90957

Fixes rust-lang#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 rust-lang#127853, if both accepted, I think that resolves all steps from the tracking issue.

r? `@Amanieu`

try-job: x86_64-gnu-nopt
  • Loading branch information
jieyouxu authored Dec 9, 2024
2 parents f6cb952 + 9eb35a0 commit 4e3d98c
Show file tree
Hide file tree
Showing 20 changed files with 641 additions and 107 deletions.
11 changes: 10 additions & 1 deletion compiler/rustc_attr/src/builtin.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,12 +49,21 @@ pub enum InlineAttr {
Never,
}

#[derive(Clone, Encodable, Decodable, Debug, PartialEq, Eq, HashStable_Generic)]
#[derive(Copy, Clone, Encodable, Decodable, Debug, PartialEq, Eq, HashStable_Generic)]
pub enum InstructionSetAttr {
ArmA32,
ArmT32,
}

impl InstructionSetAttr {
pub fn as_str(self) -> &'static str {
match self {
Self::ArmA32 => sym::a32.as_str(),
Self::ArmT32 => sym::t32.as_str(),
}
}
}

#[derive(Clone, Encodable, Decodable, Debug, HashStable_Generic)]
pub enum OptimizeAttr {
None,
Expand Down
7 changes: 7 additions & 0 deletions compiler/rustc_codegen_gcc/src/asm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
8 changes: 8 additions & 0 deletions compiler/rustc_codegen_llvm/src/asm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>(
Expand Down
14 changes: 3 additions & 11 deletions compiler/rustc_codegen_llvm/src/attributes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down
11 changes: 7 additions & 4 deletions compiler/rustc_codegen_ssa/src/codegen_attrs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -542,6 +542,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;
Expand Down Expand Up @@ -626,10 +633,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.
Expand Down
6 changes: 6 additions & 0 deletions compiler/rustc_codegen_ssa/src/mir/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ mod coverageinfo;
pub mod debuginfo;
mod intrinsic;
mod locals;
mod naked_asm;
pub mod operand;
pub mod place;
mod rvalue;
Expand Down Expand Up @@ -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");
Expand Down
257 changes: 257 additions & 0 deletions compiler/rustc_codegen_ssa/src/mir/naked_asm.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,257 @@
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 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.
// None means the default, which corresponds to internal linkage
let linkage = match item_data.linkage {
Linkage::External => Some(".globl"),
Linkage::LinkOnceAny => Some(".weak"),
Linkage::LinkOnceODR => Some(".weak"),
Linkage::WeakAny => Some(".weak"),
Linkage::WeakODR => Some(".weak"),
Linkage::Internal => None,
Linkage::Private => None,
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")
}
};

let mut begin = String::new();
let mut end = String::new();
match AsmBinaryFormat::from_target(&tcx.sess.target) {
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();
if let Some(linkage) = linkage {
writeln!(begin, "{linkage} {asm_name}").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();
if let Some(linkage) = linkage {
writeln!(begin, "{linkage} {asm_name}").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();
if let Some(linkage) = linkage {
writeln!(begin, "{linkage} {asm_name}").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)
}
9 changes: 8 additions & 1 deletion compiler/rustc_codegen_ssa/src/mono_item.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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(..) => {}
}
Expand Down
Loading

0 comments on commit 4e3d98c

Please sign in to comment.