联合拆分:它是啥,以及你为什么应该关心

2018年8月9日 | Tim Holy

在密切关注 Julia 开发的人群中,有一项(众多)备受期待的新功能叫做“联合拆分”。在 2018 年的 JuliaCon 上,我发现自己反复解释这个特性,因此我决定撰写这篇博文,以便更广泛地传播这些重要信息。首先我要说,我不是这方面的专家——这项功能是由 Jameson Nash 和 Jacob Quinn 添加的,并由 Keno Fisher 的优化器改进进行了增强——但我确实是众多对这项功能感到兴奋的人之一,因为它已经改变了我编写 Julia 代码的方式。

背景是这样的:在过去,你必须非常小心地确保你编写的几乎每个函数都返回可预测的类型。经验丰富的 Julia 程序员经常使用一个工具 @code_warntype 来检查代码是否存在令人恐惧的“类型不稳定性”。这种代码返回(由 Julia 的推理引擎确定)类型为 Any 或类型为 Union{Type1, Type2, ...} 的对象。前者意味着推理引擎无法对返回类型做出任何具体陈述;后者意味着推理引擎能够确定可能的返回类型的一个特定列表。不幸的是,编译器在利用这些部分知识方面并不是特别擅长,因此在实践中,这两种结果都预示着你的代码性能会非常糟糕。

快进到 0.7 和 1.0 版本,情况既“相同”又“完全不同”。我的意思是,Any 仍然表示存在潜在问题,因为编译器无法对代码进行任何优化。但是,通常情况下,Union{Type1, Type2, ...} 没什么好担心的,因为它几乎不会造成任何性能损失。

这种魔法是如何运作的?简单来说,假设你在一个函数内部有一段代码块,如下所示

ret1 = function1(args...)
ret2 = function2(ret1, ...)
ret3 = function3(ret1, ret2, ...)
...

假设 ret1 可以是两种类型之一,AB(即 Union{A,B})。在旧版本的 Julia 中,会发生以下情况:从 function1 开始,编译器会说“我无法确定应该使用 function2 的哪种方法”。因此,它不会对后续的任何代码进行专门化;相反,每次执行此代码块时,它都会获取 ret1 的实际类型,并开始遍历方法表,执行类型交集以尝试找到 function2 的适用编译版本。类型交集涉及的计算虽然经过了很好的优化,但仍然非常耗费资源,因此“方法查找”步骤相当缓慢(特别是对于具有数十或数百种方法的函数)。

在 Julia 0.7 和 1.0 中,编译器会做一些完全不同的事情:它会自动(无需你做任何努力)将上面的代码块编译成类似于以下内容:

ret1 = function1(args...)    # ret1 isa Union{A,B}
if ret1 isa A
    ret2 = function2_specialized_for_A(ret1, ...)
    ret3 = function3_specialized_for_A(ret1, ret2, ...)
    ...
else
    ret2 = function2_specialized_for_B(ret1, ...)
    ret3 = function3_specialized_for_B(ret1, ret2, ...)
    ...
end

这里的区别非常大。虽然 Julia 无法提前知道 ret1 的精确类型,但在第一个代码块中它绝对是 A 类型(因为它已经检查过了),在第二个代码块中它绝对是 B 类型(因为这是唯一其他选项)。因此,Julia 可以**在编译时而不是运行时查找 function2function3 的适当编译方法**,这使得它在实际运行时速度非常快。

现在,我听到一些人说,“但是那里有一个分支,与许多其他 CPU 指令相比,分支速度较慢”。的确如此。但是单个分支与方法查找相比几乎可以忽略不计;此外,在使用联合拆分的情况下,通常情况下你无论如何都需要该分支。在这种情况下,成本实际上为零。

为了说明原因,请考虑对数组 A 执行 findfirst(isequal(7), A) 操作,该操作以前总是返回一个整数,表示在 A 中找到值 7 的第一个索引。出现了一个有问题的案例:如果 A 不包含任何 7 会怎样?以前,我们返回 0,接收方必须检查 if ret1 == 0 以确定是否需要将执行转移到错误处理代码。因此,在编写正确的代码时,无法避免需要该分支。更糟糕的是,如果你忘记检查,并且当你向 function2 传递 0 时它没有报错,那么你很可能会得到一个毫无意义的答案。这比得到错误要糟糕得多,因为跟踪错误结果的来源要困难得多。

在 Julia 0.7 和 1.0 中,Milan Bouchet-Valat 重写了我们所有的 find* 函数,其中一项更改(在众多更改中)是 findfirst 现在在找不到所需值时返回 nothing。与返回 0 的旧方法不同,此返回值是真正明确的,并且对索引的泛化(其中 0 可能是完全有效的数组索引)具有鲁棒性。它还为你提供了更可靠的代码,因为如果你忘记检查,那么你实际上无法对 nothing 做任何事情而不会触发(非常受欢迎的)错误。并且由于联合拆分,它不会造成任何性能损失。

因此,虽然联合拆分最初听起来像是只对编译器专家感兴趣的某种神秘功能,但实际上它改变了你应该如何设计代码,并且它允许创建更易于理解和更强大的软件。这是每个人都能欣赏的功能。