跳到正文
Hacker News· xnacly·· 6 天前AI 评分12

purple-garden 解释器如何用 StackAllocator 优化寄存器溢出栈

Interpreter spill stack on the process stack

AI 导读

purple-garden 解释器将寄存器溢出栈从 Vec 迁移到基于 Rust 1.100 稳定 Allocator trait 的 StackAllocator，把溢出数据直接放在进程栈上。

正文 · AI 翻译

我目前仍在开发我的脚本语言 purple-garden，并希望将编译管道的所有阶段都移植到新的分配器 trait 上（原因有很多，例如我在所有阶段都使用暂存存储进行临时分配，并希望既重用又优化这块内存区域，以提高局部性和降低分配开销）。这些受 zig 风格启发的分配器之一就是 StackAllocator，我将在本文中解释它，并用它来优化寄存器溢出和恢复到虚拟机栈的操作（其他几个是 PageAlloc、BumpAlloc<A> 和 MetricAllocator<A>）。

这些改动实际上不再影响 purple garden 的通用性能，因为大部分速度提升都来自运行时的 JIT（目前仅支持 x86）

一些前置知识

trait::Allocator

trait::Allocator 提案最近在 Rust 1.100 版本中 稳定，它引入了一个新的 unsafe trait 以及一些接受这些分配器的函数，例如 {Vec,Box,HashMap,HashSet}::new_in()。这些函数随后将内存分配到传入的分配器中，因此现在可以存在于该分配器所依托的任何内存区域中：

RUST

 1pub unsafe trait Allocator {
 2    fn allocate(
 3        &self,
 4        layout: Layout,
 5    ) -> Result<NonNull<[u8]>, AllocError>;
 6
 7    unsafe fn deallocate(
 8        &self,
 9        ptr: NonNull<u8>,
10        layout: Layout,
11    );
12
13    fn allocate_zeroed(
14        &self,
15        layout: Layout,
16    ) -> Result<NonNull<[u8]>, AllocError>
17
18    unsafe fn grow(
19        &self,
20        ptr: NonNull<u8>,
21        old_layout: Layout,
22        new_layout: Layout,
23    ) -> Result<NonNull<[u8]>, AllocError>
24
25    unsafe fn grow_zeroed(
26        &self,
27        ptr: NonNull<u8>,
28        old_layout: Layout,
29        new_layout: Layout,
30    ) -> Result<NonNull<[u8]>, AllocError>
31
32    unsafe fn shrink(
33        &self,
34        ptr: NonNull<u8>,
35        old_layout: Layout,
36        new_layout: Layout,
37    ) -> Result<NonNull<[u8]>, AllocError>
38}

Purple garden 虚拟机

Purple garden 的解释器以及虚拟机是基于寄存器的（出于性能原因），因此可能会耗尽寄存器空间。字节码寄存器分配器会针对寄存器耗尽以及将值的生命周期延长到函数之外的情况进行指示，例如考虑这个简单的 fib 函数：

GARDEN

#! Compute Fibonacci numbers with the classic doubly-recursive definition.
fn fib(n:Int) Int {
    match {
        n == 0 { 0 }
        n == 1 { 1 }
        { fib(n - 1) + fib(n - 2) }
    }
}

以及它的 garden 反汇编：

ASM

 1; purple-garden --no-jit -D examples/fib.garden
 200000000 <fib>:
 3  ; 8: fn fib(n:Int) Int {
 4  0000:    push r1
 5  ; 10: n == 0 { 0 }
 6  0001:    jmpne_imm r0, #0, 0006 <fib.bb_0006>
 7  0002:    load_imm r1, #0
 8  0003:    mov r0, r1
 9  0004:    pop r1
10  0005:    ret
11
1200000006 <fib.bb_0006>:
13  ; 11: n == 1 { 1 }
14  0006:    jmpne_imm r0, #1, 000b <fib.bb_000b>
15  0007:    load_imm r1, #1
16  0008:    mov r0, r1
17  0009:    pop r1
18  000a:    ret
19
200000000b <fib.bb_000b>:
21  ; 12: { fib(n - 1) + fib(n - 2) }
22  000b:    isub_imm r1, r0, #1
23  000c:    push r0
24  000d:    mov r0, r1
25  000e:    call 0000 <fib>
26  000f:    mov r1, r0
27  0010:    pop r0
28  0011:    isub_imm r0, r0, #2
29  0012:    call 0000 <fib>
30  0013:    iadd r0, r1, r0
31  0014:    pop r1
32  0015:    ret

在这个例子中，函数序言（push r1）和尾声（pop r1）的作用是保存和恢复被递归调用破坏的被调用者保存寄存器。

信息 - Purple garden 调用约定简述

寄存器被分为两个区域（类似于 x86 和 aarch64）：

  • r0..r{n} 被视为参数区，而 r0 同时也是返回值槽位，这些是调用者保存的，因此调用者若在此范围内有活跃值，必须在调用前后进行溢出和恢复
  • r{n+1}..r{max} 是被调用者保存寄存器，意味着被调用者不会破坏它们

fib 示例如果能把常量放入更合适的寄存器会更好，例如放入 r0 而不是通过 r1 间接访问，并且只在真正破坏 r1 的分支中压栈，但正如德国人所说，这是 zukunftsmusik :)

之前的 Vm 栈、{Push,Pop}{2,3} 及其处理程序

虚拟机的栈交互操作码定义如下：

RUST

1Push { src: u8, },
2Push2 { a: u8, b: u8, },
3Push3 { a: u8, b: u8, c: u8, },
4Pop { dst: u8, },
5Pop2 { a: u8, b: u8, },
6Pop3 { a: u8, b: u8, c: u8, },

在实现虚拟机时，我能想到的最简单的栈是使用 Vec<Value> 作为 Vm::spilled，预分配 4096 个值，并通过未检查的可变写入和读取进行压栈和弹栈：

RUST

 1Op::Push { src } => {
 2    // whole handler basically
 3    // equivalent to self.spilled.push(*r!(src))
 4
 5    let len = self.spilled.len();
 6
 7    if std::hint::unlikely(
 8        len + 1 > self.spilled.capacity()) {
 9        self.spilled.reserve(1);
10    }
11
12    unsafe {
13        // use mut pointer, advance to length
14        let dst = self.spilled.as_mut_ptr().add(len);
15        // write value in register src into spilled[len]
16        dst.write(*r!(src));
17        self.spilled.set_len(len + 1);
18    }
19}
20
21// ...
22
23Op::Pop { dst } => unsafe {
24    // whole handler basically equivalent to
25    // *r_mut!(dst) = self.spilled.pop();
26
27    let len = self.spilled.len();
28    let ptr = self.spilled.as_ptr();
29    debug_assert!(len >= 1);
30    r_mut!(dst) = ptr.add(len - 1).read();
31    self.spilled.set_len(len - 1);
32}

然后我还添加了融合的 {Push,Pop}{2,3}，将两次和三次栈交互合并为单个操作，且不进行边界检查，例如 Pop3：

RUST

1Op::Pop3 { a, b, c } => unsafe {
2    let len = self.spilled.len();
3    let ptr = self.spilled.as_ptr();
4    debug_assert!(len >= 3);
5    r_mut!(a) = ptr.add(len - 1).read();
6    r_mut!(b) = ptr.add(len - 2).read();
7    r_mut!(c) = ptr.add(len - 3).read();
8    self.spilled.set_len(len - 3);
9}

StackAllocator

从技术上讲，它只是一个分配到已分配缓冲区中的分配器，而这个缓冲区恰好位于栈上 :O（我想把它命名为 BufferAllocator 或类似的名字会更有用，但我用它来在栈上分配，所以也没问题）。例如，由（我们）调用者在栈上分配一个缓冲区，像这样：

RUST

1let mut buf = [MaybeUninit::<u8>::uninit(), 1024];
2let alloc = StackAlloc::new(&mut spillstack);
3let mut vec_on_the_stack = Vec::new_in(&alloc);
4vec_on_the_stack.push(787);
5vec_on_the_stack.push(161);
6vec_on_the_stack.push(256);
7let box_on_the_top = Box::new_in(0xDEADAFFE, &alloc)

StackAlloc 除了将传入的缓冲区保留作为生命周期标记之外，并不消耗或使用它，这样编译器就能确保数据不会比其中存放的数据活得更久。然后它有一个基址，称为 block，一个用于边界检查的 cap，以及用于推进的 top。

RUST

 1pub struct StackAlloc<'buf> {
 2    /// base ptr
 3    block: NonNull<u8>,
 4    /// buffer size
 5    cap: usize,
 6    /// next free byte
 7    top: Cell<usize>,
 8    /// marker for lifetime >:(
 9    _buf: PhantomData<&'buf mut [MaybeUninit<u8>]>,
10}

下一章将解释 Allocator trait 的实现，但概括来说，分配器的工作方式是：

  1. 找到 start，即 block+top；向上取整到对齐边界（16 字节，见下方“提示”）
  2. 找到分配所需的 end；将请求大小向上取整到对齐边界，再加上 start
  3. 检查缓冲区是否有足够空间；end > cap
  4. 将 top 推进到 end
  5. 返回 block+start，并切片到请求大小以省略对齐填充

提示 - 为什么对齐是 16 字节

StackAllocator 使用 16 字节对齐来解决多次对齐留下填充的问题，例如：偏移 0 处的 3 个字节，然后偏移 8 处对齐到 8 的 8 个字节。释放这 8 个字节会将 top 移回 8，但 3 个字节结束于 3，因此它们永远不会再被视为顶部分配，直到分配器重置前都会丢失。通过向上取整到 16 字节对齐，每个对齐到 16 的分配都从 top 开始，并且按相反顺序释放。

支持虚拟机栈的方法不使用这个抽象，因为虚拟机中的所有值都是统一的 8 字节。

实现 trait::Allocator

RUST

 1impl<'buf> StackAlloc<'buf> {
 2    pub fn new(buf: &'buf mut [MaybeUninit<u8>]) -> Self {
 3        let skip = buf.as_ptr().align_offset(UPROUND).min(buf.len());
 4        let block = NonNull::from(&mut buf[skip..]).cast::<u8>();
 5        Self {
 6            block,
 7            cap: (buf.len() - skip) / UPROUND * UPROUND,
 8            top: Cell::new(0),
 9            _buf: PhantomData,
10        }
11    }
12
13    #[must_use]
14    pub fn depth(&self) -> usize {
15        self.top.get()
16    }
17
18    pub fn reset(&mut self) {
19        self.top.set(0);
20    }
21
22    fn offset(&self, ptr: NonNull<u8>) -> usize {
23        ptr.as_ptr() as usize - self.block.as_ptr() as usize
24    }
25
26    fn is_top(&self, ptr: NonNull<u8>, layout: Layout) -> bool {
27        self.offset(ptr) + granules(layout.size()) == self.top.get()
28    }
29}
30
31// singular buffer borrow, so sending is safe
32unsafe impl Send for StackAlloc<'_> {}

depth、reset、offset 和 is_top 在分配器 trait 实现和虚拟机栈方法中使用：

RUST

 1
 2pub const UPROUND: usize = 16;
 3
 4#[inline(always)]
 5fn granules(size: usize) -> usize {
 6    size.next_multiple_of(UPROUND)
 7}
 8
 9unsafe impl Allocator for StackAlloc<'_> {
10    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
11        let base = self.block.as_ptr() as usize;
12        let start = (base + self.top.get()).next_multiple_of(layout.align()) - base;
13        let end = start
14            .checked_add(granules(layout.size()))
15            .filter(|&end| end <= self.cap)
16            .ok_or(AllocError)?;
17        self.top.set(end);
18        let ptr = unsafe { self.block.add(start) };
19        Ok(NonNull::slice_from_raw_parts(ptr, layout.size()))
20    }
21
22    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
23        if self.is_top(ptr, layout) {
24            self.top.set(self.offset(ptr));
25        }
26    }
27
28    unsafe fn grow(
29        &self,
30        ptr: NonNull<u8>,
31        old: Layout,
32        new: Layout,
33    ) -> Result<NonNull<[u8]>, AllocError> {
34        if self.is_top(ptr, old) && (ptr.as_ptr() as usize).is_multiple_of(new.align()) {
35            let end = self
36                .offset(ptr)
37                .checked_add(granules(new.size()))
38                .filter(|&end| end <= self.cap)
39                .ok_or(AllocError)?;
40            self.top.set(end);
41            return Ok(NonNull::slice_from_raw_parts(ptr, new.size()));
42        }
43        let moved = self.allocate(new)?;
44        unsafe {
45            ptr.copy_to_nonoverlapping(moved.cast(), old.size());
46            self.deallocate(ptr, old);
47        }
48        Ok(moved)
49    }
50
51    unsafe fn grow_zeroed(
52        &self,
53        ptr: NonNull<u8>,
54        old: Layout,
55        new: Layout,
56    ) -> Result<NonNull<[u8]>, AllocError> {
57        let block = unsafe { self.grow(ptr, old, new)? };
58        unsafe {
59            block
60                .cast::<u8>()
61                .add(old.size())
62                .write_bytes(0, new.size() - old.size());
63        }
64        Ok(block)
65    }
66
67    unsafe fn shrink(
68        &self,
69        ptr: NonNull<u8>,
70        old: Layout,
71        new: Layout,
72    ) -> Result<NonNull<[u8]>, AllocError> {
73        if !(ptr.as_ptr() as usize).is_multiple_of(new.align()) {
74            let moved = self.allocate(new)?;
75            unsafe {
76                ptr.copy_to_nonoverlapping(moved.cast(), new.size());
77                self.deallocate(ptr, old);
78            }
79            return Ok(moved);
80        }
81        if self.is_top(ptr, old) {
82            self.top.set(self.offset(ptr) + granules(new.size()));
83        }
84        Ok(NonNull::slice_from_raw_parts(ptr, new.size()))
85    }
86}

StackAllocator::{push,pop,push_all,pop_all}

现在，对于虚拟机溢出栈的交互，压入和弹出是必要的，为此我决定在分配器本身上实现这些，值得注意的是这些不对齐到 16，而是对齐到 T：

RUST

 1#[inline(always)]
 2pub fn push<T: Copy>(&self, value: T) -> bool {
 3    let top = self.top.get();
 4    debug_assert!(top.is_multiple_of(align_of::<T>()));
 5    let end = top + size_of::<T>();
 6    if std::hint::unlikely(end > self.cap) {
 7        return false;
 8    }
 9    unsafe { self.block.add(top).cast::<T>().write(value) };
10    self.top.set(end);
11    true
12}
13
14#[inline(always)]
15pub unsafe fn pop<T: Copy>(&self) -> T {
16    let top = self.top.get() - size_of::<T>();
17    self.top.set(top);
18    unsafe { self.block.add(top).cast::<T>().read() }
19}
20
21#[inline(always)]
22pub fn push_all<T: Copy, const N: usize>(&self, values: [T; N]) -> bool {
23    let top = self.top.get();
24    debug_assert!(top.is_multiple_of(align_of::<T>()));
25    let end = top + N * size_of::<T>();
26    if std::hint::unlikely(end > self.cap) {
27        return false;
28    }
29    let dst = unsafe { self.block.add(top).cast::<T>() };
30    for (i, value) in values.into_iter().enumerate() {
31        unsafe { dst.add(i).write(value) };
32    }
33    self.top.set(end);
34    true
35}
36
37#[inline(always)]
38pub unsafe fn pop_all<T: Copy, const N: usize>(&self) -> [T; N] {
39    let top = self.top.get() - N * size_of::<T>();
40    self.top.set(top);
41    let src = unsafe { self.block.add(top).cast::<T>() };
42    std::array::from_fn(|i| unsafe { src.add(i).read() })
43}

现在：带有实际栈空间的虚拟机栈

所以现在有了这些辅助方法，我可以用一个基于栈的实际溢出栈替换之前的处理器，首先分配一个栈缓冲区并在 Vm::run 中创建分配器：

RUST

1// 262144 B => 256KiB to spill results in /8 => 32768 vm values
2pub const SPILL_STACK: usize = 256 << 10;
3
4let mut spill_stack = [MaybeUninit::<u8>::uninit(); SPILL_STACK];
5let spilled = StackAlloc::new(&mut spill_stack);

然后这用于所有栈操作码处理器：

RUST

 1Op::Push { src } => unsafe {
 2    if std::hint::unlikely(!spilled.push(*r!(src))) {
 3        return Err(Anomaly::StackOverflow { pc });
 4    }
 5},
 6Op::Push2 { a, b } => unsafe {
 7    if std::hint::unlikely(!spilled.push_all([*r!(a), *r!(b)])) {
 8        return Err(Anomaly::StackOverflow { pc });
 9    }
10},
11Op::Push3 { a, b, c } => unsafe {
12    if std::hint::unlikely(!spilled.push_all([*r!(a), *r!(b), *r!(c)])) {
13        return Err(Anomaly::StackOverflow { pc });
14    }
15},
16Op::Pop { dst } => unsafe {
17    r_mut!(dst) = spilled.pop::<Value>();
18},
19Op::Pop2 { a, b } => unsafe {
20    let [below, top] = spilled.pop_all::<Value, 2>();
21    r_mut!(a) = top;
22    r_mut!(b) = below;
23},
24Op::Pop3 { a, b, c } => unsafe {
25    let [bottom, below, top] = spilled.pop_all::<Value, 3>();
26    r_mut!(a) = top;
27    r_mut!(b) = below;
28    r_mut!(c) = bottom;
29},

基准测试

对于本文所涉及的提交 1296116， 请参阅 bench/cross 获取基准测试源代码。

溢出栈更改

在 3 次预热运行后进行了超过 10 次运行，仅解释器（–no-jit）， 在 AMD Ryzen 7 3700X 上使用 -C target-cpu=native 的发布构建。之前是 提交 1296116 的父提交，之后是 1296116：

基准测试之前 (ms)之后 (ms)加速比
ackermann162.9129.11.26x / 20.8%
collatz104.389.01.17x / 14.7%
fib62.555.41.13x / 11.4%
gcd_sum93.982.71.14x / 11.9%
keyword_table19.318.01.07x / 6.5%
mandelbrot73.460.71.21x / 17.4%
n_body28.725.11.14x / 12.6%
primes_count50.946.41.10x / 8.8%
tak322.5267.51.21x / 17.0%

所以对于这些基准测试，我们快了 6-20%，对于这样一个非侵入性的更改来说，结果不错 :)

启用 JIT 并与 bun、luajit 对比：

这是为了快速比较 garden 在启用 jit 时与禁用 jit 时的速度（见上表）。我从这个表中移除了 python，因为它比 bun、luajit 和 garden 慢得多：

jit 还不支持双精度浮点数，所以未来会更快 O.O

我有跨语言基准测试可在 xnacly.github.io/purple-garden/ 获取， 下表来自提交 da1ad03：

工作负载garden (ms)bun (ms)luajit (ms)garden 对比最快
ackermann39.340.839.8最快
collatz20.722.214.1比 luajit 慢 1.47 倍
fib17.916.19.61比 luajit 慢 1.86 倍
gcd_sum27.433.745.1最快
keyword_table16.676228.4最快
mandelbrot89.114.26.45比 luajit 慢 13.81 倍
n_body49.716.45.02比 luajit 慢 9.90 倍
primes_count7.9617.76.09比 luajit 慢 1.31 倍
tak58.253.055.9比 bun 慢 1.10 倍

显示结果的网站嵌入在下方（这意味着每次新提交都会更新，不过它们是在启用 jit 的情况下测量的）：

来源：Hacker News · xnacly.me

部署工程教程实践#教程/实践#部署/工程