Progress

kifmm/examples/mpi_test_fmm.rs

@@ -17,7 +17,7 @@ fn main() {
         datatype::PartitionMut,
         traits::{Communicator, Root},
     };
-    use rlst::RawAccess;
+    use rlst::{rlst_dynamic_array2, RawAccess, RawAccessMut};
 
     let (universe, _threading) = mpi::initialize_with_threading(mpi::Threading::Funneled).unwrap();
     let world = universe.world();
@@ -45,6 +45,10 @@ fn main() {
 
         // Generate some random test data local to each process
         let points = points_fixture::<f32>(n_points, None, None, Some(world.rank() as u64));
+        let nvecs = 1;
+        let tmp = vec![1.0; n_points];
+        let mut charges = rlst_dynamic_array2!(f32, [n_points, 1]);
+        charges.data_mut().copy_from_slice(&tmp);
 
         let mut multi_fmm = MultiNodeBuilder::new(false)
             .tree(
@@ -57,7 +61,7 @@ fn main() {
                 sort_kind.clone(),
             )
             .unwrap()
-            .parameters(expansion_order, kernel.clone(), source_to_target)
+            .parameters(charges.data(), expansion_order, kernel.clone(), source_to_target)
             .unwrap()
             .build()
             .unwrap();
@@ -158,6 +162,10 @@ fn main() {
 
         // Generate some random test data local to each process
         let points = points_fixture::<f32>(n_points, None, None, Some(world.rank() as u64));
+        let nvecs = 1;
+        let tmp = vec![1.0; n_points];
+        let mut charges = rlst_dynamic_array2!(f32, [n_points, 1]);
+        charges.data_mut().copy_from_slice(&tmp);
 
         let mut multi_fmm = MultiNodeBuilder::new(false)
             .tree(
@@ -170,7 +178,7 @@ fn main() {
                 sort_kind,
             )
             .unwrap()
-            .parameters(expansion_order, kernel, source_to_target)
+            .parameters(charges.data(), expansion_order, kernel, source_to_target)
             .unwrap()
             .build()
             .unwrap();

kifmm/examples/mpi_test_layout.rs

@@ -19,7 +19,7 @@ fn main() {
         datatype::PartitionMut,
         traits::{Communicator, CommunicatorCollectives},
     };
-    use rlst::RawAccess;
+    use rlst::{rlst_dynamic_array2, RawAccess, RawAccessMut};
 
     let (universe, _threading) = mpi::initialize_with_threading(mpi::Threading::Single).unwrap();
     let world = universe.world();
@@ -40,6 +40,10 @@ fn main() {
 
         // Generate some random test data local to each process
         let points = points_fixture::<f32>(n_points, None, None, None);
+        let nvecs = 1;
+        let tmp = vec![1.0; n_points];
+        let mut charges = rlst_dynamic_array2!(f32, [n_points, 1]);
+        charges.data_mut().copy_from_slice(&tmp);
 
         ThreadPoolBuilder::new()
             .num_threads(1)
@@ -58,7 +62,7 @@ fn main() {
                 sort_kind,
             )
             .unwrap()
-            .parameters(expansion_order, kernel, source_to_target)
+            .parameters(charges.data(), expansion_order, kernel, source_to_target)
             .unwrap()
             .build()
             .unwrap();

kifmm/examples/mpi_test_p2m.rs

@@ -24,7 +24,7 @@ fn main() {
         datatype::PartitionMut,
         traits::{Communicator, Root},
     };
-    use rlst::RawAccess;
+    use rlst::{rlst_dynamic_array2, RawAccess, RawAccessMut};
 
     let (universe, _threading) = mpi::initialize_with_threading(mpi::Threading::Single).unwrap();
     let world = universe.world();
@@ -44,6 +44,9 @@ fn main() {
 
     // Generate some random test data local to each process
     let points = points_fixture::<f32>(n_points, None, None, None);
+    let tmp = vec![1.0; n_points];
+    let mut charges = rlst_dynamic_array2!(f32, [n_points, 1]);
+    charges.data_mut().copy_from_slice(&tmp);
 
     ThreadPoolBuilder::new()
         .num_threads(1)
@@ -62,7 +65,7 @@ fn main() {
             sort_kind,
         )
         .unwrap()
-        .parameters(expansion_order, kernel, source_to_target)
+        .parameters(charges.data(), expansion_order, kernel, source_to_target)
         .unwrap()
         .build()
         .unwrap();

kifmm/examples/mpi_test_trees.rs

@@ -21,7 +21,7 @@ fn main() {
         datatype::PartitionMut,
         traits::{Communicator, Root},
     };
-    use rlst::RawAccess;
+    use rlst::{rlst_dynamic_array2, RawAccess, RawAccessMut};
 
     let (universe, _threading) = mpi::initialize_with_threading(mpi::Threading::Single).unwrap();
     let world = universe.world();
@@ -41,6 +41,9 @@ fn main() {
 
     // Generate some random test data local to each process
     let points = points_fixture::<f32>(n_points, None, None, Some(world.rank() as u64));
+    let tmp = vec![1.0; n_points];
+    let mut charges = rlst_dynamic_array2!(f32, [n_points, 1]);
+    charges.data_mut().copy_from_slice(&tmp);
 
     ThreadPoolBuilder::new()
         .num_threads(1)
@@ -58,7 +61,7 @@ fn main() {
             sort_kind,
         )
         .unwrap()
-        .parameters(expansion_order, kernel, source_to_target)
+        .parameters(charges.data(), expansion_order, kernel, source_to_target)
         .unwrap()
         .build()
         .unwrap();

kifmm/examples/mpi_test_upward_pass.rs

@@ -17,7 +17,7 @@ fn main() {
         datatype::PartitionMut,
         traits::{Communicator, Root},
     };
-    use rlst::{RawAccess, RlstScalar};
+    use rlst::{rlst_dynamic_array2, RawAccess, RawAccessMut, RlstScalar};
 
     let (universe, _threading) = mpi::initialize_with_threading(mpi::Threading::Funneled).unwrap();
     let world = universe.world();
@@ -38,6 +38,9 @@ fn main() {
 
     // Generate some random test data local to each process
     let points = points_fixture::<f32>(n_points, None, None, None);
+    let tmp = vec![1.0; n_points];
+    let mut charges = rlst_dynamic_array2!(f32, [n_points, 1]);
+    charges.data_mut().copy_from_slice(&tmp);
 
     let mut fmm = MultiNodeBuilder::new(false)
         .tree(
@@ -50,7 +53,7 @@ fn main() {
             sort_kind,
         )
         .unwrap()
-        .parameters(expansion_order, kernel, source_to_target)
+        .parameters(charges.data(), expansion_order, kernel, source_to_target)
         .unwrap()
         .build()
         .unwrap();

kifmm/src/fmm.rs

@@ -3,7 +3,7 @@ mod builder;
 pub mod constants;
 mod data_access;
 mod eval;
-mod field_translation;
+pub mod field_translation;
 pub mod helpers;
 pub mod isa;
 mod kernel;

kifmm/src/fmm/builder/multi_node.rs

@@ -1,6 +1,9 @@
+use itertools::Itertools;
+
 use mpi::{
+    datatype::{Partition, PartitionMut},
     topology::{Communicator, SimpleCommunicator},
-    traits::Equivalence,
+    traits::{Equivalence, CommunicatorCollectives}, Count,
 };
 use num::Float;
 use rlst::{rlst_dynamic_array2, MatrixSvd, RlstScalar};
@@ -22,6 +25,7 @@ use crate::{
         fmm::{HomogenousKernel, Metadata, MetadataAccess},
         general::{multi_node::GlobalFmmMetadata, single_node::Epsilon},
         types::{CommunicationTime, CommunicationType, MetadataTime, MetadataType},
+        tree::MultiTree,
     },
     tree::{
         types::{Domain, SortKind},
@@ -206,6 +210,7 @@ where
     /// Parameters
     pub fn parameters(
         mut self,
+        charges: &[Scalar],
         expansion_order: usize,
         kernel: Kernel,
         source_to_target: FieldTranslation,
@@ -216,7 +221,68 @@ where
                 "Must build tree before specifying FMM parameters",
             ))
         } else {
-            // TODO: Mapping of global indices needs to happen here eventually.
+
+            let global_indices = &self
+                .tree
+                .as_ref()
+                .unwrap()
+                .source_tree
+                .unsorted_global_indices;
+
+            let destination_ranks = &self
+                .tree
+                .as_ref()
+                .unwrap()
+                .source_tree
+                .destination_ranks;
+
+            let packets= charges.iter().zip(global_indices.iter())
+                    .map(|(c, gidx)| (c, gidx)).collect_vec();
+
+            // Communicate message sizes
+            let size = self.tree.as_ref().unwrap().source_tree.communicator.size();
+            let communicator = &self.tree.as_ref().unwrap().source_tree.communicator;
+
+            let mut counts_snd = vec![0i32; size as usize];
+            for &destination_rank in destination_ranks.iter() {
+                counts_snd[destination_rank as usize] += 1
+            }
+
+            let displs_snd = counts_snd
+                .iter()
+                .scan(0, |acc, &x| {
+                    let tmp = *acc;
+                    *acc += x;
+                    Some(tmp)
+                })
+                .collect_vec();
+
+            let mut counts_recv = vec![0 as Count; size as usize];
+
+            communicator.all_to_all_into(&counts_snd, &mut counts_recv);
+
+            // Communicate packets
+            let displs_recv = counts_recv
+                .iter()
+                .scan(0, |acc, &x| {
+                    let tmp = *acc;
+                    *acc += x;
+                    Some(tmp)
+                })
+                .collect_vec();
+
+            let total = counts_recv.iter().sum::<Count>();
+
+            let total = counts_recv.iter().sum::<Count>();
+
+            let mut received = vec![T::default(); total as usize];
+            let mut partition_received =
+                PartitionMut::new(&mut received[..], counts_recv, &displs_recv[..]);
+            let partition_snd = Partition::new(&packets[..], counts_snd, &displs_snd[..]);
+
+
+
+
             self.n_coeffs_equivalent_surface = Some(ncoeffs_kifmm(expansion_order));
             self.n_coeffs_check_surface = Some(ncoeffs_kifmm(expansion_order));
             self.kernel = Some(kernel);

kifmm/src/sorting/samplesort.rs

@@ -6,7 +6,7 @@ use mpi::{
     datatype::{Partition, PartitionMut},
     topology::SimpleCommunicator,
     traits::{Communicator, CommunicatorCollectives, Equivalence},
-    Count,
+    Count, Rank,
 };
 use rand::{thread_rng, Rng};
 
@@ -21,7 +21,7 @@ pub fn samplesort<T>(
     array: &mut Vec<T>,
     communicator: &SimpleCommunicator,
     n_samples: usize,
-) -> Result<(), std::io::Error>
+) -> Result<Option<Vec<Rank>>, std::io::Error>
 where
     T: Equivalence + Ord + Default + Clone + Debug,
 {
@@ -67,7 +67,7 @@ where
             .step_by(n_samples)
             .collect_vec();
 
-        let mut ranks = Vec::new();
+        let mut destination_ranks = Vec::new();
         for item in array.iter() {
             let mut rank_index = -1i32;
             for (i, splitter) in splitters.iter().enumerate() {
@@ -80,12 +80,12 @@ where
                 }
             }
 
-            ranks.push(rank_index);
+            destination_ranks.push(rank_index);
         }
 
         let mut counts_snd = vec![0i32; size as usize];
-        for &rank in ranks.iter() {
-            counts_snd[rank as usize] += 1
+        for &destination_rank in destination_ranks.iter() {
+            counts_snd[destination_rank as usize] += 1
         }
 
         let displs_snd = counts_snd
@@ -120,7 +120,9 @@ where
         communicator.all_to_all_varcount_into(&partition_snd, &mut partition_received);
         received.sort();
         *array = received;
+
+        return Ok(Some(destination_ranks))
     }
 
-    Ok(())
+    Ok(None)
 }

kifmm/src/sorting/simplesort.rs

@@ -7,7 +7,7 @@ use mpi::{
     datatype::{Partition, PartitionMut},
     topology::SimpleCommunicator,
     traits::{Communicator, CommunicatorCollectives, Equivalence},
-    Count,
+    Count, Rank,
 };
 
 /// Simple sort is a bucket style sort specialised for octrees. In this case, the number of 'splitters'
@@ -26,7 +26,7 @@ pub fn simplesort<T>(
     array: &mut Vec<T>,
     communicator: &SimpleCommunicator,
     splitters: &mut [T],
-) -> Result<(), std::io::Error>
+) -> Result<Option<Vec<Rank>>, std::io::Error>
 where
     T: Equivalence + Ord + Default + Clone + Debug,
 {
@@ -47,7 +47,7 @@ where
 
     if size > 1 {
         // Sort local data into buckets
-        let mut ranks = Vec::new();
+        let mut destination_ranks = Vec::new();
         for item in array.iter() {
             let mut rank_index = -1i32;
             for (i, splitter) in splitters.iter().enumerate() {
@@ -60,11 +60,11 @@ where
                 }
             }
 
-            ranks.push(rank_index);
+            destination_ranks.push(rank_index);
         }
 
         let mut counts_snd = vec![0i32; size as usize];
-        for &rank in ranks.iter() {
+        for &rank in destination_ranks.iter() {
             counts_snd[rank as usize] += 1
         }
 
@@ -101,7 +101,9 @@ where
         received.sort();
 
         *array = received;
+
+        return Ok(Some(destination_ranks))
     }
 
-    Ok(())
+    Ok(None)
 }

kifmm/src/traits/tree.rs

@@ -81,7 +81,7 @@ pub trait SingleTree {
     /// - `leaf` - node being queried.
     fn global_indices(&self, leaf: &Self::Node) -> Option<&[usize]>;
 
-    /// gets all global indices (local in mult inode setting)
+    /// gets all global indices
     fn all_global_indices(&self) -> Option<&[usize]>;
 
     /// Get domain defined by the points, gets global domain in multi node setting.
@@ -217,6 +217,9 @@ pub trait MultiTree {
     /// - `idx` - Index being query.
     fn node(&self, idx: usize) -> Option<&<Self::SingleTree as SingleTree>::Node>;
 
+    /// gets all global indices
+    fn all_global_indices(&self) -> Option<&[usize]>;
+
     /// Get domain defined by the points across all nodes.
     fn domain(&self) -> &<Self::SingleTree as SingleTree>::Domain;
 }