文档HOT
下载
中文
注册

register_op_compute

函数功能

装饰器函数，注册算子compute实现。

其作用是整网运行时，支持算子做UB融合，使得算子在UB中根据UB融合规则自动与其他算子的compute进行拼接，提升算子运行效率。

函数原型

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

参数说明

参数名

类型

说明

op_type

string

注册名。

  • 如果op_mode是static，则是算子定义函数的名称。
  • 如果op_mode是dynamic，则是算子类型。

op_mode

string

算子实现类型，支持如下两种取值：

  • static：静态实现。
  • dynamic：动态实现。

默认值：dynamic

support_fusion

bool

是否支持UB融合。

  • True：支持UB融合。
  • False：不支持UB融合。

默认值：True

返回值说明

无。

约束说明

  • 此接口仅适用于使用DSL方式开发的算子。
  • 若算子实现逻辑中涉及reshape操作，则不支持UB自动融合，算子compute函数声明时不可使用此装饰器函数。

调用示例

@tbe.common.register.register_op_compute("abs", "static", True)
def abs_compute(x, y, kernel_name="abs"):
父主题： TBE Register API