Fix bug in modified locateParticles

src/Particle/ParticleSpatialLayout.h

@@ -76,7 +76,7 @@ namespace ippl {
         RegionLayout_t rlayout_m;
 
         //! The FieldLayout containing information on nearest neighbors
-        FieldLayout_t flayout_m;
+        FieldLayout_t& flayout_m;
 
         //! Type of the Kokkos view containing the local regions.
         using region_view_type = typename RegionLayout_t::view_type;
@@ -89,6 +89,14 @@ namespace ippl {
         KOKKOS_INLINE_FUNCTION constexpr static bool positionInRegion(
             const std::index_sequence<Idx...>&, const vector_type& pos, const region_type& region);
 
+        /*!
+         * Evaluates the total number of MPI ranks sharing the spatial nearest neighbors.
+         * @param neighbors structure containing, for every spatial direction, a list of
+         * MPI ranks IDs corresponding to the nearest neighbors of the current local domain section.
+         * @return The total number of the ranks.
+         */
+        size_type getNeighborSize(const neighbor_list& neighbors) const;
+
     public:
         /*!
          * For each particle in the bunch, determine the rank on which it should

src/Particle/ParticleSpatialLayout.hpp

@@ -27,8 +27,36 @@
 
 #include "Communicate/Window.h"
 
+
 namespace ippl {
 
+    /*!
+     * We need this struct since Kokkos parallel_scan only accects
+     * one variable of type ReturnType where to perform the reduction operation.
+     * For more details, see
+     * https://kokkos.github.io/kokkos-core-wiki/API/core/parallel-dispatch/parallel_scan.html.
+     */
+    struct increment_type {
+        size_t count[2];
+
+        KOKKOS_FUNCTION void init() {
+            count[0] = 0;
+            count[1] = 0;
+        }
+
+        KOKKOS_INLINE_FUNCTION increment_type& operator+=(bool* values) {
+            count[0] += values[0];
+            count[1] += values[1];
+            return *this;
+        }
+
+        KOKKOS_INLINE_FUNCTION increment_type& operator+=(increment_type values) {
+            count[0] += values.count[0];
+            count[1] += values.count[1];
+            return *this;
+        }
+    };
+
     template <typename T, unsigned Dim, class Mesh, typename... Properties>
     ParticleSpatialLayout<T, Dim, Mesh, Properties...>::ParticleSpatialLayout(FieldLayout<Dim>& fl,
                                                                               Mesh& mesh)
@@ -38,6 +66,7 @@ namespace ippl {
     template <typename T, unsigned Dim, class Mesh, typename... Properties>
     void ParticleSpatialLayout<T, Dim, Mesh, Properties...>::updateLayout(FieldLayout<Dim>& fl,
                                                                           Mesh& mesh) {
+        //flayout_m = fl;
         rlayout_m.changeDomain(fl, mesh);
     }
 
@@ -163,34 +192,153 @@ namespace ippl {
         return ((pos[Idx] > region[Idx].min()) && ...) && ((pos[Idx] <= region[Idx].max()) && ...);
     };
 
+
+    /* Helper function that evaluates the total number of neighbors for the current rank in Dim dimensions.
+    */
+    template <typename T, unsigned Dim, class Mesh, typename... Properties>
+    detail::size_type ParticleSpatialLayout<T, Dim, Mesh, Properties...>::getNeighborSize(
+        const neighbor_list& neighbors) const {
+        size_type totalSize = 0;
+
+        for (const auto& componentNeighbors : neighbors) {
+            totalSize += componentNeighbors.size();
+        }
+
+        return totalSize;
+    }
+
     template <typename T, unsigned Dim, class Mesh, typename... Properties>
     template <typename ParticleContainer>
     detail::size_type ParticleSpatialLayout<T, Dim, Mesh, Properties...>::locateParticles(
         const ParticleContainer& pc, locate_type& ranks, bool_type& invalid) const {
-        auto& positions                            = pc.R.getView();
-        typename RegionLayout_t::view_type Regions = rlayout_m.getdLocalRegions();
+        auto positions           = pc.R.getView();
+        region_view_type Regions = rlayout_m.getdLocalRegions();
 
         using mdrange_type = Kokkos::MDRangePolicy<Kokkos::Rank<2>, position_execution_space>;
 
-        int myRank = Comm->rank();
+        size_type myRank = Comm->rank();
 
         const auto is = std::make_index_sequence<Dim>{};
 
-        size_type invalidCount = 0;
-        Kokkos::parallel_reduce(
+        const neighbor_list& neighbors = flayout_m.getNeighbors();
+
+        /// outsideIds: Container of particle IDs that travelled outside of the neighborhood.
+        locate_type outsideIds("Particles outside of neighborhood", size_type(pc.getLocalNum()));
+
+        /// outsideCount: Tracks the number of particles that travelled outside of the neighborhood.
+        size_type outsideCount       = 0;
+        /// invalidCount: Tracks the number of particles that need to be sent to other ranks.
+        size_type invalidCount       = 0;
+
+        /// neighborSize: Size of a neighborhood in D dimentions.
+        const size_type neighborSize = getNeighborSize(neighbors);
+
+        /// neighbors_view: Kokkos view with the IDs of the neighboring MPI ranks.
+        locate_type neighbors_view("Nearest neighbors IDs", neighborSize);
+
+        /* red_val: Used to reduce both the number of invalid particles and the number of particles 
+         * outside of the neighborhood (Kokkos::parallel_scan doesn't allow multiple reduction values, so we
+         * use the helper class increment_type). First element updates InvalidCount, second
+         * one updates outsideCount.
+        */
+        increment_type red_val;
+        red_val.init();
+
+        auto neighbors_mirror =
+            Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), neighbors_view);
+
+        size_t k = 0;
+
+        for (const auto& componentNeighbors : neighbors) {
+            for (size_t j = 0; j < componentNeighbors.size(); ++j) {
+                neighbors_mirror(k) = componentNeighbors[j];
+                //std::cout << "Neighbor: " << neighbors_mirror(k) << std::endl;
+                k++;
+            }
+        }
+
+        Kokkos::deep_copy(neighbors_view, neighbors_mirror);
+
+        /*! Begin Kokkos loop:
+         * Step 1: search in current rank
+         * Step 2: search in neighbors
+         * Step 3: save information on whether the particle was located
+         * Step 4: run additional loop on non-located particles
+         */
+
+        Kokkos::parallel_scan(
             "ParticleSpatialLayout::locateParticles()",
-            mdrange_type({0, 0}, {ranks.extent(0), Regions.extent(0)}),
-            KOKKOS_LAMBDA(const size_t i, const size_type j, size_type& count) {
-                bool xyz_bool = positionInRegion(is, positions(i), Regions(j));
-                if (xyz_bool) {
-                    ranks(i)   = j;
-                    invalid(i) = (myRank != ranks(i));
-                    count += invalid(i);
+            Kokkos::RangePolicy<size_t>(0, ranks.extent(0)),
+            KOKKOS_LAMBDA(const size_type i, increment_type& val, const bool final) {
+                /* Step 1
+                * inCurr: True if the particle hasn't left the current MPI rank.
+                * inNeighbor: True if the particle is found in a neighboring rank.
+                * found: True either if inCurr = True or inNeighbor = True.
+                * increment: Helper variable to update red_val.
+                */
+                bool inCurr = false;
+                bool inNeighbor = false;
+                bool found = false;
+                bool increment[2];
+
+                inCurr = positionInRegion(is, positions(i), Regions(myRank)); 
+
+                ranks(i)     = inCurr * myRank;
+                invalid(i)   = !inCurr;
+                found =  inCurr || found;
+
+                /// Step 2
+                for (size_t j = 0; j < neighbors_view.extent(0); ++j) {
+                    size_type rank = neighbors_view(j);
+
+                    inNeighbor = positionInRegion(is, positions(i), Regions(rank));
+
+                    ranks(i) = !(inNeighbor) * ranks(i) + inNeighbor * rank;
+                    found =  inNeighbor || found;
+                }
+
+                /// Step 3
+                /* isOut: When the last thread has finished the search, checks whether the particle has been found 
+                 * either in the current rank or in a neighboring one.
+                 * Used to avoid race conditions when updating outsideIds.
+                 */
+                if(final && !found) {
+                    outsideIds(val.count[1]) = i;
                 }
+                //outsideIds(val.count[1]) = i * isOut;
+                increment[0] = invalid(i);
+                increment[1] = !found;
+                val += increment;
+
             },
-            Kokkos::Sum<size_type>(invalidCount));
+            red_val);
+
         Kokkos::fence();
 
+        invalidCount  = red_val.count[0];
+        outsideCount  = red_val.count[1];
+
+        /// Step 4 
+        if (outsideCount > 0) {
+            static IpplTimings::TimerRef nonNeighboringParticles = IpplTimings::getTimer("nonNeighboringParticles");
+            IpplTimings::startTimer(nonNeighboringParticles);
+
+            Kokkos::parallel_for(
+                "ParticleSpatialLayout::leftParticles()",
+                mdrange_type({0, 0}, {outsideCount, Regions.extent(0)}),
+                KOKKOS_LAMBDA(const size_t i, const size_type j) {
+                    /// pID: (local) ID of the particle that is currently being searched.
+                    size_type pId = outsideIds(i);
+
+                    /// inRegion: Checks whether particle pID is inside region j.
+                    bool inRegion = positionInRegion(is, positions(pId), Regions(j));
+                    ranks(pId) = inRegion * j + !(inRegion) * ranks(pId);
+                });
+            Kokkos::fence();
+
+            IpplTimings::stopTimer(nonNeighboringParticles);
+        }
+
         return invalidCount;
     }