CN115469879A - 一种基于多面体模型的自动调度生成方法 - Google Patents

Abstract

本发明提出了一种基于多面体模型的自动调度生成方法，涉及计算机、编译器等领域。本发明方法以算子计算描述的多面体模型的实例集和依赖关系为输入，对状态进行自动调度，生成调度树的过程是一个递归的过程，每一步可能有多种优化操作可以进行；本发明利用机器学习模型预测调度优化操作的收益，将进行的各优化操作预测到的收益的对数和作为调度收益的预测结果；根据预测的调度收益选取最优调度，输出调度树形式的多面体模型的调度，用于后端代码生成。本发明方法使用的机器学习模型只需要进行预训练，即可适用不同算子，与计算方式无关，在实际的编译期无需再调整，能加快编译时间，适应难以定义调度优化的目标或有多个调度优化目标的场景。

Description

一种基于多面体模型的自动调度生成方法

技术领域

本发明涉及计算机、编译器等领域，具体涉及一种基于多面体模型的自动调度生成方法。

背景技术

多面体模型相关的技术在编译优化领域中通常被用于程序的表示、分析和变换，比如嵌套循环的变换和优化。随着深度学习的部署等技术的发展，多面体模型也被用于深度学习编译器，用于深度学习算子的代码生成和调度优化。

如今，在深度学习、高性能计算等领域中，针对不同的硬件平台，有很多种加速计算的手段。一个常用的手段是将一些计算量大、比较常用的算子包装成独立的库，即算子库，但这样的手段需要大量的人力和时间，对算子在相应的硬件平台上的实现进行调整、优化，开发和维护的成本都很高，通用性也比较差，并且很难支持自定义算子的场景；另一种思路是将“计算”和“调度”单独抽象出来，对不同的算子描述其计算的方式，并为不同的硬件平台设计相应的“调度”；更进一步地，则是让“调度”的生成也自动化起来，为不同的硬件或平台设计通用的“调度”优化流程。TVM是一款开源的、端到端的深度学习模型编译框架，用于优化深度学习模型在CPU、GPU、ARM等任意目标环境下的推理运行速度。TVM实现从深度学习模型到算子(硬件上可执行的代码)编译。TVM中，Tensor Expression模块是为了能描述深度学习算子的计算方式；AutoTVM模块需要用户使用TVM的调度原语来描述算子调度优化方式，并提供调度优化参数的范围，比如分块大小的范围，然后AutoTVM模块通过自动调优的方式，来搜索效果最好的调度优化的参数。

“计算”库的建立相对比较轻松，并且也很适合自定义算子等场景的二次开发，因此一些项目已经有了成熟的“计算”的描述方式和“计算”库，并且通常也比较容易转换为多面体模型的实例集和依赖关系。但对于“调度”生成的自动化，相关的项目则比较少，而调度生成等相关的问题解决起来也比较困难。

发明内容

针对上述现有问题，本发明的目的是提供一种基于多面体模型的自动调度生成方法，基于多面体模型，实现一种自动的、相对通用的调度生成技术，从多面体模型的实例集、依赖关系等计算相关的信息，从而生成相应的调度。除了通常的编译优化任务会有自动调度生成的需求，在深度学习领域也有相应的需求，本发明方法可以满足这种需求。

本发明提供的基于多面体模型的自动调度生成方法，包括如下步骤：

(1)获取算子计算描述的多面体模型的实例集和依赖关系，生成状态集S和调度池；调度池初始为空；状态集中每个状态包括算子计算描述和调度树，调度树初始只包含一个根节点，算子计算描述为多面体模型的实例集和依赖关系；

所述调度树的结构和ISL(整数集库)的调度树的结构一致，叶节点的子问题的输入信息包括到上一子节点的实例集、实例之间的依赖关系、以及实例计算结果所分配到的内存；实例集使用ISL形式化的集合表示，依赖关系使用ISL形式化的二元关系表示；

(2)判断状态集S中状态是否都被枚举过，若否，取一个没有被枚举过的状态，继续执行下一步，否则，转下面步骤(4)执行；

(3)判断当前被枚举到的状态的调度是否生成完毕，即该状态所生成的调度树将实例集完全进行了调度，且后续无可进行的优化操作，若是，将该状态所生成的调度放入调度池中，转步骤(7)执行；否则，继续执行下一步；

(4)对状态根据预先定义的优化操作进行调度优化，将执行一个或多个优化操作，将进行的各优化操作预测到的收益的对数和作为调度收益的预测结果；预测的调度收益将保存在状态中；预训练一个机器学习模型，用于预测优化操作的收益，该机器学习模型的输入是量化的状态；训练该机器学习模型时，将硬件理论计算效率与实际生成的算子的计算效率的比值作为训练样本标注的收益；

(5)根据预测的调度收益，从中选择收益高的调度执行，生成新的状态加入状态集T，预先设置状态集T的大小，当有新的调度收益更高时，替换掉状态集T中收益低的调度；然后转步骤(2)执行；

(6)用状态集T更新状态集S，将状态集T初始化为空集，若状态集S为空，进入步骤(7)执行，否则转步骤(2)执行；

(7)若状态集S为空，或者调度池中的调度数量达到预期阈值，则结束调度生成，继续执行下一步，否则，转步骤(2)执行；

(8)根据用户需求，输出当前调度池中所有的调度，或者输出预测收益最高的调度；调度为调度树形式的多面体模型的调度。

本发明的优点与积极效果在于：

(1)本发明方法基于多面体模型，整体框架通用性较强；输入为多面体模型的实例集和依赖关系，输出为调度树形式的多面体模型的调度表示；对于不同的前端计算描述，以及不同的后端编译工具链、硬件平台，只需进行较少的调整即可适配。

(2)本发明方法通过机器学习等基于模型的方式，来搜寻效果最优的调度方式，相比现有技术，能够适应难以定义调度优化的目标或有多个调度优化目标的场景。

(3)本发明方法所需的机器学习模型只需要进行预训练，并且模型与计算方式无关，对于不同的算子可以使用相同的模型，在实际的编译期无需再调整模型；相比现有需要自动调优的技术能适应更多样的计算方式，编译时间相对也更短。

附图说明

图1是本发明深度学习编译部署系统结构图；

图2是本发明实施例实现的基于多面体模型的自动调度生成流程图。

具体实施方式

下面将结合附图和实施例对本发明作进一步的详细说明。

本发明提出了一种基于多面体模型的自动调度生成方法，以多面体模型的实例集和依赖关系为输入，而多面体模型可以从各种前端提取出来，经过本发明调度生成方法，最终能根据目标后端，生成优化好的多面体模型调度。

如图1所示，深度学习编译部署系统中，虚线框内的部分，就可以应用本发明所提出的方法，进行自动的调度生成。本发明所关系到的只有虚线框内的部分，主要是可以将用于算子计算描述的“多面体模型的实例集、依赖关系”，能够进行“自动调度生成”，并最终得到调度优化后的“多面体模型调度表示”。

本发明方法实现的主要思路是：主要考虑到多面体模型的调度的表示方式中，有一种树形结构的调度表示方式，被称为调度树，有ISL(Integer Set Library，整数集库)能支持调度树的表示、分析、变换和控制流AST(抽象语法树)的生成。再考虑到从多面体模型的实例集、依赖关系等计算相关的输入，“生成调度树的过程”，可以看作是在递归地生成调度。递归的每一步只生成部分调度或对输入进行较少的操作，对于后续调度的生成，则是规模更小或更受限的输入的“生成调度树的过程”。

以上述的思路，已经可以构成一个简单的递归形式的调度生成方案。但更进一步地，实际上每一步通常能有多种可能的操作可以进行，比如有若干组待分配内存的实例集的计算结果，当前剩余的内存可以分配给其中一组，这种情况就会出现多种可行的操作。

对于潜在的多种可行的操作，本发明方法可将固定规则、机器学习结合起来，对模型进行预训练，用于预测所有可行操作所带来的收益。其中，固定规则可以涉及并发度(线程数和硬件实际支持的并行数结合起来计算)、局部性(依赖到的计算、数据的分散程度)等，而机器学习模型则利用量化后的信息、固定规则所计算到的数据、所需要进行的操作，来预测各种调度优化操作能带来的收益。

同时，本发明也将上述的递归，更进一步地修改为Beam Search(束搜索)。即将“已生成的调度树、叶节点的子问题的输入信息”作为状态；对于每个状态，每次从其可进行的多种操作中，选择预测收益最高的若干个；对于所有新生成的状态中，只保留总预测收益最高的若干个，再循环继续相同的操作。

若状态的“调度”等信息已经生成完毕，此时会将生成的调度保存在“调度池”中。当“调度池”达到一定的数量时，则停止循环。后续可对“调度池”所有的调度生成代码、进行测试，保留算子效率最高的一个。

具体地，本发明的基于多面体模型的自动调度生成方法的一个实现流程，如图2所示。

图2中的①表示本发明所接收到的多面体模型实例集、依赖关系等输入，会先处理成为初始的“状态”，一般是只有一个根节点的调度树。然后用只有初始状态的单元集合初始化出“状态集”S，用空集初始化出“调度池”。

图2中的②表示最终的输出可以选择使用“调度池”中所有生成出的调度，然后使用者可以利用各个调度来进行测试，以选出效果最好的调度；同时，由于已经使用机器学习模型预测过各调度的总收益，因此可以直接提供预测的总收益最高的调度作为输出。调度池中的调度为调度树形式的多面体模型的调度。

图2中的③表示当“状态集”S非空，即还有可以继续处理的状态，并且“调度池”中调度的数量还没有达到预期的阈值，即生成的调度数量还不够，则进行图2中的④的流程，继续处理“状态集”S；当“状态集”S为空集，或是“调度池”中调度的数量已经达到预期的阈值，则进行图2中的②的输出流程。

图2中的④表示当“状态集”S还有没有被枚举过的状态，则进行图2中的⑤的流程，继续循环枚举“状态集”S中的状态；当“状态集”S中所有的状态均被枚举过，则进行图2中的⑥的流程，更换“状态集”。

图2中的⑤表示从“状态集”S中取出一个还有没有被枚举过的状态，进行后续图2中的⑦的流程，处理该状态。

图2中的⑥表示“状态集”S更新为生成出的新“状态集”T。同时，该流程也会将“状态集”T重新初始化为空集。然后重新进行图2中的③表示的流程。

图2中的⑦表示当被枚举到的状态所生成的调度树已经将实例集完全地进行了调度，并且通常也没有后续可进行的调度优化操作，则进行图2中的⑨表示的流程；当被枚举到的状态所生成的部分调度树并没有将所有的计算完整地进行调度，则进行图2中的⑧表示的流程。

图2中的⑧表示对于当前被枚举到的状态的各种可行的操作，即调度优化，预先记录优化操作的前提条件和对状态的调整方式等，对状态调度优化时，将遍历设置的优化操作，若符合执行条件即可执行相应的优化操作。状态调度优化将包含一个或多个优化操作，将进行的各优化操作预测到的收益的对数和作为调度收益的预测结果。本发明方法使用机器学习模型预测优化操作的收益，将“状态”量化为机器学习模型的输入，例如，将实例到调度的关系，使用整数集库ISL的函数isl_basic_map_inequalities_matrix转为矩阵形式后，使用预训练的循环神经网络来处理为固定大小的张量，再使用机器学习模型预测调度优化能带来的收益。其中，对机器学习模型预训练时，使用硬件理论计算效率与实际生成的算子的计算效率的比值作为优化操作收益的标注结果，对状态的一系列优化操作以进行的各种优化操作预测到的收益的对数和作为状态调度收益的预测结果。在该流程预测各操作的收益后，进行后续图2中的⑩表示的流程。对一个状态，将产生一个或者多个调度，每个调度将通过训练好的机器学习模型获得收益预测结果，收益预测结果将存储在调度生成的状态中。

图2中的⑨表示将当前被枚举到的状态所生成的调度，放入调度池中，然后进行图2中的④表示的流程。

图2中的⑩表示根据机器学习所预测的各种操作的预期收益，选择较高的几个进行，生成新的若干状态，然后进行图2中的

表示的流程。根据预测的调度的收益的大小，从调度池选择收益高的若干调度执行，生成新的状态。

图2中的

表示将生成的新状态加入到“状态集”T中，并当“状态集”T的大小超过一定阈值时，则只保留“状态集”T中预测到的收益的对数和最大的若干个，然后进行图2中的④表示的流程。状态集T的大小预先设定，当有新的调度收益更高时，替换掉状态集T中收益低的调度。当没有可以进行的优化操作，即没有新的状态生成时，状态集T为空。

细节上，本发明状态集中的状态中的已生成的调度树的结构和ISL的调度树的结构一致，但由于只是生成的部分调度，因此部分性质、不变量并不一致，因此实际的实现中，需要另外定义相应的结构，并最终能生成为ISL的调度树。状态中调度树的叶节点的子问题的输入信息，主要包括可抵达相应子节点的实例集、实例之间的依赖关系、实例计算结果所分配到的内存等信息，其中实例集可以使用的是ISL形式化的集合表示，依赖关系可以使用ISL形式化的二元关系来表示；另外，也需要根据前端，附加上计算相关的信息，以保证在调度生成后，各实例可以正确地进行后续的代码生成。

另外，对于最终完整的调度生成，由于调度本身是一个从实例集到调度空间的一个二元关系，一个可用于生成AST的调度树通常需要该二元关系为双射。本发明的调度树生成流程中，可以针对已经生成的部分调度，即非双射的二元关系，利用现有的ISL的调度树的生成，或利用其他相关工作，生成完整的、可用于生成AST的调度树。

本发明上述可行的操作，即其他项目所称的“调度优化”。可以生成通常的ISL的“整数仿射”并构建调度树的Band类型节点，比如状态允许进行“循环分块”的操作，则可以生成相应维度除以分块大小并向下取整的映射；如果情况允许，比如说循环的范围没有超过某个阈值，则可以在Band节点上进行标注，来让后续生成的循环展开，或成为并行线程的ID。或者，在已经确定部分中间计算结果的内存分配情况下，可以进行实例集分割的操作，生成Filter(过滤器)节点，并将问题规模缩减。

另外，可行的操作也可以包含“为计算结果分配内存或寄存器”等操作，即将许多项目所称的Bufferize(分配内存，并实例化内存操作)，在调度生成阶段就完成。比如已经只剩下一组计算没有分配内存，并且此时片上的内存对齐后也足够分配给该计算，则可以生成一个Mark(标注)节点，并在节点保存的值中标注相关的信息，或者标注为独立的标识符，并在特定的映射数据结构中保留相关信息；并且，在Mark节点的子节点上，生成分配好内存后的子问题的输入。这样就生成出了该操作之后的新状态。而且，根据硬件设备的不同，也可以在选择不同的对齐方式，将各种偏移量、对齐方式的组合，均作为独立的操作，并由机器学习模型来预测最优的若干种。

同时，各种操作也可以是将输入状态进行向量化、张量化等与目标硬件平台的计算加速能力相关的调度优化。比如，目标硬件支持在片上内存直接进行向量化加法的计算加速，但需要参数是对齐的指针。而此时某个叶子节点只剩下一组加法计算，并且计算结果已经分配到了片上内存，同时输入也可从片上内存中获取到，依赖关系也符合向量化加法的形式。若此时存在有某个参数并没对齐，则可以考虑将计算切分开，生成Filter节点，将能对齐和不能对齐的计算区分开，生成新的状态，在后续的循环中，完成向量化；若所有参数均对齐了，则可以生成Extension节点，将计算映射为相应的向量化操作，并使用Mark节点标注或使用映射数据结构保存信息。

总之，通过上述的流程，本发明可以逐步地将各种调度优化、加速手段，应用到生成的调度中，将多面体模型的实例集、依赖关系等计算相关的信息，生成出相应的调度树。由于多面体模型本身是形式化的表示方式，通用性较强，可以很方便地从各种已有的项目中，从原有的计算的表示方式，生成出相应的多面体模型的实例集和依赖关系；而调度树通过ISL可以生成出控制流AST，可以很方便地接入到各种后端，生成相应硬件平台的算子。从而充分利用已有的计算描述和计算库，用更通用的方式为各种硬件平台进行调度生成。

除说明书所述的技术特征外，均为本专业技术人员的已知技术。本发明省略了对公知组件和公知技术的描述，以避免赘述和不必要地限制本发明。上述实施例中所描述的实施方式也并不代表与本申请相一致的所有实施方式，在本发明技术方案的基础上，本领域技术人员不需要付出创造性的劳动即可做出的各种修改或变形仍在本发明的保护范围内。

Claims (3)

1.一种基于多面体模型的自动调度生成方法，其特征在于，包括如下步骤：

(1)获取算子计算描述的多面体模型的实例集和依赖关系，生成状态集S和调度池；调度池初始为空；状态集中每个状态包括算子计算描述和调度树，调度树初始只包含一个根节点，算子计算描述为多面体模型的实例集和依赖关系；

所述调度树和ISL的调度树的结构一致，叶节点包含的子问题的输入信息包括到上一子节点的实例集、实例之间的依赖关系、以及实例计算结果所分配到的内存；实例集使用ISL形式化的集合表示，依赖关系使用ISL形式化的二元关系表示；ISL表示整数集库；

(2)判断状态集S中状态是否都被枚举过，若否，取一个没有被枚举过的状态，继续执行下一步，若是，转下面步骤(6)执行；

(3)判断当前被枚举到的状态的调度是否生成完毕，即该状态所生成的调度树将实例集完全进行了调度，且后续无可进行的优化操作，若是，将该状态所生成的调度放入调度池中，转步骤(7)执行；否则，继续执行下一步；

(4)对状态根据预先定义的优化操作进行调度优化，将执行一个或多个优化操作，将进行的各优化操作预测到的收益的对数和作为调度收益的预测结果；预测的调度收益将保存在状态中；预训练一个机器学习模型，用于预测优化操作的收益，该机器学习模型的输入是量化的状态；训练该机器学习模型时，将硬件理论计算效率与实际生成的算子的计算效率的比值作为训练样本标注的收益；

(5)根据预测的调度收益，从中选择收益高的调度执行，生成新的状态加入状态集T，预先设置状态集T的大小，当有新的调度收益更高时，替换掉状态集T中收益低的调度；然后转步骤(2)执行；

(6)用状态集T更新状态集S，将状态集T初始化为空集，若状态集S为空，进入步骤(7)执行，否则转步骤(2)执行；

(7)若状态集S为空，或者调度池中的调度数量达到预期阈值，则结束调度生成，继续执行下一步，否则，转步骤(2)执行；

(8)根据用户需求，输出当前调度池中所有的调度，或者输出预测收益最高的调度；调度为调度树形式的多面体模型的调度。

2.根据权利要求1所述的方法，其特征在于，所述的方法，生成调度树的过程，是在递归地生成调度，包括：(a)将已生成的调度树、叶节点的子问题的输入信息作为状态，对每个状态，从可进行的操作中选择预测收益最高的一个或多个操作，从新生成的状态中保留总预测收益最高的多个状态，(b)对新生成的状态继续执行(a)。

3.根据权利要求1所述的方法，其特征在于，所述的步骤(3)中，预先定义的优化操作包括：生成ISL的整数仿射，对状态进行循环分块操作，构建调度树的Band节点；在确定中间计算结果的内存分配情况下，进行实例集分割的操作，生成过滤器节点；为计算结果分配内存或寄存器，生成Mark节点；将输入状态进行向量化或张量化。

Images