Best practice: threaded sparse matrix assembly

[j-fu]

Track this Discourse thread.

Link to the MWE

[j-fu]

The MWE has two issues:

• using threadid() to address thread local data is no longer recommended
• While the approach to use an array of matrices formally works, this has at least the following problems:
  • Indeed reduction is slow; there could be ways to sped this up along these lines - I'll think about this
  • Using a vector of ExtendableSparseMatrixCSC wastes space, as they all have the full size and need to manage this.

I will come up with an alternaive MWE later.

[tduretz]

Thanks! I am adding the MWE here for this page to be self-contained:

using Base.Threads, ExtendableSparse, LinearAlgebra

@views function serial_assembly!(M, n)
    for i=1:n
        if i==1
            M[i,i  ] =  1.0  
            M[i,i+1] = -1.0
        elseif i==n
            M[i,i-1] = -1.0  
            M[i,i  ] =  1.0
        else
            M[i,i-1] = -1.0  
            M[i,i  ] =  2.0
            M[i,i+1] = -1.0
        end
    end
end

@views function threaded_assembly!(M, M_loc, n)
    Threads.@threads for i=1:n
        tid = threadid()
        if i==1
            M_loc[tid-1][i,i  ] =  1.0  
            M_loc[tid-1][i,i+1] = -1.0
        elseif i==n
            M_loc[tid-1][i,i-1] = -1.0  
            M_loc[tid-1][i,i  ] =  1.0
        else
            M_loc[tid-1][i,i-1] = -1.0  
            M_loc[tid-1][i,i  ] =  2.0
            M_loc[tid-1][i,i+1] = -1.0
        end
    end
end

let 
    n  = 10_000_000

    # Serial
    M1 = ExtendableSparseMatrix(n,n)
    @time serial_assembly!(M1, n)

    # Threaded - version 1
    M2 = ExtendableSparseMatrix(n,n)
    M2_loc = [ExtendableSparseMatrix(n,n) for i=1:nthreads()]
    @time threaded_assembly!(M2, M2_loc, n)
    # Reduce. Looks like a conversion to sparse is needed to allow for .+
    # This makes things a slightly faster
    @time for k=1:nthreads()
        M2 .= sparse(M2) .+ sparse(M2_loc[k])
    end

    M3 = ExtendableSparseMatrix(n,n)
    M3_loc = [ExtendableSparseMatrix(n,n) for i=1:nthreads()]
    @time threaded_assembly!(M3, M3_loc, n)
    # Reduce. 
    @time begin 
        I = Int[]
        J = Int[]
        V = Float64[]
        for k=1:nthreads()
            I_loc, J_loc, V_loc = findnz(M2_loc[k])
            append!(I, I_loc)
            append!(J, J_loc)
            append!(V, V_loc)
        end
        M3 = sparse(I, J, V)
    end

    # Check
    @show norm(M1 .- M2)
    @show norm(M1 .- M3)
end

[j-fu]

Package version 2.1 fixed an allocation issue and introduced a Base.sum method for ExtendabelSparse matrices. See several modified versions of your MWE in this gist, and the corresponding comments.

Multithreaded performance is tricky and so far I see a significant gain only if assembly per item is expensive. There probably is a way to speed up the reduction by directly creating new columns from all matrices at once, but this is quite hard to realize if one cannot avoid that different submatrices can add to the same entry, as it essentially happens with finite element assembly.

[tduretz]

Would it be worth to extend if I extend the MWE to a 2D case, like a 9-point stencil?
I think it could be more representative of actual work in a production code. Maybe we better appreciate a speed up with such example? I could also throw a heavy function evaluation instead of 2 and -1.