Speed up group_by_color (#116)

* Constant number of allocations in `group_by_color` * Improvements and benchmark memory * Remove playground * Docstrings
gdalle · Oct 7, 2024 · e349f50 · e349f50
1 parent c28490d
commit e349f50
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 24 deletions.
diff --git a/.github/workflows/Benchmark.yml b/.github/workflows/Benchmark.yml
@@ -38,7 +38,7 @@ jobs:
                 benchpkg ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --url=${{ github.event.repository.clone_url }} --bench-on="${{github.event.repository.default_branch}}" --output-dir=results/ --tune
             - name: Create markdown table from benchmarks
               run: |
-                benchpkgtable ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --input-dir=results/ --ratio > table.md
+                benchpkgtable ${{ steps.extract-package-name.outputs.package_name }} --mode="time,memory" --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --input-dir=results/ --ratio > table.md
                 echo '### Benchmark Results' > body.md
                 echo '' >> body.md
                 echo '' >> body.md

diff --git a/src/decompression.jl b/src/decompression.jl
@@ -26,7 +26,7 @@ julia> A = sparse([
 
 julia> result = coloring(A, ColoringProblem(), GreedyColoringAlgorithm());
 
-julia> column_groups(result)
+julia> collect.(column_groups(result))
 3-element Vector{Vector{Int64}}:
  [1, 2, 4]
  [3, 5]
@@ -86,7 +86,7 @@ julia> A = sparse([
 
 julia> result = coloring(A, ColoringProblem(), GreedyColoringAlgorithm());
 
-julia> column_groups(result)
+julia> collect.(column_groups(result))
 3-element Vector{Vector{Int64}}:
  [1, 2, 4]
  [3, 5]
@@ -152,7 +152,7 @@ julia> A = sparse([
 
 julia> result = coloring(A, ColoringProblem(), GreedyColoringAlgorithm());
 
-julia> column_groups(result)
+julia> collect.(column_groups(result))
 3-element Vector{Vector{Int64}}:
  [1, 2, 4]
  [3, 5]
@@ -217,7 +217,7 @@ julia> A = sparse([
 
 julia> result = coloring(A, ColoringProblem(), GreedyColoringAlgorithm());
 
-julia> column_groups(result)
+julia> collect.(column_groups(result))
 3-element Vector{Vector{Int64}}:
  [1, 2, 4]
  [3, 5]

diff --git a/src/interface.jl b/src/interface.jl
@@ -156,7 +156,7 @@ julia> column_colors(result)
  2
  3
 
-julia> column_groups(result)
+julia> collect.(column_groups(result))
 3-element Vector{Vector{Int64}}:
  [1, 2, 4]
  [3, 5]

diff --git a/src/result.jl b/src/result.jl
@@ -58,23 +58,32 @@ function row_groups end
 """
     group_by_color(color::Vector{Int})
 
-Create `group::Vector{Vector{Int}}` such that `i ∈ group[c]` iff `color[i] == c`.
+Create a color-indexed vector `group` such that `i ∈ group[c]` iff `color[i] == c`.
 
 Assumes the colors are contiguously numbered from `1` to some `cmax`.
 """
 function group_by_color(color::AbstractVector{<:Integer})
     cmin, cmax = extrema(color)
-    @assert cmin == 1
-    group_sizes = zeros(Int, cmax)
+    @assert cmin >= 1
+    # Compute group sizes and offsets for a joint storage
+    group_sizes = zeros(Int, cmax)  # allocation 1, size cmax
     for c in color
         group_sizes[c] += 1
     end
-    group = [Vector{Int}(undef, group_sizes[c]) for c in 1:cmax]
-    fill!(group_sizes, 1)
+    group_offsets = cumsum(group_sizes)  # allocation 2, size cmax
+    # Concatenate all groups inside a single vector
+    group_flat = similar(color)  # allocation 3, size n
     for (k, c) in enumerate(color)
-        pos = group_sizes[c]
-        group[c][pos] = k
-        group_sizes[c] += 1
+        i = group_offsets[c] - group_sizes[c] + 1
+        group_flat[i] = k
+        group_sizes[c] -= 1
+    end
+    # Create views into contiguous blocks of the group vector
+    group = Vector{typeof(view(group_flat, 1:1))}(undef, cmax)  # allocation 4, size cmax
+    for c in 1:cmax
+        i = 1 + (c == 1 ? 0 : group_offsets[c - 1])
+        j = group_offsets[c]
+        group[c] = view(group_flat, i:j)
     end
     return group
 end
@@ -110,7 +119,7 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct ColumnColoringResult{M<:AbstractMatrix,G<:BipartiteGraph} <:
+struct ColumnColoringResult{M<:AbstractMatrix,G<:BipartiteGraph,V} <:
        AbstractColoringResult{:nonsymmetric,:column,:direct}
     "matrix that was colored"
     A::M
@@ -119,7 +128,7 @@ struct ColumnColoringResult{M<:AbstractMatrix,G<:BipartiteGraph} <:
     "one integer color for each column or row (depending on `partition`)"
     color::Vector{Int}
     "color groups for columns or rows (depending on `partition`)"
-    group::Vector{Vector{Int}}
+    group::V
     "flattened indices mapping the compressed matrix `B` to the uncompressed matrix `A` when `A isa SparseMatrixCSC`. They satisfy `nonzeros(A)[k] = vec(B)[compressed_indices[k]]`"
     compressed_indices::Vector{Int}
 end
@@ -156,12 +165,12 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct RowColoringResult{M<:AbstractMatrix,G<:BipartiteGraph} <:
+struct RowColoringResult{M<:AbstractMatrix,G<:BipartiteGraph,V} <:
        AbstractColoringResult{:nonsymmetric,:row,:direct}
     A::M
     bg::G
     color::Vector{Int}
-    group::Vector{Vector{Int}}
+    group::V
     compressed_indices::Vector{Int}
 end
 
@@ -197,12 +206,12 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct StarSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph} <:
+struct StarSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V} <:
        AbstractColoringResult{:symmetric,:column,:direct}
     A::M
     ag::G
     color::Vector{Int}
-    group::Vector{Vector{Int}}
+    group::V
     star_set::StarSet
     compressed_indices::Vector{Int}
 end
@@ -241,12 +250,12 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct TreeSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,R} <:
+struct TreeSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V,R} <:
        AbstractColoringResult{:symmetric,:column,:substitution}
     A::M
     ag::G
     color::Vector{Int}
-    group::Vector{Vector{Int}}
+    group::V
     vertices_by_tree::Vector{Vector{Int}}
     reverse_bfs_orders::Vector{Vector{Tuple{Int,Int}}}
     buffer::Vector{R}
@@ -395,12 +404,12 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct LinearSystemColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,R,F} <:
+struct LinearSystemColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V,R,F} <:
        AbstractColoringResult{:symmetric,:column,:substitution}
     A::M
     ag::G
     color::Vector{Int}
-    group::Vector{Vector{Int}}
+    group::V
     strict_upper_nonzero_inds::Vector{Tuple{Int,Int}}
     strict_upper_nonzeros_A::Vector{R}  # TODO: adjust type
     T_factorization::F  # TODO: adjust type

diff --git a/test/result.jl b/test/result.jl
@@ -0,0 +1,13 @@
+using SparseMatrixColorings: group_by_color
+using Test
+
+@testset "Group by color" begin
+    for n in 10 .^ (2, 3, 4), cmax in (1, 2, 10, 100), iteration in 1:10
+        color = rand(1:cmax, n)
+        group = group_by_color(color)
+        @test length(group) == maximum(color)
+        @test all(1:maximum(color)) do c
+            all(color[group[c]] .== c) && issorted(group[c])
+        end
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -42,6 +42,9 @@ include("utils.jl")
         @testset "Constructors" begin
             include("constructors.jl")
         end
+        @testset "Result" begin
+            include("result.jl")
+        end
         @testset "Constant coloring" begin
             include("constant.jl")
         end