mindspore.ops.Custom

class mindspore.ops.Custom(func, out_shape, out_dtype, func_type, bprop=None, reg_info=None)[source]

Custom primitive is used for user defined operators and is to enhance the expressive ability of built-in primitives. You can construct a Custom object with a predefined function, which describes the computation logic of a user defined operator. You can also construct another Custom object with another predefined function if needed. Then these Custom objects can be directly used in neural networks.

Warning

This is an experimental prototype that is subject to change.

Parameters

func (Union[function, str]) –
- function: If func is of function type, then func should be a Python function which describes the computation logic of a user defined operator. The function can be one of the following:
  1. A AKG operator implementation function, which can use ir builder/tvm compute/hybrid grammar.
  2. A TBE operator implementation function.
  3. A pure python function
- str: If func is of str type, then str should be a path of binary file along with a function name. This could only be used when func_type is “aot”. Currently “aot” supports GPU/CPU(linux only) platform. “aot” means ahead of time, in which case Custom directly launches user defined “xxx.so” file as an operator. Users need to compile a handwriting “xxx.cu”/”xxx.cc” file into “xxx.so” ahead of time, and offer the path of the file along with a function name.
  - ”xxx.so” file generation:
    
    1) GPU Platform: Given user defined “xxx.cu” file (ex. “{path}/add.cu”), use nvcc command to compile it.(ex. “nvcc –shared -Xcompiler -fPIC -o add.so add.cu”)
    
    2) CPU Platform: Given user defined “xxx.cc” file (ex. “{path}/add.cc”), use g++/gcc command to compile it.(ex. “g++ –shared -fPIC -o add.so add.cc”)
  - Define a “xxx.cc”/”xxx.cu” file:
    
    ”aot” is a cross-platform identity. The functions defined in “xxx.cc” or “xxx.cu” share the same args. Typically, the function should be as:
    int func(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra)
    Parameters:
    - nparam(int): total number of inputs plus outputs; suppose the operator has 2 inputs and 3 outputs, then nparam=5
    - params(void **): a pointer to the array of inputs and outputs’ pointer; the pointer type of inputs and outputs is void * ; suppose the operator has 2 inputs and 3 outputs, then the first input’s pointer is params[0] and the second output’s pointer is params[3]
    - ndims(int *): a pointer to the array of inputs and outputs’ dimension num; suppose params[i] is a 1024x1024 tensor and params[j] is a 77x83x4 tensor, then ndims[i]=2, ndims[j]=3.
    - shapes(int64_t **): a pointer to the array of inputs and outputs’ shapes(int64_t *); the ith input’s jth dimension’s size is shapes[i][j](0<=j<ndims[i]); suppose params[i] is a 2x3 tensor and params[j] is a 3x3x4 tensor, then shapes[i][0]=2, shapes[j][2]=4.
    - dtypes(const char **): a pointer to the array of inputs and outputs’ types(const char *); (ex. “float32”, “float16”, “float”, “float64”, “int”, “int8”, “int16”, “int32”, “int64”, “uint”, “uint8”, “uint16”, “uint32”, “uint64”, “bool”)
    - stream(void *): stream pointer, only used in cuda file
    - extra(void *): used for further extension
    Return Value(int):
    - 0: MindSpore will continue to run if this aot kernel is successfully executed
    - others: MindSpore will raise exception and exit
    Examples: see details in tests/st/ops/graph_kernel/custom/aot_test_files/
  - Use it in Custom:
    Custom(func="{dir_path}/{file_name}:{func_name}",...) (ex. Custom(func="./reorganize.so:CustomReorganize", out_shape=[1], out_dtype=mstype.float32))
out_shape (Union[function, list, tuple]) –
The output shape infer function or the value of output shape of func.

If func has single output, then the value of output shape is a list or tuple of int.

If func has multiple outputs, then the value of output shape is a tuple, each item represents the shape of each output.
out_dtype (Union[function, mindspore.dtype, tuple[mindspore.dtype]]) –
The output data type infer function or the value of output data type of func.

If func has single output, then the value of output shape is a mindspore.dtype.

If func has multiple outputs, then the value of output shape is a tuple of mindspore.dtype, each item represents the data type of each output.
func_type (str) –
The implementation type of func, should be one of [“akg”, “tbe”, “aot”, “pyfunc”]. Each func_type only supports specific platforms(targets). The supported platforms of func_type:
- ”akg”: supports [“Ascend”, “GPU”].
- ”tbe”: supports [“Ascend”].
- ”aot”: supports [“GPU”, “CPU”].
- ”pyfunc”: supports [“CPU”].
bprop (function) – The back propagation function of func. Default: None.
reg_info (Union[str, dict, list, tuple]) –
Represents the registration information(reg info) of func with json format of type str or dict. The reg info specifies supported data types and formats of inputs and outputs, attributes and target of func. Default: None.

If reg info is a list or tuple, then each item should be with json format of type str or dict, which represents the registration information of func in a specific target. You need to invoke CustomRegOp or the subclass of RegOp to generate the reg info for func. Then you can invoke custom_info_register to bind the reg info to func or just pass the reg info to reg_info parameter. The reg_info parameter takes higher priority than custom_info_register and the reg info in a specific target will be registered only once.

If reg info is not set, then we will infer the data types and formats from the inputs of Custom operator.

Please note that, if func_type is “tbe” or the func only supports some specified data types and formats, or it has attribute inputs, then you should set the reg info for func.

Inputs:

input (Union(tuple, list)) - The input tuple or list is made up of multiple tensors, and attributes value(optional).

Outputs:

Tensor or tuple[Tensor], execution results.

Raises

TypeError – If the type of func is invalid or the type of register information for func is invalid.
ValueError – If func_type is invalid.
ValueError – If the register information is invalid, including the target is not supported, the input numbers or the attributes of func differs in different targets.

Supported Platforms:: Ascend GPU CPU

Examples

>>> import mindspore.ops as ops
>>> from mindspore.ops import CustomRegOp, custom_info_register, DataType
>>> from mindspore.common import dtype as mstype
>>> from mindspore.nn import Cell
>>>
>>> # Example, func_type = "akg"
>>> def outer_product(a, b):
...     c = output_tensor(a.shape, a.dtype)
...     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
>>>
>>> class AkgNet(Cell):
...     def __init__(self):
...         super(AkgNet, self).__init__()
...         def infer_func(x, y):
...             return x
...         self.program = ops.Custom(outer_product, out_shape=infer_func, out_dtype=infer_func, \
...                                   func_type="akg")
...     def construct(self, x, y):
...         return self.program(x, y)
>>>
>>> # Example, func_type = "tbe"
>>> square_with_bias_op_info = CustomRegOp() \
...     .fusion_type("OPAQUE") \
...     .attr("bias", "required", "float") \
...     .input(0, "x") \
...     .output(0, "y") \
...     .dtype_format(DataType.F32_Default, DataType.F32_Default) \
...     .dtype_format(DataType.F16_Default, DataType.F16_Default) \
...     .target("Ascend") \
...     .get_op_info()
>>>
>>> @custom_info_register(square_with_bias_op_info)
... def square_with_bias(input_x, output_y, bias=0.0, kernel_name="square_with_bias"):
...     import te.lang.cce
...     from te import tvm
...     from topi.cce import util
...
...     shape = input_x.get("shape")
...     dtype = input_x.get("dtype").lower()
...
...     shape = util.shape_refine(shape)
...     data = tvm.placeholder(shape, name="data", dtype=dtype)
...
...     with tvm.target.cce():
...         res0 = te.lang.cce.vmul(data, data)
...         res = te.lang.cce.vadds(res0, bias)
...         sch = te.lang.cce.auto_schedule(res)
...
...     config = {"print_ir": False,
...               "name": kernel_name,
...               "tensor_list": [data, res]}
...
...     te.lang.cce.cce_build_code(sch, config)
>>>
>>> class TbeNet(Cell):
...     def __init__(self):
...         super(TbeNet, self).__init__()
...         self.square_with_bias = ops.Custom(square_with_bias, out_shape=lambda x, _: x, \
...                                            out_dtype=lambda x, _: x, func_type="tbe")
...     def construct(self, x):
...         res = self.square_with_bias(x, 1.0)
...         return res
>>>
>>> # Example, func_type = "aicpu"
>>> resize_bilinear_op_info = CustomRegOp("ResizeBilinear") \
...     .fusion_type("OPAQUE") \
...     .input(0, "input", "required") \
...     .output(1, "output", "required") \
...     .attr("align_corners", "required", "bool") \
...     .attr("cust_aicpu", "optional", "str", "aicpu_kernels") \
...     .dtype_format(DataType.F32_Default, DataType.F32_Default) \
...     .dtype_format(DataType.F16_Default, DataType.F32_Default) \
...     .target("Ascend") \
...     .get_op_info()
>>>
>>> @custom_info_register(resize_bilinear_op_info)
... def resize_bilinear_aicpu():
...     return
>>>
>>> class AicpuNet(Cell):
...     def __init__(self):
...         super(AicpuNet, self).__init__()
...         self.resize_bilinear_op = ops.Custom(resize_bilinear_aicpu, out_shape=[1, 1, 9, 9], \
...                                              out_dtype=mstype.float32, func_type="aicpu")
...     def construct(self, x):
...         res = self.resize_bilinear_op(x, True, "aicpu_kernels")
...         return res
>>>
>>> # Example, func_type = "aot"
>>> class AOTSingleOutputNet(Cell):
...     def __init__(self, out_shapes, out_types):
...         super(AOTSingleOutputNet, self).__init__()
...         self.program = ops.Custom("./reorganize.so:CustomReorganize", out_shapes, out_types, "aot")
...     def construct(self, x, y):
...         return self.program(x, y)
>>>
>>> # Example, func_type = "pyfunc"
>>> def func_multi_output(x1, x2):
...     return (x1 + x2), (x1 - x2)
>>>
>>> class PyFuncNet(Cell):
...     def __init__(self):
...         super(PyFuncNet, self).__init__()
...         self.func = ops.Custom(func_multi_output, lambda x, _: (x, x), lambda x, _: (x, x), "pyfunc")
...     def construct(self, x1, x2):
...         return self.func(x1, x2)