드미트리 트리포노프(Dmitry Trifonov)가 Emmy 0.3.22를 사용해 PyTorch 연산을 CUDA 코드로 변환하는 컴파일러 구축 과정을 설명하는 튜토리얼을 공개했다. GPU가 없어도 중간 표현(IR)을 살펴볼 수 있지만, 벤치마크 실행에는 NVIDIA GPU와 nvcc가 필요하다.
튜토리얼은 PyTorch 그래프를 포착한 뒤 Tensor IR, Loop IR, Tile IR, Kernel IR을 거쳐 CUDA 코드로 낮추는 흐름을 RMSNorm, 행렬 곱셈, 어텐션 사례로 보여준다. 각 단계는 프레임워크별 연산을 공통 표현으로 바꾸고, 연산을 융합하거나 중복 계산을 제거하며, GPU 작업 배치를 정하고 최종 코드를 생성하는 역할을 맡는다.
Tile IR에서는 계산과 실행 스케줄을 분리한다. 스케줄 코덱은 타일 크기와 워프 배치 같은 선택을 조합해 이름을 붙이며, 컴파일러는 적법한 조합을 탐색하고 GPU에서 측정할 수 있다. RMSNorm 실험에서는 RTX 5090에서 가능한 63개 스케줄 중 12개를 측정했다. 가장 빠른 측정 결과는 2.91마이크로초였고, 기본 스케줄은 같은 조건에서 3.25마이크로초로, 1.12배 차이가 났다. 이 결과는 해당 GPU와 측정 조건에 한정되며, 다른 GPU나 측정된 스케줄에 따라 달라질 수 있다.
어텐션 예제에서는 일반적인 계산이 키를 세 차례 순회하고 점곱을 반복하는 구조를 보인 뒤, 온라인 소프트맥스 방식으로 최대값·분모·출력 누산 상태를 갱신하며 키를 한 차례 순회하는 표현을 설명한다. 다만 예제에서 고정한 간단한 스케줄은 성능 최적화용이 아니다. 출력 열마다 점수를 다시 계산하고 워프의 일부 스레드가 중복 작업을 한다고 명시한다.
설치와 실행 조건, 전체 예제는 Emmy 0.3.22 패키지와 원문 튜토리얼에서 확인할 수 있다. 원문은 컴파일러의 IR 구성과 스케줄 탐색 방법을 소개하는 기술 튜토리얼이며, 제시된 벤치마크 결과가 모든 GPU나 작업에서 같은 성능을 보장하지는 않는다.
부록: 원문 예시
// Selected lines from the 7-node Tensor IR output; names shortened. let wt: f32[3,2] = load("weight").transpose(axes=(-2, -1)); let a: f32[4,3,2] = emmy::index_map_simple(inputs.input, |i, j, _k| [i, j]); let b: f32[4,3,2] = emmy::index_map_simple(wt, |_i, j, k| [j, k]); let prod: f32[4,3,2] = a * b; let red: f32[4,1,2] = sum(prod, axis=-2);
emmy compile -c 'x=torch.randn(4,3); w=torch.randn(3,2); x @ w' --ir tensor
emmy compile -c 'x=torch.randn(8); y=torch.randn(8); x+x+y' --ir loop emmy compile -c 'x=torch.randn(8); y=torch.randn(8); y+(x+x)' --ir loop
// Sanitized Tile IR; names shortened and arithmetic written as operators. === 0: k_rms_norm === place free=(a0) grid=(a0) work t256 Fold free ├─ operand[v3]: Fold free ‹computed› │ ├─ operand[acc0]: Fold[a1 in 0..2048] reduce ⟨REDUCE=coop⟩ ‹computed› │ │ ├─ operand[in2]: load x[0, a0, a1] ‹materialized› │ │ ├─ init: (0) │ │ ├─ lift: λ(a1, in2) -> (v0) │ │ │ v0 = in2 * in2 │ │ └─ combine: λ(acc0, acc0__o) -> (acc0) │ │ acc0 += acc0__o │ └─ lift: λ(acc0) -> (v3) │ v1 = acc0 / 2048 │ v2 = 1e-06 + v1 │ v3 = rsqrt(v2) └─ lift: λ(v3, a0, a2) -> (v5) in4 = load p_weight[a2] in3 = load x[0, a0, a2] v4 = in3 * v3 v5 = in4 * v4 outputs └─ sweep(a2) rms_norm[0, a0, a2] = v5
EMMY_KNOBS='WORK=w1x2,TILE=mma_m16n8k16_f16_f32/f1x1/k8, REDUCE=,STAGE=d3/smem-tma,RASTER=' \ emmy compile -c 'torch.matmul(torch.randn(128,128,dtype=torch.float16), torch.randn(128,128,dtype=torch.float16))' \ --target sm_120 --ir tile EMMY_KNOBS='WORK=w2x2,TILE=mma_m16n8k16_f16_f32/f2x2/k4, REDUCE=,STAGE=d2/smem-tma,RASTER=' \ emmy compile -c 'torch.matmul(torch.randn(128,128,dtype=torch.float16), torch.randn(128,128,dtype=torch.float16))' \ --target sm_120 --ir tile
// Sanitized excerpt: eight unrolled loads and stores grouped as loops. === 0: k_rms_norm === Tile[row, col] (N=8192) // 32 rows, 256 workers per row f32 sum = 0 f32 x_reg[8] for u in 0..8: x_reg[u] = load x[0, row, col + 256*u] sum += x_reg[u] * x_reg[u] WarpShuffle(sum, length=32) // Combine lanes within each warp. Smem f32 warp_sums[8] if lane == 0: warp_sums[warp] = sum Sync TreeHalve(sum, length=8, tid=warp) // Combine the 8 warp sums. f32 scale = rsqrt(1e-6 + sum / 2048) for u in 0..8: rms_norm[0, row, col + 256*u] = x_reg[u] * scale * load p_weight[col + 256*u]
// Sanitized Fold tree; one free output coordinate is (i,c). place free=(i,c) grid=(i,c) work t32 Fold free -> O[i,c] └─ Fold[j in 0..16] contraction <twist=softmax> ├─ Fold[r in 0..16] contraction <TILE=f1> │ ├─ load Q[i,r] │ └─ load K[j,r] ├─ load V[j,c] ├─ init (m, d, z) = (-infinity, 0, 0) └─ output z / d
댓글 (0)
로그인하면 이 기사에 내 생각을 남길 수 있어요