Fix AffineConstraints::set_zero for MemorySpace::CUDA

diff --git a/include/deal.II/lac/affine_constraints.templates.h b/include/deal.II/lac/affine_constraints.templates.h
index 5aad3b34a6d261b9f4e3f22bf858ebe3a4f41852..0bebac51bdc99e3881b9bcea34bb985a69bf811b 100644 (file)
--- a/include/deal.II/lac/affine_constraints.templates.h
+++ b/include/deal.II/lac/affine_constraints.templates.h
@@ -1904,17 +1904,14 @@ namespace internal
       size_type                                                      shift = 0)
     {
       Assert(shift == 0, ExcNotImplemented());
-      const int              n_constraints = cm.size();
-      std::vector<size_type> constrained_local_dofs_host(n_constraints);
+      std::vector<size_type> constrained_local_dofs_host;
 
-      std::transform(cm.begin(),
-                     cm.end(),
-                     constrained_local_dofs_host.begin(),
-                     [&vec](const size_type global_index) {
-                       return vec.get_partitioner()->global_to_local(
-                         global_index);
-                     });
+      for (const auto global_index : cm)
+        if (vec.in_local_range(global_index))
+          constrained_local_dofs_host.push_back(
+            vec.get_partitioner()->global_to_local(global_index));
 
+      const int  n_constraints = constrained_local_dofs_host.size();
       size_type *constrained_local_dofs_device;
       Utilities::CUDA::malloc(constrained_local_dofs_device, n_constraints);
       Utilities::CUDA::copy_to_dev(constrained_local_dofs_host,
@@ -1922,7 +1919,7 @@ namespace internal
 
       const int n_blocks = 1 + (n_constraints - 1) / CUDAWrappers::block_size;
       set_zero_kernel<<<n_blocks, CUDAWrappers::block_size>>>(
-        constrained_local_dofs_device, cm.size(), vec.get_values());
+        constrained_local_dofs_device, n_constraints, vec.get_values());
 #  ifdef DEBUG
       // Check that the kernel was launched correctly
       AssertCuda(cudaGetLastError());

diff --git a/tests/cuda/affine_constraints_set_zero.cu b/tests/cuda/affine_constraints_set_zero.cu
index 53c13c39d91657a6972f45fa7de73a4d9f98172d..9b497f2fb8963763ad7a9425331b1ebbb6d8e245 100644 (file)
--- a/tests/cuda/affine_constraints_set_zero.cu
+++ b/tests/cuda/affine_constraints_set_zero.cu
@@ -48,30 +48,53 @@ test()
   local_active.set_size(2 * numproc);
   local_active.add_range(myid * numproc, (myid + 1) * numproc);
 
-  LinearAlgebra::distributed::Vector<double, MemorySpace::CUDA> v;
-  v.reinit(local_active, complete_index_set(2 * numproc), MPI_COMM_WORLD);
-
-  const int n_blocks = 1 + v.size() / CUDAWrappers::block_size;
-  initialize_vector<<<n_blocks, CUDAWrappers::block_size>>>(v.get_values(),
-                                                            numproc,
-                                                            myid * numproc);
-  v.compress(VectorOperation::insert);
-
   AffineConstraints<double> cm;
-  cm.add_line(numproc * myid + 1);
+  cm.add_line(1);
+  cm.add_line(2);
   cm.close();
 
-  deallog << "vector before:" << std::endl;
-  v.print(deallog.get_file_stream());
-
-  deallog << std::endl;
   deallog << "CM:" << std::endl;
   cm.print(deallog.get_file_stream());
 
-  cm.set_zero(v);
+  LinearAlgebra::distributed::Vector<double, MemorySpace::CUDA> ghosted;
+  {
+    ghosted.reinit(local_active,
+                   complete_index_set(2 * numproc),
+                   MPI_COMM_WORLD);
+
+    const int n_blocks = 1 + ghosted.size() / CUDAWrappers::block_size;
+    initialize_vector<<<n_blocks, CUDAWrappers::block_size>>>(
+      ghosted.get_values(), numproc, myid * numproc);
+    ghosted.compress(VectorOperation::insert);
+
+    deallog << "ghosted vector before:" << std::endl;
+    ghosted.print(deallog.get_file_stream());
+
+    cm.set_zero(ghosted);
+
+    deallog << "ghosted vector after:" << std::endl;
+    ghosted.print(deallog.get_file_stream());
+  }
+
+  LinearAlgebra::distributed::Vector<double, MemorySpace::CUDA> distributed;
+  {
+    distributed.reinit(local_active,
+                       complete_index_set(2 * numproc),
+                       MPI_COMM_WORLD);
+
+    const int n_blocks = 1 + distributed.size() / CUDAWrappers::block_size;
+    initialize_vector<<<n_blocks, CUDAWrappers::block_size>>>(
+      distributed.get_values(), numproc, myid * numproc);
+    distributed.compress(VectorOperation::insert);
+
+    deallog << "distributed vector before:" << std::endl;
+    distributed.print(deallog.get_file_stream());
+
+    cm.set_zero(distributed);
 
-  deallog << "vector after:" << std::endl;
-  v.print(deallog.get_file_stream());
+    deallog << "distributed vector after:" << std::endl;
+    distributed.print(deallog.get_file_stream());
+  }
 
   deallog << "OK" << std::endl;
 }

diff --git a/tests/cuda/affine_constraints_set_zero.mpirun=2.output b/tests/cuda/affine_constraints_set_zero.mpirun=2.output
index 15ecc47505e39bf920baba6902fd7083dea24d40..deae0a430a91948defc6fd2d04cda365c7558012 100644 (file)
--- a/tests/cuda/affine_constraints_set_zero.mpirun=2.output
+++ b/tests/cuda/affine_constraints_set_zero.mpirun=2.output
@@ -1,31 +1,51 @@
 
-DEAL:0::vector before:
+DEAL:0::CM:
+    1 = 0
+    2 = 0
+DEAL:0::ghosted vector before:
 Process #0
 Local range: [0, 2), global size: 4
 Vector data:
 1.000e+00 2.000e+00 
-DEAL:0::
-DEAL:0::CM:
-    1 = 0
-DEAL:0::vector after:
+DEAL:0::ghosted vector after:
+Process #0
+Local range: [0, 2), global size: 4
+Vector data:
+1.000e+00 0.000e+00 
+DEAL:0::distributed vector before:
+Process #0
+Local range: [0, 2), global size: 4
+Vector data:
+1.000e+00 2.000e+00 
+DEAL:0::distributed vector after:
 Process #0
 Local range: [0, 2), global size: 4
 Vector data:
 1.000e+00 0.000e+00 
 DEAL:0::OK
 
-DEAL:1::vector before:
+DEAL:1::CM:
+    1 = 0
+    2 = 0
+DEAL:1::ghosted vector before:
 Process #1
 Local range: [2, 4), global size: 4
 Vector data:
 3.000e+00 4.000e+00 
-DEAL:1::
-DEAL:1::CM:
-    3 = 0
-DEAL:1::vector after:
+DEAL:1::ghosted vector after:
+Process #1
+Local range: [2, 4), global size: 4
+Vector data:
+0.000e+00 4.000e+00 
+DEAL:1::distributed vector before:
+Process #1
+Local range: [2, 4), global size: 4
+Vector data:
+3.000e+00 4.000e+00 
+DEAL:1::distributed vector after:
 Process #1
 Local range: [2, 4), global size: 4
 Vector data:
-3.000e+00 0.000e+00 
+0.000e+00 4.000e+00 
 DEAL:1::OK