종종 기발한 아이디어를 떠올리다 보면,
"이걸 구현하기 위한 레이어가 파이토치에 있었나?" 라는 의구심이 든다.
다행히 세상이 좋아져서 대부분의 멋진 아이디어들은 쓰기 좋고 깔끔한 파이토치 구현체가 있지만,
정 안 되는 경우 내가 직접 만들어야 하는 경우가 생긴다.
파이토치의 custom layer는 만들기 어려운 편은 아니지만,
효율 등의 문제로 CUDA 프로그래밍이 필요하다면 문제가 많이 복잡해진다.
지난 CVPR을 준비하며 CUDA 프로그래밍을 동반한 custom layer를 만들어보았는데,
언젠가 또 비슷한 걸 만들지 모르니 정리해두는 게 좋겠다.
Toy example로 vector x와 y를 더하는 vecadd의 CUDA custom layer를 만들어보자.
(즉, z = x + y)
vecadd를 동작하게 만들 수 있다면, 훨씬 더 복잡한 연산도 문제없이 구현이 가능하다.
물론 연산 자체를 C++로 어떻게 짜야하는지 모르겠다면 그건 어쩔 수 없지만...
아무튼 전체 프로젝트는 아래 다이어그램처럼 구성된다.
뭐가 이렇게 많나 생각이 들지만, 구조를 이해하면 나름 합리적이다.
(그렇다고 내가 각각의 디테일을 다 이해한 것은 아니다.)
가장 먼저 하면 좋은 것은 CUDA 커널을 작성하는 것이다. (사실 가장 어려운 일이기도 하다.)
직접 해 보면 나머지 파일들은 대부분 boilerplate 코드를 적절히 수정하는 것이지만
CUDA 커널은 layer의 핵심 로직이 담겨있는 것과 동시에 상당히 낯선 프로그래밍 스킬을 요구한다.
조금 복잡한 연산을 구현하고자 하는 경우에는,
파이썬 등으로 해당 로직을 충분히 검증한 뒤 CUDA 커널을 작성하는 게 시간을 아끼는 길이다.
안타깝게도 나는 CUDA 프로그래밍을 지난 CVPR을 준비하며 처음 해 보아,
이것저것 다양한 잡기술들을 적어놓을 재주가 없다.
그렇기에 일단 최대한 작동하는 무언가를 만들어놓고,
거기서 하면 안 되는 것, 주의할 것 등을 적어놓는 게 나중에 참조하는 데에 있어 훨씬 유용할 것 같다.
아 참, 일단 CUDA 프로그래밍을 하려면 시스템에 CUDA가 설치되어 있어야 한다.
다른 사람들이 열심히 정리해 놓았으니 그 부분은 각자 알아서.
일단 아무 이름으로 프로젝트 폴더를 만들고, 그 아래에 cuda라는 이름의 폴더를 추가로 만들자.
CUDA 커널 관련한 모든 코드는 그 폴더에 넣는다.
CUDA 프로그래밍을 처음 접하면 조금 헷갈리는 것이,
CUDA 커널 자체와 CUDA 커널을 호출하는 함수 2개를 짜야한다.
두 함수의 구현은 .cu 파일에 넣고, CUDA 커널을 호출하는 함수의 정의만 .cuh 파일에 넣는다.
따라서 vecadd_kernel.cuh 파일은 아래와 같이 작성할 수 있다.
// vecadd_kernel.cuh
#ifndef VECADD
#define VECADD
void VecAddCuda(
const float* x,
const float* y,
float* z,
const int n
);
#endif //VECADD기본적으로 파이토치 텐서도 메모리 레이아웃상에서는 C의 배열과 동일하기에 float* 형태로 인수를 전달받는 게 편하다.
PyTorch의 C++ wrapper를 사용하는 방법도 있지만, 밑준비가 배로 늘어나므로 개인적으로는 이 방법이 훨씬 쉬운 것 같다.
그리고 z = x + y를 구현하고 싶기 때문에, z만 const를 붙이지 않는다.
주의해야 할 것은 C의 배열 크기 정보에 직접 접근할 수 있는 방법이 없기 때문에, n을 같이 넘겨줘야 한다.
조금 더 정교한 프로그래밍을 하고 싶으면 x, y, z의 배열 크기가 다를 때, 자료형이 다를 때 적절한 예외 처리가 필요하지만
그건 조금 더 high-level에서 다루고, 일단 CUDA 코드 상에서는 모든 게 잘 갖춰져 있다고 가정하고 진행하는 게 속 편하다.
다음은 실제 함수를 구현하는 vecadd_kernel.cu 파일이다.
// vecadd_kernel.cu
#ifndef NUM_THREADS
#define NUM_THREADS 1024
#endif
#include <cuda.h>
#include <cuda_runtime.h>
#include "vecadd_kernel.cuh"
__global__ void VecAddCudaKernel(
const float* __restrict__ x,
const float* __restrict__ y,
float* __restrict__ z,
const int n
)
{
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
z[idx] = x[idx] + y[idx];
}
return;
}
void VecAddCuda(
const float* x,
const float* y,
float* z,
const int n
)
{
int n_blocks = int((n + NUM_THREADS - 1) / NUM_THREADS);
VecAddCudaKernel<<<n_blocks, NUM_THREADS>>>(x, y, z, n);
return;
}Line 2~9는 선언 부분이다.
특히 CUDA 프로그램 빌드를 위해서는 Line 6, 7의 cuda.h와 cuda_runtime.h를 include 해야 한다.
해당 파일들이 어디에 있는지 몰라도, 시스템에 CUDA가 설치되어 있으면 상관없다.
NUM_THREADS는 디바이스 특성인데 현재 딥러닝에 정상적으로 사용할 수 있는 디바이스들이면
Line 11~23은 CUDA 커널이다.
GPU에서 돌아가는 코드이기 때문에, 기존 C++에서 볼 수 없었던 예약어들이 여럿 보인다.
잠시 뒤로 하고 Line 25~35를 보면 아까 헤더 파일에서 선언한 VecAddCuda의 실질적인 구현이 있다.
내용물을 보면 별거 없는 게, 단순히 CUDA 커널 VecAddCudaKernel을 호출하는 인터페이스 역할만을 한다.
여기서 조금 더 깔끔하고 이해하기 쉬운 프로그래밍을 하려면, grid - block - thread의 hierarchy를 알아야 한다.
warp라는 개념도 있는 것 같은데, 아주 빡센 최적화를 할 때 고려하는 요소 같으니 패스.
어차피 자꾸 봐도 헷갈리니까 여기에 정리해 놓자.
쉽게 말해 grid는 하나의 CUDA 커널마다 하나씩 할당되는 block들의 집합,
block은 같은 CUDA 코드를 실행시키는 thread의 집합이다.
더 자세한 설명은 stackoverflow나 구글이 알려줄 것.
그림에서는 block과 thread가 2D 배열처럼 놓여있지만
이는 일단 프로그래밍의 편의를 위한 것으로 1D, 2D, 3D 중 원하는 인덱싱 방법을 고를 수 있다.
(여기에서는 1D로 놓고 했다. 2D, 3D도 써봤는데 나중에 정리해야지.)
다만 thread의 경우 한 block 안에 NUM_THREADS (=1024) 개가 최대이니 잘 배치해보자.
잠시 코드로 돌아가면, Line 32의 n_blocks는 우리가 실행시킬 block의 개수이다.
하나의 block에 NUM_THREADS개의 커널을 실행시키고
총 n번의 커널 호출 (element-wise 덧셈)이 필요하므로 이에 맞춰서 필요한 block 수를 계산한다.
이를 이용해서 Line 33에서 <<< >>> 표현을 통해 CUDA 커널을 호출하는데,
총 n_blocks * NUM_THREADS 개의 커널이 호출된다.
n이 NUM_THREADS의 배수가 아닌 경우 노는 커널이 있을 수 있다는 걸 알아놓자.
Line 11에서 __global__로 시작하는 부분은 해당 함수가 GPU 상에서 실행된다는 걸 명시한다.
또한 인수 부분의 __restrict__는 __restrict__가 붙은 인수들이 메모리상에서 독립적이라는 걸 의미하며,
프로그램 최적화에 도움을 준다고 한다.
막상 작성할 때는 습관적으로 넣었지만, 정리하며 다시 생각해보니 z = x + y에서
x, y, z가 항상 서로 다른 메모리 주소를 가리킨다는 보장이 없으므로 빼는 게 맞는 것 같다.
(예: x = x + x)
그냥 이런 것도 있구나 하고 넘어가면 될 듯.
그다음은 현재 커널의 인덱스를 구해줘야 한다.
표현이 조금 헷갈리기 쉬운데,
blockIdx.x는 이게 x 방향 (1D)으로 몇 번째 block인지 (block마다 다름),
blockDim.x는 한 block 안에 x 방향으로 총 몇 개의 thread가 있는지 (하나의 커널에서는 항상 동일),
마찬가지로 threadIdx.x는 해당 block에서 x 방향으로 몇 번째의 thread인지를 나타낸다.
쉽게 말해 Idx는 해당 방향에서의 절대적 위치, Dim은 하위 요소가 해당 방향으로 몇 개나 들어가 있는지를 의미한다.
따라서 현재 커널의 인덱스는 Line 18처럼 구할 수 있다.
만약 block 안에 64 * 16같이 2D로 thread를 배치하였으면 아래와 같이 커널 인덱스를 구하면 될 것 같다.
int idx = blockIdx.x * (blockDim.x * blockDim.y) + threadIdx.x * blockDim.x + threadIdx.y;여기에 조금 더 구체적인 예시가 있으니 참고.
Line 19에서는, 해당 커널이 유효한 커널인지를 확인 후에 Line 20에서 실제 연산이 이루어진다.
총 n_blocks * NUM_THREADS 개의 커널이 호출되므로, 이게 n보다 크다면 위에 말한 대로 남는 커널은 그냥 return.
하나의 커널이 x, y, z 텐서의 하나의 element를 담당하며 덧셈을 수행한다.
물론 최적화나 기타 구현 상의 이유로 하나의 커널에서 여러 개의 element에 접근해도 전혀 문제가 없다.
중요한 것은 n개의 커널이 n개의 결과를 온전히 만드는 것.
일단 이렇게 하면 가장 중요한 부분인 CUDA 프로그래밍은 어느 정도 완료되었다.
다음에는 이렇게 작성한 커널들을 빌드하고 PyTorch에 연결하는 작업이다.
다음 글:
2021.04.30 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (2) - 커널 빌드 및 Python 바인딩
파이토치 PyTorch CUDA custom layer 만들기 (2) - 커널 빌드 및 Python 바인딩
이전 글: 2021.04.28 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (1) - CUDA 커널 파이토치 PyTorch CUDA custom layer 만들기 (1) - CUDA 커널 종종 기발한 아이디어를 떠올리다 보면, "이걸..
sanghyun.tistory.com