MLIR용 Python 프런트엔드
지난번에는 MLIR을 GPU 백엔드로 가져가서, 그 결과로 나온 코드를 실행하려면 무엇이 필요한지 살펴보았습니다. 이제는 컴파일러의 반대편 끝에 놓일 것이 필요합니다. 바로 Python으로 텐서 계산을 기술하고 그 백엔드들이 소비하는 IR을 구성하는 방법입니다.
이 부분은 MLIR의 Python 바인딩을 중심으로 구축하겠습니다. 이 바인딩은 지금까지 텍스트 형식으로 작성해 온 것과 같은 객체들, 즉 타입, 값, 연산, 영역, 블록을 노출합니다. Python 프로그램은 linalg.matmul 연산을 구성하고, 그 결과를 다른 연산에 전달하고, 그렇게 만들어진 모듈에 컴파일러 패스를 실행할 수 있습니다. 출력된 MLIR은 그 모듈을 살펴보는 한 가지 방법입니다.
함께 제공되는 프로젝트는 tiny-mlir입니다. 이 프로젝트에는 남은 장들에서 사용하는 전체 구현이 들어 있습니다. 저장소를 클론하고 환경을 설치하십시오.
git clone https://github.com/sdiehl/tiny-mlir.git
cd tiny-mlir
uv sync
uv run pytest
이 프로젝트는 Python 3.12와 특정 MLIR 22 개발 빌드를 고정해 두며, 여기에는 Python 바인딩과 실행 엔진이 포함됩니다. 이 예제들은 그 번들된 툴체인을 사용합니다. 앞선 장들에서 설치한 명령줄 LLVM은 별도로 두어도 됩니다. 그 라이브러리를 Python 환경에 섞을 필요는 없습니다.
컴파일할 함수
우리가 사용할 인터페이스는 다음과 같습니다.
import numpy as np
from tinymlir import jit
@jit
def scale_bias(x, bias):
return x * 2.0 + bias
x = np.arange(8, dtype=np.float32)
bias = np.ones(8, dtype=np.float32)
print(scale_bias(x, bias))
첫 번째 호출은 인자의 형상과 dtype에 맞게 함수를 특수화하고, MLIR을 구성하고, 이를 하강시키고, 실행 엔진을 생성합니다. 이후 같은 타입으로 호출하면 그 특수화를 재사용합니다. x의 내용을 바꾸는 것은 컴파일을 요구하지 않습니다. 형상을 바꾸면 컴파일이 필요합니다.
이 함수는 memref 디스크립터나 실행 엔진에 대해 아무것도 알 필요가 없습니다. 함수는 텐서 표현식을 기술할 뿐입니다. 표현과 런타임 경계는 컴파일러가 맡습니다.
Python 트레이싱
프런트엔드를 만드는 합리적인 방법은 여러 가지가 있습니다. 6부에서는 Python 구문을 파싱하고 변환했습니다. 여기서는 트레이싱을 사용하겠습니다. 즉, 배열 대신 기호적 텐서 인자로 Python 함수를 호출합니다. 이 인자들의 산술 메서드가 표현식 노드를 구성합니다.
Python이 x * 2.0을 평가할 때, 우리 기호적 값의 __mul__ 메서드를 호출합니다. 덧셈을 평가할 때는 그 메서드가 곱셈 노드를 왼쪽 피연산자로 받습니다. 그 결과는 배열 원소를 하나도 읽지 않고 계산을 기술하는 그래프입니다.
이것은 작은 임베디드 언어입니다. 일반적인 Python 함수 호출은 트레이싱 중에 자연스럽게 합성됩니다. 정적 형상 정보를 순회하는 루프도 경계가 알려져 있으므로 트레이싱 중에 실행할 수 있습니다. 텐서 내용에 의존하는 분기는 그렇게 할 수 없습니다. 기호 값의 __bool__이 오류를 발생시키기 때문입니다. 임의의 Python을 컴파일하거나 Python 객체에 대한 변경을 포착하지는 않습니다.
우리의 타입은 특수화에 필요한 정보를 정확히 담고 있습니다.
@dataclass(frozen=True)
class TensorType:
shape: tuple[int, ...]
dtype: np.dtype
구현은 형상과 dtype을 정규화하고, 양수인 정적 차원을 요구하며, float32, float64, int32를 허용합니다. 정수 타입은 토큰 인덱스에 사용됩니다. 부동소수점 연산은 정밀도를 명시적으로 지닙니다.
각 표현식은 연산, 결과 타입, 피연산자, 그리고 정적 속성을 저장합니다.
@dataclass(eq=False)
class Expr:
op: Op
type: TensorType
args: tuple[Expr, ...] = ()
attr: object = None
Op는 프런트엔드와 빌더가 공유하는 열거형이므로 연산 이름은 한곳에서 정의됩니다. 이 노드들은 구조적 동등성이 아니라 동일성을 사용합니다. 계산에서 같은 노드를 재사용하면 로어링도 그 노드의 MLIR 값을 캐시해 두었다가 재사용합니다. 이후 MLIR의 공통 하위 표현식 제거 패스가 독립적으로 구성된 동등한 연산들을 식별할 수 있습니다.
코드 생성 이전의 형상
덧셈은 먼저 피연산자들이 브로드캐스트될 수 있는지 확인합니다. 형상 (3, 7)과 (7,)의 경우 결과의 형상은 (3, 7)입니다. (3, 7)과 (3, 1)의 경우에도 결과의 형상은 같습니다. (3, 7)과 (5,)의 경우 컴파일이 실패합니다.
프런트엔드는 또한 공통 dtype을 결정하고 필요할 때 명시적인 캐스트 노드를 삽입합니다. Python 숫자 리터럴은 주변 텐서의 dtype을 사용합니다. NumPy 스칼라 값은 자신의 dtype을 유지합니다. 이 구분은 NumPy GPT-2 구현을 재현할 때 중요해집니다. 해당 구현은 일부 공식에서 float64 값을 도입하기 때문입니다.
이 단계 이후에는 백엔드가 덧셈이 스칼라인지, 원소별인지, 브로드캐스트인지 추측할 필요가 없습니다. 피연산자 형상과 결과 형상이 필요한 인덱싱을 이미 기술하고 있습니다.
형상 특수화는 정적 정보를 런타임 데이터와 분리하는 유용한 사례입니다. 이는 6부의 부분 평가 개념과 관련이 있습니다. 타입과 정적 매개변수를 고정하고, 배열 내용은 컴파일된 프로그램의 입력으로 남겨 둡니다. 우리의 구현은 Python 인터프리터를 위한 범용 특수화기가 아니라 전용 텐서 컴파일러입니다.
모듈 구성하기
MLIR 바인딩은 컨텍스트, 모듈, 삽입 지점을 제공합니다.
from mlir import ir
from mlir.dialects import func
with ir.Context(), ir.Location.unknown():
module = ir.Module.create()
with ir.InsertionPoint(module.body):
f32 = ir.F32Type.get()
buffer = ir.MemRefType.get([8], f32)
function = func.FuncOp("kernel", ([buffer, buffer], []))
block = function.add_entry_block()
with ir.InsertionPoint(block):
func.ReturnOp([])
module.operation.verify()
이것은 흥미롭지는 않지만 완전한 함수로, 두 개의 버퍼를 받고 아무것도 반환하지 않습니다. 컨텍스트는 타입과 같은 고유화된(uniqued) 객체를 소유합니다. 삽입 지점은 새로 생성된 연산이 어디에 들어갈지를 결정합니다. 블록은 자신의 인자와 연산을 소유합니다. 이러한 관계 덕분에 바인딩은 프로그램을 구성하는 동안 객체들이 서로 어떻게 맞물리는지 검사할 수 있습니다.
실제 빌더는 각 입력마다 memref 인자를 하나씩 만들고, 호출자의 출력 버퍼를 위한 memref 인자도 하나 만듭니다. 빌더는 bufferization.to_tensor로 입력을 텐서 값으로 변환하고, 텐서 위에서 계산을 구성한 뒤, bufferization.materialize_in_destination으로 마무리합니다. 후자는 결과가 호출자가 제공한 출력 버퍼에 최종적으로 담겨야 한다는 것을 나타냅니다.
이 경계 덕분에 외부 호출 규약은 단순하게 유지되면서도 본문은 함수형 텐서 계산으로 남을 수 있습니다. 어떤 내부 텐서에 저장 공간이 필요하고 어떤 버퍼를 재사용할 수 있는지는 버퍼화가 결정합니다.
원소별 영역
원소별 표현식은 linalg.generic이 됩니다. 빌더는 입력 값, 빈 목적지 텐서, 인덱싱 맵, 반복자 타입을 제공합니다. 그런 다음 연산의 영역을 만들고 스칼라 계산을 채워 넣습니다:
op = linalg.GenericOp(
[output.type], inputs, [output], indexing_maps,
["parallel"] * rank,
)
block = op.regions[0].blocks.append(
*(value.type.element_type for value in [*inputs, output])
)
with ir.InsertionPoint(block):
product = arith.MulFOp(block.arguments[0], block.arguments[1])
linalg.YieldOp([product.result])
이 코드 조각은 빌더 안의 곱셈 경우입니다. output, inputs, 맵, 랭크는 타입이 지정된 표현식에서 가져옵니다. 영역 인자는 스칼라 원소이므로, 그 본문은 두 개의 부동소수점 수를 곱하여 하나의 부동소수점 수를 산출합니다. 이를 감싸는 linalg.generic은 그 스칼라 연산이 텐서 전체에 걸쳐 어떻게 적용되는지를 기술합니다.
브로드캐스팅은 인덱싱 맵에서 다룹니다. 행렬에 행 벡터를 더하는 경우, 그 맵은 다음과 같습니다:
matrix: (i, j) -> (i, j)
vector: (i, j) -> (j)
output: (i, j) -> (i, j)
일곱 개의 열에 걸쳐 브로드캐스트되는 (3, 1) 텐서의 경우, 그 맵은 (i, j) -> (i, 0)이 됩니다. 스칼라 텐서는 빈 인덱스 목록을 가집니다. 이 맵들은 AffineDimExpr, AffineConstantExpr, AffineMap을 사용하여 구성하며, 텍스트 형식의 아핀 표현식 파서는 전혀 관여하지 않습니다.
스칼라 영역과 그 반복 공간을 구분하는 것이 바로 우리가 최적화기에게 보여 주고자 하는 구조입니다. 최적화기는 루프 스케줄을 확정하기 전에 맵과 영역을 합성하여 원소별 연산들을 융합할 수 있습니다.
합성
헬퍼 함수는 단지 같은 그래프를 더 많이 구성할 뿐입니다:
from tinymlir import tanh
def activation(x):
return 0.5 * x * (1 + tanh(x))
@jit
def activated_scale(x, bias):
return activation(x * 2.0 + bias)
컴파일러에는 특별한 activated_scale 커널이 없습니다. activation도 별도의 백엔드 구현이 필요하지 않습니다. 컴파일러에 필요한 것은 구성 요소가 되는 텐서 연산들에 대한 규칙뿐입니다. 공식을 바꾸면 그래프와 생성된 모듈도 함께 바뀝니다.
중첩된 @jit 함수 호출도 트레이싱 중에는 같은 방식으로 동작합니다. 인수가 심볼릭 값이면, 호출은 둘러싼 그래프에 표현식을 추가합니다. 인수가 배열이면, 컴파일된 특수화를 선택하여 실행합니다.
NumPy 호출
연산자 오버로딩으로 x * 2.0 + bias는 처리할 수 있지만, NumPy 코드는 np.sqrt, np.tanh, np.mean, np.split 같은 함수도 호출합니다. 심볼릭 값은 두 가지 NumPy 디스패치 프로토콜을 구현하므로 이러한 호출도 트레이싱됩니다. __array_ufunc__는 원소별 유니버설 함수를 처리하므로, 트레이싱된 값에 대한 np.tanh(x)는 우리가 만든 tanh(x)와 같은 노드를 만듭니다. __array_function__는 mean, var, sum, max, split, reshape, transpose로 이루어진 짧은 배열 함수 표를 처리합니다. NumPy는 인수 중 하나라도 우리 타입이면 이 메서드들을 호출합니다.
@jit
def normalize(x):
centered = x - np.mean(x, axis=-1, keepdims=True)
return centered / np.sqrt(np.var(x, axis=-1, keepdims=True) + 1e-5)
이것은 평범한 NumPy 코드입니다. 배열에 대해 호출하면 NumPy에서 실행되고, JIT를 통해 호출하면 그래프로 트레이싱됩니다. 표에 없는 함수는 NotImplemented를 반환하며, NumPy는 우리가 컴파일할 수 없는 것을 조용히 평가하는 대신 TypeError를 발생시킵니다. 덕분에 14부에서는 참조 GPT-2 구현을 우리 라이브러리에 맞게 다시 작성하지 않고도 트레이싱할 수 있습니다.
트레이싱되는 함수는 마스크나 Python에서 만든 표처럼 바깥 스코프에 있는 배열을 사용할 수도 있습니다. 이렇게 캡처된 NumPy 배열은 상수 노드가 됩니다. JIT는 배열의 내용을 IR에 임베딩하는 대신, 캡처된 각 배열을 컴파일된 함수의 뒤쪽 인수로 끌어올려 매 호출마다 전달합니다. 특수화는 배열의 값이 아니라 shape과 dtype에 따라 달라집니다.
결과 살펴보기
컴파일된 객체는 세 개의 모듈을 보관합니다.
compiled = scale_bias.compile(x, bias)
print(compiled.module)
print(compiled.optimized)
print(compiled.lowered)
첫 번째 모듈에는 빌더가 생성한 텐서 프로그램이 들어 있습니다. 두 번째는 정규화, 공통 하위 표현식 제거, 원소별 융합을 거친 결과를 보여 줍니다. 세 번째 모듈에는 실행 엔진이 사용하는 LLVM 다이얼렉트 프로그램이 들어 있습니다. 이들은 별개의 모듈 객체이므로, 패스를 실행해도 우리가 살펴보려는 버전이 망가지지 않습니다.
런타임은 출력 배열을 할당하고, 인수에 대한 ranked memref 디스크립터를 구성한 다음, 생성된 C 인터페이스를 호출합니다. 컴파일된 객체는 자신의 컨텍스트와 실행 엔진을 계속 살려 둡니다. 컴파일이 실패하더라도 NumPy 실행으로 대체하는 폴백은 없습니다.
이로써 작은 컴파일러의 윤곽이 갖추어집니다. 제한된 Python 계산을 트레이싱하고, 텐서 타입을 검사하고, 구조화된 MLIR을 구성하고, 변환을 실행한 뒤, 로워링된 결과를 실행합니다. 다음으로는 소프트맥스와 레이어 정규화에 필요한 리덕션을 추가하겠습니다.
외부 자료
이전: 8부 - MLIR을 이용한 GPU 컴파일 · 시리즈 목차 · 다음: 10부 - 트랜스포머 커널 컴파일