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::Kokkos
      15             : {
      16             : 
      17             : /**
      18             :  * The base class for a user to derive their own Kokkos kernels where the residual is of the form
      19             :  *
      20             :  * $(\dots, \psi_i)$
      21             :  *
      22             :  * i.e. the test function $(\psi_i)$ can be factored out for optimization.
      23             :  *
      24             :  * The user should now define precomputeQpResidual(), precomputeQpJacobian(), and
      25             :  * precomputeQpOffDiagJacobian() which do not include the test function, instead of the original
      26             :  * hooks. The signature of precomputeQpResidual() expected to be defined in the derived class is as
      27             :  * follows:
      28             :  *
      29             :  * @tparam Derived The object type
      30             :  * @param qp The local quadrature point index
      31             :  * @param datum The AssemblyDatum object of the current thread
      32             :  * @returns The component of the residual contribution that will be multiplied by the test function
      33             :  *
      34             :  * template <typename Derived>
      35             :  * KOKKOS_FUNCTION Real precomputeQpResidual(const unsigned int qp,
      36             :  *                                           AssemblyDatum & datum) const;
      37             :  *
      38             :  * The signature of precomputeQpJacobian() and precomputeQpOffDiagJacobian() can be found in the
      39             :  * code below.
      40             :  */
      41             : class KernelValue : public Kernel
      42             : {
      43             : public:
      44             :   static InputParameters validParams();
      45             : 
      46             :   /// KernelValue hooks factor out the test function
      47             :   static constexpr bool use_precompute_hooks = true;
      48             : 
      49             :   /**
      50             :    * Constructor
      51             :    */
      52             :   KernelValue(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             :    * @tparam Derived The object type
      62             :    * @param j The trial function DOF index
      63             :    * @param qp The local quadrature point index
      64             :    * @param datum The AssemblyDatum object of the current thread
      65             :    * @returns The component of the Jacobian contribution that will be multiplied by the test
      66             :    * function
      67             :    */
      68             :   template <typename Derived>
      69           0 :   KOKKOS_FUNCTION Real precomputeQpJacobian(const unsigned int /* j */,
      70             :                                             const unsigned int /* qp */,
      71             :                                             AssemblyDatum & /* datum */) const
      72             :   {
      73           0 :     ::Kokkos::abort("Default precomputeQpJacobian() should never be called. Make sure you properly "
      74             :                     "redefined this method in your class without typos.");
      75             : 
      76             :     return 0;
      77             :   }
      78             :   /**
      79             :    * Compute off-diagonal Jacobian contribution on a quadrature point
      80             :    * @tparam Derived The object type
      81             :    * @param j The trial function DOF index
      82             :    * @param jvar The variable number for column
      83             :    * @param qp The local quadrature point index
      84             :    * @param datum The AssemblyDatum object of the current thread
      85             :    * @returns The component of the off-diagonal Jacobian contribution that will be multiplied by the
      86             :    * test function
      87             :    */
      88             :   template <typename Derived>
      89           0 :   KOKKOS_FUNCTION Real precomputeQpOffDiagJacobian(const unsigned int /* j */,
      90             :                                                    const unsigned int /* jvar */,
      91             :                                                    const unsigned int /* qp */,
      92             :                                                    AssemblyDatum & /* datum */) const
      93             :   {
      94           0 :     ::Kokkos::abort(
      95             :         "Default precomputeQpOffDiagJacobian() should never be called. Make sure you properly "
      96             :         "redefined this method in your class without typos.");
      97             : 
      98             :     return 0;
      99             :   }
     100             :   ///@}
     101             : 
     102             :   /**
     103             :    * Functions used to check if users have overriden the hook methods, whose calculations can be
     104             :    * skipped when not overriden
     105             :    * @returns The function pointer of the default hook method
     106             :    */
     107             :   ///@{
     108             :   template <typename Derived>
     109      289878 :   static auto defaultJacobian()
     110             :   {
     111      289878 :     return &KernelValue::precomputeQpJacobian<Derived>;
     112             :   }
     113             :   template <typename Derived>
     114      289878 :   static auto defaultOffDiagJacobian()
     115             :   {
     116      289878 :     return &KernelValue::precomputeQpOffDiagJacobian<Derived>;
     117             :   }
     118             :   ///@}
     119             : 
     120             :   /**
     121             :    * The parallel computation bodies that hide the base class methods to optimize for factoring
     122             :    * out the test function
     123             :    */
     124             :   ///@{
     125             :   template <typename Derived>
     126             :   KOKKOS_FUNCTION void computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const;
     127             :   template <typename Derived>
     128             :   KOKKOS_FUNCTION void computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const;
     129             :   template <typename Derived>
     130             :   KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived & kernel,
     131             :                                                       AssemblyDatum & datum) const;
     132             :   ///@}
     133             : };
     134             : 
     135             : template <typename Derived>
     136             : KOKKOS_FUNCTION void
     137     1885579 : KernelValue::computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const
     138             : {
     139     1885579 :   ResidualObject::computeResidualInternal(
     140             :       datum,
     141     3771158 :       [&](Real * local_re, const unsigned int ib, const unsigned int ie)
     142             :       {
     143    10073903 :         for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     144             :         {
     145     8188324 :           Real value = datum.JxW(qp) * kernel.template precomputeQpResidual<Derived>(qp, datum);
     146             : 
     147    42286964 :           for (unsigned int i = ib; i < ie; ++i)
     148    34098640 :             local_re[i] += value * _test(datum, i, qp);
     149             :         }
     150             :       });
     151     1885579 : }
     152             : 
     153             : template <typename Derived>
     154             : KOKKOS_FUNCTION void
     155          24 : KernelValue::computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
     156             : {
     157          24 :   ResidualObject::computeJacobianInternal(
     158             :       datum,
     159         192 :       [&](Real * local_ke, const unsigned int ib, const unsigned int ie, const unsigned int j)
     160             :       {
     161         480 :         for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     162             :         {
     163         384 :           Real value = datum.JxW(qp) * kernel.template precomputeQpJacobian<Derived>(j, qp, datum);
     164             : 
     165        1920 :           for (unsigned int i = ib; i < ie; ++i)
     166        1536 :             local_ke[i] += value * _test(datum, i, qp);
     167             :         }
     168             :       });
     169          24 : }
     170             : 
     171             : template <typename Derived>
     172             : KOKKOS_FUNCTION void
     173        1250 : KernelValue::computeOffDiagJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
     174             : {
     175        1250 :   ResidualObject::computeJacobianInternal(
     176             :       datum,
     177       10000 :       [&](Real * local_ke, const unsigned int ib, const unsigned int ie, const unsigned int j)
     178             :       {
     179       25000 :         for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     180             :         {
     181       20000 :           Real value = datum.JxW(qp) * kernel.template precomputeQpOffDiagJacobian<Derived>(
     182             :                                            j, datum.jvar(), qp, datum);
     183             : 
     184      100000 :           for (unsigned int i = ib; i < ie; ++i)
     185       80000 :             local_ke[i] += value * _test(datum, i, qp);
     186             :         }
     187             :       });
     188        1250 : }
     189             : 
     190             : } // namespace Moose::Kokkos