1 // --------------------------------------------------------------------- 2 // 3 // Copyright (C) 2023 by the deal.II authors 4 // 5 // This file is part of the deal.II library. 6 // 7 // The deal.II library is free software; you can use it, redistribute 8 // it, and/or modify it under the terms of the GNU Lesser General 9 // Public License as published by the Free Software Foundation; either 10 // version 2.1 of the License, or (at your option) any later version. 11 // The full text of the license can be found in the file LICENSE.md at 12 // the top level directory of deal.II. 13 // 14 // Authors: Peter Munch, Martin Kronbichler 15 // 16 // --------------------------------------------------------------------- 17 18 #pragma once 19 #ifndef bps_kokkos_h 20 # define bps_kokkos_h 21 22 // deal.II includes 23 # include <deal.II/dofs/dof_tools.h> 24 25 # include <deal.II/fe/mapping.h> 26 27 # include <deal.II/lac/la_parallel_vector.h> 28 29 # include <deal.II/matrix_free/fe_evaluation.h> 30 # include <deal.II/matrix_free/matrix_free.h> 31 # include <deal.II/matrix_free/shape_info.h> 32 # include <deal.II/matrix_free/tools.h> 33 34 // local includes 35 # include "bps.h" 36 37 using namespace dealii; 38 39 40 41 template <int dim, int fe_degree, int n_q_points_1d, int n_components, typename Number> 42 class OperatorDealiiMassQuad 43 { 44 public: 45 DEAL_II_HOST_DEVICE void 46 operator()(Portable::FEEvaluation<dim, fe_degree, n_q_points_1d, n_components, Number> *fe_eval, 47 const int q_point) const 48 { 49 fe_eval->submit_value(fe_eval->get_value(q_point), q_point); 50 } 51 }; 52 53 54 55 template <int dim, int fe_degree, int n_q_points_1d, int n_components, typename Number> 56 class OperatorDealiiLaplaceQuad 57 { 58 public: 59 DEAL_II_HOST_DEVICE void 60 operator()(Portable::FEEvaluation<dim, fe_degree, n_q_points_1d, n_components, Number> *fe_eval, 61 const int q_point) const 62 { 63 fe_eval->submit_gradient(fe_eval->get_gradient(q_point), q_point); 64 } 65 }; 66 67 68 69 template <int dim, int fe_degree, int n_q_points_1d, int n_components, typename Number> 70 class OperatorDealiiMassLocal 71 { 72 public: 73 DEAL_II_HOST_DEVICE void 74 operator()(const typename Portable::MatrixFree<dim, Number>::Data *data, 75 const Portable::DeviceVector<Number> &src, 76 Portable::DeviceVector<Number> &dst) const 77 { 78 Portable::FEEvaluation<dim, fe_degree, n_q_points_1d, n_components, Number> fe_eval(data); 79 fe_eval.read_dof_values(src); 80 fe_eval.evaluate(EvaluationFlags::values); 81 fe_eval.apply_for_each_quad_point( 82 OperatorDealiiMassQuad<dim, fe_degree, n_q_points_1d, n_components, Number>()); 83 fe_eval.integrate(EvaluationFlags::values); 84 fe_eval.distribute_local_to_global(dst); 85 } 86 87 static const unsigned int n_local_dofs = Utilities::pow(fe_degree + 1, dim) * n_components; 88 static const unsigned int n_q_points = Utilities::pow(n_q_points_1d, dim); 89 }; 90 91 92 93 template <int dim, int fe_degree, int n_q_points_1d, int n_components, typename Number> 94 class OperatorDealiiLaplaceLocal 95 { 96 public: 97 DEAL_II_HOST_DEVICE void 98 operator()(const typename Portable::MatrixFree<dim, Number>::Data *data, 99 const Portable::DeviceVector<Number> &src, 100 Portable::DeviceVector<Number> &dst) const 101 { 102 Portable::FEEvaluation<dim, fe_degree, n_q_points_1d, n_components, Number> fe_eval(data); 103 fe_eval.read_dof_values(src); 104 fe_eval.evaluate(EvaluationFlags::gradients); 105 fe_eval.apply_for_each_quad_point( 106 OperatorDealiiLaplaceQuad<dim, fe_degree, n_q_points_1d, n_components, Number>()); 107 fe_eval.integrate(EvaluationFlags::gradients); 108 fe_eval.distribute_local_to_global(dst); 109 } 110 111 static const unsigned int n_local_dofs = Utilities::pow(fe_degree + 1, dim) * n_components; 112 static const unsigned int n_q_points = Utilities::pow(n_q_points_1d, dim); 113 }; 114 115 116 117 /** 118 * Operator GPU implementation using deal.II. 119 */ 120 template <int dim, typename Number> 121 class OperatorDealii : public OperatorBase<Number, MemorySpace::Default> 122 { 123 public: 124 using VectorType = typename OperatorBase<Number, MemorySpace::Default>::VectorType; 125 126 /** 127 * Constructor. 128 */ 129 OperatorDealii(const Mapping<dim> &mapping, 130 const DoFHandler<dim> &dof_handler, 131 const AffineConstraints<Number> &constraints, 132 const Quadrature<dim> &quadrature, 133 const BPType &bp) 134 : mapping(mapping) 135 , dof_handler(dof_handler) 136 , constraints(constraints) 137 , quadrature(quadrature) 138 , bp(bp) 139 { 140 reinit(); 141 } 142 143 /** 144 * Destructor. 145 */ 146 ~OperatorDealii() = default; 147 148 /** 149 * Initialized internal data structures, particularly, MatrixFree. 150 */ 151 void 152 reinit() override 153 { 154 // configure MatrixFree 155 typename Portable::MatrixFree<dim, Number>::AdditionalData additional_data; 156 157 if (bp <= BPType::BP2) // mass matrix 158 additional_data.mapping_update_flags = update_JxW_values | update_values; 159 else 160 additional_data.mapping_update_flags = update_JxW_values | update_gradients; 161 162 // create MatrixFree 163 AssertThrow(quadrature.is_tensor_product(), ExcNotImplemented()); 164 matrix_free.reinit( 165 mapping, dof_handler, constraints, quadrature.get_tensor_basis()[0], additional_data); 166 } 167 168 /** 169 * Matrix-vector product. 170 */ 171 void 172 vmult(VectorType &dst, const VectorType &src) const override 173 { 174 dst = 0.0; 175 176 const unsigned int n_components = dof_handler.get_fe().n_components(); 177 const unsigned int fe_degree = dof_handler.get_fe().tensor_degree(); 178 const unsigned int n_q_points_1d = quadrature.get_tensor_basis()[0].size(); 179 180 if (n_components == 1 && fe_degree == 1 && n_q_points_1d == 2) 181 this->vmult_internal<1, 1, 2>(dst, src); 182 else if (n_components == 1 && fe_degree == 2 && n_q_points_1d == 3) 183 this->vmult_internal<1, 2, 3>(dst, src); 184 else if (n_components == dim && fe_degree == 1 && n_q_points_1d == 2) 185 this->vmult_internal<dim, 1, 2>(dst, src); 186 else if (n_components == dim && fe_degree == 2 && n_q_points_1d == 3) 187 this->vmult_internal<dim, 2, 3>(dst, src); 188 else if (n_components == 1 && fe_degree == 1 && n_q_points_1d == 3) 189 this->vmult_internal<1, 1, 3>(dst, src); 190 else if (n_components == 1 && fe_degree == 2 && n_q_points_1d == 4) 191 this->vmult_internal<1, 2, 4>(dst, src); 192 else if (n_components == dim && fe_degree == 1 && n_q_points_1d == 3) 193 this->vmult_internal<dim, 1, 3>(dst, src); 194 else if (n_components == dim && fe_degree == 2 && n_q_points_1d == 4) 195 this->vmult_internal<dim, 2, 4>(dst, src); 196 else 197 AssertThrow(false, ExcInternalError()); 198 199 matrix_free.copy_constrained_values(src, dst); 200 } 201 202 /** 203 * Initialize vector. 204 */ 205 void 206 initialize_dof_vector(VectorType &vec) const override 207 { 208 matrix_free.initialize_dof_vector(vec); 209 } 210 211 /** 212 * Compute inverse of diagonal. 213 */ 214 void 215 compute_inverse_diagonal(VectorType &diagonal) const override 216 { 217 this->initialize_dof_vector(diagonal); 218 219 const unsigned int n_components = dof_handler.get_fe().n_components(); 220 const unsigned int fe_degree = dof_handler.get_fe().tensor_degree(); 221 const unsigned int n_q_points_1d = quadrature.get_tensor_basis()[0].size(); 222 223 if (n_components == 1 && fe_degree == 1 && n_q_points_1d == 2) 224 this->compute_inverse_diagonal_internal<1, 1, 2>(diagonal); 225 else if (n_components == 1 && fe_degree == 2 && n_q_points_1d == 3) 226 this->compute_inverse_diagonal_internal<1, 2, 3>(diagonal); 227 else if (n_components == dim && fe_degree == 1 && n_q_points_1d == 2) 228 this->compute_inverse_diagonal_internal<dim, 1, 2>(diagonal); 229 else if (n_components == dim && fe_degree == 2 && n_q_points_1d == 3) 230 this->compute_inverse_diagonal_internal<dim, 2, 3>(diagonal); 231 else if (n_components == 1 && fe_degree == 1 && n_q_points_1d == 3) 232 this->compute_inverse_diagonal_internal<1, 1, 3>(diagonal); 233 else if (n_components == 1 && fe_degree == 2 && n_q_points_1d == 4) 234 this->compute_inverse_diagonal_internal<1, 2, 4>(diagonal); 235 else if (n_components == dim && fe_degree == 1 && n_q_points_1d == 3) 236 this->compute_inverse_diagonal_internal<dim, 1, 3>(diagonal); 237 else if (n_components == dim && fe_degree == 2 && n_q_points_1d == 4) 238 this->compute_inverse_diagonal_internal<dim, 2, 4>(diagonal); 239 else 240 AssertThrow(false, ExcInternalError()); 241 } 242 243 private: 244 /** 245 * Templated vmult function. 246 */ 247 template <int n_components, int fe_degree, int n_q_points_1d> 248 void 249 vmult_internal(VectorType &dst, const VectorType &src) const 250 { 251 if (bp <= BPType::BP2) // mass matrix 252 { 253 OperatorDealiiMassLocal<dim, fe_degree, n_q_points_1d, n_components, Number> mass_operator; 254 matrix_free.cell_loop(mass_operator, src, dst); 255 } 256 else 257 { 258 OperatorDealiiLaplaceLocal<dim, fe_degree, n_q_points_1d, n_components, Number> 259 local_operator; 260 matrix_free.cell_loop(local_operator, src, dst); 261 } 262 } 263 264 /** 265 * Templated compute_inverse_diagonal function. 266 */ 267 template <int n_components, int fe_degree, int n_q_points_1d> 268 void 269 compute_inverse_diagonal_internal(VectorType &diagonal) const 270 { 271 if (bp <= BPType::BP2) // mass matrix 272 { 273 OperatorDealiiMassQuad<dim, fe_degree, n_q_points_1d, n_components, Number> op_quad; 274 275 MatrixFreeTools::compute_diagonal<dim, fe_degree, n_q_points_1d, n_components, Number>( 276 matrix_free, diagonal, op_quad, EvaluationFlags::values, EvaluationFlags::values); 277 } 278 else 279 { 280 OperatorDealiiLaplaceQuad<dim, fe_degree, n_q_points_1d, n_components, Number> op_quad; 281 282 MatrixFreeTools::compute_diagonal<dim, fe_degree, n_q_points_1d, n_components, Number>( 283 matrix_free, diagonal, op_quad, EvaluationFlags::gradients, EvaluationFlags::gradients); 284 } 285 286 287 Number *diagonal_ptr = diagonal.get_values(); 288 289 Kokkos::parallel_for( 290 "lethe::invert_vector", 291 Kokkos::RangePolicy<MemorySpace::Default::kokkos_space::execution_space>( 292 0, diagonal.locally_owned_size()), 293 KOKKOS_LAMBDA(int i) { diagonal_ptr[i] = 1.0 / diagonal_ptr[i]; }); 294 } 295 296 /** 297 * Mapping object passed to the constructor. 298 */ 299 const Mapping<dim> &mapping; 300 301 /** 302 * DoFHandler object passed to the constructor. 303 */ 304 const DoFHandler<dim> &dof_handler; 305 306 /** 307 * Constraints object passed to the constructor. 308 */ 309 const AffineConstraints<Number> &constraints; 310 311 /** 312 * Quadrature rule object passed to the constructor. 313 */ 314 const Quadrature<dim> &quadrature; 315 316 /** 317 * Selected BP. 318 */ 319 const BPType bp; 320 321 /** 322 * MatrixFree object. 323 */ 324 Portable::MatrixFree<dim, Number> matrix_free; 325 }; 326 327 #endif 328