KokkosLinearSystemContributionObject.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#pragma once

#include "KokkosDispatcher.h"
#include "KokkosMesh.h"
#include "KokkosSystem.h"
#include "KokkosVariable.h"

#include "MooseObject.h"
#include "SetupInterface.h"
#include "FunctionInterface.h"
#include "UserObjectInterface.h"
#include "TransientInterface.h"
#include "PostprocessorInterface.h"
#include "VectorPostprocessorInterface.h"
#include "Restartable.h"
#include "MeshChangedInterface.h"
#include "TaggingInterface.h"
#include "LinearSystemContributionObject.h"

class InputParameters;
class FEProblemBase;

namespace Moose::Kokkos
{

class LinearSystemContributionObject : public MooseObject,
                                       public SetupInterface,
                                       public FunctionInterface,
                                       public UserObjectInterface,
                                       public TransientInterface,
                                       public PostprocessorInterface,
                                       public VectorPostprocessorInterface,
                                       public Restartable,
                                       public MeshChangedInterface,
                                       public TaggingInterface,
                                       public MeshHolder,
                                       public SystemHolder
{
public:
  static InputParameters validParams();

  LinearSystemContributionObject(const InputParameters & parameters);
  LinearSystemContributionObject(const LinearSystemContributionObject & object);

  Variable variable() const { return _var; }

  struct RightHandSideLoop
  {
  };
  struct MatrixLoop
  {
  };

  virtual void computeRightHandSide() = 0;
  virtual void computeMatrix() = 0;

protected:
  enum class AccumulationMode
  {
    Atomic,
    NonAtomic
  };

  template <AccumulationMode mode = AccumulationMode::Atomic>
  KOKKOS_FUNCTION void accumulateTaggedVector(Real value, dof_id_type row) const;
  template <AccumulationMode mode = AccumulationMode::Atomic>
  KOKKOS_FUNCTION void accumulateTaggedMatrix(Real value, dof_id_type row, dof_id_type col) const;

  FEProblemBase & _fe_problem;

  Variable _var;

  Array<TagID> _vector_tags;
  Array<TagID> _matrix_tags;

  std::unique_ptr<DispatcherBase> _rhs_dispatcher;
  std::unique_ptr<DispatcherBase> _matrix_dispatcher;

  Scalar<Real> _t;
};

template <LinearSystemContributionObject::AccumulationMode mode>
KOKKOS_FUNCTION inline void
LinearSystemContributionObject::accumulateTaggedVector(const Real value,
                                                       const dof_id_type row) const
{
  if (!value)
    return;

  KOKKOS_ASSERT(_var.components() == 1);
  auto & sys = kokkosSystem(_var.sys());
  const MOOSE_KOKKOS_INDEX_TYPE tags_size = _vector_tags.size();
  for (MOOSE_KOKKOS_INDEX_TYPE index = 0; index < tags_size; ++index)
  {
    const auto tag = _vector_tags[index];
    if (sys.isResidualTagActive(tag))
    {
      if constexpr (mode == AccumulationMode::Atomic)
        ::Kokkos::atomic_add(&sys.getVectorDofValue(row, tag), value);
      else
        sys.getVectorDofValue(row, tag) += value;
    }
  }
}

template <LinearSystemContributionObject::AccumulationMode mode>
KOKKOS_FUNCTION inline void
LinearSystemContributionObject::accumulateTaggedMatrix(const Real value,
                                                       const dof_id_type row,
                                                       const dof_id_type col) const
{
  if (!value)
    return;

  KOKKOS_ASSERT(_var.components() == 1);
  auto & sys = kokkosSystem(_var.sys());
  const MOOSE_KOKKOS_INDEX_TYPE tags_size = _matrix_tags.size();
  for (MOOSE_KOKKOS_INDEX_TYPE index = 0; index < tags_size; ++index)
  {
    const auto tag = _matrix_tags[index];
    if (sys.isMatrixTagActive(tag))
    {
      if constexpr (mode == AccumulationMode::Atomic)
        ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), value);
      else
        sys.getMatrixValue(row, col, tag) += value;
    }
  }
}

} // namespace Moose::Kokkos