自定义算子（基于Custom表达） — MindSpore master 文档

link管理
链接快照平台
输入网页链接，自动生成快照
标签化管理网页链接
自定义算子入门 

为了帮助用户快速入门自定义算子，这里以pyfunc类型自定义算子为例帮助用户理解自定义算子的定义流程。下面基于pyfunc模式定义一个实现sin计算的自定义算子。pyfunc类型的自定义算子使用原生Python语法定义算子实现函数，描述算子内部计算逻辑的实现。网络运行时框架会自动调用此函数。为了表达自定义算子的计算，我们写一个基于numpy的计算正弦函数的Python原生函数。
import numpy as np
def sin_by_numpy(x):
    return np.sin(x)
然后我们要定义两个函数，一个是张量形状的推导函数（infer_shape），另一个是张量数据类型的推导函数（infer_dtype）。这里要注意：
张量形状的推导函数是输入张量的形状；
张量数据类型的推导函数是输入张量的数据类型。
def infer_shape(x):
    #    1. 这里的输入x是算子输入张量的形状
    #    2. sin函数是逐元素计算，输入的形状和输出的一样
    return x
def infer_dtype(x):
    #    1. 这里的输入x是算子输入张量的数据类型
    #    2. sin函数输入的数据类型和输出的一样
    return x
下面我们用上面的函数自定义一个算子，其输入包括
func：自定义算子的函数表达，这里我们用sin_by_numpy函数；
out_shape: 输出形状的推导函数，这里我们用infer_shape函数；
out_dtype: 输出数据类型的推导函数，这里我们用infer_dtype函数；
func_type: 自定义算子类型，这里我们用"pyfunc"。
from mindspore import ops
sin_by_numpy_op = ops.Custom(func=sin_by_numpy,     # 这里填入自定义算子的函数表达
                             out_shape=infer_shape, # 这里填入输出形状的推导函数
                             out_dtype=infer_dtype, # 这里填入输出数据类型的推导函数
                             func_type="pyfunc"     # 这里填入自定义算子类型
加上其他环境依赖依赖和算子调用语句，我们获得完整的自定义算子用例如下。
import numpy as np
import mindspore as ms
from mindspore import ops
ms.set_context(mode=ms.GRAPH_MODE, device_target="CPU")
def sin_by_numpy(x):
    return np.sin(x)
def infer_shape(x):
    return x
def infer_dtype(x):
    return x
sin_by_numpy_op = ops.Custom(func=sin_by_numpy,
                             out_shape=infer_shape,
                             out_dtype=infer_dtype,
                             func_type="pyfunc")
input_tensor = ms.Tensor([0, 1, 0.2, 0.3, 0.4], dtype=ms.float32)
result_cus = sin_by_numpy_op(input_tensor)
print(result_cus)
我们可以得到结果为，即上面输入对应的sin值。
如此我们完成一个pyfunc类型自定义算子的定义。对于更多完整的pyfunc类型自定义算子的例子，参见MindSpore源码中的用例。
采用JIT编译的自定义算子
JIT（Just In Time）指算子在网络编译或运行期间被框架直接编译。用户可以直接用Python脚本在网络脚本中直接定义此种类型的自定义算子，然后根据算子和后端类型调用对应算子编译器自动编译。此种类型的自定义算子定义方便，而且有着更好的后端适应性。
Hybrid类型的自定义算子开发
Hybrid类型的自定义算子是自定义算子的默认定义类型。通过使用Hybrid类型的自定义算子，用户可以用类Python的语法描述算子计算逻辑，且无需关注MindSpore框架对于算子定义的工程细节，让用户专注于算法本身。
Hybrid类型的自定义算子使用MindSpore Hybrid DSL描述算子内部计算逻辑的实现。用MindSpore Hybrid DSL定义的函数可以被AKG算子编译器解析进行JIT编译生成高效算子，在大规模模型的训练推理中使用。同时，用MindSpore Hybrid DSL定义的函数可以当做一个numpy函数直接调用，方便用户调试的同时也可以灵活的切换到pyfunc
类型的自定义算子，做到一次开发，多个模式多个平台多个场景复用的自定义算子表达。
下面用例(test_custom_hybrid.py)介绍hybrid类型的自定义算子开发流程，其中自定义算子实现两个输入张量相加的功能。 值得注意的是，Hybrid类型的自定义算子采取源码变换的方式打通MindSpore的图编译器和算子编译器，用户可以直接使用MindSpore Hybrid DSL提供的关键词，例如下面的output_tensor，而无需引入对应Python函数。更多MindSpore Hybrid DSL关键词的介绍，参见MindSpore Hybrid DSL关键词。
import numpy as np
from mindspore import ops
import mindspore as ms
from mindspore.ops import kernel
ms.set_context(device_target="CPU")
# 算子实现，Hybrid DSL
@kernel
def add(a, b):
    c = output_tensor(a.shape, a.dtype)
    for i0 in range(a.shape[0]):
        for i1 in range(a.shape[1]):
            c[i0, i1] = a[i0, i1] + b[i0, i1]
    return c
if __name__ == "__main__":
    # 定义hybrid类型的自定义算子(Custom的默认模式)
    op = ops.Custom(add)
    x0 = np.array([[0.0, 0.0], [1.0, 1.0]]).astype(np.float32)
    x1 = np.array([[2.0, 2.0], [3.0, 3.0]]).astype(np.float32)
    output = op(ms.Tensor(x0), ms.Tensor(x1))
    print(output)
Hybrid类型是Custom的默认类型。
Hybrid类型自定义算子的输入必须是一个带有[@kernel](https://www.mindspore.cn/docs/zh-CN/br_base/api_python/ops/mindspore.ops.kernel.html)的函数。
Hybrid类型自定义算子定义时可以使用自带的自动shape/dtype推导函数，也可以手动输入shape/dtype推导函数。
执行用例：
python test_custom_hybrid.py
执行结果：
[[2. 2.]
 [4. 4.]]
对于更多完整的hybrid类型自定义算子的例子，参见MindSpore源码中的用例。
akg类型的自定义算子开发
akg类型的自定义算子使用MindSpore AKG算子DSL，描述算子内部计算逻辑的实现。MindSpore AKG是基于TVM（Tensor Virtual Machine）和Polyhedral技术的算子开发和编译框架，支持Hybrid、IR builder和TVM compute等多种类型的算子DSL。
算子输出shape和数据类型推理可以通过定义Python函数实现，描述算子输出shape和数据类型的推导逻辑。
若算子包含属性或者只支持特定的输入输出数据类型或数据格式，则需要注册算子信息，算子信息生成方式请参考算子信息注册。若未注册算子信息，在后端做算子选择和映射的时候，将会从当前算子的输入中推导算子信息。
下面以test_custom_akg.py为例介绍akg类型的自定义算子开发流程，其中自定义算子实现两个输入张量相加的功能。
test_custom_akg.py内容：
import numpy as np
import mindspore as ms
import mindspore.ops as ops
ms.set_context(device_target="CPU")
# 算子实现，Hybrid DSL
def add(a, b):
    c = output_tensor(a.shape, a.dtype)
    for i0 in range(a.shape[0]):
        for i1 in range(a.shape[1]):
            c[i0, i1] = a[i0, i1] + b[i0, i1]
    return c
if __name__ == "__main__":
    # 定义akg类型的自定义算子
    op = ops.Custom(add, out_shape=lambda x, _: x, out_dtype=lambda x, _: x, func_type="akg")
    x0 = np.array([[0.0, 0.0], [1.0, 1.0]]).astype(np.float32)
    x1 = np.array([[2.0, 2.0], [3.0, 3.0]]).astype(np.float32)
    output = op(ms.Tensor(x0), ms.Tensor(x1))
    print(output)
set_context(device_target="GPU")表示算子运行在GPU平台。




    

用Python lambda函数定义输出shape和数据类型推理函数，并分别传给Custom原语的out_shape和out_dtype参数。本例中lambda函数表明输出shape和数据类型和第一个输入张量的信息相同。
未注册算子信息，所以自定义算子的算子信息将会从算子输入中推理。
执行用例：
python test_custom_akg.py
执行结果：
[[2. 2.]
 [4. 4.]]
对于更多完整的akg类型自定义算子的例子，参见MindSpore源码中的用例。
采用AOT编译的自定义算子
AOT类型的自定义算子指用户事先把算子编译成二进制文件后接入网络。通常用户通过C/C++/CUDA等编程语言手工优化算子实现，并把算子以动态库的形式接入MindSpore加速网络。如此，用户可以针对算子进行极致优化，发挥对应后端硬件的极致性能。这里我们会介绍AOT类型自定义算子的一些基础知识，对于AOT类型自定义算子的更多用法和功能，请参见AOT类型自定义算子进阶用法
aot类型的自定义算子开发
aot类型的自定义算子采用AOT编译方式，要求网络开发者基于特定接口，手写算子实现函数对应的源码文件，并提前将源码文件编译为动态链接库，然后在网络运行时框架会自动调用执行动态链接库中的函数。在算子实现的开发语言方面，GPU平台支持CUDA，CPU平台支持C和C++。源码文件中的算子实现函数的接口规范如下：
extern "C" int CustomFunc(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra);
其中，函数名CustomFunc可替换成任意有效函数名。返回值为int类型，约定0表示正常退出，非0表示发生异常。参数列表的含义如下：
nparam (int): 输入输出总数。比如算子有2个输入，1个输出，则nparam的值为3。
params (void **): 输入输出指针数组。比如算子有2个输入，1个输出，params[0]指向第一个输入数据，params[1]指向第二个输入数据，params[2]指向输出数据。
ndims (int *): 输入输出shape维度数组。比如params[i]是个shape[1024, 1024]的张量，则ndims[i]的值为2。
shapes (int64_t **): 输入输出shape数组。比如params[i]是个shape[1024, 1024]的张量，则shapes[i][0]的值为1024，shapes[i][1]的值为1024。
dtypes (const char **): 输入输出数据类型数组。dtypes里的元素取值可为：“float32”, “float16”, “float”, “float64”, “int”, “int8”, “int16”, “int32”, “int64”, “uint”, “uint8”, “uint16”, “uint32”, “uint64”, “bool”。
stream (void *): CUDA流指针，仅定义GPU算子实现时需要。
extra (void *): 用于后续扩展。
在Python脚本中，Custom接口中的func输入的格式为Path_To_Func:CustomFunc，其中CustomFunc为上面函数的名字，而Path_To_Func为对应函数源文件或者二进制库的地址。
MindSpore识别自动编译的方式为文件名后缀。为了使用自动编译功能，请使用后缀为cpp、cc或者cu的源文件。其他情况MindSpore将处理为二进制库的路径；
为了防止恶意第三方库篡改，请在环境变量MS_CUSTOM_AOT_WHITE_LIST设置合法第三方库的路径。只有在MS_CUSTOM_AOT_WHITE_LIST设置的目录及其子目录下文件才会被自定义算子调用。
算子输出shape和数据类型推理可以通过定义Python函数实现，描述算子输出shape和数据类型的推导逻辑。
若自定义算子只支持特定的输入输出数据类型，则需要定义算子信息，算子信息生成方式请参考算子信息注册。
下面通过例子介绍GPU平台和CPU平台上aot类型的自定义算子开发流程，其中自定义算子实现两个输入张量相加的功能。
GPU示例
使用CUDA语言，编写算子实现的源码文件add.cu：
#define THREADS 1024
__global__ void CustomAddKernel(float *input1, float *input2, float *output, size_t size) {
  auto idx = blockIdx.x * THREADS + threadIdx.x;
  if (idx < size) {
    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);
  if (nparam != 3) return 1;
  void *input1 = params[0];
  void *input2 = params[1];
  void *output = params[2];
  size_t size = 1;
  for (int i = 0; i < ndims[2]; i++) {
    size *= shapes[2][i];
  int n = size / THREADS;
  for (int i = 0; i < nparam; i++) {
    if (strcmp(dtypes[i], "float32") != 0) {
      return 2;
  CustomAddKernel<<<n + 1, THREADS, 0, custream>>>(static_cast<float *>(input1), static_cast<float *>(input2),
                                                   static_cast<float *>(output), size);
  return 0;
将add.cu编译成动态库add.so：
nvcc --shared -Xcompiler -fPIC -o add.so add.cu
编写测试用例test_custom_aot.py：
import numpy as np
import mindspore as ms
import mindspore.ops as ops
ms.set_context(device_target="GPU")
if __name__ == "__main__":
    # 定义aot类型的自定义算子
    op = ops.Custom("./add.so:CustomAdd", out_shape=lambda x, _: x, out_dtype=lambda x, _: x, func_type="aot")
    x0 = np.array([[0.0, 0.0], [1.0, 1.0]]).astype(np.float32)
    x1 = np.array([[2.0, 2.0], [3.0, 3.0]]).astype(np.float32)
    output = op(ms.Tensor(x0), ms.Tensor(x1))
    print(output)
本例中，有如下几点需要说明：
本例中需要将test_custom_aot.py和add.so放置在同一目录下，若add.so在其他目录，则需要将Custom第一个参数里路径修改为add.so的绝对路径。
用Python lambda函数定义输出shape和数据类型推理函数，并分别传给Custom原语的out_shape和out_dtype参数。本例中lambda函数表明输出shape和数据类型和第一个输入张量的信息相同。
未注册算子信息，所以自定义算子的算子信息将会从算子输入中推理。
执行用例：
python test_custom_aot.py
执行结果：
[[2. 2.]
 [4. 4.]]
#include <string.h>
using size_t = decltype(sizeof(int));
using int64_t = decltype(sizeof(long));
extern "C" int CustomAdd(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra) {
  if (nparam != 3) return 1;
  float *input1 = static_cast<float *>(params[0]);
  float *input2 = static_cast<float *>(params[1]);
  float *output = static_cast<float *>(params[2]);
  size_t size = 1;
  for (int i = 0; i < ndims[2]; i++) {
    size *= shapes[2][i];
  for (int i = 0; i < nparam; i++) {
    if (strcmp(dtypes[i], "float32") != 0) {
      return 2;
  for (int i = 0; i < size; i++) {
    output[i] = input1[i] + input2[i];
  return 0;
将add.cc编译成动态库add.so：
g++ --shared -fPIC -o add.so add.cc
编写测试用例test_custom_aot.py：
import numpy as np
import mindspore as ms
import mindspore.ops as ops
ms.set_context(device_target="CPU")
if __name__ == "__main__":
    # 定义aot类型的自定义算子
    op = ops.Custom("./add.so:CustomAdd", out_shape=lambda x, _: x, out_dtype=lambda x, _: x, func_type="aot")
    x0 = np.array([[0.0, 0.0], [1.0, 1.0]]).astype(np.float32)
    x1 = np.array([[2.0, 2.0], [3.0, 3.0]]).astype(np.float32)
    output = op(ms.Tensor(x0), ms.Tensor(x1))
    print(output)
本例中，有如下几点需要说明：
本例中需要将test_custom_aot.py和add.so放置在同一目录下，若add.so在其他目录，则需要将Custom第一个参数里路径修改为add.so的绝对路径。
用Python lambda函数定义输出shape和数据类型推理函数，并分别传给Custom原语的out_shape和out_dtype参数。本例中lambda函数表明输出shape和数据类型和第一个输入张量的信息相同。
未注册算子信息，所以自定义算子的算子信息将会从算子输入中推理。
执行用例：
python test_custom_aot.py
执行结果：
[[2. 2.]
 [4. 4.]]
对于更多完整的aot类型自定义算子的例子，参见MindSpore源码中的用例。
自定义算子接入第三方前端
作为MindSpore未来的发展方向之一，AI和科学计算的融合越来越受到业界的重视。MindSpore自定义算子基于自身表达的灵活性，也在科学计算方面做出了探索：把面向HPC的编程前端以自定义算子的方式接入MindSpore。
julia类型的自定义算子开发
Julia是一种速度快且使用简单的高级通用编程语言，最初设计用于科学计算领域，而由于其高效而实用的特性，近些年来越来越受到用户的青睐，逐步迈向主流编程语言。 julia类型的自定义算子使用Julia语法定义算子实现函数，描述算子内部计算逻辑的实现。网络运行时框架会自动调用执行相应的Julia函数。
算子输出shape和数据类型推导可以通过定义Python函数实现，描述算子输出shape和数据类型的推导逻辑。
若自定义算子只支持特定的输入输出数据类型，则需要定义算子信息，算子信息生成方式请参考算子信息注册。
下面以两个输入张量相加为例，介绍julia类型的自定义算子开发流程:
首先，用户需要通过单独文件实现Julia函数，如(add.jl)：
# add.jl
module Add
# inputs: x, y, output: z, output should use .= to inplace assign
function add(x, y, z)
    z .= x + y
其次，在网络脚本中通过自定义算子方式引用上面所写的Julia函数，以test_custom_julia.py为例：
import numpy as np
import mindspore as ms
import mindspore.ops as ops
ms.set_context(device_target="CPU")
if __name__ == "__main__":
    # 定义julia类型的自定义算子
    op = ops.Custom("./add.jl:Add:add", out_shape=lambda x, _: x, out_dtype=lambda x, _: x, func_type="julia")
    x0 = np.array([[0.0, 0.0], [1.0, 1.0]]).astype(np.float32)
    x1 = np.array([[2.0, 2.0], [3.0, 3.0]]).astype(np.float32)
    output = op(ms.Tensor(x0), ms.Tensor(x1))
    print(output)
本例中，有如下几点需要说明：
用Python lambda函数定义输出shape和数据类型推理函数，并分别传给Custom原语的out_shape和out_dtype参数。本例中lambda函数表明输出shape和数据类型和第一个输入张量的信息相同。
未注册算子信息，所以自定义算子的算子信息将会从算子输入中推理。
执行用例：
python test_custom_julia.py
执行结果：
[[2. 2.]
 [4. 4.]]
注意事项：
用户需确保下载正确版本的Julia，即version>=1.6.0。
由于运行时调用的Julia C api是从libjulia.so中获取的，因此需要用户设置julia/lib到LD_LIBRARY_PATH，以julia-1.6.5为例:
# download julia-1.6.5
wget https://julialang-s3.julialang.org/bin/linux/x64/1.6/julia-1.6.5-linux-x86_64.tar.gz
# for arm server
# wget https://julialang-s3.julialang.org/bin/linux/aarch64/1.6/julia-1.6.5-linux-aarch64.tar.gz
# extract file
tar xvf julia-1.6.5-linux-x86_64.tar.gz
# if $JULIA_DIR not exist
export LD_LIBRARY_PATH=$PWD/julia-1.6.5/lib:$LD_LIBRARY_PATH
# else
export LD_LIBRARY_PATH=$JULIA_DIR/lib:$LD_LIBRARY_PATH
Custom 第一个入参指定用户书写的Julia函数需按照file_name:module_name:func_name格式指定，file_name需包含文件路径，建议使用绝对路径。
Julia代码文件需包含module, module内包含function，且module/function都以end结束。
Julia函数的输入输出顺序需与算子的输入输出顺序一致。
Julia函数的最终输出，即kernel output的赋值需要使用.=，否则结果无法写入内存。
Julia代码支持Julia的常用语法，用户需自行保证语法正确，函数可正确执行。
用户想在Julia文件内使用Julia的第三方软件包，需自行下载对应软件以确保能正确调用，可以通过 import pkg; pkg.add("somepkg")进行安装。
julia array在内存上是column major排列的，而numpy array是row major排列的，如果Julia和numpy做比较，非elemwise计算需考虑内存排布。在Julia函数中，可以通过如下代码示例进行numpy array和julia array的相互转换:
function change_input_to_row_major(x)
    return permutedims(reshape(x, reverse(size(x))), length(size(x)):-1:1)
function change_output_to_row_major(x)
    return reshape(permutedims(x, length(size(x)):-1:1), size(x))
以矩阵乘为例：
# julia array is column-major, numpy array is row-major
# user should change julia or numpy's layout to keep same behavior
#= EXAMPLE
A[2,3]               B[3,4]               C[2,4]
NUMPY:
[[1, 2, 3]       [[1, 2, 3, 4]         [[38, 44, 50,  56]
 [4, 5, 6]]       [5, 6, 7, 8]          [83, 98, 113,128]]
                  [9,10,11,12]]
JULIA:
change_input_to_row_major:
1.inputs read numpy data from memory:
[[1, 3, 5]       [[1, 4, 7,10]
 [2, 4, 6]]       [2, 5, 8,11]
                  [3, 6, 9,12]]
2.inputs after reshape(reverse(shape)):
[[1, 4]          [[1, 5, 9]
 [2, 5]           [2, 6,10]
 [3, 6]]          [3, 7,11]
                  [4, 8,12]]
3.inputs after transpose/permutedims:
[[1, 2, 3]       [[1, 2, 3, 4]         [[38, 44, 50,  56]
 [4, 5, 6]]       [5, 6, 7, 8]          [83, 98, 113,128]]
                  [9,10,11,12]]
change_output_to_row_major:
1.output after transpose/permutedims:
                                       [[38, 83]
                                        [44, 98]
                                        [50,113]
                                        [56,128]
2.output after reshape:
                                       [[38, 50, 83, 113]
                                        [44, 56, 98, 128]]
3.output read numpy data from memory:
                                       [[38, 44, 50,  56]
                                        [83, 98,113, 128]]
function foo!(x, y, z)
    x = change_input_to_row_major(x)
    y = change_input_to_row_major(y)
    z .= gemm(x, y, z)
    z .= change_output_to_row_major(z)