【CUDA编程--编程模型简介&算子开发流程】

【CUDA编程--编程模型简介&算子开发流程】_第1张图片

官方文档:https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html

什么是CUDA

  • CUDA全称(Compute Unified Device Architecture)统一计算架构,是NVIDIA推出的并行计算平台
  • 深度学习加速:对于神经网络,无论是离线训练还是在线推理,都有巨量的矩阵、归一化、softmax等运算,且其中有非常多的并行计算,非常适合用GPU来进行运算加速
    【CUDA编程--编程模型简介&算子开发流程】_第2张图片
  • 一般来说,应用程序混合有并行部分和顺序部分,因此系统设计时混合使用 GPU 和 CPU,以最大限度地提高整体性能。具有高度并行性的应用程序可以利用 GPU 的大规模并行特性来实现比 CPU 更高的性能

CUDA编程模型

  • 多核CPU和众核GPU的出现意味着主流处理器芯片现在都是并行系统

kernel 核

  • 不同于C语言中函数的调用,CUDA的内核函数调用时需要指定总的线程数量,以及相应的线程布局(grid和block维度配置)
// C函数
function_name (argument list);
// CUDA kernel call
kernel_name<<<4, 8>>>(argument list);  // 这里执行有grid中有4个block, 以及每个block中有8个线程运行

【CUDA编程--编程模型简介&算子开发流程】_第3张图片

限定符

因为数据在全局内存中是线性存储的,所以可以通过blockIdx.x和threadIdx.x来标识grid中的线程,建立线程和数据之间的映射关系
核函数限定符的意义如下

限定符 执行 调用 备注
global Device执行 Host调用/Device调用 必须有一个void的返回类型
device Device执行 Device调用
host Host执行 Host调用

举例

  • 实现的功能是两个长度为的tensor相加,每个block有1024个线程,一共有n/1024
    个block
    cudademo.cu
#include   
#include 

// 代码的核心诉求(Cuda上运行):
// 输入a: 0,1,2,3,4.....
// 输入b: 0,2,4,6,8.....
// 输出c: 0,3,6,9,12.....

__global__ void my_add_kernel(float* c,  
                            const float* a,  
                            const float* b,  
                            int n) {  
    // 定义核函数 add  
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;  
            i < n; i += gridDim.x * blockDim.x) {  
        c[i] = a[i] + b[i];  
    }  
}

int main() {
    int n = 1000; // 元素数量

    // 分配主机内存用于输入和输出数组
    float* host_a = new float[n];
    float* host_b = new float[n];
    float* host_c = new float[n];

    // 在主机上填充输入数组 a 和 b
    for (int i = 0; i < n; i++) {
        host_a[i] = i;
        host_b[i] = i * 2;
    }

    // 在设备上分配内存用于输入和输出数组
    float* device_a;
    float* device_b;
    float* device_c;
    cudaMalloc((void**)&device_a, n * sizeof(float));
    cudaMalloc((void**)&device_b, n * sizeof(float));
    cudaMalloc((void**)&device_c, n * sizeof(float));

    // 将输入数组从主机内存复制到设备内存
    cudaMemcpy(device_a, host_a, n * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(device_b, host_b, n * sizeof(float), cudaMemcpyHostToDevice);

    // 定义 CUDA 核函数的执行配置
    int block_size = 256;
    int grid_size = (n + block_size - 1) / block_size;

    // 调用 CUDA 核函数
    my_add_kernel<<<grid_size, block_size>>>(device_c, device_a, device_b, n);

    // 将输出数组从设备内存复制到主机内存
    cudaMemcpy(host_c, device_c, n * sizeof(float), cudaMemcpyDeviceToHost);

    // 打印输出数组
    for (int i = 0; i < n; i++) {
        std::cout << host_c[i] << " ";
    }
    std::cout << std::endl;

    // 释放主机和设备内存
    delete[] host_a;
    delete[] host_b;
    delete[] host_c;
    cudaFree(device_a);
    cudaFree(device_b);
    cudaFree(device_c);

    return 0;
}
  • 直接编译执行
[]# nvcc cudademo.cu -o cudademo
[]# cudademo
0 3 6 9 12 15 18 21 24 27 30 33 36 39 42 45 48 51 54 57 60 63 .....
如何准备安装cuda运行环境
  • 确认驱动版本等信息
[]:~$ nvidia-smi
Mon Nov 13 11:22:17 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.65.01    Driver Version: 515.65.01    CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:35:00.0 Off |                  N/A |
| 30%   23C    P8    15W / 350W |    807MiB / 24576MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  Off  | 00000000:36:00.0 Off |                  N/A |
| 30%   26C    P8    13W / 350W |      2MiB / 24576MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  NVIDIA GeForce ...  Off  | 00000000:39:00.0 Off |                  N/A |
| 75%   67C    P2   251W / 350W |   9306MiB / 24576MiB |     85%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   3  NVIDIA GeForce ...  Off  | 00000000:3D:00.0 Off |                  N/A |
| 30%   25C    P8    13W / 350W |      2MiB / 24576MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
  • 也就需要我们使用cuda11.7版本
  • cuda与pytorch 软件版本对应关系以及安装:https://pytorch.org/get-started/previous-versions/
  • cuda官方镜像网站:https://hub.docker.com/r/nvidia/cuda/tags
  • 当前根据版本对应关系,直接可以使用 docker pull nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04 该版本
version: '3'
services:
  my_container:
    image: harbor.uat.enflame.cc/library/enflame.cn/nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04
    runtime: nvidia  # 指定使用NVIDIA GPU运行时
    devices:
      - /dev/nvidia0  # 将主机的NVIDIA GPU设备映射到容器
      - /dev/nvidia1
      - /dev/nvidia2
      - /dev/nvidia3
      - /dev/nvidia4
      - /dev/nvidia5
      - /dev/nvidia6
      - /dev/nvidia7
    network_mode: host
    command: sleep 10000000000000000
    shm_size: '8gb'
    privileged: true

数据处理方式

佛林分类法Flynn’s Taxonomy,根据指令和数据进入CPU的方式对计算机架构进行分类,分为以下四类
- 单指令单数据 (SISD):传统的单核处理数据方式
- 单指令多数据(SIMD):单核执行一条指令完成多数据处理(游戏中向量、矩阵)
- 多指令单数据 (MISD):多核执行不同的指令处理单个数据(少见)
- 多指令多数据 (MIMD):多核执行不同的指令处理多个数据
【CUDA编程--编程模型简介&算子开发流程】_第4张图片
为了提高并行的计算能力,架构上实现下面这些性能提升:

  • 降低延迟(latency):指操作从开始到结束所需要的时间,一般用微秒计算,延迟越低越好
  • 增加带宽(bandwidth):单位时间内处理的数据量,一般用MB/s或者GB/s表示
  • 增加吞吐(throughput):单位时间内成功处理的运算数量

内存划分

  • 分布式内存的多节点系统

    • 集群,各个机器之前通过网络进行数据交互
    • 传统的比如redis集群这种通信等
  • 共享内存的多处理器系统

    • 包括单片多核,多片多核,主要是针对同设备多核进行数据通信,GPU是众核架构,表述为Single Instruction, Multiple Thread (SIMT),不同于SIMD,SIMT是真正的启动了多个线程,执行相同的指令,去完成数据的并行运算
    • 3090显卡拥有10496个CUDA核心,相比上一代2080Ti显卡的4352个CUDA核心数量增加了一倍

编程结构

  • CUDA编程让你可以在CPU-GPU的异构计算系统上高效执行应用程序,语法只是在C语言的基础上做了简单的扩展,CUDA C++ 通过允许程序员定义称为内核的 C++ 函数来扩展 C++,这些函数在调用时由 N 个不同的CUDA 线程并行执行 N 次,而不是像常规 C++ 函数那样只执行一次,在开始编程前,我们首先得理清Host和Device的概念

    • Host:CPU及其内存
    • Device:GPU及其内存
  • 运行在GPU设备上的代码我们称为kernel

  • 典型的CUDA程序处理流程

    • 分配内存,数据初始化
      将数据从Host拷贝到Device
    • 调用kernels处理数据,然后存在GPU内存(Device)
    • 将数据从Device拷贝到Host
    • 内存释放
      【CUDA编程--编程模型简介&算子开发流程】_第5张图片

内存管理

标准C函数 CUDA 函数 CUDA函数说明
malloc cudaMalloc GPU内存分配
memcpy cudaMemcpy 用于Host和Device之间数据传输
memset cudaMemset 设定数据填充到GPU内存中
free cudaFree 释放GPU内存

【CUDA编程--编程模型简介&算子开发流程】_第6张图片
【CUDA编程--编程模型简介&算子开发流程】_第7张图片

  • CUDA 线程在执行期间可以访问多个内存空间中的数据
  • 每个线程都有私有本地内存
  • 每个线程块都有对该块的所有线程可见的共享内存,并且与该块具有相同的生命周期
  • 线程块簇中的线程块可以对彼此的共享内存执行读、写和原子操作。所有线程都可以访问相同的全局内存。
  • 还有两个可供所有线程访问的附加只读内存空间:常量内存空间和纹理内存空间

Grid&&Block

【CUDA编程--编程模型简介&算子开发流程】_第8张图片

  • 一个Kernel所launch的所有线程称为grid,他们共享相同的全局内存空间(global memory space)
  • 一个grid由多个block(线程块)组成,block内部的线程可以通过以下两点进行协作(不同block间的线程不能协作)
    • block本地同步(synchronization)
    • block本地共享内存(sharedmemory)
  • 一个线程通过blockIdx(grid内的index)和threadIdx(block内的index)这两个坐标变量(三维类型unit3)来唯一标识(线程运行的时候这两个变量会被CUDA赋上相应的坐标值,可以直接使用)
  • grid和block的维度信息通过gridDim和blockDim(dim3)来表示
    • gridDim:表示一个grid里面有多少个blocks
    • blockDim:表示一个block里面有多少个threads
线程块结构

为了方便起见,threadIdx是一个3分量向量,因此可以使用一维、二维或三维线程索引 来标识线程,形成一维、二维或三维线程块,称为线程块。这提供了一种自然的方式来调用域中元素(例如向量、矩阵或体积)的计算
【CUDA编程--编程模型简介&算子开发流程】_第9张图片
每个块的线程数量是有限的,因为块中的所有线程都应驻留在同一个流式多处理器核心上,并且必须共享该核心的有限内存资源。在当前的 GPU 上,一个线程块最多可以包含 1024 个线程

然而,一个内核可以由多个形状相同的线程块来执行,因此线程总数等于每个块的线程数乘以块数
块被组织成一维、二维或三维线程块网格,如图4所示。网格中线程块的数量通常由正在处理的数据的大小决定,该数据通常超过系统中处理器的数量

// Kernel definition
__global__ void MatAdd(float A[N][N], float B[N][N],
                       float C[N][N])
{
    int i = threadIdx.x;
    int j = threadIdx.y;
    C[i][j] = A[i][j] + B[i][j];
}

int main()
{
    ...
    // Kernel invocation with one block of N * N * 1 threads
    int numBlocks = 1;
    dim3 threadsPerBlock(N, N);
    MatAdd<<<numBlocks, threadsPerBlock>>>(A, B, C);
    ...
}

线程块簇

【CUDA编程--编程模型简介&算子开发流程】_第10张图片

  • 随着 NVIDIA计算能力 9.0的推出,CUDA 编程模型引入了一个可选的层次结构级别,称为由线程块组成的线程块集群。与如何保证线程块中的线程在流式多处理器上共同调度类似,集群中的线程块也保证在 GPU 中的 GPU 处理集群 (GPC) 上共同调度
小结
  • 如何要使用cuda进行并行计算,使用cuda函数进行数据等操作
  • cuda的线程结构将cuda编程结构分为块与线程,都是可以由一维、二维或三维唯一索引来标识,如代码
{
    int i = threadIdx.x;
    int j = threadIdx.y;
    C[i][j] = A[i][j] + B[i][j];
}

GPU架构

阅读官方文档:https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html, 一句话:太多了,不如实操,还是简单总结下

GPU架构就是由可扩展的流式多处理器(Streaming Multiprocessors简称SM)阵列所构建,整个硬件的并行就是不断复制这种架构实现的。通常每个GPU都有多个SM,每个SM都支持上百个线程的并行,所以GPU可以支持上千个线程的并行

SM中的核心部件

  • CUDA Cores:核心,是最小的执行单元
  • Shared Memory/L1 Cache:共享内存和L1缓存,他们共用64KB空间,根据Bl
  • Register File:寄存器,根据线程划分
  • Load/Store Units:16个数据读写单元,支持16个线程一起从Cache/DRAM存取数据
  • Special Function Units:4个特殊函数处理单元,用于sin/cos这类指令计算
  • Warp Scheduler:Warp调度器,所谓Warp就是32个线程组成的线程束,是最小的调度单元

GPU内存

【CUDA编程--编程模型简介&算子开发流程】_第11张图片

  • 寄存器:GPU上访问最快的存储空间,是SM中的稀缺资源,对于每个线程是私有的,Fermi架构中每个线程最多63个,Kepler结构扩展到255个。如果变量太多寄存器不够,会发生寄存器溢出,此时本地内存会存储多出来的变量,这种情况对性能影响较大。
  • 本地内存:本质上是和全局内存放在同一块存储区域中(compute capability 2.0以上的设备,会放在SM的一级缓存,或者设备的二级缓存上)具有高延迟、低带宽,编译器可能会将以下变量存放于本地内存:
    • 编译时期无法确定索引引用的本地数组
    • 可能会消耗大量寄存器的较大本地数组/结构体
    • 任何不满足核函数寄存器限定条件的变量
  • 共享内存:因为是片上内存,所以相比全局内存和本地内存,具有较高的带宽和较低的延迟
    • SM中的一级缓存,和共享内存共享一个64k的片上内存,L1不可编程,共享内存可以
    • 切勿过度使用共享内存,导致部分线程块无法被SM启动,影响Warp调度
    • 可以使用__syncthreads()来实现Block内线程的同步
  • 常量内存:驻留在设备内存中,每个SM都有专用的常量内存缓存
    • 常量内存在核函数外,全局范围内声明,对于所有设备,只可以声明64k的常量内存
    • 核函数无法修改,Host端使用cudaMemcpyToSymbol接口初始化
  • 纹理内存:驻留在设备内存中,在每个SM的只读缓存中缓存,对于2D数据的访问性能较好
  • 全局内存:GPU上最大的内存空间,延迟最高,使用最常见的内存,访问是对齐访问,也就是一次要读取指定大小(32,64,128)整数倍字节的内存,所以当线程束执行内存加载/存储时,需要满足的传输数量通常取决与以下两个因素:
  • 跨线程的内存地址分布
  • 内存事务的对齐方式。
修饰符 变量名称 存储器 作用域 生命周期
float var 寄存器 线程 线程
float var[100] 本地 线程 线程
share float var* 共享
device float var* 全局 全局 应用程序
constant float var* 常量 全局 应用程序
存储器 缓存 存取 范围 生命周期
寄存器 R/W 一个线程 线程
本地 1.0以上有 R/W 一个线程 线程
共享 R/W 块内所有线程
全局 1.0以上有 R/W 所有线程+主机 主机配置
常量 R 所有线程+主机 主机配置
纹理 R 所有线程+主机 主机配置
  • GPU缓存
    与CPU缓存类似,GPU缓存不可编程,其行为出厂是时已经设定好了。GPU上有4种缓存:
    • 一级缓存:每个SM都有一个一级缓存,与共享内存公用空间
    • 二级缓存:所有SM公用一个二级缓存
    • 只读常量缓存:每个SM有
    • 只读纹理缓存:每个SM有

案例

不说概念了,直接肝

获取GPU信息

#include
#include
#include
int main() {
    int dev = 0;
    cudaDeviceProp devProp;
    cudaGetDeviceProperties(&devProp, dev);
    std::cout << "GPU Device Name" << dev << ": " << devProp.name << std::endl;
    std::cout << "SM Count: " << devProp.multiProcessorCount << std::endl;
    std::cout << "Shared Memory Size per Thread Block: " << devProp.sharedMemPerBlock / 1024.0 << " KB" << std::endl;
    std::cout << "Threads per Thread Block: " << devProp.maxThreadsPerBlock << std::endl;
    std::cout << "Threads per SM: " << devProp.maxThreadsPerMultiProcessor << std::endl;
    std::cout << "Warps per SM: " << devProp.maxThreadsPerMultiProcessor / 32 << std::endl;
    return 0;
}
[]# nvcc checkDeviceInfor.cu -o checkDeviceInfor
GPU Device Name0: NVIDIA GeForce RTX 3090
SM Count: 82
Shared Memory Size per Thread Block: 48 KB
Threads per Thread Block: 1024
Threads per SM: 1536
Warps per SM: 48

实现CUDA算子

下面的案例你将学习到:

  • 最简单的CUDA算子的写法。
  • 最简洁的PyTorch和TensorFlow封装CUDA算子的方法。
  • 几种编译CUDA算子的方法。
  • python调用CUDA算子的几种方式。
  • python中统计CUDA算子运行时间的正确方法。
  • PyTorch和TensorFlow自定义算子梯度的方法

代码结构

├── include
│   └── add2.h # cuda算子的头文件
├── kernel
│   ├── add2_kernel.cu # cuda算子的具体实现
│   └── add2.cpp # cuda算子的cpp torch封装
├── CMakeLists.txt
├── LICENSE
├── README.md
├── setup.py
├── time.py # 比较cuda算子和torch实现的时间差异
└── train.py # 使用cuda算子来训练模型

代码结构还是很清晰的。include文件夹用来放cuda算子的头文件(.h文件),里面是cuda算子的定义。kernel文件夹放cuda算子的具体实现(.cu文件)和cpp torch的接口封装(.cpp文件)。

最后是python端调用,我实现了两个功能。一是比较运行时间,上一篇教程详细讲过了;二是训练一个PyTorch模型

头文件
void launch_add2(float *c,
                 const float *a,
                 const float *b,
                 int n);
算子核函数
  • kernel/add2_kernel.cu
__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);
}
  • kernel/add2.cpp
#include 
#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");
} // cpp端用的是pybind11进行封装,适用python

TORCH_LIBRARY(add2, m) {
    m.def("torch_launch_add2", torch_launch_add2);
} // cpp端用的是TORCH_LIBRARY进行封装,适用于

编译
JIT

JIT就是just-in-time,也就是即时编译,或者说动态编译,就是说在python代码运行的时候再去编译cpp和cuda文件。

import torch
from torch.utils.cpp_extension import load
cuda_module = load(name="add2",
                   extra_include_paths=["include"],
                   sources=["kernel/add2.cpp", "kernel/add2_kernel.cu"],
                   verbose=True)
cuda_module.torch_launch_add2(c, a, b, n)

需要注意的就是两个参数,extra_include_paths表示包含的头文件目录,sources表示需要编译的代码,一般就是.cpp和.cu文件
运行成功可以看到

[1/2] nvcc -c add2_kernel.cu -o add2_kernel.cuda.o
[2/3] c++ -c add2.cpp -o add2.o
[3/3] c++ add2.o add2_kernel.cuda.o -shared -o add2.so
Setuptools

编译的方式是通过Setuptools,也就是编写setup.py,具体代码如下

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

setup(
    name="add2",
    include_dirs=["include"],
    ext_modules=[
        CUDAExtension(
            "add2",
            ["kernel/add2.cpp", "kernel/add2_kernel.cu"],
        )
    ],
    cmdclass={
        "build_ext": BuildExtension
    }
)

执行python3 setup.py install
这样就能生成动态链接库,同时将add2添加为python的模块了,可以直接import add2来调用

import torch
import add2
add2.torch_launch_add2(c, a, b, n)
cmake

最后就是cmake编译的方式了,要编写一个CMakeLists.txt文件,代码如下

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
# 修改为你自己的nvcc路径,或者删掉这行,如果能运行的话。
set(CMAKE_CUDA_COMPILER "/usr/local/cuda/bin/nvcc")
project(add2 LANGUAGES CXX CUDA)

find_package(Torch REQUIRED)
find_package(CUDA REQUIRED)
find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib")

# 修改为你自己的python路径,或者删掉这行,如果能运行的话。
include_directories(/usr/include/python3.7)
include_directories(include)

set(SRCS kernel/add2.cpp kernel/add2_kernel.cu)
add_library(add2 SHARED ${SRCS})

target_link_libraries(add2 "${TORCH_LIBRARIES}" "${TORCH_PYTHON_LIBRARY}")

编译

mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH="$(python3 -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ../
make

最后会在build目录下生成一个libadd2.so,通过如下方式在python端调用:

import torch
torch.ops.load_library("build/libadd2.so")
torch.ops.add2.torch_launch_add2(c, a, b, n)
将上诉代码汇总
  • timer.py
import time
import argparse
import numpy as np
import torch

# 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 = None
    # GPU warm up
    for _ in range(10):
        res = func()
    for _ in range(ntest):
        # sync the threads to get accurate cuda running time
        torch.cuda.synchronize(device="cuda:0")
        start_time = time.time()
        func()
        torch.cuda.synchronize(device="cuda:0")
        end_time = time.time()
        times.append((end_time-start_time)*1e6)
    return times, res

def run_cuda():
    if args.compiler == 'jit':
        cuda_module.torch_launch_add2(cuda_c, a, b, n)
    elif args.compiler == 'setup':
        add2.torch_launch_add2(cuda_c, a, b, n)
    elif args.compiler == 'cmake':
        torch.ops.add2.torch_launch_add2(cuda_c, a, b, n)
    else:
        raise Exception("Type of cuda compiler must be one of jit/setup/cmake.")

    return cuda_c

def run_torch():
    c = a + b
    return c.contiguous()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--compiler', type=str, choices=['jit', 'setup', 'cmake'], default='jit')
    args = parser.parse_args()

    if args.compiler == 'jit':
        from torch.utils.cpp_extension import load
        cuda_module = load(name="add2",
                           extra_include_paths=["include"],
                           sources=["pytorch/add2_ops.cpp", "kernel/add2_kernel.cu"],
                           verbose=True)
    elif args.compiler == 'setup':
        import add2
    elif args.compiler == 'cmake':
        torch.ops.load_library("build/libadd2.so")
    else:
        raise Exception("Type of cuda compiler must be one of jit/setup/cmake.")

    print("Running cuda...")
    cuda_time, cuda_res = show_time(run_cuda)
    print("Cuda time:  {:.3f}us".format(np.mean(cuda_time)))

    print("Running torch...")
    torch_time, torch_res = show_time(run_torch)
    print("Torch time:  {:.3f}us".format(np.mean(torch_time)))

    torch.allclose(cuda_res, torch_res)
    print("Kernel test passed.")
  • train.py
import argparse
import numpy as np
import torch
from torch import nn
from torch.autograd import Function

class AddModelFunction(Function):
    @staticmethod
    def forward(ctx, a, b, n):
        c = torch.empty(n).to(device="cuda:0")

        if args.compiler == 'jit':
            cuda_module.torch_launch_add2(c, a, b, n)
        elif args.compiler == 'setup':
            add2.torch_launch_add2(c, a, b, n)
        elif args.compiler == 'cmake':
            torch.ops.add2.torch_launch_add2(c, a, b, n)
        else:
            raise Exception("Type of cuda compiler must be one of jit/setup/cmake.")

        return c

    @staticmethod
    def backward(ctx, grad_output):
        return (grad_output, grad_output, None)


class AddModel(nn.Module):
    def __init__(self, n):
        super(AddModel, self).__init__()
        self.n = n
        self.a = nn.Parameter(torch.Tensor(self.n))
        self.b = nn.Parameter(torch.Tensor(self.n))
        self.a.data.normal_(mean=0.0, std=1.0)
        self.b.data.normal_(mean=0.0, std=1.0)

    def forward(self):
        a2 = torch.square(self.a)
        b2 = torch.square(self.b)
        c = AddModelFunction.apply(a2, b2, self.n)
        return c

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--compiler', type=str, choices=['jit', 'setup', 'cmake'], default='jit')
    args = parser.parse_args()

    if args.compiler == 'jit':
        from torch.utils.cpp_extension import load
        cuda_module = load(name="add2",
                           extra_include_paths=["include"],
                           sources=["pytorch/add2_ops.cpp", "kernel/add2_kernel.cu"],
                           verbose=True)
    elif args.compiler == 'setup':
        import add2
    elif args.compiler == 'cmake':
        torch.ops.load_library("build/libadd2.so")
    else:
        raise Exception("Type of cuda compiler must be one of jit/setup/cmake.")

    n = 1024

    print("Initializing model...")
    model = AddModel(n)
    model.to(device="cuda:0")

    print("Initializing optimizer...")
    opt = torch.optim.SGD(model.parameters(), lr=0.01)

    print("Begin training...")
    for epoch in range(500):
        opt.zero_grad()
        output = model()
        loss = output.sum()
        loss.backward()
        opt.step()
        if epoch % 25 == 0:
            print("epoch {:>3d}: loss = {:>8.3f}".format(epoch, loss))

你可能感兴趣的:(Ai,机器学习,深度学习,人工智能,cuda,nvcc)