DiscourseMWE.jl

DiscourseMWE.jl

module DiscourseMWE

    using Base.Threads, ExtendableSparse, LinearAlgebra, SparseArrays
    using ChunkSplitters

    #
    # As long as we don't work wit ranges like M[:,i], we do not need @vies
    #
    function serial_assembly!(M, n)
        for i in 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
        #
        # In ExtendableSparse, the flush! method turns the internal linked list structure
        # into an intearnal SparseMatrixCSC structure, this needs to be included in the timing.
        # Typically, this is automatically called within any operation like '\', `sparse` etc.
        flush!(M)
        return
    end

    #
    # Corrected to include flush! which can be run in parallel.
    # Also, we do not need M.
    # More importantly, this method now avoids the use of Threads.threadid().
    # The main point here is that Julia's parallelism is based on Tasks, and the thread
    # a task is running on can change unless we choose :static scheduling.
    function threaded_assembly!(M_loc::Vector{<:AbstractMatrix}, n, nt)
        chunks = ChunkSplitters.chunks(1:n, n = nt)
        Threads.@threads for ichunk in 1:length(chunks)
            for i in chunks[ichunk]
                if i == 1
                    M_loc[ichunk][i, i] = 1.0
                    M_loc[ichunk][i, i + 1] = -1.0
                elseif i == n
                    M_loc[ichunk][i, i - 1] = -1.0
                    M_loc[ichunk][i, i] = 1.0
                else
                    M_loc[ichunk][i, i - 1] = -1.0
                    M_loc[ichunk][i, i] = 2.0
                    M_loc[ichunk][i, i + 1] = -1.0
                end
            end
        end
        Threads.@threads for i in 1:length(M_loc)
            flush!(M_loc[i])
        end
        return
    end

    #
    # This is the slightly modified original MWE
    #
    function nearlyoriginal_mwe(; n = 10_000_000)

        println("nearlyoriginal_mwe, $(Threads.nthreads()) threads")

        println("Serial:")
        M1 = ExtendableSparseMatrix(n, n)
        ts = @elapsed @time serial_assembly!(M1, n)

        println("Threaded:")
        M2 = ExtendableSparseMatrix(n, n)
        M2_loc = [ExtendableSparseMatrix(n, n) for i in 1:nthreads()]

        tt = @elapsed begin
            @time threaded_assembly!(M2_loc, n, Threads.nthreads())

            println("Reduce:")
            # Reduce. Looks like a conversion to sparse is needed to allow for .+
            # This makes things a slightly faster
            # JF: it is not clear if broadcasting is implemented well. `sparse` is fast
            # as it just returns the internal SparseMatrixCSC object created by `flush!` which is always needed anyway.
            @time for k in 1:nthreads()
                M2 .= sparse(M2) .+ sparse(M2_loc[k])
            end
        end
        @show tt / ts
        @show norm(sparse(M1) - sparse(M2))
        return
    end


    function modified_mwe1(; n = 10_000_000)
        println("modified_mwe1, $(Threads.nthreads()) threads")

        println("Serial:")
        M1 = ExtendableSparseMatrix(n, n)
        ts = @elapsed @time serial_assembly!(M1, n)

        println("Threaded:")
        M2_loc = [ExtendableSparseMatrix(n, n) for i in 1:nthreads()]
        tt = @elapsed begin
            @time threaded_assembly!(M2_loc, n, nthreads())

            println("Reduce:")
            # There was no sum function defined in ExtendableSparse before v2.1.
            # I think it fell back to Base.+ defined in https://github.com/WIAS-PDELib/ExtendableSparse.jl/blob/4de82cee639d967bc0a90e98a8911e854e802d51/src/abstractextendablesparsematrixcsc.jl#L308
            # Now, this calls a more efficient sum method.
            @time M2 = sum(M2_loc)
        end
        @show tt / ts
        @show norm(sparse(M1) - sparse(M2))
        return
    end


    #
    # Threaded assembly for MTExtendableSparseMatrixCSC
    # This type keeps a vector of internal linked list like matrices
    #
    function threaded_assembly!(M::AbstractSparseMatrix, n, nchunks)
        chunks = ChunkSplitters.chunks(1:n, n = nchunks)
        @time "    assemble" Threads.@threads for ichunk in 1:nchunks
            for i in chunks[ichunk]
                if i == 1
                    rawupdateindex!(M, +, 1.0, i, i, ichunk)
                    rawupdateindex!(M, +, -1.0, i, i + 1, ichunk)
                elseif i == n
                    rawupdateindex!(M, +, -1.0, i, i - 1, ichunk)
                    rawupdateindex!(M, +, 1.0, i, i, ichunk)
                else
                    rawupdateindex!(M, +, -1.0, i, i - 1, ichunk)
                    rawupdateindex!(M, +, 2.0, i, i, ichunk)
                    rawupdateindex!(M, +, -1.0, i, i + 1, ichunk)
                end
            end
        end
        # Internally, this uses a similar reduction as in Base.sum
        # which however can ignore the redundant colptr entries, as we
        # store the column start indices a dict which in itself has a price
        @time "flush+reduce:" flush!(M)
        return
    end

    #
    # Use multithreaded sparse matrices
    #
    function modified_mwe2(; n = 10_000_000)

        println("modified_mwe2, $(Threads.nthreads()) threads")

        println("Serial:")
        M1 = ExtendableSparseMatrix(n, n)
        # The threaded_assembly! method also works for ExtendableSparseMatrix when using
        # a single thread
        ts = @elapsed threaded_assembly!(M1, n, 1)

        println("Threaded:")
        # Here, we need to pass a number of chunks. This can be large
        # then the number of threads.
        nchunks = 2 * Threads.nthreads()
        M2 = MTExtendableSparseMatrixCSC(n, n, nchunks)
        tt = @elapsed threaded_assembly!(M2, n, nchunks)

        @show tt / ts
        @show norm(sparse(M1) - sparse(M2))
        return
    end

    #
    # Thereaded assembly with an expensive function
    #
    function threaded_assembly!(M::AbstractSparseMatrix, n, nchunks, f::F) where {F}
        chunks = ChunkSplitters.chunks(1:n, n = nchunks)
        @time "    assemble" Threads.@threads for ichunk in 1:nchunks
            for i in chunks[ichunk]
                if i == 1
                    rawupdateindex!(M, +, 1.0 * f(i), i, i, ichunk)
                    rawupdateindex!(M, +, -1.0 * f(i), i, i + 1, ichunk)
                elseif i == n
                    rawupdateindex!(M, +, -1.0 * f(i), i, i - 1, ichunk)
                    rawupdateindex!(M, +, 1.0 * f(i), i, i, ichunk)
                else
                    rawupdateindex!(M, +, -1.0 * f(i), i, i - 1, ichunk)
                    rawupdateindex!(M, +, 2.0 * f(i), i, i, ichunk)
                    rawupdateindex!(M, +, -1.0 * f(i), i, i + 1, ichunk)
                end
            end
        end
        @time "flush+reduce" flush!(M)
        return
    end

    #
    # Case with expensive assembly
    #
    function modified_mwe3(; n = 10_000_000)

        println("modified_mwe3 $(Threads.nthreads()) threads")
        f(i) = exp(sin(exp(sin(exp(sin(exp(cos(i))))))))
        println("Serial:")
        M1 = ExtendableSparseMatrix(n, n)
        ts = @elapsed threaded_assembly!(M1, n, 1, f)

        println("Threaded:")
        nchunks = 2 * Threads.nthreads()
        M2 = MTExtendableSparseMatrixCSC(n, n, nchunks)
        tt = @elapsed threaded_assembly!(M2, n, nchunks, f)

        @show tt / ts
        @show norm(sparse(M1) - sparse(M2))
        return
    end

    function runall(; n = 10_000_000)
        nearlyoriginal_mwe(; n)
        println()
        modified_mwe1(; n)
        println()
        modified_mwe2(; n)
        println()
        modified_mwe3(; n)
        return
    end

    function cli(ARGS)
        println("compiling")
        runall(; n = 100)
        println("\n#################################\nfull run")
        return runall()
    end
end

#
# If the file is run as script, call the cli method with the command line arguments
#
abspath(PROGRAM_FILE) == @__FILE__() &&  DiscourseMWE.cli(ARGS)