KokkosDatum.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 "KokkosTypes.h"
#include "KokkosAssembly.h"
#include "KokkosSystem.h"
#include "KokkosVariable.h"

namespace Moose
{
namespace Kokkos
{

class Datum
{
public:
  KOKKOS_FUNCTION
  Datum(const ContiguousElementID elem,
        const unsigned int side,
        const Assembly & assembly,
        const Array<System> & systems)
    : _assembly(assembly),
      _systems(systems),
      _elem(assembly.kokkosMesh().getElementInfo(elem)),
      _side(side),
      _neighbor(_side == libMesh::invalid_uint ? libMesh::DofObject::invalid_id
                                               : assembly.kokkosMesh().getNeighbor(_elem.id, side)),
      _n_qps(side == libMesh::invalid_uint ? assembly.getNumQps(_elem)
                                           : assembly.getNumFaceQps(_elem, side)),
      _qp_offset(side == libMesh::invalid_uint ? assembly.getQpOffset(_elem)
                                               : assembly.getQpFaceOffset(_elem, side))
  {
  }
  KOKKOS_FUNCTION
  Datum(const ContiguousNodeID node, const Assembly & assembly, const Array<System> & systems)
    : _assembly(assembly), _systems(systems), _node(node)
  {
  }

  KOKKOS_FUNCTION const Assembly & assembly() const { return _assembly; }
  KOKKOS_FUNCTION const System & system(unsigned int sys) const { return _systems[sys]; }

  KOKKOS_FUNCTION const ElementInfo & elem() const { return _elem; }
  KOKKOS_FUNCTION ContiguousSubdomainID subdomain() const { return _elem.subdomain; }
  KOKKOS_FUNCTION unsigned int side() const { return _side; }
  KOKKOS_FUNCTION ContiguousNodeID node() const { return _node; }
  KOKKOS_FUNCTION unsigned int n_qps() const { return _n_qps; }
  KOKKOS_FUNCTION dof_id_type qpOffset() const { return _qp_offset; }
  KOKKOS_FUNCTION bool hasNeighbor() const { return _neighbor != libMesh::DofObject::invalid_id; }
  KOKKOS_FUNCTION bool isNodal() const { return _node != libMesh::DofObject::invalid_id; }

  KOKKOS_FUNCTION const Real33 & J(const unsigned int qp)
  {
    if (!isNodal())
      reinitTransform(qp);
    else
      _J.identity(_assembly.getDimension());

    return _J;
  }
  KOKKOS_FUNCTION Real JxW(const unsigned int qp)
  {
    if (!isNodal())
      reinitTransform(qp);
    else
      _JxW = 1;

    return _JxW;
  }
  KOKKOS_FUNCTION Real3 q_point(const unsigned int qp)
  {
    if (!isNodal())
      reinitTransform(qp);
    else
      _xyz = _assembly.kokkosMesh().getNodePoint(_node);

    return _xyz;
  }

  KOKKOS_FUNCTION void reinit() { _transform_reinit = false; }

protected:
  const Assembly & _assembly;
  const Array<System> & _systems;
  const ElementInfo _elem;
  const unsigned int _side = libMesh::invalid_uint;
  const ContiguousNodeID _node = libMesh::DofObject::invalid_id;
  const ContiguousElementID _neighbor = libMesh::DofObject::invalid_id;
  const unsigned int _n_qps = 1;
  const dof_id_type _qp_offset = libMesh::DofObject::invalid_id;

private:
  KOKKOS_FUNCTION void reinitTransform(const unsigned int qp);

  bool _transform_reinit = false;
  Real33 _J;
  Real _JxW;
  Real3 _xyz;
};

KOKKOS_FUNCTION inline void
Datum::reinitTransform(const unsigned int qp)
{
  if (_transform_reinit)
    return;

  if (_side == libMesh::invalid_uint)
  {
    _J = _assembly.getJacobian(_elem, qp);
    _JxW = _assembly.getJxW(_elem, qp);
    _xyz = _assembly.getQPoint(_elem, qp);
  }
  else
    _assembly.computePhysicalMap(_elem, _side, qp, &_J, &_JxW, &_xyz);

  _transform_reinit = true;
}

class ResidualDatum : public Datum
{
public:
  KOKKOS_FUNCTION
  ResidualDatum(const ContiguousElementID elem,
                const unsigned int side,
                const Assembly & assembly,
                const Array<System> & systems,
                const Variable & ivar,
                const unsigned int jvar,
                const unsigned int comp = 0)
    : Datum(elem, side, assembly, systems),
      _tag(ivar.tag()),
      _ivar(ivar.var(comp)),
      _jvar(jvar),
      _ife(systems[ivar.sys(comp)].getFETypeID(_ivar)),
      _jfe(systems[ivar.sys(comp)].getFETypeID(_jvar)),
      _n_idofs(assembly.getNumDofs(_elem.type, _ife)),
      _n_jdofs(assembly.getNumDofs(_elem.type, _jfe))
  {
  }
  KOKKOS_FUNCTION
  ResidualDatum(const ContiguousNodeID node,
                const Assembly & assembly,
                const Array<System> & systems,
                const Variable & ivar,
                const unsigned int jvar,
                const unsigned int comp = 0)
    : Datum(node, assembly, systems),
      _tag(ivar.tag()),
      _ivar(ivar.var(comp)),
      _jvar(jvar),
      _ife(systems[ivar.sys(comp)].getFETypeID(_ivar)),
      _jfe(systems[ivar.sys(comp)].getFETypeID(_jvar))
  {
  }

  KOKKOS_FUNCTION unsigned int n_dofs() const { return _n_idofs; }
  KOKKOS_FUNCTION unsigned int n_idofs() const { return _n_idofs; }
  KOKKOS_FUNCTION unsigned int n_jdofs() const { return _n_jdofs; }
  KOKKOS_FUNCTION unsigned int var() const { return _ivar; }
  KOKKOS_FUNCTION unsigned int ivar() const { return _ivar; }
  KOKKOS_FUNCTION unsigned int jvar() const { return _jvar; }
  KOKKOS_FUNCTION unsigned int fe() const { return _ife; }
  KOKKOS_FUNCTION unsigned int ife() const { return _ife; }
  KOKKOS_FUNCTION unsigned int jfe() const { return _jfe; }

protected:
  const TagID _tag;
  const unsigned int _ivar, _jvar;
  const unsigned int _ife, _jfe;
  const unsigned int _n_idofs = 1, _n_jdofs = 1;
};

} // namespace Kokkos
} // namespace Moose

using Datum = Moose::Kokkos::Datum;
using ResidualDatum = Moose::Kokkos::ResidualDatum;