From 50112fffc76d2639330233f744fdba84e2cff071 Mon Sep 17 00:00:00 2001 From: Bruno Turcksin Date: Fri, 30 Dec 2022 15:00:13 -0500 Subject: [PATCH] Use Kokkos in CUDAWrappers::MatrixFree::set_constrained_values() --- .../deal.II/matrix_free/cuda_matrix_free.h | 28 ------ .../matrix_free/cuda_matrix_free.templates.h | 87 ++++--------------- 2 files changed, 19 insertions(+), 96 deletions(-) diff --git a/include/deal.II/matrix_free/cuda_matrix_free.h b/include/deal.II/matrix_free/cuda_matrix_free.h index fca3ad059b..23fc9cc34d 100644 --- a/include/deal.II/matrix_free/cuda_matrix_free.h +++ b/include/deal.II/matrix_free/cuda_matrix_free.h @@ -482,34 +482,6 @@ namespace CUDAWrappers const LinearAlgebra::CUDAWrappers::Vector &src, LinearAlgebra::CUDAWrappers::Vector & dst) const; - /** - * Helper function. Set the constrained entries of @p dst to @p val. This - * function is used when MPI is not used. - */ - template - void - serial_set_constrained_values(const Number val, VectorType &dst) const; - - /** - * Helper function. Set the constrained entries of @p dst to @p val. This - * function is used when MPI is used. - */ - void - distributed_set_constrained_values( - const Number val, - LinearAlgebra::distributed::Vector &dst) const; - - /** - * This function should never be called. Calling it results in an internal - * error. This function exists only because set_constrained_values needs - * distributed_set_constrained_values() to exist for - * LinearAlgebra::CUDAWrappers::Vector. - */ - void - distributed_set_constrained_values( - const Number val, - 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 7ab5b105f8..f74752fdcb 100644 --- a/include/deal.II/matrix_free/cuda_matrix_free.templates.h +++ b/include/deal.II/matrix_free/cuda_matrix_free.templates.h @@ -36,6 +36,7 @@ # include # include +# include # include # include @@ -545,26 +546,6 @@ namespace CUDAWrappers - template - __global__ void - set_constrained_dofs( - const dealii::types::global_dof_index *constrained_dofs, - const unsigned int n_constrained_dofs, - const unsigned int size, - Number val, - 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 set the values of the - // constrained dofs of non-ghosted vectors. - if ((dof < n_constrained_dofs) && (constrained_dofs[dof] < size)) - dst[constrained_dofs[dof]] = val; - } - - - template __global__ void apply_kernel_shmem(Functor func, @@ -789,10 +770,24 @@ namespace CUDAWrappers static_assert( std::is_same::value, "VectorType::value_type and Number should be of the same type."); - if (partitioner) - distributed_set_constrained_values(val, dst); - else - serial_set_constrained_values(val, dst); + Number *dst_ptr = dst.get_values(); + // 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; + // When working with distributed vectors, the constrained dofs are + // computed for ghosted vectors but we want to set the values of the + // constrained dofs of non-ghosted vectors. + const unsigned int size = + partitioner ? dst.locally_owned_size() : dst.size(); + Kokkos::parallel_for( + "set_constrained_values", + Kokkos::RangePolicy< + ::dealii::MemorySpace::Default::kokkos_space::execution_space>( + 0, n_constrained_dofs), + KOKKOS_LAMBDA(int dof) { + if (constr_dofs[dof] < size) + dst_ptr[constr_dofs[dof]] = val; + }); } @@ -1358,50 +1353,6 @@ namespace CUDAWrappers Assert(false, ExcInternalError()); } - - - template - template - void - MatrixFree::serial_set_constrained_values(const Number val, - VectorType & dst) const - { - internal::set_constrained_dofs - <<>>(constrained_dofs, - n_constrained_dofs, - dst.size(), - val, - dst.get_values()); - AssertCudaKernel(); - } - - - - template - void - MatrixFree::distributed_set_constrained_values( - const Number val, - LinearAlgebra::distributed::Vector &dst) const - { - internal::set_constrained_dofs - <<>>(constrained_dofs, - n_constrained_dofs, - dst.locally_owned_size(), - val, - dst.get_values()); - AssertCudaKernel(); - } - - - - template - void - MatrixFree::distributed_set_constrained_values( - const Number, - LinearAlgebra::CUDAWrappers::Vector &) const - { - Assert(false, ExcInternalError()); - } } // namespace CUDAWrappers DEAL_II_NAMESPACE_CLOSE -- 2.39.5