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

官方文档：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   
#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)：多核执行不同的指令处理多个数据

为了提高并行的计算能力，架构上实现下面这些性能提升：

• 降低延迟（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
#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))