강의 번호
L062
Halide 가 깐 algorithm 과 schedule 의 분리 를 — 그 schedule 자체를 사용자 정의 가능한 라이브러리 로 만든 MIT 의 DSL. Exo 의 핵심 아이디어는 “scheduling language 는 fixed 가 아니라 grown 된다” 는 것 — 새 hardware 가 등장하면 컴파일러를 갱신하지 않고 스케줄과 instruction 을 사용자가 추가한다. Yuka Ikarashi 가 Exo 2 에서 깐 새 기능과 검증 모델, 그리고 Halide / TVM 과의 차이를 정리한 노트.
2024–25 시점의 hardware 풍경 — NVIDIA H100/B200, AMD MI300, Apple AMX, Google TPU, AWS Trainium, Meta MTIA, 그리고 D-Matrix / Groq / Cerebras / Tenstorrent 의 dataflow 칩. 각자 자기 ISA 와 자기 memory hierarchy 를 갖는다. 컴파일러가 모두 따라가는 건 불가능. Exo 의 답 — schedule 과 instruction 을 사용자가 라이브러리로 추가할 수 있게 한다.
강의 transcript 가 비어 있으므로 본 노트는 Exo 의 PLDI 2022 paper, Exo 2 의 후속 paper (Ikarashi et al. 2024+), exo-lang.dev 자료, github.com/exo-lang/exo repo 위에서 재구성한다. 강의에서 Yuka 가 강조한 정확한 demo / 채택 사례는 영상 검증 필요.
기존 컴파일러 (Halide, TVM, MLIR) 의 schedule 과 lowering 은 컴파일러 코드 안에 박혀 있다. Exo 는 — 그 schedule rewrite 를 Python library 로 노출. 사용자가 자기 hardware 에 맞는 새 transformation 을 정의하고, 새 ISA instruction 을 추가하고, 그 결합으로 자기 커널을 짠다. 컴파일러 자체가 “growable” 한 것.
Halide (2012, MIT/Stanford) 가 깐 큰 idea — image / tensor 코드의 알고리즘 정의 (무엇을 계산하는가) 와 schedule (어떻게 도는가) 을 분리. tile size, loop order, vectorization, parallelization 을 하나의 큰 transformation 라이브러리로 표현. 같은 알고리즘이 — schedule 만 바꾸면 — CPU 의 SIMD, GPU 의 thread, ARM 의 NEON 을 다 쓸 수 있다.
Halide 의 한계는 — schedule 의 set 이 컴파일러에 박혀 있다는 점. split, tile, vectorize, parallel, fuse 정도. 새 hardware 의 특수 instruction (예: GEMMINI 의 systolic array load) 을 표현하려면 컴파일러를 수정해야 한다. Exo 는 이 자리를 그 transformation 도 라이브러리화하는 식으로 푼다.
Exo 의 입력은 Procedure (proc) — Python decorator 안의 함수가 ImperativeIR 로 lowering. 처음에는 단순한 reference algorithm — naive nested loop 의 GEMM, conv. 그 다음 같은 proc 위에 schedule transformation 을 chain. 각 transformation 이 새 proc 를 반환 (functional).
# Exo 의 reference GEMM
@proc
def gemm(M:size, N:size, K:size,
A:f32[M,K], B:f32[K,N], C:f32[M,N]):
for i in seq(0, M):
for j in seq(0, N):
for k in seq(0, K):
C[i,j] += A[i,k] * B[k,j]# schedule 의 chain — 각 단계가 새 proc 반환
g = gemm
g = divide_loop(g, "i", 16, ["io", "ii"])
g = divide_loop(g, "j", 16, ["jo", "ji"])
g = reorder_loops(g, "ii jo")
g = stage_mem(g, "C[ii,ji]", "Cb")
g = call_eqv(g, "Cb", gemm_tile_inst)각 함수 — divide_loop, reorder_loops, stage_mem — 가 Exo 의 표준 transformation 라이브러리에 있는 한 종류. 모두 입력 proc 을 받아 출력 proc 을 돌려주는 순수 함수.
Exo 2 의 새 기능 — cursor. 코드의 특정 위치 (예: 안 쪽 loop, 특정 statement) 를 가리키는 1급 객체. transformation 이 cursor 를 받고 cursor 를 반환. 사용자 정의 schedule 라이브러리를 만들 때 이 cursor 가 필요한 robust composition 을 제공.
O(MNK) 의 단순 reference. 어떤 hw 도 비효율.
divide_loop(i, 16, ...), divide_loop(j, 16, ...) 로 16×16 tile 구조. cache 친화.
stage_mem(g, "C[ii,ji]", "Cb").
call_eqv(g, ..., gemm_tile_inst).
Exo 1 (PLDI 2022) 이 깐 idea — algorithm/schedule 분리 + 사용자 정의 instruction. Exo 2 는 한 단계 더 — schedule transformation 자체도 사용자가 정의하는 함수 가 되게 한다. 즉 “tile + stage + replace_with_inst” 같은 자주 쓰이는 패턴을 사용자가 자기 라이브러리로 묶고, 다른 사용자가 import 해서 쓴다.
divide_loop, reorder_loops, stage_mem 같은 ~20 개의 built-in. 사용자는 이걸 chain 만 한다. 새 패턴이 생기면 결국 컴파일러를 fork.
한계 — 같은 8 step 의 chain 을 매번 반복. 추상화 안 됨.
cursor 추상 + 사용자 정의 schedule 함수. 자기 hw 의 conventional pattern (예: “tile-stage-load-vec-mma”) 을 한 함수로 묶음.
새 가용성 — 같은 hw 그룹 안에서 schedule lib 공유. community / vendor 가 자기 hw 의 schedule 을 sub-library 로 publish.
cursor 는 proc 안의 특정 위치 (loop, statement, expression 의 특정 instance) 를 가리킨다. transformation 이 cursor 를 받아 — 그 자리에서 rewrite. 같은 변환을 여러 자리에 반복 적용할 때, naming 만으로 파악되지 않는 위치를 정확히 가리킨다.
# Exo 2 — cursor 를 받는 사용자 schedule
def tile_and_vectorize(p, loop_cursor, tile_size):
p = divide_loop(p, loop_cursor, tile_size,
["o", "i"])
p = vectorize(p, p.find_loop("i"))
return p
# 같은 schedule 을 여러 loop 에 적용
g = gemm
g = tile_and_vectorize(g, g.find_loop("i"), 16)
g = tile_and_vectorize(g, g.find_loop("j"), 8)이 추상화 덕에 — 한 hw 그룹의 community 가 모여 “ARM Neon 의 conv2d 표준 schedule”, “GEMMINI 의 dense GEMM schedule” 같은 라이브러리를 만들 수 있다. 컴파일러를 수정하지 않고.
Exo 의 다른 핵심 기능 — instruction 도 사용자 정의. 어떤 hw 의 inner-loop tile 처리 instruction 을 — Exo 안에서 reference Python 함수로 정의 + 그 instruction 의 C/asm 표현을 명시. schedule 은 이 instruction 을 호출하는 형태로 lowering.
# 예 — 가상의 16×16 systolic GEMM tile instruction
@instr("my_tile_mma_16x16(&{C_data}, &{A_data}, &{B_data})")
def gemm_tile_inst(C: f32[16,16] @ Mem,
A: f32[16,16] @ Mem,
B: f32[16,16] @ Mem):
# reference semantics — 이게 hw 가 하는 일
for i in seq(0, 16):
for j in seq(0, 16):
for k in seq(0, 16):
C[i,j] += A[i,k] * B[k,j]이 정의는 두 가지를 한다 — (1) semantics: instruction 이 정확히 어떤 결과를 만드는지. Exo 가 schedule 의 정확성을 검증할 때 이 reference 를 쓴다. (2) codegen: 실제 emit 시 instruction 의 textual form (C 호출 / asm) 으로 나간다.
Exo 는 buffer 가 어디에 살고 있는지 (HBM, SRAM, register) 를 type 차원에서 추적. instruction 이 “이 16×16 buffer 가 systolic 의 register 에 있어야 한다” 같은 제약을 표현 가능. schedule 의 stage_mem 이 이 위치를 정확히 만들어준다. 잘못된 위치에 있으면 컴파일이 거부.
그래서 새 hw 를 지원할 때 사용자가 하는 일 — (a) 그 hw 의 instruction 을 @instr 로 모두 정의, (b) 그 hw 의 memory hierarchy (HBM, SRAM, scratchpad) 를 annotation 으로 정의, (c) 자기 알고리즘을 그 instruction 으로 lowering 시키는 schedule chain 을 짠다. 컴파일러는 손 안 댄다.
Exo 의 정확성 보장 — 모든 schedule transformation 이 semantic-preserving. divide_loop, reorder_loops, stage_mem 같은 각 primitive 가 정형적으로 검증된 rewrite. 사용자가 잘못된 위치에 transformation 을 걸면 컴파일러가 그 자리에서 거부.
예를 들어 — reorder_loops 가 두 loop 의 순서를 바꾸려는데 — 그 순서 교환이 dependency 를 깨면 (예: 안쪽이 바깥쪽의 update 에 의존), Exo 가 거부한다. “이 reorder 는 dependence 를 위반함” 이라는 정확한 에러 메시지.
Halide / TVM 의 schedule 도 어느 정도 검증을 하지만 — 사용자 정의 transformation 이 들어가는 순간 그 보장이 약해진다. Exo 는 새 transformation 을 정의할 때도 같은 정확성 framework 안에서 정의된다. 즉 community / vendor 가 라이브러리를 publish 해도 그게 wrong code 를 만들 수 없다. scaling trust.
이 framework 는 PLDI 22 paper 의 본격 contribution. “effects” 라는 abstraction 이 핵심 — 각 instruction 이 어떤 buffer 의 어떤 영역을 read/write 하는지를 명시. transformation 이 이 effects 를 보존하는지 자동 검증.
Exo paper 들이 보여준 실험 / production 사례 (확인 필요 — 강의에서 다룬 정확한 채택 list).
Exo 의 niche 는 — GPU 보다 작고, GPU 보다 다양한 hardware. accelerator / DSP / 새 ISA. 그 영역에서 커널 한 개 짜는 게 큰 일이 되는 자리에서 Exo 가 효율을 낸다.
Exo 가 자동 search 를 안 한다는 점은 의도적 — Exo 는 전문가가 정확히 통제하는 도구. 자동 search 는 다른 layer (예: TVM 의 Ansor 처럼) 에서 한다는 분리. 강의 후속의 L063 (Search-based compilers) 와 정반대 방향에 있다는 점이 흥미로움.
“같은 알고리즘에 대해 모든 hw 가 자기만의 optimal schedule 을 갖는다. 그 schedule 은 자동 search 보다 — 그 hw 를 잘 아는 사람의 통제 하에 — 더 빠르고, 더 maintain 가능하게 만들 수 있다.” 이 frame 이 search-based 와 정반대.
그럼에도 불구하고 — “새 hw 가 등장할 때 컴파일러를 fork 하지 않고 schedule 라이브러리만 추가” 라는 자리는 Exo 가 거의 유일. accelerator 가 늘어나는 시대에서 의미가 커지는 영역.
pip install exo-lang. @proc 으로 naive nested-loop GEMM. Python 으로 직접 호출.@instr 로 가상의 16×16 GEMM tile inst. call_eqv 로 schedule 의 안쪽 loop 를 그 inst 로 치환.본 노트는 Exo 의 PLDI 2022 paper 와 공식 자료의 일반적 framework 를 기반으로 재구성. 강의에서 Yuka 가 보여준 정확한 demo / channel 별 디테일 / 청중 Q&A 는 직접 영상 시청으로 검증 필요. 특히 cursor 의 정확한 type 시스템과 schedule library 의 표준화 디테일은 paper 의 detail 을 직접 읽기.