mindspore.ops.Custom
- class mindspore.ops.Custom(func, bprop=None, out_dtype=None, func_type='hybrid', out_shape=None, reg_info=None)[源代码]
Custom 算子是MindSpore自定义算子的统一接口。用户可以利用该接口自行定义MindSpore内置算子库尚未包含的算子。 根据输入函数的不同,你可以创建多个自定义算子,并且把它们用在神经网络中。 关于自定义算子的详细说明和介绍,包括参数的正确书写,见 自定义算子教程 。
警告
这是一个实验性API,后续可能修改或删除。
说明
不同自定义算子的函数类型(func_type)支持的平台类型不同。每种类型支持的平台如下:
"hybrid": ["GPU", "CPU"].
"akg": ["GPU", "CPU"].
"aot": ["GPU", "CPU","ASCEND"].
"pyfunc": ["CPU"].
"julia": ["CPU"].
- 参数:
func (Union[function, str]) - 自定义算子的函数表达。
function:如果 func 是函数类型,那么 func 应该是一个Python函数,它描述了用户定义的操作符的计算逻辑。该函数可以是以下之一:
AKG操作符实现函数,可以使用ir builder/tvm compute/hybrid语法。
纯Python函数。
使用Hybrid DSL编写的带有装饰器的内核函数。
字符串:如果 func 是字符串类型,那么 str 应该是包含函数名的文件路径。当 func_type 是"aot"或"julia"时,可以使用这种方式。
对于"aot":
GPU/CPU(仅Linux)平台
"aot"意味着提前编译,在这种情况下,Custom直接启动用户定义的"xxx.so"文件作为操作符。用户需要提前将手写的"xxx.cu"/"xxx.cc"文件编译成"xxx.so",并提供文件路径和函数名。
"xxx.so"文件生成:
GPU平台:给定用户定义的"xxx.cu"文件(例如"{path}/add.cu"),使用nvcc命令进行编译(例如"nvcc –shared -Xcompiler -fPIC -o add.so add.cu")。
CPU平台:给定用户定义的"xxx.cc"文件(例如"{path}/add.cc"),使用g++/gcc命令进行编译(例如"g++ –shared -fPIC -o add.so add.cc")。
定义"xxx.cc"/"xxx.cu"文件:
"aot"是一个跨平台的标识符。"xxx.cc"或"xxx.cu"中定义的函数具有相同的参数。通常,该函数应该像这样:
int func(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra)
参数:
nparam(int) : 输入和输出的总数;假设操作符有2个输入和3个输出,那么 nparam=5 。
params(void **) : 输入和输出指针的数组指针;输入和输出的指针类型为 void * ;假设操作符有2个输入和3个输出,那么第一个输入的指针是 params[0] ,第二个输出的指针是 params[3] 。
ndims(int *) : 输入和输出维度数的数组指针;假设 params[i] 是一个1024x1024的张量, params[j] 是一个77x83x4的张量,那么 ndims[i]=2 , ndims[j]=3 。
shapes(int64_t **) : 输入和输出形状( int64_t * )的数组指针;第 i 个输入的第 j 个维度的大小是 shapes[i][j] (其中 0<=j<ndims[i] );假设 params[i] 是一个2x3的张量, params[j] 是一个3x3x4的张量,那么 shapes[i][0]=2 , shapes[j][2]=4 。
dtypes(const char **) : 输入和输出类型( const char * )的数组指针;(例如:"float32"、"float16"、"float"、"float64"、"int"、"int8"、"int16"、"int32"、"int64"、"uint"、"uint8"、"uint16"、"uint32"、"uint64"、"bool")
stream(void *) : 流指针,仅在CUDA文件中使用。
extra(void *) : 用于进一步扩展。
返回值(int):
0: 如果这个AOT内核成功执行,MindSpore将继续运行。
其他值: MindSpore将引发异常并退出。
示例:详见 tests/st/ops/graph_kernel/custom/aot_test_files/ 中的详细信息。
在Custom中使用:
Custom(func="{dir_path}/{file_name}:{func_name}", ...)
例如:Custom(func="./reorganize.so:CustomReorganize", out_shape=[1], out_dtype=mstype.float32, "aot")
ASCEND平台
在ASCEND平台使用Custom算子之前,用户首先需要基于Ascend C开发自定义算子并编译。算子开发可参考 快速上手端到端算子开发 ,自定义算子编译可使用工具 Ascend C自定义算子离线编译。 在入参 func 中传入算子的名字, 以自定义算子实现中命名为 AddCustom 为例,存在以下几种使用方式:
算子底层调用TBE: func="AddCustom"
算子底层调用AclNN: func="aclnnAddCustom"
算子的infer shape通过c++推导: func="infer_shape.cc:aclnnAddCustom",其中 infer_shape.cc 为c++实现的shape推导。
对于"julia":
目前,"julia"仅支持CPU(仅限Linux平台)。对于julia,它使用JIT编译器(即时编译器),并且julia支持C API来调用julia代码。自定义功能可以直接将用户定义的"xxx.jl"文件作为一个操作符来启动。用户需要编写一个包含模块和函数的"xxx.jl"文件,并提供该文件的路径以及模块名称和函数名称。
示例:详情见 tests/st/ops/graph_kernel/custom/julia_test_files/
在Custom中使用:
Custom(func="{dir_path}/{file_name}:{module_name}:{func_name}",...)
例如:Custom(func="./add.jl:Add:add", out_shape=[1], out_dtype=mstype.float32, "julia")
out_shape (Union[function, list, tuple]) - 自定义算子的输入的形状或者输出形状的推导函数。默认值:
None
。out_dtype (Union[function,
mindspore.dtype
, tuple[mindspore.dtype
]]) - 自定义算子的输入的数据类型或者输出数据类型的推导函数。默认值:None
。func_type (str) - 自定义算子的函数类型,必须是[
"hybrid"
,"akg"
,"aot"
,"pyfunc"
,"julia"
]中之一。默认值:"hybrid"
。bprop (function) - 自定义算子的反向函数。默认值:
None
。reg_info (Union[str, dict, list, tuple]) - 自定义算子的算子注册信息。默认值:
None
。
- 输入:
input (Union(tuple, list)) - 输入要计算的Tensor。
- 输出:
Tensor。自定义算子的计算结果。
- 异常:
TypeError - 如果输入 func 不合法,或者 func 对应的注册信息类型不对。
ValueError - func_type 的值不在列表内。
ValueError - 算子注册信息不合法,包括支持平台不匹配,算子输入和属性与函数不匹配。
- 支持平台:
GPU
CPU
ASCEND
样例:
>>> import numpy as np >>> from mindspore import Tensor, ops >>> from mindspore.ops import CustomRegOp, custom_info_register, DataType, kernel >>> from mindspore import dtype as mstype >>> from mindspore.nn import Cell >>> input_x = Tensor(np.ones([16, 16]).astype(np.float32)) >>> input_y = Tensor(np.ones([16, 16]).astype(np.float32)) >>> >>> # Example, func_type = "hybrid" >>> # This is the default func_type in Custom, >>> # and both out_shape and out_dtype can be None(default value). >>> # In this case, the input func must be a function written in the Hybrid DSL >>> # and decorated by @kernel. >>> @kernel ... def add_script(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 >>> >>> test_op_hybrid = ops.Custom(add_script) >>> output = test_op_hybrid(input_x, input_y) >>> # the result will be a 16 * 16 tensor with all elements 2 >>> print(output.shape) (16, 16) >>> # Example, func_type = "aot" >>> def test_aot(x, y, out_shapes, out_types): ... program = ops.Custom("./reorganize.so:CustomReorganize", out_shapes, out_types, "aot") ... out = program(x, y) ... return out >>> >>> # Example, func_type = "pyfunc" >>> def func_multi_output(x1, x2): ... return (x1 + x2), (x1 - x2) >>> >>> test_pyfunc = ops.Custom(func_multi_output, lambda x, _: (x, x), lambda x, _: (x, x), "pyfunc") >>> output = test_pyfunc(input_x, input_y) >>> >>> # Example, func_type = "julia" >>> # julia code: >>> # add.jl >>> # module Add >>> # function add(x, y, z) >>> # z .= x + y >>> # return z >>> # end >>> # end >>> def test_julia(x, y, out_shapes, out_types): ... program = ops.Custom("./add.jl:Add:add", out_shapes, out_types, "julia") ... out = program(x, y) ... return out