MindSpore自定义算子中的张量维度问题

技术背景

在前面的几篇博客中，我们介绍了MindSpore框架下使用CUDA来定义本地算子的基本方法，以及配合反向传播函数的使用，这里主要探讨一下MindSpore框架对于CUDA本地算子的输入输出的规范化形式。

测试思路

MindSpore使用的CUDA算子规范化接口形式为：

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

也就是说，我们在一个.cu文件中按照这种形式写好函数接口，其中主要是规范化输入输出的形式，然后再将各项输入传给写好的CUDA Kernel函数进行计算并获得返回值。

我们可以使用一个Kernel打印函数的测试案例来说明MindSpore对于输入输出的处理：

#include <iostream>
#define THREADS 1024

__global__ void OpsKernel(const int shape0, const int *input){
    auto i = blockIdx.x * THREADS + threadIdx.x;
    if (i < shape0){
        printf("%d\n", input[i]);
    }
}

在这个函数体内，会把指定大小范围内的input的内容打印出来。

常数输入

首先我们来看一下最简单的常数输入，可以用一个最简单的整数来测试，对应的CUDA算子代码为：

// nvcc --shared -Xcompiler -fPIC -o test_shape.so test_shape.cu

#include <iostream>
#define THREADS 1024

__global__ void OpsKernel(const int shape0, const int *input){
    auto i = blockIdx.x * THREADS + threadIdx.x;
    if (i < shape0){
        printf("%d\n", input[i]);
    }
}

extern "C" int CustomOps(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes,
                         void *stream, void *extra){
    int *input = static_cast<int*>(params[0]);
    OpsKernel<<<1, THREADS>>>(shapes[0][0], input);
    return 0;
}

调用CUDA算子的Python代码为：

import os
import numpy as np
import mindspore as ms
from mindspore import ops, Tensor, context

context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

CURRENT_PATH = os.path.abspath(__file__)
CustomOps = ops.Custom(CURRENT_PATH.replace(".py", ".so:CustomOps"),
                       out_shape=lambda x:x,
                       out_dtype=ms.int32,
                       func_type="aot")
T0 = Tensor([7], ms.int32)
print (T0)
CustomOps(T0)

运行的指令为：

$ nvcc --shared -Xcompiler -fPIC -o test_shape.so test_shape.cu && python3 test_shape.py
[7]
7

需要注意的是，这里只能给MindSpore内置的几种Tensor变量，如果是直接调用CustomOps(7)会报一个段错误。

高维张量输入

这里一维的张量输入我们就不做讨论了，因为跟前面用到的常数输入本质上是一样的形式。这里我们用一个二维的张量来做一个测试，CUDA代码保持不动，只修改Python代码中的输入：

import os
import numpy as np
import mindspore as ms
from mindspore import ops, Tensor, context

context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

CURRENT_PATH = os.path.abspath(__file__)
CustomOps = ops.Custom(CURRENT_PATH.replace(".py", ".so:CustomOps"),
                       out_shape=lambda x:x,
                       out_dtype=ms.int32,
                       func_type="aot")
T0 = Tensor(np.arange(12).reshape((4, 3)), ms.int32)
print (T0)
CustomOps(T0)

运行结果为：

$ nvcc --shared -Xcompiler -fPIC -o test_shape.so test_shape.cu && python3 test_shape.py
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]
0
1
2
3

需要注意的是，我们在CUDA的打印函数中设置的打印输出大小是输入张量的第一个维度的大小，我们给的是一个(4,3)大小的张量，因此会顺序打印4个数出来。这里我们也能够发现MindSpore在进行输入的规范化的时候，会自动压平输入的张量变成一个维度。因此这里的调用代码等价于先对输入张量做一个reshape，然后再把第一个维度对应大小的张量元素打印出来。如果要打印所有的元素也很简单，可以修改一下CUDA代码：

// nvcc --shared -Xcompiler -fPIC -o test_shape.so test_shape.cu

#include <iostream>
#define THREADS 1024

__global__ void OpsKernel(const int shape0, const int *input){
    auto i = blockIdx.x * THREADS + threadIdx.x;
    if (i < shape0){
        printf("%d\n", input[i]);
    }
}

extern "C" int CustomOps(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes,
                         void *stream, void *extra){
    int *input = static_cast<int*>(params[0]);
    int elements = 1;
    for (int i=0; i<ndims[0]; i++){
        elements *= shapes[0][i];
    }
    OpsKernel<<<1, THREADS>>>(elements, input);
    return 0;
}

通过定义一个elements变量用于存储对应张量的元素数量，然后再逐一打印出来即可，执行结果为：

$ nvcc --shared -Xcompiler -fPIC -o test_shape.so test_shape.cu && python3 test_shape.py
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]
0
1
2
3
4
5
6
7
8
9
10
11

输出规范化

当我们使用ops.Custom算子时，如果指定了out_dtype和out_shape，那么算子会自动帮我们分配好相应的device memory空间。那么我们在CUDA计算的时候可以直接修改对应的内存空间：

// nvcc --shared -Xcompiler -fPIC -o test_shape.so test_shape.cu

#include <iostream>
#define THREADS 1024

__global__ void OpsKernel(const int shape0, const int *input, float *output){
    auto i = blockIdx.x * THREADS + threadIdx.x;
    if (i < shape0){
        output[i] = input[i] * 0.5;
    }
}

extern "C" int CustomOps(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes,
                         void *stream, void *extra){
    int *input = static_cast<int*>(params[0]);
    float *output = static_cast<float*>(params[1]);
    int elements = 1;
    for (int i=0; i<ndims[0]; i++){
        elements *= shapes[0][i];
    }
    OpsKernel<<<1, THREADS>>>(elements, input, output);
    return 0;
}

这里我们对算子的功能做了一点调整，我们输出的结果是整个张量的元素值乘以0.5，同时也把一个整形变量转化成了一个浮点型变量。其运行Python代码也要做一点调整：

import os
import numpy as np
import mindspore as ms
from mindspore import ops, Tensor, context

context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

CURRENT_PATH = os.path.abspath(__file__)
CustomOps = ops.Custom(CURRENT_PATH.replace(".py", ".so:CustomOps"),
                       out_shape=lambda x:x,
                       out_dtype=ms.float32,
                       func_type="aot")
T0 = Tensor(np.arange(12).reshape((4, 3)), ms.int32)
print (T0)
output = CustomOps(T0)
print (output)

这里主要是修改了out_dtype为浮点型，这里如果写错了，会直接导致内存溢出。上述代码的运行结果为：

$ nvcc --shared -Xcompiler -fPIC -o test_shape.so test_shape.cu && python3 test_shape.py
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]
[[0.  0.5 1. ]
 [1.5 2.  2.5]
 [3.  3.5 4. ]
 [4.5 5.  5.5]]

可以看到这里输出的张量形状是跟输入保持一致的，即时这个输入张量在经过MindSpore的Custom算子接口时已经被压平成一个一维张量，但是因为我们设置了out_shape=lambda x:x，这表示输出的张量shape跟输入的张量shape一致，当然，直接用Python的列表来给out_shape赋值也是可以的。例如我们写一个输入输出不同shape的案例：

// nvcc --shared -Xcompiler -fPIC -o test_shape.so test_shape.cu

#include <iostream>
#define THREADS 1024

__global__ void OpsKernel(const int shape0, const int *input, int *output){
    auto i = blockIdx.x * THREADS + threadIdx.x;
    if (i < shape0){
        atomicAdd(&output[0], input[i]);
    }
}

extern "C" int CustomOps(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes,
                         void *stream, void *extra){
    int *input = static_cast<int*>(params[0]);
    int *output = static_cast<int*>(params[1]);
    int elements = 1;
    for (int i=0; i<ndims[0]; i++){
        elements *= shapes[0][i];
    }
    OpsKernel<<<1, THREADS>>>(elements, input, output);
    return 0;
}

这个Kernel函数的主要功能是通过一个atomicAdd函数，把输入张量的所有元素做一个求和，这样输出的张量的shape只有[1]，对应的Python调用形式也要做一定的调整：

import os
import numpy as np
import mindspore as ms
from mindspore import ops, Tensor, context

context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

CURRENT_PATH = os.path.abspath(__file__)
CustomOps = ops.Custom(CURRENT_PATH.replace(".py", ".so:CustomOps"),
                       out_shape=[1],
                       out_dtype=ms.int32,
                       func_type="aot")
T0 = Tensor(np.arange(12).reshape((4, 3)), ms.int32)
print (T0)
output = CustomOps(T0)
print (output)

由于atomicAdd(addr, element)原子操作要求输入输出的类型要一致，因此这里我们还是使用的int类型的output，输出结果如下所示：

$ nvcc --shared -Xcompiler -fPIC -o test_shape.so test_shape.cu && python3 test_shape.py
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]
[66]

总结概要

当我们使用GPU进行快速运算时，虽然可以用成熟的深度学习框架如MindSpore和PyTorch等进行实现，但其实从速度上来说，最快不过直接使用C/C++的CUDA来实现。也正是因为如此，在MindSpore框架中支持了对CUDA实现的算子的直接调用，只是在格式规范上有一定的要求。本文主要介绍MindSpore调用本地CUDA算子的一些规范化和技巧。

版权声明

本文首发链接为：https://www.cnblogs.com/dechinphy/p/custom-ops-shape.html

作者ID：DechinPhy

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

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