LCOV - code coverage report
Current view: top level - include/kokkos/auxkernels - KokkosAuxKernel.h (source / functions) Hit Total Coverage
Test: idaholab/moose framework: #31706 (f8ed4a) with base bb0a08 Lines: 52 53 98.1 %
Date: 2025-11-03 17:23:24 Functions: 15 26 57.7 %
Legend: Lines: hit not hit

          Line data    Source code
       1             : //* This file is part of the MOOSE framework
       2             : //* https://mooseframework.inl.gov
       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 "KokkosDispatcher.h"
      13             : #include "KokkosVariableValue.h"
      14             : #include "KokkosMaterialPropertyValue.h"
      15             : #include "KokkosAssembly.h"
      16             : #include "KokkosSystem.h"
      17             : 
      18             : #include "AuxKernelBase.h"
      19             : 
      20             : namespace Moose
      21             : {
      22             : namespace Kokkos
      23             : {
      24             : 
      25             : /**
      26             :  * The base class for a user to derive their own Kokkos auxiliary kernels.
      27             :  *
      28             :  * The user should define computeValue() as inlined public method in their derived class (not
      29             :  * virtual override). The signature of computeValue() expected to be defined in the derived
      30             :  * class is as follows:
      31             :  *
      32             :  * @param qp The local quadrature point index
      33             :  * @param datum The AssemblyDatum object of the current thread
      34             :  * @returns The value at the quadrature point
      35             :  *
      36             :  * KOKKOS_FUNCTION Real computeValue(const unsigned int qp, AssemblyDatum & datum) const;
      37             :  */
      38             : class AuxKernel : public ::AuxKernelBase,
      39             :                   public MeshHolder,
      40             :                   public AssemblyHolder,
      41             :                   public SystemHolder
      42             : {
      43             : public:
      44             :   static InputParameters validParams();
      45             : 
      46             :   /**
      47             :    * Constructor
      48             :    */
      49             :   AuxKernel(const InputParameters & parameters);
      50             : 
      51             :   /**
      52             :    * Copy constructor for parallel dispatch
      53             :    */
      54             :   AuxKernel(const AuxKernel & object);
      55             : 
      56             :   // Unused for Kokkos auxiliary kernels because all elements are computed in parallel
      57           0 :   virtual void subdomainSetup() override final {}
      58             : 
      59             :   /**
      60             :    * Dispatch calculation
      61             :    */
      62             :   virtual void compute() override;
      63             : 
      64             :   /**
      65             :    * Get whether this auxiliary kernel is nodal
      66             :    * @returns Whether this auxiliary kernel is nodal
      67             :    */
      68         488 :   KOKKOS_FUNCTION bool isNodal() const { return _nodal; }
      69             : 
      70             :   /**
      71             :    * Kokkos function tags
      72             :    */
      73             :   ///@{
      74             :   struct ElementLoop
      75             :   {
      76             :   };
      77             :   struct NodeLoop
      78             :   {
      79             :   };
      80             :   ///@}
      81             : 
      82             :   /**
      83             :    * The parallel computation entry functions called by Kokkos
      84             :    */
      85             :   ///@{
      86             :   template <typename Derived>
      87             :   KOKKOS_FUNCTION void operator()(ElementLoop, const ThreadID tid, const Derived & auxkernel) const;
      88             :   template <typename Derived>
      89             :   KOKKOS_FUNCTION void operator()(NodeLoop, const ThreadID tid, const Derived & auxkernel) const;
      90             :   ///@}
      91             : 
      92             :   /**
      93             :    * The parallel computation bodies that can be customized in the derived class by defining
      94             :    * them in the derived class with the same signature.
      95             :    * Make sure to define them as inlined public methods if to be defined in the derived class.
      96             :    */
      97             :   ///@{
      98             :   /**
      99             :    * Compute an element
     100             :    * @param kernel The auxiliary kernel object of the final derived type
     101             :    * @param datum The AssemblyDatum object of the current thread
     102             :    */
     103             :   template <typename Derived>
     104             :   KOKKOS_FUNCTION void computeElementInternal(const Derived & auxkernel,
     105             :                                               AssemblyDatum & datum) const;
     106             :   /**
     107             :    * Compute a node
     108             :    * @param kernel The auxiliary kernel object of the final derived type
     109             :    * @param datum The AssemblyDatum object of the current thread
     110             :    */
     111             :   template <typename Derived>
     112             :   KOKKOS_FUNCTION void computeNodeInternal(const Derived & auxkernel, AssemblyDatum & datum) const;
     113             :   ///@}
     114             : 
     115             : protected:
     116             :   /**
     117             :    * Retrieve the old value of the variable that this kernel operates on
     118             :    * @returns The old variable value object
     119             :    */
     120             :   VariableValue uOld() const;
     121             :   /**
     122             :    * Retrieve the older value of the variable that this kernel operates on
     123             :    * @returns The older variable value object
     124             :    */
     125             :   VariableValue uOlder() const;
     126             : 
     127             :   /**
     128             :    * Set element values to the auxiliary solution vector
     129             :    * @param values The array containing the solution values of the element
     130             :    * @param datum The AssemblyDatum object of the current thread
     131             :    * @param comp The variable component
     132             :    */
     133             :   KOKKOS_FUNCTION void setElementSolution(const Real * const values,
     134             :                                           const AssemblyDatum & datum,
     135             :                                           const unsigned int comp = 0) const;
     136             :   /**
     137             :    * Set node value to the auxiliary solution vector
     138             :    * @param values The node solution value
     139             :    * @param datum The AssemblyDatum object of the current thread
     140             :    * @param comp The variable component
     141             :    */
     142             :   KOKKOS_FUNCTION void
     143             :   setNodeSolution(const Real value, const AssemblyDatum & datum, const unsigned int comp = 0) const;
     144             : 
     145             :   /**
     146             :    * Flag whether this kernel is nodal
     147             :    */
     148             :   const bool _nodal;
     149             : 
     150             :   /**
     151             :    * Kokkos variable
     152             :    */
     153             :   Variable _kokkos_var;
     154             :   /**
     155             :    * Kokkos functor dispatchers
     156             :    */
     157             :   ///@{
     158             :   std::unique_ptr<DispatcherBase> _element_dispatcher;
     159             :   std::unique_ptr<DispatcherBase> _node_dispatcher;
     160             :   ///@}
     161             : 
     162             :   /**
     163             :    * Current test function
     164             :    */
     165             :   const VariableTestValue _test;
     166             :   /**
     167             :    * Current solution
     168             :    */
     169             :   const VariableValue _u;
     170             : 
     171             :   /**
     172             :    * TODO: Move to TransientInterface
     173             :    */
     174             :   ///@{
     175             :   /**
     176             :    * Time
     177             :    */
     178             :   Scalar<Real> _t;
     179             :   /**
     180             :    * Old time
     181             :    */
     182             :   Scalar<const Real> _t_old;
     183             :   /**
     184             :    * The number of the time step
     185             :    */
     186             :   Scalar<int> _t_step;
     187             :   /**
     188             :    * Time step size
     189             :    */
     190             :   Scalar<Real> _dt;
     191             :   /**
     192             :    * Size of the old time step
     193             :    */
     194             :   Scalar<Real> _dt_old;
     195             :   ///@}
     196             : 
     197             : private:
     198             :   /**
     199             :    * Override of the MaterialPropertyInterface function to throw on material property request for
     200             :    * nodal kernels
     201             :    */
     202             :   void getKokkosMaterialPropertyHook(const std::string & prop_name_in,
     203             :                                      const unsigned int state) override final;
     204             : };
     205             : 
     206             : template <typename Derived>
     207             : KOKKOS_FUNCTION void
     208      210176 : AuxKernel::operator()(ElementLoop, const ThreadID tid, const Derived & auxkernel) const
     209             : {
     210      210176 :   auto elem = kokkosBlockElementID(tid);
     211             : 
     212      420352 :   AssemblyDatum datum(elem,
     213             :                       libMesh::invalid_uint,
     214             :                       kokkosAssembly(),
     215             :                       kokkosSystems(),
     216      210176 :                       _kokkos_var,
     217             :                       _kokkos_var.var());
     218             : 
     219      210176 :   auxkernel.computeElementInternal(auxkernel, datum);
     220      210176 : }
     221             : 
     222             : template <typename Derived>
     223             : KOKKOS_FUNCTION void
     224     2081260 : AuxKernel::operator()(NodeLoop, const ThreadID tid, const Derived & auxkernel) const
     225             : {
     226     2081260 :   auto node = _bnd ? kokkosBoundaryNodeID(tid) : kokkosBlockNodeID(tid);
     227     2081260 :   auto & sys = kokkosSystem(_kokkos_var.sys());
     228             : 
     229     2081260 :   if (!sys.isNodalDefined(node, _kokkos_var.var()))
     230     1505280 :     return;
     231             : 
     232      575980 :   AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());
     233             : 
     234      575980 :   auxkernel.computeNodeInternal(auxkernel, datum);
     235             : }
     236             : 
     237             : template <typename Derived>
     238             : KOKKOS_FUNCTION void
     239      209576 : AuxKernel::computeElementInternal(const Derived & auxkernel, AssemblyDatum & datum) const
     240             : {
     241             :   Real x[MAX_CACHED_DOF];
     242             :   Real b[MAX_CACHED_DOF];
     243             :   Real A[MAX_CACHED_DOF * MAX_CACHED_DOF];
     244             : 
     245     1104152 :   for (unsigned int i = 0; i < datum.n_dofs(); ++i)
     246             :   {
     247      894576 :     x[i] = 0;
     248      894576 :     b[i] = 0;
     249             : 
     250     5899152 :     for (unsigned int j = 0; j < datum.n_dofs(); ++j)
     251     5004576 :       A[j + datum.n_dofs() * i] = 0;
     252             :   }
     253             : 
     254     2023184 :   for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     255             :   {
     256     1813608 :     const auto value = auxkernel.computeValue(qp, datum);
     257             : 
     258     1813608 :     datum.reinit();
     259             : 
     260     9792216 :     for (unsigned int i = 0; i < datum.n_dofs(); ++i)
     261             :     {
     262     7978608 :       const auto t = datum.JxW(qp) * _test(datum, i, qp);
     263             : 
     264     7978608 :       b[i] += t * value;
     265             : 
     266    52947216 :       for (unsigned int j = 0; j < datum.n_dofs(); ++j)
     267    44968608 :         A[j + datum.n_dofs() * i] += t * _test(datum, j, qp);
     268             :     }
     269             :   }
     270             : 
     271      209576 :   if (datum.n_dofs() == 1)
     272       72576 :     x[0] = b[0] / A[0];
     273             :   else
     274      137000 :     Utils::choleskySolve(A, x, b, datum.n_dofs());
     275             : 
     276      209576 :   setElementSolution(x, datum);
     277      209576 : }
     278             : 
     279             : template <typename Derived>
     280             : KOKKOS_FUNCTION void
     281      575054 : AuxKernel::computeNodeInternal(const Derived & auxkernel, AssemblyDatum & datum) const
     282             : {
     283      575054 :   auto value = auxkernel.computeValue(0, datum);
     284             : 
     285      575054 :   setNodeSolution(value, datum);
     286      575054 : }
     287             : 
     288             : KOKKOS_FUNCTION inline void
     289      209576 : AuxKernel::setElementSolution(const Real * const values,
     290             :                               const AssemblyDatum & datum,
     291             :                               const unsigned int comp) const
     292             : {
     293      209576 :   auto & sys = kokkosSystem(_kokkos_var.sys(comp));
     294      209576 :   auto var_num = _kokkos_var.var(comp);
     295      209576 :   auto tag = _kokkos_var.tag();
     296      209576 :   auto elem = datum.elem().id;
     297             : 
     298     1104152 :   for (unsigned int i = 0; i < datum.n_dofs(); ++i)
     299      894576 :     sys.getVectorDofValue(sys.getElemLocalDofIndex(elem, i, var_num), tag) = values[i];
     300      209576 : }
     301             : 
     302             : KOKKOS_FUNCTION inline void
     303      575054 : AuxKernel::setNodeSolution(const Real value,
     304             :                            const AssemblyDatum & datum,
     305             :                            const unsigned int comp) const
     306             : {
     307      575054 :   auto & sys = kokkosSystem(_kokkos_var.sys(comp));
     308      575054 :   auto var_num = _kokkos_var.var(comp);
     309      575054 :   auto tag = _kokkos_var.tag();
     310      575054 :   auto node = datum.node();
     311             : 
     312      575054 :   sys.getVectorDofValue(sys.getNodeLocalDofIndex(node, 0, var_num), tag) = value;
     313      575054 : }
     314             : 
     315             : } // namespace Kokkos
     316             : } // namespace Moose

Generated by: LCOV version 1.14