mindspore.ops.ms_kernel

mindspore.ops.ms_kernel(fn=None, reg_info=None, compile_attrs=None)[源代码]

用于MindSpore Hybrid DSL函数书写的装饰器。给用MindSpore Hybrid DSL书写的函数加上此装饰器后，它可以用作一个普通的Python函数。与此同时，他可以用于自定义算子 mindspore.ops.Custom 的输入，其对应的 func_type 可以设置为 hybrid 或者 pyfunc 。使用 hybrid 类型的 mindspore.ops.Custom 自定义算子可以自动推导数据类型和形状。

参数：

fn (Function) - 将被作为自定义算子运行的Python函数。默认值：None。
reg_info (tuple[str, dict]) - 包含算子注册信息的dict或json字符串。默认值：None。
compile_attrs (Dict) - 算子编译信息。默认值：None。

返回：

Function，如果 fn 不是None，那么返回一个用Hybrid DSL写的可执行函数；如果 fn 是None，则返回一个装饰器，该装饰器只有 fn 一个参数。

支持平台：: Ascend GPU CPU

样例：

>>> import numpy as np
>>> from mindspore import ops, Tensor
>>> from mindspore.ops import ms_kernel, DataType, CustomRegOp
...
>>> # Create a dict for the compile flags.
>>> attrs = {
...     "test1": True,
...     "test2": "good",
...     "test3": 12,
... }
>>> # Create the reg info json string.
>>> op_gpu_info = CustomRegOp() \
...     .input(0, "a") \
...     .input(0, "b") \
...     .output(0, "y") \
...     .dtype_format(DataType.F32_None, DataType.F32_None, DataType.F32_None) \
...     .target("GPU") \
...     .get_op_info()
>>>
>>> # Create inputs for the custom op.
>>> input_x = np.ones([4, 4]).astype(np.float32)
>>> input_y = np.ones([4, 4]).astype(np.float32)
...
>>> # Write a Hybrid DSL function through the decorator @ms_kernel.
>>> # We can also pass the compile attrs and the reg info through the decorator.
>>> @ms_kernel(reg_info=op_gpu_info, compile_attrs=attrs)
... def outer_product(a, b):
...     c = output_tensor(a.shape, a.dtype)
...
...     with block_realize(c):
...         for i0 in range(a.shape[0]):
...             for i1 in range(b.shape[1]):
...                 c[i0, i1] = 0.0
...                 for i2 in range(a.shape[1]):
...                     c[i0, i1] = c[i0, i1] + (a[i0, i2] * b[i2, i1])
...     return c
...
>>> # We can use the function directly as a python function.
>>> # In this case, the inputs should be numpy arrays.
>>> result = outer_product(input_x, input_y)
...
>>> # Create a custom op with mode "hybrid" (default value) by the Hybrid DSL function.
>>> # In this case, we will enjoy the automatic dtype/shape infer for free.
>>> # The inputs should be mindspore tensors.
>>> test_op_hybrid = ops.Custom(outer_product)
>>> output = test_op_hybrid(Tensor(input_x), Tensor(input_y))