MLIR 입문 — AI 추론 최적화의 메커니즘을 저수준부터 이해하기 (vec_add Lowering 체험)

서론

llama.cpp, ONNX Runtime, TensorFlow와 같은 주요 AI 추론 런타임(Runtime)에서는 연산을 하드웨어에 맞춰 최적화하고 있습니다. 행렬 곱의 타이링(Tiling), SIMD 명령어로의 매핑, 메모리 계층 활용과 같은 최적화를 통해, 동일한 모델이라도 추론 속도는 크게 달라집니다.

이러한 최적화를 실현하는 프로세스 중 하나가 바로 MLIR을 이용한 컴파일입니다. MLIR (Multi-Level Intermediate Representation)은 텐서 연산의 고수준 표현을 하드웨어 고유의 명령어로 단계적으로 변환하는 컴파일러 기반이며, TensorFlow XLA, IREE, torch-mlir 등에 채택되어 있습니다.

이 기사의 큰 목적은 AI 추론 최적화가 어떤 프로세스로 이루어지는지 이해하는 것입니다.

그 입문으로서, 이번에는 MLIR의 기본인 Dialect를 이용한 기술과 Lowering (단계적 변환)을 벡터 가산 (vec_add)을 예로 들어 체험합니다. linalg dialect로 작성한 연산이 memref (메모리 버퍼)를 거치면서 scf → cf → llvm dialect로 단계적으로 변환되어, 최종적으로 LLVM IR이 될 때까지 직접 손을 움직여 확인합니다.

동작 환경

이 기사의 커맨드는 LLVM/MLIR 17.0.6에서 동작을 확인했습니다. mlir-opt와 mlir-translate가 로컬에서 동작하는 상태라면, 버전이 다소 다르더라도 기본적인 흐름은 동일합니다.

Dialect란 무엇인가

MLIR에서의 Dialect란 특정 추상 수준이나 대상 도메인에 대응하는 연산(op)과 타입의 집합입니다. 예를 들어 linalg dialect는 텐서 연산을 고추상도로 표현하며, llvm dialect는 LLVM IR에 직접 대응하는 저수준 명령어를 표현합니다.

MLIR의 중요한 설계 사상 중 하나가 progressive lowering (단계적 낮추기)입니다. "한 번에 하나의 추상화 단계만 낮춘다"라는 규칙에 따라, 고수준 표현에서 저수준으로 조금씩 변환해 나갑니다. 각 단계에서 변환/최적화 패스(pass)를 삽입할 수 있기 때문에 최적화 기회를 놓치지 않을 수 있습니다.

이번 변환 흐름은 다음과 같습니다.

linalg (텐서 연산)
↓ bufferize
memref (메모리 버퍼) + linalg
...

vec_add를 linalg dialect로 작성하기

먼저, linalg dialect를 사용한 vec_add의 MLIR 코드를 작성합니다.

// kernels/vec_add.mlir
module {
func.func @vec_add(
...

tensor를 사용하는 이유

tensor 타입은 값 세맨틱스(Value Semantics)를 가지며 부작용(Side Effect)이 없습니다. 후술할 bufferize 패스가 tensor를 memref (메모리 버퍼)로 변환하며, memref.alloc 등의 명시적인 메모리 할당 명령어를 삽입합니다. 고수준에서는 값으로 취급하고, 메모리의 상세 사항은 하위 변환에 맡기는 분리가 가능합니다.

tensor.empty()의 역할

tensor.empty()는 출력을 위한 "타입 정보만을 가진 placeholder"입니다. 값을 확보하는 것이 아니라, linalg.add의 outs 오퍼랜드(operand) 타입을 맞추기 위해 사용합니다. 실제 메모리 확보는 bufferize 이후에 memref.alloc으로 나타납니다.

파싱(Parsing)이 올바른지 확인합니다.

mlir-opt kernels/vec_add.mlir -o /dev/null

에러가 발생하지 않는다면 구문은 문제없습니다.

mlir-opt로 단계적 lowering 관찰

이 부분이 이 기사의 메인 파트입니다. mlir-opt에 pass를 전달하여 단계적으로 변환하고, 각 스텝에서 어떻게 변화하는지 확인합니다.

pass의 순서

pass의 순서

# 00: tensor.empty → bufferization.alloc_tensor
mlir-opt --empty-tensor-to-alloc-tensor kernels/vec_add.mlir \
-o build/00_alloc_tensor.mlir
...

각 단계의 변화 추적

01_bufferized.mlir — tensor → memref

func.func @vec_add(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) -> memref<1024xf32> {
%alloc = memref.alloc() {alignment = 64 : i64} : memref<1024xf32>
linalg.add ins(%arg0, %arg1 : memref<1024xf32>, memref<1024xf32>) outs(%alloc : memref<1024xf32>)
...

tensor가 memref로 바뀌었고, memref.alloc이 등장했습니다. tensor.empty()는 memref.alloc으로 대체되었습니다. linalg.add는 아직 이 단계에서는 남아 있으며, 추상적인 텐서 연산을 표현한 채입니다.

02_loops.mlir — linalg.add → scf.for

func.func @vec_add(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) -> memref<1024xf32> {
%c0 = arith.constant 0 : index
%c1024 = arith.constant 1024 : index
...

linalg.add가 scf.for 루프로 전개되었습니다. arith.addf라는 스칼라 덧셈 op이 처음 모습을 드러냈습니다. scf.for는 아직 '구조화된 루프(structured loop)' 표현으로, 루프의 의도가 인간에게도 읽기 쉬운 형태를 유지하고 있습니다.

03_cf.mlir — scf.for → cf의 기본 블록

func.func @vec_add(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) -> memref<1024xf32> {
%c0 = arith.constant 0 : index
%c1024 = arith.constant 1024 : index
...

scf.for가 cf.br / cf.cond_br을 이용한 기본 블록(basic block) 간의 분기로 변환되었습니다. ^bb1이 루프 헤더, ^bb2가 루프 본문, ^bb3이 탈출점입니다. 구조화된 루프 개념이 사라지고, 레이블과 점프의 조합으로 표현되고 있습니다. memref나 arith op은 아직 남아 있으며, 이 단계는 '제어 흐름의 비구조화(unstructured control flow)'만을 담당하고 있습니다.

04_affine.mlir — affine의 해소

--lower-affine은 affine dialect의 식(affine.apply 등)을 arith 연산으로 전개하는 pass입니다. 이번 vec_add에서는 affine dialect op이 직접 나타나지 않았기 때문에, 03_cf.mlir로부터의 차이는 거의 없습니다. affine dialect를 사용한 타일링(tiling)이나 루프 변환을 삽입했을 경우, 이 pass에서 arith / cf 레벨로 떨어지게 됩니다.

05_llvm.mlir — llvm dialect로 통일

module attributes {llvm.data_layout = ""} {
llvm.func @malloc(i64) -> !llvm.ptr
llvm.func @vec_add(%arg0: !llvm.ptr, %arg1: !llvm.ptr) -> !llvm.ptr {
...

모든 op이 llvm.* 접두사로 통일되었습니다. scf.for는 llvm.cond_br을 이용한 기본 블록 간의 분기 제어로 변환되었고, 스칼라 덧셈은 llvm.fadd

되어 있습니다. 함수 시그니처(Function Signature)도 memref<1024xf32>에서 !llvm.ptr로 변경되어 있으며, LLVM IR로의 변환 준비가 완료된 상태입니다.

mlir-translate로 LLVM IR 생성

마지막으로, mlir-translate를 사용하여 llvm dialect에서 LLVM IR(텍스트 형식)로 변환합니다.

mlir-translate --mlir-to-llvmir build/05_llvm.mlir -o build/vec_add.ll

생성된 vec_add.ll의 도입부는 다음과 같습니다 (opt가 적용되지 않은 스칼라 상태).

define ptr @vec_add(ptr %0, ptr %1) {
...
br label %24
...

LLVM IR 레벨에서는 아직 스칼라 루프입니다. phi 명령어로 루프 카운터를 관리하고, icmp slt로 종료 판정을 하며, fadd로 요소를 하나씩 가산하고 있습니다. 벡터 타입(예: <4 x float>)은 아직 나타나지 않았지만, SIMD 명령어로의 변환은 이 이후의 이야기입니다.

Lowering의 각 단계가 최적화의 삽입 지점이 된다

지금까지의 절차를 통해 텐서 연산(Tensor Operation) → 루프(Loop) → 기본 블록(Basic Block) → LLVM이라는 단계적인 변환을 확인할 수 있었습니다. 이러한 다단계 구성이야말로 MLIR을 추론 최적화에 유용하게 만드는 설계의 핵심입니다.

각 단계는 독립적인 최적화 삽입 지점(Insertion Point)이 됩니다.

• linalg (텐서 연산) — 루프 타이링(Tiling)이나 퓨전(Fusion)을 적용하기 쉽습니다. 예를 들어 linalg 단계에서 linalg.tiling 패스를 삽입하면, scf.for로 전개되기 전에 루프 분할을 수행하여 캐시에 적재될 수 있는 크기로 계산을 나눌 수 있습니다.
• scf (구조화 루프) — 루프 병렬화(scf 단계의 scf.parallel)나 언롤링(Unroll)을 적용할 수 있습니다. 루프 구조가 아직 명시되어 있기 때문에 변환 의도를 추적하기 쉬운 단계입니다.
• llvm (LLVM dialect) — 메모리 액세스 패턴의 최적화나 타겟 고유 명령(SIMD 등)으로의 매핑을 여기서 수행합니다. SIMD 명령어로의 벡터화(Vectorize)는 이 단계 이후의 작업입니다.

수동 커널(Hand-written Kernel)로 유사한 최적화를 수행할 경우, 타이링과 벡터화를 모두 동일한 소스 코드 상에서 관리해야 합니다. MLIR의 접근 방식은 각 최적화를 담당하는 변환 패스(Transformation Pass)를 단계별로 분리할 수 있기 때문에, "어느 추상화 레벨에서 어떤 최적화를 수행할 것인가"를 정리하며 구현 및 디버깅할 수 있습니다.

요약

이번 lowering을 통해 확인한 내용을 정리합니다.

• linalg (텐서 연산) → scf (구조화 루프) → cf (기본 블록) → llvm (LLVM dialect)라는 단계를 거치며, 각 변환 패스는 "하나의 추상화 단계만을 낮춘다"는 책임을 가집니다.
• 각 단계는 최적화의 삽입 지점이 되어 타이링, 병렬화, 메모리 최적화, SIMD 매핑을 각각 적절한 추상화 레벨에서 적용할 수 있습니다.
• mlir-translate로 생성한 LLVM IR은 스칼라 루프 상태이지만, 이는 출발점일 뿐입니다. SIMD 명령어로의 벡터화는 이후의 변환 패스에서 수행됩니다.

AI 추론의 가속화는 이러한 다단계 변환 파이프라인 위에 구축됩니다. MLIR의 Lowering을 직접 따라가 봄으로써, 프레임워크가 "텐서 연산을 어떻게 하드웨어로 낮추는가"라는 질문에 대한 구체적인 이미지를 얻을 수 있습니다.

다음 글에서는 여기서 생성한 스칼라 LLVM IR(vec_add.ll)을 기점으로, opt로 벡터화를 적용하고 llc로 NEON 명령어를 포함한 어셈블리를 생성합니다. "SIMD 명령어가 나오는 조건"과 "실제로 얼마나 빨라지는지"를 실측을 통해 확인하겠습니다.

Discussion