mindspore.ops.TBERegOp

class mindspore.ops.TBERegOp(op_name)

Class for TBE operator information registration. TBE (Tensor Boost Engine) is the Ascend operator development tool, which is extended on the basis of the TVM framework to develop custom operators.

Parameters:

op_name (str) – Name of operator.

Examples

>>> from mindspore.ops import TBERegOp, DataType
>>> op_name_op_info = TBERegOp("OpName") \
...    .fusion_type("ELEMWISE") \
...    .async_flag(False) \
...    .binfile_name("op_name.so") \
...    .compute_cost(10) \
...    .kernel_name("op_name") \
...    .partial_flag(True) \
...    .op_pattern("formatAgnostic") \
...    .need_check_supported(True) \
...    .dynamic_shape(True) \
...    .dynamic_rank_support(True) \
...    .dynamic_compile_static(True) \
...    .attr("format", "required", "str", "all") \
...    .input(0, "x1", None, "required", None) \
...    .input(0, "x2", None, "required", None) \
...    .input(1, "axis", None, "required", None) \
...    .output(0, "y", True, "required", "all") \
...    .real_input_index([1, 0]) \
...    .input_to_attr_index([2]) \
...    .unknown_shape_formats(["ND", "ND", "ND", "ND"]) \
...    .reshape_type("NC") \
...    .is_dynamic_format(True) \
...    .dtype_format(DataType.F16_None, DataType.F16_None, DataType.F16_None, DataType.F16_None) \
...    .dtype_format(DataType.F32_None, DataType.F32_None, DataType.F32_None, DataType.F32_None) \
...    .dtype_format(DataType.I32_None, DataType.I32_None, DataType.I32_None, DataType.I32_None) \
...    .get_op_info()
>>>

async_flag(async_flag=False)

Define the calculation efficiency of the operator, whether the asynchronous calculation is supported.

Parameters:

async_flag (bool, optional) – Value of async flag. Default: False .

attr(name=None, param_type=None, value_type=None, value=None, default_value=None, **kwargs)

Register TBE op attribute information.

Parameters:

• name (str, optional) – Name of the attribute. Default: None .

• param_type (str, optional) – Param type of the attribute. Default: None .

• value_type (str, optional) – Type of the attribute. Default: None .

• value (str, optional) – Value of the attribute. Default: None .

• default_value (str, optional) – Default value of attribute. Default: None .

• kwargs (dict) – Other information of the attribute.

binfile_name(binfile_name)

Set the binary file name of the operator, it is optional.

Parameters:

binfile_name (str) – The binary file name of the operator.

compute_cost(compute_cost=10)

Define the calculation efficiency of operator, which refers to the value of the cost model in the tiling module.

Parameters:

compute_cost (int, optional) – Value of compute cost. Default: 10 .

dynamic_compile_static(dynamic_compile_static=False)

Whether the operator supports dynamic compile static.

Parameters:

dynamic_compile_static (bool, optional) – An identifier that indicates whether the operator supports dynamic compile static. Default: False .

dynamic_rank_support(dynamic_rank_support)

Description whether the operator supports dynamic rank (dynamic dimension).

Parameters:

dynamic_rank_support (bool, optional) – Description whether the operator supports dynamic rank (dynamic dimension). True: indicates that dynamic rank is supported. False: indicates that the operator does not support dynamic rank. Default: False .

dynamic_shape(dynamic_shape=False)

Whether the operator supports dynamic shape.

Parameters:

dynamic_shape (bool, optional) – Value of dynamic shape. Default: False .

input(index=None, name=None, need_compile=None, param_type=None, shape=None, value_depend=None, **kwargs)

Register TBE op input information.

Parameters:

• index (int, optional) – Order of the input. Default: None .

• name (str, optional) – Name of the input. Default: None .

• need_compile (bool, optional) – Whether the input needs to be compiled or not. Default: None .

• param_type (str, optional) – Type of the input. Default: None .

• shape (str, optional) – Shape of the input. Default: None .

• value_depend (str, optional) – Whether the input is constant value depend. Default: None .

• kwargs (dict) – Other information of the input.

input_to_attr_index(input_to_attr_index)

Description the index of input need to cast to attr.

Parameters:

input_to_attr_index (list) – Value of input_to_attr_index. Default: () .

is_dynamic_format(is_dynamic_format=False)

Whether the operator needs calop_select_format api.

Parameters:

is_dynamic_format (bool, optional) – Value of is_dynamic_format. Default: False .

kernel_name(kernel_name)

The name of operator kernel.

Parameters:

kernel_name (str) – Name of operator kernel.

need_check_supported(need_check_supported=False)

Whether the operator needs check supports.

Parameters:

need_check_supported (bool, optional) – Value of need_check_supported. Default: False .

op_pattern(pattern=None)

The behavior type of operator, such as broadcast, reduce and so on.

Parameters:

pattern (str, optional) – Value of op pattern, e.g. "broadcast", "reduce". Default: None .

output(index=None, name=None, need_compile=None, param_type=None, shape=None, **kwargs)

Register TBE op output information.

Parameters:

• index (int, optional) – Order of the output. Default: None .

• name (str, optional) – Name of the output. Default: None .

• need_compile (bool, optional) – Whether the output needs to be compiled or not. Default: None .

• param_type (str, optional) – Type of the output. Default: None .

• shape (str, optional) – Shape of the output. Default: None .

• kwargs (dict) – Other information of the output.

partial_flag(partial_flag=True)

Define the calculation efficiency of operator, whether the partial calculation is supported.

Parameters:

partial_flag (bool, optional) – Value of partial flag. Default: True .

real_input_index(real_input_index)

Description operator front end and tbe operator input mapping.

Parameters:

real_input_index (list) – Value of real_input_index. Default: () .

reshape_type(reshape_type)

Reshape type of operator.

Parameters:

reshape_type (str) – Value of reshape type. For example, if the input shape is \((2, 3)\) and reshape_type is set to "CH", then the new shape is \((1, 2, 3, 1)\). "CH" means the C and H dimensions are kept and new dimensions are added for N and W dimension.

unknown_shape_formats(unknown_shape_formats)

Description data arrangement of operator input / output tensor in dynamic shape scene.

Parameters:

unknown_shape_formats (list) – Description data arrangement of operator input / output tensor in dynamic shape scene.