Julia code optimization : is this the time to use SIMD?

user_15519366 0
Viewed 997

I am struggling to optimize my Julia code and make it run faster.

I abstracted a part of my whole code and I want you to evaluate if there are some bottlenecks which make the process slower than that of optimized one.

A brief explanation of the code:

  • Let array1 n x n array ,which is initialize with random entries with randn
  • array1 would be updated under a certain algorithm in the loop until all the entries of array1 would be bigger than -0.8
  • If the loop does not end within 100000 loops, the loop ends forcibly.

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()

I think I can save time using @simd at some for statement.

Is there better way of writing or superior algorithm? If you have any useful information other than SIMD, I would be happy to hear it.

Any information would be appreciated.

julia simd optimization arrays algorithm
1 Answers

Yes, especially using LoopVectorization.jl -- though there are a few other optimizations we have to do first before we're in the ballpark where this is the limiting factor.

First, some initial timings of your code on my system (an oldish laptop) using 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.

There seems to be some pronounced bimodality between cases where the loop short circuits in only a few tens of evaluations vs cases where it runs the full 100000. As an aside, this is a classic example of why despite some conventional wisdom, minimum times are not a great metric for benchmarking, and either means or (even better) full histograms are much better -- so the histograms from BenchmarkTools.jl or BenchmarkHistograms.jl are great.

Let's start with a few general optimizations (reducing unnecessary allocations, etc.)

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

As you can see, I moved the allocations for bar and adding outside of the loop -- since they don't change size, we can just allocate them once and then just refill with zeros if necessary. Another change involves unnecessary allocation on assignment when you are adding and subtracting arrays. The original array1 = array1 - bar + adding for example is allocating, but @. array1 = array1 - bar + adding, which is just a convenient shorthand for array1 .= array1 .- bar .+ adding, makes explicit that we want element-wise operations, and will ensure that they all occur in-place (i.e., without triggering new heap allocations). I also put an @inbounds on the for loops and refactored the summation in the innermost loop a bit (while sum is fast, allocating a brand new array for it to add up on each iteration of the loop is not, as mentioned in the comments).

This seems to give us about a 10x mean speedup, and substantially reduces our allocations and memory usage.

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.

Now for the SIMD. It turns out that just adding @inbounds @simd in front of the for loops doesn't appear from the timings to enable any optimizations that Julia's compiler isn't already finding. LoopVectorization's @turbo does, however:

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

Note that you can use @turbo to SIMD-vectorize both for loops and @. broadcasts, as seen here.

This gives us about another 1.5x speedup, for a total of nearly 15x over the original implementation

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.

The advantage for LoopVectorization.jl would likely be even greater on a system with AVX512 vector registers (I only have AVX2 on this laptop).

user_136397340
Related