Parallel for abstraction

CMakeLists.txt

@@ -16,13 +16,17 @@ include(CTest)
 option(IS_TESTING "Build for CTest" OFF)
 message(STATUS "IS_TESTING: ${IS_TESTING}")
 
-add_executable(SliceWrapper test/SliceWrapper.cpp)
+add_executable(SliceWrapper test/testSliceWrapper.cpp)
 target_include_directories(SliceWrapper PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/src) #hack - this should be a header-only library
 target_link_libraries(SliceWrapper Omega_h::omega_h)
 target_link_libraries(SliceWrapper Cabana::cabanacore)
 target_compile_definitions(SliceWrapper PUBLIC ENABLE_CABANA)
 
-add_test(sliceWrapper10 ./SliceWrapper 10)
-add_test(sliceWrapper32 ./SliceWrapper 32)
-add_test(sliceWrapper33 ./SliceWrapper 33)
+add_executable(MeshField test/testMeshField.cpp)
+target_include_directories(MeshField PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/src) #hack - this should be a header-only library
+target_link_libraries(MeshField Omega_h::omega_h)
+target_link_libraries(MeshField Cabana::cabanacore)
+target_compile_definitions(MeshField PUBLIC ENABLE_CABANA)
+
 add_test(sliceWrapper50 ./SliceWrapper 50)
+add_test(meshField ./MeshField)

src/MeshField.hpp

@@ -0,0 +1,61 @@
+#ifndef meshfield_hpp
+#define meshfield_hpp
+
+#include "SliceWrapper.hpp"
+
+namespace MeshField {
+
+template <class Slice>
+class Field {
+
+  Slice slice;
+
+public:
+  Field(Slice s) : slice(s) {}
+
+  KOKKOS_INLINE_FUNCTION
+  auto& operator()(int s, int a) const {
+    return slice.access(s,a);
+  }
+
+  KOKKOS_INLINE_FUNCTION
+  auto& operator()(int s, int a, int i) const {
+    return slice.access(s,a,i);
+  }
+
+  KOKKOS_INLINE_FUNCTION
+  auto& operator()(int s, int a, int i, int j) const {
+    return slice.access(s,a,i,j);
+  }
+
+  KOKKOS_INLINE_FUNCTION
+  auto& operator()(int s, int a, int i, int j, int k) const {
+    return slice.access(s,a,i,j,k);
+  }
+};
+
+template <class Controller>
+class MeshField {
+
+  Controller sliceController;  
+
+public:
+  MeshField(Controller sc) : sliceController(std::move(sc)) {}
+
+  template <std::size_t index>
+  auto makeField() {
+    auto slice = sliceController.template makeSlice<index>();
+    return Field(std::move(slice));
+  }
+
+  template<typename FunctorType>
+  void parallel_for(int lower_bound, int upper_bound,
+		    FunctorType vector_kernel,
+		    std::string tag) {
+    sliceController.parallel_for(lower_bound, upper_bound, vector_kernel, tag);
+  }
+
+};
+
+}
+#endif

src/SliceWrapper.hpp

@@ -9,6 +9,8 @@ struct SliceWrapper {
   SliceType st_; //store the underlying instance
 
   SliceWrapper(SliceType st) : st_(st)  {}
+
+  SliceWrapper() {}
 
   KOKKOS_INLINE_FUNCTION
   T& access(int s, int a) const {
@@ -25,14 +27,13 @@ struct SliceWrapper {
   KOKKOS_INLINE_FUNCTION
   auto& access(int s, int a, int i, int j, int k) const {
     return st_.access(s,a,i,j,k);
-  }
-
+  }  
 };
 
 using namespace Cabana;
 
 template <class ExecutionSpace, class MemorySpace, class... Ts>
-class CabSliceFactory {
+class CabSliceController {
   using TypeTuple = std::tuple<Ts...>;
   using DeviceType = Kokkos::Device<ExecutionSpace, MemorySpace>;
   using DataTypes = Cabana::MemberTypes<Ts...>;
@@ -60,15 +61,23 @@ class CabSliceFactory {
   Cabana::AoSoA<DataTypes, DeviceType, vecLen> aosoa; 
 
 public:
+  template<typename FunctorType>
+  void parallel_for(int lower_bound, int upper_bound, FunctorType vectorKernel, std::string tag) {
+    Cabana::SimdPolicy<vecLen, ExecutionSpace> simd_policy(lower_bound, upper_bound);
+    Cabana::simd_parallel_for(simd_policy, vectorKernel, tag);
+  }
+
   template <std::size_t index>
-  auto makeSliceCab() {
+  auto makeSlice() {
     using type = std::tuple_element_t<index, TypeTuple>;
     const int stride = sizeof(soa_t) / sizeof(member_value_t<index>);
     auto slice = Cabana::slice<index>(aosoa);
     return wrapper_slice_t< type, stride >(std::move(slice));
   }
+
+  CabSliceController() {}
 
-  CabSliceFactory(int n) : aosoa("sliceAoSoA", n) {
+  CabSliceController(int n) : aosoa("sliceAoSoA", n) {
     if (sizeof...(Ts) == 0) {
       throw std::invalid_argument("Must provide at least one member type in template definition");
     }

test/testMeshField.cpp

@@ -0,0 +1,262 @@
+#include "MeshField.hpp"
+#include "SliceWrapper.hpp"
+
+#include <Cabana_Core.hpp>
+
+using ExecutionSpace = Kokkos::DefaultExecutionSpace;
+using MemorySpace = ExecutionSpace::memory_space;
+
+
+void single_type(int num_tuples) {
+  using Controller = CabSliceController<ExecutionSpace, MemorySpace, double>;
+
+  // Slice Wrapper Controller
+  Controller c(num_tuples);
+  MeshField::MeshField<Controller> cabMeshField(c);
+
+  auto field0 = cabMeshField.makeField<0>();
+
+  auto vector_kernel = KOKKOS_LAMBDA(const int s, const int a)
+  {
+   double d0 = 10;
+   field0(s,a) = d0;
+   assert(field0(s,a) == d0);
+  };
+
+  cabMeshField.parallel_for(0,num_tuples,vector_kernel,"single_type_pfor");
+
+}
+
+void multi_type(int num_tuples) {
+  using Controller = CabSliceController<ExecutionSpace, MemorySpace, double,
+					double, float, int, char>;
+
+  // Slice Wrapper Controller
+  Controller c(num_tuples);
+  MeshField::MeshField<Controller> cabMeshField(c);
+
+  auto field0 = cabMeshField.makeField<0>();
+  auto field1 = cabMeshField.makeField<1>();
+  auto field2 = cabMeshField.makeField<2>();
+  auto field3 = cabMeshField.makeField<3>();
+  auto field4 = cabMeshField.makeField<4>();
+
+  auto vector_kernel = KOKKOS_LAMBDA(const int s, const int a)
+  {
+   double d0 = 10.456;
+   field0(s,a) = d0;
+   double d1 = 43.973234567;
+   field1(s,a) = d1;
+   float f0 = 123.45;
+   field2(s,a) = f0;
+   int i0 = 22;
+   field3(s,a) = i0;
+   char c0 = 'a';
+   field4(s,a) = c0;
+
+   assert(field0(s,a) == d0);
+   assert(field1(s,a) == d1);
+   assert(field2(s,a) == f0);
+   assert(field3(s,a) == i0);
+   assert(field4(s,a) == c0);
+  };
+
+  cabMeshField.parallel_for(0,num_tuples,vector_kernel,"multi_type_pfor");
+}
+
+void many_type(int num_tuples) {
+
+  using Controller = CabSliceController<ExecutionSpace, MemorySpace,
+					double, double, float, float, int,
+					short int, char>;
+
+  // Slice Wrapper Controller
+  Controller c(num_tuples);
+  MeshField::MeshField<Controller> cabMeshField(c);
+
+  auto field0 = cabMeshField.makeField<0>();
+  auto field1 = cabMeshField.makeField<1>();
+  auto field2 = cabMeshField.makeField<2>();
+  auto field3 = cabMeshField.makeField<3>();
+  auto field4 = cabMeshField.makeField<4>();
+  auto field5 = cabMeshField.makeField<5>();
+  auto field6 = cabMeshField.makeField<6>();
+
+  auto vector_kernel = KOKKOS_LAMBDA(const int s, const int a)
+  {
+   double d0 = 10.456;
+   field0(s,a) = d0;
+   double d1 = 43.973234567;
+   field1(s,a) = d1;
+   float f0 = 123.45;
+   field2(s,a) = f0;
+   float f1 = 543.21;
+   field3(s,a) = f1;
+   int i0 = 222;
+   field4(s,a) = i0;
+   short int i1 = 50;
+   field5(s,a) = i1;
+   char c0 = 'h';
+   field6(s,a) = c0;
+
+   assert(field0(s,a) == d0);
+   assert(field1(s,a) == d1);
+   assert(field2(s,a) == f0);
+   assert(field3(s,a) == f1);
+   assert(field4(s,a) == i0);
+   assert(field5(s,a) == i1);
+   assert(field6(s,a) == c0);
+  };
+
+  cabMeshField.parallel_for(0,num_tuples,vector_kernel,"many_type_pfor");
+}
+
+void rank1_arr(int num_tuples) {
+  const int width = 3;
+  using Controller = CabSliceController<ExecutionSpace, MemorySpace, double[width]>;
+
+  // Slice Wrapper Controller
+  Controller c(num_tuples);
+  MeshField::MeshField<Controller> cabMeshField(c);
+
+  auto field0 = cabMeshField.makeField<0>();
+
+
+  auto vector_kernel = KOKKOS_LAMBDA(const int s, const int a)
+  {
+   for (int i = 0; i < width; i++)
+   {
+    double d0 = 10+i;
+    field0(s,a,i) = d0;
+    assert(field0(s,a,i) == d0);
+   }
+  };
+
+  cabMeshField.parallel_for(0,num_tuples,vector_kernel,"rank1_arr_pfor");
+}
+
+void rank2_arr(int num_tuples) {
+  const int width = 3;
+  const int height = 4;
+  using Controller = CabSliceController<ExecutionSpace, MemorySpace,
+					double[width][height]>;
+
+  // Slice Wrapper Controller
+  Controller c(num_tuples);
+  MeshField::MeshField<Controller> cabMeshField(c);
+
+  auto field0 = cabMeshField.makeField<0>();
+
+  auto vector_kernel = KOKKOS_LAMBDA(const int s, const int a)
+  {
+   for (int i = 0; i < width; i++)
+   {
+    for (int j = 0; j < height; j++)
+    {
+     double d0 = (10+i)/(j+1);
+     field0(s,a,i,j) = d0;
+     assert(field0(s,a,i,j) == d0);
+    }
+   }
+  };
+
+  cabMeshField.parallel_for(0,num_tuples,vector_kernel,"rank2_arr_pfor");
+}
+
+void rank3_arr(int num_tuples) {
+  const int width = 3;
+  const int height = 4;
+  const int depth = 2;
+  using Controller = CabSliceController<ExecutionSpace, MemorySpace,
+					double[width][height][depth]>;
+
+  // Slice Wrapper Controller
+  Controller c(num_tuples);
+  MeshField::MeshField<Controller> cabMeshField(c);
+
+  auto field0 = cabMeshField.makeField<0>();
+
+  auto vector_kernel = KOKKOS_LAMBDA(const int s, const int a)
+  {
+   for (int i = 0; i < width; i++)
+   {
+    for (int j = 0; j < height; j++)
+    {
+     for (int k = 0; k < depth; k++)
+     {
+      double d0 = ((10+i)*(k+1))/(j+1);
+      field0(s,a,i,j,k) = d0;
+      assert(field0(s,a,i,j,k) == d0);
+     }
+    }
+   }
+  };
+
+  cabMeshField.parallel_for(0,num_tuples,vector_kernel,"rank3_arr_pfor");
+
+}
+
+void mix_arr(int num_tuples) {
+  const int width = 3;
+  const int height = 4;
+  const int depth = 2;
+  using Controller = CabSliceController<ExecutionSpace, MemorySpace,
+					double[width][height][depth],
+				        float[width][height], int[width],
+					char>;
+
+  // Slice Wrapper Controller
+  Controller c(num_tuples);
+  MeshField::MeshField<Controller> cabMeshField(c);
+
+  auto field0 = cabMeshField.makeField<0>();
+  auto field1 = cabMeshField.makeField<1>();
+  auto field2 = cabMeshField.makeField<2>();
+  auto field3 = cabMeshField.makeField<3>();
+
+  auto vector_kernel = KOKKOS_LAMBDA(const int s, const int a)
+  {
+   float f0;
+   int i0;
+   char c0 = 's';
+   field3(s,a) = c0;
+
+   for (int i = 0; i < width; i++)
+   {
+    i0 = i+s+a;
+    field2(s,a,i) = i0;
+    for (int j = 0; j < height; j++)
+    {
+     f0 = i0 / (i+j+1.123);
+     field1(s,a,i,j) = f0;
+     for (int k = 0; k < depth; k++)
+     {
+      double d0 = ((10+i)*(k+1))/(j+1);
+      field0(s,a,i,j,k) = d0;
+      assert(field0(s,a,i,j,k) == d0);
+     }
+     assert(field1(s,a,i,j) == f0);
+    }
+    assert(field2(s,a,i) == i0);
+   }
+   assert(field3(s,a) == c0);
+  };
+
+  cabMeshField.parallel_for(0,num_tuples,vector_kernel,"mix_arr_pfor");
+
+}
+
+int main(int argc, char* argv[]) {
+  int num_tuples = (argc < 2) ? (1000) : (atoi(argv[1]));
+  Kokkos::ScopeGuard scope_guard(argc, argv);
+
+  single_type(num_tuples);
+  multi_type(num_tuples);
+  many_type(num_tuples);
+  rank1_arr(num_tuples);
+  rank2_arr(num_tuples);
+  rank3_arr(num_tuples);
+  mix_arr(num_tuples);
+
+  return 0;
+}