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