이전 글:

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

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

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

 

직접 구현한 PyTorch 연산에서 backward를 구현하기 위해서는 두 가지가 필요한데,

  1. back-propagation 되는 미분 값을 계산하는 코드
  2. autograd.Function을 사용한 wrapping

이다.

 

다행히 두 가지 모두 어렵지 않게 처리가 가능하다.

 

일단 chain rule을 사용하면, 임의의 변수 $x$와 $z$에 대해 아래와 같은 관계가 성립한다.

 

$\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial z} \frac{\partial z}{\partial x}$.

 

구현한 연산 $z = x + y$에서 $\frac{\partial z}{\partial x} = 1$이므로,

 

$\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial z}$이고 $x$ 대신 $y$를 넣어도 성립한다.

 

Backpropagation이란 $\frac{\partial \mathcal{L}}{\partial z}$를 알 때 $\frac{\partial \mathcal{L}}{\partial x}$와 $\frac{\partial \mathcal{L}}{\partial y}$를 구하는 것이므로 계산 끝.

 

CVPR에 사용하기 위해 구현한 코드는 몇 배 이상 복잡하고 잔머리도 많이 굴려야했는데,

 

가능하다면 Python 등을 활용해 미리 내 로직이 맞는지 점검하는 과정이 있어야 한다.

 

 

아무튼 이제 이걸 CUDA 커널에 적당히 프로그래밍하면 된다.

 

귀찮지만, 총 3개의 파일을 건드려야 한다.

 

우선 vecadd_kernel.cuh에 다음 Line들을 추가한다.

 //vecadd_kernel.cuh에 추가
void VecAddBackwardCuda(
    float* dx,
    float* dy,
    const float* dz,
    const int n
);

이번에는 dz를 받아 dx, dy를 구하는 것이 목적이므로 (1)과 비교하면 const의 위치가 바뀐 것을 주의해야겠다.

 

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

 

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

종종 기발한 아이디어를 떠올리다 보면, "이걸 구현하기 위한 레이어가 파이토치에 있었나?" 라는 의구심이 든다. 다행히 세상이 좋아져서 대부분의 멋진 아이디어들은 쓰기 좋고 깔끔한 파이토

sanghyun.tistory.com

 

다음으로 vecadd_kernel.cuCUDA 커널을 호출하는 함수실제 CUDA 커널을 작성한다.

// vecadd_kernel.cu에 추가
__global__ void VecAddBackwardCudaKernel(
    float* __restrict__ dx,
    float* __restrict__ dy,
    const float* __restrict__ dz,
    const int n
)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        dx[idx] = dz[idx];
        dy[idx] = dz[idx];
    }
    return;
}

void VecAddBackwardCuda(
    float* dx,
    float* dy,
    const float* dz,
    const int n
)
{
    int n_blocks = int((n + NUM_THREADS - 1) / NUM_THREADS);
    VecAddBackwardCudaKernel<<<n_blocks, NUM_THREADS>>>(dx, dy, dz, n);
    return;
}

실제 back-propagation이 구현이 CUDA 커널에서 어떻게 구현되었는지 참고하면 된다.

 

지난번에도 언급했지만, __restrict__는 in-place operation의 가능성이 있을 때는 사용하면 안 된다.

 

여기서도 지우는 게 맞는데 참고.

 

여담이지만 이 과정에서 문제 아닌 문제는 각 함수들의 이름을 짓는 게 썩 쉽지 않다는 것...

 

나는 Forward/Backward, Cuda, Kernel을 잘 조합해서 지었는데 (Forward는 생략했지만)

 

나중에 어떻게든 구분이 가능하게 재주껏 지으면 된다.

 

 

Makefile은 수정할 필요 없고, vecadd.cpp에도 backward를 추가해준다.

// vecadd.cpp에 추가
void VecAddBackward(
    torch::Tensor dx,
    torch::Tensor dy,
    const torch::Tensor dz
)
{
    int n = dx.numel();
    VecAddBackwardCuda(
        dx.data_ptr(),
        dy.data_ptr(),
        dz.data_ptr(),
        n
    );
    return;
}

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

새롭게 vecadd_backward 함수가 바인딩되는 것을 알 수 있다.

 

이렇게 해서 일단 구현 자체는 완료했고,

 

실제로 사용하려면 cuda_extension 바로 아래에 있는 Makefile을 통해서 빌드하고 호출하면 된다.

 

다만 이렇게 하는 경우 PyTorch의 가장 큰 강점 중 하나인 자동 미분이 안되므로

 

autograd.Function과 결합해주는 과정이 필요하다.

 

 

링크의 공식 documentation을 보면 잘 설명이 되어있는데, 적당히 필요한 정보를 취했다.

 

추가적으로, setup.py를 통해 작성한 autograd.Function을 패키지 안에 포함시켜

 

개별적인 파일 관리 없이 해당 패키지가 설치된 environment 내에서 자유롭게 import 할 수 있도록 했다.

 

이 과정은 이전 setup.py 작성 시에 다 고려를 했으므로 별도로 건드릴 내용은 없지만.

 

(setup.py의 packages=setuptools.find_packages() 부분)

 

 

아무튼 cuda_extension 폴더 아래에 (setup.py가 있는 위치) pyvecadd라는 폴더를 새로 만들고,

 

내용 없는 __init__.py를 해당 폴더에 추가하여 해당 폴더가 Python 패키지임을 명시해준다.

 

그다음 __init__.py가 있는 위치에  layer.py라는 파일을 아래와 같이 작성한다.

현재 폴더 구조. build, dist, *-info는 자동적으로 생성된 파일들.

# layer.py
import typing

import torch
from torch import autograd
from torch import nn

import vecadd_cuda


class VecAddFunction(autograd.Function):

    @staticmethod
    def forward(
            ctx: autograd.function._ContextMethodMixin,
            x: torch.Tensor,
            y: torch.Tensor) -> torch.Tensor:

        z = torch.empty_like(x)
        vecadd_cuda.vecadd(x, y, z)
        return z

    @staticmethod
    def backward(
            ctx: autograd.function._ContextMethodMixin,
            dz: torch.Tensor) -> typing.Tuple[torch.Tensor, torch.Tensor]:

        dx = torch.empty_like(dz)
        dy = torch.empty_like(dz)
        vecadd_cuda.vecadd_backward(dx, dy, dz)
        return dx, dy


vecadd = VecAddFunction.apply


class VecAddLayer(nn.Module):

    def __init__(self) -> None:
        super().__init__()
        return

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        z = vecadd(x, y)
        return z

요즘 Python 코드를 작성할 때 가능하면 type hinting (참고)을 사용하려고 하는데, 좋은지는 잘 모르겠다.

 

절대 필수가 아니므로 귀찮으면 다 빼버려도 된다 (훨씬 깔끔하게 보일 수도?).

 

 

결코 코드가 긴 것은 아니지만, 어디서부터 정리해야 할지 제법 막막하다.

 

Line 8은 우리가 setup.py로 빌드한 vecaddvecadd_backward를 포함하고 있는 vecadd_cuda를 import 한다.

 

VecAddFunction은 실제 자동 미분 로직을 구현하는 class인데,

 

공식 사이트의 documentation이 잘 되어있어 그대로 따라 하면 된다.

 

몇 가지만 가볍게(?) 정리해놓자면,

  • Line 20과 30이 실제 우리가 C++에서 바인딩하고, setup.py를 통해 빌드한 함수를 호출하는 부분이다.
  • Line 15와 25의 ctx는 forward시에 사용한 값이 backward시에도 필요한 경우, 이를 저장해두는 버퍼이다.
    아마 나중에 실제 구현한 코드를 정리하며 다시 등장할 수도 있는데,
    지금은 굳이 백업이 필요 없으므로 사용되지는 않았다.
  • Line 19, 28, 29는 return에 사용되는 값들을 위한 메모리를 할당하는 부분이다.
    C++ 나 CUDA 레벨에서 메모리 관리를 하기 싫기 때문에 PyTorch 레벨에서 관리하도록 했다.
    따라서 forward 로직은 이제 input x, y 2개 만을 인수로 받아 새로운 output z를 return 한다.
    이전처럼 모양 빠지게 z를 미리 할당하고, vecadd_cuda.vecadd(x, y, z) 같이 할 필요가 없다는 뜻.
  • backward에서는 forward에서 ctx를 제외하고 입력받은 모든 인수에 대해 gradient를 계산해야 한다.
    따라서 x와 y에 대한 gradient인 dx, dy를 순서대로 return 하는데,
    텐서가 아니거나 상수 인수를 forward에서 전달받은 경우에는 해당 순서에 None을 return 하면 된다.
  • VecAddFunction은 추상적인 개념이고, 실제로는 Line 34에 정의한 vecadd를 호출하여 사용한다.
    그 아래의 VecAddLayer는 해당 연산을 custom layer로 구현한 class인데,
    vecadd <-> torch.nn.functional.conv1d
    VecAddLayer <-> torch.nn.Conv1d
    같은 관계로 생각할 수 있다.
    실제로 VecAddLayer 내부에서 vecadd를 호출하여 계산을 진행하는데,
    torch.nn.Sequential 같은 모듈을 활용하려면 이런 layer 형태를 쓰는 것이 유리하다.

 

이제 진짜 완성이다!

 

마지막으로 cuda_extension 아래에 있는 setup.py를 다시 실행해서

  • 새로 구현한 backward CUDA 커널 및 C++ 인터페이스
  • 해당 함수의 Python 바인딩
  • 해당 함수의 자동 미분을 구현하는 vecadd 폴더 (패키지) 아래에 포함된 layer.py 모듈

들을 현재 environment에 설치한다.

 

 

사용법은 아래와 같다.

  • vecadd.cpp 파일에 구현되어 Python에 바인딩된 함수를 바로 사용 (이제는 그럴 일이 거의 없지만)
>>> import vecadd_cuda
>>> vecadd_cuda.vecadd(x, y, z)
>>> vecadd_cuda.vecadd_backward(dx, dy, dz)

import 하는 패키지의 이름은 우리가 setuptools.setup에 지정했던 ext_modules의 name이다.

 

당연히 자동 미분을 지원하지 않으며, z와 dz에 대한 사전 할당이 필요하다.

  • 자동 미분이 지원되는 PyTorch (Python)식 함수 사용 (이걸로 쓰자)
>>> from pyvecadd import layer
>>> z = layer.vecadd(x, y)
>>> z.mean().backward()

그냥 우리가 cuda_extension 폴더 아래에 만든 pyvecadd 폴더가

프로젝트 폴더에도 똑같이 있다고 생각하고 사용해주면 된다 (그래서 이름을 잘 지어놓아야 한다...).

실제로는 environment에 설치되어 있기 때문에,

 

파일을 매번 옮겨놓지 않아도 되는 점이 내가 가장 마음에 드는 부분이다.

 

조금 구체적인 예시는 아래와 같다.

>>> import torch
>>> from pyvecadd import layer
>>> x = torch.randn(4, 4, requires_grad=True, device=torch.device('cuda'))
>>> y = torch.randn(4, 4, requires_grad=True, device=torch.device('cuda'))
>>> print(x)
tensor([[-0.3005, -1.2514,  0.0675, -1.6077],
        [-1.1321,  0.8361, -0.7227,  1.4445],
        [ 1.1019, -0.2202, -0.3690,  0.3014],
        [ 0.6847, -0.7965, -0.7134,  2.0782]], device='cuda:0',
       requires_grad=True)
>>> print(y)
tensor([[ 1.4736, -0.5728,  0.4896, -0.9496],
        [-0.0979,  0.0057,  0.7292, -0.5978],
        [ 1.5153, -0.2581,  0.0160,  0.8342],
        [ 0.1993,  0.2244,  0.0465, -0.7861]], device='cuda:0',
       requires_grad=True)
>>> z = layer.vecadd(x, y)
>>> print(z)
tensor([[ 1.1731, -1.8242,  0.5571, -2.5573],
        [-1.2300,  0.8418,  0.0065,  0.8467],
        [ 2.6172, -0.4784, -0.3530,  1.1356],
        [ 0.8840, -0.5722, -0.6669,  1.2921]], device='cuda:0',
       grad_fn=<VecAddFunctionBackward>)
>>> print(x + y)
tensor([[ 1.1731, -1.8242,  0.5571, -2.5573],
        [-1.2300,  0.8418,  0.0065,  0.8467],
        [ 2.6172, -0.4784, -0.3530,  1.1356],
        [ 0.8840, -0.5722, -0.6669,  1.2921]], device='cuda:0',
       grad_fn=<AddBackward0>)
>>> z.mean().backward()
>>> print(x.grad)
tensor([[0.0625, 0.0625, 0.0625, 0.0625],
        [0.0625, 0.0625, 0.0625, 0.0625],
        [0.0625, 0.0625, 0.0625, 0.0625],
        [0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0')
>>> print(y.grad)
tensor([[0.0625, 0.0625, 0.0625, 0.0625],
        [0.0625, 0.0625, 0.0625, 0.0625],
        [0.0625, 0.0625, 0.0625, 0.0625],
        [0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0')

본래의 목적인 custom layer로 사용하는 것도 가능.

>>> vecadd_layer = layer.VecAddLayer()
>>> zz = vecadd_layer(x, y)

 

아무튼 이렇게 해서 PyTorch CUDA custom layer 구현에 대한 설명은 정리가 끝난 것 같다.

 

라고 말하고 싶은데, 사실 구현 대상 함수가 굉장히 쉬웠기 때문에 고려하지 않은 많은 요인이 있다.

 

이후의 내용은 코딩 실력이 완벽하여 에러를 절대 내지 않는다면 굳이 참조할 필요는 없지만,

 

나는 아쉽게도 그렇지 못해서 조금 더 정리를 해둬야 할 것 같다.

 

To be continued...

이전 글:

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

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

 

이전 포스트에 이어 PyTorch (C++)와 PyTorch (Python)을 연결하는 방법을 정리해야겠다.

 

이 부분은 의외로 간단한것이, setup.py 템플릿을 조금만 손보면

 

PyTorch의 cpp_extension이라는 모듈로 어렵지 않게 해결이 가능하다.

 

 

원리는 생략하고 아래와 같이 setup.py를 작성하자.

 

참고로 나는 setup.py의 위치를 cuda 폴더 바깥으로 했다.

 

전체 폴더 구조는 대충 아래와 같다.

cuda_extension 폴더 아래 setup.py를 놓는다.

# setup.py
from os import path
import setuptools

from torch.utils import cpp_extension

def main() -> None:
    setuptools.setup(
        name='ext_vecadd',
        version='1.0.0',
        author='Your name',
        author_email='your@email.com',
        packages=setuptools.find_packages(),
        ext_modules=[cpp_extension.CppExtension(
            name='vecadd_cuda',
            sources=[path.join('cuda', 'vecadd.cpp')],
            libraries=[
                'vecadd_kernel',
            ],
            library_dirs=[path.join('.', 'cuda')],
            extra_compile_args=['-g', '-fPIC'],
        )],
        cmdclass={'build_ext': cpp_extension.BuildExtension},
    )
    return

if __name__ == '__main__':
    main()

당연한 소리지만 version, author, author_email은 취향대로.

 

 

여기서 중요하게 볼 부분은 setup 함수의 ext_modules 부분이다.

 

링크에서 공식 documentation을 보면 CppExtension 말고도 CUDAExtension이라는 함수도 존재하는데,

 

막상 써보니까 잘 동작하지 않았고 돌아돌아 결론을 내린게 CppExtension이었다.

 

CUDAExtension을 쓰면 별도의 Makefile로 shared object (.so)를 빌드하지 않아도 될 것 같은데,

 

잘 안되니까 나중에 다시 해봐야겠다.

 

CppExtension의 name 인수에는 나중에 Python에서 해당 모듈을 불러올때 사용하는 이름을 지정한다.

 

예를 들어, 위의 setup.py로 빌드한 모듈은 아래와 같이 불러올 수 있다.

import torch
import vecadd_cuda

반드시 Line 1이 선행되어야 하는 점에 주의하면, 큰 문제는 없다.

  • sources 인수는 빌드할 cpp 파일(들)을 지정한다.
    이미 vecadd.cpp에서 CUDA 커널을 호출하도록 되어있기 때문에, .cu 파일을 포함할 필요는 없다.
  • 대신에 우리가 빌드한 .so 파일을 libraries라는 인수에 포함시키는데,
    vecadd_kernel을 지정하면 자동으로 libvecadd_kernel.so가 빌드에 포함된다고 생각하면 된다.
  • library_dirs에는 해당 .so 파일이 포함되어 있는 폴더를 지정해주고,
  • extra_complie_args는 -g와 -fPIC를 넣어주면 완성.

이제 터미널에서 아래와 같은 커맨드로 전체 프로젝트를 빌드한다.

$ python setup.py install

여기서 주의할 점은, conda 등의 가상 환경을 사용 중인 경우 해당 environment에만 패키지가 설치된다는 것.

 

뭔가 더 해야할 것 같은 느낌이 들지만 의외로 지금 당장 작성한 CUDA 커널을 PyTorch (Python)에서 사용 가능하다.

 

우선 Python을 실행시키고, 아래와 같이 하면 테스트 가능.

>>> import torch
>>> import vecadd_cuda
>>> x = torch.randn(4, 4, device=torch.device('cuda'))
>>> y = torch.randn(4, 4, device=torch.device('cuda'))
>>> z = torch.zeros(4, 4, device=torch.device('cuda'))
>>> print(x)
tensor([[ 0.7651,  0.2747, -0.3275, -1.0251],
        [-0.5340, -0.8416,  0.2952,  0.7431],
        [-0.2698, -0.3798, -0.0225,  0.0245],
        [ 1.2056,  1.9831,  0.1669, -0.2347]], device='cuda:0')
>>> print(y)
tensor([[-1.0483,  0.6131, -0.2479,  0.3100],
        [ 0.6286,  0.4175, -0.4093, -1.5263],
        [ 0.6393,  1.3988, -0.0907,  0.1472],
        [-1.0677,  1.0088, -1.4845, -0.7323]], device='cuda:0')
>>> print(z)
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]], device='cuda:0')
>>> vecadd_cuda.vecadd(x, y, z)
>>> print(z)
tensor([[-0.2832,  0.8878, -0.5754, -0.7151],
        [ 0.0947, -0.4241, -0.1141, -0.7832],
        [ 0.3695,  1.0189, -0.1133,  0.1716],
        [ 0.1379,  2.9919, -1.3177, -0.9670]], device='cuda:0')
>>> print(x + y)
tensor([[-0.2832,  0.8878, -0.5754, -0.7151],
        [ 0.0947, -0.4241, -0.1141, -0.7832],
        [ 0.3695,  1.0189, -0.1133,  0.1716],
        [ 0.1379,  2.9919, -1.3177, -0.9670]], device='cuda:0')

random 수치들은 테스트할 때마다 다르게 나오겠지만, 아무튼 z = x + y가 잘 작동한다.

 

코드를 잘(?) 만들어놔서인지, 2D 텐서를 처리하는 데에도 큰 문제가 없다.

 

 

그런데 현재 상태로는 CUDA 커널이 수정될때마다 Makefile 후 setup.py를 install 하는 두 단계를 거쳐야 해서 조금 번거롭다.

 

기왕 하는거 두 개를 한 번에 하는 스크립트도 작성해놨다.

 

아래 내용을 setup.py가 있는 폴더에 Makefile 파일을 만들고 붙여넣으면 된다.

SUBDIRS = cuda

all: vecadd
    python setup.py install

vecadd:
    $(MAKE) -C cuda

clean:
    rm -rf build dist ./*.egg-info
    for dir in $(SUBDIRS); do $(MAKE) -C $$dir Makefile $@; done

역시나 야매지만 $(MAKE)와 -C라는 문법을 통해 sub directory의 Makefile을 실행하는 형식인가 보다.

 

clean의 경우 저것 말고도 여러 가지 방법이 있는데 내가 보기에 깔끔한 걸로 골랐다.

 

아무튼 이제 setup.py를 실행시키는 대신에, Makefile을 한 번만 실행하면 된다!

$ make

여기서 끝?

 

이면 참 좋겠지만, 몇 가지 불편한 사실들이 남아있다.

  1. z가 자동으로 할당되지 않는다.
    보통 우리는 x와 y를 가지고 있을 때, x + y가 z를 return 하기를 원하지
    z를 미리 할당하고 x + y를 채우는 것을 별로 선호하지 않는다.
  2. 예외 처리가 되어있지 않다.
    작성한 커널은 CUDA 텐서에 대해서만 작동하지만 CPU 텐서가 들어와도 아무 경고가 없다 (내가 가장 삽질한 부분).
    x + y가 정상적으로 계산되지 않는 것은 당연하고.
  3. backward()가 작동하지 않는다. = backpropagation이 불가능하다.
    사실 가장 중요한 부분이다.
    달리 말하면, custom operation으로의 기능은 하지만 custom layer로의 기능은 전혀 할 수가 없다.

1과 2는 그러려니 해도, 3은 조금 치명적이다.

 

조금 불행한 소식은 3 또한 CUDA 커널을 직접 작성해서 구현해야 한다는 점,

 

그나마 다행인 소식은 CUDA 커널만 추가로 작성하고

 

Makefile이나 vecadd.cpp, setup.py를 크게 고칠 필요는 없다는 점이다.

 

또한 backward가 사용 가능하게 하려면 autograd.Function을 상속받아야 하는데, 이를 통해 1과 2를 함께 해결 가능하다.

 

 

이전에 backward를 미리 작성해놓지 않은 이유는 사실 간단한데,

 

이 단계에서 구현한 연산이 정확히 동작하지 않는다면 backward 또한 제대로 동작 할리 만무하기 때문이다.

 

그러니까 일단 CUDA 커널을 제대로 작성했는지 PyTorch (Python) 인터페이스를 사용하여 편하게 확인하고,

 

문제가 있다면 고쳐서 완벽한 동작을 보장한 다음 나머지 잡일 (좀 많지만)을 처리하러 가면 된다.

 

다음 글:

2021.05.01 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (4) - autograd.Function과 backward 구현

 

파이토치 PyTorch CUDA custom layer 만들기 (4) - autograd.Function과 backward 구현

이전 글: 2021.04.28 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (1) - CUDA 커널 2021.04.30 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (2) - 커널 빌드 및 Python 바..

sanghyun.tistory.com

 

이전 글:

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

 

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

종종 기발한 아이디어를 떠올리다 보면, "이걸 구현하기 위한 레이어가 파이토치에 있었나?" 라는 의구심이 든다. 다행히 세상이 좋아져서 대부분의 멋진 아이디어들은 쓰기 좋고 깔끔한 파이토

sanghyun.tistory.com

기본적인 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

이전 글: 2021.04.28 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (1) - CUDA 커널 2021.04.30 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (2) - 커널 빌드 및 Python 바..

sanghyun.tistory.com

 

종종 기발한 아이디어를 떠올리다 보면,

 

"이걸 구현하기 위한 레이어가 파이토치에 있었나?" 라는 의구심이 든다.

 

다행히 세상이 좋아져서 대부분의 멋진 아이디어들은 쓰기 좋고 깔끔한 파이토치 구현체가 있지만,

 

정 안 되는 경우 내가 직접 만들어야 하는 경우가 생긴다.

 

파이토치의 custom layer는 만들기 어려운 편은 아니지만,

 

효율 등의 문제로 CUDA 프로그래밍이 필요하다면 문제가 많이 복잡해진다.

 

지난 CVPR을 준비하며 CUDA 프로그래밍을 동반한 custom layer를 만들어보았는데,

 

언젠가 또 비슷한 걸 만들지 모르니 정리해두는 게 좋겠다.

 

 

Toy example로 vector x와 y를 더하는 vecadd의 CUDA custom layer를 만들어보자.

 

(즉, z = x + y)

 

vecadd를 동작하게 만들 수 있다면, 훨씬 더 복잡한 연산도 문제없이 구현이 가능하다.

 

물론 연산 자체를 C++로 어떻게 짜야하는지 모르겠다면 그건 어쩔 수 없지만...

 

 

아무튼 전체 프로젝트는 아래 다이어그램처럼 구성된다.

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

뭐가 이렇게 많나 생각이 들지만, 구조를 이해하면 나름 합리적이다.

 

(그렇다고 내가 각각의 디테일을 다 이해한 것은 아니다.)

 

 

가장 먼저 하면 좋은 것은 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 - block - thread hierarchy.

쉽게 말해 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

 

+ Recent posts