모델을 훈련하거나 서빙한다면, 여러분이 MLIR 코드를 단 한 줄이라도 작성했는지 여부와 상관없이 MLIR에 의존하고 있는 것입니다. XLA는 이를 통해 하위 단계로 내려가고(lowers), Triton은 그 위에 구축되었으며, Mojo는 MLIR 네이티브(MLIR-native)입니다. 또한 Torch-MLIR, IREE, ONNX-MLIR은 각각의 프론트엔드(frontend)를 MLIR로 전달하기 위해 존재합니다. 단일 인프라가 서로 경쟁 관계에 있는 수많은 스택의 하단에 자리 잡게 된 이유는 이해할 가치가 있습니다. 왜냐하면 이는 현대적인 ML 컴파일러(ML compilers)가 실제로 어떻게 구축되는지, 그리고 그 경계(seams)가 어디인지에 대해 많은 것을 설명해주기 때문입니다.

이 포스트는 MLIR에 대한 투어입니다. MLIR이 무엇인지, MLIR을 다르게 만드는 Dialect(방언) 개념은 무엇인지, 텐서 연산(tensor operation)이 어떻게 점진적으로 머신 코드(machine code)로 낮아지는지, 그리고 이 인프라가 무엇을 제공하고 무엇을 제공하지 않는지에 대해 다룹니다.

흔한 오해는 MLIR이 "LLVM IR와 같은 또 다른 IR"이라고 생각하는 것입니다. MLIR은 IR *구성 키트(construction kit)*라고 설명하는 것이 더 적절합니다. LLVM IR은 단일하고 고정된 저수준 표현(low-level representation)으로, 대략 SSA(Static Single Assignment) 값을 가진 타입화된 어셈블리(typed assembly)입니다. 이는 머신 코드로 가는 마지막 단계(last mile)에는 적합한 추상화이지만, 텐서 상의 행렬 곱셈(matrix multiply)에는 잘못된 추상화입니다. 역사적으로 더 높은 수준의 표현이 필요한 모든 도메인 컴파일러(domain compiler)는 각자 처음부터 자신만의 표현을 만들어냈습니다. XLA는 HLO를, Halide는 자체 IR을, TensorFlow는 그래프를 가졌으며, 각각 별도의 패스 매니저(pass manager), 직렬화 형식(serialization format), 검증기(verifier), 그리고 버그 더미를 가지고 출시되었습니다.1

MLIR의 전제는 이러한 표현들이 서로 공통점이 매우 많으며, 공통된 부분은 한 번만 구축하여 공유할 수 있다는 것입니다. 즉, SSA, 리전(regions)과 블록(blocks)으로 구성된 CFG(Control Flow Graph), 패스 인프라(pass infrastructure), 패턴 재작성기(pattern rewriter), 위치 추적(location tracking), 그리고 검증(verification) 등이 공유됩니다. 도메인 간에 차이가 나는 부분은 그 후 *Dialect(방언)*로 표현됩니다.

MLIR의 모든 단위는 **Operation(연산)**입니다. 하나의 연산(op)은 피연산자(operands)와 결과(results, SSA 값), 타입이 지정된 Attributes(속성) 세트(형태(shapes)나 스트라이드(strides)와 같은 컴파일 타임 상수), 그리고 0개 이상의 **Regions(리전)**를 가집니다. 이 리전들은 그 자체로 더 많은 연산의 블록을 포함합니다. 이 마지막 속성이 IR을 진정으로 멀티 레벨(multi-level)로 만드는 요소입니다. 단일 연산이 전체 중첩된 계산을 담을 수 있으므로, 고수준의 linalg.generic...

and 저수준의 llvm.add는 서로 다른 고도(altitude)에 있는 동일한 종류의 객체입니다. 모든 연산(op)은 방언(dialect)에 속하며, 방언은 단순히 관련된 연산(ops), 타입(types), 속성(attributes) 제품군을 위한 네임스페이스(namespace)입니다.

문법(grammar)으로 표현하자면, 이 관계는 작고 재귀적입니다. 방언은 연산 이름, 타입, 속성이라는 *어휘(vocabulary)*를 제공하는 반면, 연산의 *형태(shape)*는 보편적입니다. 아래 형식은 필수적인 구조를 보여주기 위해 MLIR의 텍스트 문법(textual grammar)을 단순화한 것이며, 권위 있는 생성 규칙(productions)은 Language Reference2에 기술되어 있습니다:

; 방언은 연산(ops), 타입(types), 속성(attributes) 제품군을 제공하는 네임스페이스입니다.
dialect ::= (operation-def | type-def | attribute-def)*
; 연산의 문법은 모든 방언에 걸쳐 동일합니다.
...

여기서 두 가지 결과가 도출됩니다. 첫째, 연산 이름, 타입, 속성은 모두 방언에 의해 네임스페이스화되어 있으므로 (linalg.matmul, tensor<128x256xf32>), "방언을 추가하는 것"은 문법을 건드리지 않고 어휘를 확장하는 것이며, 이것이 바로 주변 인프라가 방언에 무관(dialect-agnostic)할 수 있는 정확한 이유입니다. 둘째, operation이 region을 포함할 수 있고, region이 더 많은 operation들로 구성된 block들을 포함하기 때문에, 구조는 무제한으로 중첩될 수 있습니다. 그리고 이러한 재귀성이 단일 연산이 하나의 명령어가 아닌 전체 계산을 담을 수 있게 해줍니다.

결정적인 특징은 방언들이 *공존(coexist)*한다는 점입니다. 컴파일 중간 단계의 모듈(module)은 통상적으로 여러 방언의 연산들을 동시에 보유하며, 로워링(lowering)은 타겟 방언(target dialect)만 남을 때까지 고수준 연산을 저수준 연산으로 점진적으로 교체하는 과정입니다. ML 파이프라인이 고수준에서 저수준으로 거치게 되는 방언들은 다음과 같습니다:

고수준 (의도하는 바): stablehlo 및 tosa (전체 텐서 연산 세트), linalg (텐서/버퍼에 대한 구조화된 연산), tensor (값 의미론적 텐서 조작).

중수준 (구조화 방식): memref (레이아웃/스트라이드(layout/strides)를 가진 버퍼), affine 및 scf (루프 중첩 및 구조화된 제어 흐름), vector (SIMD), arith

(스칼라 연산 (scalar math)). 4저수준 (실행되는 단계): llvm (CPU를 위한 LLVM IR로 변환), gpu 및 nvvm / rocdl (GPU 타겟), spirv (Vulkan/compute).

MLIR이 담고 있는 기술은 다음 단계로 언제 고도를 낮출지(drop)를 결정하는 것입니다. 너무 오랫동안 높은 고도에 머물면 하드웨어 특화된 스케줄을 표현할 수 없고, 너무 일찍 낮은 단계로 내려가면 최적화 도구(optimizer)가 필요로 하는 구조를 버리게 됩니다.

*공존(coexist)*이라는 단어는 쉽게 간과하기 쉬우므로, 여기 네 개의 Dialect를 동시에 사용하는 단일 함수를 소개합니다. 아직 아무것도 하위 단계로 변환(lowered)되지 않았습니다. 이 연산(ops)들은 하나의 SSA 리전(region) 내에서 단순히 나란히 존재하며, 검증기(verifier)가 이들을 함께 체크합니다:

// 하나의 리전, 네 개의 Dialect; 각 라인은 해당 Dialect 태그가 붙어 있습니다.
func.func @scale_in_place(%buf: memref<1024xf32>, %a: f32) {
%c0 = arith.constant 0 : index // 1. arith: 루프 경계는 스칼라 인덱스 상수입니다
%c1 = arith.constant 1 : index // 2. arith
%n = arith.constant 1024 : index // 3. arith
scf.for %i = %c0 to %n step %c1 { // 4. scf: 리전을 포함하는 구조화된 루프
%x = memref.load %buf[%i] : memref<1024xf32> // 5. memref: 명시적 버퍼로부터 읽기
%y = arith.mulf %x, %a : f32 // 6. arith: 스칼라 곱셈
memref.store %y, %buf[%i] : memref<1024xf32> // 7. memref: 결과를 다시 쓰기
}
return // 8. func: 종료 연산(terminator)
}

func, arith, scf, memref라는 네 개의 Dialect가 임피던스 불일치(impedance mismatch) 없이 하나의 리전에 나타납니다. 이후의 패스(pass)가 scf.for를 cf 분기(branches)로 다시 쓰거나, 루프 본문을 vector Dialect로 벡터화할 수도 있지만, 이 고도에서는 이들이 단순히 결합(compose)됩니다. 그 결합 가능성(composability)이 핵심입니다. Dialect는 서로 간에 번역해야 하는 별개의 IR이 아니라, 단일 프로그램 내에서 혼합하여 사용하는 어휘(vocabularies)입니다.

graph TD
subgraph FE["Frontends"]
SH["StableHLO / TOSA"]
...

단일 행렬 곱셈(matrix multiply)을 예로 들어보겠습니다. 최상위 단계에서 이것은 (값 의미론 (value semantic)을 가진) 텐서(tensors)에 대한 하나의 구조화된 **연산 (op)**입니다. 형태(Shapes)는 타입(type)에 존재하며, 메모리나 루프에 대해서는 아직 아무것도 언급되지 않았습니다.

1
2
3
4
5
6
7
8
9
10
11
12
13

// 형태(Shapes)는 타입(type)의 일부입니다. 이것은 텐서(tensor)에 대한 순수한 데이터 흐름(dataflow)입니다.
func.func @matmul(%A: tensor<128x256xf32>,
%B: tensor<256x64xf32>) -> tensor<128x64xf32> {
%zero = arith.constant 0.0 : f32 // 1. 누산기(accumulator)를 위한 스칼라 항등원(scalar identity)
%init = tensor.empty() : tensor<128x64xf32> // 2. 초기화되지 않은 결과 값
%acc = linalg.fill ins(%zero : f32)
outs(%init : tensor<128x64xf32>)
-> tensor<128x64xf32> // 3. 누산기를 0으로 초기화
%C = linalg.matmul // 4. 하나의 연산(op)으로서의 전체 곱셈
ins(%A, %B : tensor<128x256xf32>, tensor<256x64xf32>)
outs(%acc : tensor<128x64xf32>) -> tensor<128x64xf32>
return %C : tensor<128x64xf32>
}

두 가지 변환(transformations)을 통해 이것은 기계어(machine code)에 가까운 무언가로 변합니다. 첫째, **버퍼화 (bufferization)**는 값 의미론(value-semantic)을 가진 tensor를 주소, 레이아웃(layout), 수명(lifetime)을 가진 명시적 버퍼인 memref로 변환합니다. 이 시점에서 프로그램은 순수한 데이터 흐름(dataflow) 상태를 벗어나 에일리어싱(aliasing), 할당(allocation), 그리고 새로운 값을 반환하는 대신 호출자가 제공한 출력 버퍼에 결과를 기록해야 하는 특성을 갖게 됩니다. 둘째, 구조화된 linalg.matmul은 명시적인 루프 중첩(loop nest)으로 낮아집니다(lowered):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

// 텐서는 이제 memref입니다; 결과는 %C에 제자리에서(in place) 기록됩니다.
func.func @matmul(%A: memref<128x256xf32>,
%B: memref<256x64xf32>,
%C: memref<128x64xf32>) {
affine.for %i = 0 to 128 { // 1. matmul이 암시했던 루프 중첩(loop nest)
affine.for %j = 0 to 64 {
affine.for %k = 0 to 256 {
%a = affine.load %A[%i, %k] : memref<128x256xf32> // 2. 버퍼로부터의 명시적 로드(loads)
%b = affine.load %B[%k, %j] : memref<256x64xf32>
%c = affine.load %C[%i, %j] : memref<128x64xf32>
%p = arith.mulf %a, %b : f32 // 3. 이제 스칼라가 된 산술 연산(arithmetic)
%s = arith.addf %c, %p : f32
affine.store %s, %C[%i, %j] : memref<128x64xf32> // 4. 버퍼로 다시 누적
}
}
}
return
}

여기서부터 affine / scf 루프는 타일링(tiled)될 수 있고, vector 디렉트(dialect)로 벡터화(vectorized)될 수 있으며, 마지막으로 llvm으로 낮아질(lowered) 수 있습니다.

(CPU용) dialect 또는 gpu/nvvm 경로입니다. 본질적으로, 서로 다른 낮추기 패스(lowering passes)를 선택함으로써 동일한 고수준 시작 지점에서 이러한 모든 타겟 경로로 내려갈 수 있습니다. 하나의 프론트엔드 표현(frontend representation)을 재사용하여 다양한 타겟으로 낮출 수 있다는 점이 바로 MLIR이 가진 경제적 논거의 핵심입니다.

이것을 실용적으로 만드는 것은 디렉트(dialect) 그 자체라기보다 그 아래에 있는 공유 인프라입니다. MLIR은 디렉트 변환 (dialect conversion) 프레임워크를 제공합니다. 사용자는 변환 대상(최종적으로 어떤 디렉트/연산(ops)이 "합법적(legal)"인지)을 선언하고, 불법적인 연산을 합법적인 연산으로 변환하는 재작성 패턴(rewrite patterns)을 제공하면, 프레임워크가 이를 고정점(fixed point)까지 적용하며 그 과정에서 타입 변환(type conversions)과 피연산자 재매핑(operand remapping)을 처리합니다.2 패턴 재작성(Pattern rewriting), 검증기(verifier), 위치/디버그 정보 전파(location/debug-info propagation), 그리고 패스 매니저(pass manager)는 모두 디렉트와 무관하게(dialect-agnostic) 동작합니다. 대부분의 변환은 수동으로 코딩된 IR 순회(traversal) 방식이 아니라 로컬 재작성 패턴으로 작성됩니다. 여기에는 각 디렉트가 자체 연산을 접고 단순화하기 위해 등록하는 정규화(canonicalization) 패턴이 포함되며, 이는 선언적(PDL/DRR을 통해)으로 선언되어 공유 드라이버에 의해 적용될 수 있습니다.6 새로운 추상화란 일련의 연산 정의(op definitions)와 패턴 세트이며, 그 주변의 모든 것은 상속됩니다.

이것이 생태계가 통합된 이유입니다. 각 프로젝트가 맞춤형 컴파일러 미들엔드(middle-end)를 유지하는 대신, 하나의 미들엔드를 공유합니다:

graph LR
MLIR["MLIR core<br/>ops · dialects · passes · rewriter"]
MLIR --- XLA["XLA / StableHLO"]
...

또한 이것은 새로운 하드웨어가 타겟팅되는 방식이기도 합니다. 벤더(vendor)는 자신의 장치 연산과 메모리를 모델링하는 디렉트를 추가하고, 표준 중간 단계 디렉트(linalg, memref, vector)를 해당 장치로 낮추는 패스들을 추가합니다. 그러면 프론트엔드, 최적화 도구, 툴링은 거저 얻게 됩니다. 벤더는 오직 자신들의 실리콘(silicon)에 진정으로 특화된 부분만을 작성하면 됩니다.7

MLIR을 생산적으로 만드는 이 메커니즘에는 명확히 짚고 넘어갈 만한 날카로운 측면(sharp edges)들이 있습니다.

올바른 검증(verification)을 정의하는 것은 실제적인 작업입니다. 모든 Dialect는 자신의 Op와 Type에 대한 불변성(invariants)을 강제하는 검증기(verifier)를 포함하며, 공유 프레임워크는 이들을 모두 함께 실행합니다. 이를 통해 잘못된 IR을 조기에 잡아낼 수 있지만, 올바른 검증 규칙을 작성해야 하는 부담은 Dialect를 정의하는 사람에게 돌아갑니다. Type 불변성을 잘못 정의하거나 누락하면, 잘못된 IR이 검증기를 통과하게 되고, 실제 오류가 발생한 지점에서 한참 떨어진 이후의 Pass에서 크래시(crash)로 나타나게 됩니다.8

디버깅은 다양한 고도(altitudes)에 걸쳐 이루어집니다. 하나의 모듈에 여러 Dialect가 있고 긴 Pass 시퀀스가 존재할 경우, 컴파일 오류나 성능 저하(performance regression)를 고수준의 Tensor Op에서부터 특정 LLVM 또는 PTX 명령어로까지 추적해야 합니다. 위치 추적(Location tracking) 기능이 lowering 과정에서 소스 위치를 전달하여 도움을 주기는 하지만, 버그는 종e종 원인이 된 Pass로부터 여러 Pass가 지난 후에 나타나곤 합니다. 실질적인 도구로는 Pass 사이의 IR을 덤프하는 방식(mlir-opt --mlir-print-ir-after-all)과 자동 크래시 재현기(automatic crash reproducer)가 있습니다. 이러한 도구들이 있더라도, 다층적(multi-level) 디버깅은 그 자체로 하나의 기술이며, 이를 개선하는 것은 MLIR 커뮤니티의 반복적인 주제입니다.9

Dialect는 안정적이지만, 파이프라인(pipelines)은 그렇지 않습니다. StableHLO는 버전 관리가 되는 이식 가능한 교환 형식(interchange)이며, 내장된 Dialect들은 상당히 안정적입니다. 하지만 변환(transformations)들을 하나로 묶는 Pass 파이프라인은 안정적이지 않습니다. Torch-MLIR에 맞춰 조정된 lowering 파이프라인이 IREE에서 수정 없이 그대로 실행되지는 않으며, 한 릴리스에서 작동하던 시퀀스가 다음 릴리스에서 깨질 수도 있습니다. “MLIR 기반(Built on MLIR)”이라는 말은 Dialect들이 결합(compose) 가능하다는 의미이지, 두 개의 MLIR 기반 스택이 엔드 투 엔드(end to end)로 상호 운용 가능하다는 의미는 아닙니다.10