Merge RedistributeCPU and RedistributeGPU into one implementation

Src/Base/AMReX_GpuAtomic.H

@@ -18,8 +18,7 @@ namespace Gpu::Atomic {
 // For LogicalOr and LogicalAnd, the data type is int.
 // For Exch and CAS, the data type is generic.
 // All these functions are non-atomic in host code!!!
-// If one needs them to be atomic in host code, use HostDevice::Atomic::*.  Currently only
-// HostDevice::Atomic::Add is supported.  We could certainly add more.
+// If one needs them to be atomic in host code, use HostDevice::Atomic::*.
 // If we add more types for atomicAdd, we also need to update HasAtomicAdd in AMReX_TypeTraits.H.
 
 /// \cond DOXYGEN_IGNORE
@@ -617,6 +616,21 @@ namespace HostDevice::Atomic {
         *sum += value;
     }
 
+    template <class T>
+    AMREX_FORCE_INLINE
+    T FetchAdd_Host (T* const sum, T const value) noexcept
+    {
+        T old;
+#ifdef AMREX_USE_OMP
+#pragma omp atomic capture
+#endif
+        {
+            old = *sum;
+            *sum += value;
+        }
+        return old;
+    }
+
     template <class T>
     AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     void Add (T* const sum, T const value) noexcept
@@ -625,6 +639,14 @@ namespace HostDevice::Atomic {
         AMREX_IF_ON_HOST((Add_Host(sum,value);))
     }
 
+    template <class T>
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    T FetchAdd (T* const sum, T const value) noexcept
+    {
+        AMREX_IF_ON_DEVICE((return Gpu::Atomic::Add(sum,value);))
+        AMREX_IF_ON_HOST((return FetchAdd_Host(sum,value);))
+    }
+
 }
 
 #ifdef AMREX_USE_GPU

Src/Particle/AMReX_NeighborParticles.H

@@ -16,6 +16,7 @@ namespace amrex {
   struct NeighborCode
   {
       int grid_id;
+      Box grid_box;
       IntVect periodic_shift;
   };
 

Src/Particle/AMReX_NeighborParticlesGPUImpl.H

@@ -99,6 +99,7 @@ buildNeighborMask ()
             {
                 NeighborCode code;
                 code.grid_id        = nbor_grid.grid_id;
+                code.grid_box       = ba[nbor_grid.grid_id];
                 code.periodic_shift = nbor_grid.periodic_shift;
                 h_code_arr.push_back(code);
                 h_isec_boxes.push_back(nbor_grid.box);
@@ -170,6 +171,8 @@ buildNeighborCopyOp (bool use_boundary_neighbor)
         auto p_code_array   = m_code_array[gid].dataPtr();
         auto p_isec_boxes   = m_isec_boxes[gid].dataPtr();
         const int nisec_box = m_isec_boxes[gid].size();
+        const bool do_tiling = this->do_tiling;
+        const IntVect tile_size = this->tile_size;
         // auto p_code_offsets = m_code_offsets[gid].dataPtr();
 
         AMREX_FOR_1D ( np, i,
@@ -194,12 +197,14 @@ buildNeighborCopyOp (bool use_boundary_neighbor)
         Gpu::dtoh_memcpy_async(&num_copies, offsets.data()+np, sizeof(int));
         Gpu::streamSynchronize();
 
-        neighbor_copy_op.resize(gid, lev, num_copies);
+        neighbor_copy_op.resize(gid, tid, lev, num_copies);
 
-        auto p_boxes = neighbor_copy_op.m_boxes[lev][gid].dataPtr();
-        auto p_levs = neighbor_copy_op.m_levels[lev][gid].dataPtr();
-        auto p_src_indices = neighbor_copy_op.m_src_indices[lev][gid].dataPtr();
-        auto p_periodic_shift = neighbor_copy_op.m_periodic_shift[lev][gid].dataPtr();
+        auto tile_index = std::make_pair(gid, tid);
+        auto p_boxes = neighbor_copy_op.m_boxes[lev][tile_index].dataPtr();
+        auto p_levs = neighbor_copy_op.m_levels[lev][tile_index].dataPtr();
+        auto p_tiles = neighbor_copy_op.m_tiles[lev][tile_index].dataPtr();
+        auto p_src_indices = neighbor_copy_op.m_src_indices[lev][tile_index].dataPtr();
+        auto p_periodic_shift = neighbor_copy_op.m_periodic_shift[lev][tile_index].dataPtr();
 
         Gpu::streamSynchronize();
         AMREX_FOR_1D ( np, i,
@@ -213,6 +218,9 @@ buildNeighborCopyOp (bool use_boundary_neighbor)
             for (int j=0; j<nisec_box; ++j) {
                 if (p_isec_boxes[j].contains(iv)) {
                     p_boxes[k]          = p_code_array[j].grid_id;
+                    Box tbx;
+                    p_tiles[k]          = getTileIndex(iv, p_code_array[j].grid_box,
+                                                       do_tiling, tile_size, tbx);
                     p_levs[k]           = 0;
                     p_periodic_shift[k] = p_code_array[j].periodic_shift;
                     p_src_indices[k]    = pid;
@@ -235,6 +243,9 @@ fillNeighborsGPU ()
 
     AMREX_ASSERT(numParticlesOutOfRange(*this, 0) == 0);
 
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(this->do_tiling == 0,
+        "Tiling on the GPU is not supported for neighbor particles.");
+
     buildNeighborMask();
     this->defineBufferMap();
 

Src/Particle/AMReX_ParticleBufferMap.H

@@ -13,39 +13,45 @@ namespace amrex {
 struct GetPID
 {
     const int* m_bucket_to_pid;
-    const int* m_lev_gid_to_bucket;
+    const int* m_lev_gid_tid_to_bucket;
     const int* m_lev_offsets;
+    const int* m_gid_offsets;
 
     GetPID (const Gpu::DeviceVector<int>& bucket_to_pid,
-            const Gpu::DeviceVector<int>& lev_gid_to_bucket,
-            const Gpu::DeviceVector<int>& lev_offsets)
+            const Gpu::DeviceVector<int>& lev_gid_tid_to_bucket,
+            const Gpu::DeviceVector<int>& lev_offsets,
+            const Gpu::DeviceVector<int>& gid_offsets)
         : m_bucket_to_pid(bucket_to_pid.dataPtr()),
-          m_lev_gid_to_bucket(lev_gid_to_bucket.dataPtr()),
-          m_lev_offsets(lev_offsets.dataPtr())
+          m_lev_gid_tid_to_bucket(lev_gid_tid_to_bucket.dataPtr()),
+          m_lev_offsets(lev_offsets.dataPtr()),
+          m_gid_offsets(gid_offsets.dataPtr())
         {}
 
     AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
-    int operator() (const int lev, const int gid) const noexcept
+    int operator() (const int lev, const int gid, const int tid) const noexcept
     {
-        return m_bucket_to_pid[m_lev_gid_to_bucket[m_lev_offsets[lev]+gid]];
+        return m_bucket_to_pid[m_lev_gid_tid_to_bucket[m_gid_offsets[m_lev_offsets[lev]+gid] + tid]];
     }
 };
 
 struct GetBucket
 {
-    const int* m_lev_gid_to_bucket;
+    const int* m_lev_gid_tid_to_bucket;
     const int* m_lev_offsets;
-
-    GetBucket (const int* lev_gid_to_bucket_ptr,
-               const int* lev_offsets_ptr)
-        : m_lev_gid_to_bucket(lev_gid_to_bucket_ptr),
-          m_lev_offsets(lev_offsets_ptr)
+    const int* m_gid_offsets;
+
+    GetBucket (const int* lev_gid_tid_to_bucket_ptr,
+               const int* lev_offsets_ptr,
+               const int* gid_offsets_ptr)
+        : m_lev_gid_tid_to_bucket(lev_gid_tid_to_bucket_ptr),
+          m_lev_offsets(lev_offsets_ptr),
+          m_gid_offsets(gid_offsets_ptr)
         {}
 
     AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
-    int operator() (const int lev, const int gid) const noexcept
+    int operator() (const int lev, const int gid, const int tid) const noexcept
     {
-        return m_lev_gid_to_bucket[m_lev_offsets[lev]+gid];
+        return m_lev_gid_tid_to_bucket[m_gid_offsets[m_lev_offsets[lev]+gid] + tid];
     }
 };
 
@@ -56,28 +62,40 @@ class ParticleBufferMap
     Vector<DistributionMapping> m_dm;
 
     Vector<int> m_bucket_to_gid;
+    Vector<int> m_bucket_to_tid;
     Vector<int> m_bucket_to_lev;
     Vector<int> m_bucket_to_pid;
 
-    Vector<int> m_lev_gid_to_bucket;
+    Vector<int> m_lev_gid_tid_to_bucket;
     Vector<int> m_lev_offsets;
+    Vector<int> m_gid_offsets;
 
     Vector<int> m_proc_box_counts;
     Vector<int> m_proc_box_offsets;
 
+    bool m_do_tiling{false};
+    IntVect m_tile_size{AMREX_D_DECL(1024000, 1024000, 1024000)};
+
     Gpu::DeviceVector<int> d_bucket_to_pid;
-    Gpu::DeviceVector<int> d_lev_gid_to_bucket;
+    Gpu::DeviceVector<int> d_lev_gid_tid_to_bucket;
     Gpu::DeviceVector<int> d_lev_offsets;
+    Gpu::DeviceVector<int> d_gid_offsets;
 
 public:
     ParticleBufferMap () = default;
 
     ParticleBufferMap (const ParGDBBase* a_gdb);
 
+    ParticleBufferMap (const ParGDBBase* a_gdb, bool a_do_tiling, const IntVect& a_tile_size);
+
     void define (const ParGDBBase* a_gdb);
 
+    void define (const ParGDBBase* a_gdb, bool a_do_tiling, const IntVect& a_tile_size);
+
     bool isValid (const ParGDBBase* a_gdb) const;
 
+    bool isValid (const ParGDBBase* a_gdb, bool a_do_tiling, const IntVect& a_tile_size) const;
+
     [[nodiscard]] AMREX_FORCE_INLINE
     int numLevels () const
     {
@@ -99,6 +117,13 @@ public:
         return m_bucket_to_gid[bid];
     }
 
+    [[nodiscard]] AMREX_FORCE_INLINE
+    int bucketToTile (int bid) const
+    {
+        AMREX_ASSERT(m_defined);
+        return m_bucket_to_tid[bid];
+    }
+
     [[nodiscard]] AMREX_FORCE_INLINE
     int bucketToLevel (int bid) const
     {
@@ -115,9 +140,16 @@ public:
 
     [[nodiscard]] AMREX_FORCE_INLINE
     int gridAndLevToBucket (int gid, int lev) const
+    {
+        AMREX_ASSERT(!m_do_tiling);
+        return gridAndTileAndLevToBucket(gid, 0, lev);
+    }
+
+    [[nodiscard]] AMREX_FORCE_INLINE
+    int gridAndTileAndLevToBucket (int gid, int tid, int lev) const
     {
         AMREX_ASSERT(m_defined);
-        return m_lev_gid_to_bucket[m_lev_offsets[lev] + gid];
+        return m_lev_gid_tid_to_bucket[m_gid_offsets[m_lev_offsets[lev] + gid] + tid];
     }
 
     [[nodiscard]] AMREX_FORCE_INLINE
@@ -148,14 +180,21 @@ public:
 
     [[nodiscard]] AMREX_FORCE_INLINE
     int procID (int gid, int lev) const
+    {
+        AMREX_ASSERT(!m_do_tiling);
+        return procID(gid, 0, lev);
+    }
+
+    [[nodiscard]] AMREX_FORCE_INLINE
+    int procID (int gid, int tid, int lev) const
     {
         AMREX_ASSERT(m_defined);
-        return m_dm[lev][gid];
+        return m_bucket_to_pid[gridAndTileAndLevToBucket(gid, tid, lev)];
     }
 
-    [[nodiscard]] GetPID getPIDFunctor () const noexcept { return GetPID(d_bucket_to_pid, d_lev_gid_to_bucket, d_lev_offsets);}
-    [[nodiscard]] GetBucket getBucketFunctor () const noexcept { return GetBucket(d_lev_gid_to_bucket.data(), d_lev_offsets.data());}
-    [[nodiscard]] GetBucket getHostBucketFunctor () const noexcept { return GetBucket(m_lev_gid_to_bucket.data(), m_lev_offsets.data());}
+    [[nodiscard]] GetPID getPIDFunctor () const noexcept { return GetPID(d_bucket_to_pid, d_lev_gid_tid_to_bucket, d_lev_offsets, d_gid_offsets);}
+    [[nodiscard]] GetBucket getBucketFunctor () const noexcept { return GetBucket(d_lev_gid_tid_to_bucket.data(), d_lev_offsets.data(), d_gid_offsets.data());}
+    [[nodiscard]] GetBucket getHostBucketFunctor () const noexcept { return GetBucket(m_lev_gid_tid_to_bucket.data(), m_lev_offsets.data(), m_gid_offsets.data());}
 
 };