Merge pull request #182 from gridap/expand-ghost

CartesianDiscreteModels with multiple ghost layers

Author: JordiManyer

NEWS.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+
+- Added support for multiple ghost layers on cartesian models. Since PR[#182](https://github.com/gridap/GridapDistributed.jl/pull/182).
+
 ## [0.4.9] - 2025-08-08
 
 ### Added

Project.toml

@@ -5,6 +5,7 @@ version = "0.4.9"
 
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+CircularArrays = "7a955b69-7140-5f4e-a0ed-f168c5e2e749"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Gridap = "56d4f2e9-7ea1-5844-9cf6-b9c51ca7ce8e"
@@ -17,6 +18,7 @@ WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [compat]
 BlockArrays = "1"
+CircularArrays = "1.4.0"
 FillArrays = "1"
 ForwardDiff = "0.10, 1"
 Gridap = "0.19.4"

src/Geometry.jl

@@ -191,19 +191,20 @@ function _setup_face_gids!(dmodel::GenericDistributedDiscreteModel{Dc},dim) wher
 end
 
 # CartesianDiscreteModel
-struct DistributedCartesianDescriptor{A,B,C}
+struct DistributedCartesianDescriptor{A,B,C,D}
   ranks::A
   mesh_partition::B
   descriptor::C
+  ghost::D
   function DistributedCartesianDescriptor(
     ranks::AbstractArray{<:Integer},
     mesh_partition::NTuple{Dc,<:Integer},
-    descriptor::CartesianDescriptor{Dc}
+    descriptor::CartesianDescriptor{Dc},
+    ghost = map(i -> true, mesh_partition)
   ) where Dc
-    A = typeof(ranks)
-    B = typeof(mesh_partition)
-    C = typeof(descriptor)
-    new{A,B,C}(ranks,mesh_partition,descriptor)
+    A, B = typeof(ranks), typeof(mesh_partition)
+    C, D = typeof(descriptor), typeof(ghost)
+    new{A,B,C,D}(ranks,mesh_partition,descriptor,ghost)
   end
 end
 
@@ -223,44 +224,47 @@ function emit_cartesian_descriptor(
   new_mesh_partition
 ) where Dc
   f(a) = Tuple(PartitionedArrays.getany(emit(a)))
-  a, b, c, d = map(new_ranks) do rank
+  a, b, c, d, e = map(new_ranks) do rank
     if rank == 1
       desc = pdesc.descriptor
       @assert desc.map === identity
-      Float64[desc.origin.data...], Float64[desc.sizes...], Int[desc.partition...], Bool[desc.isperiodic...]
+      Float64[desc.origin.data...], Float64[desc.sizes...], Int[desc.partition...], Bool[desc.isperiodic...], Int16[pdesc.ghost...]
     else
-      Float64[], Float64[], Int[], Bool[]
+      Float64[], Float64[], Int[], Bool[], Int16[]
     end
   end |> tuple_of_arrays
-  origin, sizes, partition, isperiodic = VectorValue(f(a)...), f(b), f(c), f(d)
+  origin, sizes, partition, isperiodic, ghost = VectorValue(f(a)...), f(b), f(c), f(d), f(e)
   new_desc = CartesianDescriptor(origin,sizes,partition;isperiodic)
-  return DistributedCartesianDescriptor(new_ranks,new_mesh_partition,new_desc)
+  return DistributedCartesianDescriptor(new_ranks,new_mesh_partition,new_desc,ghost)
 end
 
 const DistributedCartesianDiscreteModel{Dc,Dp,A,B,C} =
   GenericDistributedDiscreteModel{Dc,Dp,<:AbstractArray{<:CartesianDiscreteModel},B,<:DistributedCartesianDescriptor}
 
 function Geometry.CartesianDiscreteModel(
-  ranks::AbstractArray{<:Integer},parts::NTuple{N,<:Integer},args...;kwargs...
+  ranks::AbstractArray{<:Integer}, # Distributed array with the rank IDs
+  parts::NTuple{N,<:Integer},      # Number of ranks (parts) in each direction
+  args...; ghost = map(i -> true, parts), kwargs...
 ) where N
   desc = CartesianDescriptor(args...;kwargs...)
   @check N == length(desc.partition)
   @check prod(parts) == length(ranks)
-  pdesc = DistributedCartesianDescriptor(ranks,parts,desc)
+  pdesc = DistributedCartesianDescriptor(ranks,parts,desc,ghost)
   return CartesianDiscreteModel(pdesc)
 end
 
 function Geometry.CartesianDiscreteModel(pdesc::DistributedCartesianDescriptor)
   desc = pdesc.descriptor
   isperiodic = desc.isperiodic
   if any(isperiodic)
+    @notimplementedif pdesc.ghost != map(i->true,pdesc.mesh_partition)
     models, cell_indices = _cartesian_model_with_periodic_bcs(pdesc)
   else
     nc = desc.partition
     ranks = pdesc.ranks
     parts = pdesc.mesh_partition
-    ghost = map(i->true,parts)
-    cell_indices = uniform_partition(ranks,parts,nc,ghost,isperiodic)
+    ghost = pdesc.ghost
+    cell_indices = _uniform_partition(ranks,parts,nc,ghost,isperiodic)
     gcids  = CartesianIndices(nc)
     models = map(cell_indices) do cell_indices
       cmin = gcids[first(cell_indices)]

src/GridapDistributed.jl

@@ -25,6 +25,7 @@ using FillArrays
 using BlockArrays
 using LinearAlgebra
 using ForwardDiff
+using CircularArrays
 
 import Gridap.TensorValues: inner, outer, double_contraction, symmetric_part
 import LinearAlgebra: det, tr, cross, dot, ⋅, diag
@@ -43,6 +44,8 @@ export with_ghost, no_ghost
 
 export redistribute
 
+include("PArraysExtras.jl")
+
 include("BlockPartitionedArrays.jl")
 
 include("Algebra.jl")

src/PArraysExtras.jl

@@ -0,0 +1,84 @@
+
+# This is copied from PArrays v0.5, which is not yet supported. It is necessary to get 
+# more than one layers of ghosts.
+
+function _uniform_partition(rank,np,n,args...)
+  @assert prod(np) == length(rank)
+  indices = map(rank) do rank
+    _block_with_constant_size(rank,np,n,args...)
+  end
+  if length(args) == 0
+    map(indices) do indices
+      cache = assembly_cache(indices)
+      copy!(cache,empty_assembly_cache())
+    end
+  else
+    assembly_neighbors(indices;symmetric=true)
+  end
+  indices
+end
+
+function _block_with_constant_size(rank,np,n,ghost,periodic=map(i->false,ghost))
+  N = length(n)
+  p = CartesianIndices(np)[rank]
+  own_ranges = map(_local_range,Tuple(p),np,n)
+  local_ranges = map(_local_range,Tuple(p),np,n,ghost,periodic)
+  owners = map(Tuple(p), np, n, local_ranges) do p, np, n, lr
+    myowners = zeros(Int32,length(lr))
+    i = 1
+    for p in Iterators.cycle(1:np)
+      plr = _local_range(p, np, n)
+      while mod(lr[i]-1, n)+1 in plr
+        myowners[i] = p
+        (i += 1) > length(myowners) && return myowners
+      end
+    end
+  end
+  n_local = prod(map(length, local_ranges))
+  n_own = prod(map(length, own_ranges))
+  n_ghost = n_local - n_own
+
+  ghost_to_global = zeros(Int,n_ghost)
+  ghost_to_owner = zeros(Int32,n_ghost)
+  perm = zeros(Int32,n_local)
+  i_ghost = 0
+  i_own = 0
+
+  cis = CartesianIndices(map(length,local_ranges))
+  lis = CircularArray(LinearIndices(n))
+  local_cis = CartesianIndices(local_ranges)
+  owner_lis = LinearIndices(np)
+  for (i,ci) in enumerate(cis)
+    flags = map(Tuple(ci), own_ranges, local_ranges) do i, or, lr
+      i in (or .- first(lr) .+ 1)
+    end
+    if !all(flags)
+      i_ghost += 1
+      ghost_to_global[i_ghost] = lis[local_cis[i]]
+      o = map(getindex,owners,Tuple(ci))
+      o_ci = CartesianIndex(o)
+      ghost_to_owner[i_ghost] = owner_lis[o_ci]
+      perm[i] = i_ghost + n_own
+    else
+      i_own += 1
+      perm[i] = i_own
+    end
+  end
+  ghostids = GhostIndices(prod(n),ghost_to_global,ghost_to_owner)
+  ids = PartitionedArrays.LocalIndicesWithConstantBlockSize(p,np,n,ghostids)
+  PartitionedArrays.PermutedLocalIndices(ids,perm)
+end
+
+function _local_range(p,np,n,ghost=false,periodic=false)
+  l, rem = divrem(n, np)
+  offset = l * (p-1)
+  if rem >= (np-p+1)
+    l += 1
+    offset += p - (np-rem) - 1
+  end
+  start = 1+offset-ghost
+  stop = l+offset+ghost
+
+  periodic && return start:stop
+  return max(1, start):min(n,stop)
+end

test/GeometryTests.jl

@@ -109,6 +109,9 @@ function main(distribute,parts)
   Λs = Skeleton(Ωs)
   Γs = Boundary(Ωs)
 
+  # Multiple ghost layers
+  model = CartesianDiscreteModel(ranks,parts,domain,cells;ghost=map(i->2,parts))
+
 end
 
 function test_local_part_face_labelings_consistency(lmodel::CartesianDiscreteModel{D},gids,gmodel) where {D}