Or a loveletter? Read this while listening to Ode to Joy, also look at this awesome pixelart i did just for this article:
its analyzing a life 🤓
I am currently working on register allocation for both my register based bytecode virtual machine and my x86 just in time compiler, both for my programming language purple-garden.
Register allocation is assigning the unlimited amount of virtual registers to the limited set of storage resources hardware has, specifically the gprs (general purpose registers), the fprs (floating point registers) and the stack (see CPU Registers x86 for an examplary overview). The first step in this odeal is livelyness analsis.
A garden(ing) example
THe simplest, yet somewhat non trivial example would be:
1fn f(a:Int n:Int acc:Int) Int {
2 match {
3 n < 1 { a + acc }
4 { f(a n-1 acc*n) }
5 }
6}This is equiv to the following c code, which i “disassembled”
from the very well written Linear scan register allocation on
SSA article by Max
Bernstein1, in which i didnt really understand the
IR display, before translating it to c and then to my own
language (purple garden does not have loops 🤓 so this is only
equiv when called with f(a n 1)):
1long f(long a, long n) {
2 long acc = 1;
3 while (n >= 1) {
4 acc = acc * n;
5 n = n - 1;
6 }
7 return a + acc;
8}Feeding f(a,n,acc: Int): Int into purple gardens pipeline
produces the following, optimised immediate representation
(using -I), this means some values were dead and thus removed 2:
1fn f1(%v0:Int, %v1:Int, %v2:Int) -> Int {
2 b0(%v0, %v1, %v2):
3 %v4:Bool = ILt %v1, 1
4 br %v4, b1(%v0, %v1, %v2), b2(%v0, %v1, %v2)
5 b1(%v5, %v6, %v7):
6 %v8:Int = IAdd %v5, %v7
7 ret %v8
8 b2(%v9, %v10, %v11):
9 %v13:Int = ISub %v10, 1
10 %v14:Int = IMul %v11, %v10
11 tail f1(%v9, %v13, %v14)
12}PUrple garden IR structure and notation
Purple gardens immediate representation is static single assigment
(SSA) form based,
meaning each value / virtual register is assigned exactly once and all are
immutable. They are notated with %vN, for virtual register or value, whatever
you (the reader) like more.
Compared to many other ssa irs, pgir does not use phi nodes. I opted for block params and terminators instead. It is made up of blocks, where each block has a list of parameters, a terminator and a body of instructions, where each instruction is in ssa form. All values are block local and the only interblock dataflow is facilitated via block params and block terminators.
1; block with id N
2; |
3; v
4 b2(%v0, %v1, %v2):
5; ^^^^^^^^^^^^^ parameters 0..2
6
7; value / virtual reg def
8; | instr kind
9; | |
10; v v
11 %v7:Int = ISub %v1, 1 ; <-- imm value
12; ^ ^ argument
13; |
14; instr result type
15
16; terminator with params
17 tail f1(%v0, %v7, %v8)Or in rust terms:
1pub struct Block<'b, A: Allocator = Global> {
2 pub id: Id,
3 pub instructions: Vec<Instr<'b>, A>,
4 pub params: ParamsId,
5 pub term: Option<Terminator>,
6 // [...]
7}
8
9pub enum Instr<'i> {
10 BinImm {
11 op: BinOp,
12 dst: TypeId<'i>,
13 lhs: Id,
14 imm: i32,
15 },
16 // [...]
17 Noop,
18}
19
20pub enum Terminator {
21 Return {
22 value: Option<Id>,
23 },
24 // [...]
25 Tail {
26 func: Id,
27 args: ParamsId,
28 },
29 // [...]
30}Pretty easy right?
Virtual registers / Values and their lifetimes
LEts look at the previous IR example, but this time with virtual register
lifetime annotations (when the cli is invoked with -II):
1fn f1(%v0:Int, %v1:Int, %v2:Int) -> Int { %v0 %v1 %v2 %v4 %v5 %v7 %v8 %v9 %v10 %v11 %v13 %v14
2 b0(%v0, %v1, %v2): | | |
3 %v4:Bool = ILt %v1, 1 | | | |
4 br %v4, b1(%v0, %v1, %v2), b2(%v0, %v1, %v2) | | | |
5 b1(%v5, %v6, %v7): | |
6 %v8:Int = IAdd %v5, %v7 | | |
7 ret %v8 |
8 b2(%v9, %v10, %v11): | | |
9 %v13:Int = ISub %v10, 1 | | | |
10 %v14:Int = IMul %v11, %v10 | | | | |
11 tail f1(%v9, %v13, %v14) | | |
12}This display enables a visual understanding of the livetimes of each value, encoded as lo for start, and hi for last, the following enables a non visual understanding: (The positions are linearised: block params are at the block base, block instructions are at block base+1+instr pos and terminators are at block base+instr count+1)
| value | lo | hi |
|---|---|---|
| %v0 | 0 | 2 |
| %v1 | 0 | 2 |
| %v2 | 0 | 2 |
| %v4 | 1 | 2 |
| %v5 | 3 | 4 |
| %v7 | 3 | 4 |
| %v8 | 4 | 5 |
| %v9 | 6 | 9 |
| %v10 | 6 | 8 |
| %v11 | 6 | 8 |
| %v13 | 7 | 9 |
| %v14 | 8 | 9 |
A short but sweet implementation:
With that out of the way the algorithm amounts to a simple:
- walk blocks
- visit terminator, observe all virtual register usages
- visit instructions in reverse, observe all virtual register definition and usages
- visit parameters to block, observe each
Observing
Where observing is defined as:
- with tuple
(lo,hi)andpos - if
hi==0; sethi=pos lo=pos
In oxidised terms:
1/// observer Id := (lo,hi), always updates lo, only fills hi, if previously unset
2fn observe<V: Allocator>(out: &mut Vec<(u32, u32), V>, id: Id, pos: u32) {
3 let idx = id.0 as usize;
4 let e = &mut out[idx];
5 if e.1 == 0 {
6 e.1 = pos;
7 }
8 e.0 = pos;
9}Helpers to help with defs and usages:
Asking for defs and uses in instructions and terminators:
Extracting definitions
What it says, Func::def_of returns a definition of a virtual register, if
instr has one:
1impl Func<'_> {
2 #[must_use]
3 pub fn def_of(instr: &Instr<'_>) -> Option<Id> {
4 match instr {
5 Instr::Alloc { dst, .. }
6 | Instr::Load { dst, .. }
7 | Instr::AddrOf { dst, .. }
8 | Instr::Bin { dst, .. }
9 | Instr::BinImm { dst, .. }
10 | Instr::LoadConst { dst, .. }
11 | Instr::Call { dst, .. }
12 | Instr::Sys { dst, .. }
13 | Instr::Cast { dst, .. }
14 | Instr::Lookup { dst, .. }
15 | Instr::Assert { dst, .. } => Some(dst.id),
16 Instr::Store { .. } | Instr::Noop => None,
17 }
18 }
19}Calling f for each usage of value in an instruction
What it says, calling f for each value usage in an instruction
1pub fn for_each_use_of_instr(&self, instr: &Instr<'_>, mut f: impl FnMut(Id)) {
2 match instr {
3 Instr::Bin { lhs, rhs, .. } => {
4 f(*lhs);
5 f(*rhs);
6 }
7 Instr::BinImm { lhs, .. } => f(*lhs),
8 Instr::Call { args, .. } | Instr::Sys { args, .. } => {
9 for &a in self.params(*args) {
10 f(a);
11 }
12 }
13 Instr::Cast { from, .. } => f(from.id),
14 Instr::Lookup { subject, .. } => f(*subject),
15 Instr::Store { src, base, .. } => {
16 f(*src);
17 f(*base);
18 }
19 Instr::Assert { cond, .. } => f(*cond),
20 Instr::Load { base, .. } | Instr::AddrOf { base, .. } => f(*base),
21 Instr::LoadConst { .. } | Instr::Noop | Instr::Alloc { .. } => {}
22 }
23}Calling f for each usage of value in a terminator
What it says, calling f for each value usage in a terminator
1pub fn for_each_use_of_term(&self, term: &Terminator, mut f: impl FnMut(Id)) {
2 match term {
3 Terminator::Return {
4 value: Some(id), ..
5 } => f(*id),
6 Terminator::Return { value: None, .. } => {}
7 Terminator::Jump { params, .. } => {
8 for &p in self.params(*params) {
9 f(p);
10 }
11 }
12 Terminator::Tail { args, .. } => {
13 for &a in self.params(*args) {
14 f(a);
15 }
16 }
17 Terminator::Switch {
18 subject,
19 cases,
20 default: (_, default_params),
21 ..
22 } => {
23 f(*subject);
24 for case in self.cases(*cases) {
25 for &p in self.params(case.target.1) {
26 f(p);
27 }
28 }
29 for &p in self.params(*default_params) {
30 f(p);
31 }
32 }
33 Terminator::Branch {
34 cond,
35 yes: (_, yes_params),
36 no: (_, no_params),
37 ..
38 } => {
39 f(*cond);
40 for &p in self.params(*yes_params) {
41 f(p);
42 }
43 for &p in self.params(*no_params) {
44 f(p);
45 }
46 }
47 Terminator::BranchCmpImm {
48 lhs,
49 yes: (_, yes_params),
50 no: (_, no_params),
51 ..
52 }
53 | Terminator::BranchCmp {
54 lhs,
55 yes: (_, yes_params),
56 no: (_, no_params),
57 ..
58 } => {
59 f(*lhs);
60 if let Terminator::BranchCmp { rhs, .. } = term {
61 f(*rhs);
62 }
63 for &p in self.params(*yes_params) {
64 f(p);
65 }
66 for &p in self.params(*no_params) {
67 f(p);
68 }
69 }
70 }
71}Fully rusted together:
A few quick context injectors:
- out is a reusable Vec allocated in the V allocator, thats why its reset and resized
- UNSET is used as the default empty value
- Func::next_id() returns the highest id defined in the function (self)+1
1impl<A: Allocator> Func<'_, A> {
2 pub fn live_set_into<V: Allocator>(&self, out: &mut Vec<(u32, u32), V>) {
3 const UNSET: (u32, u32) = (u32::MAX, 0u32);
4
5 out.clear();
6 out.resize(self.next_id() as usize, UNSET);
7
8 // observe defined as before
9
10 let mut base = 0;
11 for block in self.blocks.iter().filter(|b| !b.tombstone) {
12 let block_offset = block.instructions.len();
13 let offset = base;
14 if let Some(term) = &block.term {
15 self.for_each_use_of_term(term, |id| {
16 // term is instruction count +1
17 observe(out, id, (offset + block_offset + 1) as u32)
18 });
19 }
20
21 for (i, instr) in block.instructions.iter().enumerate().rev() {
22 // +1, so it cant start at 0, since thats where the params start :)
23 let pos = (base + 1 + i) as u32;
24 self.for_each_use_of_instr(instr, |id| observe(out, id, pos));
25 if let Some(def) = Func::def_of(instr) {
26 observe(out, def, pos)
27 }
28 }
29
30 for &p in self.params(block.params) {
31 observe(out, p, base as u32);
32 }
33
34 // params + instructions + terminator
35 base += block_offset + 2;
36 }
37 }
38}Wiring liveset analysis into register allocation for the bytecode backend
Since we now know where what value is alive and where it is to be considered dead we can pass this information to the register allocator and ask it for the register or stack slot of any value we have to fit into hardware or virtual machine constraints.
Prep and register allocation prepass
When compiling the IR to bytecode for the purple garden vm, the first step is
to compute the liveset, with the algorithm introduced before and then
performing the preallocation pass with the Cc::prealloc method:
1impl<'cc, S: Allocator + Clone, A: Allocator + Clone> Cc<'cc, S, A> {
2 fn cc<F: Allocator>(
3 &mut self,
4 fun: &'cc Func<'cc, F>,
5 native: Option<&mut purple_garden_jit::CodeArena<impl Pages>>,
6 ) {
7 // [...]
8
9 // compute liveset
10 fun.live_set_into(&mut self.live_set);
11
12 // walk the IR once before lowering, after this its just a simple
13 // lookup
14 self.prealloc(fun);
15
16 // [...] iterate blocks
17
18 // [...] iterate instructions
19
20 // secret postprocessing of instructions which dst points to the stack
21 if let Some(Location::Stack(slot)) =
22 Func::def_of(instruction).and_then(|dst| self.regalloc.location(dst))
23 {
24 self.emit(Op::StoreSlot { slot, src: TMP });
25 }
26 }
27}Lowering Bin ir ops
Bin is the baseline operator. The BinImm ops we see in our example (ISub %v10, 1) are optimisations (one less dispatch, less memory interactions) compiled in
the same way, just with an immediate rhs:
1impl<'cc, S: Allocator + Clone, A: Allocator + Clone> Cc<'cc, S, A> {
2 fn instr<F: Allocator>(&mut self, owner: &'cc Func<'cc, F>, i: &'cc ir::Instr<'cc>) {
3 match i {
4 // [...] all the other IR instructions
5 ir::Instr::Bin {
6 op, dst, lhs, rhs, ..
7 } => {
8 // interacting with the register alloc prepass
9 let dst = self.def_register(dst.id);
10 let lhs = self.ensure_register(*lhs);
11 let rhs = self.ensure_register(*rhs);
12
13 macro_rules! emit_bins {
14 ($($name:ident),*) => {
15 match op {
16 $(
17 ir::BinOp::$name => {
18 Op::$name {
19 dst,
20 lhs,
21 rhs,
22 }
23 },
24 )*
25 }
26 };
27 }
28
29 self.emit(emit_bins! {
30 IAdd, ISub, IMul, IDiv, IMod, IAnd, ILt, IGt, IEq,
31 DAdd, DSub, DMul, DDiv, DLt, DGt,
32 BEq, SEq
33 });
34 }
35 ir::Instr::BinImm {
36 op, dst, lhs, imm, ..
37 } => {
38 let dst = self.def_register(dst.id);
39 let lhs = self.ensure_register(*lhs);
40 let imm = *imm;
41
42 macro_rules! emit_bin_imms {
43 ($($ir:ident => $op:ident),* $(,)?) => {
44 match op {
45 $(ir::BinOp::$ir => Op::$op { dst, lhs, imm },)*
46 _ => unreachable!("only integer immediate binops are represented as BinImm"),
47 }
48 };
49 }
50
51 self.emit(emit_bin_imms! {
52 IAdd => IAddI,
53 ISub => ISubI,
54 IMul => IMulI,
55 IDiv => IDivI,
56 IMod => IModI,
57 IAnd => IAndI,
58 IEq => IEqI,
59 IGt => IGtI,
60 ILt => ILtI,
61 });
62 }
63 }
64 }
65}This is so clean and a nice use of rust macros, i love this, dont @me; lowering all binary and binary with immediate value operations in 57 lines is where its at.
Ensuring and definining a register
Two helpers for making sure a requested virtual register is accessible to a point where operations can be applied on it, the first one being ensure_register:
1/// Translate virtual register into vm register, abstracts spilling, restoring and register
2/// allocation in general away
3fn ensure_register(&mut self, id: Id) -> u8 {
4 match self.regalloc.location(id) {
5 Some(Location::Reg(r)) => return r.num(),
6 // requested value lives on the stack, we need to restore it and we use a temp
7 // register
8 Some(Location::Stack(slot)) => {
9 let dst = self.next_tmp();
10 self.emit(Op::LoadSlot { dst, slot });
11 dst
12 }
13 None => unreachable!("%v{} has no location", id.0),
14 }
15}
16
17fn next_tmp(&self) -> u8 {
18 match self.buf.last() {
19 Some(Op::LoadSlot { dst: SCRATCH, .. }) => TMP,
20 _ => SCRATCH,
21 }
22}This means if a value is requested and:
- the prepass assigned it to a register, thats perfect and we can simply use it
- or if its on the stack we have to ask the compiler state for a temporary register and materialise the value into it
Temporary registers in the virtual machine are defined as
SCRATCHandTMP, respectively in the preamble of the virtual machines source code:RUST1pub const REGISTER_COUNT: usize = 32; 2const SCRATCH: u8 = REGISTER_COUNT as u8 - 1; 3const TMP: u8 = SCRATCH - 1; 4 5/// Every vm register but [`SCRATCH`] and [`TMP`], r0 last since xralloc2 hands 6/// out the lowest free one first, which also puts the caller-saved registers 7/// before the callee-saved ones. 8const POOL: [u8; TMP as usize] = { 9 let mut pool = [0; TMP as usize]; 10 let mut i = 0; 11 while i < TMP as usize { 12 pool[i] = (TMP as usize - 1 - i) as u8; 13 i += 1; 14 } 15 pool 16};
Thusfar ensure is simple to understand, the other helper is def_register, which means we, again, take the registers number, if it has one, otherwise we take TMP, which our previously introduced (and probably not noticed by anyone) snipped at the end of the codeblock in Prep and register allocation prepass then “spills” into its slot on the stack with the respective bytecode instruction.
1/// The register an op writes `id` into, [`TMP`] if `id` lives in a slot: the block loop stores
2/// it there after the op.
3fn def_register(&self, id: Id) -> u8 {
4 match self.regalloc.location(id) {
5 Some(Location::Reg(r)) => r.num(),
6 Some(Location::Stack(_)) => TMP,
7 None => unreachable!("%v{} has no location", id.0),
8 }
9}All the bytes and the codes together:
Swapping -I for -D (with --no-jit, otherwise we would get machinecode 🤓), results in the following
virtual machine registers assigned to the ssa values:
100000000 <f>:
2 0000: ilt_imm r3, r1, #1 ; 4: n < 1 { a + acc }
3 0001: jmpf r3, 0004 <f.bb_0004>
4 0002: iadd r0, r0, r2
5 0003: ret
600000004 <f.bb_0004>:
7 0004: isub_imm r3, r1, #1 ; 5: { f(a n-1 acc*n) }
8 0005: imul r2, r2, r1
9 0006: mov r1, r3
10 0007: tail 0000 <f>JIT :O
For the crazy hackermans, here is the same example with the jit on, jit and bc share IR, liveness analysis, register allocation and standard libary, aint that fun?
10000000000000000 <jit_f>:
2 0: 48 81 fa 01 00 00 00 cmp $0x1,%rdx
3 7: 49 c7 c0 00 00 00 00 mov $0x0,%r8
4 e: 41 0f 9c c0 setl %r8b
5 12: 4d 85 c0 test %r8,%r8
6 15: 0f 84 0b 00 00 00 je 26 <jit_f+0x26>
7 1b: 48 01 ce add %rcx,%rsi
8 1e: 48 89 f0 mov %rsi,%rax
9 21: c3 ret
10 22: cc int3
11 23: cc int3
12 24: cc int3
13 25: cc int3
14 26: 49 89 d0 mov %rdx,%r8
15 29: 49 81 e8 01 00 00 00 sub $0x1,%r8
16 30: 49 89 c9 mov %rcx,%r9
17 33: 4c 0f af ca imul %rdx,%r9
18 37: 4c 89 c2 mov %r8,%rdx
19 3a: 4c 89 c9 mov %r9,%rcx
20 3d: e9 be ff ff ff jmp 0 <jit_f>
21 42: c3 retDont be mean to the x86 output, i know this could be shorter, the compare and jmp could be a single instr and there is a lot of dead code and long encodings, but this is my first jit and im glad it works at all :)
Future work
The attentive reader may have figured out how i have described the inputs and ouputs of the register allocator, respectively the liveset (and the pool of regs) and assignment of virtual to physical register, but have not actually shown or introduced the register allocator itself, or the results of the IR design which forces the regalloc to treat all block transitions (jumps, entries, returns) as edge moves and has to do extra work to resolve these.
If my motivation holds up i will write a follow up article about these topics, but right now i am still in the process of “perfecting” the interactions between my ir, my regalloc and my backends, so maybe stay tuned? I have an rss feed if anyone is actually at the end of the abyss im screaming my incohesive blabbering into.