pytorch的c++/cuda扩展，CUDA编程

xx_xjm

已于 2024-09-01 11:38:23 修改

阅读量2.6k

点赞数 5

分类专栏： pytorch扩展文章标签： pytorch 人工智能 python

于 2023-04-25 18:48:40 首次发布

本文链接：https://blog.csdn.net/xx_xjm/article/details/130367955

版权

本文介绍了CUDA编程的基本概念，包括CUDA编程的目的、.cu文件的构成和编译、核函数的定义与调用，以及主机和设备间的数据传输。同时，详细阐述了如何将CUDA核函数应用于PyTorch的C++扩展，涉及Aten、c10和torch命名空间的理解以及扩展流程。

摘要由CSDN通过智能技术生成

1：首先，我们写一个单纯的cpp代码，用于计算两个相同长度的一维数组对应元素之和，代码如下：

2：接下来，我们将上面的cpp代码转化为.cu代码：

3：cuda相关api函数讲解：

a: cudaMalloc

b: cudaFree

c：cudaMemcpy:主机和设备之间的数据传输

一：三大命名空间Aten，c10，torch

二：扩展流程

主要分两部分来讲，一是cuda编程，二是pytorch的c++/cuda扩展；

先说cuda编程

一：什么是cuda编程

我们知道C++，C这类的编程语言是为了让计算机执行我们的指令，确切一点是让计算机的cpu执行我们的执行，现在cuda编程则是要让显卡中的计算核心执行我们的指令；

所以，cuda编程其实就是编写显卡中计算核心执行指令。为了区别于.cpp，.c这样的文件，我们取.cu后缀来指明当前的代码文件是给显卡用的；

二：.cu文件的构成

（这部分基于个人理解，没有严格证明，仅方便理解，如有不正之处，欢迎指正）

因为显卡上面大量的都是计算单元，只有少量的控制单元，我们是无法像C++，C直接指挥CPU那样直接把代码编译成显卡能识别的东西，简单点说，就是，我们不能像C++，C这样写完代码后，编译链接完CPU就可以直接运行了。为此，我们是需要通过CPU转达我们的指令给到显卡；所以说到底，.cu文件也是写给CPU的，只不过你告诉了CPU要怎么去指挥显卡工作，从这个角度来看，其实.cu文件和.cpp等文件一样，都是C++的代码；

因此，.cu文件中包含两部分的内容，一是告诉CPU怎么传递信息和指挥显卡工作，二是指明具体显卡上面应该怎么操作；

三：.cu文件的编译

上面说了，.cu文件包含两部分内容，一部分是写给CPU的，一部分是写给GPU的，所以编译的时候自然需要先区分开这两部分代码，分别编译，然后再合起来一起链接，形成相应的可调用的库（动态库或者静态库），.cu文件的编译需要用到nvcc，这部分详细的可以参考：

(3条消息) CUDA学习（一）-NVCC的编译过程_nvcc编译_Scott f的博客-CSDN博客

nvcc编译网上有很多教程，这一片看不懂的话，大家也可以自查

四：.cu文件的编写

这部分内容主要参考《CUDA编程基础与实践》---清华大学出版社，讲的非常好，强推！

我们前面说了，.cu文件中需要实现cpu对显卡（也就是设备端，后面若无特指，均以设备代替gpu）进行调用，这一步是通过核函数来实现的；

一个典型的，简单的cuda程序结构如下：

注：本文所用代码实例无特殊说明均来自《CUDA编程基础与实践》

int main(void)
{
    主机代码
    核函数调用
    主机代码
    return 0;
}

1：什么是核函数？

我们上面说了，.cu本质上还是写给CPU的，所以核函数其实也是C++函数的一种，只不过有一个特殊的限定词"__global__"，用以指明“这个C++函数，是用来调用显卡的！”。

一个简单的核函数如下，其中__global__和void的顺序可以互换，另外，核函数的返回值必须是空类型：

__global__ void hello_world()
{
    printf("hello world")
}

2：核函数的调用

核函数虽然是C++函数的一种，但是它的调用有一点区别，举例如下：

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

__global__ void hello_world()
{
    printf("hello world");
}

int main(void)
{
    hello_world<<<1, 1>>>();
    cudaDeviceSynchronize();
    return 0;
}

可以看到，我们这里调用核函数的时候有一个<<<1,1>>>,这其实是指明核函数中的线程数目和排列情况，具体的可以自行百度或者参考书籍，简单理解就是分配多少计算资源来进行运算；

五：主机和设备之间的数据传输

在第四节中，我们简单用主机通过核函数来指挥设备进行工作，但没有进行数据的传输，在这一节，我们更细致的讲解主机和设备间的工作流程；

1：首先，我们写一个单纯的cpp代码，用于计算两个相同长度的一维数组对应元素之和，代码如下：

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

const double EPSILON = 1.0e-15;
const double a = 1.23;
const double b = 2.34;
const double c = 3.57;

void add(con