Implemented a first repartition strategy

Author: fverdugo

extensions/PartitionedSolvers/src/amg.jl

@@ -260,10 +260,33 @@ function omega_for_1d_laplace(invD,A)
     2/3
 end
 
+function enhance_coarse_partition(A,Ac,Bc,R,P,cache,repartition_threshold)
+    Ac,Bc,R,P,cache,repartition_threshold
+end
+
+function enhance_coarse_partition(A::PSparseMatrix,Ac,Bc,R,P,cache,repartition_threshold)
+    # TODO now we redistribute to a single proc when the threshold is reached
+    # We need to redistributed to fewer but several processors
+    # TODO a way of avoiding the extra rap ?
+    n_coarse = size(Ac,1)
+    if n_coarse <= repartition_threshold
+        repartition_threshold = 0 # do not repartition again
+        parts = linear_indices(partition(Ac))
+        coarse_partition = trivial_partition(parts,n_coarse)
+        P = repartition(P,partition(axes(P,1)),coarse_partition) |> fetch
+        Bc = map(b->repartition(b,partition(axes(P,2)))|>fetch,Bc)
+        R = transpose(P)
+        Ac,cache = rap(R,A,P;reuse=true)
+    end
+    Ac,Bc,R,P,cache,repartition_threshold
+end
+
 function smoothed_aggregation(;
     epsilon = 0,
     approximate_omega = omega_for_1d_laplace,
-    tentative_prolongator = tentative_prolongator_for_laplace)
+    tentative_prolongator = tentative_prolongator_for_laplace,
+    repartition_threshold = 2000,
+    )
     function coarsen(operator)
         A = matrix(operator)
         B = nullspace(operator)
@@ -275,7 +298,7 @@ function smoothed_aggregation(;
         P = smoothed_prolongator(A,P0,diagA;approximate_omega)
         R = transpose(P)
         Ac,cache = rap(R,A,P;reuse=true)
-        # TODO enhance partitioning for Ac,R,P
+        Ac,Bc,R,P,cache,repartition_threshold = enhance_coarse_partition(A,Ac,Bc,R,P,cache,repartition_threshold)
         coarse_operator = attach_nullspace(Ac,Bc)
         coarse_operator,R,P,cache
     end
@@ -289,7 +312,7 @@ function smoothed_aggregation(;
 end
 
 function amg_level_params(;
-    pre_smoother = additive_schwarz(gauss_seidel(;iters=1)),
+    pre_smoother = richardson(additive_schwarz(gauss_seidel(;iters=1));iters=1),
     coarsening = smoothed_aggregation(;),
     cycle = v_cycle,
     pos_smoother = pre_smoother,
@@ -345,7 +368,7 @@ function amg_setup(x,operator,b,amg_params)
             done = true
         end
         r = similar(b)
-        rc = similar(r,axes(Ac,1))
+        rc = similar(r,axes(Ac,2)) # we need ghost ids for the mul!(rc,R,r)
         e = similar(x)
         ec = similar(e,axes(Ac,2))
         level_setup = (;R,P,r,rc,e,ec,operator,coarse_operator,pre_setup,pos_setup,coarse_operator_setup)

extensions/PartitionedSolvers/test/amg_tests.jl

@@ -69,7 +69,7 @@ cg!(y,A,b;Pl,verbose=true)
 
 # Now in parallel
 
-parts_per_dir = (4,4)
+parts_per_dir = (1,2)
 np = prod(parts_per_dir)
 parts = DebugArray(LinearIndices((np,)))
 
@@ -100,6 +100,7 @@ finalize!(solver)(S)
 
 level_params = amg_level_params(;
     pre_smoother = jacobi(;iters=1,omega=2/3),
+    coarsening = smoothed_aggregation(;repartition_threshold=10000000)
    )
 
 fine_params = amg_fine_params(;
@@ -112,15 +113,8 @@ Pl = preconditioner(solver,y,A,b)
 y .= 0
 cg!(y,A,b;Pl,verbose=true)
 
-level_params = amg_level_params(;
-    pre_smoother = richardson(additive_schwarz(gauss_seidel(;iters=1)),iters=1),
-   )
-
-fine_params = amg_fine_params(;
-    level_params,
-    n_fine_levels=5)
 
-solver = amg(;fine_params)
+solver = amg()
 
 Pl = preconditioner(solver,y,A,b)
 y .= 0

src/p_sparse_matrix.jl

@@ -1817,6 +1817,7 @@ function LinearAlgebra.mul!(c::PVector,at::Transpose{T,<:PSparseMatrix} where T,
     a = at.parent
     @assert a.assembled
     map(ghost_values(c),own_ghost_values(a),own_values(b)) do ch,aoh,bo
+        fill!(ch,zero(eltype(ch)))
         atoh = transpose(aoh)
         mul!(ch,atoh,bo,α,1)
     end