rat의 작은 레지스터 할당기

AI 요약

소형 컴파일러 rat의 백엔드가 중간 표현(IR)의 가상 레지스터(vreg)를 x86-64 물리 레지스터나 스택 슬롯에 배치하는 과정을 소개합니다. 기존 선형 스캔 할당기는 작동했지만 수정이 쌓이며 1,392줄로 커졌습니다. 저자는 기존 기능의 효과를 측정한 뒤 상당수를 걷어내고 584줄짜리 우선순위 기반 빈 패킹 할당기를 새로 작성했습니다. LLVM의 greedy allocator와 같은 계열이지만 복잡한 기능 대부분을 빼고도 더 나은 코드를 생성합니다.

레지스터 할당의 기본

값은 기록된 지점부터 마지막으로 읽히는 지점까지 살아 있습니다. 두 값의 생존 구간이 겹치지 않으면 같은 레지스터를 공유합니다. 한 시점에 살아 있는 값이 레지스터 수보다 많으면 일부 값을 메모리로 내보내는데, 이를 spill이라고 합니다. spill에는 저장과 재로드가 따릅니다. 최선의 배치를 찾는 문제는 NP-hard이므로 실제 할당기는 휴리스틱을 사용합니다.

호출 규약도 배치에 영향을 줍니다. 호출 중 caller-saved 레지스터는 덮어쓸 수 있고, 함수는 반환 전에 callee-saved 레지스터를 복구해야 합니다. 예시 함수 h는 g를 호출한 뒤에도 y를 사용합니다. 새 할당기는 y를 callee-saved인 rbx에 넣고 함수 앞뒤에 저장·복구 명령을 배치합니다. t는 rax에 남습니다. 복사 다섯 개 중 다섯 개가 사라집니다. 호출을 가로지르는 값을 callee-saved 레지스터에 넣으라는 별도 규칙은 없으며, 할당기 설계의 결과로 그렇게 배치됩니다.

다섯 단계로 진행하는 할당

함수마다 생존 구간 계산, 고정 레지스터 표시, 복사 합치기, 레지스터 선택, spill의 다섯 단계를 한 번씩 수행합니다. 할당기는 이미 배정한 레지스터를 빼앗지 않고, 생존 구간을 레지스터와 메모리로 나누지 않으며, 단계를 다시 실행하지 않습니다. 각 묶음은 생존 기간 전체에 하나의 레지스터나 스택 슬롯을 사용합니다.

명령어 i에 읽기 슬롯 2i와 쓰기 슬롯 2i+1을 둡니다. 복사 명령은 원본이 읽기 슬롯에서 끝나고 대상은 쓰기 슬롯에서 시작하므로 두 값이 레지스터를 공유할 수 있습니다. 그 복사는 나중에 제거할 수 있습니다. 일반 명령의 피연산자는 쓰기 슬롯까지 살아 있어 결과가 다른 피연산자를 덮어쓰지 않습니다. x86의 add처럼 첫 번째 피연산자에 결과를 덮어쓰는 2주소 명령은 먼저 복사 명령을 만들고, 이후 병합으로 불필요한 복사를 없앱니다.

생존 구간은 코드 순서대로 번호를 매긴 기본 블록을 따라 계산합니다. rat은 블록마다 비트셋을 두고 고정점 반복을 수행하는 대신 vreg 하나씩 처리합니다. 해당 값을 읽고 다시 쓰기 전에 정의하지 않는 블록에서 역방향으로 선행 블록을 따라가며 live-out을 표시하고, 값을 정의하는 지점에서 탐색을 멈춥니다. 이후 각 블록을 뒤에서 앞으로 훑으며 생존 구간을 만들고 spill 비용 가중치를 더합니다. 루프 깊이를 d라 할 때 각 정의와 사용은 3의 d제곱을 가중치에 더합니다. 따라서 일반 코드의 기여도는 1, 루프 안은 3, 중첩 루프 안은 9입니다. 루프 깊이는 최대 11까지 반영합니다.

생존 구간에는 구멍도 생깁니다. 코드에서 블록을 건너뛰어 값이 사용되지 않는 구간입니다. 이 구간을 생존 구간에서 빼면 다른 값이 해당 슬롯의 레지스터를 사용할 수 있습니다.

고정 레지스터와 복사 병합

rat은 레지스터 40개를 1부터 40까지 번호로 나타내고, U64 비트셋으로 슬롯마다 사용 중인 레지스터를 기록합니다. 호출 명령은 caller-saved 레지스터를 사용 중으로 표시합니다. 인자 전달, 반환값, 나눗셈처럼 코드에서 물리 레지스터를 직접 쓰는 경우도 생존 구간과 같은 역방향 순회 중에 표시합니다. 슬롯 64개마다 요약 마스크를 만들어 긴 생존 구간에서 사용 가능성을 빠르게 확인합니다.

복사로 연결된 vreg는 생존 구간이 겹치지 않으면 하나의 묶음(bundle)으로 합칩니다. 묶음의 생존 구간은 합쳐지고 가중치도 더해집니다. 묶음 안에서 일어나는 복사는 삭제 대상입니다. 복사 명령은 루프 깊이가 큰 순서로 처리하므로 뜨거운 루프의 복사를 먼저 합칩니다. 병합에는 union-find를 사용합니다. 두 묶음의 구간이 겹치는지 확인할 때 둘을 합쳐 256개가 넘는 구간을 만들면 병합을 건너뜁니다. 물리 레지스터와 vreg 사이의 복사는 병합 대신 해당 레지스터를 선호한다는 힌트로 기록합니다.

우선순위에 따른 레지스터 선택

묶음의 우선순위는 가중치를 슬롯 길이의 제곱근으로 나눈 값입니다. 짧고 자주 쓰이는 구간이 먼저 배정됩니다. 긴 구간은 뒤로 밀려 spill될 가능성이 커집니다. 제곱근을 사용해 긴 루프 카운터가 지나치게 불리해지는 현상을 줄입니다. 각 묶음의 생존 구간에서 사용 중인 레지스터를 모두 모은 뒤, 힌트 레지스터가 비어 있으면 먼저 선택합니다. 그렇지 않으면 caller-saved를 먼저, callee-saved를 나중에 살펴봅니다.

예시 함수 h에서는 x가 rdi, t가 rax에 배치됩니다. y는 호출 시점에 caller-saved 레지스터가 모두 사용 중이므로 rbx를 받습니다. rat은 실제로 사용한 callee-saved 레지스터만 함수 진입부에서 저장하고 반환 전에 복구합니다. Linux에서는 callee-saved XMM 레지스터가 없으므로 호출을 가로지르는 부동소수점 값은 스택으로 spill됩니다.

Spill과 코드 재작성

레지스터를 얻지 못한 묶음은 생존 구간 전체에서 스택 슬롯을 사용합니다. rat은 spill된 묶음을 시작점 순으로 정렬하고, 앞 묶음의 생존이 끝난 스택 슬롯을 재사용합니다. 명령을 다시 작성할 때 spill된 피연산자는 명령 전에 임시 레지스터로 불러오고, spill된 결과는 명령 뒤에 저장합니다. 정수에는 r10이나 r11, 부동소수점에는 xmm14나 xmm15를 임시 레지스터로 씁니다. 둘 다 사용 중이면 해당 명령에서 비어 있는 레지스터를 찾습니다. 레지스터와 spill 값 사이의 복사는 로드나 저장 명령 자체로 처리하며, 스택 인자는 스택 슬롯에서 직접 읽습니다.

14개 값이 동시에 살아 있는 예시에서는 사용할 수 있는 레지스터가 11개라 세 값이 spill됩니다. 코드 후반의 곱셈에서 x9를 다시 불러오지만, 저장과 재로드 사이에 r10을 덮어쓰는 명령이 없습니다. peephole 최적화가 재로드를 삭제하고, 그 스택 슬롯을 읽는 명령이 사라지면서 저장도 삭제합니다.

단순한 설계의 한계와 측정 결과

레지스터를 되찾아 재배치하지 않으므로 일찍 내린 선택이 나중까지 유지됩니다. 예시 sum 함수에서는 루프에서 자주 쓰는 주소 계산 값이 rax를 먼저 차지합니다. 그 결과 누적값 s는 rax 힌트를 잃고 rsi를 받으며, 배열 주소 a와 길이 n도 기존 힌트 대신 다른 레지스터를 받습니다. 저자는 이 사례를 가장 아쉬운 경우로 꼽습니다.

기존 할당기와 비교한 결과 할당기 코드는 58% 줄었습니다. 생성 명령어는 3.6%, 저장 명령어는 22% 감소했습니다. SQLite를 -O0으로 컴파일할 때 할당기 실행 시간은 77%, 전체 컴파일 시간은 43% 줄었습니다.

기존 기능을 하나씩 끄며 측정한 결과, 복사 병합과 복사 힌트는 명령어를 약 3분의 1 줄여 새 할당기에 남겼습니다. 생존 구간의 구멍은 10%, 사용 가중치에 따른 spill 선택은 5.5%, 물리 레지스터 힌트는 1%를 줄였습니다. 반면 spill된 구간을 두 번째로 다시 시도하는 기능은 할당기 시간의 37%를 썼지만 명령어 감소 효과가 0.01%에 그쳐 제거했습니다. 재계산으로 재로드를 대체하는 rematerialization과 spill 슬롯 캐시도 측정 가능한 효과가 없어 제외했습니다. 저자가 얻은 교훈은 기존 코드를 옮기기 전에 기능별 효과를 먼저 측정하라는 것입니다.

Lobsters 반응

• @xnacly — 제가 직접 만든 JIT의 레지스터 할당기는 머리를 아프게 했습니다. 특히 liveness 분석을 제대로 구현하기가 정말 어려웠습니다. 이 글은 제가 막연하게만 알던 개념 몇 가지를 잘 정리해 주는 것 같습니다. 제 할당기는 기능이 훨씬 적고 컴파일 도중에 할당만 해서 필요할 때 레지스터를 빼앗고 spill해야 합니다. 나중에 개선해 볼지도 모르지만, 지금은 일단 작동한다는 사실만으로도 기쁩니다 :|

원문: hexrat.cc / 번역·요약: Trawling