Skip to content

Commit c637bda

Browse files
folkertdevfolkertdev
committed
support wasm inline assembly in naked functions
1 parent bcd0683 commit c637bda

File tree

2 files changed

+335
-4
lines changed

2 files changed

+335
-4
lines changed

‎compiler/rustc_codegen_ssa/src/mir/naked_asm.rs‎

Lines changed: 136 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
1+
use rustc_abi::{BackendRepr, Float, Integer, Primitive, RegKind};
12
use rustc_attr_parsing::InstructionSetAttr;
23
use rustc_middle::mir::mono::{Linkage, MonoItem, MonoItemData, Visibility};
34
use rustc_middle::mir::{Body, InlineAsmOperand};
4-
use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutOf};
5-
use rustc_middle::ty::{Instance, TyCtxt};
5+
use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, LayoutOf};
6+
use rustc_middle::ty::{Instance, Ty, TyCtxt};
67
use rustc_middle::{bug, ty};
78
use rustc_span::sym;
9+
use rustc_target::callconv::{ArgAbi, FnAbi, PassMode};
810

911
use crate::common;
1012
use crate::traits::{AsmCodegenMethods, BuilderMethods, GlobalAsmOperandRef, MiscCodegenMethods};
@@ -32,7 +34,8 @@ pub(crate) fn codegen_naked_asm<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
3234

3335
let item_data = cx.codegen_unit().items().get(&MonoItem::Fn(instance)).unwrap();
3436
let name = cx.mangled_name(instance);
35-
let (begin, end) = prefix_and_suffix(cx.tcx(), instance, &name, item_data);
37+
let fn_abi = cx.fn_abi_of_instance(instance, ty::List::empty());
38+
let (begin, end) = prefix_and_suffix(cx.tcx(), instance, &name, item_data, fn_abi);
3639

3740
let mut template_vec = Vec::new();
3841
template_vec.push(rustc_ast::ast::InlineAsmTemplatePiece::String(begin.into()));
@@ -103,6 +106,7 @@ enum AsmBinaryFormat {
103106
Elf,
104107
Macho,
105108
Coff,
109+
Wasm,
106110
}
107111

108112
impl AsmBinaryFormat {
@@ -111,6 +115,8 @@ impl AsmBinaryFormat {
111115
Self::Coff
112116
} else if target.is_like_osx {
113117
Self::Macho
118+
} else if target.is_like_wasm {
119+
Self::Wasm
114120
} else {
115121
Self::Elf
116122
}
@@ -122,6 +128,7 @@ fn prefix_and_suffix<'tcx>(
122128
instance: Instance<'tcx>,
123129
asm_name: &str,
124130
item_data: &MonoItemData,
131+
fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
125132
) -> (String, String) {
126133
use std::fmt::Write;
127134

@@ -169,7 +176,7 @@ fn prefix_and_suffix<'tcx>(
169176
}
170177
Linkage::LinkOnceAny | Linkage::LinkOnceODR | Linkage::WeakAny | Linkage::WeakODR => {
171178
match asm_binary_format {
172-
AsmBinaryFormat::Elf | AsmBinaryFormat::Coff => {
179+
AsmBinaryFormat::Elf | AsmBinaryFormat::Coff | AsmBinaryFormat::Wasm => {
173180
writeln!(w, ".weak {asm_name}")?;
174181
}
175182
AsmBinaryFormat::Macho => {
@@ -264,7 +271,132 @@ fn prefix_and_suffix<'tcx>(
264271
writeln!(end, "{}", arch_suffix).unwrap();
265272
}
266273
}
274+
AsmBinaryFormat::Wasm => {
275+
let section = link_section.unwrap_or(format!(".text.{asm_name}"));
276+
277+
writeln!(begin, ".section {section},\"\",@").unwrap();
278+
writeln!(begin, ".balign {align}").unwrap();
279+
write_linkage(&mut begin).unwrap();
280+
if let Visibility::Hidden = item_data.visibility {
281+
writeln!(begin, ".hidden {asm_name}").unwrap();
282+
}
283+
writeln!(begin, ".type {asm_name}, @function").unwrap();
284+
if !arch_prefix.is_empty() {
285+
writeln!(begin, "{}", arch_prefix).unwrap();
286+
}
287+
writeln!(begin, "{asm_name}:").unwrap();
288+
writeln!(begin, ".functype {asm_name} {}", wasm_functype(tcx, fn_abi)).unwrap();
289+
290+
writeln!(end).unwrap();
291+
// .size is ignored for function symbols, so we can skip it
292+
writeln!(end, "end_function").unwrap();
293+
}
267294
}
268295

269296
(begin, end)
270297
}
298+
299+
/// The webassembly type signature for the given function.
300+
///
301+
/// Used by the `.functype` directive on wasm targets.
302+
fn wasm_functype<'tcx>(tcx: TyCtxt<'tcx>, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> String {
303+
let mut signature = String::with_capacity(64);
304+
305+
let ptr_type = match tcx.data_layout.pointer_size.bits() {
306+
32 => "i32",
307+
64 => "i64",
308+
other => bug!("wasm pointer size cannot be {other} bits"),
309+
};
310+
311+
let hidden_return =
312+
matches!(fn_abi.ret.mode, PassMode::Indirect { .. } | PassMode::Pair { .. });
313+
314+
signature.push('(');
315+
316+
if hidden_return {
317+
signature.push_str(ptr_type);
318+
if !fn_abi.args.is_empty() {
319+
signature.push_str(", ");
320+
}
321+
}
322+
323+
let mut it = fn_abi.args.iter().peekable();
324+
while let Some(arg_abi) = it.next() {
325+
wasm_type(&mut signature, arg_abi, ptr_type);
326+
if it.peek().is_some() {
327+
signature.push_str(", ");
328+
}
329+
}
330+
331+
signature.push_str(") -> (");
332+
333+
if !hidden_return {
334+
wasm_type(&mut signature, &fn_abi.ret, ptr_type);
335+
}
336+
337+
signature.push(')');
338+
339+
signature
340+
}
341+
342+
fn wasm_type<'tcx>(signature: &mut String, arg_abi: &ArgAbi<'_, Ty<'tcx>>, ptr_type: &'static str) {
343+
match arg_abi.mode {
344+
PassMode::Ignore => { /* do nothing */ }
345+
PassMode::Direct(_) => {
346+
let direct_type = match arg_abi.layout.backend_repr {
347+
BackendRepr::Scalar(scalar) => wasm_primitive(scalar.primitive(), ptr_type),
348+
BackendRepr::Vector { .. } => "v128",
349+
other => unreachable!("unexpected BackendRepr: {:?}", other),
350+
};
351+
352+
signature.push_str(direct_type);
353+
}
354+
PassMode::Pair(_, _) => match arg_abi.layout.backend_repr {
355+
BackendRepr::ScalarPair(a, b) => {
356+
signature.push_str(wasm_primitive(a.primitive(), ptr_type));
357+
signature.push_str(", ");
358+
signature.push_str(wasm_primitive(b.primitive(), ptr_type));
359+
}
360+
other => unreachable!("{other:?}"),
361+
},
362+
PassMode::Cast { pad_i32, ref cast } => {
363+
// For wasm, Cast is used for single-field primitive wrappers like `struct Wrapper(i64);`
364+
assert!(!pad_i32, "not currently used by wasm calling convention");
365+
assert!(cast.prefix[0].is_none(), "no prefix");
366+
assert_eq!(cast.rest.total, arg_abi.layout.size, "single item");
367+
368+
let wrapped_wasm_type = match cast.rest.unit.kind {
369+
RegKind::Integer => match cast.rest.unit.size.bytes() {
370+
..=4 => "i32",
371+
..=8 => "i64",
372+
_ => ptr_type,
373+
},
374+
RegKind::Float => match cast.rest.unit.size.bytes() {
375+
..=4 => "f32",
376+
..=8 => "f64",
377+
_ => ptr_type,
378+
},
379+
RegKind::Vector => "v128",
380+
};
381+
382+
signature.push_str(wrapped_wasm_type);
383+
}
384+
PassMode::Indirect { .. } => signature.push_str(ptr_type),
385+
}
386+
}
387+
388+
fn wasm_primitive(primitive: Primitive, ptr_type: &'static str) -> &'static str {
389+
match primitive {
390+
Primitive::Int(integer, _) => match integer {
391+
Integer::I8 | Integer::I16 | Integer::I32 => "i32",
392+
Integer::I64 => "i64",
393+
Integer::I128 => "i64, i64",
394+
},
395+
Primitive::Float(float) => match float {
396+
Float::F16 | Float::F32 => "f32",
397+
Float::F64 => "f64",
398+
Float::F128 => "i64, i64",
399+
},
400+
Primitive::Pointer(_) => ptr_type,
401+
}
402+
}

‎tests/assembly/wasm32-naked-fn.rs‎

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
//@ revisions: wasm32-unknown wasm64-unknown wasm32-wasip1
2+
//@ add-core-stubs
3+
//@ assembly-output: emit-asm
4+
//@ [wasm32-unknown] compile-flags: --target wasm32-unknown-unknown
5+
//@ [wasm64-unknown] compile-flags: --target wasm64-unknown-unknown
6+
//@ [wasm32-wasip1] compile-flags: --target wasm32-wasip1
7+
//@ [wasm32-unknown] needs-llvm-components: webassembly
8+
//@ [wasm64-unknown] needs-llvm-components: webassembly
9+
//@ [wasm32-wasip1] needs-llvm-components: webassembly
10+
11+
#![crate_type = "lib"]
12+
#![feature(no_core, naked_functions, asm_experimental_arch, f128, linkage, fn_align)]
13+
#![no_core]
14+
15+
extern crate minicore;
16+
use minicore::*;
17+
18+
// CHECK: .section .text.nop,"",@
19+
// CHECK: .p2align 2
20+
// CHECK: .globl nop
21+
// CHECK-LABEL: nop:
22+
// CHECK: .functype nop () -> ()
23+
// CHECK-NOT: .size
24+
// CHECK: end_function
25+
#[no_mangle]
26+
#[naked]
27+
unsafe extern "C" fn nop() {
28+
naked_asm!("nop")
29+
}
30+
31+
// CHECK: .section .text.weak_aligned_nop,"",@
32+
// CHECK: .p2align 5
33+
// CHECK: .weak weak_aligned_nop
34+
// CHECK-LABEL: nop:
35+
// CHECK: .functype weak_aligned_nop () -> ()
36+
// CHECK-NOT: .size
37+
// CHECK: end_function
38+
#[no_mangle]
39+
#[naked]
40+
#[linkage = "weak"]
41+
#[repr(align(32))]
42+
unsafe extern "C" fn weak_aligned_nop() {
43+
naked_asm!("nop")
44+
}
45+
46+
// CHECK-LABEL: fn_i8_i8:
47+
// CHECK-NEXT: .functype fn_i8_i8 (i32) -> (i32)
48+
//
49+
// CHECK-NEXT: local.get 0
50+
// CHECK-NEXT: local.get 0
51+
// CHECK-NEXT: i32.mul
52+
//
53+
// CHECK-NEXT: end_function
54+
#[no_mangle]
55+
#[naked]
56+
unsafe extern "C" fn fn_i8_i8(num: i8) -> i8 {
57+
naked_asm!("local.get 0", "local.get 0", "i32.mul")
58+
}
59+
60+
// CHECK-LABEL: fn_i8_i8_i8:
61+
// CHECK: .functype fn_i8_i8_i8 (i32, i32) -> (i32)
62+
#[no_mangle]
63+
#[naked]
64+
unsafe extern "C" fn fn_i8_i8_i8(a: i8, b: i8) -> i8 {
65+
naked_asm!("local.get 1", "local.get 0", "i32.mul")
66+
}
67+
68+
// CHECK-LABEL: fn_unit_i8:
69+
// CHECK: .functype fn_unit_i8 () -> (i32)
70+
#[no_mangle]
71+
#[naked]
72+
unsafe extern "C" fn fn_unit_i8() -> i8 {
73+
naked_asm!("i32.const 42")
74+
}
75+
76+
// CHECK-LABEL: fn_i8_unit:
77+
// CHECK: .functype fn_i8_unit (i32) -> ()
78+
#[no_mangle]
79+
#[naked]
80+
unsafe extern "C" fn fn_i8_unit(_: i8) {
81+
naked_asm!("nop")
82+
}
83+
84+
// CHECK-LABEL: fn_i32_i32:
85+
// CHECK: .functype fn_i32_i32 (i32) -> (i32)
86+
#[no_mangle]
87+
#[naked]
88+
unsafe extern "C" fn fn_i32_i32(num: i32) -> i32 {
89+
naked_asm!("local.get 0", "local.get 0", "i32.mul")
90+
}
91+
92+
// CHECK-LABEL: fn_i64_i64:
93+
// CHECK: .functype fn_i64_i64 (i64) -> (i64)
94+
#[no_mangle]
95+
#[naked]
96+
unsafe extern "C" fn fn_i64_i64(num: i64) -> i64 {
97+
naked_asm!("local.get 0", "local.get 0", "i64.mul")
98+
}
99+
100+
// CHECK-LABEL: fn_i128_i128:
101+
// wasm32-unknown,wasm32-wasip1: .functype fn_i128_i128 (i32, i64, i64) -> ()
102+
// wasm64-unknown: .functype fn_i128_i128 (i64, i64, i64) -> ()
103+
#[allow(improper_ctypes_definitions)]
104+
#[no_mangle]
105+
#[naked]
106+
unsafe extern "C" fn fn_i128_i128(num: i128) -> i128 {
107+
naked_asm!(
108+
"local.get 0",
109+
"local.get 2",
110+
"i64.store 8",
111+
"local.get 0",
112+
"local.get 1",
113+
"i64.store 0",
114+
)
115+
}
116+
117+
// CHECK-LABEL: fn_f128_f128:
118+
// wasm32-unknown,wasm32-wasip1: .functype fn_f128_f128 (i32, i64, i64) -> ()
119+
// wasm64-unknown: .functype fn_f128_f128 (i64, i64, i64) -> ()
120+
#[no_mangle]
121+
#[naked]
122+
unsafe extern "C" fn fn_f128_f128(num: f128) -> f128 {
123+
naked_asm!(
124+
"local.get 0",
125+
"local.get 2",
126+
"i64.store 8",
127+
"local.get 0",
128+
"local.get 1",
129+
"i64.store 0",
130+
)
131+
}
132+
133+
#[repr(C)]
134+
struct Compound {
135+
a: u16,
136+
b: i64,
137+
}
138+
139+
// CHECK-LABEL: fn_compound_compound:
140+
// wasm32-wasip1: .functype fn_compound_compound (i32, i32) -> ()
141+
// wasm32-unknown: .functype fn_compound_compound (i32, i32, i64) -> ()
142+
// wasm64-unknown: .functype fn_compound_compound (i64, i64) -> ()
143+
#[no_mangle]
144+
#[naked]
145+
unsafe extern "C" fn fn_compound_compound(_: Compound) -> Compound {
146+
// this is the wasm32-unknown-unknown assembly
147+
naked_asm!(
148+
"local.get 0",
149+
"local.get 2",
150+
"i64.store 8",
151+
"local.get 0",
152+
"local.get 1",
153+
"i32.store16 0",
154+
)
155+
}
156+
157+
#[repr(C)]
158+
struct WrapperI32(i32);
159+
160+
// CHECK-LABEL: fn_wrapperi32_wrapperi32:
161+
// CHECK: .functype fn_wrapperi32_wrapperi32 (i32) -> (i32)
162+
#[no_mangle]
163+
#[naked]
164+
unsafe extern "C" fn fn_wrapperi32_wrapperi32(_: WrapperI32) -> WrapperI32 {
165+
naked_asm!("local.get 0")
166+
}
167+
168+
#[repr(C)]
169+
struct WrapperI64(i64);
170+
171+
// CHECK-LABEL: fn_wrapperi64_wrapperi64:
172+
// CHECK: .functype fn_wrapperi64_wrapperi64 (i64) -> (i64)
173+
#[no_mangle]
174+
#[naked]
175+
unsafe extern "C" fn fn_wrapperi64_wrapperi64(_: WrapperI64) -> WrapperI64 {
176+
naked_asm!("local.get 0")
177+
}
178+
179+
#[repr(C)]
180+
struct WrapperF32(f32);
181+
182+
// CHECK-LABEL: fn_wrapperf32_wrapperf32:
183+
// CHECK: .functype fn_wrapperf32_wrapperf32 (f32) -> (f32)
184+
#[no_mangle]
185+
#[naked]
186+
unsafe extern "C" fn fn_wrapperf32_wrapperf32(_: WrapperF32) -> WrapperF32 {
187+
naked_asm!("local.get 0")
188+
}
189+
190+
#[repr(C)]
191+
struct WrapperF64(f64);
192+
193+
// CHECK-LABEL: fn_wrapperf64_wrapperf64:
194+
// CHECK: .functype fn_wrapperf64_wrapperf64 (f64) -> (f64)
195+
#[no_mangle]
196+
#[naked]
197+
unsafe extern "C" fn fn_wrapperf64_wrapperf64(_: WrapperF64) -> WrapperF64 {
198+
naked_asm!("local.get 0")
199+
}

0 commit comments

Comments
 (0)