协调抽象和高性能: MetaOCaml方法

摘要

生成式编程的一个常见应用是生成高度特化于手头问题的高性能计算内核. 一个典型的线性代数内核会针对数值域 (有理数, 浮点数, 双精度浮点数, 等等), 循环展开因子, 数组布局和先验知识 (例如, 矩阵是正定的) 进行特化. 手动特化 (编写相同算法的诸多变种) 是乏味无聊且容易出错的.

广泛使用的生成器, 例如ATLAS和SPIRAL, 能够可靠地产生高度特化的代码, 但是很难扩展. 对于ATLAS而言, 其使用printf生成代码, 甚至连括号匹配都成为了挑战. 根据ATLAS的作者所言, debug如同梦魇.

诸如MetaOCaml这样的类型化多阶段编程语言, 允许我们先陈述一个通用的显然正确的算法, 然后以模块化的方式逐层添加特化处理. 由于MetaOCaml能够确保生成的代码总是可以编译, 并且让我们能够快速测试, 它使得编写生成器变得不那么令人畏惧, 也更加高效.

读者将在这篇实践教程中亲身体会到这一点. 本教程不假设读者具有MetaOCaml的预备知识, 只需对函数式编程有基本的了解. 我们最终将实现一个简单的线性代数DSL, 其中包含针对矩阵和向量的稀疏性, 内存布局及其代数性质的多层优化. 我们将生成最优的BLAS内核. 我们将体会到无罪恶感的抽象的魅力.

第1章引论

第1.1节为什么元编程?

在编程的所有领域中, 一个令人痛苦的权衡始终存在: 一边是可维护可复用, 易于阅读和理解, 显然正确的教科书式代码, 另一边则是性能优异的代码. 这种权衡在高性能计算 (HPC) 中尤为突出. 将矩阵和向量的乘法简单地写成a * v, 既清晰可移植, 又能自我描述. 然而, 典型的高性能代码在执行整数矩阵与整数向量的乘法时, 往往需要非常多的代码行, 且完全不是一目了然的. 它看起来与a * v毫无相似之处. 它同样与单精度浮点矩阵乘以浮点向量的代码截然不同. 而后者与稀疏矩阵乘以向量的高性能代码之间, 也几乎看不出任何相似性.

早在上世纪末, 人们就已经认识到, 我们不能再依赖优化编译器来将高级代码转换为高性能代码 (参见Cohen et al. (2006) 中的参考文献): 许多有价值的优化是领域特定的, 而且往往适用范围很窄, 因此不太可能被通用编译器所支持. 即使是最简单的将0*e替换为0的操作, 在一般情况下也是不正确的: 试想e调用了外部函数或返回了NaN的情形. 一位领域专家, 凭借对输入数据和整个算法的了解, 能够判断e的副作用可以忽略, 或者NaN不会出现——因此对于特定表达式中的特定乘法, 该优化应当被执行.

让通用编译器为程序员提供如此精细的优化控制, 无论在认知上还是经济上都是不可承受的. 因此, 专家们手工编写计算内核是非常普遍的做法——并且不得不反复重写, 以适应新的硬件架构或输入数据中的新模式.

元编程——具体来说是代码生成——提供了一条出路: 我们不直接编写程序, 而是编写程序生成器, 将领域特定知识融入其中, 输出若干底层的特化的高性能程序. 这正是以下广为人知且被广泛使用的项目所采取的方法: 快速傅里叶变换生成器FFTW (Frigo和Johnson, 2005), 基础线性代数 (BLAS) 生成器ATLAS (Whaley和Petitet, 2005), DSP与线性代数生成器SPIRAL (Püschel et al., 2005), 以及图像滤波器生成器Halide (Ragan-Kelley et al., 2013).

上述项目同时也表明, 编写一个优秀的生成器仍然非常困难: 其成果足以在顶级学术会议上发表论文. 例如, ATLAS——使用C语言以字符串形式生成C代码——在编写, 调试和扩展方面一直以困难著称. 我们需要代码生成方面的辅助工具——MetaOCaml, Scala中的轻量级模块化分阶段框架 (Rompf和Odersky, 2012) 或Template Haskell (Sheard和Peyton Jones, 2002) 正是提供这种帮助的. 我们需要多层次的抽象.

理想情况下, 最终用户只需将矩阵-向量乘法生成器写成a * v即可. 这个(领域特定的)运算*将使用一个不同领域MapReduce的语汇来实现 (可能由另一位程序员, 比如算法设计者完成):

let dot v1 v2 = reduce add zero (zip_with mul v1 v2)
let (*) a v = map (dot v) a

其中reduce, add等生成器将由其他领域专家, 即数据布局方面的专家来提供. 而矩阵和向量所定义的域上的专家, 则会提供一个代数法则库, 用于简化标量表达式. 最终由MetaOCaml负责生成OCaml代码, 或者通过离岸 (offshoring)机制生成C代码或LLVM代码. 这一理想是可以实现的! 事实上, 在本教程结束时, 我们将恰好实现这样一个分层的, 面向简单线性代数的领域特定语言. Rompf et al. (2013) 以及FEniCS项目 (Markall et al., 2013) 展示了更多此类通过组合逐步细化的抽象层而构建的生成器DSL的实例, 以及对它们的实验评估.

总而言之, 我们确实让最终用户以他们最熟悉的领域词汇, 以对他们来说最清晰的形式编写程序——同时获得针对各种领域调优的高性能代码. 借用Ken Kennedy的话来说, 元编程赋予了我们无罪恶感的抽象.

第1.2节为什么这个教程?

本教程的目标是教授如何编写类型化的代码生成器, 如何使其模块化, 以及如何借助MetaOCaml逐步引入领域特定的优化. 在教程结束时, 我们将实现一个简单的线性代数DSL, 其中包含针对矩阵和向量的内存布局, 稀疏性及代数性质的多层优化. 我们将生成最优的基础线性代数 (BLAS) 内核. 希望读者能够看到, 编写生成器并没有那么复杂, 而且(分阶段的)类型系统能提供极大的帮助.

读者无需事先了解MetaOCaml, 但应对某种现代函数式语言有一定的熟悉. 即使只有ML语言家族的简短使用经验, 也会大有裨益. 不过, Scala或Haskell等语言的程序员也不必觉得被排除在外.

本教程基本上是此前多次现场教程的书面记录 (首次于2013年CUFP——函数式编程商业用户大会上讲授). 它继承了那些教程的动手实践风格, 围绕现场编码展开, 与MetaOCaml及其类型检查器——以及听众——进行互动. 我们将逐步开发代码, 通过向MetaOCaml解释器提交小段代码片段, 修复其指出的类型问题, 生成示例代码并进行测试, 注意可以改进之处并相应调整生成器. 本教程包含许多练习和课后项目, 供读者独立完成或分组合作.

第1.3节为什么MetaOCaml

我们将使用BER MetaOCaml (Kiselyov, 2017, 2014), 它是对由Walid Taha, Cristiano Calcagno及其合作者 (Calcagno et al., 2003) 开发的, 现已不可用的原始MetaOCaml的一次完整重新实现.

BER MetaOCaml是OCaml的一个保守扩展, 用于编写生成程序的程序. BER MetaOCaml在OCaml的基础上增加了代码值类型 (表示程序代码, 即未来阶段的计算), 以及两个用于构建代码值的基本构造: 引用 (quoting) 和拼接 (splicing) . 生成的代码可以被打印存储到文件中, 也可以被编译并链接回正在运行的程序, 从而实现运行时代码优化. 不含分阶段注解的MetaOCaml代码, 或者擦除注解后的代码, 就是普通的OCaml代码.

MetaOCaml已被成功应用于最优流融合 (Kiselyov et al., 2017), 数值和动态规划算法的特化, 构建FFT内核, 图像处理和数据库查询DSL的编译器, OCaml服务器页面, 生成特化的基础线性代数和高斯消元例程, 以及高性能模板计算 (Aktemur et al., 2013). 关于MetaOCaml的应用集锦, 请参阅Lengauer和Taha (2006).

用MetaOCaml这样的类型化分阶段语言编写代码生成器有诸多好处. 首先, 生成的代码将是良构的, 所有括号都能正确匹配. 这种保证在使用printf编写C代码 (如ATLAS中所做的那样) 或使用Matlab编写C++代码时是一个迫切的愿望. MetaOCaml确保生成的代码是良类型的, 编译时不会出错. 不再需要在生成的代码中费力排查编译错误——生成的代码通常体量庞大, 晦涩难懂且变量名毫无帮助. 更重要的是, 代码生成错误是基于生成器而非生成的代码报告的. 本教程将提供大量机会来体会良好错误报告的重要性.

MetaOCaml的生成器是卫生的 (hygienic), 能够产生作用域良好的代码, 不存在未绑定变量. 否则, 卫生性违规在实践中难以检测, 并可能导致意外绑定变量这一隐蔽的错误. 虽然生成代码中的未绑定变量在编译时会暴露出来, 但确定其成因在实践中被证明极其困难, 正如Ofenbeck et al. (2016) 所报告的那样. 该文作者编写了一个新的编译器测试框架, 专门用于检测在重构生成器过程中引入的未绑定变量及其他类似问题. MetaOCaml的设计从源头上就防止了这类问题代码的生成.

最重要的是, MetaOCaml是类型化的. 类型, 尤其是分阶段类型, 确实有助于编写代码. 在整个教程中, 我们将在与类型检查器的实时交互中编写代码——不把类型错误视为惩罚, 而是将其视为有价值的提示. 我们将在许多场合看到, 一旦确定了类型签名, 生成器几乎就能自动写出来. 类型检查器会告诉我们应该在哪里放置分阶段标注.

MetaOCaml是纯生成式的: 生成的代码被视为黑盒, 无法被检查. 人们可以将代码组合在一起, 但无法将其拆开. 纯生成性显著简化了类型系统, 并增强了静态保证. 这似乎也意味着纯生成性排除了代码优化的可能. 但幸运的是, 我们很快就会看到, 事实并非如此.

MetaOCaml的分阶段标注就像是元编程的汇编指令. 我们需要更高层次的抽象. MetaOCaml相较于camlp4或ppx等预处理器的最后一个好处在于, 它是OCaml自身的一部分, 因此能够充分利用OCaml的抽象和组合机制, 从高阶函数到模块系统. 构建优化库以及组合生成器是本教程的重点.

第1.4节概览

本教程基于一系列循序渐进的问题, 除了引导性的第一个问题外, 其余都是实际问题的略微简化版本:

分阶段计算与MetaOCaml的入门
数字滤波器
复数向量乘法: 变化的数据表示 (结构体的数组 vs. 数组的结构体)
简单线性代数的系统化优化: 构建广泛特化的通用BLAS
从解释器到编译器: 图像处理DSL
进阶挑战 (课后作业)

事实上, 问题3, 4和6是由高性能计算 (HPC) 研究人员作为对程序生成社区的挑战而提出的 (Shonan挑战). 贯穿其中的共同主题是构建针对特定问题高度调优的高性能计算核. 因此, 大多数问题围绕简单线性代数展开——这是HPC中最典型且执行最频繁的部分.

对高性能应用以及模块化优化和生成器的强调, 使本教程有别于Taha极具亲和力的, 对经典部分求值和分阶段计算的温和介绍. Taha的介绍侧重于将一个通常为高阶语言的解释器转化为编译器 (Taha, 2004, 2008). 我们在第6章也会涉及这一经典领域; 不过, 我们对lambda演算的关注较少, 而更多地关注图像处理. 此外, 本教程还提及了MetaOCaml的近期新增功能, 如离岸编译 (offshoring) 和let插入 (let-insertion).

第1.5节获得MetaOCaml

第2章第一步

第2.1节 Now or later

第2.2节幂

尽管已经成为陈词滥调, power这个例子仍然是介绍元编程的最快方式之一. 以下的power函数可以计算 $x^{n}$ :

let square x = x * x
let rec power : int -> int -> int = fun n x ->
  if n = 0 then 1
  else if n mod 2 = 0 then square (power (n/2) x)
  else x * (power (n-1) x)

我们已经显式地附加了类型签名, 尽管并无必要: 类型可以推断出来. 不过, 给top-level的函数添加签名一般而言是个好的想法: 它可以使得代码更容易编写, 之后更容易理解, 并且能够使得错误信息更加清晰. 在元编程中, 编写签名有助于放置分阶段标注 (有时甚至可以达到自动化的程度), 我们很快就会看到这一点. 作为一个快速验证,

power 7 2;;
⇝ - : int = 128

设我们的程序需要计算 $x^{7}$ 很多次, 那么我们可以定义

let power7 x = power 7 x

来命名和分享这部分求值.

在MetaOCaml里, 我们也可以特化幂函数于一个特定的值 $n$ , 得到之后可以接受 $x$ 并计算 $x^{n}$ 的代码. 我们将重写power n x, 并将表达式标注为现在计算 (当n已知时) 或之后计算 (当x给出时). 让我们逐步进行这一过程.

作为第一次尝试, 我们复制粘贴之前的power代码, 将函数名改为spower, 并修改类型签名以表明哪些参数是现在已知的 (它们具有普通类型, 如int), 哪些参数要到稍后才能知道. 后者的类型形如int code, 无法被立即参与计算. spower的返回结果也将具有int code类型.

练习4. 为什么结果类型为int code?

除了将power改名为spower, 我们没有对于函数体进行任何调整:

let rec spower : int -> int code -> int code = fun n x ->
  if n = 0 then 1
  else if n mod 2 = 0 then square (spower (n/2) x)
  else x * (spower (n-1) x)

无怪乎类型检查器会抱怨, 因为此时签名和函数体不再对得上了:

⇝    Characters 74-75:
  if n = 0 then 1
                ^
Error: This expression has type int but an expression
was expected of type int code

我们应当认真对待这个报错: 它表明类型检查器期望在1的位置看到int code类型的东西. 为了满足这一期望, 我们只需将1用括号 (brackets) 括起来, 这样就得到了int code类型的表达式. 因此, 代码的第二个版本如下 (我们高亮了修改之处):

let rec spower : int -> int code -> int code = fun n x ->
  if n = 0 then .<1>.
  else if n mod 2 = 0 then square (spower (n/2) x)
  else x * (spower (n-1) x)

⇝    Characters 114-130:
  else if n mod 2 = 0 then square (spower (n/2) x)
                                  ^^^^^^^^^^^^^^^^
Error: This expression has type int code
but an expression was expected of type int

类型检查器再次报错, 不过这次是在另一个位置: spower会生成之后将给出一个int的代码, 但square现在就需要一个整数.

练习5. 你觉得我们应该做什么?

以下是解决这一类型不匹配的一种方法: 既然我们无法现在就给square一个int, 那么我们就必须将整个应用推迟到之后:

let rec spower : int -> int code -> int code = fun n x ->
  if n = 0 then .<1>.
  else if n mod 2 = 0 then .<square .~(spower (n/2) x)>.
  else x * (spower (n-1) x)

⇝    Characters 146-147:
  else x ∗ (spower (n-1) x)
       ^
Error: This expression has type int code
but an expression was expected of type int

类型检查器指出, 我们同样需要延迟与x的乘法运算, 因为x的int code类型表明它在当前阶段不可用. 代码的下一个版本如下:

let rec spower : int -> int code -> int code = fun n x ->
  if n = 0 then .<1>.
  else if n mod 2 = 0 then .<square .~(spower (n/2) x)>.
  else .<.~x * (spower (n-1) x)>.

        Characters 154-170:
  else .<.~x ∗ (spower (n-1) x)>.
               ^^^^^^^^^^^^^^^^
Error: This expression has type int code
but an expression was expected of type int

这段代码仍有不足之处: 我们忘了去escape spower的调用. 如果e具有类型int code, 那么.~e显然具有类型int (escape会消解括号, 还记得吗?), 然后类型检查器就应该得到满足了:

let rec spower : int -> int code -> int code = fun n x ->
  if n = 0 then .<1>.
  else if n mod 2 = 0 then .<square .~(spower (n/2) x)>.
  else .<.~x * .~(spower (n-1) x)>.

确实如此. 我们已经成功地将power转变为一个代码生成器: 仅仅通过调整签名, 并在类型检查器报错的位置添加标注即可. 类型确实帮助我们编写了代码.

然而, 我们还没有完全完成. 考虑特化spower 7, 它的类型是int code -> int code. 这是一个代码模板的类型. 而一个描述未来阶段函数的代码值——该函数将接收 $x$ 并计算 $x^{7}$ ——具有不同的类型, 即(int -> int) code. 不过, 将前一种类型的值转换为后一种类型是相当容易的. 类型再次提供了帮助: 如果我们需要类型为t code的东西, 就必须找到一个类型为t的表达式, 并将其用括号括起来. 因此, 作为第一步近似:

let spowern : int -> (int -> int) code = fun n -> .<fun x -> spower n x>.

它看上去就像一个eta-扩展, 不是吗? 类型检查器很快就指出了其缺失的分阶段注解的位置. 以下是最终的版本:

let spowern : int -> (int -> int) code =
  fun n -> .<fun x -> .~(spower n .<x>.)>.

这即是我们所想要的幂. 让我们检查一下:

spowern 7;;
⇝ - : (int -> int) code = .<
fun x_1 → x_1 ∗ ((∗ CSP square ∗) (x_1 ∗ ((∗ CSP square ∗) (x_1 ∗ 1))))>.

我们已经看到, 代码值是可以打印出来的. 事实上, 即使是函数的代码也可以打印出来, 因此我们总是能够看到自己生成了什么.

我们还注意到了神秘的 (∗ CSP square ∗). CSP代表跨阶段持久化 (cross-stage persistence). 回想一下, square是一个当前已定义且可用的函数. 然而我们把它放在了括号内, 假定它在将来也同样可用. 可以把(∗ CSP ... ∗)看作一种容器, 它将当前阶段的值保存起来以供未来使用.

练习6. 乘法也是一个CSP吗?

我们可以安全地假定, 标准库 (当前可用的) 在运行任何生成的代码时也同样可用. 因此, 标准库的CSP值不需要容器. 将square变成一个标准库函数是很简单的: 只需将它的代码放入一个单独的文件中. 让我们来做这件事: 创建文件square.ml, 其中包含一行代码

let square x = x * x

将其编译为ocamlc -c square.ml. 接下来重新定义spower, 使刚刚编译好的square可用 (注意, 要打开的模块名称是去掉.ml后缀并将首字母大写的文件名: 详见OCaml参考手册):

open Square;;
#load "square.cmo";;

let rec spower : int -> int code -> int code = fun n x ->
  if n = 0 then .<1>.
  else if n mod 2 = 0 then .<square .~(spower (n/2) x)>.
  else .<.~x * .~(spower (n-1) x)>.
;;
⇝ val spower : int -> int code -> int code = <fun>
let spowern : int -> (int -> int) code =
  fun n -> .<fun x -> .~(spower n .<x>.)>.;;
⇝ val spowern : int -> (int -> int) code = <fun>
let spower7c = spowern 7;;
⇝ - : (int -> int) code = .<
fun x_2 -> x_2 ∗ (Square.square (x_2 ∗ (Square.square (x_2 ∗ 1))))>.

练习7. 生成的特化代码看起来已经相当优化了, 除了乘以一这个操作. 你认为我们应该如何消除这个无用的操作?

打印出的spower7c代码中的变量很显眼: MetaOCaml会系统性地重命名所有绑定变量, 方法是在用户给定的名称后附加数字后缀. 为了理解其原因, 请考虑以下代码及其生成的代码值:

.<fun x -> .~(let body = .<x>. in .<fun x -> .~body>.)>.
⇝ .<fun x_1 -> fun x_2 -> x_1>.

如果我们将外层变量重命名为y:

.<fun y -> .~(let body = .<y>. in .<fun x -> .~body>.)>.
⇝ .<fun y_1 -> fun x_2 -> y_1>.

生成的代码与之前表示的是同一个函数 (K组合子). 这种对绑定变量的系统性重命名 (

α

-变换) 保持不变的性质, 称为卫生性 (hygiene).

练习8. 如果MetaOCaml像Lisp一样不对绑定变量进行重命名, 那么上面两个例子生成的代码会是什么? 它所表示的还会是同一个函数吗?

第2.3节离线代码生成

我们已经对幂函数进行了特化: 我们编写了生成器spowern, 它能生成一段 (相当优化的) 代码, 该代码接受 $x$ 并返回 $x$ 的 $n$ 次方, 其中 $n$ 是固定的. 我们可以打印生成的代码. 同样, 我们也应该能够轻松地将生成的代码保存到文件中. 事实上确实可以. MetaOCaml甚至为此提供了一个专用函数: Print_code.format_code可以将代码值写入指定的输出流 (严格来说是Format.formatter). 保存后的spower7c代码文件看起来与其他任何OCaml代码文件无异: 它可以用字节码编译器或优化OCaml编译器进行编译, 并链接到用户程序中.

我们刚刚看到了使用MetaOCaml的一种方式: 生成特化代码以及由特化函数组成的完整库, 以供后续链接使用. 这些特化代码也可以立即执行.

第2.4节运行时特化与基准测试

为了测试生成的代码, 我们可以按照刚才描述的方法将其保存到文件中, 然后编译并与测试驱动程序进行链接. 幸运的是, 还有一种更简便的方法. 生成的代码可以立即运行: MetaOCaml函数Runcode.run : 'a code -> 'a接受一个代码值, 对其进行编译, 将其链接回程序中, 然后求值并返回结果. 它确实是在一个正在运行的程序内部调用了编译器:

open Runcode;;
let spower7 = run spower7c;;
⇝ val spower7 : int -> int = <fun>

现在我们可以测试一下:

spower7 2;;
⇝ - : int = 128

特化的spower7和之前我们看过的部分应用的power7有着相同的类型 (以下为了参考重复一遍):

let power7 = power 7;;
⇝ val power7 : int -> int = <fun>

两者都返回其参数的

7

次幂. 然而, power7 2实质上是power 7 2的伪装, 它会递归地检查第一个参数的奇偶性. 相比之下, 特化后的spower7c没有递归 (从其代码中可以看出). 所有关于n的操作在生成spower7c时就已经完成了; 生成的代码仅仅是对x进行乘法运算.

特化究竟能带来多大的帮助? 让我们通过对power7和spower7进行基准测试来一探究竟. 为了追求性能, 我们将使用优化(Meta)OCaml编译器, 而不是基于字节码的metaocaml解释器. 在一个新文件power_rts.ml中, 我们放入之前的power, spower和spowern代码, 以及计时代码

第2.5节回顾

有两个重要的洞察值得强调. 第一, 一个代码值可以表示开放代码, 即包含自由变量的代码. 例如, 在.<fun x -> .~(spower n .<x>.)>.中, 代码值.<x>.就是未来阶段的变量x本身. 我们可以将这样的变量存储在引用单元中, 也可以将它们作为参数传递或作为函数的返回值. 因此, 分阶段计算让我们能够以符号化的方式操作(未来阶段的)变量. 然而, 我们无法确定一个未来阶段变量的真实名称, 更不用说对其进行篡改了. 代码值不能被拆解和检查. MetaOCaml的这种纯生成性有助于维护卫生性: 开放代码可以被操作, 但词法作用域仍然得以保留.

第二个洞察是关于使用MetaOCaml的两种不同方式. 第一种是运行时特化: 根据运行时获取的数据 (例如来自用户输入的数据) 对常用函数进行特化. 特化后的代码可以立即使用, 即直接运行. 第二种方式是将生成的代码保存到文件中. 由于它就是普通的OCaml代码, 因此可以用普通的OCaml编译器进行编译, 并在之后链接到各种普通的OCaml应用程序中. 在这种方式下, MetaOCaml运行代码值的功能并不是必需的: 运行仅用于测试或经验性的代码搜索.

因此, MetaOCaml不仅可以用于运行时特化, 还可以用于离线生成特化的库代码, 例如BLAS和Linpack库. 我们可以对生成的代码进行基准测试, 并选出在给定平台上性能最佳的版本. 这样一来, MetaOCaml就可以像ATLAS和其他类似的元编程系统一样使用. 本教程后续章节将进一步强调这一点.

最后, 我们利用类型检查器来判断分阶段标注应该放在哪里, 这不禁让人思考这个过程是否可以自动化. 我们能否只提供期望的签名, 就获得带有正确标注的特化代码吗? 这种类型驱动的绑定时间分析 (BTA) 在简单情况下确实是可行的 (Asai, 2016). 然而, 我们很快就会看到, 很多时候需要创造性的步骤 (§3.1).

到目前为止, 我们一直在以纯函数式 (应用式) 风格生成无副作用的纯函数式 (应用式) 代码. 在§3中, 我们将遇到使用可变单元的生成器, 这是特化命令式代码的自然结果. 之后, §4将展示 (实际上是纯函数式的) 生成器如何产生命令式的高性能数值代码.

第2.6节历史插曲

幂函数的例子是那篇开创了部分求值这一研究领域的论文中的贯穿示例: 即A.P. Ershov 1977年的论文 (Ershov, 1977a). 该例子出现在论文的图1中, 可以在以下手稿扫描件中看到: http://ershov.iis.nsk.su/files/archive/fold0241/241_75.gif. 有时人们会看到一个线性时间的版本, 其中 $x$ 的 $n$ 次方是通过反复乘以 $x$ 来计算的. 但那并不是原始的例子: Ershov和当时大多数俄罗斯程序员一样, 对效率有着近乎痴迷的追求. 他们绝不会考虑一个低效的算法.

那篇广为人知的Ershov 1977年论文是Ershov (1977b)的英文版, 后者发表在Doklady AN SSSR (苏联科学院院报) 上. 由于篇幅限制, 幂函数的例子在Doklady AN SSSR的论文中被省略了, 但可以在1976年11月草稿的注释中看到: http://ershov.iis.nsk.su/files/archive/fold0241/241_17.gif. 这或许是幂函数特化的最早出现.

关于这篇论文的其他草稿, 请参见The Ershov Archive for the history of computing.

第3章滤波

滤波器在无线电以及各种语音和声音应用中无处不在. 当我们想要收听90.5 MHz的电台时, 需要从传入的无线电信号中滤除除以90.5 MHz为中心的窄带之外的所有信号. 二维滤波同样广泛应用于相机, 打印机和通用图像处理中. 滤波器通常是模拟的. 随着专用数字信号处理器以及如今高速通用CPU的出现, 以软件方式对信号进行数字滤波成为可能. Richard Hamming的经典著作Hamming (1997, 1983, 1989)对数字滤波器作了清晰而引人入胜的介绍. 另一份有用的在线参考资料是Smith (2007).

本节从入门过渡到分阶段计算的实际应用. 我们这里的数字滤波器仍然不够优化. 我们将进一步的优化步骤推迟到学习系统化方法之后, 即第4章和第5章.

与之前一样, 我们将在与MetaOCaml的交互中编写代码, 向MetaOCaml解释器提交(试探性的)定义, 检查其响应, 并立即测试所定义的函数和生成器. 让我们启动metaocaml并开始一个新的会话. (本节的完整代码在附带的文件filter.ml中.)

数字滤波对于表示为采样或者说一致采样信号值的序列的数字化了的信号进行变换: $x_{i}$ 代表在第 $i$ 次采样的时刻的信号幅度. 这个教程处理的是所谓有限冲激响应 (FIR) 滤波, 具有以下的一般形式: $y_{i} = \sum_{k = 0}^{m - 1} b_{k} x_{i - k}$ 输出信号的采样 $y_{i}$ 是 $m$ 个最近的输入信号采样 $x_{i}, x_{i - 1}, \dots, x_{i - m + 1}$ 的加权和. 权重 $b_{k}$ 是滤波器的系数. 数字 $m - 1$ 被称为滤波器的阶. 既然系数的数目 (一般还有系数的值) 往往是已知的, 那么特化滤波器的计算是很合理的. 滤波器往往需要处理很长的信号, 或者需要在线处理信号.

这里我们遇到了我们之前已经说过太多了的一般性和性能之间的取舍, 而且我们还要再遇到无数次. 我们想要以全然的一般性一劳永逸地写完滤波器的代码, 并且希望其具有对于领域专家而言熟悉的形式. 那么, 我们想要借助于特定的情形 (例如系数已知) 来特化或者优化代码. 我们不会在本章立刻达到理想的境地.

上述的滤波器等式可以按照以下习用方式在OCaml之中进行编码:

let filter : float array -> float array -> float array = fun b x ->
  let m = Array.length b in
  let y i =
    if i < m-1 then x.(i) else
    let sum = ref 0.0 in
    for k = 0 to m-1 do
      sum := !sum +. b.(k) *. x.(i-k)
    done;
    !sum
  in
  Array.init (Array.length x) y        (* essentially, for-loop *)

函数filter接受系数的数组b和输入采样的数组x, 返回经过滤波处理的(新鲜分配的)信号采样数组. 在OCaml之中, 取一个数组x的第i个元素记作x.(i). 为了简单起见, 我们没有考虑原地滤波或者类似于流的滤波. 然而, 本章我们所学到的

第4章线性代数DSL: 复向量运算和数据布局

现在我们转向生成优化代码的系统性方法. 我们的目的在于以清晰和显然正确的方式编写代码, 并且记号近于领域专家所使用的——然后将这种记号解释为一个代码生成器, 然后逐步添加各种优化. 每个优化都体现了某种领域知识, 由领域专家独立开发. 和前一章作为对比的是, 我们特化和优化代码而不需要不断重写和注解.

使得优化模块化和可复合的是一个共同的领域特定语言 (DSL), 或者更准确地说, 是一种特定的DSL实现技巧: 所谓的tagless-final风格. 当在类似于OCaml这样的宿主语言之中嵌入一个DSL时, 一般会将DSL表达式表示为某个代数数据类型的值: 每个DSL语言形式, 诸如数值字面量或者加法, 都对应于这个数据类型的一个数据构造子. 这是所谓的深嵌入. tagless-final方法则是一种替代. 其被定义为从DSL表达式到宿主语言 (OCaml) 的映射, 因而可以视为表达式的意义, 或者说指称. 这个映射是复合性的: 每种DSL表达式形式都对应于一个宿主函数, 其根据其直接组件的意义来计算形式的意义. 因此, tagless-final方法将一个DSL表示为一个代数, 也称为一个代数结构. 我们可以选取各种OCaml的值集合作为DSL表达式的意义. 换言之, 我们可以以不同的方式解释相同的表达式, 这取决于手头上的任务. 一个代码生成器只是其中一种解释. tagless-final风格的DSL是可扩展的: 特性可以依照愿望增加, 而不需要破坏之前所写过的任何代码. 解释也是可扩展的, 可以根据需要添加(领域特定的)优化. 这种DSL是有类型的. 解释 (代码生成器) 也是 (和MetaOCaml本身一样) 有类型的. 因此, 生成的代码总是良形式, 卫生, 且良类型的. 类型不仅阻止了显然错误的生成器, 而且还能够帮助我们编写生成器. 因此, 在接下来的章节里, 我们将会开发以tagless-final风格嵌入(Meta)OCaml的DSL.

我们的主题是线性代数. 本章处理的向量-向量运算 (也就是所谓的BLAS Level 1) 和数据布局. 第5章继续处理矩阵-向量运算 (BLAS Level 2), 代数化简, 以及领域特定的稀疏矩阵优化. 我们的DSL代码看起来会像是

let vmult vout v1 v2 = vout := v1 *. v2
let mvmult vout a v = vout := a * v

~~{译注: 内心小小的疑问, 这里的a * v真不是a *. v吗? 抑或是甚至v1 *. v2其实是v1 * v2? 总之感觉不是很协调.}~~ 这里的记号类似于Matlab, 还能生成特化于特定数据布局的命令式循环. 顺便一说, 这里也没有见到bracket和escape {译注: 用于元编程的}. 的确如此, 以精心择取的抽象, 面向用户的代码 (以及许多优化层) 并不会出现台面上的staging注解. 它们只会出现在非常低层次的生成器里.

第4.1节数据布局问题

这个问题是对于一种特定的数据表示进行实际的特化. 某些计算机架构偏爱所谓的结构的数组, 而GPU和向量处理器则偏爱数组的结构. 将一个程序从GPU移植到超级计算机上或者反过来一般需要改变数据表示. 我们的目的在于自动化这样的改变, 仅仅实现算法一次, 然后机械化地使其适应于不同的数据布局.

具体来说, 当前的问题是实现两个复向量的逐点乘法, 以允许轻易特化于不同复向量表示的方式:

float cmplx array
一个数组, 其元素都是复数: 浮点数序对, 分别代表 $R$ 部和 $I$ 部. 这种所谓的结构的数组表示流行于基于缓存的机器之中.
单独的float array ( $R$ 部) 和float array ( $I$ 部)
向量处理器和GPU更青睐于这种数组的结构表示.
一个单独的float array
前半部分 $R$ 部, 后半部分 $I$ 部, 前一种表示的变种.

v: float cmplx array $\underset{\underset{v_{0}}{⏟}}{\begin{matrix} v_{0} . R & v_{0} . I \end{matrix}} \underset{\underset{v_{1}}{⏟}}{\begin{matrix} v_{1} . R & v_{1} . I \end{matrix}} \dots$
v': float array cmplx $\underset{\underset{v^{'} . R}{⏟}}{\begin{matrix} v_{0}^{'} . R & v_{1}^{'} . R & \dots \end{matrix}} \underset{\underset{v^{'} . I}{⏟}}{\begin{matrix} v_{0}^{'} . I & v_{1}^{'} . I & \dots \end{matrix}}$
图4.1: 数据布局: 结构的数组v和数组的结构v'

不同布局在图4.1之中进行了系统地比较. 为了使得区别更加明显,

第4.2节抽象算术

我们开始教授系统性的方法, 通过为能够相加和相乘的东西设计一个DSL. 这看起来像是经典开局, 特别是对于HPC实践者而言: 算术在HPC之中是如此基础, 以至于任何的运算开销都是不可忍受的. 犬儒主义者或许会将我们的DSL当作沉溺于抽象废话 (abstract nonsense) 的借口. 很快我们将会看到, 我们的抽象不仅是实际的, 而且最终还是无开销的.

我们将再次以与MetaOCaml顶层环境交互的方式开发代码, 提交定义, 查看MetaOCaml的响应, 并立即测试所定义的函数和生成器. 作为准备, 让我们启动metaocaml并开始一个新的会话. (本节的完整代码位于附带的文件ring.ml, vector.ml, complex.ml中.)

我们从最小的DSL开始, 只有目前所需的运算. 更多的东西总是可以之后再加: 可扩展性是tagless-final方法的优点. 暂时我们需要的是一集我们可以做加法, 减法, 乘法的对象, 带有两个突出元素zero和one. 这集操作可以组织为所谓的(模块)签名:

module type RING = sig
  type t
  val zero : t
  val one  : t
  val add  : t -> t -> t
  val sub  : t -> t -> t
  val mul  : t -> t -> t
end

声明type t给我们的对象的类型赋予了一个(非常没有想象力的)名字. 这个类型是抽象的: 我们对其一无所知, 除了它的名字t. 具体的实现可以实现t为浮点数或者复数之类的东西. 然而, DSL代码并不能知晓或者依赖于t的具体实现 (这是OCaml的类型系统所控制的), 因而可以应用于任何可能的实现. 我们也可以选取更加fancy的名字, 比如说+%甚至是+而非add.

暂时我们仅仅是定义了t上的运算的名字和它们的元数. 这些运算被隐式地假定满足我们所熟悉的性质: add是交换且结合的, 并且具有单位元zero; mul是结合的, 且具有单位元one和零元zero {译注: zero是所谓的零元的事实可由环的公理导出}; mul对于add分配. 换言之, 我们是要定义一个数学上的环. 尽管如此, 我们还没有在地方陈述过这些代数律. 第5章里我们将会接触它们.

我们可以将签名RING视为定义了一个迷你DSL, 其具有两个常量和三个运算. (实际上, 这就是tagless-final的观点.) 这个DSL可以立刻使用:

module ExArith(R:RING) = struct
  open R
  let xsq1 x = add (mul x x) one
end
(* module ExArith : functor (R : RING) -> sig val xsq1 : R.t -> R.t end *)

例如, 编写一个对其参数求平方再加一的函数. 第一行表明RING的实现被从后续的DSL表达式xsq1中抽象了出来. 从技术上讲, 一个tagless-final的DSL表达式采用函子 (functor) 的形式, 以某个RING实现——任意实现——作为参数. 除let xsq1 x = ...一行外, 其余各行均为样板代码, 几乎不会改变, 且可以自动化 (例如, 如Yallop (2016)中所示, 使用White et al. (2014)的技术). 然而, 在本教程中我们将不使用语法糖, 以保持内容的广泛可理解性和可移植性.

为了尝试刚才定义的DSL函数, 我们需要RING的一个实现.

第4.3节对于向量进行抽象

特化于各种各样的向量内存布局呼吁着找到一个合适的对于向量的抽象表示. 一个常规的浮点数数组或许可以想成是一个函数int -> float, 当然取值是一些特定的整数. 如果推广这个想法, 我们可以定义(输入)向量为一个函数, 其取一个索引, 返回这个索引位置的值; 向量也应该携带关于向量定义域范围的最小和最大索引的信息.

type ('i,'a) vec = Vec of 'i * ('i -> 'a)

索引

ι

被假定是一个可枚举的类型. 为了简单起见, 我们默认索引范围的下界是零值, 而其只需携带上界的信息: 索引范围的长度. 这种vec有时也被称为pull vector. 数组是可变的数据结构. 为了表示数组赋值, 我们引入输出向量:

type ('i,'a,'w) ovec = OVec of 'i * ('i -> 'a -> 'w)

其是向量长度和变动函数 (mutation function) 的序对. 后者取一个索引和一个存储于该索引位置的值, 然后返回更新过了的向量的某种表示. 对于破坏性的复制, 结果

ω

可以是unit.

为了获得对于抽象向量更为真切的感受,

第4.4节向量算术DSL

第4.5节编译向量DSL

第5章线性代数DSL: 矩阵-向量运算和模块化优化

程序生成的价值在于: 从高层次的显然正确的规范出发, 通过让领域专家能够方便地在高抽象层次上编写和应用优化pass, 从而生成高效的代码.

以高性能计算领域向staging提出的第一个湘南挑战 (Aktemur et al., 2013) 为例, 我们证明了这一价值得以实现. 我们使用领域专家所熟悉的类Matlab记法编写矩阵-向量乘法, 最终得到了可确保类型正确且针对给定的 (大部分为稀疏的) 矩阵进行了最优特化的代码. 由此, 我们扩展了前一节所介绍的程序特化的系统化方法, 不仅处理特化问题, 还处理模块化的优化.

MetaOCaml为了维持静态保证, 不允许我们检查甚至重新排列所生成的代码. 尽管如此, 优化仍然是可行的; 而且这些优化是模块化的, 并且可确保其合理性. 它们在远高于抽象语法树或源代码字符串上的搜索替换的层次上执行. 这些优化以独立的小模块形式开发, 逐个混入, 每一小步之后都可以检查其效果——并且直接应用于原始代码, 无需重写甚至标注. 它们是有类型的且保持类型不变, 从而确保优化后的程序仍然是类型正确的. 尽管类型保持本身并不能保证完全的正确性, 但它保证了合理性. 当代码的优化者是高级领域专家而非专业编译器开发者时, 防止用户搬起石头砸自己的脚是明智之举.

最终成果是一个嵌入在MetaOCaml中的小型领域特定语言, 它允许我们编写高层次的基本线性代数表达式, 并应用若干轮领域特定的优化, 例如选择性循环展开, 稀疏矩阵表示以及代数化简.

第5.1节湘南挑战1

这一挑战来源于生物信息学中的隐马尔可夫模型 (HMM) 计算, 其目标是获得矩阵-向量乘法 $v_{i}^{'} = \sum_{k = 0}^{n - 1} a_{i, k} v_{k}$ 的高效代码, 其中矩阵 $a$ 是静态已知的, 并且几乎稀疏, 例如: $[\begin{matrix} 0.5 & 0. & 0. & 0.5 & 0. \\ 0. & 0. & 1. & 0. & 0. \\ 0. & 0. & 0.2 & 0.3 & 0.5 \end{matrix}]$ 在HMM中, 矩阵的元素是转移概率. (严格来说, 湘南挑战 (Aktemur et al., 2013) 将问题简化为了二值情形: $a_{i, j}$ 的所有元素要么为零要么为一, 即它是一个邻接矩阵.) 矩阵的部分行 (而非全部) 是稀疏的. 在生物信息学的HMM中, 行数和列数可高达 $30,000$ .

挑战在于:

构建一个DSL以将问题简洁而清晰地表达为诸如vout := a * v;
生成针对静态已知矩阵a进行特化的高效代码. 该代码应经过代数化简, 即不包含形如0 * e, 1 * e, 0 + e等的表达式. 矩阵行上的循环必须完全展开; 计算v与矩阵某一行内积的内层循环, 仅当该行足够稀疏 (非零元素个数少于某一固定阈值) 时才予以展开. 对于上述示例矩阵, 假设稀疏阈值为 $3$ , 期望的代码应如下所示:

主要的挑战在于以模块化的方式添加各项优化——外层循环展开, 选择性内层循环展开, 代数化简. 正是这种模块化使得下面的解决方案区别于 (Aktemur et al., 2013) 中所描述的那些方案.

第5.2节 BLAS 2 DSL

第6章从解释器到编译器: 图像操作DSL

staging的一个常见应用是构建DSL编译器. 首先, 将DSL实现为一个解释器; 对其进行staging即可得到一个编译器. 一个计算密集型的图像处理DSL显然能从编译中获益. 我们的教学问题是Herrmann和Langhammer (2006) 中所描述的图像处理DSL编译器的一个非常简化的版本.

第6.1节图像处理DSL

我们的灰度图像处理DSL相当精简, 常规且无类型. 它由以下模块签名定义, 以tagless-final风格嵌入OCaml中 (代码位于imgdsl.mli):

`v: float cmplx array`	$\underset{\underset{v_{0}}{⏟}}{\begin{matrix} v_{0} . R & v_{0} . I \end{matrix}} \underset{\underset{v_{1}}{⏟}}{\begin{matrix} v_{1} . R & v_{1} . I \end{matrix}} \dots$
`v': float array cmplx`	$\underset{\underset{v^{'} . R}{⏟}}{\begin{matrix} v_{0}^{'} . R & v_{1}^{'} . R & \dots \end{matrix}} \underset{\underset{v^{'} . I}{⏟}}{\begin{matrix} v_{0}^{'} . I & v_{1}^{'} . I & \dots \end{matrix}}$

协调抽象和高性能: MetaOCaml方法

摘要

第1章 引论

第1.1节 为什么元编程?

第1.2节 为什么这个教程?

第1.3节 为什么MetaOCaml

第1.4节 概览

第1.5节 获得MetaOCaml

第2章 第一步

第2.1节 Now or later

第2.2节 幂

第2.3节 离线代码生成

第2.4节 运行时特化与基准测试

第2.5节 回顾

第2.6节 历史插曲

第3章 滤波

第4章 线性代数DSL: 复向量运算和数据布局

第4.1节 数据布局问题

第4.2节 抽象算术

第4.3节 对于向量进行抽象

第4.4节 向量算术DSL

第4.5节 编译向量DSL

第5章 线性代数DSL: 矩阵-向量运算和模块化优化

第5.1节 湘南挑战1