原理

概述

算子原型定义规定了在昇腾AI处理器上可运行算子的约束,主要体现算子的数学含义,包含定义算子输入、输出和属性信息,基本参数的校验和shape的推导,原型定义的信息会被注册到GE的算子原型库中。网络模型生成时,GE会调用算子原型库的校验接口进行基本参数的校验,校验通过后,会根据原型库中的推导函数推导每个节点的输出shape与dtype,进行输出tensor的静态内存的分配。

算子原型库在整个网络模型生成流程的作用如图1所示。

图1 算子原型在GE注册的流程

其中算子注册包括OP注册、InferShape与Verify的注册,注册的时候算子类型(OpType)作为Key。

  1. 首先GE接收到第三方框架的原始网络模型,并进行初始化,网络模型的拓扑图简称为图。
  2. 算子原型库管理模块从CANN算子库的“opp/built-in/op_proto/lib”目录下加载对应操作系统和架构的算子原型库so文件。
  3. 算子原型库管理模块根据so文件中的信息在OperatorFactory中进行算子信息注册,包括算子基本信息注册、InferShape注册、Verify注册,这三部分分别以算子类型(OpType)作为key注册到三个map文件中进行保存。
  4. 图准备阶段,GE向Graph发送调用InferShape函数与Verify函数的请求。其中InferShape函数用于进行输出shape的推导,从而进行静态内存的分配;Verify函数用户进行参数的基本校验。
  5. Graph会遍历图中所有节点。
  6. 每个节点都会向OpDesc发送调用InferShape与Verify函数的请求。
  7. OpDesc从OperatorFactory中根据OpType取出对应的InferShape函数与Verify函数。
  8. OpDesc执行Verify函数进行校验,如果校验成功,则继续往下执行;如果校验失败,则直接返回。
  9. OpDesc执行InferShape函数,进行输出tensor的shape推导。
  10. OpDesc向GE返回InferShape的结果,GE后续根据InferShape结果分配输出tensor的静态内存。
  11. GE进行其他处理。

InferShape流程介绍

介绍InferShape流程前先了解如下背景知识。

由上述背景介绍可知,只要全图所有首节点的TensorDesc确定了,就可以逐个向下传播,再由算子自身实现的Shape推导能力,就可以将全图所有OP的输入输出TensorDesc推导出来,这就是InferShape的推导流程。

InferShape有以下注意点:

实现文件说明

开发者进行算子原型定义开发时需要实现如下两个文件:
  • 在头文件中进行算子IR原型的注册
  • 在cc文件中进行校验函数与shape推导函数的实现