From 5cb0fe6b2dcc72ded68ec91519e388382c184fe0 Mon Sep 17 00:00:00 2001
From: Bruno Turcksin <bruno.turcksin@gmail.com>
Date: Fri, 8 Feb 2019 14:47:25 +0000
Subject: [PATCH] Add new cuda kernel function to execute a binary operation on
 a subset of the vector

---
 include/deal.II/lac/cuda_kernels.h           | 17 +++++++++++++++++
 include/deal.II/lac/cuda_kernels.templates.h | 18 ++++++++++++++++++
 source/lac/cuda_kernels.cu                   | 20 ++++++++++++++++++++
 3 files changed, 55 insertions(+)

diff --git a/include/deal.II/lac/cuda_kernels.h b/include/deal.II/lac/cuda_kernels.h
index 963a946170..c3470a5749 100644
--- a/include/deal.II/lac/cuda_kernels.h
+++ b/include/deal.II/lac/cuda_kernels.h
@@ -181,6 +181,23 @@ namespace LinearAlgebra
 
 
 
+      /**
+       * Apply the functor @tparam Binop to the elements of @p v1 that have
+       * indices in @p mask and @p v2. The size of @p mask should be greater
+       * than the size of @p v1. @p mask and @p v2 should have the same size @p
+       * N.
+       *
+       * @ingroup CUDAWrappers
+       */
+      template <typename Number, template <typename> class Binop>
+      __global__ void
+      masked_vector_bin_op(const unsigned int *mask,
+                           Number *            v1,
+                           const Number *      v2,
+                           const size_type     N);
+
+
+
       /**
        * Structure implementing the functions used to add elements when
        * using a reduction.
diff --git a/include/deal.II/lac/cuda_kernels.templates.h b/include/deal.II/lac/cuda_kernels.templates.h
index cd522131ba..cadeb4d11a 100644
--- a/include/deal.II/lac/cuda_kernels.templates.h
+++ b/include/deal.II/lac/cuda_kernels.templates.h
@@ -60,6 +60,24 @@ namespace LinearAlgebra
 
 
 
+      template <typename Number, template <typename> class Binop>
+      __global__ void
+      masked_vector_bin_op(const unsigned int *mask,
+                           Number *            v1,
+                           const Number *      v2,
+                           const size_type     N)
+      {
+        const size_type idx_base =
+          threadIdx.x + blockIdx.x * (blockDim.x * chunk_size);
+        for (unsigned int i = 0; i < chunk_size; ++i)
+          {
+            const size_type idx = idx_base + i * block_size;
+            if (idx < N)
+              v1[mask[idx]] = Binop<Number>::operation(v1[mask[idx]], v2[idx]);
+          }
+      }
+
+
       template <typename Number>
       __device__ Number
                  ElemSum<Number>::reduction_op(const Number a, const Number b)
diff --git a/source/lac/cuda_kernels.cu b/source/lac/cuda_kernels.cu
index a1c83d5935..c7bfb9b020 100644
--- a/source/lac/cuda_kernels.cu
+++ b/source/lac/cuda_kernels.cu
@@ -37,6 +37,16 @@ namespace LinearAlgebra
       vector_bin_op<float, Binop_Subtraction>(float *         v1,
                                               const float *   v2,
                                               const size_type N);
+      template __global__ void
+      masked_vector_bin_op<float, Binop_Addition>(const unsigned int *mask,
+                                                  float *             v1,
+                                                  const float *       v2,
+                                                  const size_type     N);
+      template __global__ void
+      masked_vector_bin_op<float, Binop_Subtraction>(const unsigned int *mask,
+                                                     float *             v1,
+                                                     const float *       v2,
+                                                     const size_type     N);
       template struct ElemSum<float>;
       template struct L1Norm<float>;
       template struct LInfty<float>;
@@ -137,6 +147,16 @@ namespace LinearAlgebra
       vector_bin_op<double, Binop_Subtraction>(double *        v1,
                                                const double *  v2,
                                                const size_type N);
+      template __global__ void
+      masked_vector_bin_op<double, Binop_Addition>(const unsigned int *mask,
+                                                   double *            v1,
+                                                   const double *      v2,
+                                                   const size_type     N);
+      template __global__ void
+      masked_vector_bin_op<double, Binop_Subtraction>(const unsigned int *mask,
+                                                      double *            v1,
+                                                      const double *      v2,
+                                                      const size_type     N);
       template struct ElemSum<double>;
       template struct L1Norm<double>;
       template struct LInfty<double>;
-- 
2.39.5