LCOV - code coverage report
Current view: top level - include/kokkos/kernels - KokkosKernel.h (source / functions) Hit Total Coverage
Test: idaholab/moose framework: 863ef6 Lines: 51 56 91.1 %
Date: 2025-10-15 18:16:15 Functions: 59 109 54.1 %
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 user should define computeQpResidual(), computeQpJacobian(), and computeQpOffDiagJacobian()
      23             :  * as inlined public methods in their derived class (not virtual override). The signature of
      24             :  * computeQpResidual() expected to be defined in the derived class is as follows:
      25             :  *
      26             :  * @param i The element-local DOF index
      27             :  * @param qp The local quadrature point index
      28             :  * @param datum The ResidualDatum object of the current thread
      29             :  * @returns The residual contribution
      30             :  *
      31             :  * KOKKOS_FUNCTION Real computeQpResidual(const unsigned int i,
      32             :  *                                        const unsigned int qp,
      33             :  *                                        ResidualDatum & datum) const;
      34             :  *
      35             :  * The signatures of computeQpJacobian() and computeOffDiagQpJacobian() can be found in the code
      36             :  * below, and their definition in the derived class is optional. If they are defined in the derived
      37             :  * class, they will hide the default definitions in the base class.
      38             :  */
      39             : 
      40             : class Kernel : public KernelBase
      41             : {
      42             : public:
      43             :   static InputParameters validParams();
      44             : 
      45             :   /**
      46             :    * Constructor
      47             :    */
      48             :   Kernel(const InputParameters & parameters);
      49             : 
      50             :   /**
      51             :    * Dispatch residual calculation
      52             :    */
      53             :   virtual void computeResidual() override;
      54             :   /**
      55             :    * Dispatch diagonal and off-diagonal Jacobian calculation
      56             :    */
      57             :   virtual void computeJacobian() override;
      58             : 
      59             :   /**
      60             :    * Default methods to prevent compile errors even when these methods were not defined in the
      61             :    * derived class
      62             :    */
      63             :   ///@{
      64             :   /**
      65             :    * Compute diagonal Jacobian contribution on a quadrature point
      66             :    * @param i The test function DOF index
      67             :    * @param j The trial function DOF index
      68             :    * @param qp The local quadrature point index
      69             :    * @param datum The ResidualDatum object of the current thread
      70             :    * @returns The diagonal Jacobian contribution
      71             :    */
      72           0 :   KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
      73             :                                          const unsigned int /* j */,
      74             :                                          const unsigned int /* qp */,
      75             :                                          ResidualDatum & /* datum */) const
      76             :   {
      77           0 :     return 0;
      78             :   }
      79             :   /**
      80             :    * Compute off-diagonal Jacobian contribution on a quadrature point
      81             :    * @param i The test function DOF index
      82             :    * @param j The trial function DOF index
      83             :    * @param jvar The variable number for column
      84             :    * @param qp The local quadrature point index
      85             :    * @param datum The ResidualDatum object of the current thread
      86             :    * @returns The off-diagonal Jacobian contribution
      87             :    */
      88           0 :   KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
      89             :                                                 const unsigned int /* j */,
      90             :                                                 const unsigned int /* jvar */,
      91             :                                                 const unsigned int /* qp */,
      92             :                                                 ResidualDatum & /* datum */) const
      93             :   {
      94           0 :     return 0;
      95             :   }
      96             :   /**
      97             :    * Get the function pointer of the default computeQpJacobian()
      98             :    * @returns The function pointer
      99             :    */
     100      411146 :   static auto defaultJacobian() { return &Kernel::computeQpJacobian; }
     101             :   /**
     102             :    * Get the function pointer of the default computeQpOffDiagJacobian()
     103             :    * @returns The function pointer
     104             :    */
     105      485868 :   static auto defaultOffDiagJacobian() { return &Kernel::computeQpOffDiagJacobian; }
     106             :   ///@}
     107             : 
     108             :   /**
     109             :    * The parallel computation entry functions called by Kokkos
     110             :    */
     111             :   ///@{
     112             :   template <typename Derived>
     113             :   KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const;
     114             :   template <typename Derived>
     115             :   KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const;
     116             :   template <typename Derived>
     117             :   KOKKOS_FUNCTION void
     118             :   operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const;
     119             :   ///@}
     120             : 
     121             :   /**
     122             :    * The parallel computation bodies that can be customized in the derived class by defining
     123             :    * them in the derived class with the same signature.
     124             :    * Make sure to define them as inlined public methods if to be defined in the derived class.
     125             :    */
     126             :   ///@{
     127             :   /**
     128             :    * Compute residual
     129             :    * @param kernel The kernel object of the final derived type
     130             :    * @param datum The ResidualDatum object of the current thread
     131             :    */
     132             :   template <typename Derived>
     133             :   KOKKOS_FUNCTION void computeResidualInternal(const Derived & kernel, ResidualDatum & datum) const;
     134             :   /**
     135             :    * Compute diagonal Jacobian
     136             :    * @param kernel The kernel object of the final derived type
     137             :    * @param datum The ResidualDatum object of the current thread
     138             :    */
     139             :   template <typename Derived>
     140             :   KOKKOS_FUNCTION void computeJacobianInternal(const Derived & kernel, ResidualDatum & datum) const;
     141             :   /**
     142             :    * Compute off-diagonal Jacobian
     143             :    * @param kernel The kernel object of the final derived type
     144             :    * @param datum The ResidualDatum object of the current thread
     145             :    */
     146             :   template <typename Derived>
     147             :   KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived & kernel,
     148             :                                                       ResidualDatum & datum) const;
     149             :   ///@}
     150             : 
     151             : protected:
     152             :   /**
     153             :    * Current test function
     154             :    */
     155             :   const VariableTestValue _test;
     156             :   /**
     157             :    * Gradient of the current test function
     158             :    */
     159             :   const VariableTestGradient _grad_test;
     160             :   /**
     161             :    * Current shape function
     162             :    */
     163             :   const VariablePhiValue _phi;
     164             :   /**
     165             :    * Gradient of the current shape function
     166             :    */
     167             :   const VariablePhiGradient _grad_phi;
     168             :   /**
     169             :    * Current solution at quadrature points
     170             :    */
     171             :   const VariableValue _u;
     172             :   /**
     173             :    * Gradient of the current solution at quadrature points
     174             :    */
     175             :   const VariableGradient _grad_u;
     176             : };
     177             : 
     178             : template <typename Derived>
     179             : KOKKOS_FUNCTION void
     180     2449070 : Kernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const
     181             : {
     182     2449070 :   auto elem = kokkosBlockElementID(tid);
     183             : 
     184     4898140 :   ResidualDatum datum(elem,
     185             :                       libMesh::invalid_uint,
     186             :                       kokkosAssembly(),
     187             :                       kokkosSystems(),
     188     2449070 :                       _kokkos_var,
     189             :                       _kokkos_var.var());
     190             : 
     191     2449070 :   kernel.computeResidualInternal(kernel, datum);
     192     2449070 : }
     193             : 
     194             : template <typename Derived>
     195             : KOKKOS_FUNCTION void
     196      421696 : Kernel::operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const
     197             : {
     198      421696 :   auto elem = kokkosBlockElementID(tid);
     199             : 
     200      843392 :   ResidualDatum datum(elem,
     201             :                       libMesh::invalid_uint,
     202             :                       kokkosAssembly(),
     203             :                       kokkosSystems(),
     204      421696 :                       _kokkos_var,
     205             :                       _kokkos_var.var());
     206             : 
     207      421696 :   kernel.computeJacobianInternal(kernel, datum);
     208      421696 : }
     209             : 
     210             : template <typename Derived>
     211             : KOKKOS_FUNCTION void
     212        1250 : Kernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const
     213             : {
     214        1250 :   auto elem = kokkosBlockElementID(_thread(tid, 1));
     215             : 
     216        1250 :   auto & sys = kokkosSystem(_kokkos_var.sys());
     217        1250 :   auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];
     218             : 
     219        1250 :   if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
     220           0 :     return;
     221             : 
     222        1250 :   ResidualDatum datum(
     223        1250 :       elem, libMesh::invalid_uint, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);
     224             : 
     225        1250 :   kernel.computeOffDiagJacobianInternal(kernel, datum);
     226             : }
     227             : 
     228             : template <typename Derived>
     229             : KOKKOS_FUNCTION void
     230     1454133 : Kernel::computeResidualInternal(const Derived & kernel, ResidualDatum & datum) const
     231             : {
     232     1454133 :   ResidualObject::computeResidualInternal(
     233             :       datum,
     234     2908266 :       [&](Real * local_re, const unsigned int ib, const unsigned int ie)
     235             :       {
     236     7118833 :         for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     237             :         {
     238     5664700 :           datum.reinit();
     239             : 
     240    29591164 :           for (unsigned int i = ib; i < ie; ++i)
     241    23926464 :             local_re[i] += datum.JxW(qp) * kernel.computeQpResidual(i, qp, datum);
     242             :         }
     243             :       });
     244     1454133 : }
     245             : 
     246             : template <typename Derived>
     247             : KOKKOS_FUNCTION void
     248      229598 : Kernel::computeJacobianInternal(const Derived & kernel, ResidualDatum & datum) const
     249             : {
     250      229598 :   ResidualObject::computeJacobianInternal(
     251             :       datum,
     252      506812 :       [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
     253             :       {
     254     1258582 :         for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     255             :         {
     256     1005176 :           datum.reinit();
     257             : 
     258    16796152 :           for (unsigned int ij = ijb; ij < ije; ++ij)
     259             :           {
     260    15790976 :             unsigned int i = ij % datum.n_jdofs();
     261    15790976 :             unsigned int j = ij / datum.n_jdofs();
     262             : 
     263    15790976 :             local_ke[ij] += datum.JxW(qp) * kernel.computeQpJacobian(i, j, qp, datum);
     264             :           }
     265             :         }
     266             :       });
     267      229598 : }
     268             : 
     269             : template <typename Derived>
     270             : KOKKOS_FUNCTION void
     271        1250 : Kernel::computeOffDiagJacobianInternal(const Derived & kernel, ResidualDatum & datum) const
     272             : {
     273        1250 :   ResidualObject::computeJacobianInternal(
     274             :       datum,
     275        2500 :       [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
     276             :       {
     277        6250 :         for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     278             :         {
     279        5000 :           datum.reinit();
     280             : 
     281       85000 :           for (unsigned int ij = ijb; ij < ije; ++ij)
     282             :           {
     283       80000 :             unsigned int i = ij % datum.n_jdofs();
     284       80000 :             unsigned int j = ij / datum.n_jdofs();
     285             : 
     286       80000 :             local_ke[ij] +=
     287       80000 :                 datum.JxW(qp) * kernel.computeQpOffDiagJacobian(i, j, datum.jvar(), qp, datum);
     288             :           }
     289             :         }
     290             :       });
     291        1250 : }
     292             : 
     293             : } // namespace Kokkos
     294             : } // namespace Moose

Generated by: LCOV version 1.14