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 "KokkosDispatcher.h"
      13             : #include "KokkosMesh.h"
      14             : #include "KokkosSystem.h"
      15             : #include "KokkosVariable.h"
      16             : 
      17             : #include "MooseObject.h"
      18             : #include "SetupInterface.h"
      19             : #include "FunctionInterface.h"
      20             : #include "UserObjectInterface.h"
      21             : #include "TransientInterface.h"
      22             : #include "PostprocessorInterface.h"
      23             : #include "VectorPostprocessorInterface.h"
      24             : #include "Restartable.h"
      25             : #include "MeshChangedInterface.h"
      26             : #include "TaggingInterface.h"
      27             : #include "LinearSystemContributionObject.h"
      28             : 
      29             : class InputParameters;
      30             : class FEProblemBase;
      31             : 
      32             : namespace Moose::Kokkos
      33             : {
      34             : 
      35             : /**
      36             :  * Base class for Kokkos objects that contribute to a linear system, i.e. the linear finite volume
      37             :  * kernels and boundary conditions. It provides the device-side helpers for accumulating into the
      38             :  * tagged vectors and matrices.
      39             :  */
      40             : class LinearSystemContributionObject : public MooseObject,
      41             :                                        public SetupInterface,
      42             :                                        public FunctionInterface,
      43             :                                        public UserObjectInterface,
      44             :                                        public TransientInterface,
      45             :                                        public PostprocessorInterface,
      46             :                                        public VectorPostprocessorInterface,
      47             :                                        public Restartable,
      48             :                                        public MeshChangedInterface,
      49             :                                        public TaggingInterface,
      50             :                                        public MeshHolder,
      51             :                                        public SystemHolder
      52             : {
      53             : public:
      54             :   static InputParameters validParams();
      55             : 
      56             :   LinearSystemContributionObject(const InputParameters & parameters);
      57             :   LinearSystemContributionObject(const LinearSystemContributionObject & object);
      58             : 
      59             :   /**
      60             :    * Get the Kokkos variable this object contributes to
      61             :    * @returns The Kokkos variable
      62             :    */
      63         340 :   Variable variable() const { return _var; }
      64             : 
      65             :   /// Tag dispatch type for the right-hand side computation loop
      66             :   struct RightHandSideLoop
      67             :   {
      68             :   };
      69             :   /// Tag dispatch type for the matrix computation loop
      70             :   struct MatrixLoop
      71             :   {
      72             :   };
      73             : 
      74             :   /**
      75             :    * Compute the right-hand side contributions of this object
      76             :    */
      77             :   virtual void computeRightHandSide() = 0;
      78             :   /**
      79             :    * Compute the matrix contributions of this object
      80             :    */
      81             :   virtual void computeMatrix() = 0;
      82             : 
      83             : protected:
      84             :   /// Whether tagged vector/matrix accumulation must use atomics
      85             :   enum class AccumulationMode
      86             :   {
      87             :     Atomic,
      88             :     NonAtomic
      89             :   };
      90             : 
      91             :   template <AccumulationMode mode = AccumulationMode::Atomic>
      92             :   KOKKOS_FUNCTION void accumulateTaggedVector(Real value, dof_id_type row) const;
      93             :   template <AccumulationMode mode = AccumulationMode::Atomic>
      94             :   KOKKOS_FUNCTION void accumulateTaggedMatrix(Real value, dof_id_type row, dof_id_type col) const;
      95             : 
      96             :   /// Reference to the finite element problem
      97             :   FEProblemBase & _fe_problem;
      98             : 
      99             :   /// Kokkos variable
     100             :   Variable _var;
     101             : 
     102             :   /// Vector (residual) tags this object contributes to
     103             :   Array<TagID> _vector_tags;
     104             :   /// Matrix tags this object contributes to
     105             :   Array<TagID> _matrix_tags;
     106             : 
     107             :   /// Dispatcher for the right-hand side computation loop
     108             :   std::unique_ptr<DispatcherBase> _rhs_dispatcher;
     109             :   /// Dispatcher for the matrix computation loop
     110             :   std::unique_ptr<DispatcherBase> _matrix_dispatcher;
     111             : 
     112             :   /**
     113             :    * TODO: Move to TransientInterface
     114             :    */
     115             :   ///@{
     116             :   /// Current time
     117             :   Scalar<Real> _t;
     118             :   ///@}
     119             : };
     120             : 
     121             : template <LinearSystemContributionObject::AccumulationMode mode>
     122             : KOKKOS_FUNCTION inline void
     123       12208 : LinearSystemContributionObject::accumulateTaggedVector(const Real value,
     124             :                                                        const dof_id_type row) const
     125             : {
     126       12208 :   if (!value)
     127         222 :     return;
     128             : 
     129             :   KOKKOS_ASSERT(_var.components() == 1);
     130       11986 :   auto & sys = kokkosSystem(_var.sys());
     131       11986 :   const MOOSE_KOKKOS_INDEX_TYPE tags_size = _vector_tags.size();
     132       23972 :   for (MOOSE_KOKKOS_INDEX_TYPE index = 0; index < tags_size; ++index)
     133             :   {
     134       11986 :     const auto tag = _vector_tags[index];
     135       11986 :     if (sys.isResidualTagActive(tag))
     136             :     {
     137             :       if constexpr (mode == AccumulationMode::Atomic)
     138        1376 :         ::Kokkos::atomic_add(&sys.getVectorDofValue(row, tag), value);
     139             :       else
     140       10610 :         sys.getVectorDofValue(row, tag) += value;
     141             :     }
     142             :   }
     143             : }
     144             : 
     145             : template <LinearSystemContributionObject::AccumulationMode mode>
     146             : KOKKOS_FUNCTION inline void
     147       81314 : LinearSystemContributionObject::accumulateTaggedMatrix(const Real value,
     148             :                                                        const dof_id_type row,
     149             :                                                        const dof_id_type col) const
     150             : {
     151       81314 :   if (!value)
     152        1517 :     return;
     153             : 
     154             :   KOKKOS_ASSERT(_var.components() == 1);
     155       79797 :   auto & sys = kokkosSystem(_var.sys());
     156       79797 :   const MOOSE_KOKKOS_INDEX_TYPE tags_size = _matrix_tags.size();
     157      159594 :   for (MOOSE_KOKKOS_INDEX_TYPE index = 0; index < tags_size; ++index)
     158             :   {
     159       79797 :     const auto tag = _matrix_tags[index];
     160       79797 :     if (sys.isMatrixTagActive(tag))
     161             :     {
     162             :       if constexpr (mode == AccumulationMode::Atomic)
     163       79797 :         ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), value);
     164             :       else
     165           0 :         sys.getMatrixValue(row, col, tag) += value;
     166             :     }
     167             :   }
     168             : }
     169             : 
     170             : } // namespace Moose::Kokkos