Merge pull request #105 from PFLAREProject/pmisr_kokkos_optim

Pmisr kokkos optim

Author: stevendargaville

src/PMISR_DDCk.kokkos.cxx

@@ -88,23 +88,26 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
 
    // Device memory for the global variable cf_markers_local_d - be careful these aren't petsc ints
    cf_markers_local_d = intKokkosView("cf_markers_local_d", local_rows);
-   PetscScalar *cf_markers_local_real_d_ptr = NULL, *cf_markers_nonlocal_real_d_ptr = NULL;
-   // The real equivalents so we can use the existing petscsf from the input matrix
-   PetscScalarKokkosView cf_markers_local_real_d("cf_markers_local_real_d", local_rows);
-   PetscScalarKokkosView cf_markers_nonlocal_real_d;
-   cf_markers_local_real_d_ptr = cf_markers_local_real_d.data();
+   // Can't use the global directly within the parallel 
+   // regions on the device so just take a shallow copy
+   intKokkosView cf_markers_d = cf_markers_local_d;    
+
+   intKokkosView cf_markers_nonlocal_d;
+   int *cf_markers_d_ptr = NULL, *cf_markers_nonlocal_d_ptr = NULL;
+   cf_markers_d_ptr = cf_markers_d.data();
 
    // Host and device memory for the measure
    PetscScalarKokkosViewHost measure_local_h(measure_local, local_rows);
    PetscScalarKokkosView measure_local_d("measure_local_d", local_rows);   
    PetscScalar *measure_local_d_ptr = NULL, *measure_nonlocal_d_ptr = NULL;
    measure_local_d_ptr = measure_local_d.data();
    PetscScalarKokkosView measure_nonlocal_d;
+
    if (mpi) {
       measure_nonlocal_d = PetscScalarKokkosView("measure_nonlocal_d", cols_ao);   
       measure_nonlocal_d_ptr = measure_nonlocal_d.data();
-      cf_markers_nonlocal_real_d = PetscScalarKokkosView("cf_markers_nonlocal_real_d", cols_ao); 
-      cf_markers_nonlocal_real_d_ptr = cf_markers_nonlocal_real_d.data();
+      cf_markers_nonlocal_d = intKokkosView("cf_markers_nonlocal_d", cols_ao); 
+      cf_markers_nonlocal_d_ptr = cf_markers_nonlocal_d.data();
    }
 
    // Device memory for the mark
@@ -153,43 +156,39 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
 
    // Initialise the set
    PetscInt counter_in_set_start = 0;
-   // Count how many in the set to begin with
+   // Count how many in the set to begin with and set their CF markers
    Kokkos::parallel_reduce ("Reduction", local_rows, KOKKOS_LAMBDA (const PetscInt i, PetscInt& update) {
-      if (Kokkos::abs(measure_local_d[i]) < 1) update++;
-   }, counter_in_set_start);
-
-   Kokkos::parallel_for(
-      Kokkos::RangePolicy<>(0, local_rows), KOKKOS_LAMBDA(PetscInt i) {
-
-      if (Kokkos::abs(measure_local_d(i)) < 1)
+      if (Kokkos::abs(measure_local_d[i]) < 1) 
       {
          if (zero_measure_c_point_int == 1) {
             if (pmis_int == 1) {
                // Set as F here but reversed below to become C
-               cf_markers_local_real_d(i) = -1;
+               cf_markers_d(i) = -1;
             }
             else {
                // Becomes C
-               cf_markers_local_real_d(i) = 1;
+               cf_markers_d(i) = 1;
             }  
          }
          else {
             if (pmis_int == 1) {
                // Set as C here but reversed below to become F
                // Otherwise dirichlet conditions persist down onto the coarsest grid
-               cf_markers_local_real_d(i) = 1;
+               cf_markers_d(i) = 1;
             }
             else {
                // Becomes F
-               cf_markers_local_real_d(i) = -1;
+               cf_markers_d(i) = -1;
             }
-         }
+         }         
+         // Count
+         update++;
       }
       else
       {
-         cf_markers_local_real_d(i) = 0;
-      }
-   });  
+         cf_markers_d(i) = 0;
+      }      
+   }, counter_in_set_start);
 
    // Check the total number of undecided in parallel
    PetscInt counter_undecided, counter_parallel;
@@ -238,24 +237,17 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
       // Start the async scatter of the nonlocal cf_markers
       // ~~~~~~~~~
       if (mpi) {
-         // We can't overwrite any of the values in cf_markers_local_real_d while the forward scatter is still going
-         PetscSFBcastWithMemTypeBegin(mat_mpi->Mvctx, MPIU_SCALAR,
-                     mem_type, cf_markers_local_real_d_ptr,
-                     mem_type, cf_markers_nonlocal_real_d_ptr,
+         // We can't overwrite any of the values in cf_markers_d while the forward scatter is still going
+         // Be careful these aren't petscints
+         PetscSFBcastWithMemTypeBegin(mat_mpi->Mvctx, MPI_INT,
+                     mem_type, cf_markers_d_ptr,
+                     mem_type, cf_markers_nonlocal_d_ptr,
                      MPI_REPLACE);
       }
 
-      // This keeps track of which of the candidate nodes can become in the set
+      // mark_d keeps track of which of the candidate nodes can become in the set
       // Only need this because we want to do async comms so we need a way to trigger
       // a node not being in the set due to either strong local neighbours *or* strong offproc neighbours
-      Kokkos::deep_copy(mark_d, true);   
-
-      // Any that aren't zero cf marker are already assigned so set to to false
-      Kokkos::parallel_for(
-         Kokkos::RangePolicy<>(0, local_rows), KOKKOS_LAMBDA(PetscInt i) {
-
-            if (cf_markers_local_real_d(i) != 0) mark_d(i) = false;
-      });
 
       // ~~~~~~~~
       // Go and do the local component
@@ -269,9 +261,8 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
             PetscInt strong_neighbours = 0;
 
             // Check this row isn't already marked
-            if (cf_markers_local_real_d(i) == 0)
+            if (cf_markers_d(i) == 0)
             {
-               const PetscInt i = t.league_rank();
                const PetscInt ncols_local = device_local_i[i + 1] - device_local_i[i];
 
                // Reduce over local columns to get the number of strong neighbours
@@ -281,7 +272,7 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
 
                   // Have to only check active strong neighbours
                   if (measure_local_d(i) >= measure_local_d(device_local_j[device_local_i[i] + j]) && \
-                           cf_markers_local_real_d(device_local_j[device_local_i[i] + j]) == 0)
+                           cf_markers_d(device_local_j[device_local_i[i] + j]) == 0)
                   {
                      strong_count++;
                   }
@@ -292,7 +283,22 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
                // Only want one thread in the team to write the result
                Kokkos::single(Kokkos::PerTeam(t), [&]() {                  
                   // If we have any strong neighbours
-                  if (strong_neighbours > 0) mark_d(i) = false;     
+                  if (strong_neighbours > 0) 
+                  {
+                     mark_d(i) = false;     
+                  }
+                  else
+                  {
+                     mark_d(i) = true;  
+                  }
+               });
+            }
+            // Any that aren't zero cf marker are already assigned so set to to false
+            else
+            {
+               // Only want one thread in the team to write the result
+               Kokkos::single(Kokkos::PerTeam(t), [&]() {                  
+                  mark_d(i) = false;
                });
             }
       });
@@ -303,7 +309,8 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
       if (mpi) {
 
          // Finish the async scatter
-         PetscSFBcastEnd(mat_mpi->Mvctx, MPIU_SCALAR, cf_markers_local_real_d_ptr, cf_markers_nonlocal_real_d_ptr, MPI_REPLACE);
+         // Be careful these aren't petscints
+         PetscSFBcastEnd(mat_mpi->Mvctx, MPI_INT, cf_markers_d_ptr, cf_markers_nonlocal_d_ptr, MPI_REPLACE);
 
          Kokkos::parallel_for(
             Kokkos::TeamPolicy<>(PetscGetKokkosExecutionSpace(), local_rows, Kokkos::AUTO()),
@@ -314,9 +321,8 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
                PetscInt strong_neighbours = 0;
 
                // Check this row isn't already marked
-               if (cf_markers_local_real_d(i) == 0)
+               if (mark_d(i))
                {
-                  const PetscInt i = t.league_rank();
                   PetscInt ncols_nonlocal = device_nonlocal_i[i + 1] - device_nonlocal_i[i];
 
                   // Reduce over nonlocal columns to get the number of strong neighbours
@@ -325,7 +331,7 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
                      [&](const PetscInt j, PetscInt& strong_count) {     
 
                      if (measure_local_d(i) >= measure_nonlocal_d(device_nonlocal_j[device_nonlocal_i[i] + j])  && \
-                              cf_markers_nonlocal_real_d(device_nonlocal_j[device_nonlocal_i[i] + j]) == 0)
+                              cf_markers_nonlocal_d(device_nonlocal_j[device_nonlocal_i[i] + j]) == 0)
                      {
                         strong_count++;
                      }
@@ -335,25 +341,28 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
 
                   // Only want one thread in the team to write the result
                   Kokkos::single(Kokkos::PerTeam(t), [&]() {                  
-                     // If we have any strong neighbours
-                     if (strong_neighbours > 0) mark_d(i) = false;     
+                     // If we don't have any strong neighbours
+                     if (strong_neighbours == 0) cf_markers_d(i) = loops_through;
                   });
                }
          });
       }
+      // This cf_markers_d(i) = loops_through happens above in the case of mpi, saves a kernel launch
+      else
+      {
+         // The nodes that have mark equal to true have no strong active neighbours in the IS
+         // hence they can be in the IS
+         Kokkos::parallel_for(
+            Kokkos::RangePolicy<>(0, local_rows), KOKKOS_LAMBDA(PetscInt i) {
 
-      // The nodes that have mark equal to true have no strong active neighbours in the IS
-      // hence they can be in the IS
-      Kokkos::parallel_for(
-         Kokkos::RangePolicy<>(0, local_rows), KOKKOS_LAMBDA(PetscInt i) {
-
-            if (mark_d(i)) cf_markers_local_real_d(i) = double(loops_through);
-      });      
+               if (mark_d(i)) cf_markers_d(i) = loops_through;
+         });      
+      }
 
       if (mpi) 
       {
          // We're going to do an add reverse scatter, so set them to zero
-         Kokkos::deep_copy(cf_markers_nonlocal_real_d, 0.0);  
+         Kokkos::deep_copy(cf_markers_nonlocal_d, 0.0);  
 
          Kokkos::parallel_for(
             Kokkos::TeamPolicy<>(PetscGetKokkosExecutionSpace(), local_rows, Kokkos::AUTO()),
@@ -363,27 +372,27 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
                const PetscInt i = t.league_rank();
 
                // Check if this node has been assigned during this top loop
-               if (cf_markers_local_real_d(i) == loops_through)
+               if (cf_markers_d(i) == loops_through)
                {
-                  const PetscInt i = t.league_rank();
                   PetscInt ncols_nonlocal = device_nonlocal_i[i + 1] - device_nonlocal_i[i];
 
                   // For over nonlocal columns
                   Kokkos::parallel_for(
                      Kokkos::TeamThreadRange(t, ncols_nonlocal), [&](const PetscInt j) {
 
                         // Needs to be atomic as may being set by many threads
-                        Kokkos::atomic_store(&cf_markers_nonlocal_real_d(device_nonlocal_j[device_nonlocal_i[i] + j]), 1.0);     
+                        Kokkos::atomic_store(&cf_markers_nonlocal_d(device_nonlocal_j[device_nonlocal_i[i] + j]), 1.0);     
                   });     
                }
          });
 
          // We've updated the values in cf_markers_nonlocal
          // Calling a reverse scatter add will then update the values of cf_markers_local
          // Reduce with a sum, equivalent to VecScatterBegin with ADD_VALUES, SCATTER_REVERSE
-         PetscSFReduceWithMemTypeBegin(mat_mpi->Mvctx, MPIU_SCALAR,
-            mem_type, cf_markers_nonlocal_real_d_ptr,
-            mem_type, cf_markers_local_real_d_ptr,
+         // Be careful these aren't petscints
+         PetscSFReduceWithMemTypeBegin(mat_mpi->Mvctx, MPI_INT,
+            mem_type, cf_markers_nonlocal_d_ptr,
+            mem_type, cf_markers_d_ptr,
             MPIU_SUM);
       }
 
@@ -396,25 +405,27 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
             const PetscInt i = t.league_rank();
 
             // Check if this node has been assigned during this top loop
-            if (cf_markers_local_real_d(i) == loops_through)
+            if (cf_markers_d(i) == loops_through)
             {
-               const PetscInt i = t.league_rank();
                const PetscInt ncols_local = device_local_i[i + 1] - device_local_i[i];
 
                // For over nonlocal columns
                Kokkos::parallel_for(
                   Kokkos::TeamThreadRange(t, ncols_local), [&](const PetscInt j) {
 
                      // Needs to be atomic as may being set by many threads
-                     Kokkos::atomic_store(&cf_markers_local_real_d(device_local_j[device_local_i[i] + j]), 1.0);     
+                     // Tried a version where instead of a "push" approach I tried a pull approach
+                     // that doesn't need an atomic, but it was slower
+                     Kokkos::atomic_store(&cf_markers_d(device_local_j[device_local_i[i] + j]), 1.0);     
                });     
             }
       });   
 
       if (mpi) 
       {
          // Finish the scatter
-         PetscSFReduceEnd(mat_mpi->Mvctx, MPIU_SCALAR, cf_markers_nonlocal_real_d_ptr, cf_markers_local_real_d_ptr, MPIU_SUM);         
+         // Be careful these aren't petscints
+         PetscSFReduceEnd(mat_mpi->Mvctx, MPI_INT, cf_markers_nonlocal_d_ptr, cf_markers_d_ptr, MPIU_SUM);         
       }
 
       // We've done another top level loop
@@ -427,7 +438,7 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
 
          counter_undecided = 0;  
          Kokkos::parallel_reduce ("ReductionCounter_undecided", local_rows, KOKKOS_LAMBDA (const PetscInt i, PetscInt& update) {
-            if (cf_markers_local_real_d(i) == 0) update++;
+            if (cf_markers_d(i) == 0) update++;
          }, counter_undecided); 
 
          // Parallel reduction!
@@ -442,17 +453,14 @@ PETSC_INTERN void pmisr_kokkos(Mat *strength_mat, const int max_luby_steps, cons
    // Now assign our final cf markers
    // ~~~~~~~~~
 
-   // Can't use the global directly within the parallel 
-   // regions on the device
-   intKokkosView cf_markers_d = cf_markers_local_d;   
    Kokkos::parallel_for(
       Kokkos::RangePolicy<>(0, local_rows), KOKKOS_LAMBDA(PetscInt i) {
 
-      if (cf_markers_local_real_d(i) == 0)
+      if (cf_markers_d(i) == 0)
       {
          cf_markers_d(i) = 1;
       }
-      else if (cf_markers_local_real_d(i) < 0)
+      else if (cf_markers_d(i) < 0)
       {
          cf_markers_d(i) = -1;
       }