파이토치 PyTorch CUDA custom layer 만들기 (2) - 커널 빌드 및 Python 바인딩

이전 글:

2021.04.28 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (1) - CUDA 커널

파이토치 PyTorch CUDA custom layer 만들기 (1) - CUDA 커널

기본적인 CUDA 커널을 작성했으면, 두 가지의 작업 방향이 생긴다.

 

일단 당연히 커널을 빌드하긴 해야하는데, 그다음

1. 빌드한 커널 검증.
2. 빌드한 커널을 PyTorch에 연결.

의 선택지가 생긴다.

 

개발을 하거나 복잡한 프로그램을 만들어 본 경험이 있다면 당연히 1을 하고 2를 해야 하는 게 아닌가? 생각이 들 수 있지만

 

나 같은 경우는 CUDA 프로그래밍에 대한 정확한 이해가 없는 야매다 보니까 1에 대한 오버헤드가 엄청났다.

 

조금만 조심해서 2를 마치면 1을 PyTorch API로 할 수 있기 때문에, 일단 1을 스킵했다.

 

사실 퍼포먼스 분석 등을 하려면 1이 필수적이기에 결국 1을 좀 공부했는데

 

그냥 방학숙제 밀어 두기 정도의 게으름을 피운 셈이다.

 

 

아무튼 이전 글에서 작성한 vecadd_kernel을 빌드하는 것은 간단하다.

 

아래와 같이 cuda 폴더 아래에 Makefile을 만들었다.

NVCC = /usr/local/cuda-10.1/bin/nvcc
INCFLAGS = -I /usr/local/cuda-10.1/include
CUDAFLAGS = -shared -O2 -gencode=arch=compute_75,code=sm_75 -std=c++14
CUDAFLAGSADD = --compiler-options "-fPIC"

all: libvecadd_kernel.so

libvecadd_kernel.so: vecadd_kernel.cu
    $(NVCC) $(INCFLAGS) -o $@ -c $^ $(CUDAFLAGS) $(CUDAFLAGSADD)

clean:
    rm -f ./*.so*

안타깝게도 개인적으로 학부 시절 Visual Studio만을 사용하였고 (F5 만 누르면 만사가 해결된다.),

 

직접 Makefile을 작성해야 할 만한 대규모의 프로젝트를 진행하지 않았기 때문에

 

Makefile에 대한 이해도가 몹시 떨어진다.

 

 

아무튼 굳이 정리를 하자면,

• NVCC: NVIDIA CUDA compiler가 설치된 위치.
• INCFLAGS: CUDA include 파일들이 포함된 위치.
• CUDAFLAGS: CUDA 컴파일 옵션. Compatibility 오류가 발생하면 -gencode=arch 부분을 건드리면 된다.
• CUDAFLAGSADD: 해당 옵션을 주지 않으면 PyTorch에 바인딩시킬 수 없더라. 일단 이유는 모르지만 넣으라고 해서 넣음.

CUDA 프로그래밍을 하는데 당연히 CUDA는 설치가 되어있어야 한다.

 

10.1은 버전 맞춰서 맘대로 바꾸면 되고, CUDA 설치 시에 symbolic link를 만들었다면 cuda-10.1 대신 cuda만 넣어도 된다.

 

나머지 줄들은 그냥 이렇게 정의한 인수들을 죄다 집어넣고 libvecadd_kernel.so라는 shared object(.so)를 만드는 과정이다.

 

Makefile 매크로들도 야매로 배워서 썼는데,

• $@: 현재 object 파일. 위에서는 libvecadd_kernel.so를 의미할 듯.
• $^: 현재 target (libvecadd_kernel.so)이 의존하는 전체 대상 목록.
  위에서는 vecadd_kernel.cu를 써놨기 때문에 해당 파일을 가리킬 듯.

분명 더 깔끔하고 명확하게 작성할 수 있을 텐데, 언젠가 공부할 때가 올지 안 올진 모르겠다.

 

 

빌드를 하면 cuda 폴더에 libvecadd_kernel.so 파일이 생성된다.

$ make

 

이제 이렇게 만들어진 .so 파일을 PyTorch에서 호출할 수 있게 만들어야 한다.

 

아래 다이어그램을 리뷰해보면, vecadd.cpp와 setup.py가 필요하다는 것을 알 수 있는데

파이토치 CUDA custom layer 프로젝트 다이어그램.

한 마디로 정리하면, vecadd.cpp는 우리가 만든 CUDA 커널 (.so)과 PyTorch (C++)을 연결해주기 위해

 

setup.py에 제공되는 파트이다.

 

우선 이전처럼 cuda 폴더 아래에 vecadd.cpp를 아래와 같이 작성하자.

// vecadd.cpp
#include <torch/extension.h>
#include "vecadd_kernel.cuh"

void VecAdd(
    const torch::Tensor x,
    const torch::Tensor y,
    torch::Tensor z
)
{
    // Equivalent to nelement()
    int n = x.numel();
    VecAddCuda(
        x.data_ptr(),
        y.data_ptr(),
        z.data_ptr(),
        n
    );
    return;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("vecadd", &VecAdd, "Vector add");
}

C++ 코드이지만 PyTorch 코드처럼 보이기도 하는데,

 

Line 2에서 include 한 torch/extension.h가 PyTorch (C++)의 인터페이스를 제공한다.

 

해당 헤더의 경우 별도로 설치할 필요는 없고, PyTorch (Python)을 설치하면 자동으로 딸려온다.

 

이는 그리 좋지 못한 소식일 수도 있는데, IDE를 사용하며 자동완성이나 함수 추천을 받을 수 없다는 뜻도 된다.

 

공식 사이트의 documentation이 형편없는 수준이므로, 필요한 함수가 있으면 재주껏 구글링 하자...

 

 

아무튼 vecadd.cpp에는 VecAdd라는 함수를 정의하는데,

 

인수의 자료형을 보면 알 수 있듯이 PyTorch (C++) 텐서를 받는다.

 

이 레벨부터는 상당한 추상화가 이루어진 상태이고 텐서들은 high-level 객체들이기 때문에

 

쌩 CUDA 코드에서 하기 귀찮았던 것들을 여기서 조금 만져주기 용이하다 (예외 처리라던가.).

 

다만 해당 사항들은 PyTorch (C++)와 친분이 있는 사람에 한정되므로

 

여기에서는 텐서 3개만 달랑 주어지면 CUDA 커널에 넘겨주기 위한 x 텐서의 크기를 구하는 코드만 들어갔다.

 

아쉬운 것은 C++와 Python 버전의 네이밍이 조금 다를 때가 있다는 것인데

 

예를 들어 Python 버전의 Tensor.nelement() 함수는 numel()로 바뀌어버렸다.

 

직관성은 내다 버린 디자인이라서 최소한의 전처리만 하도록 하자.

 

 

Line 5~20은 PyTorch (C++)와 CUDA 커널을 연결하는 부분인데,

 

아무래도 둘 다 C++을 기반으로 하기에 접착성이 상당히 좋은 편이다.

 

PyTorch 텐서를 data_ptr<float>()를 호출해서 float 포인터로 바꿔주면

 

바로 우리가 작성한 CUDA 커널을 호출하는 함수에 넘겨주는 것이 가능하다.

 

data_ptr()을 대신 사용하면 void 포인터를 얻을 수 있는데,

 

해당 텐서의 자료형이 C++ 기본 자료형이 아닌 경우에 유용하다.

 

(왜냐하면 vecadd.cpp는 CUDA 프로그램이 아니기 때문)

 

대표적인 예가 HalfTensor인데, CUDA에는 해당 자료형이 존재하지만 C++는 아니다.

 

이런 경우 일단 void 포인터를 CUDA 프로그램으로 넘겨주고, 이를 적당히 캐스팅해서 사용하면 된다.

 

만약 PyTorch 텐서가 GPU 상에 있으면, data_ptr()이 가리키는 주소도 GPU 상의 배열이기 때문에

 

cudaMemcpy 등의 지저분한 함수를 쓸 필요가 없는 점이 참 좋다.

 

z는 미리 할당해놓았는데 C++에서는 메모리 관리를 하나도 하지 않겠다는 의지가 들어있다.

 

나중에 PyTorch 인터페이스를 만들 때 다시 등장할 듯.

 

 

아무튼 이렇게 어렵지 않게 함수 작성이 완료된다.

 

마지막으로 이렇게 정의한 함수를 Python에서 호출할 수 있도록 바인딩해주는 과정이 필요한데,

 

Line 22~24의 형태를 그대로 사용하면 된다

 

def()의 인수들은 차례대로 우리가 Python에서 사용할 함수명, C++에 정의한 함수의 reference, 간단한 설명이다.

 

여기를 참고하면 될 듯한데, 아직까지는 봐도 별 의미가 없는 것 같다.

 

여러 개의 함수를 바인딩하려면 Line 23 아래에 같은 형식으로 추가적인 def를 넣어주면 된다.

 

이제 만든 함수를 Python에서 호출할 수 있도록 빌드해주면 되는데

 

별도의 Makefile을 작성하지 않고도 setup.py로 해결이 가능하다.

 

나에게 있어 정말 다행인 소식은 이제 CUDA와 C++ 코드를 보며 씨름하지 않고,

 

Python 레벨에서 남은 작업을 마무리하면 된다는 것이다.

 

(물론 튜닝이 필요하면 다시 돌아가야 한다...)

 

아무튼 setup.py와 추가적인 PyTorch 인터페이스 구현은 다음에 정리해야겠다.

 

다음 글:

2021.04.30 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (3) - setup.py

파이토치 PyTorch CUDA custom layer 만들기 (3) - setup.py