Line data Source code
1 : //* This file is part of the MOOSE framework 2 : //* https://www.mooseframework.org 3 : //* 4 : //* All rights reserved, see COPYRIGHT for full restrictions 5 : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT 6 : //* 7 : //* Licensed under LGPL 2.1, please see LICENSE for details 8 : //* https://www.gnu.org/licenses/lgpl-2.1.html 9 : 10 : #pragma once 11 : 12 : #include "KokkosKernel.h" 13 : 14 : namespace Moose 15 : { 16 : namespace Kokkos 17 : { 18 : 19 : /** 20 : * The base class for a user to derive their own Kokkos kernels where the residual is of the form 21 : * 22 : * $(\dots, \nabla\psi_i)$ 23 : * 24 : * i.e. the gradient of the test function $(\nabla\psi_i)$ can be factored out for optimization. 25 : * 26 : * The user is expected to define precomputeQpResidual() and precomputeQpJacobian() instead of 27 : * computeQpResidual() and computeQpJacobian(). The signature of precomputeQpResidual() expected to 28 : * be defined in the derived class is as follows: 29 : * 30 : * @param qp The local quadrature point index 31 : * @param datum The ResidualDatum object of the current thread 32 : * @returns The vector component of the residual contribution that will be multiplied by the 33 : * gradient of the test function 34 : * 35 : * KOKKOS_FUNCTION Real3 precomputeQpResidual(const unsigned int qp, 36 : * ResidualDatum & datum) const; 37 : * 38 : * The signature of precomputeQpJacobian() can be found in the code below. The definition of 39 : * computeQpOffDiagJacobian() is still the same with the original Kokkos kernel. 40 : */ 41 : template <typename Derived> 42 : class KernelGrad : public Kernel<Derived> 43 : { 44 : usingKokkosKernelMembers(Derived); 45 : 46 : public: 47 : static InputParameters validParams(); 48 : 49 : /** 50 : * Constructor 51 : */ 52 : KernelGrad(const InputParameters & parameters); 53 : 54 : /** 55 : * Default methods to prevent compile errors even when these methods were not defined in the 56 : * derived class 57 : */ 58 : ///@{ 59 : /** 60 : * Compute diagonal Jacobian contribution on a quadrature point 61 : * @param j The trial function DOF index 62 : * @param qp The local quadrature point index 63 : * @param datum The ResidualDatum object of the current thread 64 : * @returns The vector component of the Jacobian contribution that will be multiplied by the 65 : * gradient of the test function 66 : */ 67 0 : KOKKOS_FUNCTION Real3 precomputeQpJacobian(const unsigned int /* j */, 68 : const unsigned int /* qp */, 69 : ResidualDatum & /* datum */) const 70 : { 71 0 : return Real3(0); 72 : } 73 : ///@} 74 : 75 : /** 76 : * The parallel computation bodies that hide the base class methods to optimize for factoring 77 : * out the gradient of test function 78 : */ 79 : ///@{ 80 : KOKKOS_FUNCTION void computeResidualInternal(const Derived * kernel, ResidualDatum & datum) const; 81 : KOKKOS_FUNCTION void computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const; 82 : ///@} 83 : 84 : protected: 85 : /** 86 : * Get whether precomputeQpJacobian() was not defined in the derived class 87 : * @returns Whether precomputeQpJacobian() was not defined in the derived class 88 : */ 89 10 : virtual bool defaultJacobian() const override 90 : { 91 10 : return &Derived::precomputeQpJacobian == &KernelGrad::precomputeQpJacobian; 92 : } 93 : }; 94 : 95 : template <typename Derived> 96 : InputParameters 97 9132 : KernelGrad<Derived>::validParams() 98 : { 99 9132 : InputParameters params = Kernel<Derived>::validParams(); 100 9132 : return params; 101 : } 102 : 103 : template <typename Derived> 104 12 : KernelGrad<Derived>::KernelGrad(const InputParameters & parameters) : Kernel<Derived>(parameters) 105 : { 106 12 : } 107 : 108 : template <typename Derived> 109 : KOKKOS_FUNCTION void 110 96 : KernelGrad<Derived>::computeResidualInternal(const Derived * kernel, ResidualDatum & datum) const 111 : { 112 96 : ResidualObject::computeResidualInternal( 113 : datum, 114 96 : [&](Real * local_re, const unsigned int ib, const unsigned int ie) 115 : { 116 480 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp) 117 : { 118 384 : datum.reinit(); 119 : 120 384 : Real3 value = datum.JxW(qp) * kernel->precomputeQpResidual(qp, datum); 121 : 122 1920 : for (unsigned int i = ib; i < ie; ++i) 123 1536 : local_re[i] += value * _grad_test(datum, i, qp); 124 : } 125 : }); 126 96 : } 127 : 128 : template <typename Derived> 129 : KOKKOS_FUNCTION void 130 24 : KernelGrad<Derived>::computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const 131 : { 132 24 : Real3 value; 133 : 134 24 : ResidualObject::computeJacobianInternal( 135 : datum, 136 24 : [&](Real * local_ke, const unsigned int ijb, const unsigned int ije) 137 : { 138 120 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp) 139 : { 140 96 : datum.reinit(); 141 : 142 96 : unsigned int j_old = libMesh::invalid_uint; 143 : 144 1632 : for (unsigned int ij = ijb; ij < ije; ++ij) 145 : { 146 1536 : unsigned int i = ij % datum.n_jdofs(); 147 1536 : unsigned int j = ij / datum.n_jdofs(); 148 : 149 1536 : if (j != j_old) 150 : { 151 384 : value = datum.JxW(qp) * kernel->precomputeQpJacobian(j, qp, datum); 152 384 : j_old = j; 153 : } 154 : 155 1536 : local_ke[ij] += value * _grad_test(datum, i, qp); 156 : } 157 : } 158 : }); 159 24 : } 160 : 161 : } // namespace Kokkos 162 : } // namespace Moose 163 : 164 : #define usingKokkosKernelGradMembers(T) usingKokkosKernelMembers(T)