Added CachedPETScNS

gridap · Jan 3, 2024 · bb05dc4 · bb05dc4
1 parent 3ef63a7
commit bb05dc4
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 89 deletions.
diff --git a/src/LinearSolvers/GMGLinearSolvers.jl b/src/LinearSolvers/GMGLinearSolvers.jl
@@ -58,15 +58,15 @@ function Gridap.Algebra.numerical_setup(ss::GMGSymbolicSetup,mat::AbstractMatrix
   coarsest_solver = ss.solver.coarsest_solver
 
   smatrices[1] = mat
-  finest_level_cache = setup_finest_level_cache(mh,smatrices)
-  work_vectors = allocate_work_vectors(mh,smatrices)
-  pre_smoothers_caches = setup_smoothers_caches(mh,pre_smoothers,smatrices)
+  finest_level_cache = gmg_finest_level_cache(mh,smatrices)
+  work_vectors = gmg_work_vectors(mh,smatrices)
+  pre_smoothers_caches = gmg_smoothers_caches(mh,pre_smoothers,smatrices)
   if !(pre_smoothers === post_smoothers)
-    post_smoothers_caches = setup_smoothers_caches(mh,post_smoothers,smatrices)
+    post_smoothers_caches = gmg_smoothers_caches(mh,post_smoothers,smatrices)
   else
     post_smoothers_caches = pre_smoothers_caches
   end
-  coarsest_solver_cache = coarse_solver_caches(mh,coarsest_solver,smatrices)
+  coarsest_solver_cache = gmg_coarse_solver_caches(mh,coarsest_solver,smatrices,work_vectors)
 
   return GMGNumericalSetup(ss.solver,finest_level_cache,pre_smoothers_caches,post_smoothers_caches,coarsest_solver_cache,work_vectors)
 end
@@ -76,7 +76,7 @@ function Gridap.Algebra.numerical_setup!(ss::GMGNumericalSetup,mat::AbstractMatr
   ns.solver.smatrices[1] = mat
 end
 
-function setup_finest_level_cache(mh::ModelHierarchy,smatrices::Vector{<:AbstractMatrix})
+function gmg_finest_level_cache(mh::ModelHierarchy,smatrices::Vector{<:AbstractMatrix})
   cache = nothing
   parts = get_level_parts(mh,1)
   if i_am_in(parts)
@@ -87,7 +87,7 @@ function setup_finest_level_cache(mh::ModelHierarchy,smatrices::Vector{<:Abstrac
   return cache
 end
 
-function setup_smoothers_caches(mh::ModelHierarchy,smoothers::AbstractVector{<:LinearSolver},smatrices::Vector{<:AbstractMatrix})
+function gmg_smoothers_caches(mh::ModelHierarchy,smoothers::AbstractVector{<:LinearSolver},smatrices::Vector{<:AbstractMatrix})
   Gridap.Helpers.@check length(smoothers) == num_levels(mh)-1
   nlevs = num_levels(mh)
   # Last (i.e., coarsest) level does not need pre-/post-smoothing
@@ -102,64 +102,34 @@ function setup_smoothers_caches(mh::ModelHierarchy,smoothers::AbstractVector{<:L
   return caches
 end
 
-function coarse_solver_caches(mh,s,mats)
+function gmg_coarse_solver_caches(mh,s,mats,work_vectors)
   cache = nothing
   nlevs = num_levels(mh)
   parts = get_level_parts(mh,nlevs)
   if i_am_in(parts)
     mat = mats[nlevs]
+    _, _, xH, rH = work_vectors[nlevs-1]
     cache = numerical_setup(symbolic_setup(s, mat), mat)
-  end
-  return cache
-end
-
-function setup_coarsest_solver_cache(mh::ModelHierarchy,coarsest_solver::LinearSolver,smatrices::Vector{<:AbstractMatrix})
-  cache = nothing
-  nlevs = num_levels(mh)
-  parts = get_level_parts(mh,nlevs)
-  if i_am_in(parts)
-    mat = smatrices[nlevs]
-    if (num_parts(parts) == 1) # Serial
-      cache = map(own_values(mat)) do Ah
-        ss  = symbolic_setup(coarsest_solver, Ah)
-        numerical_setup(ss, Ah)
-      end
-      cache = PartitionedArrays.getany(cache)
-    else # Parallel
-      ss = symbolic_setup(coarsest_solver, mat)
-      cache = numerical_setup(ss, mat)
+    if isa(s,PETScLinearSolver)
+      cache = CachedPETScNS(cache, xH, rH)
     end
   end
   return cache
 end
 
-function setup_coarsest_solver_cache(mh::ModelHierarchy,coarsest_solver::PETScLinearSolver,smatrices::Vector{<:AbstractMatrix})
-  cache = nothing
+function gmg_work_vectors(mh::ModelHierarchy,smatrices::Vector{<:AbstractMatrix})
   nlevs = num_levels(mh)
-  parts = get_level_parts(mh,nlevs)
-  if i_am_in(parts)
-    mat   = smatrices[nlevs]
-    if (num_parts(parts) == 1) # Serial
-      cache = map(own_values(mat)) do Ah
-        rh  = convert(PETScVector,fill(0.0,size(A,2)))
-        xh  = convert(PETScVector,fill(0.0,size(A,2)))
-        ss  = symbolic_setup(coarsest_solver, Ah)
-        ns  = numerical_setup(ss, Ah)
-        return ns, xh, rh
-      end
-      cache = PartitionedArrays.getany(cache)
-    else # Parallel
-      rh = convert(PETScVector,pfill(0.0,partition(axes(mat,2))))
-      xh = convert(PETScVector,pfill(0.0,partition(axes(mat,2))))
-      ss = symbolic_setup(coarsest_solver, mat)
-      ns = numerical_setup(ss, mat)
-      cache = ns, xh, rh
+  work_vectors = Vector{Any}(undef,nlevs-1)
+  for i = 1:nlevs-1
+    parts = get_level_parts(mh,i)
+    if i_am_in(parts)
+      work_vectors[i] = gmg_work_vectors(mh,smatrices,i)
     end
   end
-  return cache
+  return work_vectors
 end
 
-function allocate_level_work_vectors(mh::ModelHierarchy,smatrices::Vector{<:AbstractMatrix},lev::Integer)
+function gmg_work_vectors(mh::ModelHierarchy,smatrices::Vector{<:AbstractMatrix},lev::Integer)
   dxh   = allocate_in_domain(smatrices[lev])
   Adxh  = allocate_in_range(smatrices[lev])
 
@@ -175,52 +145,12 @@ function allocate_level_work_vectors(mh::ModelHierarchy,smatrices::Vector{<:Abst
   return dxh, Adxh, dxH, rH
 end
 
-function allocate_work_vectors(mh::ModelHierarchy,smatrices::Vector{<:AbstractMatrix})
-  nlevs = num_levels(mh)
-  work_vectors = Vector{Any}(undef,nlevs-1)
-  for i = 1:nlevs-1
-    parts = get_level_parts(mh,i)
-    if i_am_in(parts)
-      work_vectors[i] = allocate_level_work_vectors(mh,smatrices,i)
-    end
-  end
-  return work_vectors
-end
-
-function solve_coarsest_level!(parts::AbstractArray,::LinearSolver,xh::PVector,rh::PVector,caches)
-  if (num_parts(parts) == 1)
-    map(own_values(xh),own_values(rh)) do xh, rh
-      solve!(xh,caches,rh)
-    end
-  else
-    solve!(xh,caches,rh)
-  end
-end
-
-function solve_coarsest_level!(parts::AbstractArray,::PETScLinearSolver,xh::PVector,rh::PVector,caches)
-  solver_ns, xh_petsc, rh_petsc = caches
-  if (num_parts(parts) == 1)
-    map(own_values(xh),own_values(rh)) do xh, rh
-      copy!(rh_petsc,rh)
-      solve!(xh_petsc,solver_ns,rh_petsc)
-      copy!(xh,xh_petsc)
-    end
-  else
-    copy!(rh_petsc,rh)
-    solve!(xh_petsc,solver_ns,rh_petsc)
-    copy!(xh,xh_petsc)
-  end
-end
-
 function apply_GMG_level!(lev::Integer,xh::Union{PVector,Nothing},rh::Union{PVector,Nothing},ns::GMGNumericalSetup)
   mh = ns.solver.mh
   parts = get_level_parts(mh,lev)
   if i_am_in(parts)
     if (lev == num_levels(mh)) 
       ## Coarsest level
-      #coarsest_solver = ns.solver.coarsest_solver
-      #coarsest_solver_cache = ns.coarsest_solver_cache
-      #solve_coarsest_level!(parts,coarsest_solver,xh,rh,coarsest_solver_cache)
       solve!(xh, ns.coarsest_solver_cache, rh)
     else 
       ## General case

diff --git a/src/LinearSolvers/LinearSolvers.jl b/src/LinearSolvers/LinearSolvers.jl
@@ -43,6 +43,7 @@ include("Krylov/FGMRESSolvers.jl")
 include("Krylov/MINRESSolvers.jl")
 
 include("PETSc/PETScUtils.jl")
+include("PETSc/PETScCaches.jl")
 include("PETSc/ElasticitySolvers.jl")
 include("PETSc/HipmairXuSolvers.jl")
 

diff --git a/src/LinearSolvers/PETSc/PETScCaches.jl b/src/LinearSolvers/PETSc/PETScCaches.jl
@@ -0,0 +1,33 @@
+
+"""
+ Notes on this structure: 
+
+ When converting julia vectors/PVectors to PETSc vectors, we purposely create aliasing 
+ of the vector values. This means we can avoid copying data from one to another before solving, 
+ but we need to be careful about it. 
+
+ This structure takes care of this, and makes sure you do not attempt to solve the system 
+ with julia vectors that are not the ones you used to create the solver cache.
+"""
+struct CachedPETScNS{TM,A}
+  ns     :: GridapPETSc.PETScLinearSolverNS{TM}
+  X      :: PETScVector
+  B      :: PETScVector
+  owners :: A
+  function CachedPETScNS(ns::GridapPETSc.PETScLinearSolverNS{TM},x::AbstractVector,b::AbstractVector) where TM
+    X = convert(PETScVector,x)
+    B = convert(PETScVector,b)
+    owners = (x,b)
+
+    A = typeof(owners)
+    new{TM,A}(ns,X,B,owners)
+  end
+end
+
+function Algebra.solve!(x::AbstractVector,ns::CachedPETScNS,b::AbstractVector)
+  @assert x === ns.owners[1]
+  @assert b === ns.owners[2]
+  solve!(ns.X,ns.ns,ns.B)
+  consistent!(x)
+  return x
+end