Support BOUT_FOR_RAJA GPU field operators

CMakeLists.txt

@@ -386,8 +386,15 @@ if (BOUT_GENERATE_FIELDOPS)
   if (NOT ClangFormat_FOUND)
     message(FATAL_ERROR "clang-format not found, but you have requested to generate code!")
   endif()
+  if (BOUT_ENABLE_RAJA)
+    set(GEN_LOOP_EXEC "raja")
+  elseif (BOUT_ENABLE_OPENMP)
+    set(GEN_LOOP_EXEC "openmp")
+  else()
+    set(GEN_LOOP_EXEC "serial")
+  endif()
   add_custom_command( OUTPUT ${CMAKE_CURRENT_SOURCE_DIR}/src/field/generated_fieldops.cxx
-    COMMAND ${Python3_EXECUTABLE} gen_fieldops.py --filename generated_fieldops.cxx.tmp
+    COMMAND ${Python3_EXECUTABLE} gen_fieldops.py --loop-exec ${GEN_LOOP_EXEC} --filename generated_fieldops.cxx.tmp
     COMMAND ${ClangFormat_BIN} generated_fieldops.cxx.tmp -i
     COMMAND ${CMAKE_COMMAND} -E rename generated_fieldops.cxx.tmp generated_fieldops.cxx
     DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/src/field/gen_fieldops.jinja  ${CMAKE_CURRENT_SOURCE_DIR}/src/field/gen_fieldops.py

include/bout/coordinates_accessor.hxx

@@ -31,7 +31,7 @@
 ///    -> If Coordinates data is changed, the cache should be cleared
 ///    by calling CoordinatesAccessor::clear()
 struct CoordinatesAccessor {
-  CoordinatesAccessor() = delete;
+  CoordinatesAccessor() {}
 
   /// Constructor from Coordinates
   /// Copies data from coords, doesn't modify it

include/bout/field_accessor.hxx

@@ -57,10 +57,17 @@ struct FieldAccessor {
   /// Constructor from Field3D
   ///
   /// @param[in] f    The field to access. Must already be allocated
-  explicit FieldAccessor(FieldType& f) : coords(f.getCoordinates()) {
+  explicit FieldAccessor(FieldType& f) {
     ASSERT0(f.getLocation() == location);
     ASSERT0(f.isAllocated());
 
+    if (auto* Coords = f.getCoordinates()) {
+      coords = CoordinatesAccessor{Coords};
+    }
+    else {
+      coords = CoordinatesAccessor{};
+    }
+
     data = BoutRealArray{&f(0, 0, 0)};
 
     // Field size
@@ -81,15 +88,19 @@ struct FieldAccessor {
     ddt = BoutRealArray{&(f.timeDeriv()->operator()(0, 0, 0))};
   }
 
+  explicit FieldAccessor(const FieldType& f) : FieldAccessor(const_cast<FieldType&>(f)) {}
+
   /// Provide shorthand for access to field data.
   /// Does not convert between 3D and 2D indices,
   /// so fa[i] is equivalent to fa.data[i].
   ///
   BOUT_HOST_DEVICE inline const BoutReal& operator[](int ind) const { return data[ind]; }
+  BOUT_HOST_DEVICE inline BoutReal& operator[](int ind) { return data[ind]; }
 
   BOUT_HOST_DEVICE inline const BoutReal& operator[](const Ind3D& ind) const {
     return data[ind.ind];
   }
+  BOUT_HOST_DEVICE inline BoutReal& operator[](const Ind3D& ind) { return data[ind.ind]; }
 
   // Pointers to the field data arrays
   // These are wrapped in BoutRealArray types so they can be indexed with Ind3D or int
@@ -115,6 +126,9 @@ struct FieldAccessor {
 template <CELL_LOC location = CELL_CENTRE>
 using Field2DAccessor = FieldAccessor<location, Field2D>;
 
+template <CELL_LOC location = CELL_CENTRE>
+using Field3DAccessor = FieldAccessor<location, Field3D>;
+
 /// Syntactic sugar for time derivative of a field
 ///
 /// Usage:
@@ -130,4 +144,28 @@ BOUT_HOST_DEVICE inline BoutRealArray& ddt(const FieldAccessor<location, FieldTy
   return const_cast<BoutRealArray&>(fa.ddt);
 }
 
+struct FieldPerpAccessor {
+  FieldPerpAccessor() = delete;
+
+  int nx, nz;
+  int yindex;
+  BoutReal* data;
+
+  explicit FieldPerpAccessor(const FieldPerp& f) {
+    ASSERT0(f.isAllocated());
+
+    data = BoutRealArray{const_cast<BoutReal*>(&f(0, 0, 0))};
+
+    // Field size
+    nx = f.getNx();
+    nz = f.getNz();
+
+    yindex = f.getIndex();
+  }
+
+  BOUT_HOST_DEVICE int getIndex() const { return yindex; }
+  BOUT_HOST_DEVICE inline const BoutReal& operator[](int ind) const { return data[ind]; }
+  BOUT_HOST_DEVICE inline BoutReal& operator[](int ind) { return data[ind]; }
+};
+
 #endif

include/bout/mesh.hxx

@@ -762,6 +762,11 @@ public:
     return {(indPerp.ind - jz) * LocalNy + LocalNz * jy + jz, LocalNy, LocalNz};
   }
 
+  BOUT_HOST_DEVICE int flatIndPerpto3D(const int& flatIndPerp, const int nz, int jy = 0) const {
+    int jz = flatIndPerp % nz;
+    return (flatIndPerp - jz) * LocalNy + LocalNz * jy + jz;
+  }
+
   /// Converts an Ind3D to an Ind2D representing a 2D index using a lookup -- to be used with care
   Ind2D map3Dto2D(const Ind3D& ind3D) {
     return {indexLookup3Dto2D[ind3D.ind], LocalNy, 1};

include/bout/rajalib.hxx

@@ -23,6 +23,15 @@
 
 #include "RAJA/RAJA.hpp" // using RAJA lib
 
+#if BOUT_HAS_CUDA
+// TODO: Make configurable
+const int CUDA_BLOCK_SIZE = 256;
+using EXEC_POL = RAJA::cuda_exec<CUDA_BLOCK_SIZE>;
+//using EXEC_POL = RAJA::loop_exec;
+#else  // not BOUT_USE_CUDA
+using EXEC_POL = RAJA::loop_exec;
+#endif // end BOUT_USE_CUDA
+
 /// Wrapper around RAJA::forall
 /// Enables computations to be done on CPU or GPU (CUDA).
 ///
@@ -81,7 +90,7 @@ struct RajaForAll {
     // Note: must be a local variable
     const int* _ob_i_ind_raw = &_ob_i_ind[0];
     RAJA::forall<EXEC_POL>(RAJA::RangeSegment(0, _ob_i_ind.size()),
-                           [=] RAJA_DEVICE(int id) {
+                           [=] RAJA_DEVICE(int id) mutable {
                              // Look up index and call user function
                              f(_ob_i_ind_raw[id]);
                            });
@@ -127,7 +136,7 @@ private:
 /// to create variables which shadow the class members.
 ///
 #define BOUT_FOR_RAJA(index, region, ...) \
-  RajaForAll(region) << [ =, ##__VA_ARGS__ ] RAJA_DEVICE(int index)
+RajaForAll(region) << [ =, ##__VA_ARGS__ ] RAJA_DEVICE(int index) mutable
 
 #else // BOUT_HAS_RAJA
 

include/bout/single_index_ops.hxx

@@ -7,17 +7,6 @@
 
 #include "field_accessor.hxx"
 
-#if BOUT_HAS_RAJA
-//--  RAJA CUDA settings--------------------------------------------------------start
-#if BOUT_HAS_CUDA
-const int CUDA_BLOCK_SIZE = 256; // TODO: Make configurable
-using EXEC_POL = RAJA::cuda_exec<CUDA_BLOCK_SIZE>;
-#else  // not BOUT_USE_CUDA
-using EXEC_POL = RAJA::loop_exec;
-#endif // end BOUT_USE_CUDA
-////-----------CUDA settings------------------------------------------------------end
-#endif // end BOUT_HAS_RAJA
-
 // Ind3D: i.zp():
 BOUT_HOST_DEVICE inline int i_zp(const int id, const int nz) {
   int jz = id % nz;