Merge pull request #129 from fverdugo/non_isbitstype

Support for gather/scatter with Non isbitstype objects (and more)

Author: fverdugo

CHANGELOG.md

@@ -5,10 +5,16 @@ 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).
 
+## [0.4.1] - Unreleased
+
+### Added
+
+- Gather/scatter for non isbitstype objects.
+- Function `find_local_indices`.
 
 ## [0.4.0] - 2024-01-21
 
-## Changed
+### Changed
 
 - Major refactoring in `PSparseMatrix` (and in `PVector` in a lesser extent).
 The old code is still available (but deprecated), and can be recovered applying this renaming to your code-base:
@@ -21,7 +27,7 @@ The old code is still available (but deprecated), and can be recovered applying
 The previous "monolithic" storage is not implemented anymore for the new version of `PSparseMatrix`, but can be implemented in the new setup if needed.
 - `emit` renamed to `multicast`. The former name is still available but deprecated.
 
-## Added
+### Added
 
 - Efficient re-construction of `PSparseMatrix` and `PVector` objects.
 - Functions `assemble` and `consistent` (allocating versions of `assemble!` and `consistent!` with a slightly different

Project.toml

@@ -1,7 +1,7 @@
 name = "PartitionedArrays"
 uuid = "5a9dfac6-5c52-46f7-8278-5e2210713be9"
 authors = ["Francesc Verdugo <[email protected]> and contributors"]
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 CircularArrays = "7a955b69-7140-5f4e-a0ed-f168c5e2e749"

docs/src/reference/primitives.md

@@ -16,7 +16,7 @@ scatter!
 allocate_scatter
 ```
 
-## Emit
+## Multicast
 
 ```@docs
 multicast

src/PartitionedArrays.jl

@@ -130,6 +130,7 @@ export assemble
 export consistent
 export repartition
 export repartition!
+export find_local_indices
 include("p_vector.jl")
 
 export OldPSparseMatrix

src/mpi_array.jl

@@ -281,13 +281,21 @@ function gather_impl!(
     comm = snd.comm
     if isa(destination,Integer)
         root = destination-1
-        if MPI.Comm_rank(comm) == root
-            @assert length(rcv.item) == MPI.Comm_size(comm)
-            rcv.item[destination] = snd.item
-            rcv_buffer = MPI.UBuffer(rcv.item,1)
-            MPI.Gather!(MPI.IN_PLACE,rcv_buffer,root,comm)
+        if isbitstype(T)
+            if MPI.Comm_rank(comm) == root
+                @assert length(rcv.item) == MPI.Comm_size(comm)
+                rcv.item[destination] = snd.item
+                rcv_buffer = MPI.UBuffer(rcv.item,1)
+                MPI.Gather!(MPI.IN_PLACE,rcv_buffer,root,comm)
+            else
+                MPI.Gather!(snd.item_ref,nothing,root,comm)
+            end
         else
-            MPI.Gather!(snd.item_ref,nothing,root,comm)
+            if MPI.Comm_rank(comm) == root
+                rcv.item[:] = MPI.gather(snd.item,comm;root)
+            else
+                MPI.gather(snd.item,comm;root)
+            end
         end
     else
         @assert destination === :all
@@ -330,17 +338,25 @@ end
 function scatter_impl!(
     rcv::MPIArray,snd::MPIArray,
     source,::Type{T}) where T
-    @assert source !== :all "All to all not implemented"
-    @assert rcv.comm === snd.comm
-    @assert eltype(snd.item) == typeof(rcv.item)
     comm = snd.comm
     root = source - 1
-    if MPI.Comm_rank(comm) == root
-        snd_buffer = MPI.UBuffer(snd.item,1)
-        rcv.item = snd.item[source]
-        MPI.Scatter!(snd_buffer,MPI.IN_PLACE,root,comm)
+    @assert source !== :all "All to all not implemented"
+    @assert rcv.comm === snd.comm
+    if isbitstype(T)
+        @assert eltype(snd.item) == typeof(rcv.item)
+        if MPI.Comm_rank(comm) == root
+            snd_buffer = MPI.UBuffer(snd.item,1)
+            rcv.item = snd.item[source]
+            MPI.Scatter!(snd_buffer,MPI.IN_PLACE,root,comm)
+        else
+            MPI.Scatter!(nothing,rcv.item_ref,root,comm)
+        end
     else
-        MPI.Scatter!(nothing,rcv.item_ref,root,comm)
+        if MPI.Comm_rank(comm) == root
+            rcv.item_ref[] = MPI.scatter(snd.item,comm;root)
+        else
+            rcv.item_ref[] = MPI.scatter(nothing,comm;root)
+        end
     end
     rcv
 end

src/p_vector.jl

@@ -1076,4 +1076,44 @@ function repartition!(w::PVector,v::PVector,cache;reversed=false)
     end
 end
 
+function find_local_indices(node_to_mask::PVector)
+    n_own_dofs = map(count,own_values(node_to_mask))
+    n_dofs = sum(n_own_dofs)
+    dof_partition = variable_partition(n_own_dofs,n_dofs)
+    node_partition = partition(axes(node_to_mask,1))
+    node_to_global_dof = pzeros(Int,node_partition)
+    function fill_own_dofs!(own_node_to_global_dof,own_node_to_boundary,dofs)
+        own_to_global_dof = own_to_global(dofs)
+        own_node_to_global_dof[own_node_to_boundary] = own_to_global_dof
+    end
+    map(fill_own_dofs!,own_values(node_to_global_dof),own_values(node_to_mask),dof_partition)
+    consistent!(node_to_global_dof) |> wait
+    function add_ghost_dofs(ghost_node_to_global_dof,nodes,dofs)
+        ghost_node_to_owner = ghost_to_owner(nodes)
+        free_ghost_nodes = findall(global_dof->global_dof!=0,ghost_node_to_global_dof)
+        owners = view(ghost_node_to_owner,free_ghost_nodes)
+        ghost_dofs = view(ghost_node_to_global_dof,free_ghost_nodes)
+        union_ghost(dofs,ghost_dofs,owners)
+    end
+    dof_partition = map(add_ghost_dofs,ghost_values(node_to_global_dof),node_partition,dof_partition)
+    neighbors = assembly_graph(node_partition)
+    assembly_neighbors(dof_partition;neighbors)
+    node_to_local_dof = pzeros(Int32,node_partition)
+    dof_to_local_node = pzeros(Int32,dof_partition)
+    function finalize!(local_node_to_global_dof,local_node_to_local_dof,local_dof_to_local_node,dofs)
+        global_to_local_dof = global_to_local(dofs)
+        n_local_nodes = length(local_node_to_global_dof)
+        for local_node in 1:n_local_nodes
+            global_dof = local_node_to_global_dof[local_node]
+            if global_dof == 0
+                continue
+            end
+            local_dof = global_to_local_dof[global_dof]
+            local_node_to_local_dof[local_node] = local_dof
+            local_dof_to_local_node[local_dof] = local_node
+        end
+    end
+    map(finalize!,partition(node_to_global_dof),partition(node_to_local_dof),partition(dof_to_local_node),dof_partition)
+    dof_to_local_node, node_to_local_dof
+end
 

test/primitives_tests.jl

@@ -1,6 +1,11 @@
 
 using Test
 
+struct NonIsBitsType{T}
+    data::Vector{T}
+end
+Base.:(==)(a::NonIsBitsType,b::NonIsBitsType) = a.data == b.data
+
 function primitives_tests(distribute)
 
    rank = distribute(LinearIndices((2,2)))
@@ -61,6 +66,27 @@ function primitives_tests(distribute)
        @test rcv == [[1],[1,2],[1,2,3],[1,2,3,4]]
    end
 
+   snd2 = map(rank) do rank
+       NonIsBitsType([2])
+   end
+   rcv2 = gather(snd2)
+   snd3 = scatter(rcv2)
+   map(snd2,snd3) do snd2,snd3
+       @test snd2 == snd3
+   end
+
+   np = length(rank)
+   rcv3 = map_main(rank) do rank
+       fill(NonIsBitsType([2]),np)
+   end
+   snd3 = allocate_scatter(rcv3)
+   scatter!(snd3,rcv3)
+   snd3 = scatter(rcv3)
+   rcv4 = gather(snd3)
+   map(rcv4,rcv2) do rcv4,rcv2
+       @test rcv4 == rcv2
+   end
+
    rcv = multicast(rank,source=2)
    map(rcv) do rcv
        @test rcv == 2