[论文解读] On the Representation of Partially Specified Implementations and its Application to the Optimization of Linear Algebra Kernels on GPU

Ulysse Beaugnon, Basile Clément+2 more|ArXiv.org|Apr 6, 2019

Parallel Computing and Optimization Techniques参考文献 33被引用 4

一句话总结

本文提出了一种约束编程框架，将GPU上的线性代数核优化表示为部分指定的实现形式，其中优化决策（例如，循环融合、向量化、内存分块）被建模为约束满足问题（CSP）中的相互依赖变量。通过支持双向依赖解析和分析性性能边界，该方法能够在复杂优化空间中实现高效、与顺序无关的搜索，从而实现与手工调优的GPU库相媲美甚至更优的性能。

ABSTRACT

Traditional optimizing compilers rely on rewrite rules to iteratively apply program transformations. This iterative approach hides optimization opportunities behind intermediate transformation steps. For instance, vectorization can only be applied to the innermost loop in a nest: one must first perform a loop interchange before even considering vectorization of an outer loop. In contrast, we propose an implementation framework representing programs as sets of possible implementation decisions. Specifying one decision can have an impact on others in a bidirectional manner: specifying that a loop must be vectorized prevents other loops from being nested inside it; conversely, specifying a loop as an outer loop will prevent it from being vectorized. These optimization decisions commute, obviating the pass ordering problem. We present a constraint programming system to formally define, represent and explore such implementation spaces. We also propose an exploration strategy combining tree search and branch-and-bound; the strength and novelty of this strategy reside in an analytical model of the lower bound on the execution time of a set of possible implementations. We showcase our approach on the construction and exploration of an implementation space for linear algebra kernels running on GPUs. We show this search space is expressive enough to represent complex decisions that fundamentally change the structure of the generated code. We also present preliminary results competitive with the performance of native GPU libraries.

研究动机与目标

• 为解决GPU内核生成中复杂且相互依赖的优化决策挑战，传统基于重写规则的编译器存在传递顺序问题和有限的前瞻能力。
• 将优化决策（如循环变换、数据分块、并行映射）建模为约束满足问题（CSP）中的相互依赖变量，以实现系统性探索。
• 通过执行时间的分析性下界实现早期性能估计与剪枝，即使在完整实现尚未指定时也可行。
• 开发领域特定语言和代码生成流水线，以高效操作GPU线性代数内核上下文中的部分指定实现。
• 证明该方法能够表达并探索复杂的、改变代码结构的优化（例如，分块、循环融合、向量化），其性能可与手工优化的GPU库相匹配或超越。

提出的方法

• 将程序表示为部分实例化的决策向量，其中每个决策（如向量化、循环交换）作为CSP中的变量，具有定义的取值域和约束条件。
• 将决策之间的交互建模为双向约束：例如，对一个循环进行向量化会阻止其被嵌套在另一个循环内部，而将一个循环设为外层循环则会阻止其向量化。
• 采用混合搜索策略，结合树搜索与分支定界法，并利用分析性性能下界模型引导搜索，提前剪除无前景的分支。
• 使用支持增量求值和启发式驱动探索实现候选的约束编程系统来编码优化空间。
• 从优化选择及其相互作用的高层描述中，生成用于决策向量操作的高效底层代码。
• 通过将与内存层次管理、线程映射和循环变换相关的决策编码为CSP的一部分，将该框架应用于GPU线性代数内核。

实验结果

研究问题

• RQ1 基于约束的部分指定实现表示能否克服传统基于重写规则的编译器固有的传递顺序问题？
• RQ2 在未进行完整代码生成的情况下，分析性性能边界在多大程度上能有效引导大型复杂优化空间的搜索？
• RQ3 该框架能否表达并有效探索GPU内核中的复杂、结构改变型优化（如循环融合、分块、多级并行）？
• RQ4 自动生成的内核性能与手工优化的原生GPU库相比如何？
• RQ5 该框架能否通过允许专家优先做出高影响力决策，实现灵活的、与变换顺序无关的引导式探索？

主要发现

• 所提出的框架成功将GPU内核生成中的复杂、相互依赖的优化决策建模为CSP，支持与顺序无关的探索，避免了迭代变换流水线的缺陷。
• 从子空间推导出的分析性性能边界可在完整实现指定之前有效剪枝搜索空间，显著提升搜索效率。
• 该方法支持对关键性能优化的代码结构根本性改变，包括循环融合、分块、向量化和多级并行，这些对GPU高性能至关重要。
• 初步结果表明，自动生成的线性代数内核性能可与手工调优的原生GPU库相媲美或更优。
• 该框架支持早期且精确的收益估计，使全局启发式策略能基于部分实现引导搜索，这在传统ILP或SAT求解器中难以实现。
• 领域特定语言和代码生成流水线促进了对多样化优化决策的快速原型设计与实验，展示了良好的可扩展性和表达能力。