From d12b6940df268aa91fb5b4fb058bdcd3afd80101 Mon Sep 17 00:00:00 2001 From: Timo Betcke Date: Sun, 15 Dec 2024 14:58:13 +0000 Subject: [PATCH] Cleaned up index layouts --- examples/map_index_layout.rs | 61 ++++++++++++ src/index_layout.rs | 93 +++++++++++++++++-- .../equidistributed_index_layout.rs | 65 ++----------- .../index_layout_from_local_counts.rs | 34 +++++++ src/permutation.rs | 14 +-- 5 files changed, 192 insertions(+), 75 deletions(-) create mode 100644 examples/map_index_layout.rs create mode 100644 src/index_layout/index_layout_from_local_counts.rs diff --git a/examples/map_index_layout.rs b/examples/map_index_layout.rs new file mode 100644 index 0000000..cb1d665 --- /dev/null +++ b/examples/map_index_layout.rs @@ -0,0 +1,61 @@ +//! Map betwen two index layouts + +use bempp_distributed_tools::{ + index_layout::{IndexLayout, IndexLayoutFromLocalCounts}, + EquiDistributedIndexLayout, +}; +use itertools::{izip, Itertools}; +use mpi::traits::Communicator; + +fn main() { + let universe = mpi::initialize().unwrap(); + let world = universe.world(); + + // Create an index layout with 10 indices on each rank. + + let layout1 = EquiDistributedIndexLayout::new(30, 1, &world); + + // Create a second layout with 5 indices on rank 0, 17 on rank 1 and 8 on rank 2. + + let counts = match world.rank() { + 0 => 5, + 1 => 17, + 2 => 8, + _ => panic!("This example only works with three processes."), + }; + + let layout2 = IndexLayoutFromLocalCounts::new(counts, &world); + + // Now we can map between the two layouts. + + let data = if world.rank() == 0 { + (0..10).collect_vec() + } else if world.rank() == 1 { + (10..20).collect_vec() + } else { + (20..30).collect_vec() + }; + + let mapped_data = layout1.remap(&layout2, &data); + + if world.rank() == 0 { + assert_eq!(mapped_data.len(), 5); + for (expected, &actual) in izip!(0..5, mapped_data.iter()) { + assert_eq!(expected, actual); + } + } else if world.rank() == 1 { + assert_eq!(mapped_data.len(), 17); + for (expected, &actual) in izip!(5..22, mapped_data.iter()) { + assert_eq!(expected, actual); + } + } else if world.rank() == 2 { + assert_eq!(mapped_data.len(), 8); + for (expected, &actual) in izip!(22..30, mapped_data.iter()) { + assert_eq!(expected, actual); + } + } + + let remapped_data = layout2.remap(&layout1, &mapped_data); + + assert_eq!(data, remapped_data); +} diff --git a/src/index_layout.rs b/src/index_layout.rs index 669efd8..2bcaa6a 100644 --- a/src/index_layout.rs +++ b/src/index_layout.rs @@ -10,7 +10,13 @@ /// index `last` is not contained on the process. If `first == last` then there is no index on /// the local process. mod equidistributed_index_layout; +mod index_layout_from_local_counts; pub use equidistributed_index_layout::EquiDistributedIndexLayout; +pub use index_layout_from_local_counts::IndexLayoutFromLocalCounts; +use itertools::Itertools; +use mpi::traits::{Communicator, Equivalence}; + +use crate::array_tools::redistribute; // An index layout specifying index ranges on each rank. // @@ -19,33 +25,106 @@ pub use equidistributed_index_layout::EquiDistributedIndexLayout; pub trait IndexLayout { /// MPI Communicator; type Comm: mpi::topology::Communicator; + + /// The cumulative sum of indices over the ranks. + /// + /// The number of indices on rank is is counts[1 + i] - counts[i]. + /// The last entry is the total number of indices. + fn counts(&self) -> &[usize]; + /// The local index range. If there is no local index /// the left and right bound are identical. - fn local_range(&self) -> (usize, usize); + fn local_range(&self) -> (usize, usize) { + let counts = self.counts(); + ( + counts[self.comm().rank() as usize], + counts[1 + self.comm().rank() as usize], + ) + } /// The number of global indices. - fn number_of_global_indices(&self) -> usize; + fn number_of_global_indices(&self) -> usize { + *self.counts().last().unwrap() + } /// The number of local indicies, that is the amount of indicies /// on my process. - fn number_of_local_indices(&self) -> usize; + fn number_of_local_indices(&self) -> usize { + let counts = self.counts(); + counts[1 + self.comm().rank() as usize] - counts[self.comm().rank() as usize] + } /// Index range on a given process. - fn index_range(&self, rank: usize) -> Option<(usize, usize)>; + fn index_range(&self, rank: usize) -> Option<(usize, usize)> { + let counts = self.counts(); + if rank < self.comm().size() as usize { + Some((counts[rank], counts[1 + rank])) + } else { + None + } + } /// Convert continuous (0, n) indices to actual indices. /// /// Assume that the local range is (30, 40). Then this method /// will map (0,10) -> (30, 40). /// It returns ```None``` if ```index``` is out of bounds. - fn local2global(&self, index: usize) -> Option; + fn local2global(&self, index: usize) -> Option { + let rank = self.comm().rank() as usize; + if index < self.number_of_local_indices() { + Some(self.counts()[rank] + index) + } else { + None + } + } /// Convert global index to local index on a given rank. /// Returns ```None``` if index does not exist on rank. - fn global2local(&self, rank: usize, index: usize) -> Option; + fn global2local(&self, rank: usize, index: usize) -> Option { + if let Some(index_range) = self.index_range(rank) { + if index >= index_range.1 { + return None; + } + + Some(index - index_range.0) + } else { + None + } + } /// Get the rank of a given index. - fn rank_from_index(&self, index: usize) -> Option; + fn rank_from_index(&self, index: usize) -> Option { + for (count_index, &count) in self.counts()[1..].iter().enumerate() { + if index < count { + return Some(count_index); + } + } + None + } + + /// Remap indices from one layout to another. + fn remap(&self, other: &L, data: &[T]) -> Vec { + assert_eq!(data.len(), self.number_of_local_indices()); + assert_eq!( + self.number_of_global_indices(), + other.number_of_global_indices() + ); + + let my_range = self.local_range(); + + let other_bins = (0..other.comm().size() as usize) + .map(|rank| other.index_range(rank).unwrap().0) + .collect_vec(); + + let sorted_keys = (my_range.0..my_range.1).collect_vec(); + + let counts = crate::array_tools::sort_to_bins(&sorted_keys, &other_bins) + .iter() + .map(|&key| key as i32) + .collect_vec(); + + redistribute(data, &counts, other.comm()) + } /// Return the communicator. fn comm(&self) -> &Self::Comm; diff --git a/src/index_layout/equidistributed_index_layout.rs b/src/index_layout/equidistributed_index_layout.rs index 4b55aee..ca6a76d 100644 --- a/src/index_layout/equidistributed_index_layout.rs +++ b/src/index_layout/equidistributed_index_layout.rs @@ -4,8 +4,6 @@ use mpi::traits::Communicator; /// Default index layout pub struct EquiDistributedIndexLayout<'a, C: Communicator> { - size: usize, - my_rank: usize, counts: Vec, comm: &'a C, } @@ -13,14 +11,13 @@ pub struct EquiDistributedIndexLayout<'a, C: Communicator> { impl<'a, C: Communicator> EquiDistributedIndexLayout<'a, C> { /// Crate new pub fn new(nchunks: usize, chunk_size: usize, comm: &'a C) -> Self { - let size = nchunks * chunk_size; + let nindices = nchunks * chunk_size; let comm_size = comm.size() as usize; assert!( comm_size > 0, "Group size is zero. At least one process needs to be in the group." ); - let my_rank = comm.rank() as usize; let mut counts = vec![0; 1 + comm_size]; // The following code computes what index is on what rank. No MPI operation necessary. @@ -37,10 +34,10 @@ impl<'a, C: Communicator> EquiDistributedIndexLayout<'a, C> { } for item in counts.iter_mut().take(comm_size).skip(nchunks) { - *item = size; + *item = nindices; } - counts[comm_size] = size; + counts[comm_size] = nindices; } else { // We want to equally distribute the range // among the ranks. Assume that we have 12 @@ -75,65 +72,15 @@ impl<'a, C: Communicator> EquiDistributedIndexLayout<'a, C> { } } - Self { - size, - my_rank, - counts, - comm, - } + Self { counts, comm } } } impl IndexLayout for EquiDistributedIndexLayout<'_, C> { type Comm = C; - fn index_range(&self, rank: usize) -> Option<(usize, usize)> { - if rank < self.comm.size() as usize { - Some((self.counts[rank], self.counts[1 + rank])) - } else { - None - } - } - - fn local_range(&self) -> (usize, usize) { - self.index_range(self.my_rank).unwrap() - } - - fn number_of_local_indices(&self) -> usize { - self.counts[1 + self.my_rank] - self.counts[self.my_rank] - } - - fn number_of_global_indices(&self) -> usize { - self.size - } - - fn local2global(&self, index: usize) -> Option { - if index < self.number_of_local_indices() { - Some(self.counts[self.my_rank] + index) - } else { - None - } - } - - fn global2local(&self, rank: usize, index: usize) -> Option { - if let Some(index_range) = self.index_range(rank) { - if index >= index_range.1 { - return None; - } - - Some(index - index_range.0) - } else { - None - } - } - - fn rank_from_index(&self, index: usize) -> Option { - for (count_index, &count) in self.counts[1..].iter().enumerate() { - if index < count { - return Some(count_index); - } - } - None + fn counts(&self) -> &[usize] { + &self.counts } fn comm(&self) -> &Self::Comm { diff --git a/src/index_layout/index_layout_from_local_counts.rs b/src/index_layout/index_layout_from_local_counts.rs new file mode 100644 index 0000000..31571a7 --- /dev/null +++ b/src/index_layout/index_layout_from_local_counts.rs @@ -0,0 +1,34 @@ +//! Default distributed index layout +use crate::index_layout::IndexLayout; +use mpi::traits::{Communicator, CommunicatorCollectives}; + +/// Specify an index layout from local variable counts +pub struct IndexLayoutFromLocalCounts<'a, C: Communicator> { + counts: Vec, + comm: &'a C, +} + +impl<'a, C: Communicator + CommunicatorCollectives> IndexLayoutFromLocalCounts<'a, C> { + /// Crate new + pub fn new(local_count: usize, comm: &'a C) -> Self { + let size = comm.size() as usize; + let mut counts = vec![0; size + 1]; + comm.all_gather_into(&local_count, &mut counts[1..]); + for i in 1..=size { + counts[i] += counts[i - 1]; + } + Self { counts, comm } + } +} + +impl IndexLayout for IndexLayoutFromLocalCounts<'_, C> { + type Comm = C; + + fn counts(&self) -> &[usize] { + &self.counts + } + + fn comm(&self) -> &Self::Comm { + &self.comm + } +} diff --git a/src/permutation.rs b/src/permutation.rs index ac9e0a2..6d63f39 100644 --- a/src/permutation.rs +++ b/src/permutation.rs @@ -11,7 +11,7 @@ use crate::index_layout::IndexLayout; /// Permuation of data. pub struct DataPermutation<'a, L: IndexLayout> { index_layout: &'a L, - custom_indices: &'a [usize], + nindices: usize, my_rank: usize, custom_local_indices: Vec, local_to_custom_map: Vec, @@ -21,11 +21,7 @@ pub struct DataPermutation<'a, L: IndexLayout> { impl<'a, L: IndexLayout> DataPermutation<'a, L> { /// Create a new permutation object. - pub fn new( - index_layout: &'a L, - custom_indices: &'a [usize], - comm: &C, - ) -> Self { + pub fn new(index_layout: &'a L, custom_indices: &[usize], comm: &C) -> Self { // We first need to identify which custom indices are local and which are global. let my_rank = comm.rank() as usize; @@ -72,7 +68,7 @@ impl<'a, L: IndexLayout> DataPermutation<'a, L> { Self { index_layout, - custom_indices, + nindices: custom_indices.len(), my_rank, custom_local_indices, local_to_custom_map, @@ -88,7 +84,7 @@ impl<'a, L: IndexLayout> DataPermutation<'a, L> { permuted_data: &mut [T], ) { assert_eq!(data.len(), self.index_layout.number_of_local_indices()); - assert_eq!(permuted_data.len(), self.custom_indices.len()); + assert_eq!(permuted_data.len(), self.nindices); // We first need to get the send data. This is quite easy. We can just // use the global2local method from the index layout. @@ -127,7 +123,7 @@ impl<'a, L: IndexLayout> DataPermutation<'a, L> { data: &[T], permuted_data: &mut [T], ) { - assert_eq!(data.len(), self.custom_indices.len()); + assert_eq!(data.len(), self.nindices); assert_eq!( permuted_data.len(), self.index_layout.number_of_local_indices()