Add systematic allocation and type stability tests (#84)

* Add systematic allocation and type stability tests

* Typo

Author: gdalle

src/decompression.jl

@@ -284,7 +284,6 @@ function decompress_single_color!(
 ) where {R<:Real}
     @compat (; S, group) = result
     check_same_pattern(A, S)
-    view(A, :, group[c]) .= zero(R)
     rvS = rowvals(S)
     for j in group[c]
         for k in nzrange(S, j)
@@ -347,7 +346,6 @@ function decompress_single_color!(
 ) where {R<:Real}
     @compat (; S, Sᵀ, group) = result
     check_same_pattern(A, S)
-    view(A, group[c], :) .= zero(R)
     rvSᵀ = rowvals(Sᵀ)
     for i in group[c]
         for k in nzrange(Sᵀ, i)

test/allocations.jl

@@ -0,0 +1,89 @@
+using Chairmarks
+using LinearAlgebra
+using SparseArrays
+using SparseMatrixColorings
+using SparseMatrixColorings: partial_distance2_coloring!
+using StableRNGs
+using Test
+
+rng = StableRNG(63)
+
+function test_noallocs_distance2_coloring(n)
+    bench = @be (;
+        bg=bipartite_graph(sprand(rng, n, n, 5 / n)),
+        color=Vector{Int}(undef, n),
+        forbidden_colors=Vector{Int}(undef, n),
+    ) partial_distance2_coloring!(_.color, _.forbidden_colors, _.bg, Val(1), 1:n) evals = 1
+    @test minimum(bench).allocs == 0
+end
+
+@testset "Distance-2 coloring" begin
+    test_noallocs_distance2_coloring(1000)
+end;
+
+function test_noallocs_decompression(
+    n::Integer; structure::Symbol, partition::Symbol, decompression::Symbol
+)
+    A = sparse(Symmetric(sprand(rng, n, n, 5 / n)))
+    result = coloring(
+        A, ColoringProblem(; structure, partition), GreedyColoringAlgorithm(; decompression)
+    )
+    B = compress(A, result)
+
+    @testset "Full decompression" begin
+        bench1_full = @be similar(A) decompress!(_, B, result) evals = 1
+        bench2_full = @be similar(Matrix(A)) decompress!(_, B, result) evals = 1
+        @test minimum(bench1_full).allocs == 0
+        @test minimum(bench2_full).allocs == 0
+    end
+    @testset "Single-color decompression" begin
+        if decompression == :direct
+            b = if partition == :column
+                B[:, 1]
+            else
+                B[1, :]
+            end
+            bench1_singlecolor = @be similar(A) decompress_single_color!(_, b, 1, result) evals =
+                1
+            bench2_singlecolor = @be similar(Matrix(A)) decompress_single_color!(
+                _, b, 1, result
+            ) evals = 1
+            @test minimum(bench1_singlecolor).allocs == 0
+            @test minimum(bench2_singlecolor).allocs == 0
+        end
+    end
+    @testset "Triangle decompression" begin
+        if structure == :symmetric
+            bench1_triangle = @be similar(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
+        end
+    end
+    @testset "Single-color triangle decompression" begin
+        if structure == :symmetric && decompression == :direct
+            b = B[:, 1]
+            bench1_singlecolor_triangle = @be similar(A) decompress_single_color!(
+                _, b, 1, result, :U
+            ) evals = 1
+            bench2_singlecolor_triangle = @be similar(Matrix(A)) decompress_single_color!(
+                _, b, 1, result, :U
+            ) evals = 1
+            @test minimum(bench1_singlecolor_triangle).allocs == 0
+            @test minimum(bench2_singlecolor_triangle).allocs == 0
+        end
+    end
+end
+
+@testset "Decompression" begin
+    @testset "$structure - $partition - $decompression" for (
+        structure, partition, decompression
+    ) in [
+        (:nonsymmetric, :column, :direct),
+        (:nonsymmetric, :row, :direct),
+        (:symmetric, :column, :direct),
+        (:symmetric, :column, :substitution),
+    ]
+        test_noallocs_decompression(1000; structure, partition, decompression)
+    end
+end;

test/performance.jl

@@ -59,7 +59,12 @@ include("utils.jl")
     end
     @testset verbose = true "Performance" begin
         if VERSION >= v"1.10"
-            include("performance.jl")
+            @testset "Type stability" begin
+                include("type_stability.jl")
+            end
+        end
+        @testset "Allocations" begin
+            include("allocations.jl")
         end
     end
 end

test/runtests.jl

@@ -0,0 +1,94 @@
+using ADTypes: column_coloring, row_coloring, symmetric_coloring
+using JET
+using LinearAlgebra
+using SparseArrays
+using SparseMatrixColorings
+using SparseMatrixColorings: respectful_similar
+using StableRNGs
+using Test
+
+rng = StableRNG(63)
+
+@testset "Coloring" begin
+    n = 10
+    A = sprand(rng, n, n, 5 / n)
+
+    # ADTypes
+    @testset "ADTypes" begin
+        @test_opt target_modules = (SparseMatrixColorings,) column_coloring(
+            A, GreedyColoringAlgorithm()
+        )
+        @test_opt target_modules = (SparseMatrixColorings,) row_coloring(
+            A, GreedyColoringAlgorithm()
+        )
+        @test_opt target_modules = (SparseMatrixColorings,) symmetric_coloring(
+            Symmetric(A), GreedyColoringAlgorithm()
+        )
+    end
+
+    @testset "$structure - $partition - $decompression" for (
+        structure, partition, decompression
+    ) in [
+        (:nonsymmetric, :column, :direct),
+        (:nonsymmetric, :row, :direct),
+        (:symmetric, :column, :direct),
+        (:symmetric, :column, :substitution),
+    ]
+        @test_opt target_modules = (SparseMatrixColorings,) coloring(
+            A,
+            ColoringProblem(; structure, partition),
+            GreedyColoringAlgorithm(; decompression),
+        )
+    end
+end
+
+@testset "Decompression" begin
+    n = 10
+    A0 = sparse(Symmetric(sprand(rng, n, n, 5 / n)))
+
+    @testset "$structure - $partition - $decompression" for (
+        structure, partition, decompression
+    ) in [
+        (:nonsymmetric, :column, :direct),
+        (:nonsymmetric, :row, :direct),
+        (:symmetric, :column, :direct),
+        (:symmetric, :column, :substitution),
+    ]
+        @testset "A::$(typeof(A))" for A in matrix_versions(A0)
+            result = coloring(
+                A0,
+                ColoringProblem(; structure, partition),
+                GreedyColoringAlgorithm(; decompression);
+                decompression_eltype=eltype(A),
+            )
+            B = compress(A, result)
+            @testset "Full decompression" begin
+                @test_opt compress(A, result)
+                @test_opt decompress(B, result) ≈ A0
+                @test_opt decompress!(respectful_similar(A), B, result)
+            end
+            @testset "Single-color decompression" begin
+                if decompression == :direct
+                    b = if partition == :column
+                        B[:, 1]
+                    else
+                        B[1, :]
+                    end
+                    @test_opt decompress_single_color!(respectful_similar(A), b, 1, result)
+                end
+            end
+            @testset "Triangle decompression" begin
+                if structure == :symmetric
+                    @test_opt decompress!(respectful_similar(A), B, result, :U)
+                end
+            end
+            @testset "Single-color triangle decompression" begin
+                if structure == :symmetric && decompression == :direct
+                    @test_opt decompress_single_color!(
+                        respectful_similar(A), B[:, 1], 1, result, :U
+                    )
+                end
+            end
+        end
+    end
+end;