Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

GPU Bear Teaches CUDA

MLIR을 이용한 GPU 컴파일

· 태그: ai, python, llms, compilers, mlir · 원문 · 한국어 번역(원문에서 변경됨)

7부에서는 GPT-2의 NumPy 구현을 만들었습니다. 이제 컴파일된 텐서 연산을 NVIDIA GPU에서 실행할 수 있는 백엔드가 필요합니다. 작은 요소별(elementwise) 커널에서 시작하여 이것이 MLIR에서 PTX로, 그리고 CUDA 실행(launch)에 이르기까지 따라가 보겠습니다. 소프트맥스와 어텐션에는 리덕션과 더 세심한 스레드 매핑이 필요하지만, 이를 통해 생성된 코드를 로드하고 실행할 수 있는 기반을 갖추게 됩니다.

이제 Nvidia GPU가 장착된 머신으로 옮겨 가야 합니다(죄송하지만 MacOS는 지원되지 않습니다). 따라서 Nvidia GPU가 장착된 Linux 머신을 켜거나, 여러 제공업체 중 한 곳에서 GPU를 대여하십시오.

CUDA Toolkit 설치는 번거롭기로 악명이 높으므로, 클라우드 제공업체에서 GPU를 대여한다면 이 작업은 이미 되어 있습니다. 맞춤형 인스턴스를 설정해야 한다면 다음과 같은 Nvidia의 공식 도커 이미지 중 하나를 사용하는 것이 좋습니다:

docker pull nvidia/cuda:12.1.0-devel-ubuntu22.04
docker run -it --gpus all nvidia/cuda:12.1.0-devel-ubuntu22.04 bash

하지만 직접 자신의 머신에 설치하고 싶거나 설치해야 한다면 아래에 몇 가지 안내가 있습니다. 다만 최신 안내는 온라인에 있는 여러 가이드를 참고하시기 바랍니다. CUDA Toolkit이 설치되어 있지 않은 머신이라면 다음 단계에 따라 설치할 수 있습니다.

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb
sudo dpkg -i cuda-keyring_1.0-1_all.deb
sudo apt-get update
sudo apt-get install build-essential

sudo apt-get -y install cuda-toolkit-12-x

nvcc 컴파일러가 설치되어 있는지 확인하십시오. 마찬가지로 GPU를 대여하는 경우라면 이미 설치되어 있을 가능성이 높습니다.

$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Thu_Nov_18_09:45:30_PST_2021
Cuda compilation tools, release 11.5, V11.5.119
Build cuda_11.5.r11.5/compiler.30672275_0

Nvidia GPU가 설치되어 있는지 확인하십시오:

lspci | grep -i nvidia

그런 다음 CUDA 버전과 드라이버 버전을 확인하십시오:

$ nvidia-smi
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.183.01             Driver Version: 535.183.01   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+

구체적인 CUDA 버전(이 경우 12.2)을 기록해 두십시오. 이 버전에 맞는 라이브러리를 설치해야 하므로 중요합니다.

일반적으로는 CUDA 러너를 활성화한 상태로 MLIR을 다시 컴파일해야 하는데, 이는 매우 번거로운 일입니다. 그래서 CUDA 러너가 활성화된 MLIR이 포함된 도커 이미지를 만들어 두었습니다. 따라서 RunPod이나 LambdaLabs에서 GPU를 대여한다면 이 베이스 이미지를 그대로 사용하면 됩니다. 또는 로컬 머신이 있다면 NVIDIA Container Toolkit을 사용하여 GPU를 컨테이너에 전달할 수 있습니다.

docker pull ghcr.io/sdiehl/docker-mlir-cuda:main
docker run -it ghcr.io/sdiehl/docker-mlir-cuda:main bash

그렇지 않으면 MLIR을 소스에서 다시 컴파일해야 합니다. 단계는 다음과 같습니다:

git clone https://github.com/llvm/llvm-project
mkdir llvm-project/build
cd llvm-project/build
RUN cmake -G Ninja ../llvm \
   -DCUDACXX=/usr/local/cuda/bin/nvcc \
   -DCUDA_PATH=/usr/local/cuda \
   -DCMAKE_CUDA_ARCHITECTURES="75;80;86;90" \ # Add your GPU architecture here
   -DCMAKE_C_COMPILER=clang \
   -DCMAKE_CXX_COMPILER=clang++ \
   -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \
   -DLLVM_ENABLE_PROJECTS=mlir \
   -DLLVM_BUILD_EXAMPLES=ON \
   -DLLVM_TARGETS_TO_BUILD="Native;NVPTX;AMDGPU" \
   -DCMAKE_BUILD_TYPE=Release \
   -DLLVM_ENABLE_ASSERTIONS=ON \
   -DLLVM_CCACHE_BUILD=ON \
   -DMLIR_ENABLE_CUDA_RUNNER=ON \
   -DMLIR_ENABLE_CUDA_CONVERSIONS=ON \
   -DMLIR_ENABLE_NVPTXCOMPILER=ON

RUN cmake --build . -t mlir-opt mlir-translate mlir-transform-opt mlir-runner
RUN cmake --build . -t install

미리 말씀드리자면, 이것은 다소 독특한 CUDA 프로그래밍 속성 과정이 될 것입니다. 결국 우리의 목표는 MLIR에서 GPU를 대상으로 하는 컴파일러를 만드는 것이므로, GPU 기초를 아주 빠르게 다룬 뒤 곧바로 GPU의 저수준 어셈블리 프로그래밍으로 넘어가겠습니다. 컴파일러를 다루는 사람이 아니라면, 대부분은 일상적인 작업에서 PTX를 직접 작성하지 않습니다. CUDA 코드를 처음 작성하는 것이라면, GPU 프로그래밍을 더 자세히 다루는 온라인의 여러 CUDA 프로그래밍 가이드 중 하나를 먼저 읽어 보는 것이 도움이 될 수 있습니다.

CUDA 생태계

CUDA 컴파일 과정에는 GPU 프로그래밍을 위해 이해해야 할 몇 가지 핵심 구성 요소와 중간 표현이 있습니다. 먼저 CUDA 프로그램이 어떻게 컴파일되고 실행되는지 살펴봅시다.

우선 CUDA는 GPU를 프로그래밍할 수 있게 해 주는 C++의 확장입니다. CUDA는 C++ 위에 구축된 언어로, GPU 프로그래밍에 특화된 기능을 추가합니다. 자체 컴파일러인 nvcc를 사용하여 PTX(Parallel Thread Execution, 병렬 스레드 실행)로 컴파일하며, PTX는 NVIDIA GPU를 위한 중간 어셈블리 언어입니다. PTX는 병렬 컴퓨팅을 위한 안정적인 프로그래밍 모델과 명령어 집합을 제공합니다. PTX는 아키텍처 중립적이므로 여러 세대의 NVIDIA 하드웨어를 대상으로 할 수 있습니다.

이어서 컴파일 과정은 PTX에서 CUBIN(CUDA Binary)이라고 하는 장치별 바이너리 코드로 넘어갑니다. CUBIN 파일에는 특정 GPU 아키텍처에서 실행되는 실제 기계어 코드가 들어 있습니다. 이 파일은 ELF 형식이며, 실행 코드뿐만 아니라 심볼, 재배치 정보, 디버그 정보도 포함합니다. CUBIN 파일은 일반적으로 기본적으로 호스트 실행 파일에 포함되지만, nvcc의 -cubin 옵션을 사용하여 별도로 생성할 수도 있습니다.

nvcc는 이 복잡한 컴파일 과정을 관리합니다. 여러 컴파일 단계를 조율하며, 필요에 따라 다양한 출력 형식을 생성할 수 있습니다.

  • 향후 JIT(Just-In-Time) 컴파일을 위한 PTX 코드를 생성할 수 있습니다
  • 아키텍처별 CUBIN 파일을 만들 수 있습니다
  • 여러 PTX 및 CUBIN 파일을 하나의 FATBIN 파일로 결합할 수 있습니다
  • 이러한 GPU 바이너리를 호스트 실행 파일에 포함할 수 있습니다

호스트에서 컴파일된 프로그램을 실행하면 다음과 같은 작업이 수행됩니다.

  • 사용 가능한 GPU 아키텍처를 파악합니다
  • 가능한 경우 FATBIN에서 적절한 CUBIN을 선택합니다
  • 더 새로운 아키텍처의 경우 필요하면 PTX 코드를 JIT 컴파일하는 방식으로 대체합니다
  • GPU에 코드를 로드하고 실행합니다

nvcc 컴파일러를 사용하면 이 모든 작업을 컴파일러가 최종 사용자 대신 처리합니다.

FATBIN 형식은 하나의 실행 파일이 여러 GPU 아키텍처를 지원할 수 있게 해 주므로 특히 중요합니다. FATBIN에는 PTX 코드(JIT 컴파일을 통한 향후 호환성을 위해)와 미리 컴파일된 CUBIN 파일(알려진 아키텍처에서 최적의 성능을 위해)을 모두 담을 수 있습니다. 이제 CUDA를 사용하는 간단한 C++ 코드를 살펴봅시다.

#include <stdio.h>

// CUDA kernel function
__global__ void helloKernel() {
  printf("Hello, CUDA!\n");
}

int main() {
  // Call the CUDA kernel
  helloKernel<<<1, 1>>>();

  // Wait for the kernel to finish
  cudaDeviceReset(0);

  return 0;
}

CUDA의 C++ 확장에는 커널을 **실행(launch)**하기 위한 특별한 문법이 있습니다. helloKernel 함수는 GPU에서 실행하려는 커널입니다. <<<1, 1>>>은 커널을 실행하기 위한 CUDA 문법입니다. 첫 번째 인자는 실행할 블록의 수이고, 두 번째 인자는 블록당 스레드의 수입니다.

kernel<<<blocks, threads>>>(...)

스레드와 블록에 대해서는 나중에 더 자세히 다룹니다. 이제 커널 내부에서는 특수 변수 blockDim, blockIdx, threadIdx에 접근할 수 있습니다. 간단한 예로 배열의 각 원소를 제곱하는 커널을 살펴봅시다. 여기에 제곱 커널이 있습니다. 내부 구조는 나중에 설명하겠으며, 지금은 단순히 배열의 각 원소를 제곱한다고 보면 됩니다.

__global__ void square(int* array, int n) {
    int tid = blockDim.x * blockIdx.x + threadIdx.x;
    if (tid < n)
        array[tid] = array[tid] * array[tid];
}

CUBIN 코드를 얻으려면 다음 명령을 사용할 수 있습니다:

nvcc -o square square.cu

컴파일된 CUBIN 코드를 담은 실행 코드를 실행할 수 있습니다. 이 바이너리는 호스트 로직은 CPU에서, CUBIN 코드는 GPU에서 실행합니다.

./square

nvcc에서 PTX 코드를 얻으려면 다음 옵션을 사용할 수 있습니다:

nvcc -ptx square.cu -o square.ptx

PTX 코드는 다음과 같습니다:

.visible .entry square(int*, int)(
        .param .u64 square(int*, int)_param_0,
        .param .u32 square(int*, int)_param_1
)
{

        ld.param.u64    %rd1, [square(int*, int)_param_0];
        ld.param.u32    %r2, [square(int*, int)_param_1];
        mov.u32         %r3, %ntid.x;
        mov.u32         %r4, %ctaid.x;
        mov.u32         %r5, %tid.x;
        mad.lo.s32      %r1, %r3, %r4, %r5;
        setp.ge.s32     %p1, %r1, %r2;
        @%p1 bra        $L__BB0_2;

        cvta.to.global.u64      %rd2, %rd1;
        mul.wide.s32    %rd3, %r1, 4;
        add.s64         %rd4, %rd2, %rd3;
        ld.global.u32   %r6, [%rd4];
        mul.lo.s32      %r7, %r6, %r6;
        st.global.u32   [%rd4], %r7;

$L__BB0_2:
        ret;

}

여기서 무슨 일이 일어나는지 이해하기 위해 PTX를 명령어 단위로 하나씩 분석해 봅시다. 커널에는 두 개의 매개변수가 있습니다. 첫 번째는 입력 정수 배열을 가리키는 포인터 param_0으로, 입력 정수 배열의 디바이스 주소이고, 두 번째는 입력 정수 배열의 크기인 param_1입니다.

명령어설명
ld.param.u64 %rd1, [square(int*, int)_param_0];커널 매개변수에서 입력 정수 배열의 디바이스 주소를 불러옵니다.
ld.param.u32 %r2, [square(int*, int)_param_1];커널 매개변수에서 입력 정수 배열의 크기를 불러옵니다.
mov.u32 %r3, %ntid.x;x 차원의 블록당 스레드 수를 가져옵니다.
mov.u32 %r4, %ctaid.x;x 차원에서 현재 스레드 블록의 ID를 가져옵니다.
mov.u32 %r5, %tid.x;x 차원에서 블록 내 현재 스레드의 ID를 가져옵니다.
mad.lo.s32 %r1, %r3, %r4, %r5;블록 ID와 스레드 ID를 바탕으로 현재 스레드의 선형 전역 인덱스를 계산합니다.
setp.ge.s32 %p1, %r1, %r2;계산된 인덱스가 배열 크기보다 크거나 같으면 술어 플래그 %p1을 참으로 설정합니다(경계 검사).
@%p1 bra $L__BB0_2;계산된 인덱스가 범위를 벗어나면( %p1이 참이면) 스레드는 반환문으로 분기하므로 사실상 아무 작업도 하지 않습니다.
cvta.to.global.u64 %rd2, %rd1;배열의 일반(generic) 디바이스 주소를 전역 주소로 변환합니다.
mul.wide.s32 %rd3, %r1, 4;현재 스레드가 맡은 원소의 배열 내 바이트 오프셋을 계산합니다(int 크기를 4바이트로 가정).
add.s64 %rd4, %rd2, %rd3;현재 스레드의 인덱스에 해당하는 배열 요소의 절대 전역 메모리 주소를 계산합니다.
ld.global.u32 %r6, [%rd4];계산된 전역 메모리 주소에서 32비트 정수 값을 레지스터 %r6에 로드합니다.
mul.lo.s32 %r7, %r6, %r6;%r6의 정수 값을 제곱하고 결과의 하위 32비트를 %r7에 저장합니다.
st.global.u32 [%rd4], %r7;%r7에 있는 제곱된 정수 값을 같은 전역 메모리 위치에 다시 저장하여 원래 값을 덮어씁니다.
$L__BB0_2:범위를 벗어난 스레드가 분기할 지점을 표시하는 레이블입니다.
ret;현재 스레드의 커널 실행을 마치고 반환합니다.

cuobjdump -sass 명령을 사용하여 CUBIN 코드를 SASS(“Shader ASSembly”, 즉 셰이더 어셈블리의 약자)로 디스어셈블할 수도 있습니다. -arch=sm_90에서 우리 커널은 다음과 같습니다:

square(int*, int):
 LDC R1, c[0x0][0x28]
 S2R R5, SR_CTAID.X
 ULDC UR4, c[0x0][0x0]
 S2R R0, SR_TID.X
 IMAD R5, R5, UR4, R0
 ULDC UR4, c[0x0][0x218]
 ISETP.GE.AND P0, PT, R5, UR4, PT
 @P0 EXIT
 LDC.64 R2, c[0x0][0x210]
 ULDC.64 UR4, c[0x0][0x208]
 IMAD.WIDE R2, R5, 0x4, R2
 LDG.E R0, desc[UR4][R2.64]
 IMAD R5, R0, R0, RZ
 STG.E desc[UR4][R2.64], R5
 EXIT
.L_x_0:
 BRA `(.L_x_0)
 NOP
 NOP
 NOP
 NOP
 NOP
 NOP
 NOP
 NOP
.L_x_1:

GPU 아키텍처

위의 CUDA 커널은 여러 스레드에 걸쳐 실행되면 각 스레드마다 고유한 인덱스를 계산하고, 그 인덱스가 입력 정수 배열의 범위 안에 있으면 해당 인덱스의 정수를 전역 메모리에서 로드하여 제곱한 다음 결과를 같은 메모리 위치에 다시 기록합니다. 결과적으로 배열 요소를 병렬로 제자리(in-place)에서 제곱하는 셈입니다. 계산된 인덱스가 배열 범위를 벗어나는 스레드는 어떠한 메모리 연산도 수행하지 않습니다.

sm_53과 같은 GPU 코드는 항상 sm_ 접두사로 시작합니다. 하나의 가상 GPU 아키텍처는 그에 대응하는 여러 실제 아키텍처를 포함할 수 있습니다. CUBIN은 같은 세대의 모든 GPU에서 실행할 수 있지만, 이전 세대나 이후 세대의 GPU와는 호환되지 않습니다. 예를 들어 sm_52용으로 컴파일된 CUBIN은 sm_50, sm_52, sm_53 GPU와 호환되지만 sm_60 GPU와는 호환되지 않습니다. 또한 sm_52용으로 컴파일된 CUBIN을 sm_53 GPU에서 실행하면, 같은 아키텍처에서 sm_53용으로 특별히 컴파일된 CUBIN을 실행할 때보다 성능이 떨어질 수 있습니다.

# Compile for sm_52
nvcc square.cu -o square --gpu-architecture=compute_52 --gpu-code=sm_52

# Compile for sm_52, sm_53, sm_60
nvcc square.cu -o square --gpu-architecture=compute_52 --gpu-code=sm_52,sm_53,sm_60

가장 일반적인 아키텍처는 다음과 같습니다.

아키텍처SM 버전GPU기능
Maxwellsm_50, sm_52, sm_53GTX 750 Ti, GTX 960, GTX 970, GTX 980동적 병렬성, CUDA 동적 병렬성
Pascalsm_60, sm_61, sm_62P100, GTX 1080, GTX 1080 TiHBM2 메모리, NVLink, 통합 메모리
Voltasm_70, sm_72V100, Titan V1세대 Tensor Core, 독립 스레드 스케줄링
Turingsm_75T4, RTX 2080, RTX 2080 Ti1세대 RT Core, 2세대 Tensor Core
Amperesm_80, sm_86, sm_87A100, A40, RTX 30903세대 Tensor Core, TF32 정밀도, 구조적 희소성
Ada Lovelacesm_89RTX 4090, L40, L40S3세대 RT Core, 4세대 Tensor Core, AV1 인코딩/디코딩
Hoppersm_90, sm_90aH100, H2004세대 Tensor Core, 트랜스포머 엔진, 900 GB/s NVLink
Blackwellsm_100, sm_100a, sm_101, sm_101a, sm_120, sm_120aB100, B2005세대 Tensor Core, 향상된 점유율, 강화된 공유 메모리

CUDA Python 생태계

Python Poetry 프로젝트에서 NVIDIA CUDA 라이브러리를 사용하려면 먼저 NVIDIA 패키지 인덱스를 추가하십시오. Poetry를 사용하는 경우 nvidia-pyindex 패키지를 추가하고 다음 추가 인덱스 URL을 설정하면 됩니다.

poetry add nvidia-pyindex --source https://pypi.ngc.nvidia.com

[extra]
index-url = "https://pypi.ngc.nvidia.com"

이제 다음 라이브러리를 사용할 수 있습니다. cu12 접미사는 CUDA 12.x 버전을 나타내며, 사용하는 CUDA 버전에 따라 이를 조정해야 할 수 있습니다. 설치해야 하는 유일한 핵심(Core) 라이브러리는 cuda-python이며, PyPI에서 받을 수 있습니다.

  • cuda-python - CUDA 런타임 및 기타 Core 기능에 대한 파이썬다운 접근

그다음 Core 라이브러리는 다음과 같습니다:

  • nvidia-cuda-runtime-cu12 - 필수 런타임 기능을 제공하는 Core 런타임 라이브러리입니다
  • nvidia-cublas-cu12 - GPU 가속 기본 선형대수 서브루틴(Basic Linear Algebra Subroutines)을 위한 cuBLAS 라이브러리입니다
  • nvidia-cudnn-cu12 - 딥러닝 기본 연산(스케일드 닷 프로덕트 어텐션, 합성곱, 행렬 곱셈, 소프트맥스, 풀링)을 모은 cuDNN 라이브러리입니다
  • nvidia-cudnn-frontend - cuDNN 프런트엔드 API용 Python 바인딩입니다

계측 및 컴파일러 유틸리티도 있습니다.

  • nvidia-cuda-nvrtc-cu12 - 프로그램 컴파일을 위한 런타임 컴파일 라이브러리(NVRTC)입니다
  • nvidia-nvml-dev-cu12 - GPU 모니터링 및 관리를 위한 관리 라이브러리(NVML)입니다
  • nvidia-nvtx-cu12 - 사용자 정의 프로파일링 및 트레이싱 계측을 위한 도구 확장(NVTX)입니다
  • nvidia-cuda-nvcc-cu12 - C/C++ 코드를 컴파일하기 위한 컴파일러(NVCC)입니다
  • nvidia-nvjitlink-cu12 - 디바이스 코드의 런타임 링킹을 위한 JIT 링커입니다
  • nvidia-cuda-sanitizer-api-cu12 - 애플리케이션의 메모리 오류를 탐지하기 위한 메모리 검사기입니다
  • nvidia-cuda-cupti-cu12 - 성능 분석 및 프로파일링을 위한 프로파일링 도구 인터페이스(CUPTI)입니다

나머지 라이브러리는 과학 계산, 신호 처리 등 특정 작업에 더 특화되어 있습니다. 여기서 자세히 다루지는 않겠지만, 이런 라이브러리도 있습니다.

  • nvidia-cufft-cu12 - 고속 푸리에 변환(Fast Fourier Transform)을 계산하기 위한 cuFFT 라이브러리입니다
  • nvidia-curand-cu12 - 난수 생성을 위한 cuRAND 라이브러리입니다
  • nvidia-cusolver-cu12 - 밀집 및 희소 직접 솔버를 위한 cuSOLVER 라이브러리입니다
  • nvidia-cusparse-cu12 - 희소 행렬 연산을 위한 cuSPARSE 라이브러리입니다
  • nvidia-npp-cu12 - 이미지, 비디오, 신호 처리를 위한 성능 기본 연산(NPP)입니다
  • nvidia-nvjpeg-cu12 - 하드웨어 가속 JPEG 인코딩/디코딩을 위한 nvJPEG 라이브러리입니다
  • nvidia-opencl-cu12 - GPU 컴퓨팅을 위한 OpenCL 구현체입니다

스레드, 블록, 워프

CUDA를 작성할 때 커널은 GPU에서 실행되도록 정의한 함수를 가리킵니다. 커널을 실행하면 단순히 함수 하나를 호출하는 것이 아닙니다. 그 함수를 여러 데이터 요소에 걸쳐 동시에 실행하는 수백, 심지어 수천 개의 병렬 스레드를 생성하는 것입니다.

이 실행 모델은 흔히 단일 명령 다중 스레드(Single-Instruction Multiple-Thread, SIMT) 패러다임이라고 불리며, 하나의 프로그램을 여러 데이터 조각에 걸쳐 실행하는 데 중점을 둡니다. 예를 들어 배열의 각 요소를 제곱하는 커널이 있다고 합시다. 마법처럼 보이는 blockDim.x와 blockIdx.x는 신경 쓰지 마십시오. 이에 대해서는 나중에 자세히 설명합니다.

__global__ void square(float* array, int n) {
    int tid = blockDim.x * blockIdx.x + threadIdx.x;
    if (tid < n)
        array[tid] = array[tid] * array[tid];
}

CUDA 커널에는 두 가지 핵심 속성이 있습니다. 첫째, 커널은 값을 명시적으로 반환할 수 없으므로 모든 결과 데이터를 함수에 전달된 배열에 기록해야 합니다(스칼라 계산의 경우 보통 요소가 하나인 배열을 전달합니다). 둘째, 커널은 호출할 때 스레드 블록의 수와 블록당 스레드 수를 지정하여 스레드 계층 구조를 명시적으로 선언해야 합니다. 커널은 한 번 컴파일되지만, 블록 크기나 그리드 크기를 달리하여 여러 번 실행할 수 있다는 점에 유의해야 합니다.

자, 하나씩 살펴보면, 이 코드는 각 스레드가 배열에서 작업할 때 쓸 고유한 인덱스 tid를 부여하는 것일 뿐입니다. 커널은 이 인덱스를 사용하여 각 스레드가 A와 B의 원소 하나씩을 더하도록 합니다. __global__ 키워드는 이 함수가 GPU에서 실행되는 CUDA 커널이며 CPU 쪽에서 실행을 시작할 수 있다는 것을 알려 줍니다. if (tid < n) 검사가 있는 이유는 배열에 실제로 필요한 것보다 스레드를 몇 개 더 실행하게 될 수 있기 때문입니다(보통 적당한 수로 올림하기 때문입니다). 따라서 인덱스가 범위를 벗어나는 여분의 스레드는 그냥 아무 일도 하지 않습니다.

“호스트 코드“라고 부르는 주 CPU 프로그램에서 커널을 실행하려면 사용할 스레드 수를 정의해야 합니다. 커널은 GPU에서 동작하지만, 커널을 설정하고 실행하는 CPU 코드가 반드시 있어야 합니다. 예를 들어 실행 코드는 다음과 같은 형태일 수 있습니다.

int N = 10000;
int threadsPerBlock = 256;
int numberOfBlocks = (N + threadsPerBlock - 1) / threadsPerBlock;

square<<<numberOfBlocks, threadsPerBlock>>>(d_array, N);

여기 담긴 개념을 풀어 보면, CUDA는 스레드를 일제히 실행되는 32개의 스레드로 이루어진 워프(warp)로 묶고, 이 워프들을 다시 블록으로 묶습니다. 각 블록은 스트리밍 멀티프로세서(Streaming Multiprocessor, SM)에서 동작하며, SM은 레지스터와 공유 메모리 같은 한정된 자원을 가집니다. 블록의 크기는 이러한 자원의 할당에 영향을 주고 동시에 실행될 수 있는 워프의 수를 결정하는데, 이를 **점유율(occupancy)**이라고 합니다.

효율적인 자원 관리는 GPU 성능을 극대화하는 데 매우 중요합니다. GPU 스케줄러는 사용 가능한 SM들에 블록을 분배하는 역할을 합니다. 블록 수가 SM 수를 초과하면 남는 블록은 대기열에 들어가고, 자원이 확보되는 대로 실행되도록 스케줄링됩니다. 이 스케줄링 과정은 SM당 공유 메모리의 양과 레지스터 파일의 크기를 비롯한 여러 요인의 영향을 받으며, 이 요인들이 최종적으로 동시에 실행할 수 있는 블록 수를 결정합니다.

또 다른 주요 문제는 이른바 **워프 분기(warp divergence)**입니다. 이는 if 문과 같은 조건문 때문에 워프 내 스레드들의 실행 경로가 갈라지는 상황을 말합니다. 이상적으로는 워프의 모든 스레드가 같은 명령어를 동시에 실행해야 합니다. 그러나 일부 스레드는 한 경로를 따르고 다른 스레드는 다른 경로를 따르면 실행이 직렬화되어 비효율이 발생합니다. GPU 하드웨어는 마스크 비트를 사용하여 어떤 스레드가 각 경로를 따라야 하는지 관리하며, 이를 통해 모든 스레드가 작업을 올바르게 완료하도록 보장합니다. 이 메커니즘은 정확성을 유지하지만, 분기된 실행 경로 동안 유휴 스레드가 연산 자원을 낭비하므로 성능에 악영향을 줄 수 있습니다. 그 결과, 특히 분기가 심한 경우 GPU의 전체 처리량이 떨어질 수 있습니다. 성능을 높이려면 개발자는 워프 내 조건 분기를 최소화하여 GPU 실행 효율을 최적화하는 것을 목표로 해야 합니다.

메모리 계층 구조

GPU에는 세 가지 유형의 메모리가 있습니다.

  • 레지스터 메모리: GPU 칩에 직접 위치한, 가장 빠른 유형의 GPU 메모리입니다. 개별 스레드만 접근할 수 있으며 수명이 가장 짧아 스레드가 실행되는 동안에만 유지됩니다.

  • 공유 메모리: 이 메모리는 GPU 블록 내의 모든 스레드가 공유합니다. 전역 메모리보다는 빠르지만 레지스터 메모리보다는 느립니다. 공유 메모리는 스레드 간 통신과 데이터 공유를 가능하게 하며, 블록이 실행되는 동안 유지됩니다.

  • 전역 메모리: GPU에서 가장 큰 메모리 풀로, 모든 블록의 모든 스레드가 접근할 수 있습니다. 레지스터 메모리와 공유 메모리보다 느리지만 용량은 가장 큽니다. 전역 메모리는 여러 블록이나 스레드가 공유해야 하는 데이터를 저장하는 데 사용됩니다.

가장 중요한 성능 고려 사항은 CPU(호스트)와 GPU(디바이스) 사이의 메모리 전송 오버헤드입니다. 호스트 메모리와 디바이스 메모리 사이의 데이터 전송은 (온칩 메모리 접근에 비해) 상대적으로 느리므로, 데이터를 가능한 한 오래 GPU에 유지하여 이러한 전송을 최소화하는 것이 좋습니다. 이는 대개 연산마다 데이터를 주고받는 대신, 같은 데이터에 대해 여러 연산을 수행한 뒤에 결과를 호스트로 다시 전송하는 것을 의미합니다. 가능한 한 GPU에 머무르고, 꼭 필요할 때만 CPU로 돌아오는 것이 바람직합니다.

또한 여러 메모리 접근을 하나의 트랜잭션으로 합치는 기능인 **메모리 병합(memory coalescing)**을 이해하는 것도 매우 중요합니다. 메모리 병합은 워프 내의 스레드들이 연속된 메모리 위치에 접근할 때 일어납니다. 병합되지 않은 메모리 접근은 여러 번의 메모리 트랜잭션을 필요로 하여 성능에 큰 영향을 줄 수 있으므로, 최적의 성능을 위해 개발자는 가능한 한 메모리 접근 패턴이 병합되도록 해야 합니다.

메모리 뱅크 충돌(bank conflict) 역시 성능에 영향을 줄 수 있으며, 특히 공유 메모리를 사용할 때 그렇습니다. 뱅크 충돌은 여러 스레드가 같은 메모리 뱅크의 서로 다른 주소에 동시에 접근하려 할 때 발생합니다. 뱅크 충돌을 피하려면 공유 메모리 접근 패턴을 신중하게 고려해야 하며, 필요하다면 자료 구조에 패딩을 추가하여 뱅크 전체에 걸쳐 메모리 접근이 최적으로 분산되도록 해야 합니다. 적절한 패딩과 메모리 접근 패턴을 사용하면 메모리 대역폭 활용도를 극대화하는 데 도움이 됩니다.

커널 호출과 메모리 이동

CUDA 런타임을 관리하기 위해 Python API에는 여러 Core 함수가 있습니다(이 함수들은 C API를 감싼 래퍼일 뿐입니다). 다음은 Python 바인딩의 cuda.cuda 모듈에 있는 Core 함수들입니다.

  • cuInit - CUDA 드라이버 API를 초기화합니다.
  • cuDeviceGet - 첫 번째 연산 디바이스에 대한 핸들을 가져옵니다.
  • cudaSetDevice - 현재 디바이스를 설정합니다.
  • cuCtxCreate - 지정한 디바이스에 연산 컨텍스트를 생성합니다.
  • cuModuleLoadData - PTX 문자열을 디바이스 바이너리로 JIT 컴파일합니다.
  • cuModuleGetFunction - 모듈에서 커널 함수에 대한 핸들을 가져옵니다.
  • cuLaunchKernel - 지정한 디바이스에서 커널을 실행합니다.
  • cuMemcpyHtoD - 호스트 메모리에서 디바이스 메모리로 데이터를 복사합니다.
  • cuMemcpyDtoH - 디바이스 메모리에서 호스트 메모리로 데이터를 복사합니다.
  • cuMemFree - 디바이스의 메모리를 해제합니다.
  • cuMemAlloc - 디바이스에 메모리를 할당합니다.
  • cuCtxDestroy - 연산 컨텍스트를 제거합니다.
  • cuModuleUnload - 디바이스에서 모듈을 언로드합니다.

CUDA Python API는 정수 반환 코드를 사용하는 C API를 감싼 래퍼이기 때문에 오류 처리가 다소 번거롭습니다. 따라서 이 모든 함수에 대해 반환값을 확인하고, 성공 값이 아니면 예외를 발생시켜야 합니다. 이를 쉽게 하기 위해 반환값이 성공 값이 아닐 때 예외를 발생시키는 헬퍼 함수 checkCudaErrors가 있습니다.

CUDA 컨텍스트를 설정하려면 다음과 같은 보일러플레이트 함수를 사용할 수 있습니다:

# gpu_setup.py
# Setup CUDA and create a context.

import cuda.cuda as cu  # type: ignore
import cuda.cudart as cudart  # type: ignore
import cuda.nvrtc as nvrtc  # type: ignore


def _cudaGetErrorEnum(error):
    if isinstance(error, cu.CUresult):
        err, name = cu.cuGetErrorName(error)
        return name if err == cu.CUresult.CUDA_SUCCESS else "<unknown>"
    elif isinstance(error, cudart.cudaError_t):
        return cudart.cudaGetErrorName(error)[1]
    elif isinstance(error, nvrtc.nvrtcResult):
        return nvrtc.nvrtcGetErrorString(error)[1]
    else:
        raise RuntimeError(f"Unknown error type: {error}")


def checkCudaErrors(result):
    if result[0].value:
        raise RuntimeError(
            f"CUDA error code={result[0].value}({_cudaGetErrorEnum(result[0])})"
        )
    if len(result) == 1:
        return None
    elif len(result) == 2:
        return result[1]
    else:
        return result[1:]


def findCudaDevice():
    devID = 0
    checkCudaErrors(cudart.cudaSetDevice(devID))
    return devID


def findCudaDeviceDRV():
    devID = 0
    checkCudaErrors(cu.cuInit(0))
    cuDevice = checkCudaErrors(cu.cuDeviceGet(devID))
    return cuDevice


def setup_cuda(device_id=None):
    print("Initializing CUDA...")
    # Initialize CUDA
    checkCudaErrors(cu.cuInit(0))

    # Get device
    if device_id is None:
        device = findCudaDeviceDRV()
        device_id = 0  # For printing purposes
    else:
        device = checkCudaErrors(cu.cuDeviceGet(device_id))

    # Create context
    context = checkCudaErrors(cu.cuCtxCreate(0, device))

    print(f"CUDA context created on device {device_id}.")
    return context


def cleanup_cuda(context):
    if context:
        print("Destroying CUDA context...")
        checkCudaErrors(cu.cuCtxDestroy(context))
        print("CUDA context destroyed.")

이제 CUDA Python API를 사용하여 호스트(CPU)와 디바이스(GPU) 사이의 메모리를 관리하는 방법을 살펴봅시다. 다음 예제는 GPU 프로그래밍에 필요한 핵심 메모리 연산을 보여 줍니다.

  1. 먼저 cuMemAlloc을 사용하여 GPU에 메모리를 할당합니다. 이렇게 하면 디바이스에 데이터를 저장할 공간이 확보됩니다.
  2. 다음으로 cuMemcpyHtoD(Host to Device)로 호스트에서 디바이스로 데이터를 전송합니다.
  3. 연산을 수행한 뒤에는(연산은 나중에 다룹니다) cuMemcpyDtoH(Device to Host)를 사용하여 결과를 가져옵니다.

아래 예제에서는 간단한 float 배열을 만들어 GPU로 보낸 다음, 다시 복사해 와서 모든 것이 올바르게 동작했는지 확인합니다. 이 패턴은 GPU 컴퓨팅의 기초를 이룹니다.

# gpu_memory.py
# Minimal example demonstrating CUDA context setup and basic memory movement

from gpu_setup import setup_cuda, checkCudaErrors, cleanup_cuda
import numpy as np
import ctypes
import cuda.cuda as cu  # type: ignore

# Setup CUDA context
context = setup_cuda()

try:
    # Allocate memory on GPU
    buffer_size = 5 * ctypes.sizeof(ctypes.c_float)
    device_ptr = checkCudaErrors(cu.cuMemAlloc(buffer_size))

    # Create and initialize host array
    host_array = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)
    print(f"Host array: {host_array}")

    # Copy data from host to device
    checkCudaErrors(cu.cuMemcpyHtoD(device_ptr, host_array.ctypes.data, buffer_size))

    # Create array for results
    result_array = np.zeros(5, dtype=np.float32)

    # Copy data back from device to host
    checkCudaErrors(cu.cuMemcpyDtoH(result_array.ctypes.data, device_ptr, buffer_size))
    print(f"Result array: {result_array}")

    # Free GPU memory
    checkCudaErrors(cu.cuMemFree(device_ptr))

finally:
    # Always clean up the context
    cleanup_cuda(context)

이제 기본적인 연산을 수행하는 간단한 사용자 정의 커널을 실행해 봅시다. 이 예제는 CUDA C++ 코드를 컴파일할 필요 없이 PTX 어셈블리에서 직접 GPU 코드를 실행하는 전체 워크플로를 보여 줍니다. 이 커널은 배열의 첫 번째 원소를 42로 설정할 뿐이지만, GPU 커널 실행의 핵심 구성 요소를 잘 보여 줍니다. CUDA 컨텍스트 초기화, GPU 메모리 할당, PTX 어셈블리 언어로 커널 정의, PTX 코드를 CUDA 모듈로 로드, 실행 차원 설정(단일 블록 안의 단일 스레드), 적절한 인자로 커널 실행, 완료를 보장하기 위한 동기화, 그리고 마지막으로 검증을 위해 결과를 호스트로 다시 전송하는 과정이 그것입니다. PTX 어셈블리 코드에는 첫 번째 스레드만 값을 대입하도록 스레드 ID와 블록 ID를 확인하는 부분이 들어 있습니다.

# gpu_memory.py
# Minimal example demonstraitng lauching a kernel

from gpu_setup import setup_cuda, checkCudaErrors, cleanup_cuda
import numpy as np
import cuda.cuda as cu  # type: ignore

# This is the trivial kernel that sets the first element
# of the array to 42 using a single thread

# __global__ void set_value_kernel(int *data) {
#     if (threadIdx.x == 0 && blockIdx.x == 0) {
#         *data = 42;
#     }
# }

ptx_kernel = """"
.visible .entry kernel(int*)( .param .u64 kernel(int*)_param_0) {
    ld.param.u64    %rd1, [kernel(int*)_param_0];
    mov.u32         %r1, %tid.x;
    mov.u32         %r2, %ctaid.x;
    or.b32          %r3, %r1, %r2;
    setp.ne.s32     %p1, %r3, 0;
    @%p1 bra        $L__BB0_2;
    cvta.to.global.u64      %rd2, %rd1;
    mov.u32         %r4, 42;
    st.global.u32   [%rd2], %r4;
$L__BB0_2:
    ret;
}
"""

# Setup CUDA context
context = setup_cuda()

try:
    # Allocate a single integer on the device
    data_size = np.dtype(np.int32).itemsize
    d_data = checkCudaErrors(cu.cuMemAlloc(data_size))

    # Set up grid and block dimensions - just one thread
    grid_dims = (1, 1, 1)
    block_dims = (1, 1, 1)

    # Prepare arguments for the kernel
    args = [d_data]
    arg_types = [None]  # None for pointer types

    # Load the module
    module = checkCudaErrors(cu.cuModuleLoadData(ptx_kernel.encode("utf-8")))

    # Get kernel function
    kernel_func = checkCudaErrors(
        cu.cuModuleGetFunction(module, "kernel".encode("utf-8"))
    )

    # Prepare kernel arguments
    kernel_args = (tuple(args), tuple(arg_types))

    # Launch kernel
    checkCudaErrors(
        cu.cuLaunchKernel(
            kernel_func,
            grid_dims[0],
            grid_dims[1],
            grid_dims[2],
            block_dims[0],
            block_dims[1],
            block_dims[2],
            0,  # shared memory bytes
            0,  # stream
            kernel_args,  # kernel args
            0,  # extra
        )
    )

    # Synchronize
    checkCudaErrors(cu.cuCtxSynchronize())

    # Copy result back to host
    host_data = np.zeros(1, dtype=np.int32)
    checkCudaErrors(cu.cuMemcpyDtoH(host_data.ctypes.data, d_data, host_data.nbytes))

    # Verify result
    print(f"Result: {host_data[0]}")
    assert host_data[0] == 42, f"Expected 42 but got {host_data[0]}"
    print("Success! Kernel executed correctly.")

    # Free device memory
    checkCudaErrors(cu.cuMemFree(d_data))

    # Unload module
    checkCudaErrors(cu.cuModuleUnload(module))

finally:
    cleanup_cuda(context)

스레드 블록과 그리드

앞에서 스레드가 보통 32개씩 워프로 묶이고, 이 워프들이 다시 블록으로 묶이는 방식을 살펴보았습니다. 이 블록들은 실행을 위해 SM에 할당됩니다. 이제 시야를 넓혀 봅시다. 스레드 블록은 공유 메모리와 동기화 메커니즘을 사용하여 함께 작업하며, 모두 같은 SM에서 실행되는 스레드들의 집합으로 이해할 수 있습니다.

GPU를 살펴보려면 cuda-python 라이브러리로 GPU 드라이버 API에 질의하여 사용 가능한 병렬성 속성에 관한 정보를 얻을 수 있습니다.

from cuda.cuda import CUdevice_attribute, cuDeviceGetAttribute, cuDeviceGetName, cuInit

(err,) = cuInit(0)

err, DEVICE_NAME = cuDeviceGetName(128, 0)
DEVICE_NAME = DEVICE_NAME.decode("ascii").replace("\x00", "")

err, MAX_THREADS_PER_BLOCK = cuDeviceGetAttribute(
    CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK, 0
)
err, MAX_BLOCK_DIM_X = cuDeviceGetAttribute(
    CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_X, 0
)
err, MAX_GRID_DIM_X = cuDeviceGetAttribute(
    CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X, 0
)
err, SMs = cuDeviceGetAttribute(
    CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, 0
)

print(f"GPU Device: {DEVICE_NAME}")
print(f"Number of multiprocessors: {SMs}")
print(f"Maximum number of threads per block: {MAX_THREADS_PER_BLOCK:10}")
print(f"Maximum number of blocks per grid: {MAX_BLOCK_DIM_X:10}")
print(f"Maximum number of threads per grid: {MAX_GRID_DIM_X:10}")

예를 들어 A100 GPU에서는 다음과 같은 결과를 얻습니다.

Device Name: A100
Maximum number of multiprocessors: 108
Maximum number of threads per multiprocessor: 1024
Maximum number of threads per block: 1024
Maximum number of blocks per grid: 2147483647

블록을 사용하는 데에는 현실적인 이유가 있습니다. 하나의 블록에서 실행할 수 있는 스레드 수에는 하드웨어 제한이 있으며, 최신 GPU에서는 보통 1024개로 제한됩니다. 연산에 이 제한보다 많은 스레드가 필요하다면 스레드를 여러 블록으로 나누어야 합니다. 블록은 데이터의 특정 구간을 함께 처리하는 스레드들의 팀이라고 생각하십시오. 전체 그리드는 하나의 커널이 시작한 모든 블록을 나타내며, 전체 작업을 완수하기 위해 함께 일하는 여러 팀을 상징합니다.

예를 들어 원소가 2048개인 두 배열을 더하려면 각각 1024개의 스레드로 이루어진 블록 두 개를 사용할 수 있습니다. 이때 블록 0은 인덱스 0-1023을, 블록 1은 인덱스 1024-2047을 처리합니다. 일반적으로 원소가 N개이고 블록 하나에 최대 B개의 스레드를 담을 수 있다면, 모든 원소를 처리하기 위해 ceil(N/B)개의 블록을 실행합니다.

블록은 확장과 스케줄링에도 매우 유용합니다. 이렇게 생각해 보십시오. GPU에는 SM이 여러 개(예를 들어 20개) 있고, 각 SM은 사용 가능한 자원에 따라 여러 블록을 동시에 처리할 수 있습니다. 그래서 블록 100개를 던져 주더라도 동시에 실행될 수 있는 것은 20개뿐입니다(SM마다 하나씩). GPU는 그 20개 블록을 병렬로 시작하고, 하나가 끝날 때마다 대기 중인 다른 블록을 즉시 가져와 비게 된 SM에서 실행합니다.

프로그래머의 관점에서 보면 100개의 블록 모두가 함께 답을 구하고 있으며, 마치 100 * blockSize개의 스레드가 모두 실행되고 있는 것과 같습니다. 하지만 내부적으로는 GPU가 이 블록들을 하드웨어 전반에 영리하게 분배하고 있습니다. 좋은 점은 GPU가 한 번에 물리적으로 처리할 수 있는 것보다 더 많은 스레드를 실행하는 것에 대해 걱정할 필요가 없다는 것입니다. 런타임이 필요에 따라 이 블록들을 시분할하여 실행할 뿐입니다. 또한 블록은 여러 GPU에 작업을 스프레드하거나 한 번에 수행되는 병렬 작업의 양을 줄이는 자연스러운 방법을 제공합니다(이는 공유 메모리나 레지스터 같은 자원 한계에 부딪힐 때 필수적일 수 있습니다).

같은 블록 내의 스레드들은 다른 스레드들과 구별되는 강력한 협업 능력을 가지고 있습니다. 이 스레드들은 빠른 온칩 메모리에 대한 접근을 공유하고 배리어 연산을 사용하여 동기화할 수 있으므로, 스레드 간 통신이 필요한 작업을 효율적으로 조율할 수 있습니다. 그러나 이러한 협업은 엄격하게 같은 블록 내의 스레드로 제한됩니다.-서로 다른 블록에 걸친 스레드들은 통신을 위해 더 느린 전역 메모리에 의존해야 하며, 서로 직접 동기화할 수 없습니다. 이러한 결정은 중요한 목적을 가지고 있습니다. 바로 SM 전반에 걸쳐 블록을 유연하게 스케줄링할 수 있게 하고, CUDA 프로그램이 SM을 몇 개 포함하든 관계없이 다양한 GPU 아키텍처에서 매끄럽게 확장될 수 있게 하는 것입니다.

일반적으로 GPU의 워프 크기인 32개 스레드에 맞추고 하드웨어 특성에 최적화하기 위해 2의 거듭제곱(128, 256, 512 등)인 스레드 블록 크기를 선택합니다. 이러한 크기 설정은 여러 성능 요인에 영향을 미칩니다. 자원 활용도를 극대화하고, 메모리 접근 패턴을 최적화하며, 효율적인 스레드 스케줄링을 통해 메모리 지연 시간을 효과적으로 숨기는 GPU의 능력을 향상시킵니다.

# N is size of the array to process
N = 10000

# We calculate the number of blocks needed based on the total elements and threads per block
 # Typically a power of 2 like 128, 256, or 512
threads_per_block = 256

# Ceiling division to ensure all elements are processed
num_blocks = (N + threads_per_block - 1) // threads_per_block

공유 메모리와 동기화

CUDA의 공유 메모리는 성능 향상에 사용할 수 있는 가장 강력한 최적화 도구 중 하나이며, 메모리 접근 지연 시간을 극적으로 줄여 주는 프로그래밍 가능한 캐시로 작동합니다. 각 SM 내부의 온칩에 위치한 공유 메모리는 전역 메모리보다 최대 100배 높은 대역폭을 제공하며, 지연 시간은 약 100배 낮습니다. 이러한 성능상의 이점은 연산 코어와 물리적으로 가깝다는 데서 비롯됩니다. 반면 전역 메모리는 오프칩에 위치하여 데이터가 훨씬 느린 PCIe 버스나 메모리 컨트롤러를 거쳐 이동해야 합니다. 블록 내의 스레드들이 같은 데이터를 다루는 계산에서 협력할 때, 공유 메모리는 중복된 전역 메모리 로드를 제거하여 전역 메모리 접근 비용을 블록 내 모든 스레드에 효과적으로 분산시킵니다.

공유 메모리의 동작 방식은 CUDA의 스레드 계층 구조 설계를 반영합니다. 커널이 실행되면 각 스레드 블록에는 전용 공유 메모리 영역이 할당되며, 이 영역은 블록이 존재하는 동안 유지되고 해당 블록 내의 스레드에게만 보입니다. 이 메모리는 서로 다른 스레드가 동시에 접근할 수 있는 뱅크(최신 GPU에서는 일반적으로 32개의 뱅크)로 구성되어 있어 높은 대역폭의 병렬 접근이 가능합니다. 그러나 여러 스레드가 같은 뱅크 내의 서로 다른 주소에 접근하려고 하면(즉, 뱅크 충돌) 이러한 접근이 직렬화되어 성능이 저하됩니다.

__syncthreads() 함수는 공유 메모리를 다루는 데 핵심(Core)이 되는 함수입니다. 이 함수는 블록 내의 모든 스레드가 코드의 해당 지점에 도달할 때까지 기다린 후에 실행을 계속하도록 강제합니다. 이를 통해 일부 스레드가 다른 스레드들이 아직 쓰기를 마치지 않은 데이터를 읽으려고 하는 골치 아픈 경쟁 상태를 방지할 수 있으며, 이러한 경쟁 상태는 미묘하고 디버깅하기 어려운 오류의 흔한 원인입니다.

이 예제에서는 블록 내의 모든 스레드에게 보이는 __shared__ 배열을 선언합니다. 그런 다음 전역 메모리에서 공유 메모리로 데이터를 로드하고, 모든 스레드가 데이터를 로드했는지 확인하기 위해 동기화한 후, 공유 메모리에서 데이터를 처리하고, 마지막으로 결과를 전역 메모리에 다시 씁니다.

#include <stdio.h>

__global__ void sharedMemoryExample(float* input, float* output, int n) {
    // Declare shared memory array - visible to all threads in the block
    __shared__ float sharedData[256];

    // Calculate global thread ID
    int tid = blockIdx.x * blockDim.x + threadIdx.x;

    // Load data from global memory to shared memory
    if (tid < n) {
        sharedData[threadIdx.x] = input[tid];
    }

    // Synchronize to make sure all threads have loaded their data
    __syncthreads();

    // Process data in shared memory (simple example: add 1 to each element)
    if (tid < n) {
        sharedData[threadIdx.x] += 1.0f;
    }

    // Synchronize again before writing results back
    __syncthreads();

    // Write the processed data back to global memory
    if (tid < n) {
        output[tid] = sharedData[threadIdx.x];
    }
}

공유 메모리를 사용하면 그렇지 않을 경우 실용적이지 않았을 알고리즘 기법을 사용할 수 있습니다. 타일링 알고리즘에서는 데이터를 공유 메모리에 들어갈 수 있는 작은 덩어리(타일이라고 부릅니다)로 나누어 처리하므로, 전역 메모리 트래픽을 극적으로 줄여 행렬 곱셈 같은 연산이 이론적 최대 성능에 가까운 성능을 낼 수 있습니다. 스텐실 연산에서는 공유 메모리를 통해 스레드들이 겹치는 영역에서 협업할 수 있어 중복된 전역 메모리 로드가 제거됩니다. 합계나 최댓값 계산 같은 리덕션 연산에서는 공유 메모리가 블록 내에서 스레드들이 결과를 점진적으로 결합하는 효율적인 병렬 리덕션 패턴을 가능하게 합니다. 핵심 설계 패턴은 전역 메모리에서 공유 메모리로 데이터를 로드하고, 공유 메모리에서만 연산을 수행한 다음, 결과를 전역 메모리에 다시 쓰는 것입니다. 이를 통해 그렇지 않았다면 메모리 바운드였을 연산이 컴퓨트 바운드 연산으로 바뀌며, GPU의 주요 병목인 메모리 접근 지연 시간을 최소화하면서 GPU의 막대한 연산 처리량을 효과적으로 활용할 수 있습니다.

MLIR과 연결하기

자, 이제 우리의 목표는 고수준 텐서 연산을 동적으로(즉, 런타임에) 지정하고, MLIR을 사용하여 이를 적절한 GPU 코드로 로어링할 수 있게 하는 것입니다. MLIR의 점진적 로어링 기능을 활용하여, 잘 정의된 일련의 단계를 거쳐 고수준 텐서 연산을 최적화된 GPU 코드로 변환하겠습니다.

  1. 고수준 표현: 먼저 linalg 다이얼렉트로 표현된 추상적인 텐서 계산(행렬 곱셈, 합성곱, 원소별 연산)에서 시작합니다. 이러한 연산은 구현 세부 사항 없이 순수한 수학적 의도만을 표현하므로, 코드를 도메인에 가깝게 유지할 수 있습니다.

  2. 아핀 변환: -convert-linalg-to-loops나 -convert-linalg-to-affine-loops 같은 패스를 사용하여 이러한 고수준 연산을 affine 다이얼렉트로 변환합니다. 이 다이얼렉트는 루프와 메모리 접근을 정밀한 수학적 관계로 표현합니다. 이를 통해 메모리 지역성과 계산 효율을 높이는 타일링, 융합, 루프 교환 같은 핵심 최적화의 기회가 드러납니다.

  3. GPU 매핑: 다음으로 -convert-affine-for-to-gpu와 -convert-parallel-loops-to-gpu 같은 패스를 통해 이러한 루프 구조를 GPU의 실행 모델에 매핑합니다. 이렇게 하면 코드가 gpu 다이얼렉트로 변환되며, 이 다이얼렉트는 하드웨어에 독립적인 상태를 유지하면서도 스레드 블록과 스레드 같은 개념을 명시적으로 표현합니다.

  4. 하드웨어 특화 로어링: 마지막으로 nvvm 다이얼렉트를 통해 NVIDIA 특화 기능으로 넘어가고, 최종적으로 GPU 인트린식을 포함한 LLVM IR로 전환합니다. 마지막 단계에서는 이를 NVPTX 어셈블리로, 그리고 다시 GPU에서 직접 실행되는 CUBIN 기계어 코드로 변환합니다.

이 파이프라인 전체에 걸쳐 각 로어링 단계는 프로그램 의미론을 보존하면서 대상 하드웨어에 점점 더 가까워집니다. 수학적 표현에서 명시적인 루프 중첩으로, 이어서 GPU 실행 구조로, 마지막으로 하드웨어 특화 코드로 체계적으로 변환합니다. 이 접근 방식을 사용하면 계산을 고수준에서 표현하면서도, 하드웨어의 성능을 최대한 활용하는 고도로 최적화된 GPU 코드를 생성할 수 있습니다.

gpu 다이얼렉트

MLIR의 GPU 다이얼렉트는 CUDA나 OpenCL과 유사한 프로그래밍 모델을 따르는 GPU 커널을 실행하기 위한 중간 수준의 추상화를 제공합니다. 이 다이얼렉트는 GPU 커널 실행에 필요한 장치별, 드라이버별 연산을 추상화하여, MLIR에서 GPU 실행으로 이어지는 간소화된 경로를 제공합니다.

이 다이얼렉트는 gpu를 정규 접두사로 사용하며, 일반적인 GPU 프리미티브를 감싸는 연산을 제공합니다. 이상적인 미래에는 이를 여러 GPU 백엔드를 대상으로 사용할 수도 있겠지만, 지금은 PTX로 로어링하는 것을 목표로 Nvidia 특화 연산에만 집중하겠습니다.

장치 연산

  • gpu.launch - 지정된 그리드 및 블록 차원으로 GPU 커널을 실행합니다. 블록 및 스레드 구성을 받으며 커널 본문을 포함합니다.
  • gpu.launch_func - 지정된 그리드 및 블록 차원으로 GPU 함수를 실행합니다. 함수 참조와 블록 및 스레드 구성을 받습니다.
  • gpu.barrier - 스레드 블록 내의 모든 스레드를 동기화하여, 실행을 계속하기 전에 모든 스레드가 이 지점에 도달하도록 보장합니다.
  • gpu.binary - 로드하여 실행할 수 있는 컴파일된 GPU 코드(예: PTX 또는 CUBIN 형식)를 담은 바이너리 블롭을 나타냅니다.
  • gpu.printf - C의 printf 함수와 유사하게, GPU 커널 내부에서 형식화된 출력을 인쇄합니다. GPU 코드를 디버깅할 때 유용합니다.
  • gpu.return - GPU 함수/커널에서 호스트로 값을 반환합니다.
  • gpu.terminator - GPU 실행 영역의 끝을 표시하며, gpu.launch 본문의 마지막 연산으로 반드시 있어야 합니다.
  • gpu.wait - 호스트 실행을 계속하기 전에 비동기 GPU 연산이 완료될 때까지 기다립니다.
  • gpu.yield - scf.yield와 유사하게, GPU 코드의 구조적 제어 흐름 영역에서 값을 내보냅니다.

스레드 및 블록 연산

  • gpu.block_id - 지정된 차원(x, y 또는 z)을 따라 그리드 내 현재 블록의 ID를 반환합니다
  • gpu.block_dim - 지정된 차원(x, y 또는 z)을 따라 그리드 내 블록의 차원을 반환합니다
  • gpu.block_size - 지정된 차원(x, y 또는 z)을 따라 블록 내 스레드의 총 개수를 반환합니다
  • gpu.grid_id - 지정된 차원(x, y 또는 z)을 따라 현재 실행 중인 그리드의 ID를 반환합니다
  • gpu.grid_dim - 지정된 차원(x, y 또는 z)을 따라 그리드의 차원을 반환합니다
  • gpu.grid_size - 지정된 차원(x, y 또는 z)을 따라 그리드 내 블록의 총 개수를 반환합니다
  • gpu.thread_id - 지정된 차원(x, y 또는 z)을 따라 해당 블록 내 현재 스레드의 ID를 반환합니다

메모리 연산

  • gpu.alloc - GPU 디바이스에 메모리를 할당합니다. 형태(shape)와 원소 타입을 받아 할당된 메모리를 가리키는 memref를 반환합니다.
  • gpu.memcpy - 호스트와 디바이스 사이, 또는 디바이스 내의 서로 다른 위치 사이에서 메모리를 복사합니다. 호스트에서 디바이스로, 디바이스에서 호스트로, 디바이스에서 디바이스로의 전송을 모두 처리할 수 있습니다.
  • gpu.dealloc - 이전에 gpu.alloc을 사용하여 GPU 디바이스에 할당한 메모리를 해제합니다.
  • gpu.host_register - memref를 호스트 로컬 메모리로 등록하여 호스트가 직접 접근할 수 있게 합니다.
  • gpu.host_unregister - 디바이스에서 접근할 수 있도록 등록된 memref의 등록을 해제합니다.

이 다이얼렉트는 특정 GPU 백엔드로 로워링할 수 있는 타깃 독립적인 표현으로 설계되었습니다. 이 다이얼렉트는 커널 호출을 위한 추상화를 제공하며, 향후에는 GPU용 LLVM IR 내장 함수(intrinsic)와 같은 하위 수준에는 없는 디바이스 관리 기능도 포함할 수 있습니다.

이 다이얼렉트는 GPU 코드가 gpu.module 연산 안에 구성되고, 커널은 gpu.func 연산으로 표현될 것을 전제로 합니다. 이러한 구조는 호스트 코드와 디바이스 코드를 명확하게 분리하여 컴파일 및 실행 파이프라인을 다루기 쉽게 해 줍니다. 디바이스 라이브러리 호출과 같은 비커널 함수는 func.func이나 GPU가 아닌 다른 다이얼렉트의 연산을 사용하여 정의할 수 있으므로, GPU 프로그램을 구성하는 방식에 어느 정도 유연성을 제공합니다.

module {

  // GPU module
  gpu.module @gpu_module {
    // GPU function
    gpu.func @kernel() {
        %0 = arith.addi %arg0, %arg0 : i32
        %1 = arith.addi %arg0, %arg0 : i32
        %2 = arith.addi %0, %1 : i32
        gpu.return
    }
  }

  // Host function
  func.func @main() {
      gpu.launch_func @gpu_module::@kernel
        blocks(%0, %1, %2) in (%3 = %c1, %4 = %c1, %5 = %c1)
        threads(%6, %7, %8) in (%9 = %c2, %10 = %c1, %11 = %c1)
        args()
  }
}

다른 방식으로는 gpu.launch 연산을 사용하는 것이 있으며, 이 연산은 런치 연산에 인자로 전달되는 영역(region)에 커널 함수를 내장합니다. gpu.terminator 연산은 영역의 끝을 표시하는 데 사용됩니다.

func.func @main() {
    %c2 = arith.constant 2 : index
    %c1 = arith.constant 1 : index
    gpu.launch
        blocks(%0, %1, %2) in (%3 = %c1, %4 = %c1, %5 = %c1)
        threads(%6, %7, %8) in (%9 = %c2, %10 = %c1, %11 = %c1) {
          gpu.printf "Hello from %d\n" %6 : index
          gpu.terminator
        }
    return
}

nvgpu 다이얼렉트

nvgpu 다이얼렉트는 MLIR 생태계에서 중개자 역할을 하며, gpu나 vector와 같은 상위 수준의 타깃 독립적인 다이얼렉트를 NVIDIA GPU 전용으로 설계된 하위 수준의 NVVM 다이얼렉트와 연결합니다. 이 다이얼렉트는 memref와 tensor 다이얼렉트 같은 MLIR의 고수준 추상화를 계속 사용하면서 PTX 전용 연산을 표현함으로써, 양쪽의 장점을 모두 살리는 다리 역할을 합니다.

이 다이얼렉트를 사용하면 복잡한 NVVM 내장 함수를 직접 다룰 필요 없이 NVGPU 다이얼렉트를 통해 NVIDIA 전용 하드웨어 기능에 접근할 수 있습니다. 이 중간 계층은 상위 수준 MLIR 표현의 명료성과 표현력을 유지하면서도 효율적인 GPU 코드 생성을 크게 단순화합니다.

이 다이얼렉트는 주로 전역 메모리와 공유 메모리 간의 효율적인 텐서 전송을 위한 Tensor Memory Accelerator(TMA)와 같은 고급 NVIDIA GPU 기능을 노출하는 데 중점을 둡니다. 또한 연산과 메모리 전송이 겹쳐 수행되도록 하여 지연 시간을 효과적으로 숨기는 비동기 메모리 연산 메커니즘도 제공합니다. 메모리 접근 패턴은 뱅크 충돌을 줄이는 스위즐링(swizzling) 기법을 통해 최적화되며, 메모리 배리어(mbarrier)는 스레드 간 연산을 조율하기 위한 정교한 동기화 도구를 제공합니다.

워프 수준 프로그래밍은 NVGPU 다이얼렉트에서 특별한 어텐션을 받으며, 워프 행렬 곱셈-누산 연산과 동기화 프리미티브 같은 하드웨어별 기능을 노출하는 세밀한 제어 기능을 제공합니다. 이러한 기능은 Hopper나 Blackwell 칩과 같은 최신 NVIDIA GPU의 특수 하드웨어 유닛을 활용하여 고성능 행렬 연산과 기타 연산 집약적인 워크로드를 구현하는 데 필수적입니다.

이 다이얼렉트는 컴파일 파이프라인에서 보다 범용적인 GPU 다이얼렉트와 하드웨어에 특화된 NVVM 다이얼렉트 사이에 위치하며, 이 위치는 전략적으로 중요합니다. 이러한 배치 덕분에 상위 수준의 개념을 최적화된 하드웨어 특화 연산으로 변환하면서도 개발자를 불필요한 복잡성으로부터 보호할 수 있습니다. 최종 코드를 생성할 때가 되면 -convert-nvgpu-to-nvvm 패스를 사용하여 NVGPU 다이얼렉트를 NVVM으로 로워링할 수 있으며, 이 패스는 NVGPU 연산을 대응하는 NVVM 다이얼렉트 인트린직으로 변환합니다.

이러한 기능은 이 시리즈에서 다루는 작은 백엔드의 범위를 벗어납니다. Part 12에서는 구조화된 텐서 계산과 타깃별 스케줄링 사이의 구분을 다시 다루며, 함께 제공되는 컴파일러는 CPU 실행 경로를 사용합니다.

nvvm 다이얼렉트

NVVM 다이얼렉트는 NVIDIA GPU용 LLVM IR을 표현하는 타깃 특화 다이얼렉트입니다. 이 다이얼렉트에는 스레드 및 블록 인덱싱, 동기화, 메모리 연산과 같은 GPU 특화 구성 요소를 위한 연산이 포함되어 있습니다.

NVVM IR은 LLVM IR을 기반으로 하는 컴파일러 IR(중간 표현)입니다. NVVM IR은 GPU 연산 커널(예: CUDA 커널)을 표현하도록 설계되었습니다. CUDA C 컴파일러 프런트엔드와 같은 고수준 언어 프런트엔드는 NVVM IR을 생성할 수 있습니다. (LLVM을 기반으로 하는) NVVM 컴파일러는 NVVM IR로부터 PTX 코드를 생성합니다.

CUDA 내장 함수NVVM 인트린직MLIR 연산
threadId.{x,y,z}@llvm.nvvm.read.ptx.sreg.tid.{x,y,z}gpu.thread_id {x,y,z}
blockIdx.{x,y,z}@llvm.nvvm.read.ptx.sreg.ctaid.{x,y,z}gpu.block_id {x,y,z}
blockDim.{x,y,z}@llvm.nvvm.read.ptx.sreg.ntid.{x,y,z}gpu.block_dim {x,y,z}
gridDim.{x,y,z}@llvm.nvvm.read.ptx.sreg.nctaid.{x,y,z}gpu.grid_dim {x,y,z}
__syncthreads()@llvm.nvvm.barrier0()gpu.barrier

MLIR Python 바인딩

MLIR Python 바인딩은 MLIR C++ 내부 구현에 대한 인터페이스를 제공하므로, 개발자는 명령줄 도구에 의존하지 않고도 MLIR의 중간 표현을 프로그래밍 방식으로 조작할 수 있습니다. 이 바인딩은 MLIR의 Core 기능을 Python으로 노출하므로, 사용자는 Python에서 직접 MLIR 텍스트를 파싱하고, MLIR AST를 구성 및 수정하고, 변환 패스를 적용하며, 나아가 MLIR 코드를 JIT 컴파일하여 실행할 수도 있습니다.

Python 바인딩을 사용하면 MLIR을 다루는 기존 명령줄 방식에 따르던 여러 불편함이 대부분 사라집니다. 개발자는 여러 명령줄 호출을 이어 붙이는 대신 Python 안에서 처음부터 끝까지 컴파일 파이프라인을 만들 수 있습니다. 그래서 개발 과정이 더 대화형이 되고 디버깅도 더 쉬워집니다. 하지만 이 바인딩은 LLVM과 MLIR에 대한 의존성이 복잡하기 때문에 소스에서 빌드하기가 어려울 수 있습니다. 도입을 쉽게 하기 위해 바인딩과 필요한 모든 의존성을 함께 패키징한 사전 빌드 휠이 제공됩니다. 덕분에 pip 명령 하나로 간단하게 설치할 수 있습니다.

바인딩을 설치하려면 pip를 사용하면 됩니다.

pip install mlir_python_bindings -f https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest

Poetry를 사용한다면 pyproject.toml에 다음 내용을 추가하십시오.

[tool.poetry.dependencies]
python = "^3.10"
mlir-python-bindings = { version = "*", source = "mlir-wheels"}

[[tool.poetry.source]]
name = "mlir-wheels"
url = "https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest"
priority = "supplemental"

이제 Python 바인딩을 사용하여 간단한 벡터 덧셈 커널을 PTX로 로어링해 봅시다.

module {
  func.func @square(%input: tensor<10x10xf32>, %output: tensor<10x10xf32>) -> tensor<10x10xf32> {
    %x0 = linalg.square ins(%input : tensor<10x10xf32>) outs(%output : tensor<10x10xf32>) -> tensor<10x10xf32>
    return %x0 : tensor<10x10xf32>
  }
}

여기서 패스를 적용합니다.

from mlir.ir import Context, Module
from mlir.passmanager import PassManager

mlir_module_str = open("vecadd.mlir").read()

with Context():
    # Parse the input module
    module = Module.parse(mlir_module_str)

pm = PassManager()
pm.enable_ir_printing(print_after_change=True)
pm.add("canonicalize")
pm.add(
    "one-shot-bufferize{ bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map }"
)
pm.add("canonicalize")
pm.add("convert-linalg-to-affine-loops")
pm.add("func.func(affine-loop-invariant-code-motion)")
pm.add("func.func(convert-affine-for-to-gpu)")
pm.add("gpu-kernel-outlining")
pm.add("lower-affine")
pm.add("gpu-decompose-memrefs")
pm.add("expand-strided-metadata")
pm.add("normalize-memrefs")
pm.add(
    "gpu.module(convert-gpu-to-nvvm{index-bitwidth=0 use-bare-ptr-memref-call-conv })"
)
pm.add(f"nvvm-attach-target{ {chip={chip_type} features=+ptx80 O=3} }")
pm.add("convert-nvvm-to-llvm")
pm.add("reconcile-unrealized-casts")
pm.add("gpu-to-llvm { use-bare-pointers-for-host use-bare-pointers-for-kernels }")
pm.run(module.operation)

print(module)

여기서 일어나는 일이 많으므로 명령의 각 플래그를 하나씩 살펴봅시다.

  • one-shot-bufferize{ bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map }: 단일 패스에서 텐서 연산을 버퍼 연산으로 변환합니다. 함수 경계는 항등 레이아웃 매핑으로 처리합니다.
  • convert-linalg-to-affine-loops: 고수준 linalg 연산을 텐서 원소를 명시적으로 순회하는 affine 루프 중첩으로 변환합니다.
  • func.func(affine-loop-invariant-code-motion): affine 루프에 루프 불변 코드 이동을 수행합니다. 가능하면 계산을 루프 밖으로 옮겨 중복 계산을 줄입니다.
  • func.func(convert-affine-for-to-gpu): affine 루프를 GPU 실행 모델에 매핑하여 반복을 GPU 스레드와 블록에 분배합니다.
  • gpu-kernel-outlining: GPU 커널 영역을 호스트 코드에서 실행할 수 있는 별도의 GPU 함수로 추출합니다.
  • lower-affine: affine 다이얼렉트 연산을 표준 제어 흐름 연산과 산술 연산으로 변환합니다.
  • gpu-decompose-memrefs: 복잡한 memref 타입을 GPU 백엔드가 처리할 수 있는 더 단순한 타입으로 분해합니다.
  • expand-strided-metadata: 스트라이드 메모리 접근의 메타데이터를 명시적인 계산으로 전개합니다.
  • normalize-memrefs: 메모리 참조를 GPU 백엔드가 기대하는 형태로 정규화합니다.
  • gpu.module(convert-gpu-to-nvvm{index-bitwidth=0 use-bare-ptr-memref-call-conv }): GPU 다이얼렉트 연산을 NVVM 다이얼렉트(NVIDIA의 LLVM 기반 IR)로 변환합니다. memref에는 베어 포인터 호출 규약을 사용합니다.
  • nvvm-attach-target{chip={chip_type} features=+ptx80 O=3}: NVVM 모듈에 타깃별 정보를 붙입니다. 이 정보에는 GPU 아키텍처, PTX 버전 8.0 기능, 최적화 수준 3이 포함됩니다.
  • convert-nvvm-to-llvm: 이후 처리를 위해 NVVM 다이얼렉트를 표준 LLVM 다이얼렉트로 변환합니다.
  • reconcile-unrealized-casts: 실현되지 않은 캐스트 연산을 조정하여 남은 타입 변환 문제를 해결합니다.
  • gpu-to-llvm { use-bare-pointers-for-host use-bare-pointers-for-kernels }: 남은 GPU 다이얼렉트 연산을 LLVM 다이얼렉트로 변환합니다. 호스트 코드와 디바이스 코드 모두에 베어 포인터를 사용합니다.

convert-linalg-to-affine-loops로 linalg 연산을 affine 루프로 로어링하면 다음 MLIR 모듈을 얻습니다.

module {
  func.func @square(%arg0: memref<10x10xf32>, %arg1: memref<10x10xf32>) -> memref<10x10xf32> {
    affine.for %arg2 = 0 to 10 {
      affine.for %arg3 = 0 to 10 {
        %0 = affine.load %arg0[%arg2, %arg3] : memref<10x10xf32>
        %1 = arith.mulf %0, %0 : f32
        affine.store %1, %arg1[%arg2, %arg3] : memref<10x10xf32>
      }
    }
    return %arg1 : memref<10x10xf32>
  }
}

다음은 convert-affine-for-to-gpu 패스를 적용한 결과입니다. 이 패스는 affine 루프를 GPU 커널 실행 영역으로 변환합니다.

func.func @square(%arg0: memref<10x10xf32>, %arg1: memref<10x10xf32>) -> memref<10x10xf32> {
  %c0 = arith.constant 0 : index
  %c10 = arith.constant 10 : index
  %0 = arith.subi %c10, %c0 : index
  %c1 = arith.constant 1 : index
  %c0_0 = arith.constant 0 : index
  %c10_1 = arith.constant 10 : index
  %1 = arith.subi %c10_1, %c0_0 : index
  %c1_2 = arith.constant 1 : index
  %c1_3 = arith.constant 1 : index
  gpu.launch
    blocks(%arg2, %arg3, %arg4) in (%arg8 = %0, %arg9 = %c1_3, %arg10 = %c1_3)
    threads(%arg5, %arg6, %arg7) in (%arg11 = %1, %arg12 = %c1_3, %arg13 = %c1_3) {
    %2 = arith.addi %c0, %arg2 : index
    %3 = arith.addi %c0_0, %arg5 : index
    %4 = affine.load %arg0[%2, %3] : memref<10x10xf32>
    %5 = arith.mulf %4, %4 : f32
    affine.store %5, %arg1[%2, %3] : memref<10x10xf32>
    gpu.terminator
  }
  return %arg1 : memref<10x10xf32>
}

다음은 gpu-kernel-outlining 패스를 적용한 결과입니다. 이 패스는 GPU 커널 실행 영역을 별도의 GPU 함수로 추출합니다. 이 함수에는 블록 인덱스와 스레드 인덱스를 참조하는 GPU 내장 함수가 들어 있습니다.

module attributes {gpu.container_module} {
  func.func @square(%arg0: memref<10x10xf32>, %arg1: memref<10x10xf32>) -> memref<10x10xf32> {
    %c0 = arith.constant 0 : index
    %c10 = arith.constant 10 : index
    %0 = arith.subi %c10, %c0 : index
    %c1 = arith.constant 1 : index
    %c0_0 = arith.constant 0 : index
    %c10_1 = arith.constant 10 : index
    %1 = arith.subi %c10_1, %c0_0 : index
    %c1_2 = arith.constant 1 : index
    %c1_3 = arith.constant 1 : index
    gpu.launch_func  @square_kernel::@square_kernel
        blocks in (%0, %c1_3, %c1_3)
        threads in (%1, %c1_3, %c1_3)
        args(%c0 : index, %c0_0 : index, %arg0 : memref<10x10xf32>, %arg1 : memref<10x10xf32>)
    return %arg1 : memref<10x10xf32>
  }
  gpu.module @square_kernel {
    gpu.func @square_kernel(%arg0: index, %arg1: index, %arg2: memref<10x10xf32>, %arg3: memref<10x10xf32>) kernel {
      %block_id_x = gpu.block_id  x
      %block_id_y = gpu.block_id  y
      %block_id_z = gpu.block_id  z
      %thread_id_x = gpu.thread_id  x
      %thread_id_y = gpu.thread_id  y
      %thread_id_z = gpu.thread_id  z
      %grid_dim_x = gpu.grid_dim  x
      %grid_dim_y = gpu.grid_dim  y
      %grid_dim_z = gpu.grid_dim  z
      %block_dim_x = gpu.block_dim  x
      %block_dim_y = gpu.block_dim  y
      %block_dim_z = gpu.block_dim  z
      %0 = arith.addi %arg0, %block_id_x : index
      %1 = arith.addi %arg1, %thread_id_x : index
      %2 = affine.load %arg2[%0, %1] : memref<10x10xf32>
      %3 = arith.mulf %2, %2 : f32
      affine.store %3, %arg3[%0, %1] : memref<10x10xf32>
      gpu.return
    }
  }
}

나머지 패스는 GPU 다이얼렉트를 NVVM 다이얼렉트로 변환하는 꽤 기계적인 로어링 패스입니다. 그러면 GPU 커널을 담은 gpu.module이 들어 있는 모듈을 얻습니다.

module attributes {gpu.container_module} {
  llvm.func @square(%arg0: !llvm.ptr, %arg1: !llvm.ptr) -> !llvm.ptr {
    %0 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %1 = llvm.insertvalue %arg1, %0[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %2 = llvm.insertvalue %arg1, %1[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %3 = llvm.mlir.constant(0 : index) : i64
    %4 = llvm.insertvalue %3, %2[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %5 = llvm.mlir.constant(10 : index) : i64
    %6 = llvm.insertvalue %5, %4[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %7 = llvm.mlir.constant(10 : index) : i64
    %8 = llvm.insertvalue %7, %6[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %9 = llvm.mlir.constant(10 : index) : i64
    %10 = llvm.insertvalue %9, %8[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %11 = llvm.mlir.constant(1 : index) : i64
    %12 = llvm.insertvalue %11, %10[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %13 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %14 = llvm.insertvalue %arg0, %13[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %15 = llvm.insertvalue %arg0, %14[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %16 = llvm.mlir.constant(0 : index) : i64
    %17 = llvm.insertvalue %16, %15[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %18 = llvm.mlir.constant(10 : index) : i64
    %19 = llvm.insertvalue %18, %17[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %20 = llvm.mlir.constant(10 : index) : i64
    %21 = llvm.insertvalue %20, %19[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %22 = llvm.mlir.constant(10 : index) : i64
    %23 = llvm.insertvalue %22, %21[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %24 = llvm.mlir.constant(1 : index) : i64
    %25 = llvm.insertvalue %24, %23[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %26 = llvm.mlir.constant(1 : index) : i64
    %27 = llvm.mlir.constant(0 : index) : i64
    %28 = llvm.mlir.constant(10 : index) : i64
    %29 = llvm.extractvalue %25[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %30 = llvm.extractvalue %12[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    gpu.launch_func  @square_kernel::@square_kernel blocks in (%28, %26, %26) threads in (%28, %26, %26) : i64 args(%27 : i64, %27 : i64, %29 : !llvm.ptr, %30 : !llvm.ptr)
    %31 = llvm.extractvalue %12[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    llvm.return %31 : !llvm.ptr
  }
  gpu.module @square_kernel [#nvvm.target<O = 3, chip = "sm_90", features = "+ptx80">] {
    llvm.func @square_kernel(%arg0: i64, %arg1: i64, %arg2: !llvm.ptr, %arg3: !llvm.ptr) attributes {gpu.kernel, nvvm.kernel} {
      %0 = llvm.mlir.constant(10 : index) : i64
      %1 = nvvm.read.ptx.sreg.ctaid.x : i32
      %2 = llvm.sext %1 : i32 to i64
      %3 = nvvm.read.ptx.sreg.tid.x : i32
      %4 = llvm.sext %3 : i32 to i64
      %5 = llvm.add %arg0, %2 : i64
      %6 = llvm.add %arg1, %4 : i64
      %7 = llvm.mul %5, %0 : i64
      %8 = llvm.add %7, %6 : i64
      %9 = llvm.getelementptr %arg2[%8] : (!llvm.ptr, i64) -> !llvm.ptr, f32
      %10 = llvm.load %9 : !llvm.ptr -> f32
      %11 = llvm.fmul %10, %10 : f32
      %12 = llvm.mul %5, %0 : i64
      %13 = llvm.add %12, %6 : i64
      %14 = llvm.getelementptr %arg3[%13] : (!llvm.ptr, i64) -> !llvm.ptr, f32
      llvm.store %11, %14 : f32, !llvm.ptr
      llvm.return
    }
  }
}

모듈에서 GPU 함수를 추출할 수 있습니다.

module {
      llvm.func @square_kernel(%arg0: i64, %arg1: i64, %arg2: !llvm.ptr, %arg3: !llvm.ptr) attributes {gpu.kernel, nvvm.kernel} {
        %0 = llvm.mlir.constant(10 : index) : i64
        %1 = nvvm.read.ptx.sreg.ctaid.x : i32
        %2 = llvm.sext %1 : i32 to i64
        %3 = nvvm.read.ptx.sreg.tid.x : i32
        %4 = llvm.sext %3 : i32 to i64
        %5 = llvm.add %arg0, %2 : i64
        %6 = llvm.add %arg1, %4 : i64
        %7 = llvm.mul %5, %0 : i64
        %8 = llvm.add %7, %6 : i64
        %9 = llvm.getelementptr %arg2[%8] : (!llvm.ptr, i64) -> !llvm.ptr, f32
        %10 = llvm.load %9 : !llvm.ptr -> f32
        %11 = llvm.fmul %10, %10 : f32
        %12 = llvm.mul %5, %0 : i64
        %13 = llvm.add %12, %6 : i64
        %14 = llvm.getelementptr %arg3[%13] : (!llvm.ptr, i64) -> !llvm.ptr, f32
        llvm.store %11, %14 : f32, !llvm.ptr
        llvm.return
      }
}

그런 다음 MLIR 모듈을 LLVM IR로 변환합니다.

; ModuleID = 'LLVMDialectModule'
source_filename = "LLVMDialectModule"

define ptx_kernel void @square_kernel(i64 %0, i64 %1, ptr %2, ptr %3) {
  %5 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
  %6 = sext i32 %5 to i64
  %7 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
  %8 = sext i32 %7 to i64
  %9 = add i64 %0, %6
  %10 = add i64 %1, %8
  %11 = mul i64 %9, 10
  %12 = add i64 %11, %10
  %13 = getelementptr float, ptr %2, i64 %12
  %14 = load float, ptr %13, align 4
  %15 = fmul float %14, %14
  %16 = mul i64 %9, 10
  %17 = add i64 %16, %10
  %18 = getelementptr float, ptr %3, i64 %17
  store float %15, ptr %18, align 4
  ret void
}

; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare noundef i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() #0

; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare noundef i32 @llvm.nvvm.read.ptx.sreg.tid.x() #0

attributes #0 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }

!llvm.module.flags = !{!0}

!0 = !{i32 2, !"Debug Info Version", i32 3}

그런 다음 LLVM의 nvptx 백엔드를 사용해 PTX로 컴파일합니다. 이 백엔드는 llc를 통해 호출합니다.

//
// Generated by LLVM NVPTX Back-End
//

.version 7.8
.target sm_90
.address_size 64

    // .globl   square_kernel           // -- Begin function square_kernel
                                        // @square_kernel
.visible .entry square_kernel(
    .param .u64 square_kernel_param_0,
    .param .u64 square_kernel_param_1,
    .param .u64 .ptr .align 1 square_kernel_param_2,
    .param .u64 .ptr .align 1 square_kernel_param_3
)
{
    .reg .b32   %r<3>;
    .reg .f32   %f<3>;
    .reg .b64   %rd<15>;

// %bb.0:
    ld.param.u64    %rd1, [square_kernel_param_0];
    ld.param.u64    %rd2, [square_kernel_param_3];
    cvta.to.global.u64  %rd3, %rd2;
    ld.param.u64    %rd4, [square_kernel_param_1];
    ld.param.u64    %rd5, [square_kernel_param_2];
    cvta.to.global.u64  %rd6, %rd5;
    mov.u32     %r1, %ctaid.x;
    cvt.s64.s32     %rd7, %r1;
    mov.u32     %r2, %tid.x;
    cvt.s64.s32     %rd8, %r2;
    add.s64     %rd9, %rd1, %rd7;
    add.s64     %rd10, %rd4, %rd8;
    mad.lo.s64  %rd11, %rd9, 10, %rd10;
    shl.b64     %rd12, %rd11, 2;
    add.s64     %rd13, %rd6, %rd12;
    ld.global.f32   %f1, [%rd13];
    mul.rn.f32  %f2, %f1, %f1;
    add.s64     %rd14, %rd3, %rd12;
    st.global.f32   [%rd14], %f2;
    ret;
                                        // -- End function
}

MLIR 파이프라인

이제 두 모듈 compile.py와 run.py를 구성할 수 있습니다. 이 모듈은 GPU 코드 생성 파이프라인을 구현하여 고수준 MLIR을 최적화된 PTX 어셈블리로 변환합니다.

전체 소스 코드는 Github에서 볼 수 있으며 노트북으로도 제공됩니다.

이 예제에서는 전체 컴파일 과정을 조율하는 주 함수 compile_mlir_to_ptx를 정의합니다. 이 함수는 내부에서 먼저 MLIR 문자열을 파싱하여 모듈 표현으로 만든 다음, 신중하게 순서를 정한 일련의 변환 패스를 통해 GPU 컴파일 파이프라인을 적용합니다. 이 패스들은 코드를 고수준 텐서 연산에서 GPU 전용 구성 요소로 점진적으로 낮춥니다. 여기에는 메모리 접근을 처리하는 버퍼화, linalg 연산을 affine 루프로 변환하는 작업, 그 루프를 GPU 블록과 스레드에 매핑하는 작업, 마지막으로 NVVM 다이얼렉트를 통해 NVIDIA 전용 코드를 생성하는 작업이 포함됩니다. 변환이 끝나면 함수는 GPU 모듈을 추출하고 LLVM 도구를 사용해 PTX 어셈블리로 변환합니다. 그 결과 추가 컴파일 없이 NVIDIA GPU에서 바로 실행할 수 있는 코드가 만들어집니다.

import subprocess

from mlir.ir import Context, Module
from mlir.passmanager import PassManager


def compile_mlir_to_ptx(mlir_module_str: str, chip_type="sm_75"):
    """Compiles MLIR module string to PTX code."""
    with Context():
        # Parse the input module
        module = Module.parse(mlir_module_str)

        # Apply GPU compilation pipeline
        module, gpu_module = apply_gpu_pipeline(module, chip_type)

        # Generate PTX from the GPU module
        ptx = generate_ptx(str(gpu_module), chip_type)

    return ptx


def apply_gpu_pipeline(module, chip_type="sm_75"):
    """Applies the GPU compilation pipeline to the MLIR module."""
    pm = PassManager()
    pm.enable_ir_printing(print_after_change=True)
    pm.add("canonicalize")
    pm.add(
        "one-shot-bufferize{ bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map }"
    )
    pm.add("canonicalize")
    pm.add("convert-linalg-to-affine-loops")
    pm.add("func.func(affine-loop-invariant-code-motion)")
    pm.add("func.func(convert-affine-for-to-gpu)")
    pm.add("gpu-kernel-outlining")
    pm.add("lower-affine")
    pm.add("gpu-decompose-memrefs")
    pm.add("expand-strided-metadata")
    pm.add("normalize-memrefs")
    pm.add(
        "gpu.module(convert-gpu-to-nvvm{index-bitwidth=0 use-bare-ptr-memref-call-conv })"
    )
    pm.add(f"nvvm-attach-target{ {chip={chip_type} features=+ptx80 O=3} }")
    pm.add("convert-nvvm-to-llvm")
    pm.add("reconcile-unrealized-casts")
    pm.add("gpu-to-llvm { use-bare-pointers-for-host use-bare-pointers-for-kernels }")
    pm.run(module.operation)

    # Extract the GPU module
    gpu_module = extract_gpu_module(module)

    return module, gpu_module


def extract_gpu_module(module: Module) -> Module:
    """Extracts the GPU module from a transformed MLIR module."""
    # Navigate the operation tree to find the GPU module
    # Structure: module -> region[0] -> block[0] -> operations[1] (GPU host-device code)
    # -> region[0] -> block[0] -> operations[0] (GPU module)
    try:
        main_func_op = module.operation.regions[0].blocks[0].operations[1]
        gpu_module_op = main_func_op.regions[0].blocks[0].operations[0]

        # Create a new module from the GPU module operation
        gpu_module = Module.parse(str(gpu_module_op))
        return gpu_module
    except (IndexError, AttributeError) as e:
        raise RuntimeError(f"Failed to extract GPU module: {e}") from e


def generate_ptx(gpu_module_str, chip_type="sm_75"):
    """Generates PTX from an MLIR GPU module string."""
    # First convert MLIR to LLVM IR
    llvm_ir_result = subprocess.run(
        ["mlir-translate", "--mlir-to-llvmir", "-"],
        input=gpu_module_str,
        capture_output=True,
        text=True,
    )

    if llvm_ir_result.returncode != 0:
        print("Error generating LLVM IR:")
        print(llvm_ir_result.stderr)
        return None

    llvm_ir = llvm_ir_result.stdout

    # Then convert LLVM IR to PTX
    ptx_result = subprocess.run(
        ["llc", "-march=nvptx64", f"-mcpu={chip_type}", "-"],
        input=llvm_ir,
        capture_output=True,
        text=True,
    )

    if ptx_result.returncode != 0:
        print("Error generating PTX:")
        print(ptx_result.stderr)
        return None

    return ptx_result.stdout

이제 run_kernel 함수를 정의할 수 있습니다. 이 함수는 Python 코드와 GPU 하드웨어 사이의 인터페이스 역할을 하며, 컴파일된 PTX 커널의 실행을 관리합니다. 이 함수는 PTX 코드 문자열, 커널 이름, 인자와 그 타입, 그리고 병렬화 전략을 정의하는 그리드 및 블록 차원을 받습니다. 이 함수는 먼저 CUDA 드라이버 API를 사용해 PTX 코드를 CUDA 모듈로 로드한 다음, 이름으로 지정한 커널 함수의 핸들을 가져옵니다. CUDA 런타임이 기대하는 형식으로 커널 인자를 준비한 뒤, 지정한 실행 구성으로 커널을 실행합니다. 이 구성은 몇 개의 스레드 블록과 블록당 몇 개의 스레드가 데이터를 처리할지를 결정합니다.

def run_kernel(
    ptx_code,
    kernel_name,
    args,
    arg_types,
    grid_dims,
    block_dims,
):
    """Run a PTX kernel."""
    module = checkCudaErrors(cu.cuModuleLoadData(ptx_code.encode("utf-8")))

    kernel_func = checkCudaErrors(
        cu.cuModuleGetFunction(module, kernel_name.encode("utf-8"))
    )

    kernel_args = (tuple(args), tuple(arg_types))

    checkCudaErrors(
        cu.cuLaunchKernel(
            kernel_func,
            grid_dims[0],
            grid_dims[1],
            grid_dims[2],
            block_dims[0],
            block_dims[1],
            block_dims[2],
            0,  # shared memory bytes
            0,  # stream
            kernel_args,  # kernel args
            0,  # extra
        )
    )

    checkCudaErrors(cu.cuCtxSynchronize())

    checkCudaErrors(cu.cuModuleUnload(module))

다음 코드는 이 컴파일 파이프라인을 사용하여 GPU 커널을 실행하는 완전한 엔드투엔드 예제입니다. 먼저 파일의 MLIR 코드를 PTX 어셈블리로 컴파일한 다음, CUDA 컨텍스트를 설정하고 호스트와 디바이스 양쪽에 메모리를 할당합니다. 이 예제는 각 스레드가 원소 하나를 처리하는 정방 행렬 연산을 준비하며, 그리드 차원은 행렬의 행 수에, 블록 차원은 열 수에 맞추어 계산을 구성합니다. 입력 데이터를 GPU로 복사한 뒤 적절한 인자와 차원으로 커널을 실행하고, 결과를 다시 호스트의 NumPy 배열로 가져옵니다.

import ctypes
import numpy as np
import cuda.cuda as cu
import cuda.cudart as cudart
import cuda.nvrtc as nvrtc
import subprocess
from compile import compile_mlir_to_ptx
from run import run_kernel

ptx_code = compile_mlir_to_ptx(open('square.mlir').read())

cuda_context = setup_cuda()

try:
    # Allocate device memory
    d_input = allocate_device_memory(input_data.nbytes)
    output_data = np.zeros((size, size), dtype=np.float32)
    d_output = allocate_device_memory(output_data.nbytes)

    # Copy input data to device
    copy_host_to_device(input_data, d_input)

    # Run kernel
    grid_dims = (size, 1, 1)  # One thread block per row
    block_dims = (size, 1, 1)  # One thread per column

    # Prepare arguments according to the PTX code
    # square_kernel(
    #     .param .u64 square_kernel_param_0,                // Grid dimension offset
    #     .param .u64 square_kernel_param_1,                // Block dimension offset
    #     .param .u64 .ptr .align 1 square_kernel_param_2,  // Input pointer
    #     .param .u64 .ptr .align 1 square_kernel_param_3   // Output pointer
    # )
    args = [
        0,  # Grid dimension offset
        0,  # Block dimension offset
        d_input,  # Input pointer
        d_output,  # Output pointer
    ]
    arg_types = [ctypes.c_int, ctypes.c_int, None, None]  # Using None for pointer types

    print("Running kernel on GPU...")
    run_kernel(
        ptx_code,
        "square_kernel",
        args,
        arg_types,
        grid_dims,
        block_dims,
    )

    # Copy results back to host
    copy_device_to_host(d_output, output_data)

    # Verify results
    print("Verifying results...")
    np.testing.assert_allclose(output_data, expected_output, rtol=1e-5)
    print("Success! Results verified.")

finally:
    # Clean up resources
    free_device_memory(d_input)
    free_device_memory(d_output)
    cleanup_cuda(cuda_context)

이것으로 끝입니다. 작은 컴파일러 백엔드의 뼈대가 완성되었습니다. MLIR 텐서 연산이 들어가면 커널이 GPU에서 실행됩니다. 9부에서는 Python에 임베딩된 작은 언어로부터 그 MLIR을 생성합니다. 먼저 원소별 표현식부터 시작한 다음, 트랜스포머에 필요한 리덕션과 행렬 연산을 추가하겠습니다.

바이너리 임베딩

덧붙이자면, GPU 커널을 컴파일하는 또 다른 방법은 커널을 바이너리로 컴파일하고, 그 바이너리를 MLIR 모듈에 임베딩한 다음, 호스트 코드가 바이너리를 로드하여 커널을 실행하도록 하는 것입니다.

MLIR의 gpu.binary 연산은 GPU 디바이스에 로드하여 실행할 수 있는 컴파일된 GPU 커널을 나타냅니다. 이 연산은 GPU 코드의 바이너리 표현을 담고 있으며, 그 형식은 PTX 어셈블리, CUBIN 바이너리, AMD GPU용 HSACO 등 다양할 수 있습니다. 이 연산은 미리 컴파일된 GPU 커널을 MLIR 모듈에 직접 임베딩하려 할 때 특히 유용하며, 런타임에 컴파일 파이프라인을 거치지 않고 바로 실행할 수 있게 해 줍니다.

gpu-module-to-binary 패스는 GPU 모듈을 GPU 바이너리로 변환하는 변환입니다. 이 패스는 MLIR 모듈을 훑어 중첩된 모든 GPU 모듈을 찾고, 각 모듈에 붙은 타깃 속성에 따라 이를 직렬화합니다. 이 패스는 지정된 모든 타깃 아키텍처마다 객체 하나씩을 담은 GPU 바이너리를 생성합니다. 이 패스는 오프로딩 표현, 어셈블리 코드, 바이너리, 팻바이너리 등 다양한 출력 형식을 지원합니다.

내부적으로 하는 일은 PTX 코드에 ptxas를 실행하고, gpu-module-to-binary 패스가 ptxas 명령에 전달한 인자와 함께 그 출력을 MLIR 모듈에 임베딩하는 것뿐입니다. 예를 들어 다음 MLIR 모듈은 sm_70 아키텍처용 GPU 바이너리를 담고 있습니다.

module attributes {gpu.container_module} {
  gpu.binary @vecadd_kernel  [
    #gpu.object<#nvvm.target<chip = "sm_70">,
    offload = "BC\C0\DE5\14 [... truncated ...] kernel20.1.0nvptx64-nvidia-cudaLLVMDialectModule\00\00\00\00">]
}

그런 다음 호스트 코드에서 다음과 같이 커널을 실행할 수 있습니다.

gpu.launch_func @vecadd_kernel::@vecadd_kernel
  blocks in (%0, %0, %0)
  threads in (%0, %0, %0) : i64
  args(%1 : i32, %2 : i64)

외부 자료

이전: 7부 - 트랜스포머 · 시리즈 목차 · 다음: 9부 - MLIR을 위한 Python 프런트엔드

원문 © Stephen Diehl, stephendiehl.com · 글 CC BY-SA 4.0, 코드 MIT · Stephen Diehl은 런던의 소프트웨어 엔지니어이자 작가입니다(저자가 정한 소개 문구). 이 글은 저자가 개인 자격으로 쓴 것이며 어떤 조직과의 연관도 밝히거나 암시해서는 안 됩니다 (라이선스와 고지).
이 페이지는 원문을 한국어로 옮기며 변경한 비공식 번역이며 같은 CC BY-SA 4.0으로 제공합니다 · yeokja 한국어 번역 · Anthropic 사의 claude-opus-5-5 모델을 활용하여 번역되었으며 학습을 모두 비허용한 상태로 작업하였습니다.