Replace tuple-style cast with pointer/stride access

It was noted in #1 (closed) that the get operator will not fully vectorize due to the tuple-style cast hidden in getStructMember. The pointers upfront are now cached in attempts to alleviate this. This changes the manner of data access in the AoSoA for arrays:

return getStructMember<I>(_raw_data[idx.s()]);

to:

return static_cast<struct_member_array_type<I> >( data<I>() + idx.s() * _strides[I] );

If this allows vectorization, then it may be well worth the extra arithmetic needed to index into the data block. Furthermore, if loop-invariant hoisting were done well by the compiler, idx.s() * _strides[I] would be hoisted above the inner loop over arrays as a further optimization.

Cacheing pointers and strides also had a nice side effect for future C interfaces - template-free pointer and stride access! A user can now get strides and (un-typed void*) pointers directly through the following interface:

size_t PosX = 0;
double* pos_x = (double*) aosoa.pointer(PosX);
size_t pos_x_stride = aosoa.stride(PosX);

core/src/Cabana_AoSoA.hpp

@@ -112,7 +112,6 @@ class AoSoA<MemberDataTypes<Types...>,Device,ArraySize>
         , _capacity( 0 )
         , _num_soa( 0 )
         , _managed_data( nullptr )
-        , _raw_data( nullptr )
     {}
 
     // Construct a container with n elements.
@@ -121,7 +120,6 @@ class AoSoA<MemberDataTypes<Types...>,Device,ArraySize>
         , _capacity( 0 )
         , _num_soa( 0 )
         , _managed_data( nullptr )
-        , _raw_data( nullptr )
     {
         resize( _size );
     }
@@ -192,16 +190,18 @@ class AoSoA<MemberDataTypes<Types...>,Device,ArraySize>
         if ( 0 < n % array_size ) ++num_soa_alloc;
         _capacity = num_soa_alloc * array_size;
 
-        soa_type* rp;
-        memory_policy::allocate( rp, num_soa_alloc );
+        soa_type* data_block;
+        memory_policy::allocate( data_block, num_soa_alloc );
         std::shared_ptr<soa_type> sp(
-            rp, memory_policy::template deallocate<soa_type> );
+            data_block, memory_policy::template deallocate<soa_type> );
 
-        if ( _raw_data != nullptr )
-            memory_policy::copy( rp, _raw_data, _num_soa );
+        if ( _managed_data != nullptr )
+            memory_policy::copy( data_block, _managed_data.get(), _num_soa );
 
         std::swap( _managed_data, sp );
-        std::swap( _raw_data, rp );
+
+        storePointersAndStrides(
+            std::integral_constant<std::size_t,number_of_members-1>() );
     }
 
     // Get the number of structs-of-arrays in the array.
@@ -374,51 +374,89 @@ class AoSoA<MemberDataTypes<Types...>,Device,ArraySize>
     // -------------------------------
     // Raw data access.
 
-    // Get the stride between SoA data for a given member at index I.
-    template<std::size_t I>
+    // Get the stride between SoA data for a given member at index I. Note
+    // that this strides are computed in the context of the *value_type* for
+    // each member.
     CABANA_INLINE_FUNCTION
-    constexpr std::size_t stride() const
+    std::size_t stride( const std::size_t I ) const
     {
-        static_assert( 0 ==
-                       sizeof(soa_type) % sizeof(struct_member_value_type<I>),
-                       "Stride cannont be calculated for misaligned memory!" );
-        return sizeof(soa_type) / sizeof(struct_member_value_type<I>);
+        return _strides[I];
     }
 
-    // Get a pointer to the data for a given member at index I.
-    template<std::size_t I>
+    // Get an un-typed raw pointer to the data for a given member at index
+    // I. Users will need to cast this pointer to the appropriate type for the
+    // stride associated with this member to mean anything.
     CABANA_INLINE_FUNCTION
-    struct_member_pointer_type<I> pointer()
+    void* data( const std::size_t I )
     {
-        return static_cast<struct_member_pointer_type<I> >( array<I>(0) );
+        return _pointers[I];
     }
 
-    template<std::size_t I>
     CABANA_INLINE_FUNCTION
-    struct_member_const_pointer_type<I> pointer() const
+    const void* data( const std::size_t I ) const
     {
-        return static_cast<struct_member_pointer_type<I> >( array<I>(0) );
+        return _pointers[I];
     }
 
   private:
 
-    // -------------------------------
-    // Direct array data access within a struct
+    // Get a typed pointer to the data for a given member at index I.
+    template<std::size_t I>
+    CABANA_INLINE_FUNCTION
+    struct_member_pointer_type<I> typedPointer()
+    {
+        return static_cast<struct_member_pointer_type<I> >( _pointers[I] );
+    }
 
-    // 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 )
+    struct_member_const_pointer_type<I> typedPointer() const
     {
+        return static_cast<struct_member_pointer_type<I> >( _pointers[I] );
+    }
 
-        return getStructMember<I>( _raw_data[s] );
+    // Get the array at the given struct index.
+    template<std::size_t I>
+    CABANA_INLINE_FUNCTION
+    struct_member_array_type<I> array( const std::size_t s )
+    {
+        return reinterpret_cast<struct_member_array_type<I> >(
+            typedPointer<I>() + s * _strides[I] );
     }
 
     template<std::size_t I>
     CABANA_INLINE_FUNCTION
     struct_member_const_array_type<I> array( const std::size_t s ) const
     {
-        return getStructMember<I>( _raw_data[s] );
+        return reinterpret_cast<struct_member_array_type<I> >(
+            typedPointer<I>() + s * _strides[I] );
+    }
+
+    // Store the pointers and strides for each member element.
+    template<std::size_t N>
+    void assignPointersAndStrides()
+    {
+        static_assert( 0 <= N && N < number_of_members, "Loop out of bounds!" );
+        soa_type* data_block = _managed_data.get();
+        _pointers[N] =
+            static_cast<void*>( getStructMember<N>(data_block[0]) );
+        static_assert( 0 ==
+                       sizeof(soa_type) % sizeof(struct_member_value_type<N>),
+                       "Stride cannont be calculated for misaligned memory!" );
+        _strides[N] = sizeof(soa_type) / sizeof(struct_member_value_type<N>);
+    }
+
+    // Static loop through each member element to extract pointers and strides.
+    template<std::size_t N>
+    void storePointersAndStrides( std::integral_constant<std::size_t,N> )
+    {
+        assignPointersAndStrides<N>();
+        storePointersAndStrides( std::integral_constant<std::size_t,N-1>() );
+    }
+
+    void storePointersAndStrides( std::integral_constant<std::size_t,0> )
+    {
+        assignPointersAndStrides<0>();
     }
 
   private:
@@ -432,18 +470,18 @@ class AoSoA<MemberDataTypes<Types...>,Device,ArraySize>
     // Number of structs-of-arrays in the array.
     std::size_t _num_soa;
 
-    // Structs-of-Arrays managed data. This shared pointer manages the memory
-    // pointed to by _raw_data such that the copy constructor and assignment
-    // operator for this class perform a shallow and reference counted copy of
-    // the data.
+    // Structs-of-Arrays managed data. This shared pointer manages the block
+    // of memory owned by this class such that the copy constructor and
+    // assignment operator for this class perform a shallow and reference
+    // counted copy of the data.
     std::shared_ptr<soa_type> _managed_data;
 
-    // Structs-of-Arrays raw data. This data will be allocated per the
-    // MemoryPolicy of the given device type on which the class is
-    // templated. This pointer is managed by _managed_data and will be
-    // deallocated per the MemoryPolicy when the last copy of this class
-    // instance is destroyed.
-    soa_type* _raw_data;
+    // Pointers to the first element of each member.
+    void* _pointers[number_of_members];
+
+    // Strides for each member. Note that these strides are computed in the
+    // context of the *value_type* of each member.
+    std::size_t _strides[number_of_members];
 };
 
 //---------------------------------------------------------------------------//