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 "KokkosDatum.h"
      13             : 
      14             : #include "MooseVariableBase.h"
      15             : 
      16             : namespace Moose
      17             : {
      18             : namespace Kokkos
      19             : {
      20             : 
      21             : /**
      22             :  * The Kokkos wrapper classes for MOOSE-like shape function access
      23             :  */
      24             : ///@{
      25             : class VariablePhiValue
      26             : {
      27             : public:
      28             :   /**
      29             :    * Get the current shape function
      30             :    * @param datum The AssemblyDatum object of the current thread
      31             :    * @param i The element-local DOF index
      32             :    * @param qp The local quadrature point index
      33             :    * @returns The shape function
      34             :    */
      35    14798640 :   KOKKOS_FUNCTION Real operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
      36             :   {
      37    14798640 :     auto & elem = datum.elem();
      38    14798640 :     auto side = datum.side();
      39    14798640 :     auto fe = datum.jfe();
      40             : 
      41           0 :     return side == libMesh::invalid_uint
      42    14798640 :                ? datum.assembly().getPhi(elem.subdomain, elem.type, fe)(i, qp)
      43    14798640 :                : datum.assembly().getPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp);
      44             :   }
      45             : };
      46             : 
      47             : class VariablePhiGradient
      48             : {
      49             : public:
      50             :   /**
      51             :    * Get the gradient of the current shape function
      52             :    * @param datum The AssemblyDatum object of the current thread
      53             :    * @param i The element-local DOF index
      54             :    * @param qp The local quadrature point index
      55             :    * @returns The gradient of the shape function
      56             :    */
      57    16930752 :   KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
      58             :   {
      59    16930752 :     auto & elem = datum.elem();
      60    16930752 :     auto side = datum.side();
      61    16930752 :     auto fe = datum.jfe();
      62             : 
      63           0 :     return datum.J(qp) *
      64             :            (side == libMesh::invalid_uint
      65    16930752 :                 ? datum.assembly().getGradPhi(elem.subdomain, elem.type, fe)(i, qp)
      66    33861504 :                 : datum.assembly().getGradPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp));
      67             :   }
      68             : };
      69             : 
      70             : class VariableTestValue
      71             : {
      72             : public:
      73             :   /**
      74             :    * Get the current test function
      75             :    * @param datum The AssemblyDatum object of the current thread
      76             :    * @param i The element-local DOF index
      77             :    * @param qp The local quadrature point index
      78             :    * @returns The test function
      79             :    */
      80    90181084 :   KOKKOS_FUNCTION Real operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
      81             :   {
      82    90181084 :     auto & elem = datum.elem();
      83    90181084 :     auto side = datum.side();
      84    90181084 :     auto fe = datum.ife();
      85             : 
      86           0 :     return side == libMesh::invalid_uint
      87    90181084 :                ? datum.assembly().getPhi(elem.subdomain, elem.type, fe)(i, qp)
      88    90181084 :                : datum.assembly().getPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp);
      89             :   }
      90             : };
      91             : 
      92             : class VariableTestGradient
      93             : {
      94             : public:
      95             :   /**
      96             :    * Get the gradient of the current test function
      97             :    * @param datum The AssemblyDatum object of the current thread
      98             :    * @param i The element-local DOF index
      99             :    * @param qp The local quadrature point index
     100             :    * @returns The gradient of the test function
     101             :    */
     102    42888432 :   KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
     103             :   {
     104    42888432 :     auto & elem = datum.elem();
     105    42888432 :     auto side = datum.side();
     106    42888432 :     auto fe = datum.ife();
     107             : 
     108           0 :     return datum.J(qp) *
     109             :            (side == libMesh::invalid_uint
     110    42888432 :                 ? datum.assembly().getGradPhi(elem.subdomain, elem.type, fe)(i, qp)
     111    85776864 :                 : datum.assembly().getGradPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp));
     112             :   }
     113             : };
     114             : ///@}
     115             : 
     116             : /**
     117             :  * The Kokkos wrapper classes for MOOSE-like variable value access
     118             :  */
     119             : ///@{
     120             : class VariableValue
     121             : {
     122             : public:
     123             :   /**
     124             :    * Default constructor
     125             :    */
     126          72 :   VariableValue() = default;
     127             :   /**
     128             :    * Constructor
     129             :    * @param var The Kokkos variable
     130             :    * @param dof Whether to get DOF values
     131             :    */
     132         901 :   VariableValue(Variable var, bool dof = false) : _var(var), _dof(dof) {}
     133             :   /**
     134             :    * Constructor
     135             :    * @param var The MOOSE variable
     136             :    * @param tag The vector tag name
     137             :    * @param dof Whether to get DOF values
     138             :    */
     139        3669 :   VariableValue(const MooseVariableBase & var,
     140             :                 const TagName & tag = Moose::SOLUTION_TAG,
     141         871 :                 bool dof = false)
     142        2798 :     : _var(var, tag), _dof(dof)
     143             :   {
     144        3669 :   }
     145             : 
     146             :   /**
     147             :    * Get whether the variable was coupled
     148             :    * @returns Whether the variable was coupled
     149             :    */
     150       28032 :   KOKKOS_FUNCTION operator bool() const { return _var.coupled(); }
     151             : 
     152             :   /**
     153             :    * Get the current variable value
     154             :    * @param datum The Datum object of the current thread
     155             :    * @param qp The local quadrature point index
     156             :    * @param comp The variable component
     157             :    * @returns The variable value
     158             :    */
     159             :   KOKKOS_FUNCTION Real operator()(Datum & datum, unsigned int qp, unsigned int comp = 0) const;
     160             : 
     161             : private:
     162             :   /**
     163             :    * Coupled Kokkos variable
     164             :    */
     165             :   Variable _var;
     166             :   /**
     167             :    * Flag whether DOF values are requested
     168             :    */
     169             :   bool _dof = false;
     170             : };
     171             : 
     172             : class VariableGradient
     173             : {
     174             : public:
     175             :   /**
     176             :    * Default constructor
     177             :    */
     178             :   VariableGradient() = default;
     179             :   /**
     180             :    * Constructor
     181             :    * @param var The Kokkos variable
     182             :    */
     183           0 :   VariableGradient(Variable var) : _var(var) {}
     184             :   /**
     185             :    * Constructor
     186             :    * @param var The MOOSE variable
     187             :    * @param tag The vector tag name
     188             :    */
     189        1536 :   VariableGradient(const MooseVariableBase & var, const TagName & tag = Moose::SOLUTION_TAG)
     190        1175 :     : _var(var, tag)
     191             :   {
     192        1536 :   }
     193             : 
     194             :   /**
     195             :    * Get whether the variable was coupled
     196             :    * @returns Whether the variable was coupled
     197             :    */
     198             :   KOKKOS_FUNCTION operator bool() const { return _var.coupled(); }
     199             : 
     200             :   /**
     201             :    * Get the current variable gradient
     202             :    * @param datum The Datum object of the current thread
     203             :    * @param idx The local quadrature point or DOF index
     204             :    * @param comp The variable component
     205             :    * @returns The variable gradient
     206             :    */
     207             :   KOKKOS_FUNCTION Real3 operator()(Datum & datum, unsigned int idx, unsigned int comp = 0) const;
     208             : 
     209             : private:
     210             :   /**
     211             :    * Coupled Kokkos variable
     212             :    */
     213             :   Variable _var;
     214             : };
     215             : ///@}
     216             : 
     217             : KOKKOS_FUNCTION inline Real
     218    25657910 : VariableValue::operator()(Datum & datum, unsigned int idx, unsigned int comp) const
     219             : {
     220             :   KOKKOS_ASSERT(_var.initialized());
     221             : 
     222    25657910 :   if (_var.coupled())
     223             :   {
     224    25642674 :     auto & sys = datum.system(_var.sys(comp));
     225    25642674 :     auto var = _var.var(comp);
     226    25642674 :     auto tag = _var.tag();
     227             : 
     228    25642674 :     if (_dof)
     229             :     {
     230             :       unsigned int dof;
     231             : 
     232     3824430 :       if (datum.isNodal())
     233             :       {
     234     3823830 :         auto node = datum.node();
     235     3823830 :         dof = sys.getNodeLocalDofIndex(node, 0, var);
     236             :       }
     237             :       else
     238             :       {
     239         600 :         auto elem = datum.elem().id;
     240         600 :         dof = sys.getElemLocalDofIndex(elem, idx, var);
     241             :       }
     242             : 
     243     3824430 :       return sys.getVectorDofValue(dof, tag);
     244             :     }
     245             :     else
     246             :     {
     247    21818244 :       auto & elem = datum.elem();
     248    21818244 :       auto side = datum.side();
     249    21818244 :       auto offset = datum.qpOffset();
     250             : 
     251    21818244 :       return side == libMesh::invalid_uint ? sys.getVectorQpValue(elem, offset + idx, var, tag)
     252    21818244 :                                            : sys.getVectorQpValueFace(elem, side, idx, var, tag);
     253             :     }
     254             :   }
     255             :   else
     256       15236 :     return _var.value(comp);
     257             : }
     258             : 
     259             : KOKKOS_FUNCTION inline Real3
     260    25956144 : VariableGradient::operator()(Datum & datum, unsigned int qp, unsigned int comp) const
     261             : {
     262             :   KOKKOS_ASSERT(_var.initialized());
     263             : 
     264    25956144 :   if (_var.coupled())
     265             :   {
     266             :     KOKKOS_ASSERT(!datum.isNodal());
     267             : 
     268    25956144 :     auto & elem = datum.elem();
     269    25956144 :     auto side = datum.side();
     270    25956144 :     auto offset = datum.qpOffset();
     271             : 
     272             :     return side == libMesh::invalid_uint
     273    25956144 :                ? datum.system(_var.sys(comp))
     274    25956144 :                      .getVectorQpGrad(elem, offset + qp, _var.var(comp), _var.tag())
     275           0 :                : datum.system(_var.sys(comp))
     276    25956144 :                      .getVectorQpGradFace(elem, side, datum.J(qp), qp, _var.var(comp), _var.tag());
     277             :   }
     278             :   else
     279           0 :     return Real3(0);
     280             : }
     281             : 
     282             : } // namespace Kokkos
     283             : } // namespace Moose