From feac8d56e5cc84f63f6e7be7e1668113558c9e47 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Wed, 25 Sep 2019 14:53:21 -0700
Subject: Implement `iota` on GPU

---
 Source/Particles/Sorting/Partition.cpp | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

(limited to 'Source/Particles/Sorting/Partition.cpp')

diff --git a/Source/Particles/Sorting/Partition.cpp b/Source/Particles/Sorting/Partition.cpp
index 29e187821..49d4d25a6 100644
--- a/Source/Particles/Sorting/Partition.cpp
+++ b/Source/Particles/Sorting/Partition.cpp
@@ -1,5 +1,6 @@
-#include <WarpX.H>
+#include <SortingUtils.H>
 #include <PhysicalParticleContainer.H>
+#include <WarpX.H>
 #include <AMReX_Particles.H>
 
 using namespace amrex;
@@ -43,7 +44,7 @@ PhysicalParticleContainer::PartitionParticlesInBuffers(
     Gpu::ManagedDeviceVector<int> inexflag;
     inexflag.resize(np);
 
-    // We need to partition the large buffer first
+    // Select the largest buffer first
     iMultiFab const* bmasks =
         (WarpX::n_field_gather_buffer >= WarpX::n_current_deposition_buffer) ?
         gather_masks : current_masks;
@@ -60,9 +61,10 @@ PhysicalParticleContainer::PartitionParticlesInBuffers(
          }
      );
 
-    Vector<long> pid;
+    // Partition the particles according whether they are in the large buffer or not
+    Gpu::ManagedDeviceVector<long> pid;
     pid.resize(np);
-    std::iota(pid.begin(), pid.end(), 0L);
+    fillWithConsecutiveIntegers( pid );
     auto sep = std::stable_partition(pid.begin(), pid.end(),
                                      [&inexflag](long id) { return inexflag[id]; });
 
-- 
cgit v1.2.3