技术背景

在前面的几篇
博客
中,我们介绍了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

标签: none

添加新评论