Julia 并行计算 for 循环

Question

我想从一个大的上三角矩阵计算元素的总和。常规的 Julia 代码如下。

function upsum(M); n = size(M)[1]; sum = 0
    for i = 1:n-1 for j = i+1:n
        sum = sum + M[i,j]
        end
    end
    return sum
end

R = randn(10000,10000)
upsum(R)

由于矩阵很大，我想知道有没有办法提高速度。如何在这里使用并行计算？

Answer 1

在这种情况下，我会使用线程而不是并行处理。这是一个示例代码：

using Base.Threads

function upsum_threads(M)
    n = size(M, 1)
    chunks = nthreads()
    sums = zeros(eltype(M), chunks)
    chunkend = round.(Int, n * sqrt.((1:chunks) ./ chunks))
    @assert minimum(diff(chunkend)) > 0
    chunkstart = [2; chunkend[1:end-1] .+ 1]
    @threads for job in 1:chunks
        s = zero(eltype(M))
        for i in chunkstart[job]:chunkend[job]
            @simd for j in 1:(i-1)
                @inbounds s += M[j, i]
            end
        end
        sums[job] = s
    end
    return sum(sums)
end

R = randn(10000,10000)
upsum_threads(R)

它应该会给你一个显着的加速（即使你删除 @threads 它应该会快得多）。

您可以通过设置 JULIA_NUM_THREADS 环境变量来选择 Julia 使用的线程数。