FastSparsityPattern (#1302)

This patch adds a faster internal path for allocation of matrices using only cell coupling via allocate_matrix. It could be extended to support more cases or merged with SparsityPattern in the future, see #1302 for further details and discussions.

Author: KnutAM

benchmark/Project.toml

@@ -1,4 +1,6 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Ferrite = "c061ca5d-56c9-439f-9c0e-210fe06d3992"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
+SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"

benchmark/benchmark-sparsitypattern.jl

@@ -0,0 +1,82 @@
+using Ferrite
+using SparseMatricesCSR: SparseMatrixCSR
+using SparseArrays: SparseMatrixCSC
+using DataFrames
+using Ferrite: FastSparsityPattern
+
+function dh_scalar(grid)
+    CT = getcelltype(grid)
+    RS = getrefshape(CT)
+    return close!(add!(DofHandler(grid), :u, Lagrange{RS, 1}()))
+end
+
+function dh_complex(grid::Grid{sdim}) where {sdim}
+    CT = getcelltype(grid)
+    RS = getrefshape(CT)
+    dh = DofHandler(grid)
+    add!(dh, :u, Lagrange{RS, 2}())
+    add!(dh, :v, Lagrange{RS, 1}()^sdim)
+    return close!(dh)
+end
+
+create_sp(dh) = add_sparsity_entries!(init_sparsity_pattern(dh), dh)
+create_fsp(dh) = FastSparsityPattern(dh)
+
+grid_2d = generate_grid(Triangle, 1000 .* (1, 1))
+grid_3d = generate_grid(Hexahedron, 80 .* (1, 1, 1))
+
+dofhandlers = [
+    "dh_2d_scalar" => dh_scalar(grid_2d),
+    "dh_3d_scalar" => dh_scalar(grid_3d),
+    "dh_2d_complex" => dh_complex(grid_2d),
+    "dh_3d_complex" => dh_complex(grid_3d),
+]
+
+function timef(f::F, ::Type{MatrixType}, args...; kwargs...) where {F, MatrixType}
+    sp0 = f(args...; kwargs...)       # Compile
+    allocate_matrix(MatrixType, sp0)  # Compile
+    GC.gc()
+    sp_allocs = @allocations (sp_runtime = @elapsed (sp = f(args...; kwargs...)))
+    m_allocs = @allocations (m_runtime = @elapsed allocate_matrix(MatrixType, sp))
+    return (; sp_t = sp_runtime, sp_a = sp_allocs, m_t = m_runtime, m_a = m_allocs)
+end
+
+function fmt_time(t::Number; digits = 2)
+    units = ["ns", "μs", "ms", "s", "min", "h"]
+    values = [1.0e-9, 1.0e-6, 1.0e-3, 1.0, 60.0, 3600]
+    idx = findfirst(v -> v > t, values)
+    i = max(idx === nothing ? length(values) : idx - 1, 1)
+    return string(round(t / values[i]; digits)) * " " * units[i]
+end
+
+function fmt_count(n::Integer; digits = 2)
+    units = ["k", "M", "G"]
+    values = [1.0e3, 1.0e6, 1.0e9]
+    n < 1000 && return string(n)
+    i = findlast(v -> v ≤ n, values)
+    return string(round(n / values[i]; digits)) * units[i]
+end
+
+function make_timings(MatrixType)
+    return map([create_sp, create_fsp]) do f
+        [key => timef(f, MatrixType, dh) for (key, dh) in dofhandlers]
+    end
+end
+
+function make_df(timings)
+    _getdata(v, k) = getindex.(last.(v), k)
+    return DataFrame(
+        "case" => first.(dofhandlers),
+        "t (sp) [s]" => fmt_time.(_getdata(timings[1], :sp_t)),
+        "t (fsp) [s]" => fmt_time.(_getdata(timings[2], :sp_t)),
+        "allocs (sp)" => fmt_count.(_getdata(timings[1], :sp_a)),
+        "allocs (fsp)" => fmt_count.(_getdata(timings[2], :sp_a)),
+        "t (K,sp) [s]" => fmt_time.(_getdata(timings[1], :m_t)),
+        "t (K,fsp) [s]" => fmt_time.(_getdata(timings[2], :m_t)),
+        "allocs (K,sp)" => fmt_count.(_getdata(timings[1], :m_a)),
+        "allocs (K,fsp)" => fmt_count.(_getdata(timings[2], :m_a))
+    )
+end
+
+display(make_df(make_timings(SparseMatrixCSC{Float64, Int})))
+# display(make_df(make_timings(SparseMatrixCSR)))

docs/src/devdocs/special_datastructures.md

@@ -14,3 +14,8 @@ Ferrite.ArrayOfVectorViews
 Ferrite.ConstructionBuffer
 Ferrite.push_at_index!
 ```
+
+## `FastSparsityPattern`
+```@docs
+Ferrite.FastSparsityPattern
+```

ext/FerriteSparseMatrixCSR.jl

@@ -1,7 +1,7 @@
 module FerriteSparseMatrixCSR
 
 using Ferrite, SparseArrays, SparseMatricesCSR
-import Ferrite: AbstractSparsityPattern, CSRAssembler
+import Ferrite: AbstractSparsityPattern, CSRAssembler, FastSparsityPattern, getnrows, getncols
 import Base: @propagate_inbounds
 
 # Could be generalized if https://github.com/JuliaSparse/SparseArrays.jl/pull/546 is merged
@@ -113,4 +113,41 @@ function _allocate_matrix(::Type{SparseMatrixCSR{1, Tv, Ti}}, sp::AbstractSparsi
     return S
 end
 
+## ================= ##
+# FastSparsityPattern #
+## ================= ##
+
+function _allocate_matrix(::Type{SparseMatrixCSR{1, Tv, Ti}}, sp::FastSparsityPattern{Ti}, sym::Bool) where {Tv, Ti}
+    sym && throw(ArgumentError("FastSparsityPattern does not support symmetric matrices yet"))
+    sp.is_colidx_sorted || sort_rows_threaded!(sp) # Require sorted rows
+    nzval = zeros(Tv, length(sp.colidx))
+    return SparseMatrixCSR{1}(getnrows(sp), getncols(sp), sp.rowptr, sp.colidx, nzval)
+end
+
+function sort_rows!(sp::FastSparsityPattern, rowrange::UnitRange)
+    @inbounds for row in rowrange
+        i1 = sp.rowptr[row]
+        i2 = sp.rowptr[row + 1] - 1
+        if i1 < i2
+            sort!(view(sp.colidx, i1:i2); alg = QuickSort)
+        end
+    end
+    return sp
+end
+
+function sort_rows_threaded!(
+        sp::FastSparsityPattern, # Default ΔN ≥ 1000 and `n_tasks ≥ 1`
+        ntasks = max(min(Threads.nthreads() * 100, getnrows(sp) ÷ 1000), 1)
+    )               # Otherwise, 100 per thread for load balancing
+    nrows = getnrows(sp)
+    ΔN = cld(nrows, ntasks)
+    Threads.@threads for taskid in 1:ntasks
+        first_idx = 1 + ΔN * (taskid - 1)
+        last_idx = min(first_idx + ΔN - 1, nrows)
+        sort_rows!(sp, first_idx:last_idx)
+    end
+    sp.is_colidx_sorted = true
+    return sp
+end
+
 end

src/Dofs/sparsity_pattern.jl

@@ -413,6 +413,12 @@ arguments `args` and keyword arguments `kwargs`.
     copy(K)`) instead.
 """
 function allocate_matrix(::Type{MatrixType}, dh::DofHandler, args...; kwargs...) where {MatrixType}
+    _get_Ti(::Type{<:AbstractMatrix}) = Int
+    _get_Ti(::Type{<:AbstractSparseMatrix{<:Any, Ti}}) where {Ti} = Ti
+    if _can_use_fastsp(MatrixType, args...; kwargs...)
+        fsp = FastSparsityPattern(_get_Ti(MatrixType), dh, args...; kwargs...)
+        return allocate_matrix(MatrixType, fsp)
+    end
     sp = init_sparsity_pattern(dh)
     add_sparsity_entries!(sp, dh, args...; kwargs...)
     return allocate_matrix(MatrixType, sp)
@@ -681,3 +687,198 @@ function _allocate_matrix(::Type{SparseMatrixCSC{Tv, Ti}}, sp::AbstractSparsityP
     S = SparseMatrixCSC(getnrows(sp), getncols(sp), colptr, rowval, nzval)
     return S
 end
+
+## ================= ##
+# FastSparsityPattern #
+## ================= ##
+
+"""
+    FastSparsityPattern([Ti = Int64], dh::DofHandler)
+
+This sparsity does not currently support the full `AbstractSparsityPattern` interface,
+but is used as an internal fast-path for `allocate_matrix(MatrixType, dh)` for some
+supported `MatrixType`s. It can be extended in the future or potentially be merged
+with `SparsityPattern`.
+See [#1302](https://github.com/Ferrite-FEM/Ferrite.jl/pull/1302) for details.
+
+!!! warning "Internal"
+    `FastSparsityPattern` is strictly internal and its interface and implementation
+    may change at any time.
+
+"""
+mutable struct FastSparsityPattern{Ti} <: AbstractSparsityPattern
+    const rowlen::Vector{Ti} # Number of stored entries in each row
+    const marker::Vector{Ti} # Marker if column has been "visited" by certain row
+    const rowptr::Vector{Ti} # Index of stored entries at the start of each row
+    const colidx::Vector{Ti} # colidx[i] gives the column number of the ith stored entry
+    is_colidx_sorted::Bool   # Is colidx sorted for each row
+end
+function FastSparsityPattern(::Type{Ti}, ncols, nrows) where {Ti <: Integer}
+    rowlen = zeros(Ti, nrows)
+    marker = zeros(Ti, ncols)
+    rowptr = Vector{Ti}(undef, nrows + 1)
+    colidx = Vector{Ti}(undef, 0) # To be resized later
+    return FastSparsityPattern(rowlen, marker, rowptr, colidx, false)
+end
+
+# _can_use_fastsp(MatrixType, args...; kwargs...) where args and kwargs are those passed to
+# `allocate_matrix`. See `add_sparsity_entries!` for a description of args/kwargs.
+function _can_use_fastsp(
+        ::Type{MatrixType},
+        ch = nothing;
+        topology = nothing,
+        keep_constrained = true,
+        coupling = nothing,
+        interface_coupling = nothing
+    ) where {MatrixType}
+    if ch === topology === coupling === interface_coupling === nothing
+        if MatrixType <: AbstractSparseMatrix # Symmetric/Block matrices not supported
+            return keep_constrained
+        end
+    end
+    return false
+end
+
+FastSparsityPattern(dh::DofHandler) = FastSparsityPattern(Int, dh)
+function FastSparsityPattern(::Type{Ti}, dh::DofHandler) where {Ti}
+    sp = FastSparsityPattern(Ti, ndofs(dh), ndofs(dh))
+    # Step 1: Create cell_dof_views::ArrayOfVectorViews (would be nice in `DofHandler` directly)
+    cell_dofs_views = create_celldofs(dh)
+    # Step 2: Define mapping rownr to cells
+    row_to_cells = create_row_to_cells(cell_dofs_views, sp)
+    # Step 3: Count how many cols stored for each row
+    count_row_sizes!(sp, row_to_cells, cell_dofs_views)
+    # Step 4: Build the rowptr (indices for s)
+    build_rowptr!(sp)
+    fill_colidx!(sp, row_to_cells, cell_dofs_views)
+    return sp
+end
+
+getncols(sp::FastSparsityPattern) = length(sp.marker)
+getnrows(sp::FastSparsityPattern) = length(sp.rowlen)
+
+function create_celldofs(dh::DofHandler)
+    isclosed(dh) || throw(ArgumentError("DofHandler must be closed"))
+    ncells = getncells(dh.grid)
+    indices = similar(dh.cell_dofs_offset, ncells + 1)
+    cell_dofs = similar(dh.cell_dofs)
+    n = 1
+    for cell_idx in 1:ncells
+        indices[cell_idx] = n
+        num = ndofs_per_cell(dh, cell_idx)
+        num == 0 && continue
+        r = n:(n + num - 1)
+        #celldofs!(view(cell_dofs, r), dh, cell_idx), but faster without view:
+        soffs = dh.cell_dofs_offset[cell_idx]
+        copyto!(cell_dofs, n, dh.cell_dofs, soffs, num)
+        n = last(r) + 1
+    end
+    indices[end] = n
+    return ArrayOfVectorViews(indices, cell_dofs, LinearIndices((ncells,)))
+end
+
+function create_row_to_cells(cell_dofs::ArrayOfVectorViews, sp)
+    nrows = getnrows(sp)
+    num_cells = zeros(Int, nrows)
+    # 1: Figure out how many cells are connected to each dof
+    n_connected = 0
+    @inbounds for rows in cell_dofs # dof = row
+        for row in rows
+            num_cells[row] += 1
+            n_connected += 1
+        end
+    end
+
+    # 2: Create the correct datastructure
+    data = Vector{Int}(undef, n_connected)
+    indices = Vector{Int}(undef, nrows + 1)
+    indices[1] = 1
+    @inbounds for row in 1:nrows
+        indices[row + 1] = indices[row] + num_cells[row]
+    end
+    fill!(num_cells, 0) # Now we use this to count how many have been added
+    @inbounds for (cellnr, rows) in enumerate(cell_dofs)
+        for row in rows
+            data[indices[row] + num_cells[row]] = cellnr
+            num_cells[row] += 1
+        end
+    end
+    return ArrayOfVectorViews(indices, data, LinearIndices((nrows,)))
+end
+
+function count_row_sizes!(sp::FastSparsityPattern, row_to_cells::AbstractVector, cell_dofs::ArrayOfVectorViews)
+    @inbounds for row in 1:getnrows(sp)
+        for cnr in row_to_cells[row]
+            for col in cell_dofs[cnr]
+                if sp.marker[col] != row
+                    sp.marker[col] = row
+                    sp.rowlen[row] += 1
+                end
+            end
+        end
+    end
+    return sp
+end
+
+function build_rowptr!(sp)
+    sp.rowptr[1] = 1
+    @inbounds for row in 1:getnrows(sp)
+        sp.rowptr[row + 1] = sp.rowptr[row] + sp.rowlen[row]
+    end
+    return sp
+end
+
+function fill_colidx!(sp::FastSparsityPattern, row_to_cells::AbstractVector, cell_dofs::ArrayOfVectorViews)
+    resize!(sp.colidx, sp.rowptr[end] - 1) # nnz
+    fill!(sp.marker, 0)
+    @inbounds for row in 1:getnrows(sp)
+        pos = sp.rowptr[row]
+        for cnr in row_to_cells[row]
+            for col in cell_dofs[cnr]
+                if sp.marker[col] != row
+                    sp.marker[col] = row
+                    sp.colidx[pos] = col
+                    pos += 1
+                end
+            end
+        end
+    end
+    return sp
+end
+
+allocate_matrix(sp::FastSparsityPattern) = allocate_matrix(SparseMatrixCSC, sp)
+allocate_matrix(::Type{SparseMatrixCSC}, sp::FastSparsityPattern{Int}) = allocate_matrix(SparseMatrixCSC{Float64, Int}, sp)
+function allocate_matrix(::Type{<:SparseMatrixCSC{Tv, Ti}}, sp::FastSparsityPattern{Ti}) where {Ti, Tv}
+    nnz = length(sp.colidx)
+    ncols = getncols(sp)
+    nrows = getnrows(sp)
+
+    # Number of stored entries per column
+    collen = zeros(Ti, ncols)
+    @inbounds for col in sp.colidx
+        collen[col] += 1
+    end
+
+    # Index of stored entries at the start of each column
+    colptr = Vector{Ti}(undef, ncols + 1)
+    colptr[1] = 1
+    @inbounds for col in 1:ncols
+        colptr[col + 1] = colptr[col] + collen[col]
+    end
+
+    # Build rowidx[i] giving the row number of the ith stored entry
+    rowidx = Vector{Ti}(undef, nnz)
+    next = copy(colptr)
+    @inbounds for row in 1:nrows
+        for p in sp.rowptr[row]:(sp.rowptr[row + 1] - 1)
+            col = sp.colidx[p]
+            q = next[col]
+            rowidx[q] = row
+            next[col] = q + 1
+            # For a given col, next[col] is increasing, and row is
+            # increasing in the outer loop -> rowidx sorted for each col
+        end
+    end
+    nzval = zeros(Tv, nnz)
+    return SparseMatrixCSC(nrows, ncols, colptr, rowidx, nzval)
+end