From fddd01474e8790f9554d3f59d6fe683ec0fa1aeb Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Fri, 22 Jul 2022 20:51:27 +0200
Subject: [PATCH] Add solve(!) methods for PartitionedArrays input.

This patch adds the following methods:

    solve( ::HYPRESolver,            ::PSparseMatrix, ::PVector)
    solve!(::HYPRESolver, ::PVector, ::PSparseMatrix, ::PVector)

which automatically converts to/from HYPRE(Matrix|Vector).
---
 src/solvers.jl   | 20 +++++++++++++++++++
 test/runtests.jl | 50 +++++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/src/solvers.jl b/src/solvers.jl
index 284293f..2eec3ef 100644
--- a/src/solvers.jl
+++ b/src/solvers.jl
@@ -31,6 +31,26 @@ See also [`solve`](@ref).
 solve!(pcg::HYPRESolver, x::HYPREVector, A::HYPREMatrix, ::HYPREVector)
 
 
+######################################
+# PartitionedArrays solver interface #
+######################################
+
+# TODO: Would it be useful with a method that copied the solution to b instead?
+
+function solve(solver::HYPRESolver, A::PSparseMatrix, b::PVector)
+    hypre_x = solve(solver, HYPREMatrix(A), HYPREVector(b))
+    # TODO: This could be a HYPREVector -> PVector conversion
+    x = copy!(copy(b), hypre_x)
+    return x
+end
+function solve!(solver::HYPRESolver, x::PVector, A::PSparseMatrix, b::PVector)
+    hypre_x = HYPREVector(x)
+    solve!(solver, hypre_x, HYPREMatrix(A), HYPREVector(b))
+    copy!(x, hypre_x)
+    return x
+end
+
+
 #############
 # BoomerAMG #
 #############
diff --git a/test/runtests.jl b/test/runtests.jl
index 263c1ef..9cf5897 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -65,7 +65,7 @@ end
 end
 
 function tomain(x)
-    g = gather(x)
+    g = gather(copy(x))
     be = get_backend(g.values)
     if be isa SequentialBackend
         return g.values.parts[1]
@@ -107,7 +107,7 @@ end
     CSC = PSparseMatrix(diag_data(backend, parts)...; ids=:global)
     CSR = PSparseMatrix(sparsecsr, diag_data(backend, parts)...; ids=:global)
 
-    @test tomain(copy(CSC)) == tomain(copy(CSR)) ==
+    @test tomain(CSC) == tomain(CSR) ==
         Diagonal([1, 2, 3, 8, 10, 12, 7, 8, 9, 10])
 
     map_parts(CSC.values, CSC.rows.partition, CSC.cols.partition,
@@ -141,7 +141,7 @@ end
     CSC = PSparseMatrix(diag_data(backend, parts)...; ids=:global)
     CSR = PSparseMatrix(sparsecsr, diag_data(backend, parts)...; ids=:global)
 
-    @test tomain(copy(CSC)) == tomain(copy(CSR)) ==
+    @test tomain(CSC) == tomain(CSR) ==
         Diagonal([1, 2, 3, 8, 10, 12, 7, 8, 9, 10])
 
     map_parts(CSC.values, CSC.rows.partition, CSC.cols.partition,
@@ -223,13 +223,13 @@ end
     add_gids!(rows, I)
     pb = PVector(I, V, rows; ids=:global)
     assemble!(pb)
-    @test tomain(copy(pb)) == [i in 4:6 ? 2x : x for (i, x) in zip(eachindex(b), b)]
+    @test tomain(pb) == [i in 4:6 ? 2x : x for (i, x) in zip(eachindex(b), b)]
     H = HYPREVector(pb)
     @test H.IJVector != HYPRE_IJVector(C_NULL)
     @test H.ParVector != HYPRE_ParVector(C_NULL)
     pbc = fill!(copy(pb), 0)
     copy!(pbc, H)
-    @test tomain(copy(pbc)) == tomain(copy(pb))
+    @test tomain(pbc) == tomain(pb)
     # MPI backend
     backend = MPIBackend()
     parts = get_part_ids(backend, 1)
@@ -240,13 +240,13 @@ end
     add_gids!(rows, I)
     pb = PVector(I, V, rows; ids=:global)
     assemble!(pb)
-    @test tomain(copy(pb)) == b
+    @test tomain(pb) == b
     H = HYPREVector(pb)
     @test H.IJVector != HYPRE_IJVector(C_NULL)
     @test H.ParVector != HYPRE_ParVector(C_NULL)
     pbc = fill!(copy(pb), 0)
     copy!(pbc, H)
-    @test tomain(copy(pbc)) == tomain(copy(pb))
+    @test tomain(pbc) == tomain(pb)
 end
 
 @testset "BoomerAMG" begin
@@ -304,3 +304,39 @@ end
     copy!(x, x_h)
     @test x ≈ A \ b atol=tol
 end
+
+function topartitioned(x::Vector, A::SparseMatrixCSC, b::Vector)
+    parts = get_part_ids(SequentialBackend(), 1)
+    rows = PRange(parts, size(A, 1))
+    cols = PRange(parts, size(A, 2))
+    II, JJ, VV, bb, xx = map_parts(parts) do _
+        return findnz(A)..., b, x
+    end
+    add_gids!(rows, II)
+    assemble!(II, JJ, VV, rows)
+    add_gids!(cols, JJ)
+    A_p = PSparseMatrix(II, JJ, VV, rows, cols; ids = :global)
+    b_p = PVector(bb, rows)
+    x_p = PVector(xx, cols)
+    return x_p, A_p, b_p
+end
+
+@testset "solve with PartitionedArrays" begin
+    # Setup
+    A = sprand(100, 100, 0.05); A = A'A + 5I
+    b = rand(100)
+    x = zeros(100)
+    x_p, A_p, b_p = topartitioned(x, A, b)
+    @test A == tomain(A_p)
+    @test b == tomain(b_p)
+    @test x == tomain(x_p)
+    # Solve
+    tol = 1e-9
+    pcg = HYPRE.PCG(; Tol = tol)
+    ## solve!
+    HYPRE.solve!(pcg, x_p, A_p, b_p)
+    @test tomain(x_p) ≈ A \ b atol=tol
+    ## solve
+    x_p = HYPRE.solve(pcg, A_p, b_p)
+    @test tomain(x_p) ≈ A \ b atol=tol
+end