cuTile.jl: NVIDIA CUDA 타일 기반 프로그래밍, 이제 Julia에서도 만나보세요

NVIDIA CUDA 타일은 NVIDIA CUDA 프로그래밍 역사에서 가장 중요한 진보 중 하나로 손꼽힙니다. 이 기술은 텐서 코어를 비롯한 특수 하드웨어에 대한 자동 접근 권한을 부여하며 개발자들에게 새로운 가능성을 열어주었습니다. 올해 초, NVIDIA는 Python 개발자들이 고성능 GPU 커널을 자연스럽게 작성할 수 있도록 cuTile for Python을 출시한 바 있습니다.

이제 동일한 프로그래밍 모델을 cuTile.jl을 통해 Julia 환경에서도 누릴 수 있습니다. 이번 포스팅에서는 cuTile.jl이 고성능 CUDA 커널 개발을 어떻게 단순화하는지 살펴보고, Julia 특유의 관용적 구문과 함께 기존 Python 구현체와 대등한 수준의 성능 지표를 심도 있게 분석합니다.

타일 기반 GPU 프로그래밍이란 무엇인가?

전통적인 CUDA 기반 GPU 프로그래밍 환경에서 개발자는 스레드(Threads), 워프(Warps), 그리고 복잡한 메모리 계층 구조를 끊임없이 고민해야 했습니다. 이러한 방식은 강력한 제어권을 제공하지만, 프로그래머가 알고리즘을 하드웨어 구조에 효율적으로 매핑해야 한다는 커다란 숙제를 안겨주기도 합니다. 반면 CUDA 타일을 사용하면 개발자는 데이터의 ‘타일(Tiles)’ 단위로 연산을 기술하기만 하면 됩니다. 나머지 하드웨어 매핑 작업은 컴파일러가 지능적으로 처리합니다.

벡터 덧셈(Vector Addition) 사례를 통해 그 차이를 확인해 보겠습니다. 기존의 CUDA.jl 프로그래밍 모델에서는 프로그래머가 다음과 같이 개별 스레드를 명시적으로 관리해야만 했습니다.

using CUDA
 
function vadd(a, b, c, n)
    i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    if i <= n
        @inbounds c[i] = a[i] + b[i]
    end
    return
end
 
threads = 512
blocks = cld(vector_size, threads)
@cuda threads blocks vadd(a, b, c, vector_size)

이제 cuTile.jl을 통한 CUDA 타일 방식에서는 동일한 연산이 타일 수준에서 표현됩니다. 인덱스 계산이나 범위를 벗어난 참조(Out-of-bounds) 체크와 같은 번거로운 세부 사항은 숨겨집니다.

import cuTile as ct
 
function vadd(a, b, c, tile_size)
    pid = ct.bid(1)
    tile_a = ct.load(a, pid, (tile_size,))
    tile_b = ct.load(b, pid, (tile_size,))
    ct.store(c, pid, tile_a + tile_b)
    return
end
 
tile_size = 1024
grid = cld(vector_size, tile_size)
ct.launch(vadd, grid, a, b, c, ct.Constant(tile_size))

이 방식과 Python 환경의 구현체를 비교해 볼 수 있습니다:

@ct.kernel
def vadd(a, b, c, tile_size: ct.Constant[int]):
    pid = ct.bid(0)
    tile_a = ct.load(a, index=(pid,), shape=(tile_size,))
    tile_b = ct.load(b, index=(pid,), shape=(tile_size,))
    ct.store(c, index=(pid,), tile=tile_a + tile_b)
 
tile_size = 1024
grid = ceil(vector_size / tile_size)
ct.launch(stream, grid, vadd, (a, b, c, tile_size))

두 코드는 놀라울 정도로 유사하며, 이는 의도된 설계의 결과입니다. cuTile.jl은 cuTile Python으로 작성된 커널과 동일한 추상화 수준을 유지하므로, 기존 코드를 포팅하거나 Python 문서의 도움을 받아 학습하기에 매우 용이합니다. 동시에 1부터 시작하는 인덱싱이나 요소별 연산을 위한 브로드캐스트 표현식 등 Julia 특유의 관용구를 최대한 활용하여 Julia 개발자들에게 직관적인 개발 경험을 선사합니다.

Julia다운 관용적 커널의 구현

이 기술의 진가는 단순한 데이터 로드와 저장을 넘어선 복잡한 커널에서 더욱 빛을 발합니다. 다음은 레이어 정규화의 핵심이자, 가중치와 편향을 제외한 형태인 행 정규화 커널의 예시입니다.

function normalize_rows(X, Y, tile_n)
    bid = ct.bid(1)
    tile = ct.load(X, (bid, 1), (1, tile_n))
    mean = sum(tile; dims=2) / size(X, 2)
    centered = tile .- mean
    var = sum(centered .^ 2.0f0; dims=2) / size(X, 2)
    ct.store(Y, (bid, 1), centered ./ sqrt.(var .+ 1f-5))
    return
end

점(dot) 기호들은 표준 Julia 브로드캐스팅 구문을 그대로 따르며, 각 연산이 요소별로 정교하게 적용됨이 예제에서 사용된 sum, size, sqrt는 타일 단위 연산이 가능하도록 확장된 표준 Julia 함수입니다. 또한 .^, .-, ./와 같은 기호는 표준 Julia 브로드캐스팅 구문을 그대로 따르며, 각 연산이 요소별로 정교하게 적용됨을 보여줍니다. 이 커널은 마치 평범한 Julia 배열 코드를 읽는 것과 같은 가독성을 제공합니다. cuTile.jl 커널이 일반적인 Julia 코드와 유사해질수록, CPU와 GPU 사이에서 코드를 공유하고 재사용하는 작업은 더욱 수월해집니다.

cuTile.jl의 성능 지표

cuTile.jl은 cuTile Python과 동일한 NVIDIA Tile IR 백엔드를 대상으로 합니다. 따라서 두 패키지 모두 동일한 종류의 GPU 기계 코드를 생성합니다. NVIDIA Blackwell 아키텍처 기반의 NVIDIA GeForce RTX 5080(컴퓨트 성능 12.0)에서 테스트한 결과, 연산 집약적인 커널 환경에서 Python 구현체와 대등한 성능을 달성했습니다.

레이어 정규화나 FFT와 같이 복잡한 제어 흐름을 가진 일부 커널의 경우, cuTile.jl 컴파일러가 아직 성숙 단계에 있어 완전한 성능 대등성을 확보하지 못했습니다. 현재 이러한 사례들은 알려진 이슈로 분류되어 있으며, 성능 최적화를 위한 개선 작업이 활발히 진행 중입니다.

cuTile.jl의 작동 원리

cuTile.jl은 커스텀 Julia 컴파일러를 활용하여 +, sum, reshape와 같은 표준 라이브러리 호출을 가로챕니다. 이후 해당 호출들을 타일 IR 연산으로 경로를 변경하며, 생성된 IR은 cuTile Python이 생성하는 것과 동일한 이진 형식인 타일 IR 바이트코드로 변환됩니다. 최종적으로는 NVIDIA의 tileiras 컴파일러가 이 바이트코드를 처리하여 GPU 기계 코드로의 마지막 컴파일 단계를 수행합니다.

개발자는 어떠한 커널에 대해서도 생성된 타일 IR을 직접 검사하고 분석할 수 있습니다.

julia> ct.@device_code_tiled ct.launch(vadd, grid, a, b, c, ct.Constant(16))
cuda_tile.module @kernels {
  entry @vadd(%arg0: tile<ptr<f32>>, %arg1: tile<i32>, ...) {
    ...
    return
  }
}

이러한 투명성은 디버깅에 도움이 될 뿐만 아니라, 고수준 Julia 코드가 타일 연산으로 어떻게 매핑되는지 이해하는 데에도 유용합니다.

cuTile.jl의 현재 상태와 발전 방향

cuTile.jl은 현재 JuliaGPU/cuTile.jl 리포지토리에서 활발히 개발 중인 실험적인 오픈소스 패키지입니다. 이 패키지는 메모리 접근, 산술 연산, 리덕션(Reductions), 스캔(Scans), 행렬 곱셈, 형상 조작(Shape manipulation) 및 아토믹(Atomics) 연산 등 광범위한 타일 작업을 이미 지원하고 있습니다. 또한 벡터 덧셈, 행렬 곱셈, 전치(Transpose), 배치 행렬 곱셈, 레이어 정규화 및 FFT를 위한 실무 예제 코드도 함께 포함하고 있습니다.

다만, 본 소프트웨어는 초기 개발 단계에 머물러 있음을 참고해 주세요.

• 미구현 기능: 모든 cuTile 기능이 완벽히 이식되지는 않았습니다.
• 언어 제약: 일부 Julia 언어 기능(특히 반복자 기반의 for 루프)은 커널 내에서 아직 지원되지 않거나 비효율적인 코드를 생성할 수 있습니다.
• 통합 과제: SIMT 커널과의 공존을 원활하게 만들기 위해 CUDA.jl과의 통합 수준을 더욱 높여야 합니다.
• API 변동성: 개발 진행 상황에 따라 API 명세가 예고 없이 변경될 가능성이 존재합니다.

이 프로젝트는 기존 Julia GPU 생태계를 기반으로 구축되었습니다. 배열 관리 및 커널 실행을 위해 CUDA.jl과 긴밀히 통합되므로, 이미 CUDA.jl로 GPU 코드를 작성해 온 사용자라면 타일 기반 프로그래밍으로의 전환이 매우 직관적임을 체감하실 것입니다.

시작하기

cuTile Python과 마찬가지로, cuTile.jl을 활용하기 위해서는 NVIDIA Blackwell GPU와 CUDA 13 이상의 버전이 설치된 NVIDIA 드라이버가 필요합니다. 또한 Julia 1.11 이상의 환경을 권장합니다.

Julia를 실행한 후, REPL에서`]`키를 눌러 통합 패키지 매니저에 진입한 뒤 아래 명령어로 cuTile.jl을 설치해 보세요.

pkg> add cuTile
 
pkg> # if you want, run the test suite
     test cuTile

cuTile.jl의 공식 GitHub 리포지토리에는 지원되는 전체 연산 목록과 함께, 본 패키지가 cuTile Python 및 표준 Julia와 차별화되는 지점에 대한 상세한 문서가 수록되어 있습니다.