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, \psi_i)$ 23 : * 24 : * i.e. the test function $(\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 component of the residual contribution that will be multiplied by the test function 33 : * 34 : * KOKKOS_FUNCTION Real precomputeQpResidual(const unsigned int qp, 35 : * ResidualDatum & datum) const; 36 : * 37 : * The signature of precomputeQpJacobian() can be found in the code below. The definition of 38 : * computeQpOffDiagJacobian() is still the same with the original Kokkos kernel. 39 : */ 40 : template <typename Derived> 41 : class KernelValue : public Kernel<Derived> 42 : { 43 : usingKokkosKernelMembers(Derived); 44 : 45 : public: 46 : static InputParameters validParams(); 47 : 48 : /** 49 : * Constructor 50 : */ 51 : KernelValue(const InputParameters & parameters); 52 : 53 : /** 54 : * Default methods to prevent compile errors even when these methods were not defined in the 55 : * derived class 56 : */ 57 : ///@{ 58 : /** 59 : * Compute diagonal Jacobian contribution on a quadrature point 60 : * @param j The trial function DOF index 61 : * @param qp The local quadrature point index 62 : * @param datum The ResidualDatum object of the current thread 63 : * @returns The component of the Jacobian contribution that will be multiplied by the test 64 : * function 65 : */ 66 0 : KOKKOS_FUNCTION Real precomputeQpJacobian(const unsigned int /* j */, 67 : const unsigned int /* qp */, 68 : ResidualDatum & /* datum */) const 69 : { 70 0 : return 0; 71 : } 72 : ///@} 73 : 74 : /** 75 : * The parallel computation bodies that hide the base class methods to optimize for factoring 76 : * out the test function 77 : */ 78 : ///@{ 79 : KOKKOS_FUNCTION void computeResidualInternal(const Derived * kernel, ResidualDatum & datum) const; 80 : KOKKOS_FUNCTION void computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const; 81 : ///@} 82 : 83 : protected: 84 : /** 85 : * Get whether precomputeQpJacobian() was not defined in the derived class 86 : * @returns Whether precomputeQpJacobian() was not defined in the derived class 87 : */ 88 10 : virtual bool defaultJacobian() const override 89 : { 90 10 : return &Derived::precomputeQpJacobian == &KernelValue::precomputeQpJacobian; 91 : } 92 : }; 93 : 94 : template <typename Derived> 95 : InputParameters 96 9132 : KernelValue<Derived>::validParams() 97 : { 98 9132 : InputParameters params = Kernel<Derived>::validParams(); 99 9132 : return params; 100 : } 101 : 102 : template <typename Derived> 103 12 : KernelValue<Derived>::KernelValue(const InputParameters & parameters) : Kernel<Derived>(parameters) 104 : { 105 12 : } 106 : 107 : template <typename Derived> 108 : KOKKOS_FUNCTION void 109 96 : KernelValue<Derived>::computeResidualInternal(const Derived * kernel, ResidualDatum & datum) const 110 : { 111 96 : ResidualObject::computeResidualInternal( 112 : datum, 113 192 : [&](Real * local_re, const unsigned int ib, const unsigned int ie) 114 : { 115 480 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp) 116 : { 117 384 : datum.reinit(); 118 : 119 384 : Real value = datum.JxW(qp) * kernel->precomputeQpResidual(qp, datum); 120 : 121 1920 : for (unsigned int i = ib; i < ie; ++i) 122 1536 : local_re[i] += value * _test(datum, i, qp); 123 : } 124 : }); 125 96 : } 126 : 127 : template <typename Derived> 128 : KOKKOS_FUNCTION void 129 24 : KernelValue<Derived>::computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const 130 : { 131 24 : Real value = 0; 132 : 133 24 : ResidualObject::computeJacobianInternal( 134 : datum, 135 48 : [&](Real * local_ke, const unsigned int ijb, const unsigned int ije) 136 : { 137 120 : for (unsigned int qp = 0; qp < datum.n_qps(); ++qp) 138 : { 139 96 : datum.reinit(); 140 : 141 96 : unsigned int j_old = libMesh::invalid_uint; 142 : 143 1632 : for (unsigned int ij = ijb; ij < ije; ++ij) 144 : { 145 1536 : unsigned int i = ij % datum.n_jdofs(); 146 1536 : unsigned int j = ij / datum.n_jdofs(); 147 : 148 1536 : if (j != j_old) 149 : { 150 384 : value = datum.JxW(qp) * kernel->precomputeQpJacobian(j, qp, datum); 151 384 : j_old = j; 152 : } 153 : 154 1536 : local_ke[ij] += value * _test(datum, i, qp); 155 : } 156 : } 157 : }); 158 24 : } 159 : 160 : } // namespace Kokkos 161 : } // namespace Moose 162 : 163 : #define usingKokkosKernelValueMembers(T) usingKokkosKernelMembers(T)