Pytorch和CUDA联合编程的基本步骤

参考链接-w3school

参考链接-CSDN

本文代码-github

背景

目前PyTorch已经提供了丰富的接口可以直接调用，但是仍存在一些高度自定义的操作无法使用PyToch或者Python高效的完成，因此PyTorch还提供了使用C++和CUDA编程的扩展接口。C++扩展主要有两种形式，一种是使用setuptools提前构建，也可以通过torch.utils.cpp_extension.load()在运行时构建。下面仅介绍第一种方法，第二种方法之后再学习。

基本步骤

Pytorch,CUDA,C++联合编程的一般步骤如下:

1. 首先需要定义一个C++文件，该文件声明了CUDA文件中定义的函数，还需要进行一些检查，并最终将其调用转发给.cu文件。此外，该文件还需要声明将在Python中调用的函数，并使用pybind11绑定到python。OpenPCDet将上述步骤划分为以下几个步骤:

   1. 首先定义一个头文件，该头文件.h中包含了.cu文件中定义的函数和.cpp文件中定义的函数.

   2. 然后定义.cpp文件，其中函数的作用是负责进行一些检查和调用.cu文件中定义的函数

   3. .cu文件是负责执行具体的CUDA编程的操作

   4. api文件是将.cpp文件中定义的函数和PYBIND11绑定，以便Python调用

2. 在setup.py文件中声明将要编译的模块名称，源文件路径等。
3. 使用import导入声明的模块，使用Python实现其前向和反向传播的计算。

举例说明

.
├── ball_query_src
│   ├── api.cpp
│   ├── ball_query.cpp
│   ├── ball_query_cuda.cu
│   ├── ball_query_cuda.h
│   └── cuda_utils.h
├── setup.py
└── test_ball_query.py
复制代码

项目的目录如上图所示，其中api.cpp文件是将ball_query.cpp声明的函数使用PYBIND11与python进行绑定。其中api.cpp的内容如下:

#include <torch/serialize/tensor.h>
#include <torch/extension.h>

#include "ball_query_cuda.h"


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    // 第一个参数表示的是在python中调用的名称，第二个参数是对应的cpp函数，第三个参数对应的是这个函数的说明
    m.def("ball_query_wrapper", &ball_query_wrapper_fast, "ball_query_wrapper_fast");
}

复制代码

这里，作者为了使得代码的结构更加清晰，其在ball_query.h文件中分别声明了两个函数，一个是在C++中被调用的函数，另一个是在CUDA中实现的函数。其具体内容如下:

#ifndef _BALL_QUERY_GPU_H
#define _BALL_QUERY_GPU_H

#include <torch/serialize/tensor.h>
#include <vector>
#include <cuda.h>
#include <cuda_runtime_api.h>

// 与pybind11绑定的函数，其主要作用是调用下面的cuda函数
int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample, 
	at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor);

// CUDA文件中的函数
void ball_query_kernel_launcher_fast(int b, int n, int m, float radius, int nsample, 
	const float *xyz, const float *new_xyz, int *idx);

#endif
复制代码

在.h文件中声明了上述两个文件之后，再分别再ball_query.cpp和ball_query_cuda.cu文件中完成这两个函数的具体实现。

#include <torch/serialize/tensor.h>
#include <vector>
#include <THC/THC.h>
#include <cuda.h>
#include <cuda_runtime_api.h>

#include "ball_query_cuda.h"

extern THCState *state;

// 定义检查数据类型的宏
#define CHECK_CUDA(x) do { \
	  if (!x.type().is_cuda()) { \
		      fprintf(stderr, "%s must be CUDA tensor at %s:%d\n", #x, __FILE__, __LINE__); \
		      exit(-1); \
		    } \
} while (0)
#define CHECK_CONTIGUOUS(x) do { \
	  if (!x.is_contiguous()) { \
		      fprintf(stderr, "%s must be contiguous tensor at %s:%d\n", #x, __FILE__, __LINE__); \
		      exit(-1); \
		    } \
} while (0)
#define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)

// 完成输入数据类型的检查，同时调用cu文件中定义的函数
int ball_query_wrapper_fast(int b, int n, int m, float radius, int nsample, 
    at::Tensor new_xyz_tensor, at::Tensor xyz_tensor, at::Tensor idx_tensor) {
    CHECK_INPUT(new_xyz_tensor);
    CHECK_INPUT(xyz_tensor);
    const float *new_xyz = new_xyz_tensor.data<float>();
    const float *xyz = xyz_tensor.data<float>();
    int *idx = idx_tensor.data<int>();
    
    ball_query_kernel_launcher_fast(b, n, m, radius, nsample, new_xyz, xyz, idx);
    return 1;
}
复制代码

ball_query_cuda.cu实现

#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#include "ball_query_cuda.h"
#include "cuda_utils.h"


__global__ void ball_query_kernel_fast(int b, int n, int m, float radius, int nsample, 
    const float *__restrict__ new_xyz, const float *__restrict__ xyz, int *__restrict__ idx) {
    // new_xyz: (B, M, 3)
    // xyz: (B, N, 3)
    // output:
    //      idx: (B, M, nsample)
    int bs_idx = blockIdx.y;
    int pt_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (bs_idx >= b || pt_idx >= m) return;

    new_xyz += bs_idx * m * 3 + pt_idx * 3;
    xyz += bs_idx * n * 3;
    idx += bs_idx * m * nsample + pt_idx * nsample;

    float radius2 = radius * radius;
    float new_x = new_xyz[0];
    float new_y = new_xyz[1];
    float new_z = new_xyz[2];

    int cnt = 0;
    for (int k = 0; k < n; ++k) {
        float x = xyz[k * 3 + 0];
        float y = xyz[k * 3 + 1];
        float z = xyz[k * 3 + 2];
        float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z);
        if (d2 < radius2){
            if (cnt == 0){
                for (int l = 0; l < nsample; ++l) {
                    idx[l] = k;
                }
            }
            idx[cnt] = k;
            ++cnt;
            if (cnt >= nsample) break;
        }
    }
}


void ball_query_kernel_launcher_fast(int b, int n, int m, float radius, int nsample, \
    const float *new_xyz, const float *xyz, int *idx) {
    // new_xyz: (B, M, 3)
    // xyz: (B, N, 3)
    // output:
    //      idx: (B, M, nsample)

    cudaError_t err;

    dim3 blocks(DIVUP(m, THREADS_PER_BLOCK), b);  // blockIdx.x(col), blockIdx.y(row)
    dim3 threads(THREADS_PER_BLOCK);

    ball_query_kernel_fast<<<blocks, threads>>>(b, n, m, radius, nsample, new_xyz, xyz, idx);
    // cudaDeviceSynchronize();  // for using printf in kernel function
    err = cudaGetLastError();
    if (cudaSuccess != err) {
        fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
        exit(-1);
    }
}
复制代码

至此，ball_query的核心功能已经完成，然后我们需要使用setup.py文件来编译上述文件。setup.py文件的具体实现如下:

import os
import subprocess

from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension


def make_cuda_ext(name, module, sources):
    cuda_ext = CUDAExtension(
        name='%s.%s' % (module, name),
        sources=[os.path.join(module.split('.')[-1], src) for src in sources]
    )
    print([os.path.join(*module.split('.'), src) for src in sources])
    return cuda_ext

if __name__ == '__main__':
    setup(
        name='ballquery',
        packages=find_packages(),
        ext_modules=[
            CUDAExtension('ball_query_cuda',[
                'ball_query_src/api.cpp',
                'ball_query_src/ball_query.cpp',
                'ball_query_src/ball_query_cuda.cu',  
            ])
        ],
        cmdclass={
            'build_ext': BuildExtension
        }
    )
复制代码

至此，我们已经生成了上述代码的链接库，但是如果需要将其嵌入到神经网络中，还需要定义其前向传播和反向传播方法。这里我们在test_ball_query.py文件中完成其前向传播和反向传播。

import torch
import torch.nn as nn
from torch.autograd import Function, Variable
import math

import ball_query_cuda

# 定义该方法的前向传播和反向传播方法
class BallQuery(Function):
    
    @staticmethod
    def forward(ctx, radius: float, nsample: int, xyz: torch.Tensor, new_xyz: torch.Tensor) -> torch.Tensor:
        """
        :param ctx:
        :param radius: float, radius of the balls
        :param nsample: int, maximum number of features in the balls
        :param xyz: (B, N, 3) xyz coordinates of the features
        :param new_xyz: (B, npoint, 3) centers of the ball query
        :return:
            idx: (B, npoint, nsample) tensor with the indicies of the features that form the query balls
        """
        assert new_xyz.is_contiguous()
        assert xyz.is_contiguous()

        B, N, _ = xyz.size()
        npoint = new_xyz.size(1)
        idx = torch.cuda.IntTensor(B, npoint, nsample).zero_()
        
        ball_query_cuda.ball_query_wrapper(B, N, npoint, radius, nsample, new_xyz, xyz, idx)
        return idx

    @staticmethod
    def backward(ctx, a=None):
        return None, None, None, None


ball_query = BallQuery.apply

xyz = torch.randn(2, 128, 3).cuda()
new_xyz = xyz

result = ball_query(0.8, 3, xyz, new_xyz)

print(result.shape)
复制代码