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 "KokkosIntegratedBCBase.h"
      13             : 
      14             : namespace Moose::Kokkos
      15             : {
      16             : 
      17             : /**
      18             :  * The base class for a user to derive their own Kokkos integrated boundary conditions using
      19             :  * automatic differentiation (AD).
      20             :  *
      21             :  * The user should define computeQpResidual() as inlined public methods in their derived class (not
      22             :  * virtual override). The signature of computeQpResidual() expected to be defined in the derived
      23             :  * class is as follows:
      24             :  *
      25             :  * @tparam Derived The object type
      26             :  * @param i The element-local DOF index
      27             :  * @param qp The local quadrature point index
      28             :  * @param datum The AssemblyDatum object of the current thread
      29             :  * @returns The residual contribution
      30             :  *
      31             :  * template <typename Derived>
      32             :  * KOKKOS_FUNCTION Moose::Kokkos::ADReal computeQpResidual(const unsigned int i,
      33             :  *                                                         const unsigned int qp,
      34             :  *                                                         AssemblyDatum & datum) const;
      35             :  *
      36             :  * Note that computeQpJacobian() and computeQpOffDiagJacobian() are unused for AD integrated
      37             :  * boundary conditions.
      38             :  */
      39             : class ADIntegratedBC : public IntegratedBCBase
      40             : {
      41             : public:
      42             :   static InputParameters validParams();
      43             : 
      44             :   /**
      45             :    * Constructor
      46             :    */
      47             :   ADIntegratedBC(const InputParameters & parameters);
      48             : 
      49             :   virtual void computeResidual() override;
      50             :   virtual void computeJacobian() override;
      51             :   virtual void computeResidualAndJacobian() override;
      52             : 
      53             :   /**
      54             :    * The parallel computation entry function called by Kokkos
      55             :    */
      56             :   template <typename Derived>
      57             :   KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived & bc) const;
      58             : 
      59             :   /**
      60             :    * Compute residual
      61             :    * @param bc The boundary condition object of the final derived type
      62             :    * @param datum The AssemblyDatum object of the current thread
      63             :    */
      64             :   template <typename Derived>
      65             :   KOKKOS_FUNCTION void computeResidualInternal(const Derived & bc, AssemblyDatum & datum) const;
      66             : 
      67             : protected:
      68             :   /**
      69             :    * Dispatch parallel calculation
      70             :    */
      71             :   virtual void dispatch();
      72             : 
      73             :   /**
      74             :    * Current test function
      75             :    */
      76             :   const ADVariableTestValue _test;
      77             :   /**
      78             :    * Gradient of the current test function
      79             :    */
      80             :   const ADVariableTestGradient _grad_test;
      81             :   /**
      82             :    * Current shape function
      83             :    */
      84             :   const ADVariablePhiValue _phi;
      85             :   /**
      86             :    * Gradient of the current shape function
      87             :    */
      88             :   const ADVariablePhiGradient _grad_phi;
      89             :   /**
      90             :    * Current solution at quadrature points
      91             :    */
      92             :   const ADVariableValue _u;
      93             :   /**
      94             :    * Gradient of the current solution at quadrature points
      95             :    */
      96             :   const ADVariableGradient _grad_u;
      97             :   /**
      98             :    * Whether computing residual
      99             :    */
     100             :   bool _computing_residual = false;
     101             :   /**
     102             :    * Whether computing Jacobian
     103             :    */
     104             :   bool _computing_jacobian = false;
     105             : };
     106             : 
     107             : template <typename Derived>
     108             : KOKKOS_FUNCTION void
     109        9048 : ADIntegratedBC::operator()(ResidualLoop, const ThreadID tid, const Derived & bc) const
     110             : {
     111        9048 :   auto [elem, side] = kokkosBoundaryElementSideID(_thread(tid, 1));
     112             : 
     113       18096 :   AssemblyDatum datum(
     114        9048 :       elem, side, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());
     115             : 
     116        9048 :   datum.set_local_parallel(_thread(tid, 0), _thread.size(0));
     117             : 
     118        9048 :   datum.do_derivatives(_computing_jacobian);
     119             : 
     120        9048 :   bc.computeResidualInternal(bc, datum);
     121        9048 : }
     122             : 
     123             : template <typename Derived>
     124             : KOKKOS_FUNCTION void
     125        9048 : ADIntegratedBC::computeResidualInternal(const Derived & bc, AssemblyDatum & datum) const
     126             : {
     127       17180 :   for (unsigned int i = datum.local_thread_id(); i < datum.n_dofs(); i += datum.num_local_threads())
     128             :   {
     129        8132 :     ADReal local_re = 0;
     130             : 
     131       23480 :     for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
     132       15348 :       local_re += datum.JxW(qp) * bc.template computeQpResidual<Derived>(i, qp, datum);
     133             : 
     134        8132 :     if (_computing_residual)
     135        8132 :       accumulateTaggedElementalResidual(local_re.value(), datum.elem().id, i);
     136        8132 :     if (_computing_jacobian)
     137        1580 :       accumulateTaggedElementalMatrix(local_re.derivatives(), datum, i);
     138             :   }
     139        9048 : }
     140             : 
     141             : } // namespace Moose::Kokkos