Merge pull request #213 from JordiManyer/find-rcv-ids

Setting default algorithm to discover the communication graph

Author: fverdugo

.gitignore

@@ -5,9 +5,11 @@
 *.jl.mem
 .DS_Store
 Manifest.toml
+LocalPreferences.toml
 docs/build/
 tmp/
 docs/src/examples.md
+docs/src/jacobi_tutorial.md
 
 HPCG/src/results/
 

CHANGELOG.md

@@ -5,6 +5,10 @@ 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.5.14] - 2025-09-04
+
+- Added compile-time preference to choose the algorithm used within `default_rcv_ids`.
+
 ## [0.5.13] - 2025-07-08
 
 ### Fixed

Project.toml

@@ -1,7 +1,7 @@
 name = "PartitionedArrays"
 uuid = "5a9dfac6-5c52-46f7-8278-5e2210713be9"
 authors = ["Francesc Verdugo <f.verdugo.rojano@vu.nl> and contributors"]
-version = "0.5.13"
+version = "0.5.14"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -12,6 +12,7 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -26,6 +27,7 @@ Distances = "0.10"
 FillArrays = "0.10, 0.11, 0.12, 0.13, 1"
 IterativeSolvers = "0.9"
 MPI = "0.16, 0.17, 0.18, 0.19, 0.20"
+Preferences = "1"
 SparseMatricesCSR = "0.6"
 StaticArrays = "1"
 julia = "1.1"

docs/src/reference/backends.md

@@ -5,6 +5,7 @@
 ```@autodocs
 Modules = [PartitionedArrays]
 Pages = ["mpi_array.jl"]
+Filter = t -> string(t) != "find_rcv_ids_ibarrier"
 ```
 
 ## Debug

docs/src/reference/primitives.md

@@ -8,6 +8,7 @@ gather
 gather!
 allocate_gather
 ```
+
 ## Scatter
 
 ```@docs
@@ -45,7 +46,8 @@ ExchangeGraph(snd)
 exchange
 exchange!
 allocate_exchange
+default_find_rcv_ids
+set_default_find_rcv_ids
+find_rcv_ids_gather_scatter
+find_rcv_ids_ibarrier
 ```
-
-
-

src/PartitionedArrays.jl

@@ -11,6 +11,10 @@ import IterativeSolvers
 import Distances
 using BlockArrays
 using Adapt
+using Preferences
+
+export set_default_find_rcv_ids
+include("preferences.jl")
 
 export length_to_ptrs!
 export rewind_ptrs!
@@ -54,6 +58,7 @@ export ExchangeGraph
 export exchange
 export exchange!
 export allocate_exchange
+export default_find_rcv_ids
 export find_rcv_ids_gather_scatter
 export setup_non_blocking_reduction
 export non_blocking_reduction

src/mpi_array.jl

@@ -660,14 +660,23 @@ end
 Issend(data, dest::Integer, tag::Integer, comm::MPI.Comm, req=MPI.Request()) =
     Issend(MPI.Buffer_send(data), dest, tag, comm, req)
 
-
 function default_find_rcv_ids(::MPIArray)
-    find_rcv_ids_gather_scatter
+    @static if default_find_rcv_ids_algorithm == "gather_scatter"
+        find_rcv_ids_gather_scatter
+    elseif default_find_rcv_ids_algorithm == "ibarrier"
+        find_rcv_ids_ibarrier
+    else
+        error("Unknown algorithm: $(default_find_rcv_ids_algorithm)")
+    end
 end
 
 """
- Implements Alg. 2 in https://dl.acm.org/doi/10.1145/1837853.1693476
- The algorithm's complexity is claimed to be O(log(p))
+    find_rcv_ids_ibarrier(snd_ids::MPIArray)
+
+Finds the `rcv` side of an `ExchangeGraph` out of the `snd` side information.
+
+This strategy implements Alg. 2 in https://dl.acm.org/doi/10.1145/1837853.1693476.
+The algorithm's complexity is claimed to be O(log(p)).
 """
 function find_rcv_ids_ibarrier(snd_ids::MPIArray{<:AbstractVector{T}}) where T
     comm = snd_ids.comm

src/preferences.jl

@@ -0,0 +1,28 @@
+
+"""
+    set_default_find_rcv_ids(algorithm::String)
+
+Sets the default algorithm to discover communication neighbors. The available algorithms are:
+
+- `gather_scatter`: Gathers neighbors in a single processor, builds the communications graph 
+  and then scatters the information back to all processors. See [`find_rcv_ids_gather_scatter`](@ref).
+
+- `ibarrier`: Implements Alg. 2 in https://dl.acm.org/doi/10.1145/1837853.1693476. See [`find_rcv_ids_ibarrier`](@ref).
+
+Feature only available in Julia 1.6 and later due to restrictions from `Preferences.jl`.
+"""
+function set_default_find_rcv_ids(algorithm::String)
+    if !(algorithm in ("gather_scatter", "ibarrier"))
+        throw(ArgumentError("Invalid algorihtm: \"$(algorithm)\""))
+    end
+
+    # Set it in our runtime values, as well as saving it to disk
+    @set_preferences!("default_find_rcv_ids" => algorithm)
+    @info("New default algorithm set; restart your Julia session for this change to take effect!")
+end
+
+@static if VERSION >= v"1.6"
+    const default_find_rcv_ids_algorithm = @load_preference("default_find_rcv_ids", "gather_scatter")
+else
+    const default_find_rcv_ids_algorithm = "gather_scatter"
+end

src/primitives.jl

@@ -762,6 +762,13 @@ function Base.show(io::IO,k::MIME"text/plain",data::ExchangeGraph)
     println(io,typeof(data)," with $(length(data.snd)) nodes")
 end
 
+"""
+    default_find_rcv_ids(::AbstractArray)
+
+Provides a default function to find the `rcv` side of an
+`ExchangeGraph` out of the `snd` side information.
+Its behaviour can be statically changed using [`set_default_find_rcv_ids`](@ref).
+"""
 function default_find_rcv_ids(::AbstractArray)
     find_rcv_ids_gather_scatter
 end 
@@ -779,10 +786,11 @@ are set to `snd`. Otherwise, either the optional `neighbors` or
 `neighbors` is also an `ExchangeGraph`
 that contains a super set of the outgoing and incoming neighbors
 associated with `snd`. It is used to find the incoming neighbors `rcv`
-efficiently. If `neighbors` are not provided, then `find_rcv_ids` 
+efficiently. If `neighbors` are not provided, then `find_rcv_ids`
 is used (either the user-provided or a default one).
 `find_rcv_ids` is a function that implements an algorithm to find the 
-rcv side of the exchange graph out of the snd side information.
+rcv side of the exchange graph out of the snd side information. It 
+defaults to [`default_find_rcv_ids`](@ref).
 """
 function ExchangeGraph(snd;
                        rcv=nothing,
@@ -836,8 +844,14 @@ function ExchangeGraph_impl_with_find_rcv_ids(snd_ids::AbstractArray,find_rcv_id
     ExchangeGraph(snd_ids,rcv_ids)
 end
 
-# This strategy gathers the communication graph into one process
-# and then scatters back the receivers
+"""
+    find_rcv_ids_gather_scatter(snd_ids::AbstractArray)
+
+Finds the `rcv` side of an `ExchangeGraph` out of the `snd` side information.
+
+This strategy gathers the communication graph into one process
+and then scatters back the receivers.
+"""
 function find_rcv_ids_gather_scatter(snd_ids::AbstractArray)
     snd_ids_main = gather(snd_ids)
     rcv_ids_main = map(snd_ids_main) do snd_ids_main