Optimized implementation for structured matrices V2

Project.toml

@@ -12,6 +12,9 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
@@ -21,6 +24,8 @@ MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 cuSPARSE = "b26da814-b3bc-49ef-b0ee-c816305aa060"
 
 [extensions]
+SparseMatrixColoringsBandedMatricesExt = "BandedMatrices"
+SparseMatrixColoringsBlockBandedMatricesExt = ["BlockArrays", "BlockBandedMatrices"]
 SparseMatrixColoringsCUDAExt = ["CUDA", "cuSPARSE"]
 SparseMatrixColoringsCliqueTreesExt = "CliqueTrees"
 SparseMatrixColoringsColorsExt = "Colors"
@@ -29,6 +34,9 @@ SparseMatrixColoringsJuMPExt = ["JuMP", "MathOptInterface"]
 
 [compat]
 ADTypes = "1.2.1"
+BandedMatrices = "1.9.4"
+BlockArrays = "1.6.3"
+BlockBandedMatrices = "0.13.1"
 CUDA = "6.0.0"
 CliqueTrees = "1"
 Colors = "0.12.11, 0.13"

docs/make.jl

@@ -2,7 +2,10 @@ using Documenter
 using DocumenterInterLinks
 using SparseMatrixColorings
 
-links = InterLinks("ADTypes" => "https://sciml.github.io/ADTypes.jl/stable/")
+links = InterLinks(
+    "ADTypes" => "https://sciml.github.io/ADTypes.jl/stable/",
+    "BandedMatrices" => "https://julialinearalgebra.github.io/BandedMatrices.jl/stable/",
+)
 
 cp(joinpath(@__DIR__, "..", "README.md"), joinpath(@__DIR__, "src", "index.md"); force=true)
 

docs/src/api.md

@@ -25,6 +25,7 @@ ColoringProblem
 GreedyColoringAlgorithm
 ConstantColoringAlgorithm
 OptimalColoringAlgorithm
+StructuredColoringAlgorithm
 ```
 
 ## Result analysis

docs/src/dev.md

@@ -80,3 +80,9 @@ SparseMatrixColorings.what_fig_61
 SparseMatrixColorings.efficient_fig_1
 SparseMatrixColorings.efficient_fig_4
 ```
+
+## Misc
+
+```@docs
+SparseMatrixColorings.cycle_range
+```

ext/SparseMatrixColoringsBandedMatricesExt.jl

@@ -0,0 +1,45 @@
+module SparseMatrixColoringsBandedMatricesExt
+
+using BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
+using SparseMatrixColorings:
+    BipartiteGraph,
+    ColoringProblem,
+    ColumnColoringResult,
+    StructuredColoringAlgorithm,
+    RowColoringResult,
+    column_colors,
+    cycle_range,
+    row_colors
+import SparseMatrixColorings as SMC
+
+#=
+This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
+https://github.com/JuliaArrays/ArrayInterface.jl
+https://github.com/JuliaDiff/FiniteDiff.jl
+=#
+
+function SMC.coloring(
+    A::BandedMatrix,
+    ::ColoringProblem{:nonsymmetric,:column},
+    ::StructuredColoringAlgorithm;
+    kwargs...,
+)
+    width = length(bandrange(A))
+    color = cycle_range(width, size(A, 2))
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
+end
+
+function SMC.coloring(
+    A::BandedMatrix,
+    ::ColoringProblem{:nonsymmetric,:row},
+    ::StructuredColoringAlgorithm;
+    kwargs...,
+)
+    width = length(bandrange(A))
+    color = cycle_range(width, size(A, 1))
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
+end
+
+end

ext/SparseMatrixColoringsBlockBandedMatricesExt.jl

@@ -0,0 +1,97 @@
+module SparseMatrixColoringsBlockBandedMatricesExt
+
+using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
+using BlockBandedMatrices:
+    BandedBlockBandedMatrix,
+    BlockBandedMatrix,
+    blockbandrange,
+    blockbandwidths,
+    blocklengths,
+    blocksize,
+    subblockbandwidths
+using SparseMatrixColorings:
+    BipartiteGraph,
+    ColoringProblem,
+    ColumnColoringResult,
+    StructuredColoringAlgorithm,
+    RowColoringResult,
+    column_colors,
+    cycle_range,
+    row_colors
+import SparseMatrixColorings as SMC
+
+#=
+This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
+https://github.com/JuliaArrays/ArrayInterface.jl
+https://github.com/JuliaDiff/FiniteDiff.jl
+=#
+
+function subblockbandrange(A::BandedBlockBandedMatrix)
+    u, l = subblockbandwidths(A)
+    return (-l):u
+end
+
+function blockbanded_coloring(
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix}, dim::Integer
+)
+    # consider blocks of columns or rows (let's call them vertices) depending on `dim`
+    nb_blocks = blocksize(A, dim)
+    nb_in_block = blocklengths(axes(A, dim))
+    first_in_block = blockfirsts(axes(A, dim))
+    last_in_block = blocklasts(axes(A, dim))
+    color = zeros(Int, size(A, dim))
+
+    # give a macroscopic color to each block, so that 2 blocks with the same macro color are orthogonal
+    # same idea as for BandedMatrices
+    nb_macrocolors = length(blockbandrange(A))
+    macrocolor = cycle_range(nb_macrocolors, nb_blocks)
+
+    width = if A isa BandedBlockBandedMatrix
+        # vertices within a block are colored cleverly using bands
+        length(subblockbandrange(A))
+    else
+        # vertices within a block are colored naively with distinct micro colors (~ infinite band width)
+        typemax(Int)
+    end
+
+    # for each macroscopic color, count how many microscopic colors will be needed
+    nb_colors_in_macrocolor = zeros(Int, nb_macrocolors)
+    for mc in 1:nb_macrocolors
+        largest_nb_in_macrocolor = maximum(nb_in_block[mc:nb_macrocolors:nb_blocks]; init=0)
+        nb_colors_in_macrocolor[mc] = min(width, largest_nb_in_macrocolor)
+    end
+    color_shift_in_macrocolor = vcat(0, cumsum(nb_colors_in_macrocolor)[1:(end - 1)])
+
+    # assign a microscopic color to each column as a function of its macroscopic color and its position within the block
+    for b in 1:nb_blocks
+        block_color = cycle_range(width, nb_in_block[b])
+        shift = color_shift_in_macrocolor[macrocolor[b]]
+        color[first_in_block[b]:last_in_block[b]] .= shift .+ block_color
+    end
+
+    return color
+end
+
+function SMC.coloring(
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
+    ::ColoringProblem{:nonsymmetric,:column},
+    ::StructuredColoringAlgorithm;
+    kwargs...,
+)
+    color = blockbanded_coloring(A, 2)
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
+end
+
+function SMC.coloring(
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
+    ::ColoringProblem{:nonsymmetric,:row},
+    ::StructuredColoringAlgorithm;
+    kwargs...,
+)
+    color = blockbanded_coloring(A, 1)
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
+end
+
+end

ext/SparseMatrixColoringsCUDAExt.jl

@@ -1,10 +1,31 @@
 module SparseMatrixColoringsCUDAExt
-
+using LinearAlgebra
 import SparseMatrixColorings as SMC
 using SparseArrays: SparseMatrixCSC, rowvals, nnz, nzrange
-using CUDA: CuVector, CuMatrix
+using CUDA: CuArray, CuVector, CuMatrix
 using cuSPARSE: AbstractCuSparseMatrix, CuSparseMatrixCSC, CuSparseMatrixCSR
 
+## Basic support for GPU sparsity pattern stuff
+
+SMC.SparsityPatternCSC(A::CuSparseMatrixCSC) = SMC.SparsityPatternCSC(first(A.dims), last(A.dims), A.colPtr, A.rowVal)
+
+for R in (:Diagonal, :Bidiagonal, :Tridiagonal)
+    @eval function SMC.BipartiteGraph(A::$R{T, <:CuArray}; symmetric_pattern::Bool=false) where {T}
+        return SMC.BipartiteGraph(CuSparseMatrixCSC(A); symmetric_pattern)
+    end
+end
+
+function SMC.BipartiteGraph(A::CuSparseMatrixCSC; symmetric_pattern::Bool=false)
+    S2 = SMC.SparsityPatternCSC(A)
+    if symmetric_pattern
+        checksquare(A)  # proxy for checking full symmetry
+        S1 = S2
+    else
+        S1 = transpose(S2)  # rows to columns
+    end
+    return SMC.BipartiteGraph(S1, S2)
+end
+
 ## CSC Result
 
 function SMC.ColumnColoringResult(

src/SparseMatrixColorings.jl

@@ -14,11 +14,13 @@ using Base.Iterators: Iterators
 using DocStringExtensions: README, EXPORTS, SIGNATURES, TYPEDEF, TYPEDFIELDS
 using LinearAlgebra:
     Adjoint,
+    Bidiagonal,
     Diagonal,
     Hermitian,
     LowerTriangular,
     Symmetric,
     Transpose,
+    Tridiagonal,
     UpperTriangular,
     adjoint,
     checksquare,
@@ -54,6 +56,7 @@ include("interface.jl")
 include("constant.jl")
 include("adtypes.jl")
 include("decompression.jl")
+include("structured.jl")
 include("check.jl")
 include("examples.jl")
 include("show_colors.jl")
@@ -65,7 +68,7 @@ export NaturalOrder, RandomOrder, LargestFirst
 export DynamicDegreeBasedOrder, SmallestLast, IncidenceDegree, DynamicLargestFirst
 export PerfectEliminationOrder
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
-export ConstantColoringAlgorithm
+export ConstantColoringAlgorithm, StructuredColoringAlgorithm
 export OptimalColoringAlgorithm
 export coloring, fast_coloring
 export column_colors, row_colors, ncolors

src/graph.jl

@@ -17,8 +17,8 @@ Copied from `SparseMatrixCSC`:
 struct SparsityPatternCSC{Ti<:Integer} <: AbstractMatrix{Bool}
     m::Int
     n::Int
-    colptr::Vector{Ti}
-    rowval::Vector{Ti}
+    colptr::AbstractVector{Ti}
+    rowval::AbstractVector{Ti}
 end
 
 SparsityPatternCSC(A::SparseMatrixCSC) = SparsityPatternCSC(A.m, A.n, A.colptr, A.rowval)