快捷導(dǎo)航

MindSpore導(dǎo)入CUDA算子的解決方案

更新時間：2024年05月09日 09:24:37 作者：DECHIN

本文介紹了在MindSpore標(biāo)準(zhǔn)格式下進行CUDA算子開發(fā)的方法和流程,可以讓開發(fā)者在現(xiàn)有的AI框架下仍然可以調(diào)用基于CUDA實現(xiàn)的高性能的算子,感興趣的朋友跟隨小編一起看看吧

本文介紹了在MindSpore標(biāo)準(zhǔn)格式下進行CUDA算子開發(fā)的方法和流程，可以讓開發(fā)者在現(xiàn)有的AI框架下仍然可以調(diào)用基于CUDA實現(xiàn)的高性能的算子。并且，除了常規(guī)的數(shù)值計算之外，在MindSpore框架下，我們還可以通過實現(xiàn)一個bprop函數(shù)，使得我們手寫的這個CUDA算子也可以使用MindSpore框架本身自帶的自動微分-端到端微分技術(shù)。

技術(shù)背景

當(dāng)今眾多的基于Python的AI框架（如MindSpore、PyTorch等）給了開發(fā)者非常便利的編程的條件，我們可以用Python的簡單的語法寫代碼，然后由框架在后端自動編譯成可以在GPU上高效計算的程序。而對于一些定制化比較高的算法，MindSpore也支持了相關(guān)的接口，允許開發(fā)者自己開發(fā)相應(yīng)的CUDA算子（需要統(tǒng)一接口），然后編譯成.so動態(tài)鏈接庫，再用MindSpore內(nèi)置的函數(shù)加載為本地算子。本文針對這種方案寫一個簡單的示例。

程序結(jié)構(gòu)

本地自己手寫一個CUDA算子，一般至少需要兩個文件和一個nvcc的環(huán)境，最好是在安裝完成MindSpore的GPU版本之后，再嘗試CUDA算子的引入。具體MindSpore的安裝方法，可以參考MindSpore官網(wǎng)，這里不做贅述。我這里使用的環(huán)境是10.1版本的nvcc：

$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Sun_Jul_28_19:07:16_PDT_2019
Cuda compilation tools, release 10.1, V10.1.243
$ python3 -m pip show mindspore
Name: mindspore
Version: 2.1.0
Summary: MindSpore is a new open source deep learning training/inference framework that could be used for mobile, edge and cloud scenarios.
Home-page: https://www.mindspore.cn
Author: The MindSpore Authors
Author-email: contact@mindspore.cn
License: Apache 2.0
Location: /home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages
Requires: numpy, protobuf, pillow, astunparse, scipy, psutil, asttokens, packaging
Required-by:

需要準(zhǔn)備的兩個文件，一個是CUDA算子本身的.cu文件，另一個是用來調(diào)用CUDA算子的.py文件。操作流程是：先按照自己的需求寫好CUDA算子，然后用nvcc進行編譯，編譯輸出為.so的動態(tài)鏈接庫，然后在python腳本中使用mindspore.ops.Custom生成相應(yīng)的算子。在MindSpore2.1之后的版本中，對于本地CUDA算子的調(diào)用定義了統(tǒng)一的接口，其格式為：

extern "C" int CustomFunc(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra);

具體方法可以參考官網(wǎng)的這一篇文檔說明。這樣的話，在.cu文件中至少有兩個函數(shù)，一個是原本用于計算的Kernel函數(shù)，另一個是用于統(tǒng)一標(biāo)準(zhǔn)接口的CustomFunc函數(shù)。需要說明的是，舊版本的MindSpore是沒有這樣的規(guī)范的，所以舊版本的算子沒有CustomFunc函數(shù)也能夠用nvcc編譯，但是無法在新版本的MindSpore中調(diào)用。

一維張量求和

我們用一個一維的張量求和的示例來演示一下如何在本地寫一個可以用MindSpore來調(diào)用的CUDA算子，一維張量求和的算法是比較簡單的：

那么對應(yīng)的CUDA算子的代碼如下所示：

// custom_add.cu
// nvcc --shared -Xcompiler -fPIC -o custom_add.so custom_add.cu
// 常量，一般可以放在.cuh頭文件中
constexpr int THREADS = 1024;
// 用于CUDA計算的Kernel函數(shù)
__global__ void CustomAddKernel(float *input1, float *input2, float *output, size_t size) {
    auto idx = blockIdx.x * THREADS + threadIdx.x;
    if (idx < size) {
        // 逐元素操作，CUDA算子的基本寫法
        output[idx] = input1[idx] + input2[idx];
    }
}
// 標(biāo)準(zhǔn)算子接口
extern "C" int CustomAdd(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream,
                         void *extra) {
    cudaStream_t custream = static_cast<cudaStream_t>(stream);
    // 輸出變量的位置
    constexpr int OUTPUT_INDEX = 2;
    // 傳入的參數(shù)都是指針形式
    float *input1 = static_cast<float *>(params[0]);
    float *input2 = static_cast<float *>(params[1]);
    float *output = static_cast<float *>(params[2]);
    // 獲取輸入張量的大小
    int size = shapes[OUTPUT_INDEX][0];
    // GPU運算中的block和thread，一般要求block*thread大于或者等于size即可
    int blocks = ceil(size / THREADS) + 1;
    // 調(diào)用Kernel函數(shù)
    CustomAddKernel<<<blocks, THREADS, 0, custream>>>(input1, input2, output, size);
    return 0;
}

值得注意的是，上述CustomAdd函數(shù)中的params有3個輸入，但是實際上其中一個是返回值，這也是MindSpore對于標(biāo)準(zhǔn)接口的設(shè)定，不需要我們額外傳入一個變量。保存好上述的CUDA算子代碼之后，可以用如下指令直接編譯成python可以調(diào)用的動態(tài)鏈接庫：

$ nvcc --shared -Xcompiler -fPIC -o custom_add.so custom_add.cu

編譯完成后，會在當(dāng)前目錄下生成一個新的.so文件，然后就可以在python代碼中進行調(diào)用：

# test_custom_ops.py
# python3 test_custom_ops.py
import mindspore as ms
from mindspore import ops, Tensor, context
ms.set_context(device_target="GPU", mode=context.GRAPH_MODE)
t1 = Tensor([1., 2., 3.], ms.float32)
t2 = Tensor([3., 2., 1.], ms.float32)
CustomAdd = ops.Custom("./custom_add.so:CustomAdd",
                       out_shape=[t1.shape[0]],
                       out_dtype=ms.float32,
                       func_type="aot"
                       )
res = CustomAdd(t1, t2)
ops.print_(res)

上述的CustomAdd就是我們導(dǎo)入的基于CUDA算子來寫的本地MindSpore算子，并且這種算子還可以使用MindSpore進行注冊，這樣就不需要每次使用都去加載這個動態(tài)鏈接庫，感興趣的童鞋可以自己研究一下算子注冊的方法。上述Python代碼的運行結(jié)果如下：

$ python3 test_custom_ops.py 
Tensor(shape=[3], dtype=Float32, value= [ 4.00000000e+00,  4.00000000e+00,  4.00000000e+00])

可見跟我們預(yù)期的結(jié)果是一致的，那么就完成了一個本地CUDA算子的實現(xiàn)和調(diào)用。

自定義算子反向傳播

在前面的章節(jié)里面我們已經(jīng)實現(xiàn)了一個本地的一維張量求和算子，但是這還不是一個完整的算子實現(xiàn)，因為我們在AI框架中一般要求算子可自動微分，光實現(xiàn)一個算子是不完整的，因此這里我們再通過bprop函數(shù)，來實現(xiàn)一下自定義算子的反向傳播：

# test_custom_ops.py
# python3 test_custom_ops.py
import mindspore as ms
from mindspore import ops, Tensor, context
from mindspore import nn, grad
ms.set_context(device_target="GPU", mode=context.GRAPH_MODE)
t1 = Tensor([1., 2., 3.], ms.float32)
t2 = Tensor([3., 2., 1.], ms.float32)
CustomAdd = ops.Custom("./custom_add.so:CustomAdd",
                       out_shape=[t1.shape[0]],
                       out_dtype=ms.float32,
                       func_type="aot"
                       )
為了自動微分，我們需要定義一個Cell類來封裝我們的自定義算子：
class Add(nn.Cell):
    # 反向傳播函數(shù)
    def bprop(self, x, y, out, dout):
        return (y, )
    # 計算函數(shù)
    def construct(self, x, y):
        return CustomAdd(x, y)
# 把Cell類加載為custom_add函數(shù)
custom_add = Add()
# 計算求和結(jié)果
res = custom_add(t1, t2)
# 計算自動微分結(jié)果
res_g = grad(custom_add, grad_position=(0, ))(t1, t2)
print(res)
print(res_g)

在這個代碼中，主要就是增加了一個Cell類和兩個新的函數(shù)bprop和construct，分別用于計算函數(shù)的反向傳播和正向值，代碼運行的結(jié)果如下：

$ python3 test_custom_ops.py 
[4. 4. 4.]
[3. 2. 1.]

當(dāng)然，這里我們沒有再額外寫一個用于返回反向傳播值的CUDA算子，但是原則上對于較為復(fù)雜的函數(shù)，是需要自己手動寫一個用于求微分?jǐn)?shù)值的CUDA算子的。

總結(jié)概要

本文介紹了在MindSpore標(biāo)準(zhǔn)格式下進行CUDA算子開發(fā)的方法和流程，可以讓開發(fā)者在現(xiàn)有的AI框架下仍然可以調(diào)用基于CUDA實現(xiàn)的高性能的算子。并且，除了常規(guī)的數(shù)值計算之外，在MindSpore框架下，我們還可以通過實現(xiàn)一個bprop函數(shù)，使得我們手寫的這個CUDA算子也可以使用MindSpore框架本身自帶的自動微分-端到端微分技術(shù)。

版權(quán)聲明

本文首發(fā)鏈接為：https://www.cnblogs.com/dechinphy/p/mindspore-cuda.html

作者ID：DechinPhy

更多原著文章：https://www.cnblogs.com/dechinphy/

請博主喝咖啡：https://www.cnblogs.com/dechinphy/gallery/image/379634.html

到此這篇關(guān)于MindSpore導(dǎo)入CUDA算子的文章就介紹到這了,更多相關(guān)MindSpore導(dǎo)入CUDA內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: