]> Witch of Git - ivy/blob - src/trans/x64.rs
Reduce stack traffic by finding loads through SSA
[ivy] / src / trans / x64.rs
1 use crate::ast;
2 use crate::trans::{Code, FnName, Func, Program, SsaName, Var};
3 use std::{
4 borrow::Cow,
5 collections::HashMap,
6 fmt,
7 io::{self, Write},
8 };
9
10 #[cfg(target_os = "macos")]
11 macro_rules! extern_label {
12 ($l:literal) => {
13 Label::Extern(concat!("_", $l))
14 };
15 }
16 #[cfg(not(target_os = "macos"))]
17 macro_rules! extern_label {
18 ($l:literal) => {
19 Label::Extern($l)
20 };
21 }
22
23 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
24 pub enum Label {
25 Global(u32),
26 Function(FnName),
27 Extern(&'static str),
28 Builtin(&'static str),
29 }
30
31 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
32 pub enum Addr {
33 Rip(Label),
34 Off(Reg, i32),
35 }
36
37 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
38 #[allow(unused)]
39 pub enum Reg {
40 Rax,
41 Rbx,
42 Rcx,
43 Rdx,
44 Rsi,
45 Rdi,
46 Rbp,
47 Rsp,
48 R8,
49 R9,
50 R10,
51 R11,
52 R12,
53 R13,
54 R14,
55 R15,
56 }
57
58 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
59 pub enum Inst {
60 Push(Reg),
61 Load(Reg, Addr),
62 Store(Addr, Reg),
63 Mov(Reg, Reg),
64 Imm(Reg, i64),
65 Lea(Reg, Addr),
66 Call(Label),
67 Jmp(Label),
68 Add(Reg, u32),
69 Sub(Reg, u32),
70 Pop(Reg),
71 Ret,
72 }
73
74 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
75 pub struct Definition {
76 name: Label,
77 body: Cow<'static, [Inst]>,
78 }
79
80 pub struct Global {
81 name: Label,
82 value: i64,
83 }
84
85 pub fn write_compile(
86 builtins: &HashMap<String, ast::Name>,
87 out: &mut impl Write,
88 prog: &Program,
89 ) -> io::Result<()> {
90 let (entry, globals, defns) = compile(&prog);
91 struct Entry {
92 number: u32,
93 defn: Definition,
94 params: u16,
95 }
96 let builtins = builtins
97 .iter()
98 .map(|(name, id)| {
99 let number = id.global_number().unwrap();
100 let (params, name, code) = builtin(&name);
101 Entry {
102 number,
103 defn: Definition {
104 name: Label::Builtin(name),
105 body: Cow::Borrowed(code),
106 },
107 params,
108 }
109 })
110 .collect::<Vec<_>>();
111 writeln!(out, ".data")?;
112 for global in globals {
113 writeln!(out, "{}", global)?;
114 }
115 writeln!(out)?;
116 writeln!(
117 out,
118 r#".text
119 .global _start
120 _start:
121 sub rsp, 15
122 and spl, 0xF0"#
123 )?;
124
125 for builtin in &builtins {
126 writeln!(
127 out,
128 r#"
129 lea rdi, [rip + {}]
130 mov rsi, {}
131 mov rdx, 0
132 call {}
133 mov [rip + {}], rax"#,
134 builtin.defn.name,
135 builtin.params,
136 extern_label!("ivy_make_lam"),
137 Label::Global(builtin.number),
138 )?;
139 }
140
141 writeln!(
142 out,
143 r#"
144 call {}
145 mov rdi, 0
146 call {}
147 "#,
148 entry,
149 extern_label!("ivy_exit"),
150 )?;
151 for builtin in &builtins {
152 writeln!(out, "{}", builtin.defn)?;
153 writeln!(out)?;
154 }
155 for defn in defns {
156 writeln!(out, "{}", defn)?;
157 writeln!(out)?;
158 }
159 Ok(())
160 }
161
162 pub fn compile(prog: &Program) -> (Label, Vec<Global>, Vec<Definition>) {
163 let globals = prog
164 .globals
165 .iter()
166 .map(|&global| Global {
167 name: Label::Global(global),
168 value: 0,
169 })
170 .collect();
171 let defns = prog
172 .functions
173 .values()
174 .map(|func| compile_func(func))
175 .collect();
176 (Label::Function(prog.top), globals, defns)
177 }
178
179 fn even_up(x: usize) -> usize {
180 x + (x & 1)
181 }
182
183 fn compile_func(func: &Func) -> Definition {
184 fn stack(x: &SsaName) -> Addr {
185 Addr::Off(Reg::Rsp, x.0 as i32 * 8)
186 }
187
188 const HEADER_SIZE: i32 = 0x18;
189
190 fn param(x: u16) -> Addr {
191 Addr::Off(Reg::Rbx, HEADER_SIZE + x as i32 * 8)
192 }
193
194 let upvar = move |x: u16| {
195 let params = func.params as i32 * 8;
196 Addr::Off(Reg::Rbx, HEADER_SIZE + params + x as i32 * 8)
197 };
198
199 let mut body = Vec::new();
200 let stack_slots = even_up(func.order.len());
201 body.push(Inst::Push(Reg::Rbx));
202 body.push(Inst::Mov(Reg::Rbx, Reg::Rdi));
203 body.push(Inst::Sub(Reg::Rsp, stack_slots as u32 * 8));
204 let address = |v| match v {
205 Var::Global(x) => Addr::Rip(Label::Global(x)),
206 Var::Param(x) => param(x),
207 Var::Upvar(x) => upvar(x),
208 Var::Ssa(s) => stack(&s),
209 };
210 let load = |s| {
211 match func.block[s] {
212 Code::Load(v) => address(v),
213 _ => stack(s),
214 }
215 };
216 for ssa in &func.order {
217 match &func.block[ssa] {
218 // We don't need to generate code for loads themselves, they're consulted by other accesses
219 Code::Load(_) => (),
220 Code::StoreGlobal(x, s) => {
221 body.push(Inst::Load(Reg::Rax, load(s)));
222 body.push(Inst::Store(Addr::Rip(Label::Global(*x)), Reg::Rax));
223 }
224 Code::MakeLam {
225 name,
226 upvars,
227 params,
228 } => {
229 body.push(Inst::Lea(Reg::Rdi, Addr::Rip(Label::Function(*name))));
230 body.push(Inst::Imm(Reg::Rsi, *params as i64));
231 body.push(Inst::Imm(Reg::Rdx, upvars.len() as i64));
232 body.push(Inst::Call(extern_label!("ivy_make_lam")));
233 // @TODO: Implement copying in closure captures
234 assert!(upvars.is_empty());
235 body.push(Inst::Store(stack(ssa), Reg::Rax));
236 }
237 Code::App(terms) => {
238 let mut terms = terms.iter();
239 body.push(Inst::Load(
240 Reg::Rdi,
241 load(terms.next().expect("a function to apply")),
242 ));
243 match terms.next() {
244 Some(s) => body.push(Inst::Load(Reg::Rsi, load(s))),
245 None => body.push(Inst::Imm(Reg::Rsi, 0)),
246 }
247 body.push(Inst::Call(extern_label!("ivy_app")));
248 while let Some(term) = terms.next() {
249 body.push(Inst::Mov(Reg::Rdi, Reg::Rax));
250 body.push(Inst::Load(Reg::Rsi, load(term)));
251 body.push(Inst::Call(extern_label!("ivy_app_mut")));
252 }
253 body.push(Inst::Store(stack(ssa), Reg::Rax));
254 }
255 Code::Num(n) => {
256 body.push(Inst::Imm(Reg::Rax, (n << 1) | 1i64));
257 body.push(Inst::Store(stack(ssa), Reg::Rax));
258 }
259 }
260 }
261 if let Code::Load(v) = func.block[&func.result] {
262 body.push(Inst::Load(Reg::Rax, address(v)));
263 }
264 body.push(Inst::Add(Reg::Rsp, stack_slots as u32 * 8));
265 body.push(Inst::Pop(Reg::Rbx));
266 body.push(Inst::Ret);
267 Definition {
268 name: Label::Function(func.name),
269 body: Cow::Owned(body),
270 }
271 }
272
273 impl fmt::Display for Definition {
274 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
275 match self.body.len() {
276 0 => write!(fmt, "{}:", self.name),
277 1 => write!(fmt, "{}: {}", self.name, self.body[0]),
278 len => {
279 writeln!(fmt, "{}:", self.name)?;
280 for inst in &self.body[..len - 1] {
281 writeln!(fmt, " {}", inst)?;
282 }
283 write!(fmt, " {}", self.body.last().unwrap())
284 }
285 }
286 }
287 }
288
289 impl fmt::Display for Global {
290 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
291 write!(fmt, "{}: .quad {}", self.name, self.value)
292 }
293 }
294
295 impl fmt::Display for Label {
296 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
297 match self {
298 Label::Global(g) => write!(fmt, "IVY_GLOBAL${}", g),
299 Label::Function(f) => write!(fmt, "ivy_fn${}", f.0),
300 Label::Extern(l) => write!(fmt, "{}", l),
301 Label::Builtin(l) => write!(fmt, "ivy_builtin${}", l),
302 }
303 }
304 }
305
306 impl fmt::Display for Reg {
307 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
308 match self {
309 Reg::Rax => write!(fmt, "rax"),
310 Reg::Rbx => write!(fmt, "rbx"),
311 Reg::Rcx => write!(fmt, "rcx"),
312 Reg::Rdx => write!(fmt, "rdx"),
313 Reg::Rsi => write!(fmt, "rsi"),
314 Reg::Rdi => write!(fmt, "rdi"),
315 Reg::Rbp => write!(fmt, "rbp"),
316 Reg::Rsp => write!(fmt, "rsp"),
317 Reg::R8 => write!(fmt, "r8"),
318 Reg::R9 => write!(fmt, "r9"),
319 Reg::R10 => write!(fmt, "r10"),
320 Reg::R11 => write!(fmt, "r11"),
321 Reg::R12 => write!(fmt, "r12"),
322 Reg::R13 => write!(fmt, "r13"),
323 Reg::R14 => write!(fmt, "r14"),
324 Reg::R15 => write!(fmt, "r15"),
325 }
326 }
327 }
328
329 impl fmt::Display for Addr {
330 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
331 match self {
332 Addr::Rip(l) => write!(fmt, "[rip + {}]", l),
333 Addr::Off(r, o) => write!(fmt, "[{} + {}]", r, o),
334 }
335 }
336 }
337
338 impl fmt::Display for Inst {
339 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
340 match self {
341 Inst::Push(r) => write!(fmt, "push {}", r),
342 Inst::Pop(r) => write!(fmt, "pop {}", r),
343 Inst::Mov(d, s) => write!(fmt, "mov {}, {}", d, s),
344 Inst::Load(r, a) => write!(fmt, "mov {}, {}", r, a),
345 Inst::Store(a, r) => write!(fmt, "mov {}, {}", a, r),
346 Inst::Imm(r, i) => write!(fmt, "mov {}, {}", r, i),
347 Inst::Lea(r, a) => write!(fmt, "lea {}, {}", r, a),
348 Inst::Call(l) => write!(fmt, "call {}", l),
349 Inst::Jmp(l) => write!(fmt, "jmp {}", l),
350 Inst::Add(r, x) => write!(fmt, "add {}, {}", r, x),
351 Inst::Sub(r, x) => write!(fmt, "sub {}, {}", r, x),
352 Inst::Ret => write!(fmt, "ret"),
353 }
354 }
355 }
356
357 fn builtin(name: &str) -> (u16, &'static str, &'static [Inst]) {
358 match name {
359 "debug" => (
360 1,
361 "debug",
362 &[
363 Inst::Load(Reg::Rdi, Addr::Off(Reg::Rdi, 0x18)),
364 Inst::Jmp(extern_label!("ivy_debug")),
365 ],
366 ),
367 name => panic!("Unsupported builtin '{}' for x64", name),
368 }
369 }