官方文档：https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html
什么是CUDA
- CUDA全称（Compute Unified Device Architecture）统一计算架构，是NVIDIA推出的并行计算平台
- 深度学习加速：对于神经网络，无论是离线训练还是在线推理，都有巨量的矩阵、归一化、softmax等运算，且其中有非常多的并行计算，非常适合用GPU来进行运算加速
- 一般来说，应用程序混合有并行部分和顺序部分，因此系统设计时混合使用 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个线程运行
限定符
因为数据在全局内存中是线性存储的，所以可以通过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 <iostream>
#include <cuda_runtime.h>
// 代码的核心诉求(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)：多核执行不同的指令处理多个数据
为了提高并行的计算能力，架构上实现下面这些性能提升：
- 降低延迟（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
- 内存释放
- 分配内存，数据初始化
内存管理
| 标准C函数 | CUDA 函数 | CUDA函数说明 |
|---|---|---|
| malloc | cudaMalloc | GPU内存分配 |
| memcpy | cudaMemcpy | 用于Host和Device之间数据传输 |
| memset | cudaMemset | 设定数据填充到GPU内存中 |
| free | cudaFree | 释放GPU内存 |
- CUDA 线程在执行期间可以访问多个内存空间中的数据
- 每个线程都有私有本地内存
- 每个线程块都有对该块的所有线程可见的共享内存，并且与该块具有相同的生命周期
- 线程块簇中的线程块可以对彼此的共享内存执行读、写和原子操作。所有线程都可以访问相同的全局内存。
- 还有两个可供所有线程访问的附加只读内存空间：常量内存空间和纹理内存空间
Grid&&Block
- 一个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分量向量，因此可以使用一维、二维或三维线程索引 来标识线程，形成一维、二维或三维线程块，称为线程块。这提供了一种自然的方式来调用域中元素（例如向量、矩阵或体积）的计算
每个块的线程数量是有限的，因为块中的所有线程都应驻留在同一个流式多处理器核心上，并且必须共享该核心的有限内存资源。在当前的 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);
...
}
线程块簇
- 随着 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内存
- 寄存器：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<iostream>
#include<cuda.h>
#include<cuda_runtime.h>
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 <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");
} // 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))
