Cranelift
Cranelift은 Wasmtime이라는 wasm을 부라우저 밖에서도 실행하기 위한 런타임 도구의 JIT Compiler이다. (AOT 로도 쓰일 수 있긴 하다.) MLIR을 기반으로 cranelift ir 이라는 dialect를 이용하여 컴파일을 진행한다. 최근에는 이 MLIR을 정말 많이 이용하는 것 같은데, 이를 정리한 이전 글을 참고하면 된다. 이 컴파일러는 최근 최적화 이식 연구를 하면서 찾아보게 되었다. 현재 기존 컴파일러에 존재하는 최적화를 일반화 하여 올바른지 증명하고, E-graph로 최적화 모듈을 만들어서 다른 컴파일러에 이식을 하는 계획을 세웠다. 이 cranelift는 e-graph로 이미 최적화를 잘 구현했기 때문에 찾아보고 공부하게 되었다.
aegraphs
cranelift는 e-graph를 사용한 최적화를 도입한 첫 컴파일러로 단일 최적화 pass로 최적화를 구현하기 위해 e-graph를 채택하였다. 이 최적화 시스템을 aegraphs로 부른다. LLVM와 같은 경우는 수십 개의 최적화 pass로 이루어졌지만 cranelift는 e-graph의 equality saturation을 잘 활용하여 단일 최적화 pass를 구현하였다. aegraphs는 크게 두 가지 challenge를 가지고 있다.
- CFG와 e-graph를 어떻게 결합할 것인가?
- e-graph의 cycle을 어떻게 해결할 것인가?
CFG + e-graph
기존의 e-graph는 거의 한 basic block안에서만 적용이 가능하다. e-graph를 활용한 컴파일러 최적화를 설명한 이전 글에서도 한 basic block에 존재하는 instruction에 대한 최적화만을 설명하고 있다. 하지만 컴파일러에는 한 basic block의 범위를 넘어서는 중요한 최적화들이 많다. (GVN, LICM 등등..) 그렇기에 CFG를 고려하여 여러 basic block범위를 아우르는 최적화들을 구현하고자 하는 것을 목표로 한다.
aegraphs는 먼저 함수의 CFG를 그린 후에, 이를 E-graph를 그린다. 전체 코드에서 CFG skeleton만을 남기고 e-graph로 변환을 한다. CFG skeleton은 모든 block과 block parameters, side-deffect operator, block terminator 만을 남긴 것이다. 이 들은 함부로 변형할 경우 전체 프로그램의 동작을 해치게 때문에 이 들은 남기고 다른 코드들은 전부 e-graph로 변환하게 된다. 아래와 같은 프로그램이 있다고 할 때,
이는 왼쪽의 CFG skeleton, 오른쪽의 e-graph로 쪼개지게 된다.
이 후 e-graph에 대해서 모든 최적화를 진행하게 되면 결국, 함수 전체에 대해서 최적화를 진행한 것이기 때문에 basic block을 넘어선 최적화도 수행할 수 있게 된다. 물론 현재 이 정도에 대해서는 loop 최적화 등 프로그램의 실행흐름을 변경하는 최적화는 구현할 수 없다는 한계는 존재한다.
Cycle in E-graph
e-graph를 컴파일러 최적화에 사용할 때 큰 문제점 중 하나는, e-graph가 rewrite rule의 방향까지 표현해주지 못한다는 것이다. 얘를 들면 X + 0 -> X라는 최적화가 있다고 할 때, e-graph는 equivalane relation만 표현하기 때문에 X + 0 과 X가 같다라는 것만 표현할 뿐, X가 X + 0보다 비용이 싼 프로그램이라는 것을 표현하지 못한다. 또한 아래처럼 e-graph로 표현 됐을 때, 이는 X, X + 0, X + 0 + 0, X + 0 + 0+ 0, … 와 같이 무한한 프로그램이 같다라는 cycle이 생기게 된다.
이런 문제를 해결하기 위해서, aegraphs는 Union node 라는 새로운 개념의 노드를 e-graph에 추가하였다. 이는 프로그램에 최적화를 적용할 때, Union node를 만들면서 union node가 가리키는 대상은 이미 rewrite rule이 적용됐음을 표현 하면서 다시 rewrite가 일어나지 않도록 하면서 cycle을 방지한다.
마무리
cranlift는 e-graph를 적극적으로 활용해서 최적화를 구현하였으며 이를 통해서 rewrite rule을 작성하면 자동으로 최적화가 추가되는 시스템을 만들었다. 이는 현재 우리가 매우 지향하고 싶은 시스템으로, 올바른 최적화임을 증명한 rewrite rule만을 작성해서 넣으면 최적화 버그가 없는 최적화 모듈이 만들어지는 시스템을 쉽게 구축할 수 있기 때문이다. 현재 cranelift는 링크와 같이 DSL로 최적화를 넣어주고 있다.
현재 aegraphs는 GVN, LICM, Constant Folding, Alias Analysis의 최적화를 지원하고 있다. 이는 GVN과 같이 여러 basic block을 넘어선 최적화를 지원하는 점에서 매우 주목할 만하다. 앞으로 더 발전시켜야 할 점은 side effect가 고려되는 최적화를 적용하지 못하고, e-graph에서 최종적으로 최적화된 프로그램을 추출하는 부분이 heuristic으로 동작하는 점, cyclic input에 대해서는 지원하지 못하는 점을 발전시킨다면 e-graph를 실제 현대 컴파일러에서도 사용할 수 있을 것이다.