Optimal sparse decompression (#223)

* Optimal sparse decompression

* Fix reverse mode dispatch with major-respecting similar

Author: gdalle

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -31,7 +31,7 @@ using ADTypes:
 using DocStringExtensions
 using FillArrays: OneElement
 using LinearAlgebra: Symmetric, Transpose, dot, parent, transpose
-using SparseArrays: SparseMatrixCSC, nzrange, rowvals, sparse
+using SparseArrays: SparseMatrixCSC, nonzeros, nzrange, rowvals, sparse
 
 abstract type Extras end
 

DifferentiationInterface/src/sparse/compressed_matrix.jl

@@ -1,20 +1,16 @@
 """
     CompressedMatrix{dir}
 
-Compressed representation `B` of a sparse matrix `A ∈ ℝ^{m×n}` obtained by summing some of its columns (if `dir == :col`) or rows (if `dir == :row`), grouped by color.
+Compressed representation `B` of a `(m, n)` sparse matrix `A` obtained by summing some of its columns (if `dir == :col`) or rows (if `dir == :row`) if they have the same color.
 
 # Fields
 
-- `sparsity::AbstractMatrix{Bool}`: boolean sparsity pattern of the matrix `A`
-- `colors::Vector{Int}`: vector such that
-  - if `dir == `:col`, then `colors[j] ∈ 1:c` is the color of column `j`
-  - if `dir == `:row`, then `colors[i] ∈ 1:c` is the color of row `i`
-- `groups::Vector{Vector{Int}}`: vector of length `c` such that
-  - if `dir == :col`, then `groups[k]` is the vector of column indices assigned to the same color `k ∈ 1:c`
-  - if `dir == :row`, then `groups[k]` is the vector of row indices assigned to the same color `k ∈ 1:c`
-- `aggregates::AbstractMatrix`: matrix `B` such that
-  - if `dir == :col`, then `size(B) = (m, c)` and `B[:, c] = sum(A[:, k] for k in groups[c])`
-  - if `dir == :row`, then `size(B) = (c, n)` and `B[c, :] = sum(A[k, :] for k in groups[c])`
+| field        | type                     | size                 | meaning                    | if `dir` is `:col`                | if `dir` is `:row`                |
+| :----------- | :----------------------- | :------------------- | :------------------------- | :-------------------------------- | :-------------------------------- |
+| `sparsity`   | `AbstractMatrix{Bool}`   | `(m, n)`             | sparsity pattern        | column-major                      | row-major                         |
+| `colors`     | `Vector{Int}`            | `n` or `m`           | color assignments in `1:c` | `colors[j]` of col `j`            | `colors[i]` of row `i`            |
+| `groups`     | `Vector{Vector{Int}}`    | `c`                  | groups with same color     | `groups[k] = {j : colors[j] = k}` | `groups[k] = {i : colors[i] = k}` |
+| `aggregates` | `AbstractMatrix{<:Real}` | `(m, c)` or `(c, n)` | color-summed values `B`    | `B[:, c] = sum(A[:, groups[k]])`  | `B[c, :] = sum(A[groups[k], :])`  |
 """
 mutable struct CompressedMatrix{dir,S<:AbstractMatrix{Bool},M<:AbstractMatrix}
     sparsity::S
@@ -35,39 +31,62 @@ function CompressedMatrix{dir}(sparsity, colors, groups, aggregates) where {dir}
     )
 end
 
+## Column decompression
+
 function decompress!(A::AbstractMatrix, compressed::CompressedMatrix{:col})
     (; sparsity, colors, aggregates) = compressed
     A .= zero(eltype(A))
     @views for j in axes(A, 2)
         k = colors[j]
-        nz_rows_j = (!iszero).(sparsity[:, j])
-        copyto!(A[nz_rows_j, j], aggregates[nz_rows_j, k])
+        rows_j = (!iszero).(sparsity[:, j])
+        copyto!(A[rows_j, j], aggregates[rows_j, k])
+    end
+    return A
+end
+
+function decompress!(
+    A::SparseMatrixCSC, compressed::CompressedMatrix{:col,<:SparseMatrixCSC}
+)
+    # A and compressed.sparsity have the same pattern
+    (; colors, aggregates) = compressed
+    Anz, Arv = nonzeros(A), rowvals(A)
+    Anz .= zero(eltype(A))
+    @views for j in axes(A, 2)
+        k = colors[j]
+        nzrange_j = nzrange(A, j)
+        rows_j = Arv[nzrange_j]
+        copyto!(Anz[nzrange_j], aggregates[rows_j, k])
     end
     return A
 end
 
+## Row decompression
+
 function decompress!(A::AbstractMatrix, compressed::CompressedMatrix{:row})
     (; sparsity, colors, aggregates) = compressed
     A .= zero(eltype(A))
     @views for i in axes(A, 1)
         k = colors[i]
-        nz_cols_i = (!iszero).(sparsity[i, :])
-        copyto!(A[i, nz_cols_i], aggregates[k, nz_cols_i])
+        cols_i = (!iszero).(sparsity[i, :])
+        copyto!(A[i, cols_i], aggregates[k, cols_i])
     end
     return A
 end
 
-function decompress_symmetric!(A::AbstractMatrix, compressed::CompressedMatrix{:col})
-    (; sparsity, colors, groups, aggregates) = compressed
-    @views for j in axes(A, 2)
-        k = colors[j]
-        group = groups[k]
-        for i in axes(A, 1)
-            if (!iszero(sparsity[i, j]) && count(!iszero, sparsity[i, group]) == 1)
-                A[i, j] = aggregates[i, k]
-                A[j, i] = aggregates[i, k]
-            end
-        end
+function decompress!(
+    A::Transpose{<:Any,<:SparseMatrixCSC},
+    compressed::CompressedMatrix{:row,<:Transpose{<:Any,<:SparseMatrixCSC}},
+)
+    # A and compressed.sparsity have the same pattern
+    (; colors, aggregates) = compressed
+    PA = parent(A)
+    PAnz, PArv = nonzeros(PA), rowvals(PA)
+    PAnz .= zero(eltype(A))
+    @views for i in axes(A, 1)
+        k = colors[i]
+        nzrange_i = nzrange(PA, i)
+        cols_i = PArv[nzrange_i]
+        copyto!(PAnz[nzrange_i], aggregates[k, cols_i])
     end
     return A
 end

DifferentiationInterface/src/sparse/hessian.jl

@@ -10,8 +10,9 @@ end
 ## Hessian, one argument
 
 function prepare_hessian(f, backend::AutoSparse, x)
-    sparsity = hessian_sparsity(f, x, sparsity_detector(backend))
-    colors = symmetric_coloring(sparsity, coloring_algorithm(backend))
+    initial_sparsity = hessian_sparsity(f, x, sparsity_detector(backend))
+    sparsity = col_major(initial_sparsity)
+    colors = column_coloring(sparsity, coloring_algorithm(backend))
     groups = get_groups(colors)
     seeds = map(groups) do group
         seed = zero(x)
@@ -33,7 +34,7 @@ function hessian!(f, hess, backend::AutoSparse, x, extras::SparseHessianExtras)
         hvp!(f, products[k], backend, x, seeds[k], hvp_extras)
         copyto!(view(compressed.aggregates, :, k), vec(products[k]))
     end
-    decompress_symmetric!(hess, compressed)
+    decompress!(hess, compressed)
     return hess
 end
 

DifferentiationInterface/src/sparse/jacobian.jl

@@ -23,8 +23,9 @@ end
 
 function prepare_jacobian(f, backend::AutoSparse, x)
     y = f(x)
-    sparsity = jacobian_sparsity(f, x, sparsity_detector(backend))
+    initial_sparsity = jacobian_sparsity(f, x, sparsity_detector(backend))
     if Bool(pushforward_performance(backend))
+        sparsity = col_major(initial_sparsity)
         colors = column_coloring(sparsity, coloring_algorithm(backend))
         groups = get_groups(colors)
         seeds = map(groups) do group
@@ -39,6 +40,7 @@ function prepare_jacobian(f, backend::AutoSparse, x)
         aggregates = stack(vec, products; dims=2)
         compressed = CompressedMatrix{:col}(sparsity, colors, groups, aggregates)
     else
+        sparsity = row_major(initial_sparsity)
         colors = row_coloring(sparsity, coloring_algorithm(backend))
         groups = get_groups(colors)
         seeds = map(groups) do group
@@ -77,7 +79,7 @@ function jacobian!(f, jac, backend::AutoSparse, x, extras::SparseJacobianExtras{
 end
 
 function jacobian(f, backend::AutoSparse, x, extras::SparseJacobianExtras{1})
-    jac = similar(extras.compressed.sparsity, eltype(x))
+    jac = major_respecting_similar(extras.compressed.sparsity, eltype(x))
     return jacobian!(f, jac, backend, x, extras)
 end
 
@@ -94,8 +96,9 @@ end
 ## Jacobian, two arguments
 
 function prepare_jacobian(f!, y, backend::AutoSparse, x)
-    sparsity = jacobian_sparsity(f!, y, x, sparsity_detector(backend))
+    initial_sparsity = jacobian_sparsity(f!, y, x, sparsity_detector(backend))
     if Bool(pushforward_performance(backend))
+        sparsity = col_major(initial_sparsity)
         colors = column_coloring(sparsity, coloring_algorithm(backend))
         groups = get_groups(colors)
         seeds = map(groups) do group
@@ -110,6 +113,7 @@ function prepare_jacobian(f!, y, backend::AutoSparse, x)
         aggregates = stack(vec, products; dims=2)
         compressed = CompressedMatrix{:col}(sparsity, colors, groups, aggregates)
     else
+        sparsity = row_major(initial_sparsity)
         colors = row_coloring(sparsity, coloring_algorithm(backend))
         groups = get_groups(colors)
         seeds = map(groups) do group
@@ -148,7 +152,7 @@ function jacobian!(f!, y, jac, backend::AutoSparse, x, extras::SparseJacobianExt
 end
 
 function jacobian(f!, y, backend::AutoSparse, x, extras::SparseJacobianExtras{2})
-    jac = similar(extras.compressed.sparsity, eltype(x))
+    jac = major_respecting_similar(extras.compressed.sparsity, eltype(x))
     return jacobian!(f!, y, jac, backend, x, extras)
 end
 

DifferentiationInterface/src/sparse/matrices.jl

@@ -16,6 +16,14 @@ Construct a row-major representation of the matrix `A`.
 row_major(A::M) where {M<:AbstractMatrix} = transpose(M(transpose(A)))
 row_major(A::Transpose{<:Any,M}) where {M<:AbstractMatrix} = A
 
+## Similar
+
+major_respecting_similar(A::AbstractMatrix, ::Type{T}) where {T} = similar(A, T)
+
+function major_respecting_similar(A::Transpose, ::Type{T}) where {T}
+    return transpose(similar(parent(A), T))
+end
+
 ## Generic nz
 
 function nz_in_col(A_colmajor::AbstractMatrix, j::Integer)

DifferentiationInterface/test/sparsity.jl

@@ -5,7 +5,7 @@ sparse_backends = [
     AutoSparse(AutoFastDifferentiation()),
     AutoSparse(AutoSymbolics()),
     AutoSparse(AutoForwardDiff(); sparsity_detector, coloring_algorithm),
-    AutoSparse(AutoZygote(); sparsity_detector, coloring_algorithm),
+    AutoSparse(AutoEnzyme(Enzyme.Reverse); sparsity_detector, coloring_algorithm),
 ]
 
 sparse_second_order_backends = [

DifferentiationInterfaceTest/src/tests/sparsity.jl

@@ -14,10 +14,6 @@ function test_sparsity(
     _, jac1 = value_and_jacobian(f, ba, x, extras)
     jac2 = jacobian(f, ba, x, extras)
 
-    @testset "Sparse type" begin
-        @test jac1 isa SparseMatrixCSC
-        @test jac2 isa SparseMatrixCSC
-    end
     @testset "Sparsity pattern" begin
         @test nnz(jac1) == nnz(jac_true)
         @test nnz(jac2) == nnz(jac_true)
@@ -37,10 +33,6 @@ function test_sparsity(ba::AbstractADType, scen::JacobianScenario{1,:inplace}; r
     _, jac1 = value_and_jacobian!(f, mysimilar(jac_true), ba, x, extras)
     jac2 = jacobian!(f, mysimilar(jac_true), ba, x, extras)
 
-    @testset "Sparse type" begin
-        @test jac1 isa SparseMatrixCSC
-        @test jac2 isa SparseMatrixCSC
-    end
     @testset "Sparsity pattern" begin
         @test nnz(jac1) == nnz(jac_true)
         @test nnz(jac2) == nnz(jac_true)
@@ -63,10 +55,6 @@ function test_sparsity(
     _, jac1 = value_and_jacobian(f!, mysimilar(y), ba, x, extras)
     jac2 = jacobian(f!, mysimilar(y), ba, x, extras)
 
-    @testset "Sparse type" begin
-        @test jac1 isa SparseMatrixCSC
-        @test jac2 isa SparseMatrixCSC
-    end
     @testset "Sparsity pattern" begin
         @test nnz(jac1) == nnz(jac_true)
         @test nnz(jac2) == nnz(jac_true)
@@ -87,10 +75,6 @@ function test_sparsity(ba::AbstractADType, scen::JacobianScenario{2,:inplace}; r
     _, jac1 = value_and_jacobian!(f!, mysimilar(y), mysimilar(jac_true), ba, x, extras)
     jac2 = jacobian!(f!, mysimilar(y), mysimilar(jac_true), ba, x, extras)
 
-    @testset "Sparse type" begin
-        @test jac1 isa SparseMatrixCSC
-        @test jac2 isa SparseMatrixCSC
-    end
     @testset "Sparsity pattern" begin
         @test nnz(jac1) == nnz(jac_true)
         @test nnz(jac2) == nnz(jac_true)
@@ -113,9 +97,6 @@ function test_sparsity(
 
     hess1 = hessian(f, ba, x, extras)
 
-    @testset "Sparse type" begin
-        @test hess1 isa SparseMatrixCSC
-    end
     @testset "Sparsity pattern" begin
         @test nnz(hess1) == nnz(hess_true)
     end
@@ -133,9 +114,6 @@ function test_sparsity(ba::AbstractADType, scen::HessianScenario{1,:inplace}; re
 
     hess1 = hessian!(f, mysimilar(hess_true), ba, x, extras)
 
-    @testset "Sparse type" begin
-        @test hess1 isa SparseMatrixCSC
-    end
     @testset "Sparsity pattern" begin
         @test nnz(hess1) == nnz(hess_true)
     end