Support decompression in one triangle

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]

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.

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

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

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:

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