Beyond LLVM: Why MLIR

“이거 그냥 LLVM 위에 얹으면 되는 거 아닌가?”

처음 팀에서 MLIR 얘기가 나왔을 때 딱 그 생각이었다. 이미 LLVM이라는 탄탄한 컴파일러 인프라가 있는데 굳이 새로운 걸 배워야 하나 싶었는데, 막상 AI 가속기용 컴파일러를 짜기 시작하니까 왜 MLIR이 나왔는지 피부로 느끼게 됐다.

LLVM만으로는 어디가 어색한가

하는 일을 간단히 말하면 PyTorch 모델을 받아서 자체 하드웨어에서 돌아가는 바이너리로 바꿔주는 것이다. 언뜻 보면 “그냥 컴파일러잖아"처럼 들리는데, 실제로 해보면 생각보다 레이어가 여러 개 끼어있다.

PyTorch 모델은 텐서 연산 덩어리다. matmul, conv2d, softmax 같은 고수준 연산들이 그래프 형태로 연결돼 있는데, 이걸 하드웨어에서 실행하려면 단계적으로 낮춰가야 한다.

PyTorch 모델 (고수준 텐서 그래프)
  ↓
연산 분해 (matmul → 루프 + 메모리 접근)
  ↓
하드웨어 명령어 선택
  ↓
레지스터 할당, 스케줄링
  ↓
바이너리

LLVM IR은 이 중에서 아래쪽, 명령어 선택부터 바이너리까지는 정말 잘 한다. 수십 년의 최적화가 쌓인 인프라니까. 그런데 위쪽, 텐서 연산을 분해하고 하드웨어 특성에 맞게 타일링하고 메모리 레이아웃을 결정하는 부분은 LLVM이 직접 다루는 영역이 아니다.

처음에는 PyTorch 그래프를 파싱해서 직접 C++ 코드로 변환 로직을 작성하고, 그걸 다시 LLVM IR로 내리는 방식을 썼다. 처음엔 동작했는데 최적화를 붙이기 시작하면서 문제가 생겼다.

암묵적 계약의 문제

컴파일러 최적화에는 순서가 있다. 고수준에서 해야 효과적인 최적화가 있고, 저수준에서 해야 의미 있는 최적화가 따로 있다. 예를 들어 operator fusion — relu(matmul(A, B))처럼 연산 두 개를 하나로 합치는 것 — 은 텐서 연산 수준에서 해야 한다. LLVM IR까지 내려간 다음엔 이미 루프랑 메모리 접근으로 펼쳐져 있어서 패턴 인식이 훨씬 어려워지거든.

문제는 직접 짠 변환 코드에는 “지금 이게 어느 추상화 수준에 있는 IR인지"를 명확히 구분할 방법이 없다는 것이다. 그냥 다 C++ 객체고, 패스를 붙일 때마다 “이거 지금 어느 단계에서 돌려야 해?“를 팀원끼리 구두로 공유해야 하는 상황이 된다.

버그도 생긴다. 어떤 패스가 matmul이 분해되기 전에 돌아야 하는데 이미 분해된 후에 돌아서 엉뚱한 최적화가 적용된다든지, 반대로 분해 전에 메모리 레이아웃을 가정하는 패스가 끼어들어서 잘못된 코드가 나온다든지.

MLIR의 접근: 레벨을 명시한다

MLIR은 **“추상화 레벨이 다른 IR을 여러 개 정의하고, 그 사이의 변환을 명시적으로 관리하는 프레임워크”**라고 볼 수 있겠다.

MLIR에서는 dialect라는 개념으로 IR의 추상화 수준을 나눈다. 번역 프로젝트에 여러 명의 번역가가 투입됐는데 각자 담당 언어 구간이 명확히 나뉜 것과 비슷한 구조다.

linalg dialect  → 텐서/행렬 연산 수준 (matmul, conv)
affine dialect  → 루프와 메모리 접근 수준
vector dialect  → SIMD 수준
llvm dialect    → LLVM IR과 1:1 매핑

커스텀 dialect를 추가하면 자체 하드웨어에 특화된 연산도 IR에 넣을 수 있다. custom.hardware_matmul 같은 식으로.

“그냥 단계별로 클래스 나누면 안 됨?“이라고 생각할 수 있는데, MLIR이 주는 핵심 이점은 인프라에 있다. 패스 관리, 패턴 매칭, verifier(IR이 각 dialect 규칙을 지키는지 자동 검증), IR 출력 포맷까지 다 프레임워크가 제공한다. 직접 구현할 건 dialect 정의와 변환 규칙뿐이다.

실제로 써보면

처음 dialect를 정의하는 건 진입 장벽이 꽤 있다. TableGen이라는 DSL로 연산을 정의하는 방식이 낯설고, 공식 문서가 풍부한 편도 아니어서 LLVM 프로젝트 소스를 직접 뒤지는 시간이 생긴다.

그런데 한 번 패턴을 익히고 나면 확실히 달라진다. 변환 패스를 작성할 때 “이 패스는 linalg를 affine으로 내리는 패스"라고 타입 수준에서 명시하니까, 잘못된 순서로 패스가 끼어드는 버그가 없어진다. verifier가 각 단계에서 IR이 올바른 형태인지 자동으로 체크해주니까 디버깅 시간도 줄어든다.

무엇보다 팀원 간에 “지금 이 IR이 어떤 수준"인지를 코드로 소통할 수 있게 된다는 게 크다. 구두로 “이 패스는 분해 전에 돌려야 해"를 전달하는 대신, dialect 타입이 그 계약을 명문화해준다.

마무리

MLIR이 모든 걸 해결해주진 않는다. 새로운 dialect를 설계하는 것 자체가 아키텍처 결정이고, 잘못 설계하면 나중에 뒤집기 어렵다. 커뮤니티가 LLVM에 비해 아직 작아서 레퍼런스를 찾기 어려운 경우도 많다.

그래도 “각 단계가 뭘 알고 뭘 모르는지"를 IR 수준에서 명시할 수 있다는 것만으로, 암묵적 계약 문제는 상당 부분 해소된다. 다음 글에서는 실제로 커스텀 dialect를 정의하고 간단한 변환 패스를 작성하는 과정을 코드와 함께 다뤄볼 예정이다.