整数集合和仿射函数构成了多面体模型表示程序的基础。利用整数集合,多面体模型可以表示循环嵌套的迭代空间;利用仿射函数,可以表示程序语句实例和被访问的内存地址单元之间的访存关系。通过定义整数集合和仿射函数的操作,优化编译器也可以很方便地计算出语句实例之间的依赖关系,以及语句实例的执行顺序关系,即语句实例的 调度 (schedule)。关于依赖关系,我们将在第3章进行更具体和详细的介绍。对于语句实例的执行顺序,任意一个程序至少有一个执行顺序关系,这个顺序就是根据程序在文本中出现的顺序定义的,我们将其称为 原始调度 。
优化编译器利用抽象模型对程序进行表示,并在这些抽象表示的基础上进行各种运算,就是为了要计算一个 新的调度 ,这个新的调度是以提升某种程序特征为目的而执行的运算。在多面体模型的发展过程中,有过许多表示调度的方式,一个比较直观的表示方法就是采用调度树(schedule tree),这种表示方式也被目前大量的深度学习编译器和相关项目所采用或改进。
如果只考虑一个循环嵌套,并且该循环嵌套内只有一个语句,那么显然如式(2-10)所示的仿射函数可以用来表示一段代码的执行顺序,但当程序中存在多个循环嵌套,并且循环嵌套内有多个语句的时候,利用仿射函数来表示调度并没有那么简单。考虑图2.5中的例子,这段代码由两个循环嵌套构成,并且第二个循环嵌套包含多个语句。显然,直接用仿射函数表示这段代码内语句的执行顺序并不那么容易。
图2.5 一个用于说明不同调度表示的例子
用仿射函数来表示这段代码内语句的执行顺序主要面临两个问题。首先,如果用 S 1 ( i ) → ( i )和 S 3 ( i ) → ( i )表示语句 S 1 和 S 3 所有语句实例的先后顺序,那么这两个语句之间的先后顺序无法确定,即仿射函数没有体现出这两个语句属于不同循环嵌套的信息。其次,在相同循环嵌套内的语句 S 2 和 S 3 ,其仿射函数分别可以用 S 2 ( i,j ) → ( i,j )和 S 3 ( i ) → ( i )表示,但这两个语句在同一个循环嵌套的先后顺序也不明确。
多面体模型解决这个问题的方法是在仿射函数的值域中引入标量维度,用于表示不同语句之间的先后执行顺序。对于图2.5中的代码,优化编译器可以用
[{( i ) → (0 ,i )};{( i,j ) → (1 ,i, 0 ,j )};{( i ) → (1 ,i, 1)}] (2-43)
依次表示三个语句的调度。其中,每个仿射函数值域第一个维度上的标量0或1表示 S 1 所在的循环嵌套和 S 2 、 S 3 所在循环嵌套之间的先后顺序, S 1 对应的仿射函数第一个维度为0,而 S 2 和 S 3 的仿射函数第一个维度均为1,表示 S 1 属于一个循环嵌套,而 S 2 和 S 3 属于另外一个循环嵌套。根据程序语句的先后顺序,假设总是从0开始对循环嵌套进行编号。当然,也可以从任意其他整数开始编号。我们也可以假设在该例两个循环嵌套外层还有一个虚拟的外层循环,该循环被展开,第一次迭代执行的是 S 1 所在的循环嵌套,第二次迭代执行的是 S 2 和 S 3 所在的循环嵌套。以此类推,就可以将标量维度和当前的仿射函数的定义统一起来。
S 2 和 S 3 的仿射函数第三个维度上的标量0和1则表示这两个语句在相同循环嵌套内的先后顺序,这两个语句仿射函数值域的前两个维度上的值完全一致,表明这两个语句被嵌套在相同的循环层内,即 i 循环。然而,这两个语句在第三维上是不同的标量,表明这两个语句在 i 循环内的执行顺序不同。与此同时, S 2 仿射函数的值域是一个四元组表示的时间序列, S 3 仿射函数的值域是一个三元组时间序列,这表明即使在相同循环嵌套内,不同语句的仿射函数值域的维度也可以不同。更进一步地, S 2 语句的执行顺序还会由内层 j 循环决定,而 S 3 则不会。
上述表示调度的方式是由Kelly提出的,因此也被称为Kelly 表示 (Kelly’s abstraction) [40] 。在Kelly表示中,每个语句对应的仿射函数的值域维度并不完全相同,在比较不同语句实例之间的执行顺序时很不方便。Kelly表示比较时间先后顺序采用的是非严格的字典序比较方式 。对于式(2-43)中所示调度,Kelly表示总是认为(0 ,i )的时间字典序小于(1 ,i, 0 ,j )和(1 ,i, 1),而(1 ,i, 0 ,j )的时间字典序小于(1 ,i, 1)。换句话说,这三个时间节点的先后顺序依次是(0 ,i )、(1 ,i, 0 ,j )和(1 ,i, 1)。
与Kelly表示类似的是一种被称为2 d +1 表示 (2 d +1 representation) [33] 的调度表示方式,这种表示方式可以看作Kelly表示的一种特殊形式。既然Kelly表示是一种带有标量维度的仿射函数,那么对一个 d 维语句 S ( i )的 d 元组,可以构造一个 d +1维的向量,该向量的前 d 个分量构成 i ,最后一个分量代表标量1。对于图2.5中的语句,三个语句对应的向量分别为( i, 1)、( i,j, 1)和( i, 1)。针对每个 d +1维的向量,2 d +1表示构造一个(2 d +1) × ( d +1)的矩阵,用该矩阵表示每个语句的调度。仍然考虑图2.5中的三个语句,优化编译器可以用
表示这三个语句的调度,将这些矩阵与其对应的 d +1维向量分别相乘,可以得到与Kelly表示类似的2 d +1维的仿射函数,即可以用
[{( i ) → (0 ,i, 0)};{( i,j ) → (1 ,i, 0 ,j, 0)};{( i ) → (1 ,i, 1)}] (2-45)
表示与式(2-44)等价的调度。注意式(2-45)中每个仿射函数的值域都是一个2 d +1元组整数,这也是该表示方式名称的由来。与Kelly表示相比,2 d +1表示由每个语句外层循环索引变量和标量相间的向量组成,一共有 d 个循环索引变量和 d +1个标量,除最左端的标量外,每个标量表示前一个循环索引变量对应循环内该语句的执行顺序。最左端的标量代表所有循环嵌套的执行顺序。如果我们把最外层看作一个虚拟的循环,那么所有的标量代表的含义一致。由于2 d +1表示中的矩阵的维度是固定的,所以该表示方法无法直接表达如循环分块等改变循环嵌套维度的变换,必须借助其他辅助手段,如修改每个语句在迭代空间中表示的维度等。
Kelly表示中可能存在维度不同的仿射函数。将维度较小的仿射函数进行零填充(zero padding),使Kelly表示中所有的仿射函数维度都与维度最大的那个函数的维度一致,就可以得到union map 表示 (union map representation) [69] 。例如,式(2-43)中的调度可以表示成
[ S 1 ( i ) → (0 ,i, 0,0); S 2 ( i,j ) → (1 ,i, 0 ,j ); S 3 ( i ) → (1 ,i, 1,0)] (2-46)
的形式, S 1 和 S 3 的仿射函数都需要进行零填充,以将其值域扩展成为四元组整数。由于所有语句的仿射函数值域维度相同,union map表示允许不同语句的仿射函数之间进行严格的字典序比较,这可以简化工程实现所需的工作量。此外,与式(2-43)和式(2-44)不同的是,union map表示对每个仿射函数进行了命名,这使得调度的表示在本质上与访存关系、依赖关系的表示一样,都可以用命名的仿射函数进行表示,使程序表示的抽象程度进一步提升。
还有一种表示调度的方式,也是本节将要重点介绍的 调度树 (schedule tree)表示 [37] 。这种表示方式与前面几个方式最大的不同在于,调度树将语句的执行顺序表示成一种树状结构,如图2.6所示是图2.5所示代码的调度树表示,这种表示与Kelly表示式(2-43)、2 d +1表示式(2-44)以及union map表示式(2-46)的表达能力相同,但比前面几种方式的表达更直观。调度树表示引入了一系列具有不同功能的结点,这些结点可以表示和改变程序的不同特征。例如,调度树表示并不直接表示调度的标量维度,而是通过一种被称为sequence的结点来表示仿射函数中的标量维度。下面我们来具体地介绍调度树中的结点。
图2.6 图2.5所示代码的调度树表示
调度树是一种不同类型的结点相互连接而成的树状结构表示,在多面体模型中用于表示程序语句的调度。调度树表示沿用了传统编译器中常用的中间表示抽象——抽象语法树,来表示语句之间和语句实例之间的执行顺序。调度树表示不仅比上述几种表示方式更加直观,而且通过对不同类型结点的支持,能够同时封装程序不同的属性,便于后端更底层的变换和代码生成等任务的实现。此外,调度树表示也支持一些特殊功能的实现,这在Kelly表示、2 d +1表示等抽象上实现起来并不是一件简单的事情。多面体模型既要支持对原始程序语句调度的表示,又要能够表示经过调度变换之后的语句执行顺序。对一种调度的表示方法而言,这意味着需要能够同时支持调度变换前后的输入和输出。为了便于读者对本章内容的理解,我们先简单介绍一下调度变换的过程,更详细的实现调度变换的算法将在第4章中介绍。
绝大多数的调度算法都是基于程序的依赖关系计算新的调度。在构造新的调度之前,优化编译器都会将依赖关系表示成依赖图的形式,用以表示每个语句实例之间的依赖关系。依赖图的一个结点表示一个语句,每个边表示语句之间的依赖关系。如果某个语句 S 2 依赖于语句 S 1 ,意味着在调度变换前语句 S 1 需要在 S 2 之前执行,这种先后顺序在新的调度中也必须被满足。
一个调度算法可以递归地分解依赖图,在每次递归时都将依赖子图分解为若干个强连通分量,然后分别计算每个强连通分量的内语句之间的部分或局部调度,并用偏序关系(partial relation)表示。这种局部调度一般都可以用一个仿射函数表示,调度算法将所有的局部调度联结在一起,最终形成整个程序内所有语句的全局调度。这种调度算法的计算过程本身就蕴含了一种内在的树状结构关系,从而促成了调度树表示的设计。调度树的结点类型包括以下几种。
(1)domain 结点 :domain结点通常是一个调度树的根结点,代表由该调度树表示的程序语句实例的集合,即这段程序的迭代空间。因此,一个domain结点可以用一个命名整数集合的并集表示,每个整数集合代表一个独立的语句。如图2.6中的根结点是一个domain结点,表示图2.5代码的迭代空间。
(2)context 结点 :context结点用于定义编译符号常量及其约束。一个编译符号常量可以是与当前调度树对应程序的范围内的全局符号常量,也可以是一个子树范围内的局部符号常量。当前context结点内引入的符号常量只能在其后继子树上使用。当context结点是domain结点的子结点,并且当前context结点内只包含全局符号常量时,该context结点可以不必显式地表示在调度树中。如图2.6中context结点就被省略了。如果在domain结点后插入一个根结点,那么该context结点内是一个关于编译符号常量 M 和 N 的整数集合的并集。
(3)band 结点 :band结点用一个仿射函数集合的并集表示,它总是作为一个domain结点或filter结点或extension结点或expansion结点的子结点出现在调度树中,表示其父结点内语句实例的调度。band一词源自Pluto算法 [19] ,用于指代寻找满足Pluto算法代价模型的“循环带”或多重循环,因此band结点可以对应为代码中的循环嵌套。band结点可以有多个成员(member),每个成员对应循环嵌套中的一个循环维度。一个band结点有两个属性permutable和coincident,前者为一个布尔型变量,当取值为1时表示由该band结点表示的循环嵌套每个循环维度相互交换是合法的;后者为一个维度与循环嵌套层数相同的向量,向量的每个分量为一个布尔型变量,当该分量取值为1时表示该分量对应的循环维度可以被并行执行。此外,band结点中还包含许多控制代码生成过程的选项,包括指导代码生成实现分块分离、循环展开等程序变换,这些变换都将在第4章介绍。如图2.6所示,所有以[]封装的仿射函数都是band结点。
(4)sequence 结点 :sequence结点是调度树用来显示表达传统调度表示中标量维度的结点,如图2.6中domain结点的子结点是一个sequence结点。与其他结点类型不同,sequence结点既不是整数集合,也不是仿射函数。sequence结点可以有多个子结点,其子结点从左到右依次按序执行,并且sequence结点的子结点只能是filter结点。
(5)set 结点 :set结点与sequence结点的语义类似,总是有多个filter结点作为其子结点。与sequence结点的不同点在于set结点的子结点之间可以按照任意顺序执行。换句话说,set结点并不对其子结点的先后执行顺序有所限制,但是代码生成还是会按照其子结点从左到右的顺序生成代码。由于子结点之间并没有先后顺序关系,在生成的代码中将两个子结点对应的代码顺序交换并不会违反程序的语义。
(6)filter 结点 :一个filter结点往往作为sequence或set结点的子结点出现在调度树中,如图2.6中sequence结点的两个子结点都是filter结点,每个filter结点由一个整数集合的并集表示。filter结点与domain结点都是用整数集合的并集表示,不同之处在于domain结点可以是一个调度树的根结点,但filter结点不能作为一个调度树的根结点出现。filter结点的作用是用于“过滤”由domain结点封装的不同整数集合。当两个不同的整数集合对应的语句需要采用不同的调度时,调度树内引入一个sequence或set结点来将其过滤到不同的子树中。另外,filter结点也可以用于“过滤”expansion结点、extension结点以及另外一个filter结点。
(7)mark 结点 :mark结点用于向调度树插入任意信息,这些信息可以被后面的代码生成工具解析,按照特定的模式来生成代码。例如,当面向GPU生成代码时,代码生成工具需要确定哪些band结点需要被映射到GPU的哪些硬件上,此时就可以借助mark结点指导代码生成工具进行特定的代码生成。mark结点中存储的是一个字符串信息。
(8)extension 结点 :extension结点用于在调度树内添加一些没有被domain结点涵盖的语句,使基于调度树的代码生成工具自动生成一些特殊指令。调度树和2.4.1节中介绍的几种调度表示方式都是对原始程序中包含的语句寻找一个执行先后顺序的关系,但在一些特定的应用场景中,需要优化编译器生成一些在原始程序中不存在的语句。extension结点的语义需要使用仿射函数集合的并集表示,一个最典型的应用场景就是生成显式的数据传输语句。在一些特殊的情况下,extension结点也可以作为调度树的根结点出现。
(9)expansion 结点 :与extension结点类似,expansion结点的语义也是由仿射函数集合的并集表示。expansion结点每个仿射函数的定义域是被“过滤”到当前结点的一个或多个语句实例的集合,该仿射函数将其映射到一个或多个新的语句实例集合上。expansion结点可以用于将一个循环嵌套内的语句进行组合(grouping),使得这些语句在整个调度过程中总是按照相同的方式进行调度。expansion结点也可以用于实现满射函数,这在一些特定的循环变换中发挥着重要作用。从名称上来讲,expansion结点和extension结点很容易混淆,但它们的作用完全不同。
(10)guard 结点 :guard结点与context结点的作用类似,用于描述编译符号常量和外层band结点对应循环索引变量对当前子树的约束。例如,图2.1所示循环嵌套内,当考虑内层 j 循环的代码生成时,外层 i 循环的循环索引变量可以被看作内层 j 循环对应子树的一个编译符号常量,内层循环必须要满足 i 的范围约束,而且在内层循环的执行过程中 i 的值不会发生改变。
(11)leaf 结点 :leaf结点是一个不包含任何信息的结点,其作用是表示一个调度树的分支终点,因此,leaf结点并不会被显式地出现在一个调度树的表示当中。在实际应用中,leaf结点的设计是为了遍历整个调度树,方便工程设计与实现。
上述多种不同类型的结点构成了调度树表示的基础,为调度树表示的实际应用提供了许多便利。这些不同类型的结点中,只有set和sequence结点可以有多个子结点,只有leaf结点没有子结点,其他所有类型的结点都只有一个子结点。一个调度树可以通过合并多个不同的子树来生成,也可以通过在其结点上的操作来改变表示的含义,从而实现程序变换。
经过变换之后,优化编译器需要根据编译阶段实现的优化生成代码。代码生成将在第7章进行介绍。在多面体模型中,代码生成是以迭代空间和变换后的调度为输入生成代码,多面体模型中的代码生成也被称为多面体扫描(polyhedra scanning)。调度树表示将迭代空间和调度封装在一起,使得代码生成可以通过扫描调度树表示来完成。因此,调度树的操作不仅应该支持各种程序变换,还应该为代码生成提供便利。下面列举几种调度树上的基本操作。
(1)在调度树的某个位置上插入context、filter或mark结点。插入context结点允许编译器向子树引入额外的局部符号常量,这个符号常量可以只局限在当前子树上;插入filter结点允许一个整数集合并集内的不同集合被独自处理;插入mark结点则允许编译器向调度树嵌入任何额外信息,支持定制化代码生成。值得注意的是,不能向sequence或set结点与其子结点之间插入新的结点。
(2)通过修改band结点的内容可以实现包括所有幺模变换和循环分块以及循环分段等在内的循环变换。第4章将详细介绍各种循环变换,其中大部分的循环变换可以通过修改band结点的仿射函数来完成。在band结点上的操作还可能与第4章介绍的迭代空间分裂(index set splitting)一起来实现一些更复杂的变换。此外,通过修改band结点的选项、成员的属性等还可以实现循环展开、分块分离等变换。
(3)将一个band结点分裂(split)成嵌套的两个或多个band结点,每个band结点维护分裂前band结点的部分信息。注意分裂后的band结点之前是父子关系。band结点的分裂往往发生在将软件循环映射到硬件的过程中。当软件循环嵌套的层数比硬件并行的维度大时,band结点需要分裂,以适配并行硬件维度。band结点的分裂并不会导致permutable属性的变换,coincident的向量维度会变小,但是每个向量分量的值维持不变。
(4)将两个或多个嵌套的band结点进行组合(combine),即band结点分裂的逆过程。对应地,这种操作可以是在软件循环的维度小于并行硬件的维度时采用的一种适配硬件的操作。组合后band结点的permutable属性由组合前每个band结点对应的属性决定,coincident向量是组合前每个band结点的组合。
(5)将两个或多个并列的band结点进行合并(fuse)。注意并列band结点的合并和嵌套band结点的组合之间的区别。并列band结点的合并过程一般发生在调度算法的计算过程中。一些特定的调度算法对依赖图的强连通分量分别构造band结点,并试图合并不同强连通分量的band结点,以期构成循环嵌套层数更多的band结点,这样既可以挖掘更多循环维度上的分块可能性,也实现了循环合并。对并列band结点进行合并的过程还有可能会改变原来band结点的permutable和coincident等属性,并且合并后可能还需要引入一个新的sequence结点作为其子结点。
(6)将一个band结点分布(distribute)成两个或多个并列band结点,即band结点合并的逆过程。该操作用于实现循环分布,实现过程也有可能会改变原来band结点的permutable和coincident等属性。
(7)对band结点进行分块操作。band结点的分块变换可以通过修改band结点的仿射函数实现,这与传统的调度表示方式的操作一致。此外,也可以对band结点进行分块之后,将band结点的近似仿射函数进行分段,并对band结点进行分裂,将band结点分裂成用于迭代循环分块之前的维度和用于迭代循环分块内的维度,并在此基础上将band结点的不同维度映射到不同的并行硬件上。
(8)将band结点的位置下沉(sink)到其子结点的后继位置。这种操作可以和band结点的分裂一起使用。当一个band结点的子结点是一个sequence结点,并且sequence结点后有多个filter子结点时,这种操作用于修改band结点的位置。当面向GPU生成时,优化编译器可以首先将一个band结点进行分裂,得到两个嵌套的band结点,分裂后外层band结点映射到相同的block组内,然后将分裂后得到的内层band结点下沉到每个filter结点后,将不同的filter结点对应的计算映射到不同的thread组内。
(9)对sequence结点的子结点进行重排序,以获得满足调度算法优化目标的执行顺序。
上面几种是调度树表示上的一些基本操作,这些操作都是在调度树的一个结点上实现的。在工程实现上,还允许将一个子树嫁接到另外一个调度树上,将其作为新的子树与整个调度树一起执行,这种操作可以看作多个结点操作的组合。另外,调度树上的操作也不仅局限于上述几种操作,在工程实践的过程中可以根据每个结点的语义灵活运用。
调度树表示已经成功应用于许多基于多面体模型的优化编译器和工具中。与传统的调度表示方法相比,调度树并没有在表示能力上有所提升,但在许多方面为程序变换和代码生成提供了更好的灵活性。我们从以下几方面对比当前几种调度表示的特征。
调度对象的粒度 。对于一段给定的程序片段,调度表示的粒度可以是一个程序语句,也可以对整个程序片段构建一个统一的表示。从上面的介绍和讨论中不难看出,调度树表示是面向被分析的整个程序的,因为被分析的程序片段所有的语句都被统一地表示在调度树中。相反,Kelly表示和2 d +1表示则是以程序的每个语句为基本单元来构建调度的。调度树和传统的几种调度表示之间可以相互转换,但树状结构的设计使调度树表示比传统表示更为直观。
局部调度的表示 。在利用好仿射函数的同时,几种调度表示对局部调度的表示能力也存在差异,这体现在调度表示对仿射函数的约束。2 d +1表示要求调度的每个维度是一个仿射函数,而其他几种调度表示支持分段近似仿射函数,这使得循环分块这一非常重要的循环变换在2 d +1表示中的实现并不那么简单,但包括调度树在内的其他几种表示中,这种循环变换的实现就相对容易得多。
标量维度的表示 。在实现循环合并和循环分布时,标量维度是调度表示不可或缺的一个重要因素。由于调度对象的粒度是程序的一个语句,Kelly表示和2 d +1表示引入了标量维度,来表示循环嵌套之间的相对顺序。union map表示将标量维度扩展到仿射函数集合的并集中,而调度树中则用sequence结点来显式地表示这一特征。
支持多语句的组合调度 。调度变换的过程对程序进行了抽象的同时,往往忽略了程序自身携带的一些领域特定信息。程序员在编写程序时,往往期望在一个循环嵌套内的多个语句总是能够被“绑定”在一起。由于一些调度算法的计算过程以语句为对象计算新的调度,导致这种基于语句计算调度的过程不仅可能会使编译时间复杂度过高,也可能导致这些语句在新的调度下分散在不同的循环嵌套中。调度树通过expansion结点支持将这些语句进行组合调度,这种功能在其他调度表示中是不具备的。
对单射函数的支持 。当表示语句调度的一个仿射函数是单射函数时,意味着不同的语句实例在该函数的作用下会在同一时刻被执行,即语句之间的并行。这种能力是当前几种调度表示都支持的,因为自动识别和表示并行是对调度表示的一个基本要求。
对满射函数的支持 。当表示语句调度的一个仿射函数是满射函数时,意味着相同的语句实例在该函数的作用下会在不同的时刻被执行,即一个语句实例可能被执行多次。Kelly表示和2 d +1表示并不支持满射函数,而调度树表示中可以通过expansion结点的使用来实现满射函数的构建。
偏序关系的比较 。在比较由调度表示定义的偏序关系时,更严格的定义是两个被比较的偏序关系需要具备相同的维度。union map表示仅支持严格的偏序关系比较,而其他几种调度表示对这种操作进行了扩展,允许不同维度的偏序关系之间进行比较,我们将这种比较称作一种非严格的偏序关系的比较。显然,非严格的偏序关系的比较在工程实践上更灵活一些。