KokkosIntegratedBC.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 "KokkosIntegratedBCBase.h"

namespace Moose
{
namespace Kokkos
{

template <typename Derived>
class IntegratedBC : public IntegratedBCBase
{
public:
  static InputParameters validParams();

  IntegratedBC(const InputParameters & parameters);

  virtual void computeResidual() override;
  virtual void computeJacobian() override;


  KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
                                         const unsigned int /* j */,
                                         const unsigned int /* qp */,
                                         ResidualDatum & /* datum */) const
  {
    return 0;
  }
  KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
                                                const unsigned int /* j */,
                                                const unsigned int /* jvar */,
                                                const unsigned int /* qp */,
                                                ResidualDatum & /* datum */) const
  {
    return 0;
  }

  KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid) const;
  KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid) const;
  KOKKOS_FUNCTION void operator()(OffDiagJacobianLoop, const ThreadID tid) const;


  KOKKOS_FUNCTION void computeResidualInternal(const Derived * bc, ResidualDatum & datum) const;
  KOKKOS_FUNCTION void computeJacobianInternal(const Derived * bc, ResidualDatum & datum) const;
  KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived * bc,
                                                      ResidualDatum & datum) const;

protected:
  const VariableTestValue _test;
  const VariableTestGradient _grad_test;
  const VariablePhiValue _phi;
  const VariablePhiGradient _grad_phi;
  const VariableValue _u;
  const VariableGradient _grad_u;

protected:
  virtual bool defaultJacobian() const
  {
    return &Derived::computeQpJacobian == &IntegratedBC::computeQpJacobian;
  }
  virtual bool defaultOffDiagJacobian() const
  {
    return &Derived::computeQpOffDiagJacobian == &IntegratedBC::computeQpOffDiagJacobian;
  }
};

template <typename Derived>
InputParameters
IntegratedBC<Derived>::validParams()
{
  InputParameters params = IntegratedBCBase::validParams();
  return params;
}

template <typename Derived>
IntegratedBC<Derived>::IntegratedBC(const InputParameters & parameters)
  : IntegratedBCBase(parameters, Moose::VarFieldType::VAR_FIELD_STANDARD),
    _test(),
    _grad_test(),
    _phi(),
    _grad_phi(),
    _u(_var),
    _grad_u(_var)
{
  addMooseVariableDependency(&_var);
}

template <typename Derived>
void
IntegratedBC<Derived>::computeResidual()
{
  ::Kokkos::RangePolicy<ResidualLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
      0, numKokkosBoundarySides());
  ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
  ::Kokkos::fence();
}

template <typename Derived>
void
IntegratedBC<Derived>::computeJacobian()
{
  if (!defaultJacobian())
  {
    ::Kokkos::RangePolicy<JacobianLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
        0, numKokkosBoundarySides());
    ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
    ::Kokkos::fence();
  }

  if (!defaultOffDiagJacobian())
  {
    auto & sys = kokkosSystem(_kokkos_var.sys());

    _thread.resize({sys.getCoupling(_kokkos_var.var()).size(), numKokkosBoundarySides()});

    ::Kokkos::RangePolicy<OffDiagJacobianLoop, ExecSpace, ::Kokkos::IndexType<ThreadID>> policy(
        0, _thread.size());
    ::Kokkos::parallel_for(policy, *static_cast<Derived *>(this));
    ::Kokkos::fence();
  }
}

template <typename Derived>
KOKKOS_FUNCTION void
IntegratedBC<Derived>::operator()(ResidualLoop, const ThreadID tid) const
{
  auto bc = static_cast<const Derived *>(this);
  auto [elem, side] = kokkosBoundaryElementSideID(tid);

  ResidualDatum datum(
      elem, side, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());

  bc->computeResidualInternal(bc, datum);
}

template <typename Derived>
KOKKOS_FUNCTION void
IntegratedBC<Derived>::operator()(JacobianLoop, const ThreadID tid) const
{
  auto bc = static_cast<const Derived *>(this);
  auto [elem, side] = kokkosBoundaryElementSideID(tid);

  ResidualDatum datum(
      elem, side, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var());

  bc->computeJacobianInternal(bc, datum);
}

template <typename Derived>
KOKKOS_FUNCTION void
IntegratedBC<Derived>::operator()(OffDiagJacobianLoop, const ThreadID tid) const
{
  auto bc = static_cast<const Derived *>(this);
  auto [elem, side] = kokkosBoundaryElementSideID(_thread(tid, 1));

  auto & sys = kokkosSystem(_kokkos_var.sys());
  auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];

  if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
    return;

  ResidualDatum datum(elem, side, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar);

  bc->computeOffDiagJacobianInternal(bc, datum);
}

template <typename Derived>
KOKKOS_FUNCTION void
IntegratedBC<Derived>::computeResidualInternal(const Derived * bc, ResidualDatum & datum) const
{
  ResidualObject::computeResidualInternal(
      datum,
      [&](Real * local_re, const unsigned int ib, const unsigned int ie)
      {
        for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
        {
          datum.reinit();

          for (unsigned int i = ib; i < ie; ++i)
            local_re[i] += datum.JxW(qp) * bc->computeQpResidual(i, qp, datum);
        }
      });
}

template <typename Derived>
KOKKOS_FUNCTION void
IntegratedBC<Derived>::computeJacobianInternal(const Derived * bc, ResidualDatum & datum) const
{
  ResidualObject::computeJacobianInternal(
      datum,
      [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
      {
        for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
        {
          datum.reinit();

          for (unsigned int ij = ijb; ij < ije; ++ij)
          {
            unsigned int i = ij % datum.n_jdofs();
            unsigned int j = ij / datum.n_jdofs();

            local_ke[ij] += datum.JxW(qp) * bc->computeQpJacobian(i, j, qp, datum);
          }
        }
      });
}

template <typename Derived>
KOKKOS_FUNCTION void
IntegratedBC<Derived>::computeOffDiagJacobianInternal(const Derived * bc,
                                                      ResidualDatum & datum) const
{
  ResidualObject::computeJacobianInternal(
      datum,
      [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
      {
        for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
        {
          datum.reinit();

          for (unsigned int ij = ijb; ij < ije; ++ij)
          {
            unsigned int i = ij % datum.n_jdofs();
            unsigned int j = ij / datum.n_jdofs();

            local_ke[ij] +=
                datum.JxW(qp) * bc->computeQpOffDiagJacobian(i, j, datum.jvar(), qp, datum);
          }
        }
      });
}

} // namespace Kokkos
} // namespace Moose

#define usingKokkosIntegratedBCMembers(T)                                                          \
  usingKokkosIntegratedBCBaseMembers;                                                              \
                                                                                                   \
protected:                                                                                         \
  using Moose::Kokkos::IntegratedBC<T>::_test;                                                     \
  using Moose::Kokkos::IntegratedBC<T>::_grad_test;                                                \
  using Moose::Kokkos::IntegratedBC<T>::_phi;                                                      \
  using Moose::Kokkos::IntegratedBC<T>::_grad_phi;                                                 \
  using Moose::Kokkos::IntegratedBC<T>::_u;                                                        \
  using Moose::Kokkos::IntegratedBC<T>::_grad_u;                                                   \
                                                                                                   \
public:                                                                                            \
  using Moose::Kokkos::IntegratedBC<T>::operator()