Julia 中的快速数值计算

每天都在处理数值问题，我一直梦想着有一种语言能够提供优雅的接口，同时允许我编写在大型数据集上运行速度极快的代码。Julia 是一种将这个梦想变成现实的语言。使用 Julia，您可以专注于您的问题，保持代码简洁，更重要的是，即使在性能至关重要的场合，也可以编写快速代码，而无需深入到 C 或 Fortran 等低级语言。

但是，您不应该认为这种潜在的速度是理所当然的。为了使您的代码快速运行，您应该牢记性能并遵循一般的最佳实践指南。在这里，我想与您分享我在编写高效数值计算代码方面的经验。

首先，确保代码正确

与任何语言一样，在实现算法时，首要目标是*确保其正确性*。无论算法运行速度有多快，如果不能正确工作，它都是无用的。在必要时，总是可以随后优化代码。当存在解决问题的不同方法时，您应该选择*渐进效率更高*的方法。例如，在对中等大小的数组进行排序时，未经优化的快速排序实现很容易击败经过精心优化的冒泡排序。但是，对于特定的算法选择，仔细实现它并遵守常见的性能指南仍然可以显著提高性能——我将在后续部分重点介绍这一点。

取消向量化表达式

其他高级语言（如 MATLAB® 或 Python）的用户通常建议尽可能*向量化*他们的代码以提高性能，因为在这些语言中循环速度很慢。另一方面，在 Julia 中，循环可以运行得与用 C 编写的循环一样快，您不再需要依靠向量化来提高速度。实际上，将向量化表达式转换为循环（我们称之为*取消向量化*）通常会带来更高的性能。

考虑以下示例

r = exp(-abs(x-y))

非常简单的表达式，对吧？但是，在幕后，它需要很多步骤和临时数组才能获得此表达式的结果。以下是一系列构建临时数组的操作，用于计算上述表达式

n = length(x)

tmp1 = Array(Float64, n)
for i = 1:n
    tmp1[i] = x[i]-y[i]
end

tmp2 = Array(Float64, n)
for i = 1:n
    tmp2[i] = abs(tmp1[i])
end

tmp3 = Array(Float64, n)
for i = 1:n
    tmp3[i] = -tmp2[i]
end

r = Array(Float64, n)
for i = 1:n
    r[i] = exp(tmp3[i])
end

我们可以看到，此过程创建了三个临时数组，并且需要四次遍历才能完成计算。这引入了明显的开销

为临时数组分配内存需要时间；
在垃圾回收期间回收这些数组的内存需要时间；
遍历内存需要时间——通常，遍历次数越少，效率越高。

这种开销在实践中是显著的，通常会导致速度降低 2 倍到 3 倍。为了获得最佳性能，应该像这样*取消向量化*此代码

r = similar(x)
for i = 1:length(x)
    r[i] = exp(-abs(x[i]-y[i]))
end

此版本在一遍遍历中完成计算，而无需引入任何临时数组。此外，如果 r 已预先分配，则甚至可以省略创建 r 的语句。Devectorize.jl 包提供了一个宏 @devec，可以自动将向量化表达式转换为循环

using Devectorize

@devec r = exp(-abs(x-y))

推导式语法也为取消向量化的计算提供了一种简洁的语法

r = [exp(-abs(x[i]-y[i])) for i = 1:length(x)]

请注意，推导式始终创建新的数组来存储结果。因此，要将结果写入预先分配的数组，您仍然必须手动取消向量化计算或使用 @devec 宏。

将计算合并到单个循环中

遍历数组，特别是大型数组，可能会导致缓存未命中，甚至页面错误，这两者都可能导致明显的延迟。因此，希望尽可能减少往返内存的次数。例如，您可以使用一个循环计算多个映射

for i = 1:length(x)
    a[i] = x[i] + y[i]
    b[i] = x[i] - y[i]
end

这通常比编写 a = x + y; b = x - y 更快。

以下示例展示了如何有效地计算数据集上的多个统计信息（例如，和、最大值和最小值）。

n = length(x)
rsum = rmax = rmin = x[1]
for i = 2:n
    xi = x[i]
    rsum += xi
    if xi > rmax
        rmax = xi
    elseif xi < rmin
        rmin = xi
    end
end

编写缓存友好型代码

现代计算机系统具有复杂的异构内存结构，它结合了寄存器、多级缓存和 RAM。数据通过缓存层次结构访问——一个更小、速度更快的内存，它存储常用数据的副本。

大多数系统不提供直接控制缓存系统的方法。但是，如果您编写*缓存友好型*代码，则可以采取措施使自动缓存管理系统更容易帮助您。一般来说，您不必了解缓存系统的工作原理的所有细节。遵守以下简单规则通常就足够了

以与数据在内存中驻留方式类似的模式访问数据——不要在内存中非连续的位置之间跳跃。

这有时被称为*局部性原理*。例如，如果 x 是一个连续数组，那么在读取 x[i] 后，x[i+1] 更有可能已经在缓存中，而不是 x[i+1000000]，在这种情况下，访问 x[i+1] 将比访问 x[i+1000000] *快得多*。

Julia 数组以列主序存储，这意味着列的行是连续的，但行的列通常不是连续的。因此，通常按列访问数据比按行访问数据效率更高。考虑计算矩阵中每一行的和的问题。自然地，可以按如下方式实现它

m, n = size(a)
r = Array(Float64, m)

for i = 1:m
    s = 0.
    for j = 1:n
        s += a[i,j]
    end
    r[i] = s
end

此处的循环按行访问元素，如 a[i,1], a[i,2], ..., a[i,n]。这些元素之间的间隔是 m。直观地说，它在每个内循环中以长度为 m 的步长从每一行的开头跳到末尾，然后跳回到下一行的开头。这效率不高，尤其是在 m 较大的情况下。

可以通过更改计算顺序使此过程更加缓存友好

for i = 1:m
    r[i] = a[i,1]
end

for j = 2:n, i = 1:m
    r[i] += a[i,j]
end

一些基准测试表明，对于大型矩阵，此版本的速度可以比上述版本快*5 到 10 倍*。

避免在循环中创建数组

创建数组需要内存分配，并增加垃圾回收器的负载。重用同一个数组是降低内存管理成本的好方法。

在迭代算法中，您可能需要更新数组的情况并不少见。例如，在 K 均值算法中，您可能希望在每次迭代中更新聚类均值和距离。执行此操作的一种简单方法可能是

while !converged && t < maxiter
    means = compute_means(x, labels)
    dists = compute_distances(x, means)
    labels = assign_labels(dists)
    ...
end

在此 K 均值算法的实现中，数组 means、dists 和 labels 在每次迭代中都会重新创建。每次步骤的这种内存重新分配是不必要的。这些数组的大小是固定的，并且它们的存储可以在迭代之间重用。以下替代代码是实现相同算法的更有效方法

d, n = size(x)

# pre-allocate storage
means = Array(Float64, d, K)
dists = Array(Float64, K, n)
labels = Array(Int, n)

while !converged && t < maxiter
    update_means!(means, x, labels)
    update_distances!(dists, x, means)
    update_labels!(labels, dists)
    ...
end

在此版本中，循环中调用的函数就地更新预先分配的数组。

如果您正在编写包，建议您为每个输出数组的函数提供两个版本：一个就地执行更新，另一个返回一个新数组。前者通常可以实现为后者的轻量级包装器，在修改输入数组之前复制它。一个很好的例子是 Distributions.jl 包，它同时提供 logpdf 和 logpdf!，以便在需要新数组时可以编写 lp = logpdf(d,x)，或者在 lp 已预先分配时可以编写 logpdf!(lp,d,x)。

识别使用 BLAS 的机会

Julia 包装了大量用于线性代数计算的 BLAS 例程。这些例程是世界顶级数值计算专家多年研究和优化的成果。因此，在可能的情况下使用它们可以带来几乎神奇的性能提升——BLAS 例程通常比它们替换的简单循环实现快几个数量级。

例如，考虑按如下方式累加向量的加权版本

r = zeros(size(x,1))
for j = 1:size(x,2)
    r += x[:,j] * w[j]
end

您可以将语句 r += x[:,j] * w[j] 替换为对 BLAS axpy! 函数的调用以获得更好的性能

for j = 1:size(x,2)
    axpy!(w[j], x[:,j], r)
end

但是，这仍然远非最佳。如果您熟悉线性代数，您可能已经发现这只是矩阵向量乘法，可以写成 r = x * w，它不仅比上述任何一个循环都更短、更简单、更清晰——而且它也比这两个版本都运行得更快。

我们的下一个示例是将 BLAS 例程更巧妙地应用于计算两个矩阵中列之间的成对欧几里得距离。下面是一个直接计算成对距离的简单实现

m, n = size(a)
r = Array(Float64, m, n)

for j = 1:n, i = 1:m
    r[i,j] = sqrt(sum(abs2(a[:,i] - b[:,j])))
end

这显然不是最佳的——在评估内循环中的表达式时会创建许多临时数组。为了加快速度，我们可以取消向量化内部表达式

d, m = size(a)
n = size(b,2)
r = Array(Float64, m, n)

for j = 1:n, i = 1:m
        s = 0.
        for k = 1:d
            s += abs2(a[k,i] - b[k,j])
        end
        r[i,j] = sqrt(s)
    end
end

此版本比向量化形式的性能要好得多。但这是我们所能做到的最好的吗？通过采用替代策略，我们可以编写一个更快的方法来计算成对距离。诀窍在于，两个向量之间的平方欧几里得距离可以扩展为

sum(abs2(x-y)) == sum(abs2(x)) + sum(abs2(y)) - 2*dot(x,y)

如果我们分别计算这三个项，则计算可以完美地映射到 BLAS 例程。下面，我们有一个使用仅 BLAS 例程编写的成对距离的新实现，包括由 NumericExtensions.jl 包包装的范数调用

using NumericExtensions   # for sqsum
using Base.LinAlg.BLAS    # for gemm!

m, n = size(a)

sa = sqsum(a, 1)   # sum(abs2(x)) for each column in a
sb = sqsum(b, 1)   # sum(abs2(y)) for each column in b

r = sa .+ reshape(sb, 1, n)          # first two terms
gemm!('T', 'N', -2.0, a, b, 1.0, r)  # add (-2.0) * a' * b to r

for i = 1:length(r)
    r[i] = sqrt(r[i])
end

此版本的速度比我们的原始实现快*100 倍以上*——BLAS 中的 gemm 函数在过去几十年里已被许多天才的开发人员和工程师进行了极致优化。

我们应该提到，如果您确实想要计算成对距离，则不必自己实现它：Distance.jl 包提供了各种距离度量的优化实现，包括此度量。我们介绍了这个优化技巧，以说明通过编写尽可能使用 BLAS 例程的代码可以实现的巨大性能提升。

探索可用的包

Julia 拥有一个非常活跃的开源生态系统。已经开发了各种包，提供了用于高性能计算的优化算法。在决定自己动手编写之前，先寻找一个满足您需求的包——如果您找不到您需要的包，请考虑贡献一个！以下是一些可能对那些对高性能计算感兴趣的人有用的包

NumericExtensions.jl – Julia 基础功能的扩展，为各种常用计算提供高性能支持（其中许多将逐渐迁移到基础 Julia 中）。
Devectorize.jl – 用于取消向量化向量表达式的宏和函数。使用此包，用户可以以高级向量化方式编写计算，同时享受手动编码的取消向量化循环的高运行时性能。

查看Julia 包列表以获取更多包。Julia 还自带一个采样分析器来衡量代码在哪里花费了大部分时间。如有疑问，请测量，不要猜测！

2013 年 9 月 4 日 | 林达华