下载
中文
注册

概述

加速库中定义了三种算子,即单算子(OpsOperation)、图算子(GraphOperation)和自定义插件算子(PluginOperation)。

  • 单算子(OpsOperation)为一系列基础算子,提供了如矩阵乘、转置等功能。详细信息请参考atb/infer_op_params.hatb/train_op_params.h
  • 图算子(GraphOperation)本质是多个单算子的组合,其作用是简化重复多个算子的调用步骤。 例如,A算子与B算子按顺序调用多次,即可将A算子与B算子合并成图算子G。

    用户需要自行设计并定义图的结构,即图结构中节点(Node)的组合与依赖关系,包括节点对应的单算子、节点的输入Tensor与输出Tensor,并识别这些Tensor为图的输入Tensor、输出Tensor和中间Tensor。图输入Tensor为用户使用图算子时需要从外部输入的所有Tensor。图输出Tensor为用户使用图算子时不会再进行下一步运算操作的所有Tensor。图中间Tensor为图算子运算中产生的临时Tensor。例如同一图算子中有算子A和算子X,算子A的输出Tensor为算子X所用,且该tensor对于图外部模块不可见,则该Tensor为中间tensor。

    图1所示,该示例图算子由两个节点组成,这两个节点均为Elewise_Add算子。a、b、c为图算子的三个输入Tensor,output为图算子的输出Tensor,a_add_b_output为图算子的中间Tensor。其中,节点0的输入为a和b,输出为a_add_b_output,节点1的输入为a_add_b_output和c,输出为output。

  • 自定义插件算子(PluginOperation)是一种为用户实现特定功能提供的机制。如果一些功能通过单算子或图算子无法实现,用户可以通过开发自定义插件算子实现对应的功能。

这三种算子继承自同一个父类Operation,在使用时除了构造算子的操作步骤有所区别,其他步骤都基本一致。

图1 图算子图解