Type-stable compress and random order in benchmarks (#198)

* Type-stable compress and random order in benchmarks

* Ensure zero size

* Mat

Author: gdalle

benchmark/benchmarks.jl

@@ -19,7 +19,9 @@ for structure in [:nonsymmetric, :symmetric],
     p in [2 / n, 5 / n, 10 / n]
 
     problem = ColoringProblem(; structure, partition)
-    algo = GreedyColoringAlgorithm(; decompression, postprocessing=true)
+    algo = GreedyColoringAlgorithm(
+        RandomOrder(StableRNG(0), 0); decompression, postprocessing=true
+    )
 
     # use several random matrices to reduce variance
     nb_samples = 5

src/decompression.jl

@@ -47,7 +47,11 @@ function compress end
 function compress(A, result::AbstractColoringResult{structure,:column}) where {structure}
     group = column_groups(result)
     if isempty(group)
-        B = similar(A, size(A, 1), 0)
+        # ensure we get a Matrix and not a SparseMatrixCSC
+        B_model = stack([Int[]]; dims=2) do g
+            dropdims(sum(A[:, g]; dims=2); dims=2)
+        end
+        B = similar(B_model, size(A, 1), 0)
     else
         B = stack(group; dims=2) do g
             dropdims(sum(A[:, g]; dims=2); dims=2)
@@ -59,7 +63,11 @@ end
 function compress(A, result::AbstractColoringResult{structure,:row}) where {structure}
     group = row_groups(result)
     if isempty(group)
-        B = similar(A, 0, size(A, 2))
+        # ensure we get a Matrix and not a SparseMatrixCSC
+        B_model = stack([Int[]]; dims=1) do g
+            dropdims(sum(A[g, :]; dims=1); dims=1)
+        end
+        B = similar(B_model, 0, size(A, 2))
     else
         B = stack(group; dims=1) do g
             dropdims(sum(A[g, :]; dims=1); dims=1)
@@ -74,14 +82,22 @@ function compress(
     row_group = row_groups(result)
     column_group = column_groups(result)
     if isempty(row_group)
-        Br = similar(A, 0, size(A, 2))
+        # ensure we get a Matrix and not a SparseMatrixCSC
+        Br_model = stack([Int[]]; dims=1) do g
+            dropdims(sum(A[g, :]; dims=1); dims=1)
+        end
+        Br = similar(Br_model, 0, size(A, 2))
     else
         Br = stack(row_group; dims=1) do g
             dropdims(sum(A[g, :]; dims=1); dims=1)
         end
     end
     if isempty(column_group)
-        Bc = similar(A, size(A, 1), 0)
+        # ensure we get a Matrix and not a SparseMatrixCSC
+        Bc_model = stack([Int[]]; dims=2) do g
+            dropdims(sum(A[:, g]; dims=2); dims=2)
+        end
+        Bc = similar(Bc_model, size(A, 1), 0)
     else
         Bc = stack(column_group; dims=2) do g
             dropdims(sum(A[:, g]; dims=2); dims=2)

test/type_stability.jl

@@ -11,19 +11,17 @@ rng = StableRNG(63)
 
 @testset "Sparse coloring" begin
     n = 10
-    A = sprand(rng, n, n, 5 / n)
+    A = sparse(Symmetric(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()
-        )
+        @test_opt column_coloring(A, GreedyColoringAlgorithm())
+        @test_opt row_coloring(A, GreedyColoringAlgorithm())
+        @test_opt symmetric_coloring(Symmetric(A), GreedyColoringAlgorithm())
+
+        @inferred column_coloring(A, GreedyColoringAlgorithm())
+        @inferred row_coloring(A, GreedyColoringAlgorithm())
+        @inferred symmetric_coloring(Symmetric(A), GreedyColoringAlgorithm())
     end
 
     @testset "$structure - $partition - $decompression" for (
@@ -36,7 +34,12 @@ rng = StableRNG(63)
         (:nonsymmetric, :bidirectional, :direct),
         (:nonsymmetric, :bidirectional, :substitution),
     ]
-        @test_opt target_modules = (SparseMatrixColorings,) coloring(
+        @test_opt coloring(
+            A,
+            ColoringProblem(; structure, partition),
+            GreedyColoringAlgorithm(; decompression),
+        )
+        @inferred coloring(
             A,
             ColoringProblem(; structure, partition),
             GreedyColoringAlgorithm(; decompression),
@@ -55,7 +58,12 @@ end;
         (structure, partition, decompression) in
         [(:nonsymmetric, :column, :direct), (:nonsymmetric, :row, :direct)]
 
-        @test_opt target_modules = (SparseMatrixColorings,) coloring(
+        @test_opt coloring(
+            A,
+            ColoringProblem(; structure, partition),
+            GreedyColoringAlgorithm(; decompression),
+        )
+        @inferred coloring(
             A,
             ColoringProblem(; structure, partition),
             GreedyColoringAlgorithm(; decompression),
@@ -88,15 +96,21 @@ end;
                 Br, Bc = compress(A, result)
                 @testset "Full decompression" begin
                     @test_opt compress(A, result)
-                    @test_opt decompress(Br, Bc, result) ≈ A0
+                    @test_opt decompress(Br, Bc, result)
                     @test_opt decompress!(respectful_similar(A), Br, Bc, result)
+
+                    @inferred compress(A, result)
+                    @inferred decompress(Br, Bc, result)
                 end
             else
                 B = compress(A, result)
                 @testset "Full decompression" begin
                     @test_opt compress(A, result)
-                    @test_opt decompress(B, result) ≈ A0
+                    @test_opt decompress(B, result)
                     @test_opt decompress!(respectful_similar(A), B, result)
+
+                    @inferred compress(A, result)
+                    @inferred decompress(B, result)
                 end
                 @testset "Single-color decompression" begin
                     if decompression == :direct
@@ -145,5 +159,6 @@ end;
         )
         B = compress(A, result)
         @test_opt decompress(B, result)
+        @inferred decompress(B, result)
     end
 end;