(nightly) Rust로 만든 꼬리 호출 인터프리터

지난주에 저는 become 키워드를 사용해 꼬리 호출 인터프리터를 작성했습니다. 이 키워드는 최근 nightly Rust에 추가되었고 (7개월 전이면 최근이라고 할 수 있겠죠?).

생각보다 훨씬 즐거운 경험이었고, 그 결과 나온 VM은 제가 이전에 작성한 Rust 구현 그리고 손으로 작성한 ARM64 어셈블리보다도 더 빠릅니다. 최근에는 꼬리 호출 기반 기법이 큰 화제가 되고 있으니(이 개요 참고), 이것은 단순하지만 만만치 않은 시스템을 구현해 본 제 경험담이라고 생각하시면 됩니다.

---

집에서 이 주제를 계속 따라오고 계신 분들을 위해 설명하자면, 이것은 Uxn CPU의 고성능 에뮬레이션을 탐구해 온 제 작업의 최신편입니다. 이 CPU는 Hundred Rabbits 생태계에서 다양한 애플리케이션을 실행합니다.

전체 여정을 읽고 싶다면, 목록은 다음과 같습니다.

• Raven 프로젝트 글, 제 원래 Rust 구현
• 컴파일러 이기기, Rust 코드보다 빠른 ARM64 어셈블리 구현을 작성한 글
• 우리 테스트 스위트의 아름다움에 이끌리다, 1년 뒤에 코드를 다시 살펴보며 테스트와 CI를 개선한 글
• raven-uxn을 위한 x86-64 백엔드, Claude Code의 도움을 받아 ARM64 어셈블리 구현을 x86으로 이식한 글

LLM을 실험한 일은 논란의 대상이 되기도 했지만, 그건 놀랍지 않았습니다. 이번 꼬리 호출 코드는 전부 사람이 직접 작성했다는 점을 기쁘게 말씀드릴 수 있고, 새 백엔드는 약간의 성능 손해만 감수하면 x86 어셈블리 백엔드를 대체해서 사용할 수 있습니다.

(이 블로그 글 역시 제 개인 기준에 따라 전부 사람이 직접 작성했습니다.)

다음 몇 절은 이전 작업을 요약하므로, 이미 다 읽으셨다면 가볍게 훑고 바로 Rust에서의 꼬리 호출로 넘어가셔도 됩니다.

---

Uxn 에뮬레이션의 기초

Uxn은 256개의 명령어를 가진 단순한 스택 머신입니다. CPU 전체 공간은 64K가 조금 넘고, 몇 가지 메모리로 나뉩니다.

• 각각 인덱스 바이트를 가진 256바이트 스택 두 개
• 데이터와 프로그램 텍스트를 모두 담는 65536바이트 RAM
• 2바이트 프로그램 카운터
• 주변장치에 사용하는 256바이트의 "디바이스 메모리"

가장 단순한 에뮬레이터는 프로그램 카운터가 가리키는 RAM에서 바이트 하나를 읽고, 그다음 명령어를 호출합니다(이 과정에서 프로그램 카운터가 갱신될 수 있습니다).

fn run(core: &mut Uxn, dev: &mut Device, mut pc: u16) -> u16 {
    loop {
        let op = core.next(&mut pc);
        let Some(next) = core.op(op, dev, pc) else {
            break pc;
        };
        pc = next;
    }
}

impl Uxn {
    fn op(
        &mut self, op: u8, dev: &mut Device, pc: u16
    ) -> Option<u16> {
        match op {
            op::BRK => self.brk(pc),
            op::INC => self.inc::<0b000>(pc),
            op::POP => self.pop::<0b000>(pc),
            op::NIP => self.nip::<0b000>(pc),
            op::SWP => self.swp::<0b000>(pc),
            // ... etc
            op::ORA2kr => self.ora::<0b111>(pc),
            op::EOR2kr => self.eor::<0b111>(pc),
            op::SFT2kr => self.sft::<0b111>(pc),
        }
    }
}

명령어는 256개이고, 그중 상당수는 플래그로 매개변수화됩니다. 다음은 스택의 맨 위 바이트를 증가시키는 INC 명령어입니다.

impl Uxn {
    pub fn inc<const FLAGS: u8>(&mut self, pc: u16) -> Option<u16> {
        let mut s = self.stack_view::<FLAGS>();
        let v = s.pop();
        s.push(v.wrapping_add(1));
        Some(pc)
    }
}

모든 opcode 구현은 메인 op 함수 안으로 인라인되지만, 개선의 여지가 있습니다. 일부 값은 레지스터가 아니라 메모리에 저장되고, 메인 op 선택 분기는 예측하기 어렵습니다.

어셈블리에서의 threaded code

어셈블리 구현에서는 대신 threaded code, 그중에서도 특히 token threading을 사용할 수 있습니다. CPU 상태 전체를 레지스터에 저장하고, 각 명령어의 끝에서 다음 명령어로 점프합니다.

; x0 | stack pointer
; x1 | stack index
; x4 | ram pointer
; x5 | program counter
; x8 | opcode table
_INC:
    ldrb w9, [x0, x1]           ; read the byte from the top of the stack
    add w9, w9, #1              ; increment it
    strb w9, [x0, x1]           ; write it back
    ldrb w9, [x4, x5]           ; load the next opcode from RAM
    add x5, x5, #1              ; increment the program counter
    and x5, x5, #0xffff         ; wrap the program counter
    ldr x10, [x8, x9, lsl #3]   ; load the opcode implementation address
    br x10                      ; jump to the opcode's implementation

이 방식은 디스패치 연산을 모든 opcode에 분산시켜서, 분기 예측기가 프로그램 안의 opcode 시퀀스를 더 쉽게 학습하게 해 줍니다. 전체적인 속도 향상은 상당했습니다. ARM64에서는 40-50% 빨라졌고, x86-64에서는 대략 2배 빨라졌습니다.

불행히도, 이 방식은 약 2000줄의 코드를 유지해야 하고, 엄청나게 unsafe합니다. x86 포트에서는 범위를 벗어난 쓰기 버그를 넣어 버렸는데, 그 결과 디바이스 RAM 바깥의 몇 바이트를 덮어썼습니다. 유일한 증상은 퍼저가 아주 특정한 프로그램을 실행한 뒤 종료할 때 segfault가 난다는 것이었습니다.

그렇다면 어떻게 해야 할까요?

Rust에서의 꼬리 호출

우리는 어셈블리 구현과 같은 동작, 즉 VM 상태는 레지스터에 저장하고 디스패치는 각 opcode 끝에서 수행하는 방식을 원합니다. 하지만 모든 명령어를 손으로 어셈블리로 쓰고 싶지는 않습니다. 다행히도 희망이 있습니다.

핵심 아이디어는 거의 틀림없이 여러 번 독립적으로 다시 발명되었겠지만, 저는 Massey Meta Machine 글에서 처음 접했습니다. 정말 사고의 폭을 넓혀 주는 글이었습니다.

구성 요소는 두 가지입니다.

• 프로그램 상태를 함수 인자에 저장하고, 시스템 호출 규약에 따라 이것이 레지스터에 매핑되도록 하기
• 각 함수를 다음 함수를 호출하는 것으로 끝내기

이것은 오늘 당장 Rust로도 작성할 수 있습니다. 다음은 inc 함수입니다.

const TABLE: FunctionTable = FunctionTable([
    brk,
    inc::<0b000>,
    pop::<0b000>,
    nip::<0b000>,
    swp::<0b000>,
    rot::<0b000>,
    dup::<0b000>,
    // ...etc
]);

fn inc<'a, const FLAGS: u8>(
    stack_data: &'a mut [u8; 256],
    stack_index: u8,
    rstack_data: &'a mut [u8; 256],
    rstack_index: u8,
    dev: &'a mut [u8; 256],
    ram: &'a mut [u8; 65536],
    mut pc: u16,
    vdev: &mut dyn Device,
) -> (Uxn<'a>, u16) {
    let mut core = Uxn {
        stack: Stack {
            data: stack_data,
            index: stack_index,
        },
        ret: Stack {
            data: rstack_data,
            index: rstack_index,
        },
        dev,
        ram,
    };
    match core.inc::<FLAGS>(pc) {
        Some(pc) => {
            let op = core.next(&mut pc);
            TABLE.0[op as usize](
                core.stack.data,
                core.stack.index,
                core.ret.data,
                core.ret.index,
                core.dev,
                core.ram,
                pc,
                vdev,
            )
        }
        None => (core, pc)
    }
}

기존 Uxn opcode 구현을 재사용하고 싶기 때문에, 함수 시작 부분에서 core: Uxn 객체를 다시 구성하고 inc 함수를 호출한 다음, 다음 연산을 호출할 때 다시 그것을 해체합니다. 보일러플레이트가 많고, 그냥 &mut Uxn 인자를 넘기고 싶어지지만, 그렇게 하면 "상태가 레지스터에 저장된다"는 이점을 잃게 됩니다. 이 보일러플레이트는 나중에 매크로로 없애겠습니다.

안타깝게도 이 구현에는 문제가 있습니다.

thread 'snapshots::tailcall::mandelbrot' (67889685) has overflowed its stack
fatal runtime error: stack overflow, aborting
error: test failed, to rerun pass `-p raven-varvara --test snapshots`

릴리스 빌드에서도 컴파일러는 스택을 최적화해 없애지 못했습니다. 연산을 점점 더 많이 실행할수록 스택은 더 깊어지고, 결국 피할 수 없이 overflow가 납니다.

컴파일러가 bl(branch-and-link) 명령어 대신 br(register로 분기) 명령어를 생성하도록 해야 하고, 더 중요하게는 스택에 지속적으로 남는 공간을 전혀 할당하지 않도록 해야 합니다. 다시 말해, 꼬리 호출이 필요합니다.

nightly Rust에서는 단어 하나만 바꾸면 됩니다.

match core.inc::<FLAGS>(pc) {
         Some(pc) => {
             let op = core.next(&mut pc);
-            TABLE.0[op as usize](
+            become TABLE.0[op as usize](
                 core.stack.data,
                 core.stack.index,
                 core.ret.data,
                 core.ret.index,
                 core.dev,
                 core.ram,
                 pc,
                 vdev,
             )

이 변경을 적용하면, Rust 컴파일러는 다음과 같은 보장을 제공합니다.

함수를 꼬리 호출할 때는 그 함수의 스택 프레임이 스택에 추가되는 대신, 호출자의 스택 프레임이 직접 피호출자의 것으로 대체된다.

이것으로 끝입니다. 모든 것이 잘 동작합니다! 글도 끝!

구현 세부사항

좋아요, 좋아요, 아직 조금 더 할 말이 있습니다.

먼저, 보일러플레이트를 제거하는 매크로를 약속했죠. 늘 그렇듯, 보기만 해도 무시무시한 물건입니다.

macro_rules! tail_fn {
    ($name:ident $(::<$flags:ident>)?) => {
        tail_fn!($name $(::<$flags>)?[][vdev: &mut dyn Device]);
    };
    ($name:ident $(::<$flags:ident>)?($($arg:ident: $ty:ty),*)) => {
        tail_fn!($name $(::<$flags>)?[$($arg: $ty),*][]);
    };
    ($name:ident $(::<$flags:ident>)?[$($arg0:ident: $ty0:ty),*][$($arg1:ident: $ty1:ty),*]) => {
        extern "rust-preserve-none" fn $name<'a, $(const $flags: u8)?>(
            stack_data: &'a mut [u8; 256],
            stack_index: u8,
            rstack_data: &'a mut [u8; 256],
            rstack_index: u8,
            dev: &'a mut [u8; 256],
            ram: &'a mut [u8; 65536],
            pc: u16,
            $($arg0: $ty0),*
            $($arg1: $ty1),*
        ) -> (UxnCore<'a>, u16) {
            let mut core = UxnCore {
                stack: Stack {
                    data: stack_data,
                    index: stack_index,
                },
                ret: Stack {
                    data: rstack_data,
                    index: rstack_index,
                },
                dev,
                ram,
            };
            match core.$name::<$($flags)?>(pc, $($arg0),*) {
                Some(mut pc) => {
                    let op = core.next(&mut pc);
                    become TABLE.0[op as usize](
                        core.stack.data,
                        core.stack.index,
                        core.ret.data,
                        core.ret.index,
                        core.dev,
                        core.ram,
                        pc,
                        $($arg0),*
                        $($arg1),*
                    )
                }
                None => (core, pc),
            }
        }
    };
}

(이제는 실제 구현에서 가져온 것이기 때문에, 앞서 글 초반의 단순화된 코드와 비교하면 몇몇 타입이 조금 다릅니다.)

이 매크로는 매우 어색하지만, 세 종류의 함수를 선언할 수 있게 해 줍니다.

• 전설적인 것들
• 무수히 많은 것들
• 아주 가는 낙타털 붓으로 그린 것들

잠깐, 아니, 목록이 잘못됐네요. 세 종류의 함수는 다음과 같습니다.

• 인자가 없는 기본 함수
• FLAGS가 있는 함수
• FLAGS와 추가 인자가 있는 함수

세 경우는 각각 이렇게 생겼습니다.

tail_fn!(brk); // bare function
tail_fn!(inc::<FLAGS>); // function with flags
tail_fn!(dei::<FLAGS>(dev: &mut dyn Device)); // flags and arguments

이 매크로를 오래 붙들고 고민하실 필요는 없습니다. 여기서는 완전히 "컴파일되면 잘 되는 것"의 영역에 들어와 있으니까요. 그리고 중요한 점 하나는, 이것 역시 여전히 100% safe Rust라는 것입니다. #![forbid(unsafe_code)] 속성은 여전히 한 번도 발동하지 않습니다.

코드 생성 메모

컴파일러는 함수를 인라인하고 핵심 연산만 남기도록 정리하는 일을 꽤 잘합니다. UxnCore를 만들고 다시 해체하는 보일러플레이트는 완전히 최적화되어 사라집니다.

0000000100039c50 <raven_uxn::tailcall::inc::<0>>:
    and     x8, x0, #0xff           ; mask stack pointer
    ldrb    w9, [x21, x8]           ; read byte
    add     w9, w9, #1              ; increment value
    strb    w9, [x21, x8]           ; write back value
    and     x8, x2, #0xffff         ; mask pc
    ldrb    w8, [x24, x8]           ; look up next opcode
    adrp    x9, 0x100279000         ; load table base
    add     x9, x9, #3776           ; offset??
    ldr     x3, [x9, x8, lsl  #3]   ; get jump target
    add     w2, w2, #1              ; increment program counter
    br      x3                      ; jump!

제가 보기에 손으로 작성한 구현과의 주요 차이는 두 가지입니다.

• 값을 수정한 뒤가 아니라 사용하기 전에 레지스터 마스킹을 수행합니다. 예를 들어 프로그램 카운터(x2)는 조회에 사용할 때 마스킹되지만, 증가한 뒤에는 마스킹되지 않습니다.
• 테이블 베이스를 레지스터에 저장해 두는 대신 오프셋에서 로드합니다.

후자는 테이블도 인자 목록을 통해 함께 전달하도록 만들면 고칠 수 있습니다. 이렇게 하면 x86에서는 성능이 개선되지만, ARM64에서는 별로 차이가 없는 것 같습니다.

성능 결과

성능 이야기가 나왔으니, 실제로 얼마나 잘 나올까요?

저에게는 주요 벤치마크가 두 가지 있습니다.

• 피보나치 수를 순진하게 계산하는 마이크로벤치마크 프로그램
• Mandlebrot 프랙탈을 렌더링하는 더 큰 프로그램

제 노트북(M1 Macbook)에서, 이번에는 제가 더 이상 컴파일러를 이기지 못한다는 사실을 기쁘게 보고합니다. 꼬리 호출 인터프리터는 두 벤치마크 모두에서 손으로 작성한 어셈블리를 여유 있게 앞섭니다.

(모든 시간은 밀리초이며, 작을수록 빠릅니다.)

이제 계절에 맞지 않는 차를 크게 한 모금 마시고 x86에서 시험해 봅시다—

—이럴 수가. VM보다는 빠르지만, 어셈블리 백엔드에는 눈에 띄게 뒤지고 있습니다(특히 Fibonacci 마이크로벤치마크에서 그렇습니다). 여기서는 무슨 일이 벌어지고 있는 걸까요?

가장 단순한 opcode인 INC의 생성 코드를 먼저 봅시다.

_RINvNtCs7kxjQDHw3ed_9raven_uxn8tailcall3incKh0_EB4_:
    movzx  eax,dil                  ; mask stack index
    inc    BYTE PTR [r13+rax]       ; increment byte on stack
    movzx  eax,cx                   ; mask pc register
    movzx  eax,BYTE PTR [rdx+rax]   ; read next opcode from RAM
    inc    ecx                      ; increment pc reigster
    mov    rax,QWORD PTR [r8+rax*8] ; read jump target from table
    jmp    rax                      ; jump to target
    ; followed by `int3` operations to pad to 32-byte alignment

이 구현은 대체로 괜찮아 보입니다. 필요한 최소한의 읽기와 쓰기만 수행하고 있고, 기본적으로 제가 기대하는 모습입니다. 실제로 주소를 통해 바이트를 증가시키는 방식은, 그 바이트를 따로 읽고 다시 쓰는 제 어셈블리보다 더 효율적일 수도 있습니다.

(그리고 앞서 강조하지는 않았지만, 이 함수들을 extern "rust-preserve-none"으로 선언하는 것은 x86 구현에서 매우 중요합니다. 기본 호출 규약은 우리의 모든 인자를 담기에 충분한 레지스터를 사용하지 않아서, 엄청난 오버헤드가 추가됩니다.)

어쨌든 INC는 무난합니다. 이제 스택 위의 16비트 값 두 개를 더하는 ADD2를 봅시다. 먼저 어셈블리 백엔드에서 손으로 조정한 구현을 보여드리겠습니다.

_ADD2:
    movzx  eax,BYTE PTR [rbx+r12*1]     ; read byte from data stack
    dec    r12b                         ; decrement data pointer
    movzx  ecx,BYTE PTR [rbx+r12*1]     ; read byte from data stack
    dec    r12b                         ; decrement data pointer
    shl    ecx,0x8                      ; build 16-bit value
    or     eax,ecx
    movzx  ecx,BYTE PTR [rbx+r12*1]     ; read byte from data stack
    lea    r11,[r12-0x1]                ; get next byte address
    and    r11,0xff                     ; mask index byte
    movzx  edx,BYTE PTR [rbx+r11*1]     ; read byte from data stack
    shl    edx,0x8                      ; build 16-bit value
    or     ecx,edx
    add    ecx,eax                      ; do addition
    mov    BYTE PTR [rbx+r12*1],cl      ; write byte to data stack
    shr    ecx,0x8                      ; shift 16-bit value
    mov    BYTE PTR [rbx+r11*1],cl      ; write byte to data stack
    movzx  eax,BYTE PTR [r15+rbp*1]     ; read next opcode
    inc    bp                           ; increment pc
    lea    rcx,[rip+0x22e5c7]           ; read jump address
    jmp    QWORD PTR [rcx+rax*8]        ; jump

이 구현은 79바이트이며, 읽기와 쓰기 횟수도 딱 최소한입니다. 데이터 스택에서 바이트 4번 읽기 + 바이트 2번 쓰기, RAM에서 다음 opcode를 얻기 위한 바이트 1번 읽기, 점프 대상 주소를 얻기 위한 점프 테이블에서의 qword 1번 읽기입니다.

반면, 컴파일된 꼬리 호출 구현은 이렇습니다.

_RINvNtCs7kxjQDHw3ed_9raven_uxn8tailcall3addKh1_EB4_:
    push   rbp                      ; spill rbp to the stack
    mov    QWORD PTR [rsp-0x8],r11  ; spill r11 to the stack too?
    mov    r11,r9                   ; do the Register Shuffle (??)
    mov    r9,r8
    mov    r8d,ecx
    mov    rax,r13
    movzx  r10d,dil                 ; get index as byte
    movzx  edi,BYTE PTR [r13+r10]   ; read byte from stack
    lea    ebx,[r10-0x1]            ; get next index
    movzx  ebx,bl                   ; mask index to byte
    lea    r13d,[r10-0x2]           ; precompute another index
    movzx  ebp,BYTE PTR [rax+rbx*1] ; read byte from stack
    shl    ebp,0x8                  ; build 16-bit value
    or     ebp,edi
    movzx  edi,r13b                 ; mask index byte
    movzx  r13d,BYTE PTR [rax+rdi]  ; read byte from stack
    add    r10b,0xfd                ; compute another address
    movzx  ecx,r10b                 ; mask index byte
    movzx  ebx,BYTE PTR [rax+rcx]   ; read byte from stack
    shl    ebx,0x8                  ; build 16-bit value
    or     ebx,r13d
    add    ebx,ebp                  ; do addition
    mov    BYTE PTR [rax+rcx*1],bh  ; write byte to stack
    mov    BYTE PTR [rax+rdi*1],bl  ; write byte to stack
    movzx  ecx,r8w                  ; mask pc word
    movzx  ecx,BYTE PTR [rdx+rcx*1] ; read opcode
    inc    r8d                      ; increment pc
    mov    r10,QWORD PTR [r9+rcx*8] ; get next jump address
    mov    r13,rax                  ; do the Reverse Register Shuffle (??)
    mov    ecx,r8d
    mov    r8,r9
    mov    r9,r11
    mov    r11,QWORD PTR [rsp-0x8]  ; restore spilled r11
    mov    rax,r10
    pop    rbp                      ; restore pushed rbp
    jmp    rax                      ; jump
    ; followed by `int3` padding to a 32-byte boundary

이 코드는 121바이트이고, 더 걱정스러운 점은 64비트 레지스터 두 개 전체를 스택에 spill했다가 복원한다는 것입니다. 이건 말하자면 _정말 좋지 않은 코드 생성_처럼 보입니다. 한편으로는 이 기능이 rustc의 미완성 nightly 기능이니 이해할 수 있습니다. 다른 한편으로는, 설령 rustc 쪽이 미성숙하더라도 LLVM 백엔드가 이것을 더 잘 최적화하지 못하는 점은 의외입니다.

글이 점점 길어지고 있어서 여기서 너무 많은 추측을 늘어놓지는 않겠지만, 몇 가지 관찰만 적어 두겠습니다.

• 평범한 C 코드는 Clang이 더 최적에 가까운 형태로 컴파일합니다.

(gcc가 도대체 무슨 짓을 하고 있는지는 굳이 이야기하지 맙시다.)

• 마찬가지로 평범한 Rust 코드도 제가 실제로 본 것만큼 심각하지는 않지만, 여전히 불필요하게 rbp를 스택에 spill합니다.

한 가지 더

WebAssembly도 꼬리 호출을 지원하고, Raven은 WebAssembly로 컴파일하는 것도 지원합니다. 꼬리 호출 인터프리터가 네이티브 성능과 단순 VM 인터프리터에 비해 어떤 성능을 보일지 궁금했습니다.

벤치마크할 수 있는 것은 Mandelbrot 예제뿐입니다. Fibonacci 프로그램은 웹 타이머 해상도의 한계 때문에 너무 빨리 끝납니다. 결과는 다음과 같습니다.

놀랍게도, 결과는 끔찍합니다. Firefox에서는 1.2배 느리고, Chrome에서는 3.7배 느리며, wasmtime에서는 4.6배 느립니다. 좋은 어셈블리를 만들어 내는 패턴이 WASM 스택 머신에는 잘 맞지 않고, JIT도 이를 최적의 머신 코드로 내릴 만큼 똑똑하지 못한 것 같습니다.

그래도 전통적인 VM 구현에서 wasmtime은 인상적인 성능을 보여 주었습니다. 네이티브 Rust 빌드와 비교해도 몇 퍼센트포인트 차이밖에 나지 않았습니다.

(이 테스트는 모두 제 M1 Max 노트북에서 수행했습니다. wasmtime은 e9e1665c5에서 빌드했고, Firefox는 149.0, Chrome은 146.0.7680.178이었습니다.)

결론

tailcall 인터프리터 PR은 병합되었고, 0.3.0 릴리스에 포함되어 배포되었습니다. 활성화하면 ARM64 시스템에서는 기본값이고, x86-64 시스템에서는 두 번째 선택지입니다(native 기능이 활성화되어 있지 않은 경우).

x86과 WASM 성능을 개선하는 팁이 있다면 정말 궁금합니다.

이메일이나 소셜미디어로 연락 주세요.