一种图编译方法、装置、电子设备及存储介质与流程

本发明涉及计算机，具体涉及一种图编译方法、装置、电子设备及存储介质。

背景技术：

1、部分人工智能（artificial intelligence，ai）芯片采用分块格式（tile格式）数据布局，其矩阵乘计算单元要求数据以特定大小的分块（如32×32）存储，而通用图形处理器（graphics processing unit，gpu）采用线性（行主序）数据布局。因此，数据在进入此类ai芯片的计算单元前需由线性格式转换为分块格式，计算完成后需转换回线性格式。

2、在基于transformer架构的大语言模型（large language model，llm）中，矩阵乘法是最频繁的计算操作，格式转换随之大量产生。以2层llama模型为例，18个矩阵乘操作共引入90次格式转换。此类操作仅改变数据的物理排列方式，不产生计算价值，却消耗内存带宽和计算内核启动资源，影响模型推理性能。

3、现有图编译技术难以消除上述冗余。主流ai编译器面向线性数据布局设计，不具备分块格式转换的识别与优化能力；面向专用ai芯片的编译工具链虽在算子层面实现了格式转换，但优化粒度局限于单个算子，常量权重重复转换、共享输入重复转换以及跨操作往返转换等冗余仍然存在；此外，格式转换逻辑硬编码于算子分发实现中，难以在编译器层面进行全局分析和优化。

4、因此，需要一种从计算图全局出发、系统性消除冗余格式转换的图编译方法，以解决上述问题中的至少一者。

技术实现思路

1、为了克服上述问题中的至少一项，本发明提出了一种图编译方法、装置、电子设备及存储介质。

2、为了实现上述目的，本发明采用以下技术方案。

3、第一方面，本发明提供一种图编译方法，包括：获取待编译模型对应的初始计算图，所述待编译模型部署于目标硬件，所述目标硬件采用分块格式数据布局；将所述初始计算图中的目标算子节点替换为硬件特定子图，所述硬件特定子图通过图捕获机制从所述目标硬件的算子分发逻辑提取得到，以在所述初始计算图中显式引入线性格式与分块格式之间的数据格式转换操作，得到第一计算图；对所述第一计算图执行冗余格式转换消除，得到第二计算图，所述冗余格式转换消除包括常量折叠、公共子表达式消除和跨操作往返转换消除中的至少一种；将所述第二计算图编译为所述目标硬件的可执行代码。

4、可选的一些实施例中，所述硬件特定子图通过以下方式提取得到：从所述目标硬件的算子分发逻辑中提取所述目标算子的处理流程函数；利用所述图捕获机制对所述处理流程函数进行追踪，生成所述处理流程函数对应的硬件特定子图。

5、可选的一些实施例中，所述利用所述图捕获机制对所述处理流程函数进行追踪，包括：注册自定义编译后端，所述自定义编译后端用于接收追踪过程中生成的图对象并保存所述图对象，并在保存后抛出异常以短路所述处理流程函数的实际执行。

6、可选的一些实施例中，所述将所述初始计算图中的目标算子节点替换为硬件特定子图，包括：获取所述目标算子节点的输入边；将所述硬件特定子图的输入占位符节点映射至所述目标算子节点的输入边；将所述硬件特定子图的内部节点复制至所述初始计算图中以替换所述目标算子节点的位置，并将所述硬件特定子图的输出节点的后继边重定向至所述目标算子节点的后继节点；删除所述目标算子节点。

7、可选的一些实施例中，所述方法还包括：以操作类型、各输入张量的形状和数据类型的组合为缓存键，缓存已生成的所述硬件特定子图；当存在与所述缓存键匹配的目标算子节点时，复用已缓存的所述硬件特定子图替换所述目标算子节点。

8、可选的一些实施例中，所述常量折叠包括：响应于输入占位符节点的值类型为模型参数类型且存在于模型参数列表中，将所述输入占位符节点提升为指向模型模块的常量存储的常量属性读取节点；以所述常量属性读取节点为起点，追踪依次包括填充操作节点和线性格式到分块格式转换操作节点的格式转换操作链；在编译器的宿主机端执行所述格式转换操作链，得到预计算结果，将所述预计算结果写入所述常量存储，以读取所述预计算结果的常量属性读取节点替换所述格式转换操作链，并将替换后的常量属性读取节点的输出边连接至所述格式转换操作链终点的后继节点。

9、可选的一些实施例中，所述公共子表达式消除包括：消除所述第一计算图中源数据类型与目标数据类型相同的类型转换节点；对共享同一输入数据的多个算子节点各自关联的视图变换节点执行公共子表达式消除，合并具有相同输入源和相同目标形状的视图变换节点；在所述视图变换节点合并的基础上，对所述多个算子节点各自关联的填充节点执行公共子表达式消除，合并具有相同输入源的填充节点；在所述填充节点合并的基础上，对所述多个算子节点各自关联的线性格式到分块格式转换节点执行公共子表达式消除，合并具有相同输入源的线性格式到分块格式转换节点。

10、可选的一些实施例中，所述跨操作往返转换消除包括：识别所述第一计算图中存在分块格式到线性格式再到分块格式的往返转换、且往返转换后的后继算子支持分块格式数据输入的数据流路径；删除所述数据流路径内的分块格式到线性格式转换节点和线性格式到分块格式转换节点，将所述数据流路径的入口前驱节点的分块格式数据输出直接连接至所述数据流路径的出口后继算子；在所述数据流路径的入口保留一次线性格式到分块格式的转换，在所述数据流路径的出口保留一次分块格式到线性格式的转换。

11、可选的一些实施例中，所述识别所述第一计算图中存在分块格式到线性格式再到分块格式的往返转换、且往返转换后的后继算子支持分块格式数据输入的数据流路径，包括：标记所述第一计算图中以分块格式数据进行计算的计算节点；从所述计算节点的输出出发，沿数据流方向追踪，将满足以下条件的下游节点归入同一操作链：所述下游节点为分块格式到线性格式转换节点且其输出直接或经逐元素操作后传入线性格式到分块格式转换节点，或者所述下游节点为逐元素操作节点或形状变换操作节点；以数据首次从线性格式进入分块格式的位置为链起点，以数据最后一次从分块格式转回线性格式的位置为链终点。

12、可选的一些实施例中，当所述数据流路径内存在形状变换操作节点时，所述跨操作往返转换消除还包括：判断所述形状变换操作节点是否满足在分块格式下直接执行的条件；若满足，则通过重映射所述形状变换操作节点的数据访问索引关系，在分块格式下直接完成所述形状变换，而不将数据回退至线性格式。

13、可选的一些实施例中，所述判断所述形状变换操作节点是否满足在分块格式下直接执行的条件，包括：判断所述形状变换操作节点的变换边界是否与分块格式的分块边界对齐；判断所述形状变换操作节点所涉及的维度是否为参与分块格式分块的维度，以确定在分块格式下执行所述形状变换时对块间索引与块内索引的调整方式；判断所述形状变换操作节点是否引入跨分块边界的不规则数据依赖；若所述变换边界与所述分块边界对齐、且所述形状变换操作节点未引入跨分块边界的不规则数据依赖，则判定所述形状变换操作节点满足在分块格式下直接执行的条件，并按照所确定的块间索引与块内索引的调整方式执行所述形状变换。

14、可选的一些实施例中，所述形状变换操作节点为张量形状重塑操作节点，所述通过重映射所述形状变换操作节点的数据访问索引关系，在分块格式下直接完成所述形状变换，包括：重新计算分块格式下的元素索引映射关系，使变换前后的物理存储地址等价，以对所述张量形状重塑操作节点进行零拷贝的形状重解释。

15、可选的一些实施例中，所述形状变换操作节点为张量维度转置操作节点，所述通过重映射所述形状变换操作节点的数据访问索引关系，在分块格式下直接完成所述形状变换，包括：修改数据搬运描述符的源地址步长参数和目标地址步长参数，以在数据搬运过程中完成所述张量维度转置操作节点对应的转置。

16、可选的一些实施例中，所述形状变换操作节点为逐元素操作节点，所述通过重映射所述形状变换操作节点的数据访问索引关系，在分块格式下直接完成所述形状变换，包括：确定所述逐元素操作节点无需进行数据访问索引关系的重映射，保持所述逐元素操作节点的输入数据和输出数据的分块格式标注不变，直接透传执行。

17、可选的一些实施例中，所述将所述第二计算图编译为所述目标硬件的可执行代码，包括以下至少一种：将所述第二计算图中相邻的逐元素操作融合为单个计算内核；根据所述目标硬件生成对应的计算内核代码；生成用于编排各计算内核执行顺序的调用封装代码。

18、第二方面，本发明提供一种图编译装置，包括：图获取模块，用于获取待编译模型对应的初始计算图，所述待编译模型部署于目标硬件，所述目标硬件采用分块格式数据布局；子图替换模块，用于将所述初始计算图中的目标算子节点替换为硬件特定子图，所述硬件特定子图通过图捕获机制从所述目标硬件的算子分发逻辑提取得到，以在所述初始计算图中显式引入线性格式与分块格式之间的数据格式转换操作，得到第一计算图；冗余消除模块，用于对所述第一计算图执行冗余格式转换消除，得到第二计算图，所述冗余格式转换消除包括常量折叠、公共子表达式消除和跨操作往返转换消除中的至少一种；编译模块，用于将所述第二计算图编译为所述目标硬件的可执行代码。

19、第三方面，本发明提供一种计算机可读存储介质，其上存储有计算机程序，所述计算机程序被处理器执行时实现如上述第一方面任一项所述的图编译方法。

20、第四方面，本发明提供一种电子设备，包括存储器、处理器及存储在所述存储器上并可在所述处理器上运行的计算机程序，所述处理器执行所述计算机程序时实现如上述第一方面任一项所述的图编译方法。

21、本发明的技术方案具有如下优点或有益效果：

22、（1）本发明通过将隐含于算子分发逻辑中的格式转换操作显式暴露于计算图，并对所述计算图执行常量折叠、公共子表达式消除和跨操作往返转换消除中的至少一种冗余格式转换消除，解决了采用分块格式数据布局的人工智能芯片中格式转换操作大量冗余执行、消耗内存带宽和计算内核启动资源的问题，获得了以2层llama模型为例格式转换操作由基线90次降至18次、消除比例达80%、且优化效果随模型规模线性扩展而不随模型层数衰减的技术效果，所述优化思路并可推广至任何需要非标准数据布局的人工智能芯片架构。

23、（2）本发明通过从目标硬件的算子分发逻辑中提取目标算子的处理流程函数、利用图捕获机制进行图追踪以生成硬件特定子图、以操作类型与输入形状及数据类型的组合为缓存键复用子图，解决了硬件特定子图需人工编写图变换代码、易与底层分发逻辑不一致的问题，获得了自动提取子图、保证子图与底层分发逻辑一致性、新增硬件操作支持时仅需添加分发逻辑即可自动适配、并降低编译时间与代码维护成本的效果。

24、（3）本发明通过将输入占位符节点提升为常量属性读取节点、在编译器的宿主机端执行格式转换操作链并以预计算结果替换运行期操作链，解决了常量权重在运行期重复执行格式转换的问题，获得了完全消除常量权重运行期格式转换的效果（2层模型消除28次、32层模型消除448次、80层模型消除1120次）。

25、（4）本发明通过消除源数据类型与目标数据类型相同的类型转换节点，并按照视图变换节点、填充节点、线性格式到分块格式转换节点的顺序执行多阶段公共子表达式消除，解决了多个算子共享同一输入数据时语义等价的格式转换操作被重复执行的问题，获得了复用同一转换结果、消除重复格式转换的效果（2层模型消除12次）。

26、（5）本发明通过以变换边界与分块边界是否对齐、所涉维度是否为参与分块的维度、是否引入跨分块边界的不规则数据依赖作为判断条件，判定形状变换操作节点能否在分块格式下直接执行，解决了链内形状变换操作触发全量格式转换往返的问题，并确保了在分块格式下直接执行形状变换的计算语义正确性。

27、（6）本发明通过分别对张量形状重塑操作节点重新计算元素索引映射关系、对张量维度转置操作节点修改数据搬运描述符的步长参数、对逐元素操作节点保持分块格式标注直接透传，在分块格式下直接完成相应的形状变换，解决了链内形状变换操作引入额外格式转换的问题，获得了零拷贝形状重解释、单次数据搬运完成转置、零额外格式转换开销的效果。