register_op_compute

Description

A wrapper function that registers an operator compute implementation.

Performs UB fusion on operators on the network at run time. In this way, the compute functions of custom operators can be concatenated automatically in the UB according to the UB fusion patterns, improving execution efficiency.

Prototype

def register_op_compute(op_type, op_mode="dynamic", support_fusion=True)

Parameters

Parameter	Type	Description
op_type	String	Name to be registered. Set to the name of the operator definition function if op_mode is set to static. Set to the operator type if op_mode is set to dynamic.
op_mode	String	Operator implementation type, selected from: static: static implementation. dynamic: dynamic implementation. Defaults to dynamic.
support_fusion	Bool	UB fusion support. True: supported. False: not supported. Defaults to True.

Parameter

Type

Description

op_type

String

Name to be registered.

Set to the name of the operator definition function if op_mode is set to static.
Set to the operator type if op_mode is set to dynamic.

op_mode

String

Operator implementation type, selected from:

static: static implementation.
dynamic: dynamic implementation.

Defaults to dynamic.

support_fusion

Bool

UB fusion support.

True: supported.
False: not supported.

Defaults to True.

Returns

None

Restrictions

This API applies only to operators developed in DSL mode.
If the reshape operation is involved in the operator implementation logic, automatic UB fusion is not supported. As such, do not use this wrapper when declaring the operator compute function.

Example

@tbe.common.register.register_op_compute("abs", "static", True)
def abs_compute(x, y, kernel_name="abs"):

Parent topic: TBE Register APIs