pycuda 개념 및 사용 예시

PyCUDA란?

PyCUDA는 파이썬에서 CUDA를 사용하기 위한 라이브러리

CUDA는 NVIDIA 그래픽 카드에서 병렬처리를 위해 사용되는 플랫폼으로 PyCUDA는 파이썬 개발자가 CUDA C/C++ 코드를 직접 작성하지 않고도 파이썬에서 CUDA를 사용할 수 있게 해준다. 이를 통해 파이썬에서 고성능 계산을 수행할 수 있다.

• DOCS

pycuda 2022.2.2 documentation

 

1. 초기화

Pycuda 를 사용하기 전에 import하면서 autoinit으로 자동으로 초기화 할 수 있다.

autuinit을 import하지 않으면 오류가 나고, 이 과정은 수동으로도 진행할 수 있다.

import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule

 

2. 데이터 전송

cpu에 있는 데이터를 병렬처리하기 위해서는 gpu로 보내야한다. 

 

import numpy as np

# create a numpy array
a = np.random.rand(4,4).astype(np.float32)

# create allocate memory on the device
a_gpu = cuda.mem_alloc(a.nbytes)

# transfer the data to the GPU
cuda.memcpy_htod(a_gpu, a)

numpy array a를 gpu로 보낼 건데 먼저 a의 사이즈만큼의 메모리를 gpu에 mem_alloc으로 할당한다.

a.nbytes는 a의 사이즈를 리턴하고, mem_alloc은 그 사이즈만큼 그래픽카드에 메모리를 할당하고 디바이스에 병렬 메모리를 나타내는 DeviceAllocation 객체 (a_gpu)를 리턴한다.

 

그리고 memcpy_htod를 이용해서 실제 데이터 a를 a_gpu 로 옮긴다.

memcpy_htod(dest, src)

 

3. 커널 실행

pycuda를 이용해서 c언어로 작성된 함수를 그대로 파이썬에서 이용할 수 있다.

mod = SourceModule("""
  __global__ void doublify(float *a)
  {
    int idx = threadIdx.x + threadIdx.y*4;
    a[idx] *= 2;
  }
  """)

SourceModule에 작성된 코드는 C언어로 작성된 코드로 배열의 값을 두 배로 늘리는 내용의 코드이다. 

 

func = mod.get_function("doublify")
func(a_gpu, block=(4,4,1))

a_doubled = np.empty_like(a)
cuda.memcpy_dtoh(a_doubled, a_gpu)
print(a_doubled)
print(a)

[[1.617366   0.9470439  0.6003214  0.78253263]
 [0.7782898  0.24809593 0.15962741 0.49050307]
 [0.6258923  1.8083767  0.9584732  1.4263026 ]
 [0.38570014 0.20260021 0.03355372 0.24119434]]
[[0.808683   0.47352195 0.3001607  0.39126632]
 [0.3891449  0.12404796 0.0798137  0.24525154]
 [0.31294614 0.90418833 0.4792366  0.7131513 ]
 [0.19285007 0.10130011 0.01677686 0.12059717]]

이렇게 정상적으로 배열의 값이 2배가 된 것을 확인할 수 있다.

 

파이쿠다에 대해 이정도 개념을 알면 tensorrt에서 모델 역직렬화후 gpu에서 데이터 추론을 진행하는 과정을 이해할 수 있다.