Speed up CSC substitution decompression for stars

src/decompression.jl

@@ -610,6 +610,46 @@ function decompress!(
     return A
 end
 
+@inline function set_symmetric_csc!(
+    nzA, A_colptr, lower_triangle_offsets, upper_triangle_offsets, counter, i, j, val, uplo
+)
+    #! format: off
+    # A[i,j] is in the lower triangular part of A
+    if in_triangle(i, j, :L)
+        # uplo = :L or uplo = :F
+        # A[i,j] is stored at index_ij = (A.colptr[j+1] - offset_L) in A.nzval
+        if uplo != :U
+            nzind = A_colptr[j + 1] - lower_triangle_offsets[counter]
+            nzA[nzind] = val
+        end
+
+        # uplo = :U or uplo = :F
+        # A[j,i] is stored at index_ji = (A.colptr[i] + offset_U) in A.nzval
+        if uplo != :L
+            nzind = A_colptr[i] + upper_triangle_offsets[counter]
+            nzA[nzind] = val
+        end
+
+    # A[i,j] is in the upper triangular part of A
+    else
+        # uplo = :U or uplo = :F
+        # A[i,j] is stored at index_ij = (A.colptr[j] + offset_U) in A.nzval
+        if uplo != :L
+            nzind = A_colptr[j] + upper_triangle_offsets[counter]
+            nzA[nzind] = val
+        end
+
+        # uplo = :L or uplo = :F
+        # A[j,i] is stored at index_ji = (A.colptr[i+1] - offset_L) in A.nzval
+        if uplo != :U
+            nzind = A_colptr[i + 1] - lower_triangle_offsets[counter]
+            nzA[nzind] = val
+        end
+    end
+    #! format: on
+    return nothing
+end
+
 function decompress_csc!(
     nzA::AbstractVector{R},
     A_colptr::AbstractVector{<:Integer},
@@ -621,6 +661,7 @@ function decompress_csc!(
         ag,
         color,
         reverse_bfs_orders,
+        is_star,
         tree_edge_indices,
         nt,
         diagonal_indices,
@@ -667,6 +708,27 @@ function decompress_csc!(
         first = tree_edge_indices[k]
         last = tree_edge_indices[k + 1] - 1
 
+        # A star needs no substitution or scratch-buffer updates.
+        if is_star[k]
+            for pos in first:last
+                (i, j) = reverse_bfs_orders[pos]
+                counter += 1
+                val = B[i, color[j]]
+                set_symmetric_csc!(
+                    nzA,
+                    A_colptr,
+                    lower_triangle_offsets,
+                    upper_triangle_offsets,
+                    counter,
+                    i,
+                    j,
+                    val,
+                    uplo,
+                )
+            end
+            continue
+        end
+
         # Reset the buffer to zero for all vertices in the tree (except the root)
         for pos in first:last
             (vertex, _) = reverse_bfs_orders[pos]
@@ -682,40 +744,17 @@ function decompress_csc!(
             val = B[i, color[j]] - buffer_right_type[i]
             buffer_right_type[j] = buffer_right_type[j] + val
 
-            #! format: off
-            # A[i,j] is in the lower triangular part of A
-            if in_triangle(i, j, :L)
-                # uplo = :L or uplo = :F
-                # A[i,j] is stored at index_ij = (A.colptr[j+1] - offset_L) in A.nzval
-                if uplo != :U
-                    nzind = A_colptr[j + 1] - lower_triangle_offsets[counter]
-                    nzA[nzind] = val
-                end
-
-                # uplo = :U or uplo = :F
-                # A[j,i] is stored at index_ji = (A.colptr[i] + offset_U) in A.nzval
-                if uplo != :L
-                    nzind = A_colptr[i] + upper_triangle_offsets[counter]
-                    nzA[nzind] = val
-                end
-
-            # A[i,j] is in the upper triangular part of A
-            else
-                # uplo = :U or uplo = :F
-                # A[i,j] is stored at index_ij = (A.colptr[j] + offset_U) in A.nzval
-                if uplo != :L
-                    nzind = A_colptr[j] + upper_triangle_offsets[counter]
-                    nzA[nzind] = val
-                end
-
-                # uplo = :L or uplo = :F
-                # A[j,i] is stored at index_ji = (A.colptr[i+1] - offset_L) in A.nzval
-                if uplo != :U
-                    nzind = A_colptr[i + 1] - lower_triangle_offsets[counter]
-                    nzA[nzind] = val
-                end
-            end
-            #! format: on
+            set_symmetric_csc!(
+                nzA,
+                A_colptr,
+                lower_triangle_offsets,
+                upper_triangle_offsets,
+                counter,
+                i,
+                j,
+                val,
+                uplo,
+            )
         end
     end
     return nothing

src/result.jl

@@ -387,6 +387,7 @@ struct TreeSetColoringResult{
     color::Vector{T}
     group::GT
     reverse_bfs_orders::Vector{Tuple{T,T}}
+    is_star::Vector{Bool}
     tree_edge_indices::Vector{T}
     nt::T
     diagonal_indices::Vector{T}
@@ -403,7 +404,7 @@ function TreeSetColoringResult(
     tree_set::TreeSet{<:Integer},
     decompression_eltype::Type{R},
 ) where {T<:Integer,R}
-    (; reverse_bfs_orders, tree_edge_indices, nt) = tree_set
+    (; reverse_bfs_orders, is_star, tree_edge_indices, nt) = tree_set
     (; S, nb_self_loops) = ag
     nvertices = length(color)
     group = group_by_color(T, color)
@@ -484,6 +485,7 @@ function TreeSetColoringResult(
         color,
         group,
         reverse_bfs_orders,
+        is_star,
         tree_edge_indices,
         nt,
         diagonal_indices,

test/decompression.jl

@@ -1,7 +1,7 @@
 using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings
-using SparseMatrixColorings: decompress_csc!
+using SparseMatrixColorings: SparsityPatternCSC, decompress_csc!
 using Test
 
 function test_star_csc_decompression(T, Ti, order, postprocessing)
@@ -76,3 +76,114 @@ end
         test_star_csc_decompression(T, Ti, order, postprocessing)
     end
 end
+
+function test_tree_csc_decompression(T, Ti, order, postprocessing)
+    problem = ColoringProblem(; structure=:symmetric)
+    algorithm = GreedyColoringAlgorithm(order; decompression=:substitution, postprocessing)
+    @testset "$pattern" for pattern in (
+        :blocks, :star_low, :star_middle, :star_high, :two_edges, :mixed, :diagonal, :empty
+    )
+        A = spzeros(T, Ti, 10, 10)
+        if pattern in (:blocks, :mixed)
+            clique_size = pattern == :blocks ? 4 : 3
+            for j in 1:clique_size, i in 1:(j - 1)
+                A[i, j] = A[j, i] = T(10i + j)
+            end
+            edges = if pattern == :blocks
+                [(6, 7), (6, 8), (7, 8)]
+            else
+                [(4, 5), (5, 6), (6, 7), (8, 9)]
+            end
+            for (i, j) in edges
+                A[i, j] = A[j, i] = T(10i + j)
+            end
+        end
+        if pattern in (:star_low, :star_middle, :star_high, :two_edges)
+            hub = if pattern == :star_low
+                1
+            elseif pattern == :star_high
+                9
+            else
+                5
+            end
+            leaves = pattern == :two_edges ? (2, 7) : (2, 4, 7)
+            for leaf in leaves
+                A[leaf, hub] = A[hub, leaf] = T(10leaf + hub)
+            end
+        end
+        if pattern != :empty
+            # Missing diagonals, a structural zero, and an empty last column.
+            A[2, 2] = T(2)
+            A[8, 8] = T(8)
+            A[3, 3] = one(T)
+            A[3, 3] = zero(T)
+        end
+        result = coloring(SparsityPatternCSC(A), problem, algorithm; decompression_eltype=T)
+        sizes = diff(result.tree_edge_indices)
+        if pattern == :blocks
+            @test !isempty(sizes)
+            @test all(==(1), sizes)
+            @test all(result.is_star)
+        elseif pattern in (:star_low, :star_middle, :star_high, :two_edges)
+            @test sizes == [pattern == :two_edges ? 2 : 3]
+            @test all(result.is_star)
+        elseif pattern == :mixed
+            @test any(==(1), sizes)
+            @test any(>(1), sizes)
+            @test any(!, result.is_star)
+        end
+        B = compress(A, result)
+        @testset "$uplo" for uplo in (:F, :L, :U)
+            expected = if uplo == :F
+                A
+            elseif uplo == :L
+                tril(A)
+            else
+                triu(A)
+            end
+            values = fill(T(NaN), nnz(expected))
+            # Reuse the result with changed compressed values and stale scratch data.
+            for scale in (one(T), T(-2))
+                compressed = scale .* B
+                fill!(result.buffer, T(NaN))
+                @test decompress_csc!(values, expected.colptr, compressed, result, uplo) ===
+                    nothing
+                @test values == scale .* nonzeros(expected)
+                if all(result.is_star)
+                    @test all(isnan, result.buffer)
+                end
+                @test (@allocated decompress_csc!(
+                    values, expected.colptr, compressed, result, uplo
+                )) == 0
+
+                dest = copy(expected)
+                fill!(nonzeros(dest), T(NaN))
+                @test decompress!(dest, compressed, result, uplo) === dest
+                @test dest == scale .* expected
+                @test (@allocated decompress!(dest, compressed, result, uplo)) == 0
+            end
+
+            # The CSC entry point also accepts views of its input/output buffers.
+            storage = fill(T(NaN), 2nnz(expected))
+            values_view = @view storage[1:2:end]
+            compressed_view = @view B[:, :]
+            @test decompress_csc!(
+                values_view, expected.colptr, compressed_view, result, uplo
+            ) === nothing
+            @test values_view == nonzeros(expected)
+            @test (@allocated decompress_csc!(
+                values_view, expected.colptr, compressed_view, result, uplo
+            )) == 0
+        end
+    end
+    return nothing
+end
+
+@testset "Tree CSC decompression" begin
+    @testset "$T, $Ti, $order, postprocessing=$postprocessing" for T in (Float32, Float64),
+        Ti in (Int32, Int64), order in (NaturalOrder(), LargestFirst()),
+        postprocessing in (false, true)
+
+        test_tree_csc_decompression(T, Ti, order, postprocessing)
+    end
+end