arrays - Julia 代码优化 : is this the time to use SIMD?

Question

我正在努力优化我的 Julia 代码并使其运行得更快。
我抽象了我的整个代码的一部分，我希望你评估是否存在一些瓶颈，使过程比优化的过程慢。
代码的简要说明:

让 array1 n x n 数组，用 randn 随机条目初始化

array1将在循环中的某个算法下更新，直到array1的所有条目将大于 -0.8

如果循环未在 100000 次循环内结束，则循环强制结束。

using Random;

function Test()
  
    n=5
    sqn = n^2
    foo1 = circshift(collect(1:n)', (0, 1))
    foo2 = circshift(collect(1:n)', (0, -1))
    foo3 = circshift(collect(1:n), 1)
    foo4 = circshift(collect(1:n), -1)


    # initialize array1 with random entries
    array1 = randn(n,n);
    # initialize array2 with zeros
    array2 = zeros(n,n);

    #println(array1);

    loop = 0
  
    while minimum(array1) < -0.8 && loop < 100000
                loop += 1
                bar = zeros(n, n)
    
                # Use simd?
                for i = 1:sqn
                    bar[i] = max(0, array1[i] - 0.2)
                end
    
                array2 += bar
                
                adding = zeros(n, n)
                # Use simd?
                for i = 1:n
                    for j = 1:n
                        adding[i, j] =
                            (1 / 4) * sum([
                                bar[i, foo1[j]],
                                bar[i, foo2[j]],
                                bar[foo3[i], j],
                                bar[foo4[i], j],
                            ])
                    end
                end
                array1 = array1 - bar + adding
            end
   # println(array1)
   # println(array2)
   # println(loop)

end
Test()

我想我可以使用 @simd 来节省时间在一些 for statement .
有没有更好的写作方式或更好的算法？
如果您有除 SIMD以外的任何有用信息，我会很高兴听到它。
任何信息，将不胜感激。

Answer 1

是的，尤其是使用 LoopVectorization.jl -- 尽管还有一些其他优化我们必须先做，然后才能进入球场，这是限制因素。
首先，使用 BenchmarkTools.jl 在我的系统(一台旧笔记本电脑)上的代码的一些初始计时

julia> using BenchmarkTools

julia> @benchmark Test()
BechmarkTools.Trial: 31 samples with 1 evaluations.
 Range (min … max):    5.987 μs … 257.525 ms  ┊ GC (min … max):  0.00% … 22.07%
 Time  (median):     234.396 ms               ┊ GC (median):    16.67%
 Time  (mean ± σ):   162.310 ms ± 114.023 ms  ┊ GC (mean ± σ):  18.44% ±  8.84%

  █                                                        ▂     
  █▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▅▅█▅██▁▃ ▁
  5.99 μs          Histogram: frequency by time          258 ms <

 Memory estimate: 9.84 KiB, allocs estimate: 70.

循环仅在几十次评估中短路的情况与运行完整 100000 次的情况之间似乎存在一些明显的双峰性。顺便说一句，这是一个典型的例子，说明为什么尽管有一些传统智慧，最小次数不是一个很好的基准测试指标，意味着或(甚至更好)完整直方图要好得多——所以来自 BenchmarkTools.jl 的直方图或 BenchmarkHistograms.jl很棒。
让我们从一些常规优化开始(减少不必要的分配等)

function test()

    n=5
    sqn = n^2
    r = collect(1:n)
    foo1 = circshift(r', (0, 1))
    foo2 = circshift(r', (0, -1))
    foo3 = circshift(r, 1)
    foo4 = circshift(r, -1)


    # initialize array1 with random entries
    array1 = randn(n,n);
    # initialize array2 with zeros
    array2 = zeros(n,n);
    # Initialize temporary arrays used in loops
    bar = zeros(n, n)
    adding = zeros(n, n)

    loop = 0
    while minimum(array1) < -0.8 && loop < 100000
                loop += 1
                fill!(bar, 0)

                # Use simd?
                @inbounds for i = 1:sqn
                    bar[i] = max(0, array1[i] - 0.2)
                end

                @. array2 += bar

                fill!(adding, 0)
                # Use simd?
                @inbounds for i = 1:n
                    for j = 1:n
                        s = bar[i, foo1[j]] + bar[i, foo2[j]] + bar[foo3[i], j] + bar[foo4[i], j]
                        adding[i, j] = 0.25 * s
                    end
                end
                @. array1 = array1 - bar + adding
            end
   return array1, array2, loop
end

如您所见，我移动了 bar 的分配。和 adding在循环之外——因为它们不会改变大小，我们可以只分配它们一次，然后在必要时重新填充零。另一个更改涉及在添加和减去数组时不必要的分配分配。原版 array1 = array1 - bar + adding例如正在分配，但 @. array1 = array1 - bar + adding ，这只是 array1 .= array1 .- bar .+ adding 的便捷简写, 明确表示我们需要逐元素操作，并将确保它们都就地发生(即，不会触发新的堆分配)。我也放了一个 @inbounds在 for循环并稍微重构最内层循环中的求和(虽然 sum 很快，但不是如评论中所述，分配一个全新的数组以在循环的每次迭代中相加)。
这似乎给我们带来了大约 10 倍的平均加速，并大大减少了我们的分配和内存使用。

julia> @benchmark test()
BechmarkTools.Trial: 305 samples with 1 evaluations.
 Range (min … max):   1.622 μs … 30.093 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     22.438 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):   16.429 ms ± 11.530 ms  ┊ GC (mean ± σ):  0.00% ± 0.00%

  █                                           ▁                
  █▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▇▄▃▃▄▃▅▅▂▃▃▂▃▂ ▂
  1.62 μs         Histogram: frequency by time        29.4 ms <

 Memory estimate: 1.75 KiB, allocs estimate: 9.

现在是 SIMD。原来只是加上 @inbounds @simd在 for 循环前面没有出现在启用 Julia 编译器尚未找到的任何优化的计时中。 LoopVectorization的 @turbo但是:

using Random, LoopVectorization

function test_simd_turbo()

    n=5
    sqn = n^2
    r = collect(1:n)
    foo1 = circshift(r', (0, 1))
    foo2 = circshift(r', (0, -1))
    foo3 = circshift(r, 1)
    foo4 = circshift(r, -1)


    # initialize array1 with random entries
    array1 = randn(n,n);
    # initialize array2 with zeros
    array2 = zeros(n,n);
    # Initialize temporary arrays used in loops
    bar = zeros(n, n)
    adding = zeros(n, n)

    loop = 0
    while minimum(array1) < -0.8 && loop < 100000
                loop += 1
                fill!(bar, 0)

                # Use simd?
                @turbo for i = 1:sqn
                    bar[i] = max(0, array1[i] - 0.2)
                end

                @turbo @. array2 += bar

                fill!(adding, 0)
                # Use simd?
                @turbo for i = 1:n
                    for j = 1:n
                        s = bar[i, foo1[j]] + bar[i, foo2[j]] + bar[foo3[i], j] + bar[foo4[i], j]
                        adding[i, j] = 0.25 * s
                    end
                end
                @turbo @. array1 = array1 - bar + adding
            end
   return array1, array2, loop
end

请注意，您可以使用 @turbo对两者进行 SIMD 向量化 for循环和 @.广播，如这里所见。
这为我们提供了大约 1.5 倍的加速，总共比原始实现快了 15 倍

julia> @benchmark test_simd_turbo()
BechmarkTools.Trial: 419 samples with 1 evaluations.
 Range (min … max):   1.864 μs … 28.059 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     15.486 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):   11.946 ms ±  8.034 ms  ┊ GC (mean ± σ):  0.00% ± 0.00%

  █                                                            
  █▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▇▇▆▅▄▄▃▄▄▄▃▃▂▁▂▁▁▁▂▁▁▁▁▁▁▃ ▂
  1.86 μs         Histogram: frequency by time        26.3 ms <

 Memory estimate: 1.75 KiB, allocs estimate: 9.

LoopVectorization.jl 的优势在具有 AVX512 的系统上可能会更大。矢量寄存器(我在这台笔记本电脑上只有 AVX2)。