Progressive Raising in Multi-level IR
Progressive Raising in Multi-level IR
作者/机构: Lorenzo Chelini (TU Eindhoven), Andi Drebes (Inria and Ecole Normale Supérieure), Oleksandr Zinenko (Google), Albert Cohen (Google), Nicolas Vasilache (Google), Tobias Grosser (University of Edinburgh), Henk Corporaal (TU Eindhoven)
A1 主要贡献
本文针对多级中间表示(IR)编译器在处理通用语言时面临的抽象层次失配问题,提出了一个补充现有“渐进式降低”(progressive lowering)流程的新方法。
核心问题: 多级IR框架(如MLIR)通过渐进式降低,将高层抽象逐步转换到低层IR以进行代码生成,这极大地依赖于源代码的初始高级表示。然而,通用编程语言(如C/C++)由于语义不够丰富,其代码进入编译流程时处于较低的抽象层次,导致大多数领域特定的高级优化无法应用。手动将代码提升到框架内部所需的抽象级别对用户来说既不方便,也要求用户深入了解工具内部。
研究目标: 本文旨在提出一种方法,在多级IR中实现从低层抽象到高层抽象的提升(raising),从而使通用语言编写的低级代码也能利用现有的领域特定优化。
创新点与主要贡献:
1. 提出“渐进式提升”(Progressive Raising)概念:作为多级IR中“渐进式降低”的补充技术,该方法能够将低级IR表示逐步提升到更高级的抽象层次。
2. 开发 Multi-Level Tactics 技术:这是一种用于实现渐进式提升的声明式方法。它包含一种高级声明性语言(TDL),用于定义提升转换规则,这些规则独立于其目标循环、迭代器、容器和索引抽象。
3. 在MLIR框架中实现:将 Multi-Level Tactics 集成到MLIR编译器基础设施中,并实现了一系列转换,可将基于循环的MLIR方言(如Affine)提升到线性代数方言(Linalg)。
4. 展示优化效果:通过将仿射循环嵌套提升到高级线性代数操作,并随后对矩阵链式乘法进行优化,展示了跨越多个抽象级别的渐进式提升过程,并通过实验评估证明了其带来的显著性能提升。
下图形象地展示了本文的核心思想:从代表低抽象层次的山谷开始,通过Multi-Level Tactics将代码表示提升到山腰的中间抽象层次,甚至山峰的最高抽象层次。之后,可以沿着不同的优化路径(不同的斜坡)通过渐进式降低到达山谷,生成最终的底层代码。
A3 背景知识/关键Observation/设计原则
II. MLIR 基础设施
MLIR框架简介
MLIR编译器基础设施是LLVM下的一个项目,非常适合多级IR重写【索引7,MLIR: Scaling Compiler Infrastructure for Domain Specific Computation, CGO'21】。MLIR提供了一种非强制性(non-opinionated)的中间表示(IR),只有少数概念是内置的,大部分IR都是可定制的。这种可定制性允许编译器开发者通过引入自定义的类型、操作和属性来匹配手头问题的正确抽象级别。操作(Operation)是IR的基本原子构成部分,每个操作使用并产生新的值(Value)。值代表运行时的数据,并与编译时已知的类型(Type)相关联,而类型则为值建立了编译时信息的模型。作为补充,属性(Attribute)将编译时信息附加到操作上。自定义的类型、操作和属性在逻辑上被分组为方言(Dialect)。方言是使MLIR基础设施能够实现可复用抽象栈的基本要素。每种抽象都直接在其IR中编码和保留了转换有效性的前提条件,从而降低了分析过程的复杂性和成本。
MLIR中的方言与抽象层次
图2展示了MLIR中可用的一些方言及其入口点。Linalg方言用于对张量或缓冲区操作数上的线性代数运算进行建模。在相似的抽象级别,Stencil方言表示根据给定模板模式更新数组元素的迭代核【索引8,Domain-Specific Multi-Level IR Rewriting for GPU, arXiv preprint arXiv:2005.13014, 2020】。在较低的抽象级别,Affine方言模型化了一个简化的多面体表示,而F18则模型化了FORTRAN特定的结构(例如,调度表)。SCF和Standard分别代表结构化控制流和一组杂项操作。最后,MLIR LLVM方言模型化了LLVM-IR的结构。
通过TableGen进行声明式定制
IR的定制是通过一个声明式系统实现的,该系统主要基于TableGen。TableGen是一个用于定义和维护领域特定信息记录的数据建模工具,在整个LLVM代码库中被广泛使用【索引9,Tablegen overview, https://llvm.org/docs/TableGen/】。例如,MLIR使用基于TableGen构建的操作描述规范语言(ODS)来声明式地描述操作。为新操作提供声明式规范可以加快新自定义IR的开发速度,减少代码重复量以及出错的可能性。TableGen文件本身只是领域特定信息的“容器”,没有后端就没有任何意义。在编译时,由后端来解释存储的信息并生成相应的C++声明和定义 。
渐进式降低
渐进式降低(图2中的向下箭头)能够将操作从高级抽象方言降低到低级IR。列表1展示了一个渐进式降低的例子,其中linalg.matmul操作被转换为Affine方言中的一个包含乘法和加法的三重嵌套仿射循环。后续的降低过程会将循环体内的操作融合为一个乘加累积操作(MAC)。高效的渐进式降低是通过一个模式重写基础设施实现的【索引7,MLIR: Scaling Compiler Infrastructure for Domain Specific Computation, CGO'21】。
%0 = linalg.matmul( %A, %B, %C) // linalg方言
y
affine.for %i = 0 to %N // affine方言
affine.for %j = 0 to %N
affine.for %k = 0 to %N {
%0 = affine.load %C[%i, %j] : memref<?x?xf32>
%1 = affine.load %A[%i, %k] : memref<?x?xf32>
%2 = affine.load %B[%k, %j] : memref<?x?xf32>
%3 = std.mulf %1, %2
%4 = std.addf %3, %0
affine.store %4, %C[%i, %j] : memref<?x?xf32>
}
列表1:MLIR中的渐进式降低。一个linalg matmul操作被翻译成一个嵌套循环,其中包含从二维memref加载数据、一次加法和一次乘法,这些都在Affine方言中表示。
A2 方法细节
III. MULTI-LEVEL TACTICS 概述
核心思想
Multi-Level Tactics的核心思想是允许从较低的抽象级别转换到较高的抽象级别。通过提供将通用语言的常见抽象提升到专门化高级方言的基础设施,Multi-Level Tactics使得在多级IR编译器中能够对通用代码进行领域特定的优化。
整体流程
图3展示了使用Multi-Level Tactics进行转换的整体流程。顶部的步骤(橙色框)允许进行策略(tactic)规范和生成操作实际IR的代码。而底部的步骤(蓝色框)则表示对输入程序的转换。
策略描述与代码生成
所有的策略都在一个高级的、声明式的策略描述语言(Tactics Description Language, TDL)中描述。TDL能够指定策略所应用的计算模式和一组替换表达式。TDL DSL前端处理高级声明式规范,并生成一个名为策略描述规范(Tactics Description Specification, TDS)的TableGen条目。第III-A和III-B节将详细描述这两种格式。用于匹配IR操作和内存访问以及构建替换操作的实际代码,是由Multi-Level Tactics后端生成的,并在执行时使用MLIR模式重写基础设施(更多细节见第III-C节)。
输入程序的转换流程
对输入程序的转换始于使用MLIR提取工具(MET)将C代码翻译为Affine方言。MET使得属于多面体模型的C语言子集能够从Affine级别进入MLIR流水线。在翻译过程中,C代码通过循环分发(loop distribution)进行规范化,以简化模式识别。然后,Affine表示被送入MLIR模式重写引擎,其中已经挂载了生成的策略。整个流程的输出是一个提升后的MLIR,其中的循环嵌套已被提升为高级操作。
A. 策略描述语言 - TDL
TDL简介
Multi-Level Tactics中关键的面向用户的组件是策略描述语言(TDL)。TDL中的模式和替换的语法借鉴自Tensor Comprehensions,后者本身是普遍使用的爱因斯坦张量索引表示法的一个轻微变体【索引10,Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions, arXiv preprint arXiv:1802.04730, 2018】。作为一个具体示例,我们展示了用于提升一个形式为abc-acd-bd的收缩(contraction)的顺序循环实现的模式和替换语法(如列表2所示)。我们通过将其重写为等效的Transpose-Transpose-GEMM-Transpose (TTGT)表达式来提升该收缩,利用了供应商优化库提供的高效GEMM实现。TTGT计算首先通过显式的张量转置和重塑操作将张量展平为矩阵,然后执行GEMM,最后将结果矩阵折叠回原始的张量布局。
for (int a = 0; a < N; ++a)
for (int b = 0; b < M; ++b)
for (int c = 0; c < N; ++c)
for (int d = 0; d < O; ++d)
S1: C[a][b][c] += A[a][c][d] * B[d][b];
}
列表2: 收缩 abc-acd-db
TDL示例
列表3展示了用户定义的用于检测和优化该收缩的策略。该策略的签名由关键字def后跟一个名称(TTGT)组成。主体包含由关键字pattern和builder分隔的两部分:pattern描述了要在用户代码中检测的计算主题,而builder描述了转换的配方。例如,第5和第6行将张量C和A重塑为矩阵。此外,在重塑之前对C执行了一次转置(a, b, c) → (a, c, b)。第7行指定了GEMM操作,而第8行将结果折叠回原始张量布局,在折叠后再次发出一个隐式的转置(a, c, b) → (a, b, c)。
1 def TTGT {
2 pattern
3 C(a,b,c) += A(a,c,d) * B(d,b)
4 builder
5 D(f,b) = C(a,b,c) where f = a * c
6 E(f,d) = A(a,c,d) where f = a * c
7 D(f,b) += E(f,d) * B(d,b)
8 C(a,b,c) = D(f,b) where f = a * c
9 }
列表3:用于匹配收缩abc-acd-db并在每个检测到的模式上应用TTGT优化的TDL。
B. 策略描述规范 (TDS)
TDS的角色
TDL DSL前端的角色是生成基于TableGen的策略描述规范(TDS)。每个TDS文件由一系列TableGen模板的实例化组成,并在编译时从中生成C++代码。选择两步流程而不是直接从TDL生成代码,允许将模式匹配和基础设施的通用例程作为可复用的模板在TDS中进行分解,从而降低复杂性。
TDS示例
列表4展示了为列表3中的TTGT策略生成的TDS定义。每个TDS条目都派生自一个基类Tactic,该类允许定义模式和一系列构建器(builder)。模式规范与TDL策略中的pattern条目相匹配。而构建器列表则对应一组高级操作,与高级方言(如Linalg)暴露的操作有一对一的映射关系。例如,transposeBuilder将创建一个Linalg TransposeOp或一个对BLAS转置例程的函数调用,具体取决于用户选择的降低路径。回到我们的运行示例,列表3中的第5行将映射为一个转置操作后跟一个重塑操作(列表4中的第2和3行)。类似地,第6和7行将分别映射为一个重塑操作和一个matmul操作。最后,列表3中的第8行将映射为一个重塑操作后跟一个转置操作(列表4中的第6和7行)。最终,TableGen中的每个条目都会在编译时通过Multi-Level Tactic的TableGen后端降低为C++的匹配器和构建器,我们将在下一节看到。
1 def TTGT : Tactic<C(a, b, c) += A(a, c, d) * B(d, b), [
2 transposeBuilder<In<[C]>, Out<[D]>, Expr<{0, 2, 1}>>,
3 reshapeBuilder<In<[D]>, Out<[E]>, Expr<{{0, 1}, 2}>>,
4 reshapeBuilder<In<[A]>, Out<[F]>, Expr<{{0, 1}, 2}>>,
5 matmulBuilder<In<[F, B]>, Out<[E]>>,
6 reshapeBuilder<In<[E]>, Out<[D]>, Expr<{{0, 1}, 2}>>,
7 transposeBuilder<In<[D]>, Out<[C]>, Expr<{0, 2, 1}>>,
8 ]>;
列表4:用于匹配收缩abc-acd-db并在每个检测到的模式上应用TTGT优化的TDS。
C. 匹配器和构建器
代码生成组件
从每个TDS条目生成的代码包含四个部分:结构匹配器(Structural matchers)、操作匹配器(operation matchers)、访问匹配器(access matchers)和构建器(builders)。
结构和操作匹配器
结构匹配器的作用是检测IR中的控制流模式。它本质上是复制了IR中基于控制流的结构,并增加了额外的过滤能力。一个结构匹配器由一个控制流操作类型、一个子操作列表和一个可选的过滤函数组成。例如,列表5匹配了所有由二维完美嵌套循环组成的IR子树,其中最内层循环体包含一个MAC(乘加)操作,这通过回调函数isMAC进行验证。顶层操作,被称为相对根(relative root),定义了一个结构匹配器。匹配操作从IR中指定的某个操作开始,然后递归地遍历该操作的后代以及相对根匹配器的后代。如果在遍历过程中遇到不匹配,过程将停止并立即报告失败。构建结构匹配器的API在视觉上被设计成与IR本身的结构相似。前导参数可以包括一个可选的回调函数,允许调用者更精确地控制匹配。例如,识别循环体中的MAC操作需要检查一个非结构性属性。每个匹配器必须属于一个处理内存分配和所有权的上下文。另一方面,操作匹配器用于验证算术操作的类型。我们依赖m_Op匹配器,它携带了要匹配的操作类型。m_Op可以被链接起来以检测操作序列。在我们的运行示例中,MACOp查找一个带有两个参数的Add操作,其中第二个参数是一个Mul操作。每个操作的参数都会被捕获以供后续检查。
auto isMAC = [&a, &b, &c](Body loop) {
auto MACOp = m_Op<AddOp>(a, m_Op<MulOp>(b, c));
return MACOp.match(loop);
}
/* instantiate the context */
For( For(isMAC) // filtering function
));
列表5:一个结构匹配器,用于检测二维完美嵌套循环,其中循环体是一个MAC操作。
访问匹配器
作为结构匹配器的补充,Multi-Level Tactics提供了访问模式匹配器。访问匹配器依赖于“占位符”(placeholders)的思想,模型化为m_Placeholder和m_ArrayPlaceholder。前者可以匹配任何形式为k * ι + c的归纳维度,其中k和c是构成模式的系数,而ι通过匹配代表归纳变量的底层mlir::Value来定义一个候选者。与此不同,m_ArrayPlaceholder只能匹配张量访问,并接受一个m_Placeholder列表作为输入。在这两种情况下,一次匹配就是将一个候选者赋给一个占位符。分配给不同占位符的候选者必须是不同的,而在一个匹配器表达式中对同一占位符的多次引用必须指向同一个候选者。占位符可以组合在占位符表达式中(例如 C{_i, _j}),并可用于m_Op匹配器,该匹配器允许指定特定类型的操作(例如StoreOp或LoadOp)。匹配过程从检查一个MLIR块(一个没有控制流的有序操作列表)内的最后一条store指令开始,然后沿着use-def链向后遍历。如果在向后遍历期间,一个或多个占位符找不到候选者,则立即报告不匹配,过程停止。其编程接口在精神上与结构匹配器相似,提供了一种声明式的方式来指定访问模式。同样,每个占位符必须属于一个上下文,该上下文负责协调匹配过程、处理内存分配和所有权。列表6展示了调用者如何识别从一个二维数组的加载操作。
/* instantiate the context */
auto _i = m_Placeholder(), _j = m_Placeholder();
auto _A = m_ArrayPlaceholder();
auto matcher = m_Op<LoadOp>(_A({2*_i+1, _j+5}));
列表6:一个访问匹配器,用于检测从二维数组的加载操作。
构建器
转换中匹配模式的替换是由构建器生成的,它们通过直接实例化IR操作(例如,在生成对供应商优化库的调用时),或使用MLIR生态系统中已有的基础设施(例如,EDSC构建器【索引11,Embedded domain specific constructs - EDSC, https://mlir.llvm.org/docs/EDSC/】)来完成 。
综合示例
回到我们的运行示例,列表7的第1行显示了为我们的收缩操作生成的结构匹配器。这些匹配器会匹配所有四维循环嵌套,对于这些循环嵌套,调用access_callback函数的评估结果为真。access_callback函数会调用第5到30行生成的访问匹配器。具体来说,访问匹配器会寻找一个MLIR块,该块恰好包含对不同张量的三次读取访问、一次写入访问、一个满足占位符模式[a, b, c] → [a, c, d][d, b]的索引排列,以及定义了收缩计算的两个算术操作(+/*)。
1 For(For(For(For(access_callback()))));
2
3 auto access_callback = [&a](Body loop) {
4 {
5 AccessPatternContext pctx(/* MLIR ctx */);
6
7 auto _a = m_Placeholder();
8 auto _b = m_Placeholder();
9 auto _c = m_Placeholder();
10 auto _d = m_Placeholder();
11
12 auto _C = m_ArrayPlaceholder();
13 auto _A = m_ArrayPlaceholder();
14 auto _B = m_ArrayPlaceholder();
15
16 auto var0 = m_Op<AffineStoreOp>(_C({_a, _b, _c}));
17 /* check the store is the last instruction in
18 the block */
19
20 auto var1 = m_Op<AffineLoadOp>(_C({_a, _b, _c}));
21 auto var2 = m_Op<AffineLoadOp>(_A({_a, _c, _d}));
22 auto var3 = m_Op<AffineLoadOp>(_B({_d, _b}));
23 auto body = m_Op<AddOp>(var1, m_Op<MulOp>(var2, var3));
24 /* match the body starting from the store op
25 and make sure we have only the defined
26 operations in the block */
27
28 a = pctx[_a] // read out the matched value
29 ..
30 }
31 };
列表7:Multi-Level Tactic后端为列表3中定义的策略生成的结构和访问匹配器。
IV. MULTI-LEVEL TACTIC 语法
TDL与TDS语法
图4和图5分别展示了策略描述语言(TDL)和策略描述规范(TDS)的语法。TDL通过pattern和builder关键字扩展了TC语法【索引2,The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated gpu kernels, automatically, ACM Trans. Archit. Code Optim., vol. 16, Oct. 2019】。图4展示了TDL核心语法的简化版本。
id ::= [C identifier]
binOp ::= ’+’ | ’-’ | ’*’ | ...
idList ::= [comma separated id list]
stmt ::= id ( idList ) ‘=’ id ( idList ) { binOp id ( idList ) }
stmtList ::= [whitespace separated stmt list]
pattern ::= stmt
builder ::= stmtList
图4:策略描述语言的简化EBNF语法。方括号表示可选子句,花括号表示重复,中括号内为简化描述。
图5展示了TDS的语法,其中每个条目都派生自Tactic类,该类允许使用TC语法指定模式和一系列构建器。TDS支持五种不同的构建器:reshape、transpose、matmul、matvec和convolution,每种都映射到一个高级操作(例如linalg.matmul)。除了transpose和reshape只处理单个输入外,所有构建器都支持多个输入。所有构建器都产生单个输出。transpose和reshape构建器需要一个仿射表达式作为第三个参数,用以指定维度应如何转置或重塑。
tactic ::= pattern { builder }
pattern ::= TC-expression
builderId ::= reshapeBuilder | transposeBuilder | matmulBuilder matvecBuilder | convBuilder
input ::= whitespace separated list of string
output ::= string
affineExpr ::= string
builder ::= builderId(input, output, [affineExpr])
图5:策略描述规范的简化EBNF语法。花括号表示重复,尖括号内为文本描述。
匹配器语法
图6使用扩展的巴科斯-瑙尔范式(EBNF)展示了访问和结构匹配器的语法。每个访问匹配器(占位符和数组占位符)都必须属于一个上下文。该上下文通过跟踪mlir::Value到占位符的分配来协调匹配过程。如果上下文尚未实例化,则无法构造匹配器——当上下文超出作用域时,所有内容都会被释放。占位符是调用m_Placeholder的结果。占位符支持运算符重载,以匹配任何形式为k*ι+c的仿射表达式,其中k和c是来自模式的系数,ι定义了候选者。占位符列表是占位符的有序集合。占位符在列表中的位置由其在参数列表中的位置隐式推断,并在匹配过程中加以考虑。数组占位符可以通过调用m_ArrayPlaceholder来构造,并可以为其分配一个占位符列表。类似地,只有在上下文实例化之后才能构造结构匹配器。结构匹配器可以接受一个回调函数作为可选的前导参数。回调函数允许通过测试例如访问模式属性等,对匹配进行更细粒度的控制。
AccessPatternContext ::= /*res of calling AccessCtx()*/
StructuralPatternContext ::= /*res of calling NestedPatternCtx()*/
placeholder ::= /*res of calling m_Placeholder()*/
placeholder-list ::= placeholder[{’,’ placeholder}]
array-placeholder ::= /*res of calling m_ArrayPlaceholder()*/
array-placeholder ::= array-placeholder, placeholder-list
load-matcher ::= m_Op<LoadOp>(placeholder-list)
load-matcher ::= m_Op<LoadOp>(array-placeholder)
store-matcher ::= m_Op<StoreOp>(placeholder-list)
store-matcher ::= m_Op<StoreOp>(array-placeholder)
void ::= m_Capt(Value &i)
StructuralMatcher ::= StructuralMatcher-type(StructuralMatcher-list)
StructuralMatcher-type([callback],StructuralMatcher-list)
node-type ::= ’FOR’ | ’IF’
图6:访问和结构匹配器的EBNF语法。
A4 实验环境与结果
实验环境
硬件配置:
实验在两个平台上进行,详细配置见表 I。
* 平台1: Intel Core i9-9900K CPU (Coffee Lake架构, 3.6 GHz), 64 GB RAM, L1/L2/L3缓存分别为 32/256/16384 KB。
* 平台2: AMD Threadripper 2920X CPU (4.3 GHz), 64 GB RAM, L1/L2/L3缓存分别为 1.125/6/32 MB。
软件配置:
* 操作系统: Ubuntu 18.04
* 编译器:
* Clang (release 6.0.0) with -O3作为基准。
* Pluto 源码到源码编译器,用于与多面体优化器比较。
* 本文提出的 Multi-Level Tactics (MLT) 框架,基于 MLIR 实现。
* 依赖库:
* Intel Math Kernel Library for Deep Neural Network (MKL-DNN) 用于 BLAS 调用。
* 测试基准:
* GEMM 可靠性测试: Polybench 4.2 和 Darknet 中的 GEMM 核,以不同风格的 C 代码编写。
* Affine 到 Linalg 提升测试: Polybench 4.2 中的一系列线性代数基准,以及来自耦合簇方法和化学核研究的张量收缩。
* 渐进式提升测试: 文献中的三个矩阵链式乘法问题【索引22,Performance evaluation of a parallel dynamic programming algorithm for solving the matrix chain product problem, 2014 IEEE/ACS 11th International Conference on Computer Systems and Applications (AICCSA), 2014】。
图7展示了本文评估中涉及的三种提升路径,它们扩展了标准的MLIR降低流程。
实验结果
A. 从Affine循环嵌套到Affine高级操作的提升
实验内容: 此实验旨在验证Multi-Level Tactics能否将通用的GEMM(广义矩阵乘法)循环嵌套提升为MLIR Affine方言中一个专门的高性能matmul操作。该操作内部实现了类似于OpenBLAS/Blis的优化【索引14,Analytical modeling is enough for high-performance blis, ACM Trans. Math. Softw., vol. 43, Aug. 2016】。使用的策略如列表8所示。
实验结果与分析:
* 检测可靠性 (图8): Multi-Level Tactics成功检测到Polybench中所有符合条件的GEMM调用点。然而,它未能检测到Darknet中的GEMM核,因为Darknet使用了线性化的一维数组访问,而当前策略仅匹配多维数组引用。这表明需要一个反线性化(delinearization)过程来处理更广泛的代码风格。
* 性能提升: 在AMD系统上对一个2088x2048的单精度GEMM进行测试,clang -O3编译的原始循环嵌套性能为1.76 GFLOPS/s。通过Multi-Level Tactics提升到matmul操作并进行优化后,性能达到23.59 GFLOPS/s,获得了13.4倍的加速。
B. 从Affine到Linalg的提升
实验内容: 此实验将提升范围从单一的GEMM扩展到更通用的Linalg方言,涵盖矩阵乘法、矩阵向量乘积、卷积以及张量收缩的TTGT转换。评估了两种方案:
1. MLT-Linalg: 提升到Linalg后,使用Linalg默认的降低路径(如缓存分块)。
2. MLT-Blas: 提升到Linalg后,将Linalg操作替换为对供应商优化库(MKL-DNN)的BLAS调用。
将这两种方案与Clang -O3、Pluto编译器的默认设置(Pluto-default)和最佳设置(Pluto-best)进行比较。
实验结果与分析 (图9):
* MLT-Blas性能: 在两个平台上,对于映射到三级BLAS的核(如2mm, 3mm, gemm, 卷积和张量收缩),MLT-Blas均显著优于所有其他方法,包括Pluto的最佳性能。例如,在AMD系统上,ab-cad-dcb张量收缩相对于Pluto-best获得了高达294倍的加速。
* 二级BLAS性能: 对于映射到二级BLAS的核(如atax, bicg),Pluto-best的表现优于或持平于MLT-Blas。这主要是因为MLT动态链接库引入了额外的开销,在计算量较小的核上影响更明显。
* MLT-Linalg性能: MLT-Linalg的性能通常优于Clang -O3和Pluto-default。尤其是在张量收缩上,通过TTGT转换,MLT-Linalg能够显著改善数据局部性,从而获得更好的性能。
* 编译时间: Multi-Level Tactics引入的编译时间开销很小。在16个基准测试上,从Affine提升到Linalg再降低到LLVM IR,总编译时间仅增加了12%(从0.64秒增加到0.72秒)。
C. 渐进式提升案例:重排矩阵链式乘法
实验内容: 此实验展示了渐进式提升的应用。首先,将C代码中的矩阵链式乘法循环嵌套提升到Linalg方言,生成一系列linalg.matmul操作。然后,在Linalg层级应用第二个提升(或重写)步骤,通过利用矩阵乘法的结合律来重新组织计算顺序(重新加括号),以最小化总的标量乘法次数。检测矩阵链的策略如列表9所示。
linalg.matmul(linalg.matmul(A, B), C) // detect this
列表9:用于在Linalg方言中检测三个矩阵相乘链的结构匹配器。
实验结果与分析 (表II):
* 优化效果: 对于测试的三个不同规模的矩阵链,Multi-Level Tactics成功地找到了最优的计算顺序。
* 性能提升: 执行时间的减少与标量运算次数的减少成正比。例如,对于一个4矩阵链,优化后的执行时间从1.289秒减少到0.212秒,获得了6.08倍的加速。这证明了渐进式提升能够在更高的抽象层次上应用强大的算法级优化,而这是传统编译器在低级IR上难以做到的。
A7 补充细节
VI. 相关工作
习语识别(Idiom recognition)概述
习语识别是计算机科学中一个自1990年代以来就广为人知的老问题。以往的工作大致可分为三类:基于文本的、基于语法的和基于语义的。前两类因对代码变化的鲁棒性较差而逐渐被淘汰,舞台留给了更系统化的语义方法。
与语义方法的比较
* Ginsbach等人提出的IDL【索引28,Automatic matching of legacy code to heterogeneous apis: An idiomatic approach, SIGPLAN Not., vol. 53, Mar. 2018】:IDL(Idiom Description Language)将习语描述为一组约束,并通过约束求解器进行匹配。这种方法是全自动的,并且能够检测稀疏线性代数,这是MLIR和Multi-Level Tactics目前尚不支持的。然而,依赖约束求解器发现习语会大幅增加编译时间(平均增加82%)。相比之下,Multi-Level Tactics的开销可以忽略不计。在后续工作LiLAC中【索引29,Automatically harnessing sparse acceleration, Proceedings of the 29th International Conference on Compiler Construction, CC 2020, 2020】,他们引入了DSL以简化模式规范,类似于TDL,但每个模式只能映射到单个BLAS调用,而Multi-Level Tactics允许将一个模式表示为多个库调用的组合。
* Arenaz等人提出的XARK编译器【索引30,XARK: An Extensible Framework for Automatic Recognition of Computational Kernels, ACM Trans. Program. Lang. Syst., vol. 30, Oct. 2008】:该编译器通过分析Gated Single Assignment (GSA)形式上强连通分量(SCCs)中的use-def链来进行习语识别。然而,多维数组的线性化从根本上限制了可检测模式的复杂性。相反,Multi-Level Tactics在MLIR基础设施内工作,能够将多维访问作为一等公民处理。
* Chelini等人提出的Declarative Loop Tactics【索引1,Declarative loop tactics for domain-specific optimization, ACM Trans. Archit. Code Optim., vol. 16, Dec. 2019】:Multi-Level Tactics与Loop Tactics有很多共同点,但更进一步,通过TableGen自动合成匹配器,降低了编写匹配器的门槛,提高了生产力。
与语言扩展和DSL框架的比较
* Felleisen等人提出的Racket【索引31,A programmable programming language, Commun. ACM, vol. 61, Feb. 2018】:Racket是一种语言扩展API,用于扩展宿主语言的语法和语义,允许程序员嵌入上下文敏感信息以进行优化。虽然MLT与Racket以及更广泛的面向语言编程框架【索引32,Language-oriented programming, Software - Concepts and Tools, vol. 15, no. 4, pp. 147–161, 1994】【索引33,Languages as libraries, Proceedings of the 32nd ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI ’11, 2011】有相似之处,但存在显著差异:MLT专注于IR层面而非源码层面,服务于编译器开发者而非应用开发者。
* Brown等人在Delite框架中的工作【索引5,A heterogeneous parallel framework for domainspecific languages, 2011 International Conference on Parallel Architectures and Compilation Techniques, 2011】:Delite使用重写规则来应用领域特定优化。虽然最终目标相似,但Multi-Level Tactics通过从低级抽象进行提升来实现这一目标。
* Halide和TVM等框架【索引34,Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines, Acm Sigplan Notices, vol. 48, no. 6, pp. 519–530, 2013】【索引35,TVM: An automated end-to-end optimizing compiler for deep learning, 13th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18), 2018】:这些框架将问题的函数描述与其执行策略解耦。通常,前者由领域专家编写,后者由性能专家编写。这类工具要求领域专家熟悉其工具语法(如Halide语法)来编写问题的函数描述。而使用Multi-Level Tactics,领域专家可以编写纯C++代码,从而降低了语言障碍——假设专家已经了解C++,则无需学习新语言。同时或在此之前,由于匹配器和构建器的解耦特性,性能专家可以提供一个策略来决定最佳执行策略。
A5 结论
本文介绍了Multi-Level Tactics及其在MLIR框架中的实现。论文的核心思想是提出并实现“渐进式提升”(progressive raising),作为多级IR编译器中已有的“渐进式降低”(progressive lowering)的一个补充路径。我们提出的渐进式提升方法能够提升通用语言的入口点,从而在多级IR中启用领域特定的优化。
我们展示了如何使用Multi-Level Tactics,基于一种受Tensor Comprehension启发的简洁语法来表达模式和构建器,并为两种提升路径(Affine到Affine,Affine到Linalg)演示了我们的框架。此外,我们通过在Linalg方言之上实现矩阵链式乘法重排转换,展示了一个渐进式提升的案例。据我们所知,我们是第一个在像MLIR这样的多级IR编译器中提供基础设施并演示如何实现渐进式提升的工作。
未来的工作将致力于提供更多的提升路径。
A6 附录
A. 摘要
工件目标
此工件的目标是展示Multi-Level Tactics (MLT)如何将通用语言提升到更高的抽象层次,从而通过渐进式降低实现有效的领域特定编译。该工件包含一个Docker容器和配套脚本,用于复现图8、图9和表2。运行所有实验仅需此Docker容器。生成图表和表格的脚本已包含在Docker中。
B. 工件清单(元信息)
- 算法: Multi-Level Tactics,一种在MLIR框架之上实现的、用于渐进式提升的声明式方法。
- 程序: Polybench/C 4.2.1 beta版,以及从先前关于张量收缩的研究中收集的基准。此外,还考虑了一个二维卷积。对于Polybench/C 4.2.1,我们使用了一个修改版,其中每个核都经过了循环分发并翻译为MLIR中的Affine方言。所有使用的基准都包含在MLT仓库
https://github.com/LoopTactics/mlir的cgo分支中。 - 编译: 任何兼容C++11的编译器来引导LLVM/MLIR。
- 数据集: Polybench/C 4.2.1中预定义的
LARGE_DATASET。 - 运行时环境: 任何LLVM支持的Unix系统。
- 硬件: 任何LLVM支持的平台。
- 输出: 结果是复现图8、图9和表2的PDF文件。脚本已在Docker容器中。图8和图9以GFLOP/s表示结果,而表2以秒表示。中间文件也会生成,名为
result_X.txt。所有result_X.txt文件中的结果都以秒为单位。 - 所需磁盘空间 (大约): Docker镜像为7GB,应有15GB的磁盘空间。
- 准备工作流所需时间 (大约): 主要是构建MLT的时间(超过20分钟)。
- 完成实验所需时间 (大约): 20/25分钟。
- 是否公开可用: 是,通过Github和Dockerhub。
C. 描述
1) 如何交付:
* 我们通过Docker交付工件。可在https://hub.docker.com/r/lchelini/cgo获取。
* MLT和MLT的TDL DSL可在以下地址获取:https://github.com/LoopTactics/mlir 和 https://github.com/LoopTactics/TacticsDSL。
* 在MLIR开放设计会议上关于MLT的演示文稿可在此处找到。
2) 硬件依赖:
* 任何LLVM支持的平台,详见 https://llvm.org/docs/GettingStarted.html#hardware。
3) 软件依赖:
* Docker容器已包含所有依赖项,包括:
* 编译LLVM/MLIR所需的所有要求,详见https://llvm.org/docs/GettingStarted.html#software。
* MKL-DNNL,可从https://github.com/chelini/mkl-dnn.git获取。
* MKL库。
* 对于MLT的TDL DSL,我们建议使用llvm-9.0,可以通过sudo apt-get install llvm-9-dev在您的机器上安装。无需在提供的Docker容器中安装。
D. 安装
无需安装。
E. 实验工作流
Docker容器中包含三个文件:
* experiment5.1.sh 用于复现图8
* experiment5.2.sh 用于复现图9
* experiment5.2.time.sh 用于复现MLT引入的开销
* experiment5.3.sh 用于复现表2
执行流程:
$ docker pull lchelini/cgo $ docker run -it lchelini/cgo $ ./build.sh $ ./experiment5.1.sh $ ./experiment5.2.sh $ ./experiment5.2.time.sh $ ./experiment5.3.sh
运行./experiment5.X.sh后,可以在llvm-project/mlir/benchmark/section5.X中找到一个名为main.pdf的结果文件。要打开PDF文件,请使用docker cp命令将其复制到容器外部,参见https://docs.docker.com/engine/reference/commandline/cp/。脚本./experiment5.2.time.sh的结果会打印到标准输出,不生成文件。
F. 评估和预期结果
- 实验5.1: 我们评估MLT在处理以不同风格编写但语义等价的各种GEMM时的可靠性。我们预期只会在Darknet基准上错失一次提升机会,因为我们没有为线性化访问模式生成匹配器,也没有提供反线性化过程。
- 实验5.2: 我们展示了提升到更高抽象(即MLT-Linalg或MLT-Blas)如何让我们获得比
Clang -O3更好的性能。我们预期MLT-Linalg和MLT-Blas能达到比基准Clang -O3更高的GFLOP/s。此外,我们还提供了另一个基准:Pluto。我们预期Pluto会比Clang -O3好,但效果不如(或相当)MLT-Linalg。我们预期Pluto的效果不如MLT-BLAS。注意,Pluto-best不可用,因为它依赖于昂贵的自动调优,需要数天才能收敛。 - 实验5.3: 我们通过实现矩阵链重排转换来展示一个更进一步的渐进式提升案例。我们预期
Time IP > Time OP。IP是未经重排的矩阵链的时间,而OP是重排后链的时间。
💬 评论讨论
欢迎在这里分享您的想法和见解!