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::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),
      _mesh(assembly.kokkosMesh()),
      _elem(_mesh.getElementInfo(elem)),
      _side(side),
      _neighbor(!isSide() ? libMesh::DofObject::invalid_id : _mesh.getNeighbor(_elem.id, side)),
      _n_qps(!isSide() ? assembly.getNumQps(_elem) : assembly.getNumFaceQps(_elem, side)),
      _qp_offset(!isSide() ? assembly.getQpOffset(_elem) : assembly.getQpFaceOffset(_elem, side)),
      _elem_property_idx(!isSide()
                             ? _elem.id - _mesh.getStartingContiguousElementID(_elem.subdomain)
                             : assembly.getElemFacePropertyIndex(_elem, _side))
  {
  }
  KOKKOS_FUNCTION
  Datum(const ContiguousNodeID node, const Assembly & assembly, const Array<System> & systems)
    : _assembly(assembly), _systems(systems), _mesh(assembly.kokkosMesh()), _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 Mesh & mesh() const { return _mesh; }

  KOKKOS_FUNCTION const ElementInfo & elem() const { return _elem; }
  KOKKOS_FUNCTION ContiguousElementID elemID() const { return _elem.id; }
  KOKKOS_FUNCTION dof_id_type extraElemID(unsigned int index) const
  {
    return isNodal() ? libMesh::DofObject::invalid_id : _mesh.getExtraElementID(_elem.id, index);
  }
  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 dof_id_type propertyIdx(const PropertyConstantOption constant_option,
                                          const unsigned int qp) const;
  KOKKOS_FUNCTION bool hasNeighbor() const { return _neighbor != libMesh::DofObject::invalid_id; }
  KOKKOS_FUNCTION bool isSide() const { return _side != libMesh::invalid_uint; }
  KOKKOS_FUNCTION bool isNodal() const { return _node != libMesh::DofObject::invalid_id; }
  KOKKOS_FUNCTION bool isNodalDefined(const Variable & var) const;

  KOKKOS_FUNCTION const Real33 & J(const unsigned int qp);
  KOKKOS_FUNCTION Real JxW(const unsigned int qp);
  KOKKOS_FUNCTION Real3 q_point(const unsigned int qp);
  KOKKOS_FUNCTION Real3 normals(const unsigned int qp);

  KOKKOS_FUNCTION void set_local_parallel(const unsigned int local_thread_id,
                                          const unsigned int num_local_threads)
  {
    _local_thread_id = local_thread_id;
    _num_local_threads = num_local_threads;
  }
  KOKKOS_FUNCTION unsigned int local_thread_id() const { return _local_thread_id; }
  KOKKOS_FUNCTION unsigned int num_local_threads() const { return _num_local_threads; }

protected:
  const Assembly & _assembly;
  const Array<System> & _systems;
  const Mesh & _mesh;
  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;
  const dof_id_type _elem_property_idx = libMesh::DofObject::invalid_id;

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

  unsigned int _cached_qp = libMesh::invalid_uint;
  Real33 _J;
  Real _JxW;
  Real3 _xyz;
  Real3 _normal;

  unsigned int _local_thread_id = 0;
  unsigned int _num_local_threads = 1;
};

KOKKOS_FUNCTION inline dof_id_type
Datum::propertyIdx(const PropertyConstantOption constant_option, const unsigned int qp) const
{
  dof_id_type idx = 0;

  if (constant_option == PropertyConstantOption::NONE)
    idx = _qp_offset + qp;
  else if (constant_option == PropertyConstantOption::ELEMENT)
    idx = _elem_property_idx;

  return idx;
}

KOKKOS_FUNCTION inline bool
Datum::isNodalDefined(const Variable & var) const
{
  if (!isNodal() || !var.nodal())
    return false;

  return _systems[var.sys()].isNodalDefined(_node, var.var());
}

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

  return _J;
}

KOKKOS_FUNCTION inline Real
Datum::JxW(const unsigned int qp)
{
  if (!isNodal())
    reinitTransform(qp);
  else
    _JxW = 1;

  return _JxW;
}

KOKKOS_FUNCTION inline Real3
Datum::q_point(const unsigned int qp)
{
  if (!isNodal())
    reinitTransform(qp);
  else
    _xyz = _assembly.kokkosMesh().getNodePoint(_node);

  return _xyz;
}

KOKKOS_FUNCTION inline Real3
Datum::normals(const unsigned int qp)
{
  KOKKOS_ASSERT(isSide());

  if (isSide())
    reinitTransform(qp);

  return _normal;
}

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

  if (!isSide())
  {
    _J = _assembly.getJacobian(_elem, qp);
    _JxW = _assembly.getJxW(_elem, qp);
    _xyz = _assembly.getQPoint(_elem, qp);
  }
  else
    _assembly.computePhysicalMap(_elem, _side, qp, &_J, &_JxW, &_xyz, &_normal);

  _cached_qp = qp;
}

class AssemblyDatum : public Datum
{
public:
  KOKKOS_FUNCTION
  AssemblyDatum(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()),
      _sys(ivar.sys(comp)),
      _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
  AssemblyDatum(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()),
      _sys(ivar.sys(comp)),
      _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 sys() const { return _sys; }
  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; }
  KOKKOS_FUNCTION void do_derivatives(const bool flag) { _do_derivatives = flag; }
  KOKKOS_FUNCTION bool do_derivatives() const { return _do_derivatives; }

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

} // namespace Moose::Kokkos

using Datum = Moose::Kokkos::Datum;
using AssemblyDatum = Moose::Kokkos::AssemblyDatum;