CPP, CUDA 扩展 PyTorch

• 准备工作
• 用 C++ 扩展 PyTorch
• 用 CUDA 扩展 PyTorch
• 参考

请先确定一下 PyTorch 版本：

>>> torch.__version__
‘1.9.1’
>>> torch.version.cuda
‘11.1’
>>> torch.backends.cudnn.version()
8005

实验环境：

官网上案例使用的算子大多和 python 的接口差不多，而有的时候我们需要操作 Tensor 的具体元素，这时候就要用到 Using accessors 那一小节介绍的技术

这里我们在人尽皆知的快速排序上进行实验：

首先是我们的源程序：

quickSort.cpp

#include 

void quickSort(torch::Tensor& src, int begin, int end) {
        auto src_a = src.accessor();
        if (begin < end) {
                auto key = src_a[begin];
                int i = begin;
                int j = end;
                while (i < j) {
                        while (i < j && src_a[j] > key)
                                --j;
                        if (i < j) {
                                src_a[i] = src_a[j];
                                ++i;
                        }
                        while (i < j && src_a[i] < key)
                                ++i;
                        if (i < j) {
                                src_a[j] = src_a[i];
                                --j;
                        }
                }
                src_a[i] = key;
                quickSort(src, begin, i-1);
                quickSort(src, i+1, end);
        }
}

void QuickSort(torch::Tensor& src) {
        quickSort(src, 0, src.size(0));
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
        m.def("quick_sort", &QuickSort, "cpp implementation for quick sort");
}

头文件是写 cpp 扩展所必须的，它大概包括三部分：

1. Aten 库，它是 PyTorch 用于张量计算的主要 API
2. pybind11，它是 PyTorch 为 C++ 代码创建 Python 绑定的方式
3. 一些头文件，负责管理 ATen 和 pybind11 之间的交互细节

之后重点是我们的 .accessor()
它为我们提供了一种像数组一样操作 torch::Tensor 的方法，这里我们指定 src 为一维 int 数组
然后就可以用我们熟悉的 c++ 代码去书写快速排序的算法了，具体细节这里就不介绍了

最后是 pybind11 的语法，我们用 m.def(…) 在该模块内定义了名为 quick_sort ，实际调用了 QuickSort 的方法

接着在同级目录下还需要一个 setup.py 文件：
setup.py

from setuptools import setup, Extension
from torch.utils import cpp_extension

setup(name='quickSort_cpp',
        ext_modules=[cpp_extension.CppExtension('quickSort_cpp', ['quickSort.cpp'])],
        cmdclass={'build_ext': cpp_extension.BuildExtension})

这基本上就是套用模板，注意模块名为 quickSort_cpp ，对应的 cpp 源文件就是上面的 quickSort.cpp

这是最后的目录结构：

在这里调用：

python setup.py install

等待片刻，出现：

就可以进行测试了：

只要注意类型相互对应，基本上不会有什么问题

为了利用 GPU 并行计算的优势，更多的时候我们需要为自己的扩展编写 CUDA 版本的代码

下面我们用 Hadamard 乘积来演示用 CUDA 扩展 PyTorch：

其实和 C++ 扩展 PyTorch 的方法差不多，我们只需要让 C++ 程序作为中介，即可联系起 PyTorch 和 CUDA

先用之前提到的 accessors 来写：

hadamard_cuda.cpp

#include 

// CUDA declarations

void hadamard_cuda_kernel(torch::Tensor A, torch::Tensor B, torch::Tensor C);

// C++ interface

void hadamard_cuda(torch::Tensor A, torch::Tensor B, torch::Tensor C) {
        hadamard_cuda_kernel(A, B, C);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
        m.def("hadamard_cuda", &hadamard_cuda, "hadamard product (CUDA)");
}

我们用 C++ 程序封装起 CUDA 的调用：

hadamard_cuda_kernel.cu

#include 

template 
__global__ void _hadamard(
                const torch::template PackedTensorAccessor32 A,
                const torch::template PackedTensorAccessor32 B,
                torch::template PackedTensorAccessor32 C,
                int m, int n) {
        int i = blockIdx.x * blockDim.x + threadIdx.x;
        int j = blockIdx.y * blockDim.y + threadIdx.y;

        if (i < m && j < n)
                C[i][j] = A[i][j] * B[i][j];
}

void hadamard_cuda_kernel(torch::Tensor A, torch::Tensor B, torch::Tensor C) {
        dim3 threadsPerBlock(16, 16);
        dim3 numBlocks(A.size(0) / threadsPerBlock.x, A.size(1) / threadsPerBlock.y);

        AT_DISPATCH_ALL_TYPES(A.scalar_type(), "hadamard_product_cuda", ([&] {
                                _hadamard<<>>(
                                                A.packed_accessor32(),
                                                B.packed_accessor32(),
                                                C.packed_accessor32(),
                                                A.size(0), A.size(1));
                                }));
}

这里我们只是简单地实现一下，还有很多优化方法暂且不管

然后是我们的 setup.py，这个和之前的差不多：
setup.py

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

setup(
    name='hadamard_cuda',
    ext_modules=[
        CUDAExtension('hadamard_cuda', [
            'hadamard_cuda.cpp',
            'hadamard_cuda_kernel.cu',
        ])
    ],
    cmdclass={
        'build_ext': BuildExtension
    })

为了方便之后的测试，我们再写一个：
test.py

import torch
import hadamard_cuda

print("cuda:", torch.cuda.is_available())

device = torch.device('cuda')

A = torch.randint(low=0, high=10, size=(16, 16), device=device)
B = torch.randint(low=0, high=10, size=(16, 16), device=device)

print(A[0:5, 0:5])
print(B[0:5, 0:5])

C = torch.randint(low=0, high=10, size=(16, 16), device=device)

hadamard_cuda.hadamard_cuda(A, B, C)
print("After hadamard (CUDA)")
print(C[0:5, 0:5])

下面是最后的目录结构：

在当前目录下运行：

python setup.py install

等待片刻，直至出现：

接着我们就可以进行测试了：

python test.py

结果：

注意类型的统一

这是在较新版本的 PyTorch 上使用的方法，然而，很多时候我们还会看到不少 c 风格的 CUDA 程序，下面介绍和这些程序交互的方法：

假如，我们遇到下面这样的 CUDA 程序：
foo_cuda_kernel.cu

#include 
#include 

__global__ void _foo(int* A, int* B, int* C, int m, int n) {
        
}

void foo_cuda_kernel(int* A, int* B, int* C, int m, int n) {
        dim3 threadsPerBlock(16, 16);
        dim3 numBlocks(m / threadsPerBlock.x, n / threadsPerBlock.y);
        _foo<<>>(A, B, C, m, n);
}

我们可以将封装所用的 C++ 程序重新改写成下面的样子：
foo_cuda.cpp

#include 

// CUDA declarations

void foo_cuda_kernel(int* A, int* B, int* C, int m, int n);

// C++ interface

void foo_cuda(torch::Tensor A, torch::Tensor B, torch::Tensor C) {
        foo_cuda_kernel((int *) A.data_ptr(),
                        (int *) B.data_ptr(),
                        (int *) C.data_ptr(),
                         A.size(0), A.size(1));
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
        m.def("foo_cuda", &foo_cuda, "(CUDA)");
}

PyTorch：Custom C++ and CUDA Extensions