.cu 코드의 pybinding

Pybind11 은 주로 C++ 코드를 Python 에 바인딩하기 위한 라이브러리지만

CUDA 코드와 C++ 코드를 함께 사용하여 Python 과 상호작용할 수 있도록 만들 수 있다.

Pybind11 은 CUDA 자체를 직접 지원하지 않지만, C++ 코드를 통해 CUDA 커널을 호출할 수 있는 방식을 제공한다. 이를 통해 CUDA 코드를 Python 에서 호출하고 사용할 수 있다.

 

CUDA 커널 작성

// example.cu
#include <cuda_runtime.h>

__global__ void add_kernel(float* a, float* b, float* c, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        c[idx] = a[idx] + b[idx];
    }
}

void add(float* a, float* b, float* c, int size) {
    float *d_a, *d_b, *d_c;
    
    // 메모리 할당
    cudaMalloc((void**)&d_a, size * sizeof(float));
    cudaMalloc((void**)&d_b, size * sizeof(float));
    cudaMalloc((void**)&d_c, size * sizeof(float));
    
    // 데이터 복사
    cudaMemcpy(d_a, a, size * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_b, b, size * sizeof(float), cudaMemcpyHostToDevice);
    
    // CUDA 커널 실행
    int blockSize = 256;
    int numBlocks = (size + blockSize - 1) / blockSize;
    add_kernel<<<numBlocks, blockSize>>>(d_a, d_b, d_c, size);
    
    // 결과를 호스트로 복사
    cudaMemcpy(c, d_c, size * sizeof(float), cudaMemcpyDeviceToHost);
    
    // 메모리 해제
    cudaFree(d_a);
    cudaFree(d_b);
    cudaFree(d_c);
}

C++ 파일에서 CUDA 커널 호출, Pybind11 바인딩

// example.cpp
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include "example.cu"  // CUDA 코드 포함

namespace py = pybind11;

void add(py::array_t<float> a, py::array_t<float> b, py::array_t<float> c) {
    py::buffer_info a_info = a.request();
    py::buffer_info b_info = b.request();
    py::buffer_info c_info = c.request();

    float* a_ptr = static_cast<float*>(a_info.ptr);
    float* b_ptr = static_cast<float*>(b_info.ptr);
    float* c_ptr = static_cast<float*>(c_info.ptr);
    int size = a_info.size;

    // CUDA 호출
    add(a_ptr, b_ptr, c_ptr, size);
}

PYBIND11_MODULE(example, m) {
    m.def("add", &add, "A function that adds two arrays using CUDA");
}

example.cu 에서 정의한 CUDA 커널을 add 함수에서 호출,

 

컴파일 이후 Python 에서 사용

import numpy as np
import example  # 컴파일된 Pybind11 모듈

a = np.random.rand(1000000).astype(np.float32)
b = np.random.rand(1000000).astype(np.float32)
c = np.zeros_like(a)

# CUDA 커널을 사용하여 배열 더하기
example.add(a, b, c)

print(c)  # a + b 결과 출력

 

세 단계를 통해 CUDA 파일의  실행

 

C++ 에서 .cu 파일의 함수를 호출할 때는 별도의 .exe 파일이 필요하지 않다. 대신 .cu 파일을 C++ 코드에 통합하여 함께 컴파일하는 방식으로 처리한다. 

이 둘을 함께 빌드하여 하나의 실행 파일을 생성하거나 라이브러리로 만들어야 한다.