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]=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”)

      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