From 5a464048c56734aa670bc58005311fe384d381ae Mon Sep 17 00:00:00 2001
From: Bruno Turcksin <bruno.turcksin@gmail.com>
Date: Wed, 14 Nov 2018 19:36:50 +0000
Subject: [PATCH] Template CUDA MatrixFree framework on vector

Allow the CUDA MatrixFree framework to use CUDAWrappers::Vector<Number> and
LinearAlgebra::distributed::Vector<Number, MemorySpace::CUDA>.
---
 .../deal.II/matrix_free/cuda_matrix_free.h    | 19 ++++++--------
 .../matrix_free/cuda_matrix_free.templates.h  | 25 ++++++++++++-------
 tests/cuda/matrix_free_matrix_vector_01.cu    |  6 ++++-
 tests/cuda/matrix_free_matrix_vector_02.cu    |  6 ++++-
 tests/cuda/matrix_free_matrix_vector_03.cu    |  6 ++++-
 tests/cuda/matrix_free_matrix_vector_04.cu    |  6 ++++-
 tests/cuda/matrix_free_matrix_vector_05.cu    |  6 ++++-
 tests/cuda/matrix_free_matrix_vector_06.cu    |  6 ++++-
 tests/cuda/matrix_free_matrix_vector_09.cu    |  6 ++++-
 tests/cuda/matrix_vector_common.h             | 19 ++++++++------
 tests/cuda/matrix_vector_mf.h                 | 19 ++++++++------
 11 files changed, 83 insertions(+), 41 deletions(-)

diff --git a/include/deal.II/matrix_free/cuda_matrix_free.h b/include/deal.II/matrix_free/cuda_matrix_free.h
index 5d4b5fb79a..b75ff1f9cd 100644
--- a/include/deal.II/matrix_free/cuda_matrix_free.h
+++ b/include/deal.II/matrix_free/cuda_matrix_free.h
@@ -79,10 +79,6 @@ namespace CUDAWrappers
     // TODO this should really be a CUDAWrappers::Point
     using point_type = Tensor<1, dim, Number>;
 
-    // Use Number2 so we don't hide the template parameter Number
-    template <typename Number2>
-    using CUDAVector = ::dealii::LinearAlgebra::CUDAWrappers::Vector<Number2>;
-
     /**
      * Parallelization scheme used: parallel_in_elem (parallelism at the level
      * of degrees of freedom) or parallel_over_elem (parallelism at the level of
@@ -183,18 +179,19 @@ namespace CUDAWrappers
      * This method runs the loop over all cells and apply the local operation on
      * each element in parallel. @p func is a functor which is appplied on each color.
      */
-    template <typename functor>
+    template <typename functor, typename VectorType>
     void
-    cell_loop(const functor &           func,
-              const CUDAVector<Number> &src,
-              CUDAVector<Number> &      dst) const;
+    cell_loop(const functor &   func,
+              const VectorType &src,
+              VectorType &      dst) const;
 
+    template <typename VectorType>
     void
-    copy_constrained_values(const CUDAVector<Number> &src,
-                            CUDAVector<Number> &      dst) const;
+    copy_constrained_values(const VectorType &src, VectorType &dst) const;
 
+    template <typename VectorType>
     void
-    set_constrained_values(const Number val, CUDAVector<Number> &dst) const;
+    set_constrained_values(const Number val, VectorType &dst) const;
 
     /**
      * Free all the memory allocated.
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 8655befcf5..435b53b6a0 100644
--- a/include/deal.II/matrix_free/cuda_matrix_free.templates.h
+++ b/include/deal.II/matrix_free/cuda_matrix_free.templates.h
@@ -724,11 +724,14 @@ namespace CUDAWrappers
 
 
   template <int dim, typename Number>
+  template <typename VectorType>
   void
-  MatrixFree<dim, Number>::copy_constrained_values(
-    const CUDAVector<Number> &src,
-    CUDAVector<Number> &      dst) const
+  MatrixFree<dim, Number>::copy_constrained_values(const VectorType &src,
+                                                   VectorType &      dst) const
   {
+    static_assert(
+      std::is_same<Number, typename VectorType::value_type>::value,
+      "VectorType::value_type and Number should be of the same type.");
     internal::copy_constrained_dofs<Number>
       <<<constraint_grid_dim, constraint_block_dim>>>(constrained_dofs,
                                                       n_constrained_dofs,
@@ -739,10 +742,14 @@ namespace CUDAWrappers
 
 
   template <int dim, typename Number>
+  template <typename VectorType>
   void
-  MatrixFree<dim, Number>::set_constrained_values(Number              val,
-                                                  CUDAVector<Number> &dst) const
+  MatrixFree<dim, Number>::set_constrained_values(Number      val,
+                                                  VectorType &dst) const
   {
+    static_assert(
+      std::is_same<Number, typename VectorType::value_type>::value,
+      "VectorType::value_type and Number should be of the same type.");
     internal::set_constrained_dofs<Number>
       <<<constraint_grid_dim, constraint_block_dim>>>(constrained_dofs,
                                                       n_constrained_dofs,
@@ -762,11 +769,11 @@ namespace CUDAWrappers
 
 
   template <int dim, typename Number>
-  template <typename functor>
+  template <typename functor, typename VectorType>
   void
-  MatrixFree<dim, Number>::cell_loop(const functor &           func,
-                                     const CUDAVector<Number> &src,
-                                     CUDAVector<Number> &      dst) const
+  MatrixFree<dim, Number>::cell_loop(const functor &   func,
+                                     const VectorType &src,
+                                     VectorType &      dst) const
   {
     for (unsigned int i = 0; i < n_colors; ++i)
       internal::apply_kernel_shmem<dim, Number, functor>
diff --git a/tests/cuda/matrix_free_matrix_vector_01.cu b/tests/cuda/matrix_free_matrix_vector_01.cu
index 4c9dbb4c7d..efd928d119 100644
--- a/tests/cuda/matrix_free_matrix_vector_01.cu
+++ b/tests/cuda/matrix_free_matrix_vector_01.cu
@@ -41,5 +41,9 @@ test()
   AffineConstraints<double> constraints;
   constraints.close();
 
-  do_test<dim, fe_degree, double, fe_degree + 1>(dof, constraints);
+  do_test<dim,
+          fe_degree,
+          double,
+          LinearAlgebra::CUDAWrappers::Vector<double>,
+          fe_degree + 1>(dof, constraints);
 }
diff --git a/tests/cuda/matrix_free_matrix_vector_02.cu b/tests/cuda/matrix_free_matrix_vector_02.cu
index 97b6d04369..c6388d8024 100644
--- a/tests/cuda/matrix_free_matrix_vector_02.cu
+++ b/tests/cuda/matrix_free_matrix_vector_02.cu
@@ -45,5 +45,9 @@ test()
                                            constraints);
   constraints.close();
 
-  do_test<dim, fe_degree, double, fe_degree + 1>(dof, constraints);
+  do_test<dim,
+          fe_degree,
+          double,
+          LinearAlgebra::CUDAWrappers::Vector<double>,
+          fe_degree + 1>(dof, constraints);
 }
diff --git a/tests/cuda/matrix_free_matrix_vector_03.cu b/tests/cuda/matrix_free_matrix_vector_03.cu
index cba3875261..cd556e16ee 100644
--- a/tests/cuda/matrix_free_matrix_vector_03.cu
+++ b/tests/cuda/matrix_free_matrix_vector_03.cu
@@ -68,5 +68,9 @@ test()
 
   // Skip 2D tests with even fe_degree
   if ((dim == 3) || ((fe_degree % 2) == 1))
-    do_test<dim, fe_degree, double, fe_degree + 1>(dof, constraints);
+    do_test<dim,
+            fe_degree,
+            double,
+            LinearAlgebra::CUDAWrappers::Vector<double>,
+            fe_degree + 1>(dof, constraints);
 }
diff --git a/tests/cuda/matrix_free_matrix_vector_04.cu b/tests/cuda/matrix_free_matrix_vector_04.cu
index 53db3bb6fc..61dd8b491d 100644
--- a/tests/cuda/matrix_free_matrix_vector_04.cu
+++ b/tests/cuda/matrix_free_matrix_vector_04.cu
@@ -44,5 +44,9 @@ test()
   AffineConstraints<double> constraints;
   constraints.close();
 
-  do_test<dim, fe_degree, double, fe_degree + 1>(dof, constraints);
+  do_test<dim,
+          fe_degree,
+          double,
+          LinearAlgebra::CUDAWrappers::Vector<double>,
+          fe_degree + 1>(dof, constraints);
 }
diff --git a/tests/cuda/matrix_free_matrix_vector_05.cu b/tests/cuda/matrix_free_matrix_vector_05.cu
index cdc08ea0fa..54ce17e032 100644
--- a/tests/cuda/matrix_free_matrix_vector_05.cu
+++ b/tests/cuda/matrix_free_matrix_vector_05.cu
@@ -78,5 +78,9 @@ test()
 
   // Skip 2D tests with even fe_degree
   if ((dim == 3) || ((fe_degree % 2) == 1))
-    do_test<dim, fe_degree, double, fe_degree + 1>(dof, constraints);
+    do_test<dim,
+            fe_degree,
+            double,
+            LinearAlgebra::CUDAWrappers::Vector<double>,
+            fe_degree + 1>(dof, constraints);
 }
diff --git a/tests/cuda/matrix_free_matrix_vector_06.cu b/tests/cuda/matrix_free_matrix_vector_06.cu
index 55cf8de3a9..55a4540e3d 100644
--- a/tests/cuda/matrix_free_matrix_vector_06.cu
+++ b/tests/cuda/matrix_free_matrix_vector_06.cu
@@ -74,5 +74,9 @@ test()
                                            constraints);
   constraints.close();
 
-  do_test<dim, fe_degree, double, fe_degree + 1>(dof, constraints);
+  do_test<dim,
+          fe_degree,
+          double,
+          LinearAlgebra::CUDAWrappers::Vector<double>,
+          fe_degree + 1>(dof, constraints);
 }
diff --git a/tests/cuda/matrix_free_matrix_vector_09.cu b/tests/cuda/matrix_free_matrix_vector_09.cu
index 5261136a62..1224ed4b68 100644
--- a/tests/cuda/matrix_free_matrix_vector_09.cu
+++ b/tests/cuda/matrix_free_matrix_vector_09.cu
@@ -72,5 +72,9 @@ test()
                                            constraints);
   constraints.close();
 
-  do_test<dim, fe_degree, double, fe_degree + 1>(dof, constraints);
+  do_test<dim,
+          fe_degree,
+          double,
+          LinearAlgebra::CUDAWrappers::Vector<double>,
+          fe_degree + 1>(dof, constraints);
 }
diff --git a/tests/cuda/matrix_vector_common.h b/tests/cuda/matrix_vector_common.h
index ee06444eb2..662086f595 100644
--- a/tests/cuda/matrix_vector_common.h
+++ b/tests/cuda/matrix_vector_common.h
@@ -34,6 +34,7 @@
 #include <deal.II/lac/affine_constraints.h>
 #include <deal.II/lac/cuda_vector.h>
 #include <deal.II/lac/dynamic_sparsity_pattern.h>
+#include <deal.II/lac/la_parallel_vector.h>
 #include <deal.II/lac/read_write_vector.h>
 #include <deal.II/lac/sparse_matrix.h>
 #include <deal.II/lac/vector.h>
@@ -53,7 +54,11 @@ test();
 
 
 
-template <int dim, int fe_degree, typename Number, int n_q_points_1d>
+template <int dim,
+          int fe_degree,
+          typename Number,
+          typename VectorType,
+          int n_q_points_1d>
 void
 do_test(const DoFHandler<dim> &          dof,
         const AffineConstraints<double> &constraints)
@@ -70,12 +75,12 @@ do_test(const DoFHandler<dim> &          dof,
   const QGauss<1> quad(n_q_points_1d);
   mf_data.reinit(mapping, dof, constraints, quad, additional_data);
 
-  const unsigned int                                    n_dofs = dof.n_dofs();
-  MatrixFreeTest<dim, fe_degree, Number, n_q_points_1d> mf(mf_data);
-  Vector<Number>                              in_host(n_dofs), out_host(n_dofs);
-  LinearAlgebra::ReadWriteVector<Number>      in(n_dofs), out(n_dofs);
-  LinearAlgebra::CUDAWrappers::Vector<Number> in_device(n_dofs);
-  LinearAlgebra::CUDAWrappers::Vector<Number> out_device(n_dofs);
+  const unsigned int n_dofs = dof.n_dofs();
+  MatrixFreeTest<dim, fe_degree, Number, VectorType, n_q_points_1d> mf(mf_data);
+  Vector<Number>                         in_host(n_dofs), out_host(n_dofs);
+  LinearAlgebra::ReadWriteVector<Number> in(n_dofs), out(n_dofs);
+  VectorType                             in_device(n_dofs);
+  VectorType                             out_device(n_dofs);
 
   for (unsigned int i = 0; i < n_dofs; ++i)
     {
diff --git a/tests/cuda/matrix_vector_mf.h b/tests/cuda/matrix_vector_mf.h
index 1279251019..4deba79eae 100644
--- a/tests/cuda/matrix_vector_mf.h
+++ b/tests/cuda/matrix_vector_mf.h
@@ -19,6 +19,7 @@
 // general, with and without hanging nodes).
 
 #include <deal.II/lac/cuda_vector.h>
+#include <deal.II/lac/la_parallel_vector.h>
 
 #include <deal.II/matrix_free/cuda_fe_evaluation.h>
 #include <deal.II/matrix_free/cuda_matrix_free.h>
@@ -94,7 +95,8 @@ operator()(const unsigned int                                          cell,
 template <int dim,
           int fe_degree,
           typename Number,
-          int n_q_points_1d = fe_degree + 1>
+          typename VectorType = LinearAlgebra::CUDAWrappers::Vector<Number>,
+          int n_q_points_1d   = fe_degree + 1>
 class MatrixFreeTest
 {
 public:
@@ -102,8 +104,7 @@ public:
     : data(data_in){};
 
   void
-  vmult(LinearAlgebra::CUDAWrappers::Vector<Number> &      dst,
-        const LinearAlgebra::CUDAWrappers::Vector<Number> &src);
+  vmult(VectorType &dst, const VectorType &src);
 
 private:
   const CUDAWrappers::MatrixFree<dim, Number> &data;
@@ -111,11 +112,15 @@ private:
 
 
 
-template <int dim, int fe_degree, typename Number, int n_q_points_1d>
+template <int dim,
+          int fe_degree,
+          typename Number,
+          typename VectorType,
+          int n_q_points_1d>
 void
-MatrixFreeTest<dim, fe_degree, Number, n_q_points_1d>::vmult(
-  LinearAlgebra::CUDAWrappers::Vector<Number> &      dst,
-  const LinearAlgebra::CUDAWrappers::Vector<Number> &src)
+MatrixFreeTest<dim, fe_degree, Number, VectorType, n_q_points_1d>::vmult(
+  VectorType &      dst,
+  const VectorType &src)
 {
   dst = static_cast<Number>(0.);
   HelmholtzOperator<dim, fe_degree, Number, n_q_points_1d> helmholtz_operator;
-- 
2.39.5