Enable binning for ParticleTileData (#3307)

## Summary

To use binning with ParticleTileData, replace
```
using PlasmaBins = amrex::DenseBins<PlasmaParticleContainer::ParticleType>;
```
with 
```
using PlasmaBins = amrex::DenseBins<PlasmaParticleContainer::ParticleTileType::ParticleTileDataType>;
```
and
```
bins.build(pti.numParticles(), pti.GetArrayOfStructs().begin(), …
```
with 
```
bins.build(pti.numParticles(), pti.GetParticleTile().getParticleTileData(), …
```

Using ParticleTileData is necessary for PureSoA.

## Additional background

The previous call_f
```
    template <typename F, typename T, typename I>
    AMREX_GPU_HOST_DEVICE
    auto call_f(F const& f, T* pstruct_ptr, I& index)
        noexcept -> decltype(f(pstruct_ptr,index))
    {
        return f(pstruct_ptr[index],index);
    }
```
had an issue where `pstruct_ptr` and `pstruct_ptr[index]` was mixed up.
I fixed this here which may cause some backwards compatibility issues.
Otherwise binning should work as before when using a `pstruct_ptr`.


## Checklist

The proposed changes:
- [ ] fix a bug or incorrect behavior in AMReX
- [x] add new capabilities to AMReX
- [ ] changes answers in the test suite to more than roundoff level
- [ ] are likely to significantly affect the results of downstream AMReX
users
- [ ] include documentation in the code and/or rst files, if appropriate

---------

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

Src/Particle/AMReX_BinIterator.H

@@ -9,15 +9,36 @@
 namespace amrex
 {
 
+template< class T >
+constexpr decltype(T::is_particle_tile_data) IsParticleTileData () { // NOLINT(readability-const-return-type)
+    return T::is_particle_tile_data;
+}
+
+template< class T, class...Args >
+constexpr bool IsParticleTileData (Args...) {
+    return false;
+}
+
 template <typename T>
 struct BinIterator
 {
     using index_type = unsigned int;
 
+    using const_pointer_type = typename std::conditional<IsParticleTileData<T>(),
+        T,
+        const T*
+    >::type;
+
+    using const_pointer_input_type = typename std::conditional<IsParticleTileData<T>(),
+        const T&,
+        const T*
+    >::type;
+
     struct iterator
     {
         AMREX_GPU_HOST_DEVICE
-        iterator (index_type start, index_type stop, const index_type* a_perm, const T* a_items)
+        iterator (index_type start, index_type stop,
+                  const index_type* a_perm, const_pointer_input_type a_items)
             : m_items(a_items), m_perm(a_perm), m_index(start), m_stop(stop)
         {}
 
@@ -28,13 +49,13 @@ struct BinIterator
         bool operator!= (iterator const& /*rhs*/) const { return m_index < m_stop; }
 
         [[nodiscard]] AMREX_GPU_HOST_DEVICE
-        std::pair<index_type, T> operator* () const
+        auto operator* () const
         {
             return std::make_pair(m_perm[m_index], m_items[m_perm[m_index]]);
         }
 
     private:
-        const T* m_items;
+        const_pointer_type m_items;
         const index_type* m_perm;
         index_type m_index;
         index_type m_stop;
@@ -58,15 +79,16 @@ struct BinIterator
     }
 
     AMREX_GPU_HOST_DEVICE
-    BinIterator (index_type i, const index_type *offsets_ptr, const index_type *permutation_ptr, const T* items)
+    BinIterator (index_type i, const index_type *offsets_ptr,
+                 const index_type *permutation_ptr, const_pointer_input_type items)
         : m_i(i), m_offsets_ptr(offsets_ptr), m_permutation_ptr(permutation_ptr), m_items(items)
     {}
 
 private:
     const index_type m_i;
     const index_type * m_offsets_ptr;
     const index_type * m_permutation_ptr;
-    const T * m_items;
+    const_pointer_type m_items;
 
     static constexpr index_type m_not_found = std::numeric_limits<index_type>::max();
 };

Src/Particle/AMReX_DenseBins.H

@@ -31,6 +31,11 @@ struct DenseBinIteratorFactory
 {
     using index_type = unsigned int;
 
+    using const_pointer_type = typename std::conditional<IsParticleTileData<T>(),
+        T,
+        const T*
+    >::type;
+
     DenseBinIteratorFactory (const Gpu::DeviceVector<index_type>& offsets,
                              const Gpu::DeviceVector<index_type>& permutation,
                              const T* items)
@@ -47,27 +52,9 @@ struct DenseBinIteratorFactory
 
     const index_type* m_offsets_ptr;
     const index_type* m_permutation_ptr;
-    const T* m_items;
+    const_pointer_type m_items;
 };
 
-namespace
-{
-    template <typename F, typename T, typename I>
-    AMREX_GPU_HOST_DEVICE
-    auto call_f(F const& f, T* pstruct_ptr, I& index)
-        noexcept -> decltype(f(pstruct_ptr,index))
-    {
-        return f(pstruct_ptr[index],index);
-    }
-
-    template <typename F, typename T, typename I>
-    AMREX_GPU_HOST_DEVICE
-    auto call_f(F const& f, T* pstruct_ptr, I& index)
-        noexcept -> decltype(f(pstruct_ptr[index]))
-    {
-        return f(pstruct_ptr[index]);
-    }
-}
 
 /**
  * \brief A container for storing items in a set of bins.
@@ -92,6 +79,36 @@ public:
     using BinIteratorFactory = DenseBinIteratorFactory<T>;
     using index_type = unsigned int;
 
+    using const_pointer_type = typename std::conditional<IsParticleTileData<T>(),
+        T,
+        const T*
+    >::type;
+
+    using const_pointer_input_type = typename std::conditional<IsParticleTileData<T>(),
+        const T&,
+        const T*
+    >::type;
+
+private:
+
+    template <typename F, typename I>
+    AMREX_GPU_HOST_DEVICE
+    static auto call_f(F const& f, const_pointer_input_type v, I& index)
+        noexcept -> decltype(f(v,index))
+    {
+        return f(v,index);
+    }
+
+    template <typename F, typename I>
+    AMREX_GPU_HOST_DEVICE
+    static auto call_f(F const& f, const_pointer_input_type v, I& index)
+        noexcept -> decltype(f(v[index]))
+    {
+        return f(v[index]);
+    }
+
+public:
+
     /**
      * \brief Populate the bins with a set of items.
      *
@@ -110,12 +127,12 @@ public:
      * \tparam F a function that maps items to IntVect bins
      *
      * \param nitems the number of items to put in the bins
-     * \param v pointer to the start of the items
+     * \param v ParticleTileData or pointer to the start of the items
      * \param bx the Box that defines the space over which the bins will be defined
      * \param f a function object that maps items to bins
      */
     template <typename N, typename F>
-    void build (N nitems, T const* v, const Box& bx, F&& f)
+    void build (N nitems, const_pointer_input_type v, const Box& bx, F&& f)
     {
         build(BinPolicy::Default, nitems, v, bx, std::forward<F>(f));
     }
@@ -137,12 +154,12 @@ public:
      * \tparam F a function that maps items to IntVect bins
      *
      * \param nitems the number of items to put in the bins
-     * \param v pointer to the start of the items
+     * \param v ParticleTileData or pointer to the start of the items
      * \param nbins the number of bins to use
      * \param f a function object that maps items to bins
      */
     template <typename N, typename F>
-    void build (N nitems, T const* v, int nbins, F&& f)
+    void build (N nitems, const_pointer_input_type v, int nbins, F&& f)
     {
         build(BinPolicy::Default, nitems, v, nbins, std::forward<F>(f));
     }
@@ -165,19 +182,19 @@ public:
      * \tparam F a function that maps items to IntVect bins
      *
      * \param nitems the number of items to put in the bins
-     * \param v pointer to the start of the items
+     * \param v ParticleTileData or pointer to the start of the items
      * \param bx the Box that defines the space over which the bins will be defined
      * \param f a function object that maps items to bins
      */
     template <typename N, typename F>
-    void build (BinPolicy::GPUBinPolicy, N nitems, T const* v, const Box& bx, F&& f)
+    void build (BinPolicy::GPUBinPolicy, N nitems, const_pointer_input_type v, const Box& bx, F&& f)
     {
         const auto lo = lbound(bx);
         const auto hi = ubound(bx);
         build(BinPolicy::GPU, nitems, v, bx.numPts(),
-              [=] AMREX_GPU_DEVICE (const T& t) noexcept
+              [=] AMREX_GPU_DEVICE (const_pointer_type t, index_type i) noexcept
               {
-                  auto iv = f(t);
+                  auto iv = call_f(f,t,i);
                   auto iv3 = iv.dim3();
                   int nx = hi.x-lo.x+1;
                   int ny = hi.y-lo.y+1;
@@ -206,12 +223,12 @@ public:
      * \tparam F a function that maps items to IntVect bins
      *
      * \param nitems the number of items to put in the bins
-     * \param v pointer to the start of the items
+     * \param v ParticleTileData or pointer to the start of the items
      * \param nbins the number of bins to use
      * \param f a function object that maps items to bins
      */
     template <typename N, typename F>
-    void build (BinPolicy::GPUBinPolicy, N nitems, T const* v, int nbins, F&& f)
+    void build (BinPolicy::GPUBinPolicy, N nitems, const_pointer_input_type v, int nbins, F&& f)
     {
         BL_PROFILE("DenseBins<T>::buildGPU");
 
@@ -269,19 +286,19 @@ public:
      * \tparam F a function that maps items to IntVect bins
      *
      * \param nitems the number of items to put in the bins
-     * \param v pointer to the start of the items
+     * \param v ParticleTileData or pointer to the start of the items
      * \param bx the Box that defines the space over which the bins will be defined
      * \param f a function object that maps items to bins
      */
     template <typename N, typename F>
-    void build (BinPolicy::OpenMPBinPolicy, N nitems, T const* v, const Box& bx, F&& f)
+    void build (BinPolicy::OpenMPBinPolicy, N nitems, const_pointer_input_type v, const Box& bx, F&& f)
     {
         const auto lo = lbound(bx);
         const auto hi = ubound(bx);
         build(BinPolicy::OpenMP, nitems, v, bx.numPts(),
-              [=] (const T& t) noexcept
+              [=] (const_pointer_type t, index_type i) noexcept
               {
-                  auto iv = f(t);
+                  auto iv = call_f(f,t,i);
                   auto iv3 = iv.dim3();
                   int nx = hi.x-lo.x+1;
                   int ny = hi.y-lo.y+1;
@@ -311,12 +328,12 @@ public:
      * \tparam F a function that maps items to IntVect bins
      *
      * \param nitems the number of items to put in the bins
-     * \param v pointer to the start of the items
+     * \param v ParticleTileData or pointer to the start of the items
      * \param nbins the number of bins to use
      * \param f a function object that maps items to bins
      */
     template <typename N, typename F>
-    void build (BinPolicy::OpenMPBinPolicy, N nitems, T const* v, int nbins, F&& f)
+    void build (BinPolicy::OpenMPBinPolicy, N nitems, const_pointer_input_type v, int nbins, F&& f)
     {
         BL_PROFILE("DenseBins<T>::buildOpenMP");
 
@@ -408,19 +425,19 @@ public:
      * \tparam F a function that maps items to IntVect bins
      *
      * \param nitems the number of items to put in the bins
-     * \param v pointer to the start of the items
+     * \param v ParticleTileData or pointer to the start of the items
      * \param bx the Box that defines the space over which the bins will be defined
      * \param f a function object that maps items to bins
      */
     template <typename N, typename F>
-    void build (BinPolicy::SerialBinPolicy, N nitems, T const* v, const Box& bx, F&& f)
+    void build (BinPolicy::SerialBinPolicy, N nitems, const_pointer_input_type v, const Box& bx, F&& f)
     {
         const auto lo = lbound(bx);
         const auto hi = ubound(bx);
         build(BinPolicy::Serial, nitems, v, bx.numPts(),
-              [=] (const T& t) noexcept
+              [=] (const_pointer_type t, index_type i) noexcept
               {
-                  auto iv = f(t);
+                  auto iv = call_f(f,t,i);
                   auto iv3 = iv.dim3();
                   int nx = hi.x-lo.x+1;
                   int ny = hi.y-lo.y+1;
@@ -450,12 +467,12 @@ public:
      * \tparam F a function that maps items to IntVect bins
      *
      * \param nitems the number of items to put in the bins
-     * \param v pointer to the start of the items
+     * \param v ParticleTileData or pointer to the start of the items
      * \param nbins the number of bins to use
      * \param f a function object that maps items to bins
      */
     template <typename N, typename F>
-    void build (BinPolicy::SerialBinPolicy, N nitems, T const* v, int nbins, F&& f)
+    void build (BinPolicy::SerialBinPolicy, N nitems, const_pointer_input_type v, int nbins, F&& f)
     {
         BL_PROFILE("DenseBins<T>::buildSerial");
 
@@ -519,7 +536,7 @@ public:
 
 private:
 
-    const T* m_items;
+    const_pointer_type m_items;
 
     Gpu::DeviceVector<index_type> m_bins;
     Gpu::DeviceVector<index_type> m_counts;

Src/Particle/AMReX_ParticleContainer.H

@@ -1366,7 +1366,7 @@ protected:
 
     void SetParticleSize ();
 
-    DenseBins<ParticleType> m_bins;
+    DenseBins<typename ParticleTileType::ParticleTileDataType> m_bins;
 
 private:
     virtual void particlePostLocate (ParticleType& /*p*/, const ParticleLocData& /*pld*/,

Src/Particle/AMReX_ParticleContainerI.H

@@ -1176,15 +1176,14 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
         for(MFIter mfi = MakeMFIter(lev); mfi.isValid(); ++mfi)
         {
             auto& ptile           = ParticlesAt(lev, mfi);
-            auto& aos             = ptile.GetArrayOfStructs();
-            auto *pstruct_ptr     = aos().dataPtr();
             const size_t np       = ptile.numParticles();
 
             const Box& box = mfi.validbox();
 
             int ntiles = numTilesInBox(box, true, bin_size);
 
-            m_bins.build(np, pstruct_ptr, ntiles, GetParticleBin{plo, dxi, domain, bin_size, box});
+            m_bins.build(np, ptile.getParticleTileData(), ntiles,
+                         GetParticleBin{plo, dxi, domain, bin_size, box});
             ReorderParticles(lev, mfi, m_bins.permutationPtr());
         }
     }

Src/Particle/AMReX_ParticleTile.H

@@ -37,6 +37,8 @@ struct ParticleTileData
 
     using SuperParticleType = Particle<NStructReal+NAR, NStructInt+NAI>;
 
+    static constexpr bool is_particle_tile_data = true;
+
     Long m_size;
 
     ParticleType* AMREX_RESTRICT m_aos;
@@ -120,6 +122,16 @@ struct ParticleTileData
         return this->m_idata[attribute_index];
     }
 
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    decltype(auto) operator[] (const int index) const
+    {
+        if constexpr (!ParticleType::is_soa_particle) {
+            return m_aos[index];
+        } else {
+            return SoAParticle<NAR, NAI>(*this, index);
+        }
+    }
+
     AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     void packParticleData (char* buffer, int src_index, std::size_t dst_offset,
                            const int* comm_real, const int * comm_int) const noexcept
@@ -463,6 +475,8 @@ struct ConstParticleTileData
 
     using SuperParticleType = Particle<NStructReal+NArrayReal, NStructInt+NArrayInt>;
 
+    static constexpr bool is_particle_tile_data = true;
+
     Long m_size;
     const ParticleType* AMREX_RESTRICT m_aos;
 
@@ -491,6 +505,16 @@ struct ConstParticleTileData
         return this->m_idata[attribute_index];
     }
 
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    decltype(auto) operator[] (const int index) const
+    {
+        if constexpr (!ParticleType::is_soa_particle) {
+            return m_aos[index];
+        } else {
+            return ConstSoAParticle<NAR, NAI>(*this, index);
+        }
+    }
+
     int m_num_runtime_real;
     int m_num_runtime_int;
     const ParticleReal* AMREX_RESTRICT * AMREX_RESTRICT m_runtime_rdata;

Src/Particle/AMReX_ParticleUtil.H

@@ -518,8 +518,9 @@ struct BinMapper
 
     template <typename T>
     AMREX_GPU_HOST_DEVICE
-    unsigned int operator() (const T& p, int i) const
+    unsigned int operator() (const T& ptd, int i) const
     {
+        auto p = ptd[i];
         int type   = (m_bin_type_array) ? m_bin_type_array[i] : 0;
         int offset = m_off_bins_p[type];