mindspore.ops.TBERegOp
- class mindspore.ops.TBERegOp(op_name)[source]
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)[source]
Define the calculation efficiency of the operator, whether the asynchronous calculation is supported.
- Parameters
async_flag (bool) – Value of async flag. Default:
False.
- attr(name=None, param_type=None, value_type=None, value=None, default_value=None, **kwargs)[source]
Register TBE op attribute information.
- Parameters
name (str) – Name of the attribute. Default:
None.param_type (str) – Param type of the attribute. Default:
None.value_type (str) – Type of the attribute. Default:
None.value (str) – Value of the attribute. Default:
None.default_value (str) – Default value of attribute. Default:
None.kwargs (dict) – Other information of the attribute.
- binfile_name(binfile_name)[source]
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)[source]
Define the calculation efficiency of operator, which refers to the value of the cost model in the tiling module.
- Parameters
compute_cost (int) – Value of compute cost. Default:
10.
- dynamic_compile_static(dynamic_compile_static=False)[source]
Whether the operator supports dynamic compile static.
- Parameters
dynamic_compile_static (bool) – Value of dynamic compile static. Default:
False.
- dynamic_rank_support(dynamic_rank_support)[source]
Description whether the operator supports dynamic rank (dynamic dimension).
- Parameters
dynamic_rank_support (bool) – Description whether the operator supports dynamic rank (dynamic dimension). True: indicates that dynamic rank is supported, and the operator supports shape (- 2), which is used to determine whether dynamic is performed. False: indicates that the operator does not support dynamic rank. Default:
False.
- dynamic_shape(dynamic_shape=False)[source]
Whether the operator supports dynamic shape.
- Parameters
dynamic_shape (bool) – Value of dynamic shape. Default:
False.
- input(index=None, name=None, need_compile=None, param_type=None, shape=None, value_depend=None, **kwargs)[source]
Register TBE op input information.
- Parameters
index (int) – Order of the input. Default:
None.name (str) – Name of the input. Default:
None.need_compile (bool) – Whether the input needs to be compiled or not. Default:
None.param_type (str) – Type of the input. Default:
None.shape (str) – Shape of the input. Default:
None.value_depend (str) – Whether the input is constant value depend. Default:
None.kwargs (dict) – Other information of the input.
- input_to_attr_index(input_to_attr_index)[source]
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)[source]
Whether the operator needs calop_select_format api.
- Parameters
is_dynamic_format (bool) – Value of is_dynamic_format. Default:
False.
- kernel_name(kernel_name)[source]
The name of operator kernel.
- Parameters
kernel_name (str) – Name of operator kernel.
- need_check_supported(need_check_supported=False)[source]
Whether the operator needs check supports.
- Parameters
need_check_supported (bool) – Value of need_check_supported. Default:
False.
- op_pattern(pattern=None)[source]
The behavior type of operator, such as broadcast, reduce and so on.
- Parameters
pattern (str) – Value of op pattern, e.g. "broadcast", "reduce". Default:
None.
- output(index=None, name=None, need_compile=None, param_type=None, shape=None, **kwargs)[source]
Register TBE op output information.
- Parameters
index (int) – Order of the output. Default:
None.name (str) – Name of the output. Default:
None.need_compile (bool) – Whether the output needs to be compiled or not. Default:
None.param_type (str) – Type of the output. Default:
None.shape (str) – Shape of the output. Default:
None.kwargs (dict) – Other information of the output.
- partial_flag(partial_flag=True)[source]
Define the calculation efficiency of operator, whether the partial calculation is supported.
- Parameters
partial_flag (bool) – Value of partial flag. Default:
True.
- real_input_index(real_input_index)[source]
Description operator front end and tbe operator input mapping.
- Parameters
real_input_index (list) – Value of real_input_index. Default:
().
- reshape_type(reshape_type)[source]
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.