diff --git a/CHANGELOG.md b/CHANGELOG.md index df0e631a..7f59d084 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,12 +5,18 @@ 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 + +- Function `partition_from_color`. + ## [0.3.3] - 2023-08-09 ### Added -- Added an MPI ibarrier-based (supposedly scalable) algorithm to find rcv neighbours in a sparse all-to-all communication graph given the snd neighbors. We left the previous non-scalable algorithm as default (based on gather-scatter) until we have experimental evidence on the relative performance and scalability of the former with respect to the latter and for which core ranges. -- Added a new kwarg `discover_cols=true` to the `psparse!` constructor, which allows the user to skip column index discovery. +- MPI ibarrier-based (supposedly scalable) algorithm to find rcv neighbours in a sparse all-to-all communication graph given the snd neighbors. We left the previous non-scalable algorithm as default (based on gather-scatter) until we have experimental evidence on the relative performance and scalability of the former with respect to the latter and for which core ranges. +- New kwarg `discover_cols=true` to the `psparse!` constructor, which allows the user to skip column index discovery. ### Fixed diff --git a/docs/src/reference/partition.md b/docs/src/reference/partition.md index 99cdaab3..4f7ad290 100644 --- a/docs/src/reference/partition.md +++ b/docs/src/reference/partition.md @@ -5,6 +5,7 @@ ```@docs uniform_partition variable_partition +partition_from_color ``` ## AbstractLocalIndices diff --git a/src/PartitionedArrays.jl b/src/PartitionedArrays.jl index 6db70d1e..b2029575 100644 --- a/src/PartitionedArrays.jl +++ b/src/PartitionedArrays.jl @@ -61,6 +61,7 @@ include("mpi_array.jl") export PRange export uniform_partition export variable_partition +export partition_from_color export AbstractLocalIndices export OwnAndGhostIndices export LocalIndices diff --git a/src/p_range.jl b/src/p_range.jl index 5b9616a5..795f8351 100644 --- a/src/p_range.jl +++ b/src/p_range.jl @@ -696,6 +696,43 @@ function variable_partition( indices end +""" + partition_from_color(ranks,global_to_color;multicast=false,source=MAIN) + +Build an arbitrary 1d partition by defining the parts via the argument `global_to_color` (see below). +The output is a vector of vectors containing the indices in each component of +the partition. The `eltype` of the result implements the [`AbstractLocalIndices`](@ref) +interface. + +# Arguments + +- `ranks`: Array containing the distribution of ranks. +- `global_to_color`: If `multicast==false`, `global_to_color[gid]` contains the part id that owns the global id `gid`. If `multicast==true`, then `global_to_color[source][gid]` contains the part id that owns the global id `gid`. + +# Key-word arguments +- `multicast=false` +- `source=MAIN` + +This function is useful when generating a partition using a graph partitioner such as METIS. +The argument `global_to_color` is the usual output of such tools. +""" +function partition_from_color(ranks,global_to_color;multicast=false,source=MAIN) + if multicast == true + global_to_owner = getany(emit(global_to_color;source)) + else + global_to_owner = global_to_color + end + map(ranks) do rank + nglobal = length(global_to_owner) + own_to_global = findall(owner->owner==rank,global_to_owner) + ghost_to_global = Int[] + ghost_to_owner = Int32[] + own = OwnIndices(nglobal,rank,own_to_global) + ghost = GhostIndices(nglobal,ghost_to_global,ghost_to_owner) + OwnAndGhostIndices(own,ghost,global_to_owner) + end +end + function local_range(p,np,n,ghost=false,periodic=false) l = n รท np offset = l * (p-1) @@ -1119,23 +1156,28 @@ Local indices are defined by concatenating own and ghost ones. - `own::OwnIndices`: Container for the own indices. - `ghost::GhostIndices`: Container for the ghost indices. +- `global_to_owner`: [optional: it can be `nothing`] Vector containing the owner of each global id. # Supertype hierarchy - OwnAndGhostIndices <: AbstractLocalIndices + OwnAndGhostIndices{A} <: AbstractLocalIndices + +where `A=typeof(global_to_owner)`. """ -struct OwnAndGhostIndices <: AbstractLocalIndices +struct OwnAndGhostIndices{A} <: AbstractLocalIndices own::OwnIndices ghost::GhostIndices + global_to_owner::A assembly_cache::AssemblyCache @doc """ - OwnAndGhostIndices(own::OwnIndices,ghost::GhostIndices) + OwnAndGhostIndices(own::OwnIndices,ghost::GhostIndices,global_to_owner=nothing) - Build an instance of [`OwnAndGhostIndices`](@ref) from the underlying properties `own` and `ghost`. + Build an instance of [`OwnAndGhostIndices`](@ref) from the underlying properties `own`, `ghost`, and `global_to_owner`. """ - function OwnAndGhostIndices(own::OwnIndices,ghost::GhostIndices) - new(own,ghost,AssemblyCache()) + function OwnAndGhostIndices(own::OwnIndices,ghost::GhostIndices,global_to_owner=nothing) + A = typeof(global_to_owner) + new{A}(own,ghost,global_to_owner,AssemblyCache()) end end assembly_cache(a::OwnAndGhostIndices) = a.assembly_cache @@ -1146,6 +1188,15 @@ function replace_ghost(a::OwnAndGhostIndices,ghost::GhostIndices) OwnAndGhostIndices(a.own,ghost) end +function find_owner(indices,global_ids,::Type{<:OwnAndGhostIndices{T}}) where T + if T == Nothing + error("Not enough data to perform this operation without communciation") + end + map(indices,global_ids) do indices,global_ids + indices.global_to_owner[global_ids] + end +end + part_id(a::OwnAndGhostIndices) = a.own.owner function own_to_global(a::OwnAndGhostIndices) diff --git a/test/p_range_tests.jl b/test/p_range_tests.jl index daee9f9b..b3228a61 100644 --- a/test/p_range_tests.jl +++ b/test/p_range_tests.jl @@ -245,4 +245,51 @@ function p_range_tests(distribute) display(PRange(ids4)) + n = 21 + np = 4 + rank = distribute(LinearIndices((4,))) + global_to_color = zeros(Int,n) + for p in 1:np + global_to_color[local_range(p,np,n)] .= p + end + ids_color = partition_from_color(rank,global_to_color) + ids_uniform = uniform_partition(rank,n) + map(ids_color,ids_uniform) do ids1,ids2 + @test ids1 == ids2 + end + + global_to_color = map(rank) do rank + if rank == MAIN + my_global_to_color = zeros(Int,n) + for p in 1:np + my_global_to_color[local_range(p,np,n)] .= p + end + else + my_global_to_color = Int[] + end + my_global_to_color + end + ids_color = partition_from_color(rank,global_to_color;multicast=true) + ids_color = partition_from_color(rank,global_to_color;multicast=true,source=MAIN) + map(ids_color,ids_uniform) do ids1,ids2 + @test ids1 == ids2 + end + + gids = map(rank) do part + if part == 1 + [1,4,6] + elseif part == 2 + [3,1,2,8] + elseif part == 3 + [1,9,6] + else + [3,2,8,10] + end + end + owners1 = find_owner(ids_color,gids) + owners2 = find_owner(ids_uniform,gids) + map(owners1,owners2) do ids1,ids2 + @test ids1 == ids2 + end + end