MindSpore导入CUDA算子的解决方案

$ 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: 

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

// custom_add.cu
// nvcc --shared -Xcompiler -fPIC -o custom_add.so custom_add.cu
// 常量，一般可以放在.cuh头文件中
constexpr int THREADS = 1024;
// 用于CUDA计算的Kernel函数
__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;
    // 传入的参数都是指针形式
    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;
    // 调用Kernel函数
    CustomAddKernel<<<blocks, THREADS, 0, custream>>>(input1, input2, output, size);
    return 0;
}

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

# 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)

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

# 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):
    # 反向传播函数
    def bprop(self, x, y, out, dout):
        return (y, )
    # 计算函数
    def construct(self, x, y):
        return CustomAdd(x, y)
# 把Cell类加载为custom_add函数
custom_add = Add()
# 计算求和结果
res = custom_add(t1, t2)
# 计算自动微分结果
res_g = grad(custom_add, grad_position=(0, ))(t1, t2)
print(res)
print(res_g)

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