KokkosVariableValue.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 "KokkosDatum.h"

#include "MooseVariableFieldBase.h"

namespace Moose::Kokkos
{

class VariablePhiValue
{
public:
  KOKKOS_FUNCTION Real operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = datum.jfe();

    return side == libMesh::invalid_uint
               ? datum.assembly().getPhi(elem.subdomain, elem.type, fe)(i, qp)
               : datum.assembly().getPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp);
  }
};

class VariablePhiGradient
{
public:
  KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = datum.jfe();

    return datum.J(qp) *
           (side == libMesh::invalid_uint
                ? datum.assembly().getGradPhi(elem.subdomain, elem.type, fe)(i, qp)
                : datum.assembly().getGradPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp));
  }
};

class VariableTestValue
{
public:
  KOKKOS_FUNCTION Real operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = datum.ife();

    return side == libMesh::invalid_uint
               ? datum.assembly().getPhi(elem.subdomain, elem.type, fe)(i, qp)
               : datum.assembly().getPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp);
  }
};

class VariableTestGradient
{
public:
  KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = datum.ife();

    return datum.J(qp) *
           (side == libMesh::invalid_uint
                ? datum.assembly().getGradPhi(elem.subdomain, elem.type, fe)(i, qp)
                : datum.assembly().getGradPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp));
  }
};

using ADVariablePhiValue = VariablePhiValue;
using ADVariablePhiGradient = VariablePhiGradient;
using ADVariableTestValue = VariableTestValue;
using ADVariableTestGradient = VariableTestGradient;


template <bool is_ad>
class VariableValueTempl
{
  using real_type = std::conditional_t<is_ad, ADReal, Real>;

public:
  VariableValueTempl() = default;
  VariableValueTempl(Variable var, bool dof = false) : _var(var), _dof(dof) {}
  VariableValueTempl(const MooseVariableFieldBase & var,
                     const TagName & tag = Moose::SOLUTION_TAG,
                     bool dof = false)
    : _var(var, tag), _dof(dof)
  {
  }
  VariableValueTempl(const std::vector<const MooseVariableFieldBase *> & vars,
                     const TagName & tag = Moose::SOLUTION_TAG,
                     bool dof = false)
    : _var(vars, tag), _dof(dof)
  {
  }
  VariableValueTempl(const std::vector<MooseVariableFieldBase *> & vars,
                     const TagName & tag = Moose::SOLUTION_TAG,
                     bool dof = false)
    : _var(vars, tag), _dof(dof)
  {
  }

  VariableValueTempl(const VariableValueTempl<is_ad> & object);
  VariableValueTempl<is_ad> & operator=(const VariableValueTempl<is_ad> & object);

  KOKKOS_FUNCTION operator bool() const { return _var.coupled(); }

  KOKKOS_FUNCTION auto operator()(Datum & datum, unsigned int idx, unsigned int comp = 0) const
  {
    return get(datum, idx, comp);
  }

  KOKKOS_FUNCTION auto
  operator()(AssemblyDatum & datum, unsigned int idx, unsigned int comp = 0) const;

  KOKKOS_FUNCTION const Variable & variable() const { return _var; }

private:
  KOKKOS_FUNCTION auto
  get(Datum & datum, unsigned int idx, unsigned int comp = 0, Real seed = 0) const;

  Variable _var;
  Array<Real> _seed;
  bool _dof = false;
};

template <bool is_ad>
VariableValueTempl<is_ad>::VariableValueTempl(const VariableValueTempl<is_ad> & object)
  : _var(object._var), _seed(object._seed), _dof(object._dof)
{
  if constexpr (is_ad)
    if (_var.coupled())
    {
      if (!_seed.isAlloc())
        _seed.create(_var.components());

      for (unsigned int comp = 0; comp < _var.components(); ++comp)
        _seed[comp] =
            _var.dot() ? _var.mooseVar(comp)->sys().duDotDu(_var.var(comp)) : (_var.old() ? 0 : 1);

      _seed.copyToDevice();
    }
}

template <bool is_ad>
VariableValueTempl<is_ad> &
VariableValueTempl<is_ad>::operator=(const VariableValueTempl<is_ad> & object)
{
  _var = object._var;
  _dof = object._dof;

  return *this;
}

template <bool is_ad>
KOKKOS_FUNCTION auto
VariableValueTempl<is_ad>::operator()(AssemblyDatum & datum,
                                      unsigned int idx,
                                      unsigned int comp) const
{
  if constexpr (is_ad)
  {
    Real seed =
        datum.do_derivatives() && _var.coupled() && _var.sys(comp) == datum.sys() ? _seed[comp] : 0;

    return get(datum, idx, comp, seed);
  }
  else
    return get(datum, idx, comp);
}

template <bool is_ad>
KOKKOS_FUNCTION auto
VariableValueTempl<is_ad>::get(Datum & datum,
                               unsigned int idx,
                               unsigned int comp,
                               [[maybe_unused]] Real seed) const
{
  KOKKOS_ASSERT(_var.initialized());

  real_type value;

  if (_var.coupled())
  {
    auto & sys = datum.system(_var.sys(comp));
    auto var = _var.var(comp);
    auto tag = _var.tag();

    if (_dof)
    {
      unsigned int dof;

      if (datum.isNodal())
      {
        auto node = datum.node();
        dof = sys.getNodeLocalDofIndex(node, 0, var);
      }
      else
      {
        auto elem = datum.elem().id;
        dof = sys.getElemLocalDofIndex(elem, idx, var);
      }

      if constexpr (is_ad)
        value = sys.getVectorDofADValue(dof, tag, seed);
      else
        value = sys.getVectorDofValue(dof, tag);
    }
    else
    {
      auto & elem = datum.elem();
      auto side = datum.side();

      if constexpr (is_ad)
        value = side == libMesh::invalid_uint
                    ? sys.getVectorQpADValue(elem, datum.qpOffset(), idx, var, tag, seed)
                    : sys.getVectorQpADValueFace(elem, side, idx, var, tag, seed);
      else
        value = side == libMesh::invalid_uint
                    ? sys.getVectorQpValue(elem, datum.qpOffset() + idx, var, tag)
                    : sys.getVectorQpValueFace(elem, side, idx, var, tag);
    }
  }
  else
    value = _var.value(comp);

  return value;
}

template <bool is_ad>
class VariableGradientTempl
{
  using real3_type = std::conditional_t<is_ad, ADReal3, Real3>;

public:
  VariableGradientTempl() = default;
  VariableGradientTempl(Variable var) : _var(var) {}
  VariableGradientTempl(const MooseVariableFieldBase & var,
                        const TagName & tag = Moose::SOLUTION_TAG)
    : _var(var, tag)
  {
  }
  VariableGradientTempl(const std::vector<const MooseVariableFieldBase *> & vars,
                        const TagName & tag = Moose::SOLUTION_TAG)
    : _var(vars, tag)
  {
  }
  VariableGradientTempl(const std::vector<MooseVariableFieldBase *> & vars,
                        const TagName & tag = Moose::SOLUTION_TAG)
    : _var(vars, tag)
  {
  }

  VariableGradientTempl(const VariableGradientTempl<is_ad> & object);
  VariableGradientTempl<is_ad> & operator=(const VariableGradientTempl<is_ad> & object);

  KOKKOS_FUNCTION operator bool() const { return _var.coupled(); }

  KOKKOS_FUNCTION auto operator()(Datum & datum, unsigned int qp, unsigned int comp = 0) const
  {
    return get(datum, qp, comp);
  }

  KOKKOS_FUNCTION auto
  operator()(AssemblyDatum & datum, unsigned int qp, unsigned int comp = 0) const;

  KOKKOS_FUNCTION const Variable & variable() const { return _var; }

private:
  KOKKOS_FUNCTION auto
  get(Datum & datum, unsigned int qp, unsigned int comp = 0, Real seed = 0) const;

  Variable _var;
  Array<Real> _seed;
};

template <bool is_ad>
VariableGradientTempl<is_ad>::VariableGradientTempl(const VariableGradientTempl<is_ad> & object)
  : _var(object._var), _seed(object._seed)
{
  if constexpr (is_ad)
    if (_var.coupled())
    {
      if (!_seed.isAlloc())
        _seed.create(_var.components());

      for (unsigned int comp = 0; comp < _var.components(); ++comp)
        _seed[comp] =
            _var.dot() ? _var.mooseVar(comp)->sys().duDotDu(_var.var(comp)) : (_var.old() ? 0 : 1);

      _seed.copyToDevice();
    }
}

template <bool is_ad>
VariableGradientTempl<is_ad> &
VariableGradientTempl<is_ad>::operator=(const VariableGradientTempl<is_ad> & object)
{
  _var = object._var;

  return *this;
}

template <bool is_ad>
KOKKOS_FUNCTION auto
VariableGradientTempl<is_ad>::operator()(AssemblyDatum & datum,
                                         unsigned int qp,
                                         unsigned int comp) const
{
  if constexpr (is_ad)
  {
    Real seed =
        datum.do_derivatives() && _var.coupled() && _var.sys(comp) == datum.sys() ? _seed[comp] : 0;

    return get(datum, qp, comp, seed);
  }
  else
    return get(datum, qp, comp);
}

template <bool is_ad>
KOKKOS_FUNCTION auto
VariableGradientTempl<is_ad>::get(Datum & datum,
                                  unsigned int qp,
                                  unsigned int comp,
                                  [[maybe_unused]] Real seed) const
{
  KOKKOS_ASSERT(_var.initialized());

  real3_type grad;

  if (_var.coupled())
  {
    KOKKOS_ASSERT(!datum.isNodal());

    auto & elem = datum.elem();
    auto side = datum.side();

    if constexpr (is_ad)
      grad =
          side == libMesh::invalid_uint
              ? datum.system(_var.sys(comp))
                    .getVectorQpADGrad(
                        elem, datum.J(qp), datum.qpOffset(), qp, _var.var(comp), _var.tag(), seed)
              : datum.system(_var.sys(comp))
                    .getVectorQpADGradFace(
                        elem, side, datum.J(qp), qp, _var.var(comp), _var.tag(), seed);
    else
      grad =
          side == libMesh::invalid_uint
              ? datum.system(_var.sys(comp))
                    .getVectorQpGrad(elem, datum.qpOffset() + qp, _var.var(comp), _var.tag())
              : datum.system(_var.sys(comp))
                    .getVectorQpGradFace(elem, side, datum.J(qp), qp, _var.var(comp), _var.tag());
  }

  return grad;
}

using VariableValue = VariableValueTempl<false>;
using ADVariableValue = VariableValueTempl<true>;
using VariableGradient = VariableGradientTempl<false>;
using ADVariableGradient = VariableGradientTempl<true>;

template <>
struct ArrayDeepCopy<ADVariableValue>
{
  static constexpr bool value = true;
};

template <>
struct ArrayDeepCopy<ADVariableGradient>
{
  static constexpr bool value = true;
};

} // namespace Moose::Kokkos