연구를 하면서 컴파일러는 대부분 LLVM을 대상으로 진행했기에 현재는 C 코드보다 LLVM IR이 익숙해진 정도로 LLVM IR을 많이 보고 있었다. 최근 LLVM에서 MLIR level을 추가하고, 여기에서 미리 여러 최적화를 수행하자는 토론이 이루어지는 것을 보고, MLIR에 대해서 처음 알게 되었다. 알고보니 AWS에서 자체적으로 만드는 컴파일러도 MLIR을 사용하며, cranelift도 MLIR을 사용했으며, 최근 많은 컴파일러들이 MLIR을 사용한다는 것을 알게 되었다. 왜 이제서야 알게 되었나라고 반성했지만, 이제라도 알게 되었음을 다행으로 생각한다. 최근 최적화 이식 연구를 하게 되면서 LLVM에 있는 최적화를 어떻게 최신 컴파일러에 이식할 수 있을까를 고민하면서 최신 컴파일러들의 MLIR을 다뤄볼 기회가 생겨서 다행히 쉽게 이해해볼 수 있게 되었다.
MLIR은 여러 다양한 컴파일러에서 쉽게 infrastructure를 구축할 수 있기 위해 만들어졌다. 이는 여러 현대 컴파일러들이 각자만의 IR을 사용하여 컴파일하는 것이 성공적으로 받아들여졌기 때문이다. 대표적인 컴파일러인 GCC, LLVM 역시 각자만의 IR을 만들어 미들엔드에서 여러 최적화나 코드 번역을 시도하는 계층적 구조를 가지고 있다. 이런 상황에서 매번 컴파일러마다 자신의 IR을 바닥부터 구축하는 것은 매우 귀찮기 때문에, 기본적인 뼈대를 가진 MLIR을 만들어내고 컴파일러를 만들 때 각자 자신의 환경에 맞는 dialect를 붙임으로써 쉽게 IR 계층을 만드는 것을 목표로 한다. MLIR은 기본적으로 미들엔드를 지향하지만 LLVM IR보다는 한 단계 높은 수준을 유지하려고 한다. LLVM IR의 경우, 좀 더 머신 레벨과 가까운 행동을 많이 한다. REGMEM 이라는 단계에서는 각 아키텍처에 맞는 레지스터가 할당된 상태를 표현하기도 하면서 여러 머신 레벨과 같은 표현도 할 수 있게 만들어져 있다. MLIR은 (LLVM 기준) 미들엔드의 처음에서 활약하며 중간 수준에서 여러 동작(최적화, 번역)을 유도한다. 최근에는 C/C++을 이용하여 MLIR을 쉽게 만들고 다루는 Polygeist도 LLVM에서 열심히 다루고 있다. 또 TensorFlow에서도 MLIR을 사용한 TensorFlow MLIR을 열심히 만들고 있다.
MLIR은 여러 IR과 똑같이 SSA form을 사용한다. (사실 컴파일러에서는 필수인 듯 하다.) 사실 겉보기에는 LLVM IR과 크게 다를 것이 없다. 가장 중요한 핵심 목표가 기본 뼈대만을 제공하고, 사용자(컴파일러 개발자)들이 쉽게 자신만의 구조를 구현하도록 되어있기 때문이다.
아래와 같은 예시를 보자
func.func @toy_func(%tensor: tensor<2x3xf64>) -> tensor<3x2xf64> {
%t_tensor = "toy.transpose"(%tensor) { inplace = true } : (tensor<2x3xf64>) -> tensor<3x2xf64>
return %t_tensor : tensor<3x2xf64>
}
MLIR은 attributes, operations, types 등등을 모두 사용자가 만들 수 있다. 그렇기 때문에 위처럼 머신러닝에서 tensor가 어떻게 동작할 건지 역시 마음대로 정의할 수 있다.
그렇기 때문에 사용자가 인터페이스를 위해 dialect를 아래처럼 정의해볼 수도 있다.
/// This is the definition of the Toy dialect. A dialect inherits from
/// mlir::Dialect and registers custom attributes, operations, and types. It can
/// also override virtual methods to change some general behavior, which will be
/// demonstrated in later chapters of the tutorial.
class ToyDialect : public mlir::Dialect {
public:
explicit ToyDialect(mlir::MLIRContext *ctx);
/// Provide a utility accessor to the dialect namespace.
static llvm::StringRef getDialectNamespace() { return "toy"; }
/// An initializer called from the constructor of ToyDialect that is used to
/// register attributes, operations, types, and more within the Toy dialect.
void initialize();
};
이러한 dialect를 정의했으므로 이제, 새로운 연산을 개발자가 직접 정의할 수 있다. 따라서 직접 상수를 정의하려는 연산 toy.constant 정의는
%4 = "toy.constant"() {value = dense<1.0> : tensor<2x3xf64>} : () -> tensor<2x3xf64>
처럼 정의할 수 있다. 이는 인자를 받지 않고, 그냥 상수만을 반환하는 연산자이다.
이렇게 사실상 뼈대만 존재하고, 각자 원하는 dialect를 구축하기 때문에 여러 파생형이 나올 수 있다. 유사하게, cranelift ir도 구경해보자.
function %foo() {
block0:
v0 = iconst.i64 0
v2 = f32const 0.0
v9 = f32const 0.0
v20 = fneg v2
v18 = fcmp eq v20, v20
v4 = select v18, v2, v20
v8 = iconst.i32 0
v11 = iconst.i32 1
brif v0, block2, block3
block2:
brif.i32 v8, block4(v2), block4(v9)
block4(v10: f32):
v12 = select.f32 v11, v4, v10
v13 = bitcast.i32 v12
store v13, v0
trap user1
block3:
v15 = bitcast.i32 v4
store v15, v0
return
}
역시 익숙한 x86/64 cpu 대상 프로그램이라 그런지 훨씬 보고 이해하기 편한 코드가 보인다. 이역시 iconst.i64, f32const 와 같이 컴파일러에 필요한 연산(dialect)들을 직접 정의해서 사용하고 있다. 또한 block parameter라는 개념을 훨씬 잘 사용하기 위해 고안된 것도 볼 수 있다. 이 이유는 다음 글에서 확인할 수 있다.
사실 처음 보면서 MLIR을 보면서, 컴파일러 개발자가 IR을 만드는 것이 그렇게 어려운 것인가? 라고 생각했다. 하지만 내가 최적화를 증명하고, 일반화 하는 도구를 만들면서 처음에는 단순히, s-expression으로 표현하면 될 것 이라고 생각했다. 아래는 내가 단순하게 최적화 rewrite rule을 s-expression으로 표현하던 것이다. 하지만 점점 표현할 대상이 많아지고 구조가 복잡해지면서, 나만의 IR을 만드는 것이 정말 보통 일이 아니었다. 내가 문법을 생각하더라고 lexer와 parser를 구현하는 것은 또 다른 문제였고, 이를 동시에 유지보수하는 것이 너무 벅차기에 MLIR은 정말 컴파일러 개발자들에게 큰 도움이 될 것 이라 믿어 의심치 않고 있다. 최근 많은 컴파일러들이 MLIR을 사용하는 데에는 역시 다 이유가 있었다.
(rewrite
((Store
(Trunc
(Sub i32
(i32 conv29)
(i32 1))
i8)
(ptr s15))
(Ret (ptr null)))
((Store (Add i8 (Trunc (i32 conv29) i8) (i8 -1)) (ptr s15))
(Ret (ptr null))))