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.

unknown_shape_formats(unknown_shape_formats)[source]

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.