From: Bruno Turcksin <bruno.turcksin@gmail.com>
Date: Thu, 16 Mar 2023 18:11:40 +0000 (-0400)
Subject: Fix a bug when using LA::CUDAWrappers::Vector
X-Git-Tag: v9.5.0-rc1~349^2~13
X-Git-Url: https://gitweb.dealii.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=ec57499ca89915bd4614d11297da608d3f5106f9;p=dealii.git

Fix a bug when using LA::CUDAWrappers::Vector
---

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 b3b9b13d2d..9df3cfee63 100644
--- a/include/deal.II/matrix_free/cuda_matrix_free.templates.h
+++ b/include/deal.II/matrix_free/cuda_matrix_free.templates.h
@@ -576,6 +576,38 @@ namespace CUDAWrappers
       if (cell < gpu_data.n_cells)
         func(cell, &gpu_data);
     }
+
+
+
+    template <typename VectorType>
+    struct VectorLocalSize
+    {
+      static unsigned int
+      get(const VectorType &vec)
+      {
+        return vec.locally_owned_size();
+      }
+    };
+
+    template <>
+    struct VectorLocalSize<LinearAlgebra::CUDAWrappers::Vector<double>>
+    {
+      static unsigned int
+      get(const LinearAlgebra::CUDAWrappers::Vector<double> &vec)
+      {
+        return vec.size();
+      }
+    };
+
+    template <>
+    struct VectorLocalSize<LinearAlgebra::CUDAWrappers::Vector<float>>
+    {
+      static unsigned int
+      get(const LinearAlgebra::CUDAWrappers::Vector<float> &vec)
+      {
+        return vec.size();
+      }
+    };
   } // namespace internal
 
 
@@ -738,10 +770,9 @@ namespace CUDAWrappers
     // FIXME When using C++17, we can use KOKKOS_CLASS_LAMBDA and this
     // work-around can be removed.
     types::global_dof_index *constr_dofs = constrained_dofs;
-    const unsigned int       size =
-      partitioner ? dst.locally_owned_size() : dst.size();
-    const Number *src_ptr = src.get_values();
-    Number *      dst_ptr = dst.get_values();
+    const unsigned int size = internal::VectorLocalSize<VectorType>::get(dst);
+    const Number *     src_ptr = src.get_values();
+    Number *           dst_ptr = dst.get_values();
     Kokkos::parallel_for(
       "copy_constrained_values",
       Kokkos::RangePolicy<MemorySpace::Default::kokkos_space::execution_space>(