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”].

  • “pyfunc”: [“CPU”].

  • “julia”: [“CPU”].

参数:
  • func (Union[function, str]) - 自定义算子的函数表达。

    • function:如果 func 是函数类型,那么 func 应该是一个Python函数,它描述了用户定义的操作符的计算逻辑。该函数可以是以下之一:

      1. AKG操作符实现函数,可以使用ir builder/tvm compute/hybrid语法。

      2. 纯Python函数。

      3. 使用Hybrid DSL编写的带有装饰器的内核函数。

    • 字符串:如果 func 是字符串类型,那么 str 应该是包含函数名的文件路径。当 func_type 是”aot”或”julia”时,可以使用这种方式。

      1. 对于”aot”:

        目前”aot”支持GPU/CPU(仅Linux)平台。”aot”意味着提前编译,在这种情况下,Custom直接启动用户定义的”xxx.so”文件作为操作符。用户需要提前将手写的”xxx.cu”/”xxx.cc”文件编译成”xxx.so”,并提供文件路径和函数名。

        • “xxx.so”文件生成:

          1. GPU平台:给定用户定义的”xxx.cu”文件(例如”{path}/add.cu”),使用nvcc命令进行编译(例如”nvcc –shared -Xcompiler -fPIC -o add.so add.cu”)。

          2. 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]=2ndims[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]=2shapes[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”)

      2. 对于”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

样例:

>>> 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