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.

[j-fu]

Thank you for the proposal - good idea!
However, there are the parallel test examples in tests, they could be used to create benchmarks which are closer to the way things are used in FEM assembly. In particular, they use grid partitioning and partition graph coloring such that write conflicts during assembly are prevented.
Up to now I didn't find the time to make them proper benchmarks.

[tduretz]

OK, I see, thank you. I was proposing this because I used FD in a project. But, yes, FEM-like assemblies are more general and parallel performance on such assemblies is of more general interest.