From 56623fb26672cd6e9db2f2c33b19b787de64820b Mon Sep 17 00:00:00 2001
From: Bruno Turcksin <bruno.turcksin@gmail.com>
Date: Tue, 21 Mar 2023 15:15:15 -0400
Subject: [PATCH] Move constraint_mask to Kokkos::View

---
 .../deal.II/matrix_free/cuda_matrix_free.h    | 18 ++++++++-----
 .../matrix_free/cuda_matrix_free.templates.h  | 26 ++++++++++---------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/include/deal.II/matrix_free/cuda_matrix_free.h b/include/deal.II/matrix_free/cuda_matrix_free.h
index b3d1238624..80c5218aec 100644
--- a/include/deal.II/matrix_free/cuda_matrix_free.h
+++ b/include/deal.II/matrix_free/cuda_matrix_free.h
@@ -221,7 +221,9 @@ namespace CUDAWrappers
       /**
        * Mask deciding where constraints are set on a given cell.
        */
-      dealii::internal::MatrixFreeFunctions::ConstraintKinds *constraint_mask;
+      Kokkos::View<dealii::internal::MatrixFreeFunctions::ConstraintKinds *,
+                   MemorySpace::Default::kokkos_space>
+        constraint_mask;
 
       /**
        * If true, use graph coloring has been used and we can simply add into
@@ -544,7 +546,9 @@ namespace CUDAWrappers
     /**
      * Mask deciding where constraints are set on a given cell.
      */
-    std::vector<dealii::internal::MatrixFreeFunctions::ConstraintKinds *>
+    std::vector<
+      Kokkos::View<dealii::internal::MatrixFreeFunctions::ConstraintKinds *,
+                   MemorySpace::Default::kokkos_space>>
       constraint_mask;
 
     /**
@@ -782,8 +786,9 @@ namespace CUDAWrappers
     /**
      * Mask deciding where constraints are set on a given cell.
      */
-    std::vector<dealii::internal::MatrixFreeFunctions::ConstraintKinds>
-      constraint_mask;
+    typename Kokkos::View<
+      dealii::internal::MatrixFreeFunctions::ConstraintKinds *,
+      MemorySpace::Default::kokkos_space>::HostMirror constraint_mask;
 
     /**
      * If true, use graph coloring has been used and we can simply add into
@@ -837,9 +842,8 @@ namespace CUDAWrappers
         Kokkos::deep_copy(data_host.JxW, data.JxW);
       }
 
-    data_host.constraint_mask.resize(data_host.n_cells);
-    Utilities::CUDA::copy_to_host(data.constraint_mask,
-                                  data_host.constraint_mask);
+    data_host.constraint_mask = Kokkos::create_mirror(data.constraint_mask);
+    Kokkos::deep_copy(data_host.constraint_mask, data.constraint_mask);
 
     return data_host;
   }
diff --git a/include/deal.II/matrix_free/cuda_matrix_free.templates.h b/include/deal.II/matrix_free/cuda_matrix_free.templates.h
index 5dd210b3f1..f39594ed39 100644
--- a/include/deal.II/matrix_free/cuda_matrix_free.templates.h
+++ b/include/deal.II/matrix_free/cuda_matrix_free.templates.h
@@ -240,8 +240,9 @@ namespace CUDAWrappers
       Kokkos::View<Number **, MemorySpace::Default::kokkos_space> JxW;
       Kokkos::View<Number **[dim][dim], MemorySpace::Default::kokkos_space>
         inv_jacobian;
-      std::vector<dealii::internal::MatrixFreeFunctions::ConstraintKinds>
-        constraint_mask_host;
+      Kokkos::View<dealii::internal::MatrixFreeFunctions::ConstraintKinds *,
+                   MemorySpace::Default::kokkos_space>
+        constraint_mask;
       // Local buffer
       std::vector<types::global_dof_index> local_dof_indices;
       FEValues<dim>                        fe_values;
@@ -380,7 +381,11 @@ namespace CUDAWrappers
             n_cells,
             dofs_per_cell);
 
-      constraint_mask_host.resize(n_cells);
+      // Initialize to zero, i.e., unconstrained cell
+      constraint_mask =
+        Kokkos::View<dealii::internal::MatrixFreeFunctions::ConstraintKinds *,
+                     MemorySpace::Default::kokkos_space>(
+          "constraint_mask_" + std::to_string(color), n_cells);
     }
 
 
@@ -404,6 +409,10 @@ namespace CUDAWrappers
       for (unsigned int i = 0; i < dofs_per_cell; ++i)
         lexicographic_dof_indices[i] = local_dof_indices[lexicographic_inv[i]];
 
+      // FIXME too many deep_copy
+      auto constraint_mask_host =
+        Kokkos::create_mirror_view_and_copy(MemorySpace::Host::kokkos_space{},
+                                            constraint_mask);
       const ArrayView<dealii::internal::MatrixFreeFunctions::ConstraintKinds>
         cell_id_view(constraint_mask_host[cell_id]);
 
@@ -412,6 +421,7 @@ namespace CUDAWrappers
                                       {lexicographic_inv},
                                       lexicographic_dof_indices,
                                       cell_id_view);
+      Kokkos::deep_copy(constraint_mask, constraint_mask_host);
 
       // FIXME too many deep_copy
       auto local_to_global_host =
@@ -491,11 +501,7 @@ namespace CUDAWrappers
           data->inv_jacobian[color] = inv_jacobian;
         }
 
-      alloc_and_copy(
-        &data->constraint_mask[color],
-        ArrayView<const dealii::internal::MatrixFreeFunctions::ConstraintKinds>(
-          constraint_mask_host.data(), constraint_mask_host.size()),
-        n_cells);
+      data->constraint_mask[color] = constraint_mask;
     }
 
 
@@ -718,10 +724,6 @@ namespace CUDAWrappers
   void
   MatrixFree<dim, Number>::free()
   {
-    for (auto &constraint_mask_color_ptr : constraint_mask)
-      Utilities::CUDA::free(constraint_mask_color_ptr);
-    constraint_mask.clear();
-
     internal::used_objects[my_id].store(false);
     my_id = -1;
   }
-- 
2.39.5