Enable Views for LinkedCell spatial permute

core/src/Cabana_LinkedCellList.hpp

@@ -723,51 +723,26 @@ struct is_linked_cell_list
 
 //---------------------------------------------------------------------------//
 /*!
-  \brief Given a linked cell list permute an AoSoA.
+  \brief Given a linked cell list permute positions.
 
   \tparam LinkedCellListType The linked cell list type.
 
-  \tparam AoSoA_t The AoSoA type.
+  \tparam PositionType Positions type (AoSoA or Slice or Kokkos::View).
 
-  \param linked_cell_list The linked cell list to permute the AoSoA with.
+  \param linked_cell_list The linked cell list to permute the positions with.
 
-  \param aosoa The AoSoA to permute.
+  \param positions Positions to permute.
  */
-template <class LinkedCellListType, class AoSoA_t>
+template <class LinkedCellListType, class PositionType>
 void permute(
-    LinkedCellListType& linked_cell_list, AoSoA_t& aosoa,
+    LinkedCellListType& linked_cell_list, PositionType& positions,
     typename std::enable_if<( is_linked_cell_list<LinkedCellListType>::value &&
-                              is_aosoa<AoSoA_t>::value ),
+                              ( is_aosoa<PositionType>::value ||
+                                is_slice<PositionType>::value ||
+                                Kokkos::is_view<PositionType>::value ) ),
                             int>::type* = 0 )
 {
-    permute( linked_cell_list.binningData(), aosoa );
-
-    // Update internal state.
-    linked_cell_list.update( true );
-
-    linked_cell_list.storeParticleBins();
-}
-
-//---------------------------------------------------------------------------//
-/*!
-  \brief Given a linked cell list permute a slice.
-
-  \tparam LinkedCellListType The linked cell list type.
-
-  \tparam SliceType The slice type.
-
-  \param linked_cell_list The linked cell list to permute the slice with.
-
-  \param slice The slice to permute.
- */
-template <class LinkedCellListType, class SliceType>
-void permute(
-    LinkedCellListType& linked_cell_list, SliceType& slice,
-    typename std::enable_if<( is_linked_cell_list<LinkedCellListType>::value &&
-                              is_slice<SliceType>::value ),
-                            int>::type* = 0 )
-{
-    permute( linked_cell_list.binningData(), slice );
+    permute( linked_cell_list.binningData(), positions );
 
     // Update internal state.
     linked_cell_list.update( true );

core/src/Cabana_Sort.hpp

@@ -660,6 +660,64 @@ void permute(
     Kokkos::fence();
 }
 
+//---------------------------------------------------------------------------//
+/*!
+  \brief Given binning data permute a View.
+
+  \tparam BinningDataType The binning data type.
+
+  \tparam ViewType The view type.
+
+  \param binning_data The binning data.
+
+  \param view The view to permute.
+ */
+template <class BinningDataType, class ViewType,
+          class DeviceType = typename BinningDataType::device_type>
+void permute(
+    const BinningDataType& binning_data, ViewType& view,
+    typename std::enable_if<( is_binning_data<BinningDataType>::value &&
+                              Kokkos::is_view<ViewType>::value ),
+                            int>::type* = 0 )
+{
+    Kokkos::Profiling::pushRegion( "Cabana::permute" );
+
+    auto begin = binning_data.rangeBegin();
+    auto end = binning_data.rangeEnd();
+
+    // Get the number of components (1 for View).
+    std::size_t num_comp = 1;
+    for ( std::size_t d = 1; d < view.rank; ++d )
+        num_comp *= view.extent( d );
+
+    Kokkos::View<typename ViewType::value_type**, DeviceType> scratch(
+        Kokkos::ViewAllocateWithoutInitializing( "scratch" ), end - begin,
+        num_comp );
+
+    auto permute_to_scratch = KOKKOS_LAMBDA( const std::size_t i )
+    {
+        auto permute_i = binning_data.permutation( i - begin );
+        for ( std::size_t n = 0; n < num_comp; ++n )
+            scratch( i - begin, n ) = view( permute_i, n );
+    };
+    auto policy =
+        Kokkos::RangePolicy<typename DeviceType::execution_space>( begin, end );
+    Kokkos::parallel_for( "Cabana::kokkosBinSort::permute_to_scratch", policy,
+                          permute_to_scratch );
+    Kokkos::fence();
+
+    auto copy_back = KOKKOS_LAMBDA( const std::size_t i )
+    {
+        for ( std::size_t n = 0; n < num_comp; ++n )
+            view( i, n ) = scratch( i - begin, n );
+    };
+    Kokkos::parallel_for( "Cabana::kokkosBinSort::copy_back", policy,
+                          copy_back );
+    Kokkos::fence();
+
+    Kokkos::Profiling::popRegion();
+}
+
 //---------------------------------------------------------------------------//
 
 } // end namespace Cabana

core/unit_test/tstLinkedCellList.hpp

@@ -706,6 +706,67 @@ void testLinkedCellReduce()
                                    test_data.end, test_data.test_radius );
 }
 
+//---------------------------------------------------------------------------//
+void testLinkedListView()
+{
+    LCLTestData test_data;
+    auto grid_delta = test_data.grid_delta;
+    auto grid_min = test_data.grid_min;
+    auto grid_max = test_data.grid_max;
+    auto slice = Cabana::slice<LCLTestData::Position>( test_data.aosoa );
+
+    // Copy manually into a View.
+    Kokkos::View<double**, TEST_MEMSPACE> view( "positions", slice.size(), 3 );
+    copySliceToView( view, slice, 0, slice.size() );
+
+    // Bin the particles in the grid and permute only the position slice.
+    // First do this by only operating on a subset of the particles.
+    {
+        auto begin = test_data.begin;
+        auto end = test_data.end;
+        Cabana::LinkedCellList<TEST_MEMSPACE> cell_list(
+            view, begin, end, grid_delta, grid_min, grid_max );
+        Cabana::permute( cell_list, view );
+
+        // Copy manually back into the AoSoA to check values.
+        copyViewToSlice( slice, view, 0, slice.size() );
+
+        copyListToHost( test_data, cell_list );
+
+        checkBins( test_data, cell_list );
+        checkLinkedCell( test_data, begin, end, false );
+    }
+    // Now bin and permute all of the particles.
+    {
+        Cabana::LinkedCellList<TEST_MEMSPACE> cell_list( view, grid_delta,
+                                                         grid_min, grid_max );
+
+        // Copy manually into a View.
+        Kokkos::View<double**, TEST_MEMSPACE> view( "positions", slice.size(),
+                                                    3 );
+        copySliceToView( view, slice, 0, slice.size() );
+
+        Cabana::permute( cell_list, view );
+
+        copyListToHost( test_data, cell_list );
+
+        checkBins( test_data, cell_list );
+        checkLinkedCell( test_data, 0, test_data.num_p, false );
+
+        // Rebuild and make sure nothing changed.
+        cell_list.build( view );
+        Cabana::permute( cell_list, view );
+
+        // Copy manually back into the AoSoA to check values.
+        copyViewToSlice( slice, view, 0, slice.size() );
+
+        copyListToHost( test_data, cell_list );
+
+        checkBins( test_data, cell_list );
+        checkLinkedCell( test_data, 0, test_data.num_p, false );
+    }
+}
+
 //---------------------------------------------------------------------------//
 // RUN TESTS
 //---------------------------------------------------------------------------//
@@ -724,6 +785,9 @@ TEST( TEST_CATEGORY, linked_list_parallel_test ) { testLinkedCellParallel(); }
 
 TEST( TEST_CATEGORY, linked_list_reduce_test ) { testLinkedCellReduce(); }
 
+//---------------------------------------------------------------------------//
+TEST( TEST_CATEGORY, linked_list_view_test ) { testLinkedListView(); }
+
 //---------------------------------------------------------------------------//
 
 } // end namespace Test