diff --git a/Project.toml b/Project.toml index e72441d7..d152860e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "SparseMatrixColorings" uuid = "0a514795-09f3-496d-8182-132a7b665d35" -version = "0.4.28" +version = "0.5.0" authors = ["Guillaume Dalle", "Alexis Montoison"] [deps] diff --git a/docs/src/tutorial.md b/docs/src/tutorial.md index 6e803511..09514a48 100644 --- a/docs/src/tutorial.md +++ b/docs/src/tutorial.md @@ -23,7 +23,8 @@ SparseMatrixColorings.jl is based on the combination of a coloring problem and a The problem defines what you want to solve. It is always a [`ColoringProblem`](@ref), and you can select options such as - the structure of the matrix (`:nonsymmetric` or `:symmetric`) -- the type of partition you want (`:column`, `:row` or `:bidirectional`). +- the type of partition you want (`:column`, `:row` or `:bidirectional`) +- the part of the matrix you want to recover during decompression (`:F` for the full matrix, `:L` for its lower triangle or `:U` for its upper triangle). This option is only available for `:symmetric` problems, and it is fixed once and for all here, so it never has to be repeated at decompression time. ```@example tutorial problem = ColoringProblem() @@ -211,6 +212,19 @@ and its columnwise compression B_img # hide ``` +#### Recovering a single triangle + +Since a symmetric matrix is redundant, you can ask for only one of its triangles with the `uplo` option of the problem. +The choice is made once, at coloring time, and the decompression then fills in that triangle only: + +```@example tutorial +problem_U = ColoringProblem(; structure=:symmetric, partition=:column, uplo=:U) +result_U = coloring(S, problem_U, algo) +decompress(compress(S, result_U), result_U) +``` + +A given result decompresses into exactly one triangle, so build one result per triangle if you need several. + ### Acyclic coloring Acyclic coloring is the algorithm used for symmetric matrices with decompression by substitution. diff --git a/ext/SparseMatrixColoringsAMDGPUExt.jl b/ext/SparseMatrixColoringsAMDGPUExt.jl index 58714ad2..cd8eeed2 100644 --- a/ext/SparseMatrixColoringsAMDGPUExt.jl +++ b/ext/SparseMatrixColoringsAMDGPUExt.jl @@ -32,12 +32,20 @@ function SMC.StarSetColoringResult( ag::SMC.AdjacencyGraph{T}, color::Vector{<:Integer}, star_set::SMC.StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T<:Integer} + if decompression_uplo != :F + throw( + SMC.UnsupportedDecompressionError( + "Single-triangle decompression is not supported on GPU matrices" + ), + ) + end group = SMC.group_by_color(T, color) - compressed_indices = SMC.star_csc_indices(ag, color, star_set) + compressed_indices = SMC.star_csc_indices(ag, color, star_set, decompression_uplo) additional_info = (; compressed_indices_gpu_csc=ROCVector(compressed_indices)) return SMC.StarSetColoringResult( - A, ag, color, group, compressed_indices, additional_info + A, ag, color, group, compressed_indices, decompression_uplo, additional_info ) end @@ -70,12 +78,20 @@ function SMC.StarSetColoringResult( ag::SMC.AdjacencyGraph{T}, color::Vector{<:Integer}, star_set::SMC.StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T<:Integer} + if decompression_uplo != :F + throw( + SMC.UnsupportedDecompressionError( + "Single-triangle decompression is not supported on GPU matrices" + ), + ) + end group = SMC.group_by_color(T, color) - compressed_indices = SMC.star_csc_indices(ag, color, star_set) + compressed_indices = SMC.star_csc_indices(ag, color, star_set, decompression_uplo) additional_info = (; compressed_indices_gpu_csr=ROCVector(compressed_indices)) return SMC.StarSetColoringResult( - A, ag, color, group, compressed_indices, additional_info + A, ag, color, group, compressed_indices, decompression_uplo, additional_info ) end @@ -104,15 +120,7 @@ function SMC.decompress!( A::ROCSparseMatrixCSC, B::ROCMatrix, result::SMC.StarSetColoringResult{<:ROCSparseMatrixCSC}, - uplo::Symbol=:F, ) - if uplo != :F - throw( - SMC.UnsupportedDecompressionError( - "Single-triangle decompression is not supported on GPU matrices" - ), - ) - end compressed_indices = result.additional_info.compressed_indices_gpu_csc copyto!(A.nzVal, view(B, compressed_indices)) return A @@ -122,15 +130,7 @@ function SMC.decompress!( A::ROCSparseMatrixCSR, B::ROCMatrix, result::SMC.StarSetColoringResult{<:ROCSparseMatrixCSR}, - uplo::Symbol=:F, ) - if uplo != :F - throw( - SMC.UnsupportedDecompressionError( - "Single-triangle decompression is not supported on GPU matrices" - ), - ) - end compressed_indices = result.additional_info.compressed_indices_gpu_csr copyto!(A.nzVal, view(B, compressed_indices)) return A diff --git a/ext/SparseMatrixColoringsCUDAExt.jl b/ext/SparseMatrixColoringsCUDAExt.jl index ed33eece..22ca5f21 100644 --- a/ext/SparseMatrixColoringsCUDAExt.jl +++ b/ext/SparseMatrixColoringsCUDAExt.jl @@ -32,12 +32,20 @@ function SMC.StarSetColoringResult( ag::SMC.AdjacencyGraph{T}, color::Vector{<:Integer}, star_set::SMC.StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T<:Integer} + if decompression_uplo != :F + throw( + SMC.UnsupportedDecompressionError( + "Single-triangle decompression is not supported on GPU matrices" + ), + ) + end group = SMC.group_by_color(T, color) - compressed_indices = SMC.star_csc_indices(ag, color, star_set) + compressed_indices = SMC.star_csc_indices(ag, color, star_set, decompression_uplo) additional_info = (; compressed_indices_gpu_csc=CuVector(compressed_indices)) return SMC.StarSetColoringResult( - A, ag, color, group, compressed_indices, additional_info + A, ag, color, group, compressed_indices, decompression_uplo, additional_info ) end @@ -70,12 +78,20 @@ function SMC.StarSetColoringResult( ag::SMC.AdjacencyGraph{T}, color::Vector{<:Integer}, star_set::SMC.StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T<:Integer} + if decompression_uplo != :F + throw( + SMC.UnsupportedDecompressionError( + "Single-triangle decompression is not supported on GPU matrices" + ), + ) + end group = SMC.group_by_color(T, color) - compressed_indices = SMC.star_csc_indices(ag, color, star_set) + compressed_indices = SMC.star_csc_indices(ag, color, star_set, decompression_uplo) additional_info = (; compressed_indices_gpu_csr=CuVector(compressed_indices)) return SMC.StarSetColoringResult( - A, ag, color, group, compressed_indices, additional_info + A, ag, color, group, compressed_indices, decompression_uplo, additional_info ) end @@ -104,15 +120,7 @@ function SMC.decompress!( A::CuSparseMatrixCSC, B::CuMatrix, result::SMC.StarSetColoringResult{<:CuSparseMatrixCSC}, - uplo::Symbol=:F, ) - if uplo != :F - throw( - SMC.UnsupportedDecompressionError( - "Single-triangle decompression is not supported on GPU matrices" - ), - ) - end compressed_indices = result.additional_info.compressed_indices_gpu_csc copyto!(A.nzVal, view(B, compressed_indices)) return A @@ -122,15 +130,7 @@ function SMC.decompress!( A::CuSparseMatrixCSR, B::CuMatrix, result::SMC.StarSetColoringResult{<:CuSparseMatrixCSR}, - uplo::Symbol=:F, ) - if uplo != :F - throw( - SMC.UnsupportedDecompressionError( - "Single-triangle decompression is not supported on GPU matrices" - ), - ) - end compressed_indices = result.additional_info.compressed_indices_gpu_csr copyto!(A.nzVal, view(B, compressed_indices)) return A diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl index 481c8584..c8b232c3 100644 --- a/src/SparseMatrixColorings.jl +++ b/src/SparseMatrixColorings.jl @@ -26,7 +26,9 @@ using LinearAlgebra: issymmetric, ldiv!, parent, - transpose + transpose, + tril, + triu using PrecompileTools: @compile_workload using Random: Random, AbstractRNG, default_rng, randperm using SparseArrays: diff --git a/src/adtypes.jl b/src/adtypes.jl index 7c464433..2e759308 100644 --- a/src/adtypes.jl +++ b/src/adtypes.jl @@ -2,11 +2,11 @@ function coloring( A::AbstractMatrix, - problem::ColoringProblem{structure,partition}, + problem::ColoringProblem{structure,partition,uplo}, algo::ADTypes.AbstractColoringAlgorithm; decompression_eltype::Type{R}=Float64, symmetric_pattern::Bool=false, -) where {structure,partition,R} +) where {structure,partition,uplo,R} symmetric_pattern = symmetric_pattern || A isa Union{Symmetric,Hermitian} if structure == :nonsymmetric if partition == :column diff --git a/src/decompression.jl b/src/decompression.jl index 2dcf1847..d1ca763a 100644 --- a/src/decompression.jl +++ b/src/decompression.jl @@ -175,6 +175,18 @@ function decompress(B::AbstractMatrix, result::AbstractColoringResult) return decompress!(A, B, result) end +function decompress(B::AbstractMatrix, result::AbstractColoringResult{:symmetric,:column}) + A = respectful_similar(result.A, eltype(B)) + uplo = result.decompression_uplo + if A isa SparseMatrixCSC && uplo != :F + # `similar` leaves the nonzeros undefined, so they must be initialized before + # `tril`/`triu` inspects them to build the triangular pattern + fill!(nonzeros(A), zero(eltype(B))) + A = uplo == :L ? tril(A) : triu(A) + end + return decompress!(A, B, result) +end + function decompress( Br::AbstractMatrix, Bc::AbstractMatrix, @@ -187,7 +199,7 @@ end """ decompress!( A::AbstractMatrix, B::AbstractMatrix, - result::AbstractColoringResult{_,:column/:row}, [uplo=:F] + result::AbstractColoringResult{_,:column/:row} ) decompress!( @@ -204,8 +216,8 @@ The out-of-place alternative is [`decompress`](@ref). Compression means summing either the columns or the rows of `A` which share the same color. It is done by calling [`compress`](@ref). -For `:symmetric` coloring results (and for those only), an optional positional argument `uplo in (:U, :L, :F)` can be passed to specify which part of the matrix `A` should be updated: the Upper triangle, the Lower triangle, or the Full matrix. -When `A isa SparseMatrixCSC`, using the `uplo` argument requires a target matrix which only stores the relevant triangle(s). +For `:symmetric` coloring results (and for those only), the part of `A` which gets updated (the Upper triangle, the Lower triangle, or the Full matrix) is selected once and for all with the `uplo` option of [`ColoringProblem`](@ref), and it is therefore not passed at decompression time. +When `A isa SparseMatrixCSC` and `uplo != :F`, the target matrix must only store the relevant triangle. !!! warning For some coloring variants, the `result` object is mutated during decompression. @@ -260,7 +272,7 @@ function decompress! end """ decompress_single_color!( A::AbstractMatrix, b::AbstractVector, c::Integer, - result::AbstractColoringResult, [uplo=:F] + result::AbstractColoringResult ) Decompress the vector `b` corresponding to color `c` in-place into `A`, given a `:direct` coloring `result` of the sparsity pattern of `A` (it will not work with a `:substitution` coloring). @@ -272,8 +284,8 @@ Decompress the vector `b` corresponding to color `c` in-place into `A`, given a !!! warning This function will only update some coefficients of `A`, without resetting the rest to zero. -For `:symmetric` coloring results (and for those only), an optional positional argument `uplo in (:U, :L, :F)` can be passed to specify which part of the matrix `A` should be updated: the Upper triangle, the Lower triangle, or the Full matrix. -When `A isa SparseMatrixCSC`, using the `uplo` argument requires a target matrix which only stores the relevant triangle(s). +For `:symmetric` coloring results (and for those only), the part of `A` which gets updated (the Upper triangle, the Lower triangle, or the Full matrix) is selected once and for all with the `uplo` option of [`ColoringProblem`](@ref), and it is therefore not passed at decompression time. +When `A isa SparseMatrixCSC` and `uplo != :F`, the target matrix must only store the relevant triangle. !!! warning For some coloring variants, the `result` object is mutated during decompression. @@ -445,83 +457,70 @@ end ## StarSetColoringResult -function decompress!( - A::AbstractMatrix, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F -) - (; ag, compressed_indices) = result +function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::StarSetColoringResult) + (; ag, compressed_indices, decompression_uplo) = result (; S) = ag - check_compatible_pattern(A, ag, uplo) + check_compatible_pattern(A, ag, decompression_uplo) fill!(A, zero(eltype(A))) + # `compressed_indices` only holds the coefficients of the requested triangle, + # so it is indexed by a running counter `l` mirroring `star_csc_indices`, + # not by the position `k` in the full pattern of `S`. rvS = rowvals(S) + l = 0 for j in axes(S, 2) for k in nzrange(S, j) i = rvS[k] - if in_triangle(i, j, uplo) - A[i, j] = B[compressed_indices[k]] - end + in_triangle(i, j, decompression_uplo) || continue + l += 1 + A[i, j] = B[compressed_indices[l]] end end return A end function decompress_single_color!( - A::AbstractMatrix, - b::AbstractVector, - c::Integer, - result::StarSetColoringResult, - uplo::Symbol=:F, + A::AbstractMatrix, b::AbstractVector, c::Integer, result::StarSetColoringResult ) - (; ag, compressed_indices, group) = result + (; ag, compressed_indices, group, decompression_uplo) = result (; S) = ag - check_compatible_pattern(A, ag, uplo) + check_compatible_pattern(A, ag, decompression_uplo) lower_index = (c - 1) * S.n + 1 upper_index = c * S.n rvS = rowvals(S) - for j in group[c] - for k in nzrange(S, j) - # Check if the color c is used to recover A[i,j] / A[j,i] - if lower_index <= compressed_indices[k] <= upper_index - i = rvS[k] - if i == j - # Recover the diagonal coefficients of A - A[i, i] = b[i] - else - # Recover the off-diagonal coefficients of A - if in_triangle(i, j, uplo) + if decompression_uplo == :F + # `compressed_indices` is indexed by the full pattern, so we can restrict the + # traversal to the columns of color `c`, where the hub is always `j`. + for j in group[c] + for k in nzrange(S, j) + # Check if the color c is used to recover A[i,j] / A[j,i] + if lower_index <= compressed_indices[k] <= upper_index + i = rvS[k] + if i == j + # Recover the diagonal coefficients of A + A[i, i] = b[i] + else + # Recover the off-diagonal coefficients of A A[i, j] = b[i] - end - if in_triangle(j, i, uplo) A[j, i] = b[i] end end end end - end - return A -end - -function decompress!( - A::SparseMatrixCSC, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F -) - (; ag, compressed_indices) = result - (; S) = ag - nzA = nonzeros(A) - check_compatible_pattern(A, ag, uplo) - if uplo == :F - for k in eachindex(nzA, compressed_indices) - nzA[k] = B[compressed_indices[k]] - end else - rvS = rowvals(S) - l = 0 # assume A has the same pattern as the triangle + # `compressed_indices` only holds the requested triangle, so it must be walked + # with a running counter over the whole pattern. The hub is then not necessarily + # `j`, so the spoke is decoded from the stored value `(c - 1) * S.n + spoke`. + l = 0 for j in axes(S, 2) for k in nzrange(S, j) i = rvS[k] - if in_triangle(i, j, uplo) - l += 1 - nzA[l] = B[compressed_indices[k]] + in_triangle(i, j, decompression_uplo) || continue + l += 1 + index = compressed_indices[l] + if lower_index <= index <= upper_index + A[i, j] = b[index - lower_index + 1] end end end @@ -529,15 +528,33 @@ function decompress!( return A end +function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::StarSetColoringResult) + (; ag, compressed_indices, decompression_uplo) = result + nzA = nonzeros(A) + check_compatible_pattern(A, ag, decompression_uplo) + # `A` stores exactly the requested triangle, so its nonzeros are in bijection with + # `compressed_indices`, in the same order. + for k in eachindex(nzA, compressed_indices) + nzA[k] = B[compressed_indices[k]] + end + return A +end + ## TreeSetColoringResult -function decompress!( - A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F -) - (; ag, color, reverse_bfs_orders, tree_edge_indices, nt, diagonal_indices, buffer) = - result +function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult) + (; + ag, + color, + reverse_bfs_orders, + tree_edge_indices, + nt, + diagonal_indices, + buffer, + decompression_uplo, + ) = result (; S) = ag - check_compatible_pattern(A, ag, uplo) + check_compatible_pattern(A, ag, decompression_uplo) R = eltype(A) fill!(A, zero(R)) @@ -574,10 +591,10 @@ function decompress!( val = B[i, color[j]] - buffer_right_type[i] buffer_right_type[j] = buffer_right_type[j] + val - if in_triangle(i, j, uplo) + if in_triangle(i, j, decompression_uplo) A[i, j] = val end - if in_triangle(j, i, uplo) + if in_triangle(j, i, decompression_uplo) A[j, i] = val end end @@ -591,20 +608,21 @@ end A_colptr::AbstractVector, B::AbstractMatrix{R}, result::TreeSetColoringResult, - uplo::Symbol=:F, ) where {R<:Real} Decompress the values of `B` into the vector of nonzero entries `nzA` of a sparse matrix with column pointers `colptr`. This function assumes that the row indices are sorted in increasing order and are the same as those of the sparse matrix given to the `coloring` function that returned `result`. + +The triangle that gets filled in is the one selected by the [`ColoringProblem`](@ref) +which produced `result`. """ function decompress_csc!( nzA::AbstractVector{R}, A_colptr::AbstractVector{<:Integer}, B::AbstractMatrix{R}, result::TreeSetColoringResult, - uplo::Symbol=:F, ) where {R<:Real} (; ag, @@ -617,7 +635,9 @@ function decompress_csc!( lower_triangle_offsets, upper_triangle_offsets, buffer, + decompression_uplo, ) = result + uplo = decompression_uplo if eltype(buffer) == R buffer_right_type = buffer @@ -711,27 +731,28 @@ function decompress_csc!( end function decompress!( - A::SparseMatrixCSC{R}, - B::AbstractMatrix{R}, - result::TreeSetColoringResult, - uplo::Symbol=:F, + A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::TreeSetColoringResult ) where {R<:Real} - check_compatible_pattern(A, result.ag, uplo) - decompress_csc!(nonzeros(A), A.colptr, B, result, uplo) + check_compatible_pattern(A, result.ag, result.decompression_uplo) + decompress_csc!(nonzeros(A), A.colptr, B, result) return A end ## MatrixInverseColoringResult function decompress!( - A::AbstractMatrix, - B::AbstractMatrix, - result::LinearSystemColoringResult, - uplo::Symbol=:F, + A::AbstractMatrix, B::AbstractMatrix, result::LinearSystemColoringResult ) - (; ag, color, strict_upper_nonzero_inds, M_factorization, strict_upper_nonzeros_A) = - result + (; + ag, + color, + strict_upper_nonzero_inds, + M_factorization, + strict_upper_nonzeros_A, + decompression_uplo, + ) = result S = ag.S + uplo = decompression_uplo check_compatible_pattern(A, ag, uplo) # TODO: for some reason I cannot use ldiv! with a sparse QR @@ -799,7 +820,8 @@ function decompress!( nzval = Vector{R}(undef, length(large_rowval)) A_and_noAᵀ = SparseMatrixCSC(m + n, m + n, large_colptr, large_rowval, nzval) Br_and_Bc = _join_compressed!(result, Br, Bc) - decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result, :L) + # `symmetric_result` was built with `uplo = :L`, so only the lower triangle is filled + decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result) rvA = rowvals(A_and_noAᵀ) nzA = nonzeros(A_and_noAᵀ) for j in 1:n @@ -818,8 +840,8 @@ function decompress!( m, n = size(A) # pretend A is larger A_and_noAᵀ = SparseMatrixCSC(m + n, m + n, large_colptr, large_rowval, A.nzval) - # decompress lower triangle only + # decompress lower triangle only: `symmetric_result` was built with `uplo = :L` Br_and_Bc = _join_compressed!(result, Br, Bc) - decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result, :L) + decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result) return A end diff --git a/src/interface.jl b/src/interface.jl index a219186d..934b5253 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -1,4 +1,4 @@ -function check_valid_problem(structure::Symbol, partition::Symbol) +function check_valid_problem(structure::Symbol, partition::Symbol, uplo::Symbol=:F) valid = ( (structure == :nonsymmetric && partition in (:column, :row, :bidirectional)) || (structure == :symmetric && partition == :column) @@ -10,6 +10,20 @@ function check_valid_problem(structure::Symbol, partition::Symbol) ), ) end + if !(uplo in (:F, :L, :U)) + throw( + ArgumentError( + "The setting `uplo=$(repr(uplo))` is not supported by `ColoringProblem`, it must be `:F`, `:L` or `:U`.", + ), + ) + end + if uplo != :F && structure != :symmetric + throw( + ArgumentError( + "The setting `uplo=$(repr(uplo))` is only supported for `structure=:symmetric`, not `structure=$(repr(structure))`.", + ), + ) + end end function check_valid_algorithm(decompression::Symbol) @@ -24,7 +38,7 @@ function check_valid_algorithm(decompression::Symbol) end """ - ColoringProblem{structure,partition} + ColoringProblem{structure,partition,uplo} Selector type for the coloring problem to solve, enabling multiple dispatch. @@ -32,14 +46,23 @@ It is passed as an argument to the main function [`coloring`](@ref). # Constructors + ColoringProblem{structure,partition,uplo}() ColoringProblem{structure,partition}() - ColoringProblem(; structure=:nonsymmetric, partition=:column) + ColoringProblem(; structure=:nonsymmetric, partition=:column, uplo=:F) - `structure::Symbol`: either `:nonsymmetric` or `:symmetric` - `partition::Symbol`: either `:column`, `:row` or `:bidirectional` +- `uplo::Symbol`: either `:F` (full matrix), `:L` (lower triangle) or `:U` (upper triangle) + +The `uplo` setting selects which part of the matrix the decompression will fill in. +It is only supported for `structure=:symmetric`, and it defaults to `:F`, so that +`ColoringProblem{structure,partition}()` keeps working as before. + +Since `uplo` is fixed when the problem is created, a given coloring result decompresses into +exactly one triangle: build one result per triangle if you need several. !!! warning - The second constructor (based on keyword arguments) is type-unstable. + The last constructor (based on keyword arguments) is type-unstable. # Link to automatic differentiation @@ -53,11 +76,19 @@ Matrix coloring is often used in automatic differentiation, and here is the tran | Hessian | - | `:symmetric` | `:column` | yes | | Hessian | - | `:symmetric` | `:row` | no | """ -struct ColoringProblem{structure,partition} end +struct ColoringProblem{structure,partition,uplo} end -function ColoringProblem(; structure::Symbol=:nonsymmetric, partition::Symbol=:column) - check_valid_problem(structure, partition) - return ColoringProblem{structure,partition}() +# `ColoringProblem{structure,partition}` is the UnionAll `ColoringProblem{structure,partition,uplo} where uplo`, +# so this outer constructor keeps the historical two-parameter form working. +function ColoringProblem{structure,partition}() where {structure,partition} + return ColoringProblem{structure,partition,:F}() +end + +function ColoringProblem(; + structure::Symbol=:nonsymmetric, partition::Symbol=:column, uplo::Symbol=:F +) + check_valid_problem(structure, partition, uplo) + return ColoringProblem{structure,partition,uplo}() end """ @@ -283,12 +314,12 @@ end function _coloring( speed_setting::WithOrWithoutResult, A::AbstractMatrix, - ::ColoringProblem{:symmetric,:column}, + ::ColoringProblem{:symmetric,:column,uplo}, algo::GreedyColoringAlgorithm{:direct}, decompression_eltype::Type, symmetric_pattern::Bool; forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing, -) +) where {uplo} ag = AdjacencyGraph(A; augmented_graph=false, original_size=size(A)) color_and_star_set_by_order = map(algo.orders) do order vertices_in_order = vertices(ag, order) @@ -296,7 +327,7 @@ function _coloring( end color, star_set = argmin(maximum ∘ first, color_and_star_set_by_order) if speed_setting isa WithResult - return StarSetColoringResult(A, ag, color, star_set) + return StarSetColoringResult(A, ag, color, star_set, uplo) else return color end @@ -305,11 +336,11 @@ end function _coloring( speed_setting::WithOrWithoutResult, A::AbstractMatrix, - ::ColoringProblem{:symmetric,:column}, + ::ColoringProblem{:symmetric,:column,uplo}, algo::GreedyColoringAlgorithm{:substitution}, decompression_eltype::Type{R}, symmetric_pattern::Bool, -) where {R} +) where {R,uplo} ag = AdjacencyGraph(A; augmented_graph=false, original_size=size(A)) color_and_tree_set_by_order = map(algo.orders) do order vertices_in_order = vertices(ag, order) @@ -323,7 +354,7 @@ function _coloring( color, tree_set = argmin(maximum ∘ first, color_and_tree_set_by_order) end if speed_setting isa WithResult - return TreeSetColoringResult(A, ag, color, tree_set, R) + return TreeSetColoringResult(A, ag, color, tree_set, R, uplo) else return color end @@ -368,7 +399,7 @@ function _coloring( t -> maximum(t[3]) + maximum(t[4]), outputs_by_order ) # can't use ncolors without computing the full result if speed_setting isa WithResult - symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set) + symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set, :L) return BicoloringResult( A, ag, @@ -418,7 +449,7 @@ function _coloring( t -> maximum(t[3]) + maximum(t[4]), outputs_by_order ) # can't use ncolors without computing the full result if speed_setting isa WithResult - symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R) + symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R, :L) return BicoloringResult( A, ag, diff --git a/src/precompile.jl b/src/precompile.jl index 29eebaca..c7b57d0d 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -1,13 +1,15 @@ -for (structure, partition, decompression) in [ - (:nonsymmetric, :column, :direct), - (:nonsymmetric, :row, :direct), - (:symmetric, :column, :direct), - (:symmetric, :column, :substitution), - (:nonsymmetric, :bidirectional, :direct), - (:nonsymmetric, :bidirectional, :substitution), +for (structure, partition, decompression, uplo) in [ + (:nonsymmetric, :column, :direct, :F), + (:nonsymmetric, :row, :direct, :F), + (:symmetric, :column, :direct, :F), + (:symmetric, :column, :direct, :L), + (:symmetric, :column, :substitution, :F), + (:symmetric, :column, :substitution, :L), + (:nonsymmetric, :bidirectional, :direct, :F), + (:nonsymmetric, :bidirectional, :substitution, :F), ] A = sparse(Bool[1 0; 0 1]) - problem = ColoringProblem(; structure, partition) + problem = ColoringProblem(; structure, partition, uplo) algo = GreedyColoringAlgorithm(; decompression, postprocessing=true) result = coloring(A, problem, algo) if partition == :bidirectional diff --git a/src/result.jl b/src/result.jl index 489426c9..2cff7e34 100644 --- a/src/result.jl +++ b/src/result.jl @@ -312,6 +312,7 @@ struct StarSetColoringResult{ color::CT group::GT compressed_indices::VT + decompression_uplo::Symbol additional_info::A end @@ -320,43 +321,58 @@ function StarSetColoringResult( ag::AdjacencyGraph{T}, color::Vector{<:Integer}, star_set::StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T<:Integer} group = group_by_color(T, color) - compressed_indices = star_csc_indices(ag, color, star_set) - return StarSetColoringResult(A, ag, color, group, compressed_indices, nothing) + compressed_indices = star_csc_indices(ag, color, star_set, decompression_uplo) + return StarSetColoringResult( + A, ag, color, group, compressed_indices, decompression_uplo, nothing + ) end function star_csc_indices( - ag::AdjacencyGraph{T}, color::Vector{<:Integer}, star_set + ag::AdjacencyGraph{T}, + color::Vector{<:Integer}, + star_set::StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T} (; star, hub) = star_set S = pattern(ag) edge_to_index = edge_indices(ag) n = S.n rvS = rowvals(S) - compressed_indices = zeros(T, nnz(S)) # needs to be independent from the storage in the graph, in case the graph gets reused + nb_indices = nnz(S) + if decompression_uplo != :F + nb_indices = nb_edges(ag) + ag.nb_self_loops + end + compressed_indices = zeros(T, nb_indices) # needs to be independent from the storage in the graph, in case the graph gets reused + l = 0 for j in axes(S, 2) for k in nzrange(S, j) i = rvS[k] if i == j # diagonal coefficients + l += 1 c = color[i] - compressed_indices[k] = (c - 1) * n + i + compressed_indices[l] = (c - 1) * n + i else - # off-diagonal coefficients - index_ij = edge_to_index[k] - s = star[index_ij] - h = abs(hub[s]) - - # Assign the non-hub vertex (spoke) to the correct position in spokes - if i == h - # i is the hub and j is the spoke - c = color[i] - compressed_indices[k] = (c - 1) * n + j - else # j == h - # j is the hub and i is the spoke - c = color[j] - compressed_indices[k] = (c - 1) * n + i + if in_triangle(i, j, decompression_uplo) + # off-diagonal coefficients + l += 1 + index_ij = edge_to_index[k] + s = star[index_ij] + h = abs(hub[s]) + + # Assign the non-hub vertex (spoke) to the correct position in spokes + if i == h + # i is the hub and j is the spoke + c = color[i] + compressed_indices[l] = (c - 1) * n + j + else # j == h + # j is the hub and i is the spoke + c = color[j] + compressed_indices[l] = (c - 1) * n + i + end end end end @@ -394,6 +410,7 @@ struct TreeSetColoringResult{ lower_triangle_offsets::Vector{T} upper_triangle_offsets::Vector{T} buffer::Vector{R} + decompression_uplo::Symbol end function TreeSetColoringResult( @@ -402,6 +419,7 @@ function TreeSetColoringResult( color::Vector{<:Integer}, tree_set::TreeSet{<:Integer}, decompression_eltype::Type{R}, + decompression_uplo::Symbol, ) where {T<:Integer,R} (; reverse_bfs_orders, tree_edge_indices, nt) = tree_set (; S, nb_self_loops) = ag @@ -411,7 +429,7 @@ function TreeSetColoringResult( # Vector for the decompression of the diagonal coefficients diagonal_indices = Vector{T}(undef, nb_self_loops) - diagonal_nzind = Vector{T}(undef, nb_self_loops) + diagonal_nzind = (decompression_uplo == :F) ? Vector{T}(undef, nb_self_loops) : T[] if !augmented_graph(ag) l = 0 @@ -421,7 +439,9 @@ function TreeSetColoringResult( if i == j l += 1 diagonal_indices[l] = i - diagonal_nzind[l] = k + if decompression_uplo == :F + diagonal_nzind[l] = k + end end end end @@ -429,8 +449,8 @@ function TreeSetColoringResult( # Vectors for the decompression of the off-diagonal coefficients nedges = nb_edges(ag) - lower_triangle_offsets = Vector{T}(undef, nedges) - upper_triangle_offsets = Vector{T}(undef, nedges) + lower_triangle_offsets = decompression_uplo == :U ? T[] : Vector{T}(undef, nedges) + upper_triangle_offsets = decompression_uplo == :L ? T[] : Vector{T}(undef, nedges) # Index in lower_triangle_offsets and upper_triangle_offsets index_offsets = 0 @@ -454,21 +474,29 @@ function TreeSetColoringResult( if in_triangle(i, j, :L) # uplo = :L or uplo = :F # S[i,j] is stored at index_ij = (S.colptr[j+1] - offset_L) in S.nzval - lower_triangle_offsets[index_offsets] = length(col_j) - searchsortedfirst(col_j, i) + 1 + if decompression_uplo != :U + lower_triangle_offsets[index_offsets] = length(col_j) - searchsortedfirst(col_j, i) + 1 + end # uplo = :U or uplo = :F # S[j,i] is stored at index_ji = (S.colptr[i] + offset_U) in S.nzval - upper_triangle_offsets[index_offsets] = searchsortedfirst(col_i, j)::Int - 1 + if decompression_uplo != :L + upper_triangle_offsets[index_offsets] = searchsortedfirst(col_i, j)::Int - 1 + end # S[i,j] is in the upper triangular part of S else # uplo = :U or uplo = :F # S[i,j] is stored at index_ij = (S.colptr[j] + offset_U) in S.nzval - upper_triangle_offsets[index_offsets] = searchsortedfirst(col_j, i)::Int - 1 + if decompression_uplo != :L + upper_triangle_offsets[index_offsets] = searchsortedfirst(col_j, i)::Int - 1 + end # uplo = :L or uplo = :F # S[j,i] is stored at index_ji = (S.colptr[i+1] - offset_L) in S.nzval - lower_triangle_offsets[index_offsets] = length(col_i) - searchsortedfirst(col_i, j) + 1 + if decompression_uplo != :U + lower_triangle_offsets[index_offsets] = length(col_i) - searchsortedfirst(col_i, j) + 1 + end end #! format: on end @@ -491,6 +519,7 @@ function TreeSetColoringResult( lower_triangle_offsets, upper_triangle_offsets, buffer, + decompression_uplo, ) end @@ -521,6 +550,7 @@ struct LinearSystemColoringResult{ strict_upper_nonzero_inds::Vector{Tuple{T,T}} strict_upper_nonzeros_A::Vector{R} # TODO: adjust type M_factorization::F # TODO: adjust type + decompression_uplo::Symbol end function LinearSystemColoringResult( @@ -528,6 +558,7 @@ function LinearSystemColoringResult( ag::AdjacencyGraph{T}, color::Vector{<:Integer}, decompression_eltype::Type{R}, + decompression_uplo::Symbol=:F, ) where {T<:Integer,R<:Real} group = group_by_color(T, color) C = length(group) # ncolors @@ -572,6 +603,7 @@ function LinearSystemColoringResult( strict_upper_nonzero_inds, strict_upper_nonzeros_A, M_factorization, + decompression_uplo, ) end diff --git a/test/allocations.jl b/test/allocations.jl index 31305e63..03da40e4 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -62,11 +62,20 @@ function test_noallocs_sparse_decompression( @test minimum(bench2_singlecolor).allocs == 0 end end + # the triangle is selected by the problem, so it needs its own result + result_U = if structure == :symmetric + coloring( + A, + ColoringProblem(; structure, partition, uplo=:U), + GreedyColoringAlgorithm(; decompression), + ) + else + nothing + end @testset "Triangle decompression" begin if structure == :symmetric - bench1_triangle = @be similar(triu(A)) decompress!(_, B, result, :U) evals = - 1 - bench2_triangle = @be similar(Matrix(A)) decompress!(_, B, result, :U) evals = + bench1_triangle = @be similar(triu(A)) decompress!(_, B, result_U) evals = 1 + bench2_triangle = @be similar(Matrix(A)) decompress!(_, B, result_U) evals = 1 @test minimum(bench1_triangle).allocs == 0 @test minimum(bench2_triangle).allocs == 0 @@ -76,10 +85,10 @@ function test_noallocs_sparse_decompression( if structure == :symmetric && decompression == :direct b = B[:, 1] bench1_singlecolor_triangle = @be similar(triu(A)) decompress_single_color!( - _, b, 1, result, :U + _, b, 1, result_U ) evals = 1 bench2_singlecolor_triangle = @be similar(Matrix(A)) decompress_single_color!( - _, b, 1, result, :U + _, b, 1, result_U ) evals = 1 @test minimum(bench1_singlecolor_triangle).allocs == 0 @test minimum(bench2_singlecolor_triangle).allocs == 0 diff --git a/test/constructors.jl b/test/constructors.jl index 8ce88bcb..4e3c636c 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -5,9 +5,28 @@ using Test @test ColoringProblem{:symmetric,:column}() == ColoringProblem(; structure=:symmetric, partition=:column) +# the two-parameter form still works and defaults to the full matrix +@test ColoringProblem{:nonsymmetric,:column,:F}() == + ColoringProblem{:nonsymmetric,:column}() +@test ColoringProblem{:symmetric,:column,:U}() == + ColoringProblem(; structure=:symmetric, partition=:column, uplo=:U) +@test ColoringProblem{:symmetric,:column,:L}() == + ColoringProblem(; structure=:symmetric, partition=:column, uplo=:L) + @test_throws ArgumentError ColoringProblem(; structure=:weird, partition=:column) @test_throws ArgumentError ColoringProblem(; structure=:symmetric, partition=:row) +# uplo must be one of :F, :L, :U and only makes sense for symmetric problems +@test_throws ArgumentError ColoringProblem(; + structure=:symmetric, partition=:column, uplo=:weird +) +@test_throws ArgumentError ColoringProblem(; + structure=:nonsymmetric, partition=:column, uplo=:L +) +@test_throws ArgumentError ColoringProblem(; + structure=:nonsymmetric, partition=:bidirectional, uplo=:L +) + @test GreedyColoringAlgorithm{:direct}() == GreedyColoringAlgorithm() @test GreedyColoringAlgorithm{:substitution}() == GreedyColoringAlgorithm(; decompression=:substitution) diff --git a/test/cuda.jl b/test/cuda.jl index 6ae73060..d9b8123c 100644 --- a/test/cuda.jl +++ b/test/cuda.jl @@ -52,11 +52,9 @@ end; A0 = T(sparse(Symmetric(sprand(rng, n, n, p)))) test_coloring_decompression(A0, problem, algo; gpu=true) end + # the triangle is now selected by the problem, so the error fires at coloring time A0 = T(sparse(Diagonal(ones(10)))) - result = coloring(A0, problem, algo) - B = compress(A0, result) - @test_throws SMC.UnsupportedDecompressionError decompress!( - similar(A0), B, result, :U - ) + problem_U = ColoringProblem(; structure=:symmetric, partition=:column, uplo=:U) + @test_throws SMC.UnsupportedDecompressionError coloring(A0, problem_U, algo) end end; diff --git a/test/rocm.jl b/test/rocm.jl index e2aa7840..3fcebaad 100644 --- a/test/rocm.jl +++ b/test/rocm.jl @@ -52,11 +52,9 @@ end; A0 = T(sparse(Symmetric(sprand(rng, n, n, p)))) test_coloring_decompression(A0, problem, algo; gpu=true) end + # the triangle is now selected by the problem, so the error fires at coloring time A0 = T(sparse(Diagonal(ones(10)))) - result = coloring(A0, problem, algo) - B = compress(A0, result) - @test_throws SMC.UnsupportedDecompressionError decompress!( - similar(A0), B, result, :U - ) + problem_U = ColoringProblem(; structure=:symmetric, partition=:column, uplo=:U) + @test_throws SMC.UnsupportedDecompressionError coloring(A0, problem_U, algo) end end; diff --git a/test/type_stability.jl b/test/type_stability.jl index c9ea92ae..1e6840ce 100644 --- a/test/type_stability.jl +++ b/test/type_stability.jl @@ -67,6 +67,24 @@ end end end + @testset "Triangle problems" begin + # the type-parameter constructor is used on purpose: the keyword one is type-unstable + @testset "$decompression - $uplo" for decompression in (:direct, :substitution), + uplo in (:L, :U) + + @test_opt coloring( + A, + ColoringProblem{:symmetric,:column,uplo}(), + GreedyColoringAlgorithm(; decompression), + ) + @inferred coloring( + A, + ColoringProblem{:symmetric,:column,uplo}(), + GreedyColoringAlgorithm(; decompression), + ) + end + end + @testset "Explicit decompression_eltype" begin @testset "$structure - $partition - $decompression - $R" for ( structure, partition, decompression @@ -177,13 +195,25 @@ end; end @testset "Triangle decompression" begin if structure == :symmetric - @test_opt decompress!(respectful_similar(triu(A)), B, result, :U) + result_U = coloring( + A0, + ColoringProblem{structure,partition,:U}(), + GreedyColoringAlgorithm(; decompression); + decompression_eltype=eltype(A), + ) + @test_opt decompress!(respectful_similar(triu(A)), B, result_U) end end @testset "Single-color triangle decompression" begin if structure == :symmetric && decompression == :direct + result_U = coloring( + A0, + ColoringProblem{structure,partition,:U}(), + GreedyColoringAlgorithm(; decompression); + decompression_eltype=eltype(A), + ) @test_opt decompress_single_color!( - respectful_similar(triu(A)), B[:, 1], 1, result, :U + respectful_similar(triu(A)), B[:, 1], 1, result_U ) end end diff --git a/test/utils.jl b/test/utils.jl index 98c12afa..1caf8345 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -24,13 +24,13 @@ const _ALL_ORDERS = ( function test_coloring_decompression( A0::AbstractMatrix, - problem::ColoringProblem{structure,partition}, + problem::ColoringProblem{structure,partition,uplo}, algo::GreedyColoringAlgorithm{decompression}; B0=nothing, color0=nothing, test_fast=false, gpu=false, -) where {structure,partition,decompression} +) where {structure,partition,uplo,decompression} color_vec = Vector{Int}[] @testset "$(typeof(A))" for A in matrix_versions(A0) yield() @@ -124,8 +124,30 @@ function test_coloring_decompression( end end + # The triangle is now fixed when the problem is built, so one result decompresses + # into exactly one triangle: we need a sibling result per triangle. + if structure == :symmetric && uplo == :F + result_L = coloring( + A, + ColoringProblem{structure,partition,:L}(), + algo; + decompression_eltype=Float64, + ) + result_U = coloring( + A, + ColoringProblem{structure,partition,:U}(), + algo; + decompression_eltype=Float64, + ) + end + @testset "Triangle decompression" begin - if structure == :symmetric + if structure == :symmetric && uplo == :F + # the three problems must agree on the colors, otherwise the single `B` + # computed above would not be valid for the sibling results + @test column_colors(result_L) == color + @test column_colors(result_U) == color + A3upper = respectful_similar(triu(A), eltype(B)) A3lower = respectful_similar(tril(A), eltype(B)) A3both = respectful_similar(A, eltype(B)) @@ -133,18 +155,22 @@ function test_coloring_decompression( A3lower .= zero(eltype(A)) A3both .= zero(eltype(A)) - decompress!(A3upper, B, result, :U) - decompress!(A3lower, B, result, :L) - decompress!(A3both, B, result, :F) + decompress!(A3upper, B, result_U) + decompress!(A3lower, B, result_L) + decompress!(A3both, B, result) @test A3upper ≈ triu(A0) @test A3lower ≈ tril(A0) @test A3both ≈ A0 + + # out-of-place must allocate a target with the right triangular pattern + @test decompress(B, result_U) ≈ triu(A0) + @test decompress(B, result_L) ≈ tril(A0) end end @testset "Single-color triangle decompression" begin - if structure == :symmetric && decompression == :direct + if structure == :symmetric && uplo == :F && decompression == :direct A4upper = respectful_similar(triu(A), eltype(B)) A4lower = respectful_similar(tril(A), eltype(B)) A4both = respectful_similar(A, eltype(B)) @@ -154,9 +180,9 @@ function test_coloring_decompression( for c in unique(color) c == 0 && continue - decompress_single_color!(A4upper, B[:, c], c, result, :U) - decompress_single_color!(A4lower, B[:, c], c, result, :L) - decompress_single_color!(A4both, B[:, c], c, result, :F) + decompress_single_color!(A4upper, B[:, c], c, result_U) + decompress_single_color!(A4lower, B[:, c], c, result_L) + decompress_single_color!(A4both, B[:, c], c, result) end @test A4upper ≈ triu(A0)