Particle Data Structure

Creating a general AoSoA data structure with general types.

packages/Core/src/Cabana_AoSoA.hpp

@@ -2,6 +2,8 @@
 #define CABANA_AOSOA_HPP
 
 #include "Cabana_SoA.hpp"
+#include "Cabana_MemoryPolicy.hpp"
+#include "Cabana_Macros.hpp"
 
 #include <type_traits>
 #include <cmath>
@@ -10,10 +12,12 @@
 namespace Cabana
 {
 //---------------------------------------------------------------------------//
-// Type traits.
+// Forward declaration.
 template<typename Device, std::size_t ArraySize, typename... Types>
 class AoSoA;
 
+//---------------------------------------------------------------------------//
+// Static type checker.
 template<typename >
 struct is_aosoa
     : public std::false_type {};
@@ -39,6 +43,9 @@ class AoSoA
     // Device type.
     using device_type = Device;
 
+    // Memory policy.
+    using memory_policy = MemoryPolicy<device_type>;
+
     // Inner array size (size of the arrays held by the structs).
     static constexpr std::size_t soa_array_size = ArraySize;
 
@@ -91,20 +98,24 @@ class AoSoA
       public:
 
         // Constructor.
+        CABANA_FUNCTION
         Index( const std::size_t struct_id, const std::size_t offset )
             : _s( struct_id )
             , _i( offset )
         {};
 
         // Get the struct index.
-        inline std::size_t s() const
+        CABANA_INLINE_FUNCTION
+        std::size_t s() const
         { return _s; }
 
         // Get the array offset in the struct.
-        inline std::size_t i() const
+        CABANA_INLINE_FUNCTION
+        std::size_t i() const
         { return _i; }
 
         // Prefix increment operator.
+        CABANA_FUNCTION
         Index& operator++()
         {
             _i = ( soa_array_size - 1 == _i ) ? 0 : _i + 1;
@@ -113,6 +124,7 @@ class AoSoA
         };
 
         // Postfix increment operator.
+        CABANA_FUNCTION
         Index operator++(int)
         {
             Index temp = *this;
@@ -121,12 +133,14 @@ class AoSoA
         };
 
         // Equality comparator.
+        CABANA_FUNCTION
         bool operator==( const Index& rhs ) const
         {
             return (_s == rhs._s) && (_i == rhs._i);
         }
 
         // Inequality comparator.
+        CABANA_FUNCTION
         bool operator!=( const Index& rhs ) const
         {
             return (_s != rhs._s) || (_i != rhs._i);
@@ -141,59 +155,75 @@ class AoSoA
     {
         _num_soa = std::floor(_size / soa_array_size);
         if ( 0 < _size % soa_array_size ) ++_num_soa;
-        Device::allocate( _data, _num_soa );
+        memory_policy::allocate( _data, _num_soa );
     }
 
     // Destructor.
     ~AoSoA()
     {
-        Device::deallocate( _data );
+        memory_policy::deallocate( _data );
     }
 
     // Get the requested size of the AoSoA (global number of requested
-    // elements). This may or not be equal to the total number of allocated
-    // elements in the AoSoA if the size of the struct arrays is not evenly
-    // divisible by the requested size.
+    // elements). This may or may not be equal to the total number of
+    // allocated elements in the AoSoA if the size of the struct arrays is not
+    // evenly divisible by the requested size.
+    CABANA_FUNCTION
     std::size_t size() const { return _size; }
 
     // Get the allocated size of the AoSoA (global number of allocated
-    // elements). This may or not be equal to the total number of requested
-    // elements in the AoSoA if the size of the struct arrays is not evenly
-    // divisible by the requested size.
+    // elements). This may or may not be equal to the total number of
+    // requested elements in the AoSoA if the size of the struct arrays is not
+    // evenly divisible by the requested size.
+    CABANA_FUNCTION
     std::size_t allocatedSize() const { return _num_soa * soa_array_size; }
 
     // Get the number of structs-of-arrays in the array.
+    CABANA_FUNCTION
     std::size_t numSoA() const { return _num_soa; }
 
-    // Access the data array at a given struct member index.
-    template<std::size_t I>
-    inline struct_member_array_type<I> array( const std::size_t s )
-    {
-        return getStructMember<I>( _data[s] );
-    }
 
-    template<std::size_t I>
-    inline struct_member_const_array_type<I> array( const std::size_t s ) const
-    {
-        return getStructMember<I>( _data[s] );
-    }
+    // -----------------------------
+    // Array range
 
     // Get the index at the beginning of the entire AoSoA.
+    CABANA_FUNCTION
     Index begin() const
     {
         return Index( 0, 0 );
     }
 
     // Get the index at end of the entire AoSoA.
+    CABANA_FUNCTION
     Index end() const
     {
         return Index( _num_soa - 1, _size % soa_array_size );
     }
 
+    // -------------------------------
+    // Raw array data access
+
+    // Access the data array at a given struct member index.
+    template<std::size_t I>
+    CABANA_INLINE_FUNCTION
+    struct_member_array_type<I> array( const std::size_t s )
+    {
+        return getStructMember<I>( _data[s] );
+    }
+
+    template<std::size_t I>
+    CABANA_INLINE_FUNCTION
+    struct_member_const_array_type<I> array( const std::size_t s ) const
+    {
+        return getStructMember<I>( _data[s] );
+    }
+
+    // -------------------------------
     // Access the data value at a given struct member index and array index
-    // for Rank 0 data.
+
+    // Rank 0
     template<std::size_t I>
-    inline
+    CABANA_INLINE_FUNCTION
     typename std::enable_if<(0==std::rank<struct_member_data_type<I> >::value),
                             struct_member_reference_type<I> >::type
     get( const Index& idx )
@@ -202,7 +232,7 @@ class AoSoA
     }
 
     template<std::size_t I>
-    inline
+    CABANA_INLINE_FUNCTION
     typename std::enable_if<(0==std::rank<struct_member_data_type<I> >::value),
                             struct_member_const_reference_type<I> >::type
     get( const Index& idx ) const
@@ -210,10 +240,9 @@ class AoSoA
         return array<I>(idx.s())[idx.i()];
     }
 
-    // Access the data value at a given struct member index and array index
-    // for Rank 1 data.
+    // Rank 1
     template<std::size_t I>
-    inline
+    CABANA_INLINE_FUNCTION
     typename std::enable_if<(1==std::rank<struct_member_data_type<I> >::value),
                             struct_member_reference_type<I> >::type
     get( const Index& idx,
@@ -223,7 +252,7 @@ class AoSoA
     }
 
     template<std::size_t I>
-    inline
+    CABANA_INLINE_FUNCTION
     typename std::enable_if<(1==std::rank<struct_member_data_type<I> >::value),
                             struct_member_const_reference_type<I> >::type
     get( const Index& idx,
@@ -232,10 +261,9 @@ class AoSoA
         return array<I>(idx.s())[idx.i()][d0];
     }
 
-    // Access the data value at a given struct member index and array index
-    // for Rank 2 data.
+    // Rank 2
     template<std::size_t I>
-    inline
+    CABANA_INLINE_FUNCTION
     typename std::enable_if<(2==std::rank<struct_member_data_type<I> >::value),
                             struct_member_reference_type<I> >::type
     get( const Index& idx,
@@ -246,7 +274,7 @@ class AoSoA
     }
 
     template<std::size_t I>
-    inline
+    CABANA_INLINE_FUNCTION
     typename std::enable_if<(2==std::rank<struct_member_data_type<I> >::value),
                             struct_member_const_reference_type<I> >::type
     get( const Index& idx,
@@ -256,10 +284,9 @@ class AoSoA
         return array<I>(idx.s())[idx.i()][d0][d1];
     }
 
-    // Access the data value at a given struct member index and array index
-    // for Rank 3 data.
+    // Rank 3
     template<std::size_t I>
-    inline
+    CABANA_INLINE_FUNCTION
     typename std::enable_if<(3==std::rank<struct_member_data_type<I> >::value),
                             struct_member_reference_type<I> >::type
     get( const Index& idx,
@@ -271,7 +298,7 @@ class AoSoA
     }
 
     template<std::size_t I>
-    inline
+    CABANA_INLINE_FUNCTION
     typename std::enable_if<(3==std::rank<struct_member_data_type<I> >::value),
                             struct_member_const_reference_type<I> >::type
     get( const Index& idx,