LCOV - code coverage report
Current view: top level - include/kokkos/kernels - KokkosKernel.h (source / functions) Hit Total Coverage
Test: idaholab/moose framework: 6f668f Lines: 80 85 94.1 %
Date: 2025-09-22 20:01:15 Functions: 133 207 64.3 %
Legend: Lines: hit not hit

          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 "KokkosKernelBase.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.
      21             :  *
      22             :  * The polymorphic design of the original MOOSE is reproduced statically by leveraging the Curiously
      23             :  * Recurring Template Pattern (CRTP), a programming idiom that involves a class template inheriting
      24             :  * from a template instantiation of itself. When the user derives their Kokkos object from this
      25             :  * class, the inheritance structure will look like:
      26             :  *
      27             :  * class UserKernel final : public Moose::Kokkos::Kernel<UserKernel>
      28             :  *
      29             :  * It is important to note that the template argument should point to the last derived class.
      30             :  * Therefore, if the user wants to define a derived class that can be further inherited, the derived
      31             :  * class should be a class template as well. Otherwise, it is recommended to mark the derived class
      32             :  * as final to prevent its inheritence by mistake.
      33             :  *
      34             :  * The user is expected to define computeQpResidual(), computeQpJacobian(), and
      35             :  * computeQpOffDiagJacobian() as inlined public methods in their derived class (not virtual
      36             :  * override). The signature of computeQpResidual() expected to be defined in the derived class is as
      37             :  * follows:
      38             :  *
      39             :  * @param i The element-local DOF index
      40             :  * @param qp The local quadrature point index
      41             :  * @param datum The ResidualDatum object of the current thread
      42             :  * @returns The residual contribution
      43             :  *
      44             :  * KOKKOS_FUNCTION Real computeQpResidual(const unsigned int i,
      45             :  *                                        const unsigned int qp,
      46             :  *                                        ResidualDatum & datum) const;
      47             :  *
      48             :  * The signatures of computeQpJacobian() and computeOffDiagQpJacobian() can be found in the code
      49             :  * below, and their definition in the derived class is optional. If they are defined in the derived
      50             :  * class, they will hide the default definitions in the base class.
      51             :  */
      52             : template <typename Derived>
      53             : class Kernel : public KernelBase
      54             : {
      55             : public:
      56             :   static InputParameters validParams();
      57             : 
      58             :   /**
      59             :    * Constructor
      60             :    */
      61             :   Kernel(const InputParameters & parameters);
      62             : 
      63             :   /**
      64             :    * Dispatch residual calculation
      65             :    */
      66             :   virtual void computeResidual() override;
      67             :   /**
      68             :    * Dispatch diagonal and off-diagonal Jacobian calculation
      69             :    */
      70             :   virtual void computeJacobian() override;
      71             : 
      72             :   /**
      73             :    * Default methods to prevent compile errors even when these methods were not defined in the
      74             :    * derived class
      75             :    */
      76             :   ///@{
      77             :   /**
      78             :    * Compute diagonal Jacobian contribution on a quadrature point
      79             :    * @param i The test function DOF index
      80             :    * @param j The trial function DOF index
      81             :    * @param qp The local quadrature point index
      82             :    * @param datum The ResidualDatum object of the current thread
      83             :    * @returns The diagonal Jacobian contribution
      84             :    */
      85           0 :   KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
      86             :                                          const unsigned int /* j */,
      87             :                                          const unsigned int /* qp */,
      88             :                                          ResidualDatum & /* datum */) const
      89             :   {
      90           0 :     return 0;
      91             :   }
      92             :   /**
      93             :    * Compute off-diagonal Jacobian contribution on a quadrature point
      94             :    * @param i The test function DOF index
      95             :    * @param j The trial function DOF index
      96             :    * @param jvar The variable number for column
      97             :    * @param qp The local quadrature point index
      98             :    * @param datum The ResidualDatum object of the current thread
      99             :    * @returns The off-diagonal Jacobian contribution
     100             :    */
     101           0 :   KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
     102             :                                                 const unsigned int /* j */,
     103             :                                                 const unsigned int /* jvar */,
     104             :                                                 const unsigned int /* qp */,
     105             :                                                 ResidualDatum & /* datum */) const
     106             :   {
     107           0 :     return 0;
     108             :   }
     109             :   ///@}
     110             : 
     111             :   /**
     112             :    * The parallel computation entry functions called by Kokkos
     113             :    */
     114             :   ///@{
     115             :   KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid) const;
     116             :   KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid) const;
     117             :   KOKKOS_FUNCTION void operator()(OffDiagJacobianLoop, const ThreadID tid) const;
     118             :   ///@}
     119             : 
     120             :   /**
     121             :    * The parallel computation bodies that can be customized in the derived class by defining
     122             :    * them in the derived class with the same signature.
     123             :    * Make sure to define them as inlined public methods if to be defined in the derived class.
     124             :    */
     125             :   ///@{
     126             :   /**
     127             :    * Compute residual
     128             :    * @param kernel The kernel object of the final derived type
     129             :    * @param datum The ResidualDatum object of the current thread
     130             :    */
     131             :   KOKKOS_FUNCTION void computeResidualInternal(const Derived * kernel, ResidualDatum & datum) const;
     132             :   /**
     133             :    * Compute diagonal Jacobian
     134             :    * @param kernel The kernel object of the final derived type
     135             :    * @param datum The ResidualDatum object of the current thread
     136             :    */
     137             :   KOKKOS_FUNCTION void computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const;
     138             :   /**
     139             :    * Compute off-diagonal Jacobian
     140             :    * @param kernel The kernel object of the final derived type
     141             :    * @param datum The ResidualDatum object of the current thread
     142             :    */
     143             :   KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived * kernel,
     144             :                                                       ResidualDatum & datum) const;
     145             :   ///@}
     146             : 
     147             : protected:
     148             :   /**
     149             :    * Current test function
     150             :    */
     151             :   const VariableTestValue _test;
     152             :   /**
     153             :    * Gradient of the current test function
     154             :    */
     155             :   const VariableTestGradient _grad_test;
     156             :   /**
     157             :    * Current shape function
     158             :    */
     159             :   const VariablePhiValue _phi;
     160             :   /**
     161             :    * Gradient of the current shape function
     162             :    */
     163             :   const VariablePhiGradient _grad_phi;
     164             :   /**
     165             :    * Current solution at quadrature points
     166             :    */
     167             :   const VariableValue _u;
     168             :   /**
     169             :    * Gradient of the current solution at quadrature points
     170             :    */
     171             :   const VariableGradient _grad_u;
     172             : 
     173             :   /**
     174             :    * Get whether computeQpJacobian() was not defined in the derived class
     175             :    * @returns Whether computeQpJacobian() was not defined in the derived class
     176             :    */
     177        9945 :   virtual bool defaultJacobian() const
     178             :   {
     179        9945 :     return &Derived::computeQpJacobian == &Kernel::computeQpJacobian;
     180             :   }
     181             :   /**
     182             :    * Get whether computeQpOffDiagJacobian() was not defined in the derived class
     183             :    * @returns Whether computeQpOffDiagJacobian() was not defined in the derived class
     184             :    */
     185        9965 :   virtual bool defaultOffDiagJacobian() const
     186             :   {
     187        9965 :     return &Derived::computeQpOffDiagJacobian == &Kernel::computeQpOffDiagJacobian;
     188             :   }
     189             : };
     190             : 
     191             : template <typename Derived>
     192             : InputParameters
     193      121022 : Kernel<Derived>::validParams()
     194             : {
     195      121022 :   InputParameters params = KernelBase::validParams();
     196      121022 :   return params;
     197             : }
     198             : 
     199             : template <typename Derived>
     200        1299 : Kernel<Derived>::Kernel(const InputParameters & parameters)
     201             :   : KernelBase(parameters, Moose::VarFieldType::VAR_FIELD_STANDARD),
     202             :     _test(),
     203             :     _grad_test(),
     204             :     _phi(),
     205             :     _grad_phi(),
     206         994 :     _u(_var),
     207         994 :     _grad_u(_var)
     208             : {
     209        1299 :   addMooseVariableDependency(&_var);
     210        1299 : }
     211             : 
     212             : template <typename Derived>
     213             : void
     214       54261 : Kernel<Derived>::computeResidual()
     215             : {
     216       54261 :   ::Kokkos::RangePolicy<ResidualLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
     217             :       0, numKokkosBlockElements());
     218       54261 :   ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
     219       54261 :   ::Kokkos::fence();
     220       54261 : }
     221             : 
     222             : template <typename Derived>
     223             : void
     224        9965 : Kernel<Derived>::computeJacobian()
     225             : {
     226        9965 :   if (!defaultJacobian())
     227             :   {
     228        8234 :     ::Kokkos::RangePolicy<JacobianLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
     229             :         0, numKokkosBlockElements());
     230        8234 :     ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
     231        8234 :     ::Kokkos::fence();
     232        8234 :   }
     233             : 
     234        9965 :   if (!defaultOffDiagJacobian())
     235             :   {
     236          48 :     auto & sys = kokkosSystem(_kokkos_var.sys());
     237             : 
     238          96 :     _thread.resize({sys.getCoupling(_kokkos_var.var()).size(), numKokkosBlockElements()});
     239             : 
     240          48 :     ::Kokkos::RangePolicy<OffDiagJacobianLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
     241             :         0, _thread.size());
     242          48 :     ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
     243          48 :     ::Kokkos::fence();
     244          48 :   }
     245        9965 : }
     246             : 
     247             : template <typename Derived>
     248             : KOKKOS_FUNCTION void
     249     2449070 : Kernel<Derived>::operator()(ResidualLoop, const ThreadID tid) const
     250             : {
     251     2449070 :   auto kernel = static_cast<const Derived *>(this);
     252     2449070 :   auto elem = kokkosBlockElementID(tid);
     253             : 
     254     2449070 :   ResidualDatum datum(elem, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());
     255             : 
     256     2449070 :   kernel->computeResidualInternal(kernel, datum);
     257     2449070 : }
     258             : 
     259             : template <typename Derived>
     260             : KOKKOS_FUNCTION void
     261      421696 : Kernel<Derived>::operator()(JacobianLoop, const ThreadID tid) const
     262             : {
     263      421696 :   auto kernel = static_cast<const Derived *>(this);
     264      421696 :   auto elem = kokkosBlockElementID(tid);
     265             : 
     266      421696 :   ResidualDatum datum(elem, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());
     267             : 
     268      421696 :   kernel->computeJacobianInternal(kernel, datum);
     269      421696 : }
     270             : 
     271             : template <typename Derived>
     272             : KOKKOS_FUNCTION void
     273        1250 : Kernel<Derived>::operator()(OffDiagJacobianLoop, const ThreadID tid) const
     274             : {
     275        1250 :   auto kernel = static_cast<const Derived *>(this);
     276        1250 :   auto elem = kokkosBlockElementID(_thread(tid, 1));
     277             : 
     278        1250 :   auto & sys = kokkosSystem(_kokkos_var.sys());
     279        1250 :   auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];
     280             : 
     281        1250 :   if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
     282           0 :     return;
     283             : 
     284        1250 :   ResidualDatum datum(elem, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);
     285             : 
     286        1250 :   kernel->computeOffDiagJacobianInternal(kernel, datum);
     287             : }
     288             : 
     289             : template <typename Derived>
     290             : KOKKOS_FUNCTION void
     291     1454133 : Kernel<Derived>::computeResidualInternal(const Derived * kernel, ResidualDatum & datum) const
     292             : {
     293     1454133 :   ResidualObject::computeResidualInternal(
     294             :       datum,
     295     2908266 :       [&](Real * local_re, const unsigned int ib, const unsigned int ie)
     296             :       {
     297     7118833 :         for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     298             :         {
     299     5664700 :           datum.reinit();
     300             : 
     301    29591164 :           for (unsigned int i = ib; i < ie; ++i)
     302    23926464 :             local_re[i] += datum.JxW(qp) * kernel->computeQpResidual(i, qp, datum);
     303             :         }
     304             :       });
     305     1454133 : }
     306             : 
     307             : template <typename Derived>
     308             : KOKKOS_FUNCTION void
     309      229598 : Kernel<Derived>::computeJacobianInternal(const Derived * kernel, ResidualDatum & datum) const
     310             : {
     311      229598 :   ResidualObject::computeJacobianInternal(
     312             :       datum,
     313      506812 :       [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
     314             :       {
     315     1258582 :         for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     316             :         {
     317     1005176 :           datum.reinit();
     318             : 
     319    16796152 :           for (unsigned int ij = ijb; ij < ije; ++ij)
     320             :           {
     321    15790976 :             unsigned int i = ij % datum.n_jdofs();
     322    15790976 :             unsigned int j = ij / datum.n_jdofs();
     323             : 
     324    15790976 :             local_ke[ij] += datum.JxW(qp) * kernel->computeQpJacobian(i, j, qp, datum);
     325             :           }
     326             :         }
     327             :       });
     328      229598 : }
     329             : 
     330             : template <typename Derived>
     331             : KOKKOS_FUNCTION void
     332        1250 : Kernel<Derived>::computeOffDiagJacobianInternal(const Derived * kernel, ResidualDatum & datum) const
     333             : {
     334        1250 :   ResidualObject::computeJacobianInternal(
     335             :       datum,
     336        2500 :       [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
     337             :       {
     338        6250 :         for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     339             :         {
     340        5000 :           datum.reinit();
     341             : 
     342       85000 :           for (unsigned int ij = ijb; ij < ije; ++ij)
     343             :           {
     344       80000 :             unsigned int i = ij % datum.n_jdofs();
     345       80000 :             unsigned int j = ij / datum.n_jdofs();
     346             : 
     347       80000 :             local_ke[ij] +=
     348       80000 :                 datum.JxW(qp) * kernel->computeQpOffDiagJacobian(i, j, datum.jvar(), qp, datum);
     349             :           }
     350             :         }
     351             :       });
     352        1250 : }
     353             : 
     354             : } // namespace Kokkos
     355             : } // namespace Moose
     356             : 
     357             : #define usingKokkosKernelMembers(T)                                                                \
     358             :   usingKokkosKernelBaseMembers;                                                                    \
     359             :                                                                                                    \
     360             : protected:                                                                                         \
     361             :   using Moose::Kokkos::Kernel<T>::_test;                                                           \
     362             :   using Moose::Kokkos::Kernel<T>::_grad_test;                                                      \
     363             :   using Moose::Kokkos::Kernel<T>::_phi;                                                            \
     364             :   using Moose::Kokkos::Kernel<T>::_grad_phi;                                                       \
     365             :   using Moose::Kokkos::Kernel<T>::_u;                                                              \
     366             :   using Moose::Kokkos::Kernel<T>::_grad_u;                                                         \
     367             :                                                                                                    \
     368             : public:                                                                                            \
     369             :   using Moose::Kokkos::Kernel<T>::operator()

Generated by: LCOV version 1.14