Merge pull request #3 from SCOREC/multi-slice-cab-factory

CabanaSliceFactory now works with multiple types.

Author: HazelCullom

src/SliceWrapper.hpp

@@ -11,8 +11,8 @@ struct SliceWrapper {
   SliceWrapper(SliceType st) : st_(st)  {}
 
   KOKKOS_INLINE_FUNCTION
-  T& access(const int s, const int a, int i) const {
-    return st_.access(s,a,i);
+  T& access(const int s, const int a) const {
+    return st_.access(s,a);
   }
   int arraySize(int s) {
     return st_.arraySize(s);
@@ -24,23 +24,34 @@ struct SliceWrapper {
 
 using namespace Cabana;
 
-template <class ExecutionSpace, class MemorySpace, class T, int width, int vecLen>
+template <class ExecutionSpace, class MemorySpace, class... Ts>
 class CabSliceFactory {
+  static constexpr int vecLen = Impl::PerformanceTraits<ExecutionSpace>::vector_length/8;
+  using TypeTuple = std::tuple<Ts...>;
   using DeviceType = Kokkos::Device<ExecutionSpace, MemorySpace>;
-  using DataTypes = Cabana::MemberTypes<T[width]>;
+  using DataTypes = Cabana::MemberTypes<Ts...>;
+  using soa_t = SoA<DataTypes, vecLen>;
+  
+  template <class T, int stride>
   using member_slice_t = 
-    Cabana::Slice<T[width], DeviceType, 
+    Cabana::Slice<T, DeviceType, 
 		  Cabana::DefaultAccessMemory, 
-		  vecLen, width*vecLen>;
-  using wrapper_slice_t = SliceWrapper<member_slice_t, T>;
+		  vecLen, stride>;
+
+  template <class T, int stride>
+  using wrapper_slice_t = SliceWrapper<member_slice_t<T, stride>, T>;
 
   Cabana::AoSoA<DataTypes, DeviceType, vecLen> aosoa; 
 
 public:
-  wrapper_slice_t makeSliceCab() {
-    auto slice0 = Cabana::slice<0>(aosoa);
-    return wrapper_slice_t(std::move(slice0));
+  template <std::size_t index>
+  auto makeSliceCab() {
+    using type = std::tuple_element_t<index, TypeTuple>;
+    const int stride = (vecLen * sizeof(soa_t)) / (4 * sizeof(type));
+    auto slice = Cabana::slice<index>(aosoa);
+    return wrapper_slice_t< type, stride >(std::move(slice));
   }
+  
   CabSliceFactory(int n) : aosoa("sliceAoSoA", n) {}
 };
 

test/SliceWrapper.cpp

@@ -3,34 +3,37 @@
 int main(int argc, char* argv[]) {
   // AoSoA parameters
   const int vecLen = 4;
-  const int width = 1;
   int num_tuples = 10;
 
   Kokkos::ScopeGuard scope_guard(argc, argv);
 
-  using member_type = double;
-  using DataTypes = Cabana::MemberTypes<member_type[width]>;
-  using ExecutionSpace = Kokkos::Cuda;
-  using MemorySpace = Kokkos::CudaSpace;
-
+  using ExecutionSpace = Kokkos::DefaultExecutionSpace;
+  using MemorySpace = ExecutionSpace::memory_space;
 
   // Slice Wrapper Factory
   CabSliceFactory<ExecutionSpace, MemorySpace,
-		  member_type, width, vecLen> cabSliceFactory(num_tuples);
+		  double, int, float, char> cabSliceFactory(num_tuples);
 
-  auto slice_wrapper = cabSliceFactory.makeSliceCab();
+  auto slice_wrapper0 = cabSliceFactory.makeSliceCab<0>();
+  auto slice_wrapper1 = cabSliceFactory.makeSliceCab<1>();
+  auto slice_wrapper2 = cabSliceFactory.makeSliceCab<2>();
+  auto slice_wrapper3 = cabSliceFactory.makeSliceCab<3>();
 
   // simd_parallel_for setup
   Cabana::SimdPolicy<vecLen, ExecutionSpace> simd_policy(0, num_tuples);
 
   // kernel that reads and writes
   auto vector_kernel = KOKKOS_LAMBDA(const int s, const int a) {
-    for (int i = 0; i < width; i++) {
-      printf("s: %d, a: %d, i: %d\n", s,a,i);
-      double x = 42/(s+a+1.3);
-      slice_wrapper.access(s,a,i) = x;
-      printf("value: %lf\n", slice_wrapper.access(s,a,i));
-    }
+    printf("s: %d, a: %d\n", s,a);
+    double x = 42/(s+a+1.3);
+    slice_wrapper0.access(s,a) = x;
+    slice_wrapper1.access(s,a) = s+a;
+    slice_wrapper2.access(s,a) = float(x);
+    slice_wrapper3.access(s,a) = 'a'+s+a;
+    printf("SW0 value: %lf\n", slice_wrapper0.access(s,a));
+    printf("SW1 value: %d\n", slice_wrapper1.access(s,a));
+    printf("SW2 value: %f\n", slice_wrapper2.access(s,a));
+    printf("SW3 value: %c\n", slice_wrapper3.access(s,a));
   };
 
   Cabana::simd_parallel_for(simd_policy, vector_kernel, "parallel_for_cabSliceFactory");