Pytorch-c++-CUDA - 技术栈

一、CUDA

1.1 简介：

CUDA（Compute Unified Device Architecture）是 NVIDIA 推出的并行计算平台和编程模型，它让开发者能够利用 GPU 中成百上千个核心进行通用计算，而不仅仅是图形渲染，大幅加速科学计算、深度学习、图像处理等数据密集型任务。

1.2 cuda的优势

高并行吞吐量：数千个线程同时运行，适合处理大规模数据并行问题。
细粒度的内存层次：可利用共享内存、寄存器等实现极高的访存带宽。
与 C/C++ 紧密集成：开发者可快速上手，将计算密集型部分卸载到 GPU 执行，显著缩短运算时间

1.3 自定义cuda的场景：

非平行计算(如体积渲染)，大量串列计算。

当应用的计算模式无法被现有库或简单并行化覆盖时，就需要编写自定义 CUDA 核函数。典型场景如体积渲染：虽然每条光线内部包含大量串行步进和采样计算，但不同像素、不同光线之间却高度独立。通过自定义 CUDA 核函数，可以将这种"每个元素内部串行、元素之间并行"的任务在 GPU 上高效执行，将数以千计的串列计算过程同时推进，从而获得远高于纯 CPU 串行计算的性能。

二：打通Pyton调用C++文件

2.1 使用VSCode配置环境桥梁C++环境

新建C++文件，ctrl + shift + p 弹出选项卡点击Edit Configurations UI配置如下图所示：
Edit Configuration

点击后退出此时会发现文件夹中多了一个.vscode文件夹，里面有c_cpp_properties.json文件，打开后对includePath进行编辑，根据自己的创建的python环境来定，可以在你环境文件夹的目录下问gpt，让他帮你找。我的环境名是GSRaw。

复制代码

{
    "configurations": [
        {
            "name": "Win32",
            // C++项目中编译时头文件搜索目录
            "includePath": [
                "${workspaceFolder}/**",
                // 提供 CUDA 运行时 API
                "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.8/include",
                // py文件环境下的通用 Python/C 扩展头文件，确保能与 Python 解释器交互
                "D:/biancheng/Python/Anaconda3_envs_dirs/GSRaw/include",
                // PyTorch C++ 核心库（LibTorch）头文件
                "D:/biancheng/Python/Anaconda3_envs_dirs/GSRaw/Lib/site-packages/torch/include",
                // PyTorch C++ 前端 API，提供高层模块化接口，方便构建训练/推理模型
                "D:/biancheng/Python/Anaconda3_envs_dirs/GSRaw/Lib/site-packages/torch/include/torch/csrc/api/include"
            ],
            // 相当于字典用来排除单词下面波浪线的，可不管
            "defines": [
                "_DEBUG",
                "UNICODE",
                "_UNICODE",
                "__CUDACC__"
            ]
        }
    ],
    "version": 4
}

2.2 编写桥梁的C++文件

这里只是示例：用于python调用。文件名为：interpolation.cpp

复制代码

# include <torch/extension.h>

// 根据个人使用cuda的需求自定义的函数
torch::Tensor trilinear_interpolation(
    torch::Tensor feats,
    torch::Tensor point
){
    return feats;
}

// 固定结构搭配
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    // 别名（setup.py文件加入环境后py文件的可以引用的包名） ，上面自定义的函数名
    m.def("trilinear_interpolation", &trilinear_interpolation);
}

2.3 打包C++扩展供Python调用

创建并编写setup.py文件，然后ctrl + shift + ~ 打开终端命令行(最好默认的改为cmd)。

复制代码

from setuptools import setup
from torch.utils.cpp_extension import CppExtension, BuildExtension
import platform

# 未避免一些问题，确保C++桥梁文件按 C++17 编译
extra_compile_args = ['/std:c++17', '/EHsc'] if platform.system() == 'Windows' else ['-std=c++17']

setup(
    name='cppcuda_tutorial',
    version='1.0',
    author='kwea123',
    author_email='kwea123@gmail.com',
    description='cppcuda example',
    # cppcuda_tutorial是前面name属性值，sources(可填多个值)是C++中自定义的函数的别名(一般设置和自定义函数名相同)
    ext_modules=[CppExtension('cppcuda_tutorial', sources=['interpolation.cpp'], extra_compile_args=extra_compile_args)],
    cmdclass={'build_ext': BuildExtension}
)

在python文件中打开或者conda activate 环境名，切换到c_cpp_properties.json文件设置引用的python环境中。检查并将 pip更新到最新版本。然后在项目文件夹下如下pip命令。

注：我使用下面命令就提示我环境中没有torch包，但是我有就只是使用禁止隔离的pip命令。

复制代码

pip install --no-build-isolation .
或者pip install .

出现Successfully installed cppcuda-tutorial-1.0说明安装成功了，按我配置应该是没问题的，因为我是换了个新环境重新配置的。（如果从0新建的文件记得安装cuda(我的是11.8)和相应版本的torch库）。

2.4 创建python文件调用自定义的C++函数

复制代码

# 一定要先引入torch库再引用自定义的c++函数库
import torch
import cppcuda_tutorial

feats = torch.randint(0,10,size = (3,4))
point = torch.zeros(4,3)
output = cppcuda_tutorial.trilinear_interpolation(feats, point)
print(output)

注：

cppcuda_tutorial库名如下图
库名

调用函数解析如下图，实际调用的函数名也就是setup文件中的sources里面等同

调用函数解析

三、连通CUDA

3.1 编写cu文件，实现具体功能

这里依旧是示例，这里先返回一下参数。

复制代码

#include <torch/extension.h>

torch::Tensor trilinear_fw_cu(
    torch::Tensor feats,
    torch::Tensor points
){
    return feats;
}

3.2 配置自定义C++库和修改桥梁

新建文件夹include,新建.h后缀文件，我的是utils.h

复制代码

#include <torch/extension.h>
// 以下宏用于检查张量是否在 CUDA 上、是否连续
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

// 下面是cu文件中的函数，用以声明方便c++引用
torch::Tensor trilinear_fw_cu(
    torch::Tensor feats,
    torch::Tensor points
);

修改桥梁，引入自定义的库，在函数内部检查变量是否来自GPU，否拒绝是通过，然后返回cu文件中需要的函数的结果。

复制代码

# include <torch/extension.h>
# include "utils.h"

// 根据个人使用cuda的需求自定义的函数，充当桥梁转接的
torch::Tensor trilinear_interpolation(
    torch::Tensor feats,
    torch::Tensor point
){
    CHECK_INPUT(feats);
    CHECK_INPUT(point);

    // 调用前向传播的cuda函数
    return trilinear_fw_cu(feats,point);
}

// 固定结构搭配
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    // 别名（setup.py文件加入环境后py文件的可以调用的函数名） ，自定义的函数名
    m.def("trilinear_interpolation", &trilinear_interpolation);
}

注：这里那两个CHECK可能报错还有cu文件函数未识别，没关系的，setup文件还是可以正常编译的。如下图，只能说是一个大无语。

3.3 修改setup并重新编译

还是在禁用隔离的情况下编译 pip install --no-build-isolation .

注意里面从CppExtension改为CUDAExtension，因为最后要调的是cuda文件

复制代码

import glob
import os.path as osp
import platform
from setuptools import setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension

# 未避免一些问题，确保C++桥梁文件按 C++17 编译
extra_compile_args = ['/std:c++17', '/EHsc'] if platform.system() == 'Windows' else ['-std=c++17']

ROOT_DIR = osp.dirname(osp.abspath(__file__))
include_dirs = [osp.join(ROOT_DIR, "include")]
# 不需要手动指定c++和cuda文件
sources = glob.glob('*.cpp') + glob.glob('*.cu')

setup(
    name='cppcuda_tutorial',
    version='1.0',
    author='kwea123',
    author_email='kwea123@gmail.com',
    description='cppcuda_tutorial',
    long_description='cppcuda_tutorial',
    ext_modules=[
        CUDAExtension(
            name='cppcuda_tutorial',
            sources=sources,
            include_dirs=include_dirs,
            extra_compile_args={
                'cxx': ['-O2'],
                'nvcc': ['-O2']
            }
        )
    ],
    cmdclass={
        'build_ext': BuildExtension
    }
)

3.4 运行py文件调用cu文件中的自定义函数

和之前的一样，记得将数据送到CPU上

复制代码

# 一定要先引入torch库再引用自定义的c++函数库
import torch
import cppcuda_tutorial


if __name__ == '__main__':
    # feats = torch.randint(0,10,size = (3,4)) # feats must be a CUDA tensor
    # 因为要在cu文件也就是GPU中计算所以要先送到GPU上
    feats = torch.randint(0,10,size = (3,4)).to('cuda')
    point = torch.zeros(4,3).to('cuda')
    output = cppcuda_tutorial.trilinear_interpolation(feats, point)
    print(output)

结果如下图：
运行结果图

四、添加反向传播

三线性插值对输入特征 feats 是线性操作，因此反向传播比较直接。损失对每个顶点特征的梯度，等于损失对输出的梯度乘以该顶点对应的插值权重。这里设定不计算 points 的梯度，只计算 feats 的梯度

4.1 interpolation_kernel.cu 文件添加

新增参数：dL_dfeat_interp 是 PyTorch 自动求导系统传入的梯度，不是我们手动计算的

返回值：dL_feats 是我们通过反向 kernel 计算并返回的，最终成为 feats.grad

复制代码

// 反向传播
template <typename scalar_t>
__global__ void trilinear_bw_kernel(
    // 接收形参
    const torch::PackedTensorAccessor32<scalar_t, 2,torch::RestrictPtrTraits> dL_dfeat_interp,
    const torch::PackedTensorAccessor32<scalar_t, 3,torch::RestrictPtrTraits> feats,
    const torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> points,
    torch::PackedTensorAccessor32<scalar_t, 3, torch::RestrictPtrTraits> dL_feats

){
    // 计算线程编号坐标
    const int n = blockIdx.x * blockDim.x + threadIdx.x;
    const int f = blockIdx.y * blockDim.y + threadIdx.y;
    // 检查当前线程是否需要运算
    if (n >=feats.size(0) || f >= feats.size(2)) return;
    // 定义 权重u,v,w 因为point是 N,3
    // 还需要对points值做归一化使其值属于【0，1】
    // 因为假定值在-1~1区间所以归一化为（x + 1）/ 2
    const scalar_t u = (points[n][0] + 1) / 2;
    const scalar_t v = (points[n][1] + 1) / 2;
    const scalar_t w = (points[n][2] + 1) / 2;

    const scalar_t a = (1-v)*(1-w);
    const scalar_t b = (1-v)*w;
    const scalar_t c = v*(1-w);
    const scalar_t d = 1-a-b-c;
    //转到对参数feats得梯度
    dL_feats[n][0][f] = (1-u)*a*dL_dfeat_interp[n][f];
    dL_feats[n][1][f] = (1-u)*b*dL_dfeat_interp[n][f];
    dL_feats[n][2][f] = (1-u)*c*dL_dfeat_interp[n][f];
    dL_feats[n][3][f] = (1-u)*d*dL_dfeat_interp[n][f];
    dL_feats[n][4][f] = u*a*dL_dfeat_interp[n][f];
    dL_feats[n][5][f] = u*b*dL_dfeat_interp[n][f];
    dL_feats[n][6][f] = u*c*dL_dfeat_interp[n][f];
    dL_feats[n][7][f] = u*d*dL_dfeat_interp[n][f];
}

torch::Tensor trilinear_bw_cu(
    // 损失函数 L 对前向输出 feat_interp 的梯度
    torch::Tensor dL_dfeat_interp,
    // 维度 N, 8, F
    torch::Tensor feats,
    //  维度 N, 3
    torch::Tensor points
){
    // feats 形状为 (N, 8, F)，其中 8 代表三线性插值所需的八个邻近顶点，F 为特征维度。、
    // points 形状为 (N, 3)，每个样本的三个分量分别表示在体素局部坐标系中的归一化权重（范围 [0,1]），分别对应 x, y, z 方向。
    //输出 dL_feat 形状为 (N, 8, F)，通过对八个顶点的特征损失

    const int64_t N = feats.size(0), F = feats.size(2);
    // 维度 N, 8, F
    torch::Tensor dL_feats = torch::zeros({N, 8, F}, feats.options());

    const dim3 threads(16, 16);
    // 定义线程块（Block）大小
    const dim3 blocks((N + threads.x - 1) / threads.x,
                      (F + threads.y - 1) / threads.y);

    AT_DISPATCH_FLOATING_TYPES(feats.scalar_type(), "trilinear_bw_cu",[&] {
        // scalar_t 表示形状不定
        trilinear_bw_kernel<scalar_t><<<blocks, threads>>>(
            //packed_accessor32 将Tensor形态转换，非Tensor类型参数直接传入即可。
            // 参一是类型 参二是维度数  参三是一般都是torch::RestrictPtrTraits>()
            dL_dfeat_interp.packed_accessor32<scalar_t, 2, torch::RestrictPtrTraits>(),
            feats.packed_accessor32<scalar_t, 3, torch::RestrictPtrTraits>(),
            points.packed_accessor32<scalar_t, 2, torch::RestrictPtrTraits>(),
            // 一般设定为前三个为参数，最后一个是设置的接收数
            dL_feats.packed_accessor32<scalar_t, 3, torch::RestrictPtrTraits>()
        );
    });
    return dL_feats;
}

4.2 interpolation.cpp 文件添加

复制代码

torch::Tensor trilinear_interpolation_bw(
    torch::Tensor dL_dfeat_interp,
    torch::Tensor feats,
    torch::Tensor points
){
    CHECK_INPUT(dL_dfeat_interp);
    CHECK_INPUT(feats);
    CHECK_INPUT(points);

    // 调用反向传播得cuda函数
    return trilinear_bw_cu(dL_dfeat_interp, feats, points);
}

// 固定结构搭配
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    // 别名（setup.py文件加入环境后py文件的可以调用的函数名） ，自定义的函数名
    m.def("trilinear_interpolation_fw", &trilinear_interpolation_fw);
    m.def("trilinear_interpolation_bw", &trilinear_interpolation_bw);
}

4.3 utils.h 文件添加

复制代码

torch::Tensor trilinear_bw_cu(
    torch::Tensor dL_dfeat_interp,
    torch::Tensor feats,
    torch::Tensor points
);

4.4 测试文件添加

为了让自定义 CUDA 算子参与 PyTorch 自动微分，需要继承 torch.autograd.Function，并实现 forward 和 backward

复制代码

class Trilinear_interpolation_cuda(torch.autograd.Function):
    @staticmethod
    def forward(ctx, feats, points):
        feat_interp = cppcuda_tutorial.trilinear_interpolation_fw(feats, points)
        # 保存到反向传播使用
        ctx.save_for_backward(feats, points)

        return feat_interp

    @staticmethod
    # dL_dfeat_interp对应feat_interp，前向传播返回几个数，反向传播就有几个形参
    def backward(ctx, dL_dfeat_interp):
        feats, points = ctx.saved_tensors

        dL_dfeats = cppcuda_tutorial.trilinear_interpolation_bw(dL_dfeat_interp.contiguous(),feats,points)

        # 因为调用传入了几个参数需要计算梯度这里只有dL_dfeat_interp一个
        # 需要计算梯度得就正常返回，不需要的就返回None
        return dL_dfeats, None

对比测试：

复制代码

N = 65536; F = 256
rand = torch.rand(N, 8, F, device='cuda')
feats = rand.clone().requires_grad_()
feats2 = rand.clone().requires_grad_()
points = torch.rand(N, 3, device='cuda')*2-1

t = time.time()
# 用到了反向传播，所以同时需要前向和反向传播需要集中调用
out_cuda = Trilinear_interpolation_cuda.apply(feats2, points)
torch.cuda.synchronize()
print('    cuda time', time.time()-t, 's')

t = time.time()
out_py = trilinear_interpolation_py(feats, points)
torch.cuda.synchronize()
print('pytorch time', time.time()-t, 's')
print('fw all close', torch.allclose(out_py, out_cuda))

loss = out_py.sum()
loss.backward()
loss2 = out_cuda.sum()
# 执行自定义得反向传播
loss2.backward()
print('bw all close', torch.allclose(feats.grad, feats2.grad))

时间对比如下如所示：

本文内容参考YouTube博主 AI葵，并结合自身环境配置和理解进行整改，希望能帮到大家。