开发者社区> 游客z3jcatjk57fiu> 正文
阿里云
为了无法计算的价值
打开APP
阿里云APP内打开

PyTorch自定义CUDA算子教程与运行时间分析(一)

简介: 最近因为工作需要,学习了一波CUDA。这里简单记录一下PyTorch自定义CUDA算子的方法,写了一个非常简单的example,再介绍一下正确的PyTorch中CUDA运行时间分析方法。
+关注继续查看

最近因为工作需要,学习了一波CUDA。这里简单记录一下PyTorch自定义CUDA算子的方法,写了一个非常简单的example,再介绍一下正确的PyTorch中CUDA运行时间分析方法。

所有的代码都放在了github上,地址是:https://github.com/godweiyang/torch-cuda-example

完整流程


下面我们就来详细了解一下PyTorch是如何调用自定义的CUDA算子的。

首先我们可以看到有四个代码文件:

  • main.py,这是python入口,也就是你平时写模型的地方。
  • add2.cpp,这是torch和CUDA连接的地方,将CUDA程序封装成了python可以调用的库。
  • add2.h,CUDA函数声明。
  • add2.cu,CUDA函数实现。

然后逐个文件看一下是怎么调用的。

CUDA算子实现


首先最简单的当属add2.hadd2.cu,这就是普通的CUDA实现。

void launch_add2(float *c,
                 const float *a,
                 const float *b,
                 int n);
__global__ void add2_kernel(float* c,
                            const float* a,
                            const float* b,
                            int n) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
            i < n; i += gridDim.x * blockDim.x) {
        c[i] = a[i] + b[i];
    }
}
void launch_add2(float* c,
                 const float* a,
                 const float* b,
                 int n) {
    dim3 grid((n + 1023) / 1024);
    dim3 block(1024);
    add2_kernel<<<grid, block>>>(c, a, b, n);
}

这里实现的功能是两个长度为的tensor相加,每个block有1024个线程,一共有个block。具体CUDA细节就不讲了,本文重点不在于这个。

add2_kernel是kernel函数,运行在GPU端的。而launch_add2是CPU端的执行函数,调用kernel。注意它是异步的,调用完之后控制权立刻返回给CPU,所以之后计算时间的时候要格外小心,很容易只统计到调用的时间。

Torch C++封装


这里涉及到的是add2.cpp,这个文件主要功能是提供一个PyTorch可以调用的接口。

#include <torch/extension.h>
#include "add2.h"
void torch_launch_add2(torch::Tensor &c,
                       const torch::Tensor &a,
                       const torch::Tensor &b,
                       int n) {
    launch_add2((float *)c.data_ptr(),
                (const float *)a.data_ptr(),
                (const float *)b.data_ptr(),
                n);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("torch_launch_add2",
          &torch_launch_add2,
          "add2 kernel warpper");
}

torch_launch_add2函数传入的是C++版本的torch tensor,然后转换成C++指针数组,调用CUDA函数launch_add2来执行核函数。

这里用pybind11来对torch_launch_add2函数进行封装,然后用cmake编译就可以产生python可以调用的.so库。但是我们这里不直接手动cmake编译,具体方法看下面的章节。

Python调用


最后就是python层面,也就是我们用户编写代码去调用上面生成的库了。

import time
import numpy as np
import torch
from torch.utils.cpp_extension import load
cuda_module = load(name="add2",
                   sources=["add2.cpp", "add2.cu"],
                   verbose=True)
# c = a + b (shape: [n])
n = 1024 * 1024
a = torch.rand(n, device="cuda:0")
b = torch.rand(n, device="cuda:0")
cuda_c = torch.rand(n, device="cuda:0")
ntest = 10
def show_time(func):
    times = list()
    res = list()
    # GPU warm up
    for _ in range(10):
        func()
    for _ in range(ntest):
        # sync the threads to get accurate cuda running time
        torch.cuda.synchronize(device="cuda:0")
        start_time = time.time()
        r = func()
        torch.cuda.synchronize(device="cuda:0")
        end_time = time.time()
        times.append((end_time-start_time)*1e6)
        res.append(r)
    return times, res
def run_cuda():
    cuda_module.torch_launch_add2(cuda_c, a, b, n)
    return cuda_c
def run_torch():
    # return None to avoid intermediate GPU memory application
    # for accurate time statistics
    a + b
    return None
print("Running cuda...")
cuda_time, _ = show_time(run_cuda)
print("Cuda time:  {:.3f}us".format(np.mean(cuda_time)))
print("Running torch...")
torch_time, _ = show_time(run_torch)
print("Torch time:  {:.3f}us".format(np.mean(torch_time)))

这里6-8行的torch.utils.cpp_extension.load函数就是用来自动编译上面的几个cpp和cu文件的。最主要的就是sources参数,指定了需要编译的文件列表。然后就可以通过cuda_module.torch_launch_add2,也就是我们封装好的接口来进行调用。

接下来的代码就随心所欲了,这里简单写了一个测量运行时间,对比和torch速度的代码,这部分留着下一章节讲解。

总结一下,主要分为三个模块:

  • 先编写CUDA算子和对应的调用函数。
  • 然后编写torch cpp函数建立PyTorch和CUDA之间的联系,用pybind11封装。
  • 最后用PyTorch的cpp扩展库进行编译和调用。

运行时间分析


我们知道,CUDA kernel函数是异步的,所以不能直接在CUDA函数两端加上time.time()测试时间,这样测出来的只是调用CUDA api的时间,不包括GPU端运行的时间。

所以我们要加上线程同步函数,等待kernel中所有线程全部执行完毕再执行CPU端后续指令。这里我们将同步指令加在了python端,用的是torch.cuda.synchronize函数。

具体来说就是形如下面代码:

torch.cuda.synchronize()
start_time = time.time()
func()
torch.cuda.synchronize()
end_time = time.time()

其中第一次同步是为了防止前面的代码中有未同步还在GPU端运行的指令,第二次同步就是为了等fun()所有线程执行完毕后再统计时间。

这里我们torch和cuda分别执行10次看看平均时间,此外执行前需要先执行10次做一下warm up,让GPU达到正常状态。

我们分别测试四种情况,分别是:

  • 两次同步
  • 第一次同步,第二次不同步
  • 第一次不同步,第二次同步
  • 两次不同步

这里我们采用英伟达的Nsight Systems来可视化运行的每个时刻指令执行的情况。

安装命令为:

sudo apt install nsight-systems

然后在运行python代码时,在命令前面加上nsys profile就行了:

nsys profile python3 main.py

然后就会生成report1.qdstrmreport1.sqlite两个文件,将report1.qdstrm转换为report1.qdrep文件:

QdstrmImporter -i report1.qdstrm

最后将生成的report1.qdrep文件用Nsight Systems软件打开,我这里是mac系统。

版权声明:本文内容由阿里云实名注册用户自发贡献,版权归原作者所有,阿里云开发者社区不拥有其著作权,亦不承担相应法律责任。具体规则请查看《阿里云开发者社区用户服务协议》和《阿里云开发者社区知识产权保护指引》。如果您发现本社区中有涉嫌抄袭的内容,填写侵权投诉表单进行举报,一经查实,本社区将立刻删除涉嫌侵权内容。

相关文章
Java Record 的一些思考 - 序列化相关(上)
Java Record 的一些思考 - 序列化相关(上)
16 0
【小家Java】Servlet规范之---请求(request):Servlet中如何获取POST请求参数?(使用getParameter())
【小家Java】Servlet规范之---请求(request):Servlet中如何获取POST请求参数?(使用getParameter())
212 0
JavaScript加密-MD5与Bcrypt
JavaScript加密-MD5与Bcrypt
71 0
Mongodb for PHP教程之管理工具
下面物件事情应该是PHP程序员应该了解的,更详细的参考可以百度谷歌之 mongodb PHP 文档http://us.php.net/manual/en/book.mongo.php 1.MongoDB是一个单独的服务器 2.MongoDB是基于文档,而不是基于表 3.MongoDB较少的Schema 4.你不需要另外学习一种语言 5.MongoDB有良好的PHP支持 1。
1068 0
217
文章
0
问答
文章排行榜
最热
最新
相关电子书
更多
低代码开发师(初级)实战教程
立即下载
阿里巴巴DevOps 最佳实践手册
立即下载
冬季实战营第三期:MySQL数据库进阶实战
立即下载