TBE编译运行逻辑架构
一个完整的TBE算子包含四部分：算子原型定义、对应开源框架的算子适配插件、算子信息库定义和算子实现。
算子开发完成后在昇腾AI处理器硬件平台上的编译运行的架构如图1和图2所示。
其中，TF Adapter只有在基于TensorFlow框架进行训练时使用。
上图中
代表开发人员在自定义TBE算子时需要实现的交付件。
开发交付件 |
说明 |
|---|---|
算子实现 |
算子实现的python文件，包含算子的计算实现及Schedule实现。 |
算子适配插件 |
基于第三方框架（TensorFlow/Caffe）进行自定义算子开发的场景，开发人员完成自定义算子的实现代码后，需要进行适配插件的开发，将基于第三方框架的算子映射成适昇腾AI处理器的算子，将算子信息注册到GE中。基于第三方框架的网络运行时，首先会加载并调用算子适配插件信息，将第三方框架网络中的算子进行解析并映射成昇腾AI处理器中的算子。 |
算子原型库 |
算子原型定义规定了在昇腾AI处理器上可运行算子的约束，主要体现算子的数学含义，包含定义算子输入、输出、属性和取值范围，基本参数的校验和shape的推导。网络运行时，GE会调用算子原型库的校验接口进行基本参数的校验，校验通过后，会根据原型库中的推导函数推导每个节点的输出shape与dtype，进行输出tensor的静态内存的分配。 |
算子信息库 |
算子信息库主要体现算子在昇腾AI处理器上物理实现的限制，包括算子的输入输出dtype、format以及输入shape信息。网络运行时，FE会根据算子信息库中的算子信息做基本校验，判断是否需要为算子插入合适的转换节点，并根据算子信息库中信息找到对应的算子实现文件进行编译，生成算子二进制文件进行执行。 |
算子编译流程
TBE算子的详细编译流程如下图所示。
- 将开源框架网络模型下发给GE。
若基于Tensorflow框架进行在线训练，首先会调用TF Adapter适配接口生成TF原始网络模型，然后下发给GE。
若原始网络框架为MindSpore，则直接将原始网络模型下发给GE。
若使用AscendCL应用进行模型推理，则直接将原始网络模型下发给GE。
注：网络模型的拓扑图后续简称为图。
- GE调用算子插件，将原始网络模型中的算子映射为适配昇腾AI处理器的算子，从而将原始开源框架图解析为适配昇腾AI处理器的图。
若原始网络框架为MindSpore，无需调用算子插件进行算子映射，因为MindSpore图中的算子已经为适配为昇腾AI处理器的算子。
- 调用算子原型库校验接口进行基本参数的校验，校验通过后，会根据原型库中的推导函数推导每个节点的输出shape与dtype，进行输出tensor的静态内存的分配。
- GE向FE发送图优化请求，并将图下发给FE，FE匹配融合规则进行图融合，并进行算子选择，选择优先级最高的算子类型进行算子匹配（默认自定义算子优先级最高），最后将优化后的整图返回给GE。
- GE根据图中数据将图拆分为子图并下发给FE，FE首先在子图内部插入转换算子，然后按照当前子图流程进行TBE算子预编译，对TBE算子进行UB融合（算子可以在UB中根据UB融合规则自动与其他算子的计算进行组装），并根据算子信息库中算子信息找到算子实现将其编译成算子kernel（算子的*.o与*.json），最后将优化后子图返回给GE。
- GE将子图进行合并，并对合并后的整图进一步优化。
- GE进行图编译，包含内存分配、流资源分配等，并向FE发送tasking请求，FE返回算子的taskinfo信息给GE，图编译完成后生成适配昇腾AI处理器的模型。
算子运行流程
TBE算子的运行流程如图4所示。
- GE加载编译好的模型，并下发算子执行请求。
- Runtime将对应的task请求下发给Task Schedule。
- Task Schedule进行task调度，调用算子计算接口。
- 算子在AI Core上执行。