융합, 버퍼화 그리고 실행
이제 우리 프런트엔드는 트랜스포머를 텐서 연산으로 기술합니다. 저장 공간과 실행에 관한 문제는 의도적으로 이러한 정의에서 제외했습니다. 이 장에서는 컴파일러의 나머지 절반을 따라갑니다. 즉, 텐서 프로그램을 변환하고, 버퍼를 할당하고, LLVM으로 로어링한 뒤, 그 결과를 Python에서 호출하는 과정입니다.
구현은 tinymlir/jit.py에 있습니다. 이 구현은 원본 모듈, 최적화된 모듈, 로어링된 모듈을 각각 따로 보관하므로 각 단계가 무엇을 바꾸는지 살펴볼 수 있습니다.
원소별 융합
간단한 함수를 하나 생각해 봅시다.
from tinymlir import jit
@jit
def affine_activation(x, gain, bias):
return (x * gain + bias) * 0.5
빌더는 처음에 원소별 linalg.generic 연산 세 개를 생성합니다. 각 연산에는 자체 출력 텐서가 있습니다. 텐서 수준에서 이 값들은 계산의 결과를 기술할 뿐이며, 이 값들이 존재한다고 해서 아직 세 번의 별도 힙 할당이 필요한 것은 아닙니다.
우리의 최적화 파이프라인은 다음과 같습니다.
OPTIMIZATION_PIPELINE = (
"builtin.module("
"canonicalize,cse,linalg-fuse-elementwise-ops,canonicalize,cse)"
)
융합 패스는 호환되는 생산자 영역과 소비자 영역을 결합하면서 그 인덱싱 맵을 합성합니다. 이 예제에서 최적화된 모듈에는 linalg.generic이 하나만 들어 있습니다. 그 스칼라 본문은 곱셈, 덧셈, 그리고 다시 곱셈을 수행한 뒤 결과를 내보냅니다.
테스트 스위트는 두 가지 사실을 모두 확인합니다. 연산 개수가 세 개에서 한 개로 바뀌는지, 그리고 최적화된 모듈과 최적화되지 않은 모듈을 실행했을 때 같은 배열이 나오는지입니다. 이는 따로 작성한 두 커널 구현을 비교하는 것이 아니라 실제 컴파일러 변환을 검증합니다.
패스 파이프라인 문자열은 MLIR의 패스 관리자에게 적용할 변환을 지정합니다. 변환되는 프로그램은 Python 바인딩을 통해 구성된 객체 그래프입니다.
from mlir.passmanager import PassManager
PassManager.parse(OPTIMIZATION_PIPELINE).run(module.operation)
module.operation.verify()
검증기는 변환 후의 결과를 검사합니다. 오류는 호출자에게 전파되며, 컴파일 실패를 감출 수 있는 대체 NumPy 경로는 없습니다.
융합의 한계
원소별 융합은 중간 저장, 로드, 할당을 없앨 수 있기 때문에 유용합니다. 그렇다고 해서 모든 연산 사슬이 하나의 영역이 되어야 한다는 뜻은 아닙니다. 리덕션은 반복 의미론이 다르고, 행렬 곱셈에는 리덕션 차원이 있으며, 하나의 생산자에 소비자가 여럿 있을 수도 있습니다.
예를 들어 소프트맥스의 지수 값들은 합계와 마지막 나눗셈 양쪽에 모두 입력됩니다. 올바른 변환이라면 이러한 사용 패턴을 고려해야 합니다. 스케줄에 따라 결과를 저장하는 것과 다시 계산하는 것의 비용이 달라질 수 있습니다. 텐서 그래프는 컴파일러가 이러한 결정을 내리고 검사하기에 충분한 정보를 제공합니다.
부동소수점 산술도 변환을 제한합니다. 리덕션의 결합 순서를 바꾸면 결과가 달라질 수 있습니다. 우리 프런트엔드는 공식의 dtype 승격을 그대로 보존하며, 생성된 산술은 제한 없는 fast-math 동작을 요청하지 않습니다. 그래도 NumPy의 리덕션과 BLAS 구현은 다른 순서로 누산할 수 있으므로, 여전히 수치 허용 오차를 두고 NumPy와 비교합니다.
버퍼화
텐서 변환이 끝나면 one-shot 버퍼화를 실행합니다. 이 과정은 텐서의 사용 방식을 분석하여 어디서 저장 공간을 재사용할 수 있고 어디서 별도의 할당이 필요한지 결정합니다.
외부 함수에는 이미 간단한 계약이 있습니다. memref 입력들과 호출자가 제공하는 출력 memref입니다. 함수 내부에서는 bufferization.to_tensor가 입력을 텐서 연산에서 사용할 수 있게 합니다. bufferization.materialize_in_destination은 최종 결과가 출력 버퍼에 기록되도록 요구합니다.
입력 변환에는 restrict를 표시합니다. 이는 경계에서의 앨리어싱에 관한 약속입니다. 호출자가 같은 배열을 두 번 전달하거나 겹치는 뷰를 전달하면, 런타임은 함수를 호출하기 전에 겹치는 입력을 복사합니다. 새로 할당된 출력 역시 입력과 구별됩니다. 따라서 버퍼화에 한 약속은 런타임이 실제로 제공하는 것과 일치합니다.
텐서 프로그램의 관점에서 입력은 읽기 전용입니다. 중간 결과는 버퍼화가 저장 공간을 정할 때까지 함수형 값으로 남아 있습니다. 블록을 평가하는 동안 모델의 가중치는 수정하지 않습니다.
저장 공간의 소유권은 버퍼 배정과는 별개의 문제입니다. 우리 파이프라인에는 buffer-deallocation 파이프라인이 포함되어 있어 내부에서 할당된 버퍼가 해제됩니다. 호출자의 NumPy 배열은 계속 Python이 소유합니다. 관리되지 않는 임시 버퍼의 포인터를 반환하면 작은 예제가 성공하더라도 오해를 불러일으킬 수 있습니다. 추론을 반복하려면 소유권 경계도 올바라야 합니다.
모듈 로어링하기
융합이 끝나면 동반 프로젝트는 두 단계로 로어링하며, 그 사이에 타깃별 단계가 하나 있습니다.
BUFFERIZATION_PIPELINE = (
"builtin.module("
"one-shot-bufferize{bufferize-function-boundaries},buffer-deallocation-pipeline)"
)
LOWERING_PIPELINE = (
"builtin.module("
"convert-linalg-to-loops,"
"expand-strided-metadata,lower-affine,"
"convert-scf-to-cf,convert-math-to-libm,convert-math-to-llvm,"
"convert-arith-to-llvm,convert-func-to-llvm,convert-cf-to-llvm,"
"finalize-memref-to-llvm,reconcile-unrealized-casts)"
)
버퍼화는 모든 타깃에 대해 한 번 실행됩니다. 버퍼화는 텐서 저장 공간을 확정하고 해제 연산을 삽입합니다. 그 결과는 memref 위의 linalg 연산으로 이루어진 모듈이며, 이 모듈에서 행렬 곱셈은 여전히 linalg.matmul과 linalg.batch_matmul입니다. 여기서 타깃별로 경로가 갈립니다.
PassManager.parse(BUFFERIZATION_PIPELINE).run(lowered.operation)
if target == GPU:
gpu.lower(lowered)
else:
blas.rewrite(lowered)
PassManager.parse(LOWERING_PIPELINE).run(lowered.operation)
CPU에서는 blas.rewrite가 모듈을 순회하며 행 우선 matmul을 찾아 각각을 cblas_sgemm 또는 cblas_dgemm 호출로 바꿉니다. 라이브러리는 프로세스가 불러올 수 있는 것이면 무엇이든 되며, macOS에서는 Accelerate, Linux에서는 OpenBLAS입니다. NumPy 자체의 행렬 곱셈도 같은 방식으로 BLAS를 호출합니다. 이 재작성은 memref.extract_strided_metadata로 각 memref에서 정렬된 기준 포인터, 오프셋, 크기, 스트라이드를 읽어 내고, 행 스트라이드를 leading dimension으로 전달하며, 배치 곱은 GEMM 호출을 반복하는 scf.for 루프로 바꿉니다. BLAS가 없으면 이 재작성은 아무것도 하지 않으며, matmul은 다른 모든 연산과 함께 루프로 로어링됩니다. 프런트엔드가 융합을 거치는 동안에도 matmul을 알아볼 수 있는 형태로 유지하도록 신경 쓴 이유가 바로 이것입니다.
나머지 단계는 제거되는 정보를 따라 진행됩니다. Linalg 로어링은 명시적인 루프를 선택합니다. 메타데이터 확장과 affine 로어링은 주소 계산을 드러냅니다. 구조적 제어 흐름은 일반적인 분기가 됩니다. 그다음 산술 연산, 함수, 메모리 참조가 LLVM 다이얼렉트 연산이 됩니다. 수학 연산은 이 고정된 툴체인에서 사용할 수 있는 LLVM 또는 호스트 수학 라이브러리 경로를 사용합니다. 마지막으로, reconciliation은 소스 타입과 대상 타입이 같아진 변환 잔여물을 제거합니다.
CPU 파이프라인은 Python 바인딩과 같은 프로세스에서 실행되며, 바인딩과 짝을 이루는 실행 엔진을 사용합니다. 모듈을 하위 프로세스로 직렬화하여 서로 무관한 LLVM 빌드를 결합하려 하는 것이 아닙니다. GPU 경로는 같은 버퍼화된 모듈을 다른 방향으로 가져가며, 별도의 장에서 다룹니다.
네이티브 코드 호출
이 함수에는 llvm.emit_c_interface가 붙어 있으며, 이는 C 호환 호출 규약을 따르는 래퍼를 요청합니다. 런타임은 MLIR의 랭크가 있는 memref 디스크립터 헬퍼를 사용합니다.
pointers = [
ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(array)))
for array in [*inputs, output]
]
engine.invoke("kernel", *pointers)
실행 엔진의 packed 호출 API는 인자 저장 공간에 대한 포인터를 받습니다. C 인터페이스의 memref 인자는 그 자체로 디스크립터 포인터이므로, 여기서 간접 참조가 한 단계 더 생깁니다. 디스크립터는 배열의 주소, 오프셋, 크기, 스트라이드를 기록합니다.
함수를 호출하기 전에 런타임은 특수화의 shape와 dtype을 검사하고 입력을 연속 메모리로 만듭니다. 그리고 추론된 결과 타입의 NumPy 출력 배열을 할당합니다. 배열과 디스크립터는 호출이 끝날 때까지 살아 있으며, 컴파일된 객체는 해당 특수화가 캐시되어 있는 동안 컨텍스트, 모듈, 엔진을 소유합니다.
텐서 값으로 된 gather 인덱스도 네이티브 코드에 들어가기 전에 검사합니다. shape 검사만으로는 런타임 토큰 ID가 임베딩 테이블 범위 안에 있음을 증명할 수 없습니다. 모델은 어휘 범위를 검증하고, 범용 gather 경계는 자체적으로 인덱스를 검사합니다.
특수화와 비용
JIT 캐시의 키는 인자의 구조, 정적 키워드 인자, 타깃으로 구성됩니다. 인자는 트리로 평탄화됩니다. 튜플, named tuple, 데이터클래스, 딕셔너리는 배열에 도달할 때까지 재귀적으로 분해되며, 키는 트리의 모양과 함께 모든 리프의 shape와 dtype을 기록합니다. 따라서 컴파일된 함수는 배열의 함수로 남으면서도 매개변수 레코드 전체를 인자로 받을 수 있습니다. 가중치 값이 서로 다른 트랜스포머 블록이라도 인자의 shape와 dtype이 일치하면 하나의 특수화를 공유합니다. 가중치는 컴파일된 각 모듈에 포함된 상수가 아니라 런타임 입력입니다.
시퀀스 길이가 새로 바뀌면 shape가 달라지므로 특수화가 하나 더 필요합니다. 이 예제는 토큰을 하나 생성할 때마다 접두부 전체를 다시 계산하므로, 첫 실행에는 여러 길이에 대한 컴파일이 포함됩니다. 이미 있는 shape로 반복 호출하면 엔진을 재사용합니다.
이것은 의도적으로 작게 만든 프로세스 내 캐시입니다. 영속적인 객체 캐싱이나, 끝없이 이어지는 형상들에 대한 메모리 정책은 제공하지 않습니다. 그런 기능은 유용한 확장이 되겠지만, 컴파일 경계를 이해하는 데 꼭 필요하지는 않습니다.
실행 시간을 측정할 때에는 먼저 관련 특수화를 워밍업하십시오. 그렇지 않으면 측정 시간에 추론뿐 아니라 트레이싱, IR 구성, 패스 실행, 기계어 코드 생성까지 포함됩니다. 이러한 비용이 섞인 비교로는 생성된 커널에 대해 알 수 있는 것이 거의 없습니다.
다시 GPU로
8부에서는 GPU 실행 모델과 PTX 경로를 다루었습니다. 이 텐서 컴파일러의 GPU 백엔드는 구조화된 연산을 로어링하여 없애기 전에 다른 스케줄링 결정을 내립니다. 행렬 곱셈은 타일링되고, 루프 반복은 블록과 스레드에 매핑되며, 호스트-디바이스 경계에서 런치가 나타납니다.
버퍼화된 모듈에 남아 있는 linalg.matmul과 리덕션 연산은 이러한 백엔드에 변환할 만한 유용한 대상을 제공합니다. 모든 것이 스칼라 루프와 포인터 산술로 바뀌고 나면, 그 구조를 복원하기가 더 어려워집니다.
tiny-mlir 프로젝트에 이 백엔드가 있습니다. 예제나 검증 스크립트에 --gpu를 전달하면, matmul을 공유 메모리 타일링 커널로 대체하고, 나머지 병렬 루프를 블록과 스레드에 매핑하며, 이를 GPU 커널로 분리한 뒤 그 결과를 CUDA 관리 메모리 위에서 실행하는 타깃이 선택됩니다. 이 백엔드가 내보내는 PTX는 지속적 통합에서 ptxas로 검사되며, 전체 레퍼런스 비교는 Tesla T4에서 실행되었습니다. 두 타깃 사이에서 Python 모델의 공식은 하나도 바뀌지 않습니다. 14부에서 이 백엔드를 자세히 살펴봅니다.
다음으로는 사전 학습된 체크포인트를 불러와, 생성된 토큰의 정확한 시퀀스를 포함한 전체 추론 경로를 테스트하겠습니다.