WIP: Tracking work on CUDA operators

CMakeLists.txt

@@ -51,7 +51,7 @@ message(STATUS "Configuring BOUT++ version ${BOUT_FULL_VERSION}")
 project(BOUT++
   DESCRIPTION "Fluid PDE solver framework"
   VERSION ${BOUT_CMAKE_ACCEPTABLE_VERSION}
-  LANGUAGES CXX)
+  LANGUAGES C CXX)
 
 include(CMakeDependentOption)
 

examples/elm-pb-outerloop/elm_pb_outerloop.cxx

@@ -50,6 +50,27 @@
 
 #include <bout/rajalib.hxx> // Defines BOUT_FOR_RAJA
 
+#include <nvtx3/nvToolsExt.h>
+inline void nvtxPushColor(const char* name, uint32_t argb)
+{
+  nvtxEventAttributes_t ev{};
+  ev.version = NVTX_VERSION;
+  ev.size    = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+  ev.messageType = NVTX_MESSAGE_TYPE_ASCII;
+  ev.message.ascii = name;
+  ev.colorType = NVTX_COLOR_ARGB;
+  ev.color = argb;
+
+  nvtxRangePushEx(&ev);
+}
+
+inline void nvtxPop()
+{
+  nvtxRangePop();
+}
+
+
+
 #if BOUT_HAS_HYPRE
 #include <bout/invert/laplacexy2_hypre.hxx>
 #endif
@@ -1275,13 +1296,18 @@ class ELMpb : public PhysicsModel {
   bool first_run = true; // For printing out some diagnostics first time around
 
   int rhs(BoutReal t) override {
+
+    nvtxPushColor("rhs_communicate", 0xFFFFFF00);
     // Perform communications
     mesh->communicate(comms);
+    nvtxPop();
 
     Coordinates* metric = mesh->getCoordinates();
 
     ////////////////////////////////////////////
     // Transitions from 0 in core to 1 in vacuum
+
+    nvtxPushColor("rhs_nonlinear", 0xFFFFFF00);
     if (nonlinear) {
       vac_mask = (1.0 - tanh(((P0 + P) - vacuum_pressure) / vacuum_trans)) / 2.0;
 
@@ -1302,10 +1328,12 @@ class ELMpb : public PhysicsModel {
       }
       eta = interp_to(eta, loc);
     }
+    nvtxPop();
 
     ////////////////////////////////////////////
     // Resonant Magnetic Perturbation code
 
+    nvtxPushColor("rhs_include_rmp", 0xFFFFFF00);
     if (include_rmp) {
 
       if ((rmp_ramp > 0.0) || (rmp_freq > 0.0) || (rmp_rotate != 0.0)) {
@@ -1356,16 +1384,20 @@ class ELMpb : public PhysicsModel {
 
       mesh->communicate(rmp_Psi);
     }
+    nvtxPop();
 
     ////////////////////////////////////////////
     // Inversion
 
+    nvtxPushColor("BOUT_FOR_RAJA_aparSolver", 0xFFFFFF00);
     if (evolve_jpar) {
       // Invert laplacian for Psi
       Psi = aparSolver->solve(Jpar);
       mesh->communicate(Psi);
     }
+    nvtxPop();
 
+    nvtxPushColor("rhs_phi_constraint", 0xFFFFFF00);
     if (phi_constraint) {
       // Phi being solved as a constraint
 
@@ -1417,7 +1449,9 @@ class ELMpb : public PhysicsModel {
       phi.applyBoundary();
       mesh->communicate(phi);
     }
+    nvtxPop();
 
+    nvtxPushColor("rhs_evolve_jpar", 0xFFFFFF00);
     if (!evolve_jpar) {
       // Get J from Psi
       Jpar = Delp2(Psi);
@@ -1456,7 +1490,19 @@ class ELMpb : public PhysicsModel {
       }
 
       // Get Delp2(J) from J
-      Jpar2 = Delp2(Jpar);
+      nvtxPushColor("rhs_evolve_jpar", 0xFFFFFF00);
+      mesh->communicate(Jpar);
+      Jpar.applyBoundary();
+
+      Jpar2 = 0.0;
+      auto Jpar_acc  = FieldAccessor<>(Jpar);
+      auto Jpar2_acc = FieldAccessor<>(Jpar2);
+
+      BOUT_FOR_RAJA(i, Jpar.getRegion("RGN_NOBNDRY"),
+                    CAPTURE(Jpar_acc, Jpar2_acc)) {
+        Jpar2_acc[i] = Delp2(Jpar_acc, i);
+      };
+      nvtxPop();
 
       Jpar2.applyBoundary();
       mesh->communicate(Jpar2);
@@ -1479,6 +1525,7 @@ class ELMpb : public PhysicsModel {
         }
       }
     }
+    nvtxPop();
 
     ////////////////////////////////////////////////////
     // Sheath boundary conditions
@@ -1487,6 +1534,7 @@ class ELMpb : public PhysicsModel {
     // to Jpar so that Jpar = sqrt(mi/me)/(2*pi) * phi
     //
 
+    nvtxPushColor("rhs_sheath_boundaries", 0xFFFFFF00);
     if (sheath_boundaries) {
 
       // At y = ystart (lower boundary)
@@ -1533,6 +1581,7 @@ class ELMpb : public PhysicsModel {
         }
       }
     }
+    nvtxPop();
 
     ////////////////////////////////////////////////////
     // Check that settings match compile-time switches
@@ -1604,6 +1653,7 @@ class ELMpb : public PhysicsModel {
     // Note: Capture all class member variables into local scope
     //       or an illegal memory access may occur on GPUs
 
+    nvtxPushColor("BOUT_FOR_RAJA_ddt_Jpar_acc", 0xFFFFFF00);
     BOUT_FOR_RAJA(
         i, Jpar.getRegion("RGN_NOBNDRY"),
         CAPTURE(delta_i, hyperresist, relax_j_tconst, dnorm, ehyperviscos, viscos_perp)) {
@@ -1680,6 +1730,7 @@ class ELMpb : public PhysicsModel {
       ddt(P_acc)[i] = 0.0;
 #endif
     };
+    nvtxPop();
 
       // Terms which are not yet single index operators
       // Note: Terms which are included in the single index loop
@@ -1750,45 +1801,116 @@ class ELMpb : public PhysicsModel {
     //   ddt(U) += SQ(B0) * b0xGrad_dot_Grad(rmp_Psi, J0, CELL_CENTRE);
     // }
 
-    ddt(U) += b0xcv * Grad(P); // curvature term
+    nvtxPushColor("rhs_ddt", 0xFFFFFF00);
+
+    // nvtxPushColor("rhs_ddt_U_acc", 0xFFFFFF00); 
+    // ddt(U) += b0xcv * Grad(P); // curvature term
+    // nvtxPop();
+
+    // nvtxPushColor("rhs_accessor", 0xFFFFFF00);    
+    auto region = U.getRegion("RGN_NOBNDRY");
+
+    auto b0xcv_x = Field2DAccessor<>(b0xcv.x);
+    auto b0xcv_y = Field2DAccessor<>(b0xcv.y);
+    auto b0xcv_z = Field2DAccessor<>(b0xcv.z);
+
+    BOUT_FOR_RAJA(i, region, CAPTURE(P_acc, U_acc, b0xcv_x, b0xcv_y, b0xcv_z)) {
+      const int i2d = i / P_acc.mesh_nz;   // or U_acc.mesh_nz
+
+      const BoutReal dPdx = DDX(P_acc, i);
+      const BoutReal dPdy = DDY(P_acc, i);
+      const BoutReal dPdz = DDZ(P_acc, i);
+
+      ddt(U_acc)[i] += b0xcv_x[i2d] * dPdx
+                    + b0xcv_y[i2d] * dPdy
+                    + b0xcv_z[i2d] * dPdz;
+    };
+    nvtxPop();
 
     // if (!nogradparj) { // Parallel current term
     //   ddt(U) -= SQ(B0) * Grad_parP(Jpar, CELL_CENTRE); // b dot grad j
     // }
 
+    nvtxPushColor("rhs_ddt_withflow_K_H_term", 0xFFFFFF00);  
     if (withflow && K_H_term) { // K_H_term
       ddt(U) -= b0xGrad_dot_Grad(phi, U0);
     }
+    nvtxPop();
 
     // if (diamag_phi0) { // Equilibrium flow
     //   ddt(U) -= b0xGrad_dot_Grad(phi0, U);
     // }
 
+    nvtxPushColor("rhs_ddt_withflow", 0xFFFFFF00);  
     if (withflow) { // net flow
       ddt(U) -= V_dot_Grad(V0net, U);
     }
+    nvtxPop();
 
     // if (nonlinear) { // Advection
     //   ddt(U) -= bracket(phi, U, bm_exb) * B0;
     // }
 
     // Viscosity terms
-
+    nvtxPushColor("rhs_ddt_viscos_par", 0xFFFFFF00);  
     if (viscos_par > 0.0) { // Parallel viscosity
       ddt(U) += viscos_par * Grad2_par2(U);
     }
+    nvtxPop();
+
+    nvtxPushColor("rhs_ddt_diffusion_u4", 0xFFFFFF00);  
+    // if (diffusion_u4 > 0.0) {
+    //   tmpU2 = D2DY2(U);
+    //   mesh->communicate(tmpU2);
+    //   tmpU2.applyBoundary();
+    //   ddt(U) -= diffusion_u4 * D2DY2(tmpU2);
+    // }
+    // nvtxPop();
+
 
     if (diffusion_u4 > 0.0) {
-      tmpU2 = D2DY2(U);
+
+      // Kernel 1: tmpU2 = D2DY2(U)
+      nvtxPushColor("rhs_diff_initTmpU2", 0xFFFFFF00);    
+      tmpU2 = 0.0;
+      auto tmpU2a_acc = FieldAccessor<>(tmpU2);
+      nvtxPop();
+
+      nvtxPushColor("rhs_diff_loop1", 0xFFFFFF00);    
+      BOUT_FOR_RAJA(i, region, CAPTURE(U_acc, tmpU2a_acc)) {
+        tmpU2a_acc[i] = D2DY2(U_acc, i);
+      };
+      nvtxPop();
+
+      nvtxPushColor("rhs_diff_communicate", 0xFFFFFF00);    
+      // Host: halo exchange + BCs
       mesh->communicate(tmpU2);
       tmpU2.applyBoundary();
-      ddt(U) -= diffusion_u4 * D2DY2(tmpU2);
+      nvtxPop();
+
+      nvtxPushColor("rhs_diff_acce", 0xFFFFFF00);    
+      // Kernel 2: ddt(U) -= diffusion_u4 * D2DY2(tmpU2)
+      auto tmpU2b_acc = FieldAccessor<>(tmpU2);
+
+      const BoutReal diffusion = diffusion_u4;
+
+      auto& ddtU = ddt(U);               // keep alive (avoid proxy lifetime)
+      auto ddtU_acc = FieldAccessor<>(ddtU);
+      nvtxPop();
+
+      nvtxPushColor("rhs_diff_loop2", 0xFFFFFF00);    
+      BOUT_FOR_RAJA(i, region, CAPTURE(diffusion, ddtU_acc, tmpU2b_acc)) {
+        ddtU_acc[i] -= diffusion * D2DY2(tmpU2b_acc, i);
+      };
+      nvtxPop();
     }
+    nvtxPop();
 
     // if (viscos_perp > 0.0) { // Perpendicular viscosity
     //   ddt(U) += viscos_perp * Delp2(U);
     // }
-
+    nvtxPop();
+    nvtxPushColor("rhs_hyperviscos", 0xFFFFFF00);
     // Hyper-viscosity
     if (hyperviscos > 0.0) {
       // Calculate coefficient.
@@ -1874,7 +1996,8 @@ class ELMpb : public PhysicsModel {
 
     ////////////////////////////////////////////////////
     // Pressure equation
-
+    nvtxPop();
+    nvtxPushColor("rhs_evolve_pressure", 0xFFFFFF00);
     if (evolve_pressure) {
 
       // ddt(P) -= b0xGrad_dot_Grad(phi, P0);
@@ -1921,7 +2044,8 @@ class ELMpb : public PhysicsModel {
 
     ////////////////////////////////////////////////////
     // Compressional effects
-
+    nvtxPop();
+    nvtxPushColor("rhs_compress", 0xFFFFFF00);
     if (compress) {
 
       ddt(P) -= beta * Div_par(Vpar, CELL_CENTRE);
@@ -1980,6 +2104,8 @@ class ELMpb : public PhysicsModel {
 
     first_run = false;
 
+    nvtxPop();
+
     return 0;
   }
 

include/bout/bout_types.hxx

@@ -25,6 +25,8 @@
 #include <limits>
 #include <string>
 
+#include <cuda_runtime_api.h>
+
 /// Size of real numbers
 using BoutReal = double;
 

include/bout/field.hxx

@@ -546,7 +546,6 @@ T pow(BoutReal lhs, const T& rhs, const std::string& rgn = "RGN_ALL") {
   template <typename T, typename = bout::utils::EnableIfField<T>>                       \
   inline BinaryExpr<T, T, T, bout::op::name> name(const T& f,                           \
                                                   const std::string& rgn = "RGN_ALL") { \
-    std::cout << "RUNNING " #name " with CUDA\n";                                       \
     return BinaryExpr<T, T, T, bout::op::name>{static_cast<typename T::View>(f),        \
                                                static_cast<typename T::View>(f),        \
                                                bout::op::name{},                        \

include/bout/field2d.hxx

@@ -108,7 +108,7 @@ public:
                                   || (is_expr_constant_v<L> && is_expr_field2d_v<R>)
                                   || (is_expr_field2d_v<L> && is_expr_constant_v<R>)>>
   Field2D(const BinaryExpr<ResT, L, R, Func>& expr) {
-    std::cout << "RUNNING Field2D constructor with CUDA\n";
+    //std::cout << "RUNNING Field2D constructor with CUDA\n";
     Array<BoutReal> data{expr.size()};
     expr.evaluate(&data[0]);
     *this = std::move(Field2D{std::move(data), expr.getMesh(), expr.getLocation(),
@@ -192,7 +192,7 @@ public:
   template <typename ResT, typename L, typename R, typename Func>
   std::enable_if_t<is_expr_field2d_v<L>, Field2D&>
   operator=(const BinaryExpr<ResT, L, R, Func>& expr) {
-    std::cout << "RUNNING Field2D operator= with CUDA\n";
+    //std::cout << "RUNNING Field2D operator= with CUDA\n";
     if (isAllocated()) {
       expr.evaluate(&data[0]);
     } else {

include/bout/field3d.hxx

@@ -460,7 +460,7 @@ public:
   template <typename ResT, typename L, typename R, typename Func>
   std::enable_if_t<is_expr_field3d_v<L>, Field3D&>
   operator=(BinaryExpr<ResT, L, R, Func>& expr) {
-    std::cout << "RUNNING operator= with CUDA\n";
+    //std::cout << "RUNNING operator= with CUDA\n";
     regionID = expr.getRegionID();
     if(isAllocated()) {
       expr.evaluate(&data[0]);