PolyTOPS: Reconfigurable and Flexible Polyhedral Scheduler

作者/机构: Gianpietro Consolaro∗‡, Zhen Zhang∗, Harenome Razanajato∗, Nelson Lossing∗, Nassim Tchoulak∗, Adilla Susungi∗, Artur Cesar Araujo Alves∗, Renwei Zhang†, Denis Barthou∗, Corinne Ancourt‡, Cedric Bastoul∗, ∗Huawei Technologies France, Paris, France †Huawei Technologies Co., Ltd., Beijing, China ‡Mines-Paris PSL University

A1 主要贡献

本文针对现代编译器和深度学习框架中广泛使用的多面体模型优化技术,指出现有的多面体调度器(如Feautrier、Pluto、isl、Tensor Scheduler)各自针对特定的架构或应用场景,缺乏足够的灵活性和可配置性来适应日益异构的硬件环境,特别是像NPU(神经处理单元)这样的新兴架构。例如,在华为MindSpore的AKG(自动算子生成)项目中,默认的isl调度器在某些情况下无法为昇腾(Ascend)NPU生成最优的循环变换(如代码清单1所示),导致性能不佳。通过后处理(post-processing)来修正调度结果是一种繁琐且容易出错的方法。

为了解决这一核心问题,本文提出了一种名为 PolyTOPS 的新型可配置多面体调度器,其目标是提供一个灵活、可调整的工具,以应对不同架构、编译环境和应用领域的需求。

核心创新点如下:
1. 可配置性(Configurability):PolyTOPS 提供了一个富有表现力且易于使用的高级配置接口(支持JSON和C++)。用户可以通过该接口详细配置迭代调度过程的各个方面,包括:
* 并行性控制、矢量化。
* 时间和空间局部性优化。
* 对语句的循环融合(fusion)与分裂(fission)进行细粒度控制。
* 指定部分调度策略。
通过简单的配置,PolyTOPS的行为可以模拟现有的调度器(如Pluto风格、isl风格等),或者定义全新的、针对特定场景或特定算子(kernel-specific)的调度策略。

  1. 灵活性(Flexibility):PolyTOPS 采用了一种多功能设计,摒弃了“一刀切”的调度器方法。它允许用户从一个通用的策略开始,然后针对算子中的特定循环或特定架构逐步进行调整。它为迭代调度器提供了一个可扩展的基础设施,其中约束可以从预定义策略到专用的转换启发式进行精细调整。

本文通过在不同场景下的实验验证了PolyTOPS的有效性。在MindSpore AKG中,针对昇腾NPU的自定义混合算子,PolyTOPS相比isl调度取得了7.66倍的几何平均加速比。在PolyBench基准测试中,针对不同的多核CPU架构,相比Pluto调度最高取得了1.80倍的几何平均加速比。

代码清单1:左侧为原始for循环代码,完全并行,适合GPU架构。右侧代码通过循环分发和语句0的循环交换,对NPU的矢量化数据加载和计算进行了优化。Pluto或isl或许能找到循环分发(需指定正确的融合启发式),但无法找到循环交换。
代码清单1:左侧为原始for循环代码,完全并行,适合GPU架构。右侧代码通过循环分发和语句0的循环交换,对NPU的矢量化数据加载和计算进行了优化。Pluto或isl或许能找到循环分发(需指定正确的融合启发式),但无法找到循环交换。

A3 背景知识

多面体模型基础

  1. 迭代域(Iteration Domain)定义。对于代码中的每个语句,迭代域表示该语句周围循环迭代器所取的值的范围。由这些迭代器组成的向量称为迭代向量 $\vec{it}$。迭代域被假定为迭代向量的多面体,并且可以依赖于参数。参数向量 $\vec{N}$ 由在代码执行期间保持不变的变量组成。语句S的域 $D_S$ 定义如下:

    $$ \mathcal{D}_S = \left\{ \vec{it} \middle| M_S \cdot \begin{pmatrix} \vec{it} \\ \vec{N} \\ 1 \end{pmatrix} \geq 0 \right\} $$


    其中 $M_S$ 是定义域多面体的矩阵。

  2. 依赖关系与合法性(Dependencies and Legality)定义。从语句S到语句R的依赖关系 $\delta_{S \to R}$ 意味着为了保持程序的语义,语句S需要在语句R之前执行。该依赖关系定义在S和R的一组迭代向量值上,并满足以下约束:S和R访问同一内存位置(S或R中至少有一个是写操作),且S在R之前执行。这些约束与语句的域类似,也定义了一个多面体:

    $$\delta_{S \to R} = \left\{ \left. \begin{pmatrix} \vec{it}_S \\ \vec{it}_R \end{pmatrix} \right| M_{S \to R} \cdot \begin{pmatrix} \vec{it}_S \\ \vec{it}_R \\ \vec{N} \\ 1 \end{pmatrix} \geq 0 \right\}$$
  3. 调度函数(Scheduling Function)定义。调度函数 $\Theta$ 将其域中的每个语句和迭代向量映射到一个唯一的多维时间戳。时间戳通过字典序进行全序排列。对于一个语句S,$\Theta_S$ 是一个多维函数,其每个维度由仿射形式 $\phi_{S,i}$ 定义。这些仿射形式依赖于迭代器 $\vec{it}$ 和参数 $\vec{N}$。$\Theta_S$ 可定义如下:

    $$ \Theta_S: \begin{array}{rcl} \mathcal{D}_S(\vec{N}) & \to & \mathbb{N}^m \\ \underset{\vec{it}}{\phantom{\mathcal{D}_S(\vec{N})}} & \mapsto & (\phi_{S,0}(\vec{it}) \dots \phi_{S,m-1}(\vec{it})) \end{array} $$


    其中 $m$ 是调度维度的数量,$\phi_{S,i}$ 定义为:

    $$\phi_{S,i}(\vec{it}) = T_{S,i} \cdot \begin{pmatrix} \vec{it} \\ \vec{N} \\ 1 \end{pmatrix}$$
    其中 $T_{S,i}$ 是变换向量。
    补充说明。对于一个被 $k$ 个嵌套循环包围的语句S,最多需要 $2k + 1$ 个维度【索引14,Loop Transformations: Convexity, Pruning and Optimization,2011年,ACM SIGPLAN Notices】来表达所有可能的调度变换(条带化挖掘不通过调度器表达),但实际上维度的数量没有上限。

调度器

  1. 调度问题形式化迭代调度器原理。PolyTOPS是一个迭代调度器:算法会逐步找到完整的调度变换 $\Theta$,通过构建一个整数线性规划(ILP)问题来寻找每个调度维度 $\phi_{S,i}$,从最外层维度开始直到最内层维度。算法会确保在找到足够的调度维度时终止。调度器的目标是为所有的S找到最优的系数向量 $T_{S,i}$。

  2. 有效性/合法性约束(Validity/Legality Constraint)Feautrier合法性约束。该约束由Feautrier【索引6,Some efficient solutions to the affine scheduling problem. I. One-dimensional time,1992年,International Journal of Parallel Programming】引入,是多面体调度器的核心,因为它约束变换向量 $T_{S,i}$ 的取值必须保持程序语义(确保调度的合法性)。对于每个依赖关系 $\delta_{S \to R}$,S必须在R之前执行:
    $(i\vec{t}, i\vec{t}') \in \delta_{S \to R} \Rightarrow \Theta_R(i\vec{t}') > \Theta_S(i\vec{t})$
    其中符号 $\succ$ 代表字典序上更大。
    迭代调度器的合法性定义。考虑到对于迭代调度器,每个维度 $\phi_{S,i}$ 是从外到内计算的,有效性的定义变为:
    $(\vec{it}, \vec{it'}) \in \delta_{S \to R} \Rightarrow \phi_{R,i}(\vec{it'}) \geq \phi_{S,i}(\vec{it})$
    直到依赖关系 $\delta_{S \to R}$ 被满足。这个蕴含关系可以使用法卡斯引理(Farkas Lemma)【索引15,Theory of Linear and Integer programming,1986年,A. Schrijver】进行线性化:约束随后仅在由向量 $T_{S,i}$ 和 $T_{R,i}$ 组成的变量空间中表达。

  3. 进展约束(Progression Constraint)作用与定义。在每次调度迭代中添加进展约束,以确保算法的推进。它的作用是确保调度为迭代空间定义了一个全序,并避免平凡的零解。该约束强制下一个调度解与先前在迭代空间中的解线性无关。
    具体实现。我们定义矩阵 $H_S$ 为先前调度维度解 $T_{S,i}$ 的(逐行)拼接。我们定义正交补 $H_S^{\perp}$ 如下:
    $H_S^\perp = I - H_S^T (H_S H_S^T)^{-1} H_S$
    其中 $I$ 是单位矩阵,$H_S^T$ 是 $H_S$ 的转置。
    考虑到我们将调度搜索空间限制在正象限内,进展约束定义为正交补矩阵的(逐行)和,如下所示:

    $$\forall i, H_{S,i}^{\perp} \cdot h_{S}^{*} \geq 0 \ \wedge \ \sum_{i} H_{S,i}^{\perp} \cdot h_{S}^{*} \geq 1$$


    其中 $H_{S,i}^{\perp}$ 是 $H_S^{\perp}$ 的一行,$h_S^*$ 是下一个要计算的解。

  4. 邻近性代价函数(Proximity Cost Function)目标与定义。邻近性代价函数由Bondhugula等人【索引7,A Practical Automatic Polyhedral Parallelizer and Locality Optimizer,2008年,PLDI'08】定义,旨在在合法解中找到那些优化时间局部性的解。
    核心思想。其思想是最小化对同一内存位置多次访问之间(在调度时间上)的距离。数据依赖描述了对同一内存位置的多次访问,因此邻近性目标是最小化依赖距离。对于一个依赖关系 $\delta_{S \to R}$,约束定义为:
    $(i\vec{t}, i\vec{t}') \in \delta_{S \rightarrow R} \Rightarrow \phi_{i,R}(i\vec{t}') - \phi_{i,S}(i\vec{t}) \leq u\vec{N} + w$
    其中 $\vec{u}$ 和 $w$ 是要最小化的代价函数。
    效果。邻近性精确地表示了一个有用的变换特性,并间接地有利于使最外层的维度成为并行维度,其依赖距离为0。

现有技术

  1. Feautrier调度器【索引6,Some efficient solutions to the affine scheduling problem. I. One-dimensional time,1992年,International Journal of Parallel Programming】是第一个迭代式多面体调度器。其目标是为单核SIMD CPU进行优化。它将有效性约束与Feautrier代价函数相结合,该代价函数旨在找到能够承载尽可能多依赖关系的顺序外部维度,从而可能为SIMD矢量化开发带来内层循环并行性。

  2. Pluto调度器【索引7,A Practical Automatic Polyhedral Parallelizer and Locality Optimizer,2008年,PLDI'08】是一个迭代式调度器,引入了前述的邻近性代价函数,旨在利用多核CPU等架构中的高度并行性。其更新版本Pluto+【索引16,The pluto+ algorithm: A practical approach for parallelization and locality optimization of affine loop nests,2016年,ACM Transactions on Programming Languages and Systems (TOPLAS)】扩展了功能以支持循环反转和负倾斜,并解决了Pluto无法解决的一些边界情况问题。Plutolp-dfp【索引17,Polyhedral Auto-Transformation with No Integer Linear Programming,2018年,SIGPLAN Not.】是其一个扩展,采用线性规划代替整数线性规划,将调度算法分解为一系列变换,显示了在编译时间方面的潜在优势。

  3. isl调度器【索引8,Isl: An Integer Set Library for the Polyhedral Model,2010年,ICMS'10】【索引9,Polyhedral Parallel Code Generation for CUDA,2013年,ACM Trans. Archit. Code Optim.】是一个迭代式调度器,它同时使用了对Pluto和Feautrier调度器的重新实现来最大化并行性。如果没有找到外部并行性,则应用Feautrier代价函数来消除尽可能多的依赖关系,并在后续维度中寻找并行性。

  4. Tensor调度器【索引10,Polyhedral Tensor Schedulers,2019年,HPCS】是一个应用于基于张量的应用(如AI)的迭代式调度器,这些应用通常具有高并行性和少量依赖。其重点是为缓存空间局部性定义了邻近性(Contiguity)代价函数,尝试寻找优化内存访问模式的循环排列。它取得了良好的结果,但仅限于特定领域,将调度变换限制为仅循环交换。

  5. One-shot调度器【索引14,Loop Transformations: Convexity, Pruning and Optimization,2011年,ACM SIGPLAN Notices】不是迭代式的,它是通过将整个多维变换 $\Theta(S)$ 表示为单个ILP问题来计算的。这种问题表述方式使得表示整个调度 $\Theta$ 上的全局约束和代价函数变得更容易,而迭代式调度器中的约束和代价函数通常是局限于单个调度维度 $\phi_i$ 的。然而,大量的变量和约束导致了可伸缩性问题和较长的编译时间。One-shot调度器的扩展【索引18,Model-Driven Transformations for Multi- and Many-Core CPUs,2019年,PLDI 2019】【索引19,Automatic Generation of MultiObjective Polyhedral Compiler Transformations,2020年,PACT '20】提出通过代价函数字典和先前找到的最优解的缓存机制来解决复杂性问题。

  6. 现有技术局限与新兴方法:当前研究表明,现有调度器都是为特定目标函数量身定制的,针对某些特定架构。它们的行为是预先确定的,可用的选项无法提供在不同领域、不同架构和不同编译器中取得良好结果所需的灵活性。例如,在AKG中,优化目标架构是多样的,因此任何现有调度器可能在某些场景下成功,但在其他场景下则不然,现有技术仍然缺乏可配置性和灵活性。

  7. PolyLingual【索引20,PolyLingual: a Programmable Polyhedral Scheduler,2023年】是一个正在开发中的领域特定语言(DSL),用于多面体调度器,它提供了构建块函数和类型,简化了多面体调度器的设计。
  8. Tiramisu【索引21,A Deep Learning Based Cost Model for Automatic Code Optimization,2021年,CoRR】提出直接定义代码的一组多面体变换,调度器被一个由专家给出的循环变换组合中的AI引导搜索策略所取代。
    • Clint【索引22,Clint: A direct manipulation tool for parallelizing compute-intensive program parts,2014年,IEEE Symposium on Visual Languages and Human-Centric Computing (VL/HCC)】通过图形界面允许直接对多面体进行手动变换。
  9. Chlore【索引23,Opening Polyhedral Compiler’s Black Box,2016年,CGO '16】解决了可解释性问题,给定输入内核及其转换版本,尝试恢复获得相同转换版本所需的多面体变换集。
  10. "Constraint Injection"【索引13,Optimizing GPU Deep Learning Operators with Polyhedral Scheduling Constraint Injection,2022年,CGO】提出了一种在多面体调度器中注入简单约束的方法,但它主要是为针对特定内核的优化而设计的。

A2 方法细节

PolyTOPS 是一款可配置的迭代式多面体调度器。其目标是提供一个灵活、易于调整的工具,以确保现有的策略或现有策略的混合(泛化 isl 的提议)能够以极少的努力被描述,同时在需要时仍允许专家更精确地指导调度器。为了实现这一点,PolyTOPS 提供了一个通用的迭代调度器方案,其中的策略可以通过一个配置文件来定义。

PolyTOPS 工作流程

核心组件。PolyTOPS 的工作流程如图 1 所示。其主要模块与 isl 【索引8,Isl: An Integer Set Library for the Polyhedral Model,2010年,ICMS'10】 或 Pluto 【索引7,A Practical Automatic Polyhedral Parallelizer and Locality Optimizer,2008年,PLDI'08】 等调度器相似。PolyTOPS 的输入和输出都是代码的多面体表示。输入的主要部分是初始调度和依赖关系,它们可以表示为 isl 对象或 OpenScop 格式。基于 ILP 的核心调度算法的结果可能会被进一步后处理:此阶段处理分块(tiling)、块内优化(intra-tile optimization)和用于波前并行(wavefront parallelism)的倾斜(skewing)(参见 【索引7,A Practical Automatic Polyhedral Parallelizer and Locality Optimizer,2008年,PLDI'08,5.3节】)。需要强调的是,核心调度器中没有实现块大小的决策。必须从外部提供块大小才能应用分块。最后,由于 PolyTOPS 可以输出 isl 对象或 OpenScop 表示,因此可以使用 isl 或 CLooG 【索引24,Code Generation in the Polyhedral Model Is Easier Than You Think,2004年,PACT '04】 等工具或库进行代码生成。

配置模块。PolyTOPS 的重要贡献在于其配置模块,该模块支持 JSON 和 C++ 两种接口。此功能允许指定高级优化策略。预定义或新的策略可以通过简单的关键字进行组合或扩展。此外,还支持进一步的定制,下至特定于内核的策略,如语句融合或部分调度规范。

图1:PolyTOPS工作流程表示,展示了主要模块,包括后处理和配置。我们支持Openscop和isl表示作为多面体表示。
图1:PolyTOPS工作流程表示,展示了主要模块,包括后处理和配置。我们支持Openscop和isl表示作为多面体表示。

配置特性

配置控制机制。PolyTOPS 的配置控制着调度器针对每个特定调度维度的行为。例如,我们可以为特定维度添加代价函数,在另一个维度中分发一些语句,或为再一个维度添加约束。可配置的特性可以分为两种主要类型:
* 局部配置(Local configurations):它们直接控制 ILP 的创建。可以选择预定义的代价函数,并指定它们的优先级顺序。可以定义新的变量、ILP 约束或代价函数。最后,可以为每个维度定义语句的分发/融合,这将在内部转化为强制执行分发规范的特定约束。
* 全局配置(Global Configurations):这些是更高级别的特性,不仅影响特定调度维度的 ILP 定义。这些特性需要调度算法中的多个逻辑步骤来满足。例如,指令(directives)是给调度器的建议,尝试对特定循环进行矢量化或并行化。另一个例子是自动矢量化(AutoVectorization),用于自动检测(基于内存步幅和访问模式)哪些循环应该被调度到最内层以进行可能的矢量化。

接下来将详细描述所有可能的配置,从局部配置开始,然后是全局配置。

A. 局部配置

  1. 代价函数控制配置方式。可以为每个调度维度定义或省略预定义或新代价函数的特定组合及其优先级(由文本顺序给出,从左到右)(如代码清单2第7行所示)。目标函数是变量的向量,变量的顺序很重要,因为它们是按字典序最小化的。
    代码清单2:展示PolyTOPS大部分可配置特性的JSON示例,包括代价函数控制与定义、约束定义、融合控制和指令。
    新变量与预定义代价函数。可以引入新变量(如代码清单2第3行所示),用于自定义约束定义和作为代价函数。预定义的代价函数包括来自Pluto【索引7,A Practical Automatic Polyhedral Parallelizer and Locality Optimizer,2008年,PLDI'08】的 proximity、来自Feautrier【索引6,Some efficient solutions to the affine scheduling problem. I. One-dimensional time,1992年,International Journal of Parallel Programming】的 feautrier、受Tensor调度器【索引10,Polyhedral Tensor Schedulers,2019年,HPCS】中代价函数启发的 contiguity,以及一个名为 bigLoopsFirst 的新颖简单函数,它尝试将迭代空间最大的循环调度在最外层。

  2. Contiguity 代价函数设计目标与定义contiguity 代价函数旨在按提供更好空间局部性的顺序调度迭代器。对于每个语句 S,给定迭代器集合 $\vec{it}_S$,我们将名为 contiguity 的代价函数定义如下:

    $$ \begin{aligned} \text{contiguity}(S) &= \sum_{i=0}^{|T_S^{it}|} T_{S,i}^{\vec{it}} \times c_{S,i} \\ \text{contiguity}(S) &\geq 0 \end{aligned} $$


    其中 $c_{S,i}$ 是一个描述优化内存访问模式优先顺序的支持系数。例如,关注代码清单1(左侧)的内核,两个支持系数向量 $\vec{c}_{S0}$ 和 $\vec{c}_{S1}$ 将分别为:

    $$\begin{array}{l} \vec{c}_{S0} = \left(\begin{array}{cc} 10 & 1 \end{array}\right) \\ \vec{c}_{S1} = \left(\begin{array}{cc} 1 & 10 \end{array}\right) \end{array} \text{ with } \vec{it}_{S0} = \vec{it}_{S1} = \begin{pmatrix} i \\ j \end{pmatrix}$$
    以强制调度器选择具有最小 contiguity 系数的最外层循环。

  3. bigLoopsFirst (BLF) 代价函数设计目标与定义bigLoopsFirst (BLF) 代价函数旨在将具有最大域的循环调度到最外层。其设计与 contiguity 代价函数相同,但公式5中使用的系数 $c_i$ 是基于优先考虑具有最高边界的维度。在内核具有大量并行性,而架构只能利用一个或几个层次的外部并行性的场景中,BLF可能很有用。我们尝试最大化并行迭代的次数。例如,关注代码清单1中的内核,两个向量 $\vec{c}{S0}$ 和 $\vec{c}$ 将分别为:

    $$ \begin{aligned} \vec{c}_{S0} &= \begin{pmatrix} 1 & 10 \end{pmatrix} \\ \vec{c}_{S1} &= \begin{pmatrix} 1 & 10 \end{pmatrix} \end{aligned} \text{ with } \vec{it}_{S0} = \vec{it}_{S1} = \begin{pmatrix} i \\ j \end{pmatrix}. $$
  4. 自定义约束接口与语法。通过一个简单的接口,可以定义自定义约束,这些约束是仿射不等式或等式。它们可以通过系数向量 $T_{S,i}$ 来约束由等式(1)定义的任何语句S和维度i的调度函数 $\phi_{S,i}(\vec{it})$。我们可以将这个向量分为子向量 $(T_{S,i}^{\vec{it}}, T_{S,i}^{\vec{N}}, T_{S,i}^{1})$。约束可以使用以下表示法涉及 $T_{S,i}$ 的任何系数:
    $S[stmt]_[var_type]_[idx_var],$
    其中 i 被隐式定义为当前考虑的维度(迭代调度器),并且:

    • stmt 是从 0 到 M-1 的语句编号(M是语句总数),遵循初始的文本顺序,它唯一地标识一个语句 S。
    • var_type 可以是以下关键字之一:it 指子向量 $T_{S,i}^{\vec{it}}$,par 指子向量 $T_{S,i}^{\vec{N}}$,cst 指常数项 $T_{S,i}^{1}$。
    • idx_var 是变量的索引。语句的最外层迭代器被认为是迭代器0。参数的顺序是输入程序中的文本顺序。
      此外,任何用户定义的变量都可以用于约束中。注意,用关键字 'i' 替换 stmtidx_var 代表该类型所有变量的总和。
      示例。例如,一个禁用语句3倾斜(skewing)的约束可以表示为:
      $S3_it_i \le 1.$
      这等价于:

      $$ \sum_{k} T_{\mathrm{S} 3, k}^{\vec{it}} \leq 1 . $$


      定义范围。自定义约束可以为特定的调度维度定义,也可以通过 fusion 关键字为所有维度定义(见代码清单2第10-15行)。
      约束形式。被接受的约束必须是仿射的。这意味着给定刚才描述的变量向量 $\vec{V}$,所有约束都可以定义为以下形式:

      $$ \text{constraint} = A\vec{V} + c \gtreqqless 0 $$
      其中A是整数系数矩阵,c是一个整数。
  5. 融合/分发控制功能。可以为特定内核指定自定义的循环融合决策。我们提供了控制融合的能力,为每个级别选择哪些语句要融合,哪些要分发。代码清单2的第16-22行显示了一个配置示例,指定在调度维度0,语句0和1要融合,语句2要分发。

B. 全局配置

  1. 指令(Directives)功能与用法。指令(Directives),如代码清单2第23-29行所示,用于指定某些循环应该是并行的、矢量化的(调度在最内层且不融合)或顺序的。这可以用来建议部分代码转换,而将剩余的调度转换决策留给调度器。调度器将尝试满足这些指令,除非无法保证调度的合法性。与合法性保留冲突的指令将被丢弃。

  2. 自动矢量化(Auto Vectorization)功能与原理。该选项指示调度器使用一种简单的启发式方法,为每个语句检测可以被矢量化的维度。该启发式方法寻找在内存中连续移动的维度。然后,调度器计算一个调度变换,其中:(1)可矢量化的维度被调度为最内层;(2)相应的语句在该调度维度上不进行融合。对于像CPU或NPU这样的架构,矢量化对性能至关重要。

C. 配置策略

接口类型。配置可以通过两种不同的接口来指定,每种接口更适合不同的配置场景:
1. JSON接口特点与局限。JSON接口,如代码清单2所示,允许为输入内核量身定制策略。局部配置是静态定义的,并映射到调度维度。配置指定了代价函数、额外约束和可能的循环分发。然而,这个接口没有提供定义复杂策略的自由度,这些策略需要考虑最外层的部分调度结果。
2. C++接口特点与优势。在这种配置中,策略在一个动态库中定义,该库由PolyTOPS加载并在每次调度迭代之前调用。这使得能够动态地指定每个调度策略,泛化了 isl 【索引8,Isl: An Integer Set Library for the Polyhedral Model,2010年,ICMS'10】 的策略,后者默认调用Pluto风格的调度器,并在失败时回退到Feautrier风格的调度器。这个例子在代码清单3中展示。此外,策略定义可以访问关于语句和到当前迭代为止计算出的部分调度的大量细节,这为创建更复杂的策略提供了机会。
代码清单3:一个isl风格策略的C++接口示例
表达能力。该配置具有足够的表现力,允许在不同策略之间切换,如Pluto风格、Feautrier风格、isl风格和TensorScheduler风格,并定义新的策略。唯一的限制是配置只能影响图1中的核心“调度器”模块。例如,主要的Pluto ILP策略可以很容易地通过配置复制,但其后处理和内部融合启发式则不能。

D. 通用算法结构

算法框架。PolyTOPS 依赖于一个通用的算法结构,如算法1所示,该结构与 Feautrier【索引6,Some efficient solutions to the affine scheduling problem. I. One-dimensional time,1992年,International Journal of Parallel Programming】、Pluto【索引7,A Practical Automatic Polyhedral Parallelizer and Locality Optimizer,2008年,PLDI'08】、isl-scheduler【索引8,Isl: An Integer Set Library for the Polyhedral Model,2010年,ICMS'10】和 Tensor Scheduler【索引10,Polyhedral Tensor Schedulers,2019年,HPCS】等迭代式调度器通用。这是 Pluto 算法的一个泛化,使用配置策略来驱动调度器。

算法步骤。算法的终止标准是检查迭代空间是否被完全覆盖以及所有依赖是否得到满足(第42行)。算法迭代寻找新的调度维度,直到满足终止标准(从第4行到第42行)。为了计算下一个维度 $\phi$,调度器首先验证融合启发式(或 PolyTOPS 的接口)是否为该调度维度强制要求循环分发(第8-14行)。如果不是,算法继续执行标准步骤(第16-21行),构建由代价函数和为尚未完全满足的依赖定义的约束组成的 ILP 系统。如果找不到解,算法尝试移除由前一调度维度满足的依赖,并继续构建 ILP 问题并尝试寻找解(第23-30行)。如果所有前面的步骤都失败了,则通过分析依赖图的强连通分量(SCC)并分发不同 SCC 的循环来强制执行循环分发(第32-36行)。一旦找到解,它会更新进展约束,确保下一个计算的 $\phi$ 维度与之前的维度线性无关,并且调度是一个双射变换。算法计算 BandsParallelDimensionBands 用于后处理分块,以确定哪些维度可以被分块。ParallelDimension 指示哪些调度维度是并行的。

与现有技术的联系。这个算法方案仅通过定义适当的配置就可以涵盖文献中所有的迭代式调度器。PolyTOPS 通过选择和定义代价函数及约束(第16、26行)的能力对其进行了扩展。JSON 接口是静态表达的,因此它在调度算法开始时被解析一次,而 C++ 接口允许使用到目前为止找到的调度 $\Theta$ 的信息进行基于逻辑的决策(第6行),因此它在每次调度迭代中都会更新。对于这两种情况,配置都会影响循环融合/分发的决策(第9行)。

算法1:PolyTOPS的通用算法结构
算法1:PolyTOPS的通用算法结构

合法性保证。无论提供何种配置,合法性约束(公式2)和进展约束(公式3)在计算解时总是被包含在内。这意味着调度器总能终止(证明与Pluto【索引7,A Practical Automatic Polyhedral Parallelizer and Locality Optimizer,2008年,PLDI'08】的证明类似)。此外,如果在配置中没有定义自定义约束和融合/分发控制,调度器保证能找到一个有效的调度。实际上,策略不会妨碍找到合法调度,而与合法性冲突的指令会被忽略。只有自定义约束和融合/分发可能导致空解。这与Tiramisu【索引21,A Deep Learning Based Cost Model for Automatic Code Optimization,2021年,CoRR】等不使用调度器的方法不同,因为在那些方法中,通过组合变换得到的每个调度函数都必须被证明是有效的。

A4 实验环境

  1. NPU 平台:
    • 硬件: Atlas 800 (型号 9010) 服务器,配备 8 个昇腾(Ascend)910 NPU 加速器【索引26,Chapter 3 - Hardware architecture,2020年,Ascend AI Processor Architecture and Programming】【索引27,Ascend: a Scalable and Unified Architecture for Ubiquitous Deep Neural Network Computing : Industry Track Paper,2021年,HPCA】。
    • 软件: MindSpore 框架【索引3,Deep Learning and Practice with MindSpore,2021年,Springer】,AKG (Automatic Kernel Generator)【索引2,AKG: Automatic Kernel Generation for Neural Processing Units Using Polyhedral Transformations,2021年,PLDI 2021】编译器。
  2. 数据集: MindSpore 的混合自定义算子(Hybrid Custom Operators)【索引25,4.1 Unified Cross-Platform MindSpore Hybrid DSL Expression,2022年】。

  3. 多核 CPU 平台:

  4. 硬件:
  5. AMD: AMD EPYC 7452,双路,共 32 核(每核 2 线程),256 MiB L3 缓存。
    * Intel1: Intel Xeon E5-2683 CPU (x86_64),双路,每路 16 核(每核 2 线程),80 MiB L3 缓存。
    * Intel2: Intel Xeon Silver 4215 CPU (x86_64),双路,每路 8 核(每核 2 线程),22 MiB L3 缓存。
  6. 软件:
  7. 代码生成器: CLooG【索引24,Code Generation in the Polyhedral Model Is Easier Than You Think,2004年,PACT '04】。
  8. 编译器: AMD 平台使用 gcc-11.3,Intel 平台使用 gcc-10.5。
  9. 数据集:
  10. Polybench: Polybench 基准测试套件【索引28,Polybench: The polyhedral benchmark suite,2012年】,包含来自线性代数、数据挖掘和模板计算等领域的异构内核。
  11. PolyMAGE: PolyMAGE 基准测试套件【索引30,PolyMage: Automatic Optimization for Image Processing Pipelines,2015年,ASPLOS '15】,包含 7 个来自图像处理的用例。

  12. 对比工具:

    • Pluto (最新开发版 commit eddc385),Pluto+【索引16,The pluto+ algorithm: A practical approach for parallelization and locality optimization of affine loop nests,2016年,ACM Transactions on Programming Languages and Systems (TOPLAS)】,Pluto-lp-dfp【索引17,Polyhedral Auto-Transformation with No Integer Linear Programming,2018年,SIGPLAN Not.】
    • isl-scheduler (AKG中默认版本)
  13. isl-PPCG【索引9,Polyhedral Parallel Code Generation for CUDA,2013年,ACM Trans. Archit. Code Optim.】

A4 实验结果

A. MindSpore 混合自定义算子 (NPU)

<center> (a) 输入代码,指令以红色显示。 </center>
def trsmLoffdiag(a, b):
    inverse0 = allocate(b.shape, b.dtype)
    row = b.shape[0]
    col = b.shape[1]
    for i in range(row):
        for j in range(i):
            for l in parallel(col // 16):
                for k in vectorize(16):
                    inverse0[i, l*16 + k] = a[i, j] * b[j, l*16 + k]
    b[i, l*16 + k] = b[i, l*16 + k] - inverse0[i, l*16 + k]
    return b
<center> (b) PolyTOPS优化后的代码 (分块前)。 </center>
def trsmLoffdiag(a, b):
    inverse0 = allocate(b.shape, b.dtype)
    row = b.shape[0]
    col = b.shape[1]
    for l in parallel(col // 16):
        for i in range(row):
            for j in range(i):
                for k in vectorize(16):
                    inverse0[i, l*16 + k] = a[i, j] * b[j, l*16 + k]
            for k in vectorize(16):
                b[i, l*16 + k] = b[i, l*16 + k] - inverse0[i, l*16 + k]
    return b
表I:(昇腾910 NPU)自定义算子结果,显示了每种情况下的周期数以及PolyTOPS相对于isl结果获得的加速比。
表I:(昇腾910 NPU)自定义算子结果,显示了每种情况下的周期数以及PolyTOPS相对于isl结果获得的加速比。

B. Polybench 上的调度策略比较 (CPU)

代码清单5:展示pluto-style(左)和tensor-scheduler-style(右)的JSON配置
图2:PolyTOPS(使用4种不同配置:pluto-style、tensor-scheduler-style、isl-style和kernel-specific)与Pluto相比的加速比(对数尺度)。kernel-specific配置至少与前三种一样好。结果按Intel2机器上的kernel-specific降序加速比排序。测试在AMD(上)、Intel1(中)和Intel2(下)上完成。
图3:PolyTOPS与Pluto在Jacobi-1d上使用两种不同配置和多种数据集大小的加速比。蓝色是为大尺寸设计的(最佳)专用配置。红色是pluto-style配置。

C. 与其他调度工具在 Polybench 上的比较

图4:PolyTOPS(使用内核特定配置)、Pluto-lp-dfp(最佳融合启发式)和Pluto+与Pluto(最新开发版)相比的加速比(对数尺度)。对于标有*符号的案例,Pluto-lp在[29]中提出的任何融合启发式都未找到解。
图4:PolyTOPS(使用内核特定配置)、Pluto-lp-dfp(最佳融合启发式)和Pluto+与Pluto(最新开发版)相比的加速比(对数尺度)。对于标有*符号的案例,Pluto-lp在[29]中提出的任何融合启发式都未找到解。

D. 在 PolyMAGE 上的调度工具比较

表II:PolyMage基准测试:PolyTOPS与现有技术调度器(isl-PPCG、Pluto、Pluto-lp-dfp、Pluto+)之间的时间(毫秒)和相对加速比。由于技术限制,某些情况和调度器的结果不可用(n.a.)。
表II:PolyMage基准测试:PolyTOPS与现有技术调度器(isl-PPCG、Pluto、Pluto-lp-dfp、Pluto+)之间的时间(毫秒)和相对加速比。由于技术限制,某些情况和调度器的结果不可用(n.a.)。

A5 结论

本文提出了一种新颖的多面体调度器工具 PolyTOPS,它通过提供一种简单的方式来配置和调整多面体调度,改进了现有技术的“黑盒”调度器。PolyTOPS能够适应各种应用场景,尤其是在那些因为黑盒调度器效果不佳而放弃多面体优化的领域。

核心贡献与成果
1. 高可配置性与灵活性:PolyTOPS 提供了易于使用的配置接口(JSON/C++),允许用户精细控制调度策略,可以模拟现有调度器行为(如isl, Pluto),也可以创建全新的、针对特定场景和内核的优化策略。
2. 广泛的适用性:输入和输出支持 isl 对象和 OpenScop 格式,使其可以与 isl 或 CLooG 等代码生成工具无缝集成。PolyTOPS 已成功集成到 MindSpore AKG 编译器中。
3. 显著的性能提升
* 在昇腾NPU的MindSpore混合自定义算子场景中,相比isl调度器,取得了高达34倍的加速比,几何平均加速比为7.66倍。
* 在Polybench基准测试中,通过简单的内核特定配置,在不同CPU上相比Pluto调度器取得了1.8倍的几何平均加速比。
* 在Polymage基准测试中,性能优于或持平于其他现有调度器。

未来工作展望
1. 增强融合启发式:融合启发式的设计对高性能至关重要,可以将其作为PolyTOPS配置的扩展,定义模式引导的融合启发式来丰富现有调度策略。
2. 软硬件特定配置:扩展更多针对特定软硬件的配置,例如:针对输入硬件配置自适应的融合和分块启发式,以及基于内核大小的调度决策。

A6 附录

A. 摘要

制品内容。本制品提供一个docker镜像,其中包含程序和脚本,用于生成图2、图3、图4和表II的结果。结果可能因目标架构或系统而异。
包含软件。该镜像包含PolyTOPS、Pluto、Pluto+、Pluto-lp和PPCG。此外还提供了如clan、Candl、Cloog、isl和FPL等附加软件。
制品目标。在本制品中,您将能够复现论文中展示的结果,并测试PolyTOPS及其功能。

B. 制品清单(元信息)

C. 描述

  1. 交付: 一个docker镜像可在Zenodo【索引32,PolyTOPS artifact,2023年,https://doi.org/10.5281/zenodo.10203989】找到 。
  2. 硬件依赖: 参见第四节-B,分别为Intel1、Intel2、AMD。
  3. 软件依赖: Docker v24。
  4. 数据集: Polybench、PolyMage。

D. 安装

加载镜像。该docker镜像发布在Zenodo【索引32】上,可以从polytops.tar文件加载,如下所示:
加载docker镜像命令
成功后,polytops:cgo-2024镜像将可用。

E. 实验工作流

运行容器。可以使用以下命令运行一个新的polytops:cgo-2024容器:
$ docker run -it --cap-add=SYS_NICE polytops:cgo-2024
环境配置。该镜像已设置为默认命令是/bin/bash --login。请注意,内部配置(参见/etc/profile.d)需要一个登录shell才能找到并执行附加软件。
运行完整实验。进入容器后,可以运行以下命令:
运行完整实验脚本命令
提升执行优先级。为了复现我们的结果,我们强烈建议用户将测试执行包装在以下命令中:
$ sudo --login nice -n -20 bash -c "{ cd $(pwd); <test>; }"
密码是polytops。此命令允许我们优先执行实验。例如,前面的命令将变为:
$ sudo --login nice -n -20 bash -c "{ cd $(pwd); bash ./run_complete_artifact.sh; }"
为清晰起见,后续命令将不再重写此部分。
脚本功能run_complete_artifact.sh脚本将运行所有脚本并生成时间输出,代表以下结果:
* PolyTOPS-results (图 2): $HOME/test/test_fig2_and_4/fig_2.csv
* Data-Size (图 3): $HOME/test/test_fig3/fig_3.csv
* SOTA (图 4): $HOME/test/test_fig2_and_4/fig_4.csv
* PolyMage (表 II): $HOME/test/test_polymage/times_polyma
平台特定配置。请注意,完整脚本是为Intel1配置的,而对于Intel2和AMD,您可以参考A-E1节,其中解释了如何仅运行图2中的特定测试。
单个实验执行。结果也可以为每个测试用例单独计算。
1. PolyTOPS-results: 要获得图2中描述的结果,用户可以运行以下命令:
$ cd $HOME/test/test_fig2_and_4/
$ bash test_fig2.sh -c <config_path> -n 10
其中 -c 选项指定了我们用于PolyBench案例的PolyTOPS配置文件的根路径,-n 指定了我们希望为每个最终转换运行的执行次数。对于 -c 选项,您可以根据要复现图2的哪个实验选择以下任一路径:
* 图2 Intel1机器: $HOME/test/paper_best_configs_INTEL1/
* 图2 Intel2机器: $HOME/test/paper_best_configs_INTEL2/
* 图2 AMD机器: $HOME/test/paper_best_configs_AMD/

  1. Data-Size: 要复现这些实验,只需运行以下命令:
    运行Data-Size实验脚本命令
    其中 -n 是一个指定每个程序版本执行次数的选项。输出会自动生成在$HOME/test/test_fig3/fig_3.csv

  2. SOTA: 要复现图4,您可以运行以下命令:
    $ cd $HOME/test/test_fig2_and_4/
    $ bash test_fig4.sh -n 10
    其中 -n 指定每个程序版本的执行次数。请注意,这部分实验的大部分时间都由isl-PPCG结果占用(参见图4中的巨大减速)。如果用户希望将isl从实验中排除,可以编辑$HOME/test/test_fig2_and_4/test_fig4.sh脚本的第32行,删除isl关键字。

  3. PolyMage: 要复现表II中的PolyMage实验,您可以运行以下命令:
    运行PolyMage实验脚本命令
    输出将位于文件$HOME/test/test_polymage/times_polymage.csv中。

F. 评估和预期结果

输出文件格式。生成的结果(PolyTOPS-results, Data-Size, SOTA, 和 PolyMage)是CSV文件,包含不同版本测试用例的平均执行时间和时间的标准差。
结果计算。论文中的结果对于PolyMage(表II)是等效的,而对于其他图表,我们使用以下公式计算了(与Pluto结果相比)以2为底的对数尺度加速比:
speedup = pluto_time / variant_time
其中 variant_time 代表图2、图4和图3中图表里的任何PolyTOPS变体或任何其他调度器的时间。

G. 实验定制

直接使用PolyTOPS。如果您想在自定义案例中使用我们的工具,可以直接使用polytops命令。PolyTOPS支持OpenScop作为输入(由Clan生成),最终的代码生成由Cloog完成。给定一个输入C文件input.c(必须包含适当的PRAGMA),生成优化版本out.c的一个简单流程是:
$ clan input.c | polytops --input-format=openscop --tiling=true --compute-dependencies=true --output-format=openscop | cloog stdin -openscop -o ./out.c
可用选项--help选项提供了所有可用选项的列表。
辅助脚本。此外,我们还在$HOME/test/scripts/single_case.sh中提供了另一个脚本,可用于运行类似的流程,但带有一些额外的PolyTOPS选项。此脚本包含一些额外的功能,可以通过--help选项显示。