From b48411df4d155a3444bd33f7c00c42840e832486 Mon Sep 17 00:00:00 2001 From: Bruno Turcksin Date: Fri, 30 Dec 2022 15:49:11 -0500 Subject: [PATCH] Use Kokkos in CUDAWrappers::MatrixFree::copy_constrained_values() --- .../deal.II/matrix_free/cuda_matrix_free.h | 29 ------ .../matrix_free/cuda_matrix_free.templates.h | 97 ++++--------------- 2 files changed, 21 insertions(+), 105 deletions(-) diff --git a/include/deal.II/matrix_free/cuda_matrix_free.h b/include/deal.II/matrix_free/cuda_matrix_free.h index 23fc9cc34d..078798384c 100644 --- a/include/deal.II/matrix_free/cuda_matrix_free.h +++ b/include/deal.II/matrix_free/cuda_matrix_free.h @@ -453,35 +453,6 @@ namespace CUDAWrappers const LinearAlgebra::CUDAWrappers::Vector &src, LinearAlgebra::CUDAWrappers::Vector & dst) const; - /** - * Helper function. Copy the values of the constrained entries of @p src to - * @p dst. This function is used when MPI is not used. - */ - template - void - serial_copy_constrained_values(const VectorType &src, - VectorType & dst) const; - - /** - * Helper function. Copy the values of the constrained entries of @p src to - * @p dst. This function is used when MPI is used. - */ - void - distributed_copy_constrained_values( - const LinearAlgebra::distributed::Vector &src, - LinearAlgebra::distributed::Vector &dst) const; - - /** - * This function should never be called. Calling it results in an internal - * error. This function exists only because copy_constrained_values needs - * distributed_copy_constrained_values() to exist for - * LinearAlgebra::CUDAWrappers::Vector. - */ - void - distributed_copy_constrained_values( - const LinearAlgebra::CUDAWrappers::Vector &src, - LinearAlgebra::CUDAWrappers::Vector & dst) const; - /** * Unique ID associated with the object. */ 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 f74752fdcb..b3b9b13d2d 100644 --- a/include/deal.II/matrix_free/cuda_matrix_free.templates.h +++ b/include/deal.II/matrix_free/cuda_matrix_free.templates.h @@ -526,26 +526,6 @@ namespace CUDAWrappers - template - __global__ void - copy_constrained_dofs( - const dealii::types::global_dof_index *constrained_dofs, - const unsigned int n_constrained_dofs, - const unsigned int size, - const Number * src, - Number * dst) - { - const unsigned int dof = - threadIdx.x + blockDim.x * (blockIdx.x + gridDim.x * blockIdx.y); - // When working with distributed vectors, the constrained dofs are - // computed for ghosted vectors but we want to copy the values of the - // constrained dofs of non-ghosted vectors. - if ((dof < n_constrained_dofs) && (constrained_dofs[dof] < size)) - dst[constrained_dofs[dof]] = src[constrained_dofs[dof]]; - } - - - template __global__ void apply_kernel_shmem(Functor func, @@ -753,10 +733,26 @@ namespace CUDAWrappers static_assert( std::is_same::value, "VectorType::value_type and Number should be of the same type."); - if (partitioner) - distributed_copy_constrained_values(src, dst); - else - serial_copy_constrained_values(src, dst); + Assert(src.size() == dst.size(), + ExcMessage("src and dst vectors have different size.")); + // 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(); + Kokkos::parallel_for( + "copy_constrained_values", + Kokkos::RangePolicy( + 0, n_constrained_dofs), + KOKKOS_LAMBDA(int dof) { + // When working with distributed vectors, the constrained dofs are + // computed for ghosted vectors but we want to copy the values of the + // constrained dofs of non-ghosted vectors. + if (constr_dofs[dof] < size) + dst_ptr[constr_dofs[dof]] = src_ptr[constr_dofs[dof]]; + }); } @@ -781,8 +777,7 @@ namespace CUDAWrappers partitioner ? dst.locally_owned_size() : dst.size(); Kokkos::parallel_for( "set_constrained_values", - Kokkos::RangePolicy< - ::dealii::MemorySpace::Default::kokkos_space::execution_space>( + Kokkos::RangePolicy( 0, n_constrained_dofs), KOKKOS_LAMBDA(int dof) { if (constr_dofs[dof] < size) @@ -1303,56 +1298,6 @@ namespace CUDAWrappers { Assert(false, ExcInternalError()); } - - - - template - template - void - MatrixFree::serial_copy_constrained_values(const VectorType &src, - VectorType &dst) const - { - Assert(src.size() == dst.size(), - ExcMessage("src and dst vectors have different size.")); - internal::copy_constrained_dofs - <<>>(constrained_dofs, - n_constrained_dofs, - src.size(), - src.get_values(), - dst.get_values()); - AssertCudaKernel(); - } - - - - template - void - MatrixFree::distributed_copy_constrained_values( - const LinearAlgebra::distributed::Vector &src, - LinearAlgebra::distributed::Vector &dst) const - { - Assert(src.size() == dst.size(), - ExcMessage("src and dst vectors have different local size.")); - internal::copy_constrained_dofs - <<>>(constrained_dofs, - n_constrained_dofs, - src.locally_owned_size(), - src.get_values(), - dst.get_values()); - AssertCudaKernel(); - } - - - - template - void - MatrixFree::distributed_copy_constrained_values( - const LinearAlgebra::CUDAWrappers::Vector &, - LinearAlgebra::CUDAWrappers::Vector &) const - { - Assert(false, ExcInternalError()); - } - } // namespace CUDAWrappers DEAL_II_NAMESPACE_CLOSE -- 2.39.5