可扩展广播融合

Julia 版本 0.7 带来了一个令人兴奋的新功能：定制广播融合的能力！这个最近合并的更改是长期迭代设计过程的顶峰，该过程涉及社区中的许多成员。我们已经收敛到一个高度可扩展的接口，应该满足许多用例。在这篇博文中，我将通过回顾一些关键特性来解释为什么这是一个大事，并且只是触及了这种新设计所有可能性的表面。我非常确定我们进取的社区将来会想出更多巧妙的方法来利用这种新结构。

广播是 Julia 的核心功能：它允许你通过用 . 注释运算符和函数调用，以简洁高效的方式表达对容器和标量的逐元素操作。在大小不匹配的情况下，广播将通过重复它们来填充外部形状，虚拟扩展缺失的维度或“单例”维度（仅包含一个值）。例如，表达式 ([1, 2, 3] .+ [10 20 30 40]) ./ 10 将一个 3 元素列向量、一个 1x4 矩阵和一个标量组合起来，计算出一个 3x4 的结果。我把它想象成“挤出”向量穿过一行的矩阵的列，并将除以 10 分布到整个结果中

julia> ([1, 2, 3] .+ [10 20 30 40]) ./ 10
3×4 Array{Float64,2}:
 1.1  2.1  3.1  4.1
 1.2  2.2  3.2  4.2
 1.3  2.3  3.3  4.3

从版本 0.6 开始，Julia 通过 "融合" 两个操作到单个内核中来执行此表达式。也就是说，Julia 不是先构建一个由加法产生的整数矩阵（[11 21 31 41; 12 22 ...]），然后用第二个循环将每个元素除以 10，而是同时执行加法和除法，只对输出数组进行一次遍历，并完全跳过中间数组。这种融合优化作为语法级转换发生，因此它保证会发生并且易于推理。版本 0.7 在此功能之上添加了一个可扩展的 API，允许数组精确地定制此操作方式。

Julia 现在使用一个一等公民的数据结构来“延迟”地表示一个融合的广播表达式，然后再执行它。如果你不是包开发人员，这对你来说可能没有多大意义，但你仍然可以收获许多好处。仅在 Base Julia 及其标准库中，这意味着

BitArrays 可以识别它们可以一次对 64 个布尔元素进行操作的情况，从而产生巨大的性能提升——通常是两个数量级或更多！为了便于举例，我将只使用随机数据，使用一个简单的 A 而不是 B 谓词

julia> using BenchmarkTools, Random
       srand(0)
       A = bitrand(10^6)
       B = bitrand(10^6)
       @benchmark $A .& .!$B
BenchmarkTools.Trial:
  memory estimate:  122.23 KiB
  allocs estimate:  3
  ----------
  minimum time:     7.891 μs (0.00% GC)
  median time:      13.152 μs (0.00% GC)
  mean time:        17.826 μs (9.62% GC)
  maximum time:     591.497 μs (95.73% GC)
  ----------
  samples:          10000
  evals/sample:     1

与之前在 0.6 上的结果进行比较

BenchmarkTools.Trial:
  memory estimate:  126.45 KiB
  allocs estimate:  6
  ----------
  minimum time:     3.615 ms (0.00% GC)
  median time:      3.741 ms (0.00% GC)
  mean time:        3.764 ms (0.18% GC)
  maximum time:     7.744 ms (50.18% GC)
  ----------
  samples:          1328
  evals/sample:     1

这是一个 450 倍的性能提升。

对范围的广播操作现在可以简单地重新计算一个新的范围，而不是逐元素地进行操作，如果可能的话。例如，表达式 ((1:10000) .+ 20) .* 7 不需要为 10,000 个元素分配一个向量——甚至不需要进行 10,000 次计算。相反，它可以在开始、结束和步长方面进行操作，并返回一个表示结果的新范围：147:7:70140。这个新功能允许他们将 O(N) 计算转化为 O(1)。在 0.6 版本中，范围处于一个奇怪的状态，其中 (1:10000) + 20 实施了新的范围的快速 O(1) 计算，但所有其他数组类型都已弃用对数字的加法，转而支持显式的 .+ 广播，以便更清晰的语义和更高的性能。由于这个新的 API，范围现在可以识别这些情况，并以有效的方式完全支持广播。
结构化矩阵在 LinearAlgebra 标准库中不再返回稀疏数组作为广播操作的结果。它们现在要么保持适当的结构，要么返回一个密集数组。例如

julia> using LinearAlgebra
       d = Diagonal(1:3)
3×3 Diagonal{Int64,UnitRange{Int64}}:
 1  ⋅  ⋅
 ⋅  2  ⋅
 ⋅  ⋅  3

julia> d ./ 10
3×3 Diagonal{Float64,Array{Float64,1}}:
 0.1   ⋅    ⋅
  ⋅   0.2   ⋅
  ⋅    ⋅   0.3

julia> t = d .+ LowerTriangular(rand(3,3))
3×3 LowerTriangular{Float64,Array{Float64,2}}:
 1.5446     ⋅         ⋅
 0.529211  2.26095    ⋅
 0.533674  0.373702  3.88408

julia> t .+ 100
3×3 Array{Float64,2}:
 101.327  100.0    100.0
 100.85   102.138  100.0
 100.575  100.129  103.839

以前，Julia 0.6 会为 d ./ 10 返回 SparseMatrixCSC，为 d .+ LowerTriangular(rand(3,3)) 返回 Array。

最后，全局范围内的广播现在是可以预编译的，你可以在生成的函数中使用点广播。这在过去并不是一个很大的限制，但它确实让那些计时像 @time y .*= 2 这样的就地广播并看到数千次分配的人感到惊讶

# Previously:
julia> y = rand(1000);

julia> @time y .*= 2;
  0.032775 seconds (17.43 k allocations: 947.671 KiB)

julia> @time y .*= 2;
  0.020790 seconds (4.27 k allocations: 233.457 KiB)

# Now on Julia 0.7:
julia> y = rand(1000);

julia> @time y .*= 2;
  0.060494 seconds (90.41 k allocations: 5.529 MiB)

julia> @time y .*= 2;
  0.000020 seconds (6 allocations: 208 bytes)

我现在将深入探讨这个新 API 的工作原理。

融合广播的表示

你可以使用 Meta.@lower 精确地查看融合广播是如何表示的，但用更简单的术语来说，表达式 ([1, 2, 3] .+ [10 20 30 40]) ./ 10 实际上是以下内容的语法转换

julia> using .Broadcast: materialize, broadcasted
       bc = broadcasted(/, broadcasted(+, [1, 2, 3], [10 20 30 40]), 10)
       materialize(bc)
3×4 Array{Float64,2}:
 1.1  2.1  3.1  4.1
 1.2  2.2  3.2  4.2
 1.3  2.3  3.3  4.3

在这种情况下，该 bc 对象是一个 Broadcasted 结构的实例。它只是保存函数及其参数——它的参数可能包含其他嵌套的 Broadcasted 结构。materialize 函数进行一些预处理，然后调用 copy(bc)，该函数分配结果，然后最终遍历结果并执行函数。

沿途的每一步都是可扩展的，利用了 Julia 的多重分派、内联和参数专门化带来的强大功能，从而实现了接近零的开销。有了这个基本框架，你就可以开始了解内置和标准库数组如何实现上述所有新功能

当广播到 BitArray 中时，它可以首先内省表达式树中的函数及其参数，以查看它是否可以在 UInt64 的打包 64 位块级别进行操作，而不是逐位地进行操作。它甚至可以将 ! 之类的仅限于布尔值的运算符转换为 ~ 之类的等效位运算符。
范围可以通过定义专门的 Broadcast.broadcasted 方法来“选择退出”融合，该方法会立即返回那些重新计算的范围。这意味着它们根本不会融合多个操作，但作为交换，它们获得了 O(1) 算法。
当 LinearAlgebra 的结构化矩阵被要求分配结果时，它们的专门 broadcast_similar 方法可以遍历 Broadcasted 表达式树，并识别是否会保留任何结构。
融合的广播表达式不再在你的背后构建一个匿名函数；它们只是构建 Julia 数据结构的新实例，这些实例只是调用已经定义的函数。这就是它们可以预编译并在生成的函数中起作用的原因。

展望未来

当然，所有这些都已记录在案，并可供包使用。我们没有对自己的秘密武器进行囤积。看到包生态系统中众多创意人士如何利用所有这些新功能，将会令人兴奋。我认为这将产生重大效益的一些地方包括机器学习应用和 GPU 上的计算。

历史上，一些机器学习库已经实施了黑客手段来允许它们的数组选择退出融合。新的广播 API 允许它们通过一个一等公民的、完全支持的机制来选择退出融合，这种机制在实现上要简单得多。此外，许多深度学习技术严重依赖于微分；在某些情况下，内省广播内核将允许它们使用精确的解析结果，而不是计算微分。

最后，将复杂的广播表达式融合到单个内核中，已经极大地提高了 GPU 上数组的性能。大多数 GPU 编程包不需要内省或定制广播表达式，但它们正在展望未来，以扩展到允许将reduction 与广播表达式融合在一起。虽然目前还不可行，但大部分机制已经到位，可以对 Broadcasted 延迟包装器进行直接操作，而不是在像 sum(X.^2 .+ Y.^2) 这样的表达式中分配中间数组。

2018 年 5 月 11 日 | Matt Bauman (JuliaHub)

融合广播的表示

展望未来