MindSpore導入CUDA算子的解決方案
本文介紹了在MindSpore標準格式下進行CUDA算子開發(fā)的方法和流程,可以讓開發(fā)者在現(xiàn)有的AI框架下仍然可以調(diào)用基于CUDA實現(xiàn)的高性能的算子。并且,除了常規(guī)的數(shù)值計算之外,在MindSpore框架下,我們還可以通過實現(xiàn)一個bprop函數(shù),使得我們手寫的這個CUDA算子也可以使用MindSpore框架本身自帶的自動微分-端到端微分技術。
技術背景
當今眾多的基于Python的AI框架(如MindSpore、PyTorch等)給了開發(fā)者非常便利的編程的條件,我們可以用Python的簡單的語法寫代碼,然后由框架在后端自動編譯成可以在GPU上高效計算的程序。而對于一些定制化比較高的算法,MindSpore也支持了相關的接口,允許開發(fā)者自己開發(fā)相應的CUDA算子(需要統(tǒng)一接口),然后編譯成.so
動態(tài)鏈接庫,再用MindSpore內(nèi)置的函數(shù)加載為本地算子。本文針對這種方案寫一個簡單的示例。
程序結構
本地自己手寫一個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:
需要準備的兩個文件,一個是CUDA算子本身的.cu
文件,另一個是用來調(diào)用CUDA算子的.py
文件。操作流程是:先按照自己的需求寫好CUDA算子,然后用nvcc進行編譯,編譯輸出為.so
的動態(tài)鏈接庫,然后在python腳本中使用mindspore.ops.Custom
生成相應的算子。在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)一標準接口的CustomFunc函數(shù)
。需要說明的是,舊版本的MindSpore是沒有這樣的規(guī)范的,所以舊版本的算子沒有CustomFunc函數(shù)
也能夠用nvcc編譯,但是無法在新版本的MindSpore中調(diào)用。
一維張量求和
我們用一個一維的張量求和的示例來演示一下如何在本地寫一個可以用MindSpore來調(diào)用的CUDA算子,一維張量求和的算法是比較簡單的:
那么對應的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]; } } // 標準算子接口 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對于標準接口的設定,不需要我們額外傳入一個變量。保存好上述的CUDA算子代碼之后,可以用如下指令直接編譯成python可以調(diào)用的動態(tài)鏈接庫:
$ nvcc --shared -Xcompiler -fPIC -o custom_add.so custom_add.cu
編譯完成后,會在當前目錄下生成一個新的.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就是我們導入的基于CUDA算子來寫的本地MindSpore算子,并且這種算子還可以使用MindSpore進行注冊,這樣就不需要每次使用都去加載這個動態(tài)鏈接庫,感興趣的童鞋可以自己研究一下算子注冊的方法。上述Python代碼的運行結果如下:
$ python3 test_custom_ops.py Tensor(shape=[3], dtype=Float32, value= [ 4.00000000e+00, 4.00000000e+00, 4.00000000e+00])
可見跟我們預期的結果是一致的,那么就完成了一個本地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() # 計算求和結果 res = custom_add(t1, t2) # 計算自動微分結果 res_g = grad(custom_add, grad_position=(0, ))(t1, t2) print(res) print(res_g)
在這個代碼中,主要就是增加了一個Cell類和兩個新的函數(shù)bprop
和construct
,分別用于計算函數(shù)的反向傳播和正向值,代碼運行的結果如下:
$ python3 test_custom_ops.py [4. 4. 4.] [3. 2. 1.]
當然,這里我們沒有再額外寫一個用于返回反向傳播值的CUDA算子,但是原則上對于較為復雜的函數(shù),是需要自己手動寫一個用于求微分數(shù)值的CUDA算子的。
總結概要
本文介紹了在MindSpore標準格式下進行CUDA算子開發(fā)的方法和流程,可以讓開發(fā)者在現(xiàn)有的AI框架下仍然可以調(diào)用基于CUDA實現(xiàn)的高性能的算子。并且,除了常規(guī)的數(shù)值計算之外,在MindSpore框架下,我們還可以通過實現(xiàn)一個bprop函數(shù),使得我們手寫的這個CUDA算子也可以使用MindSpore框架本身自帶的自動微分-端到端微分技術。
版權聲明
本文首發(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
到此這篇關于MindSpore導入CUDA算子的文章就介紹到這了,更多相關MindSpore導入CUDA內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關文章希望大家以后多多支持腳本之家!
相關文章
解決Python出現(xiàn)_warn_unsafe_extraction問題的方法
這篇文章主要為大家詳細介紹了解決Python出現(xiàn)'_warn_unsafe_extraction'問題的方法,感興趣的小伙伴們可以參考一下2016-03-03在python中利用最小二乘擬合二次拋物線函數(shù)的方法
今天小編就為大家分享一篇在python中利用最小二乘擬合二次拋物線函數(shù)的方法,具有很好的參考價值,希望對大家有所幫助。一起跟隨小編過來看看吧2018-12-12python循環(huán)控制之break和continue流程控制語句
這篇文章主要介紹了python循環(huán)控制之break流程控制語句,Python中提供了兩個關鍵字用來控制循環(huán)語句,分別是break和continue,本文都有介紹,需要的朋友可以參考一下2022-03-03