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 "KokkosFESystem.h"
#include "KokkosVariable.h"

namespace Moose::Kokkos
{

class MeshDatum
{
public:
  KOKKOS_FUNCTION
  MeshDatum(ContiguousElementID elem, const unsigned int side, const Mesh & mesh);

  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 bool hasNeighbor() const;

  KOKKOS_FUNCTION const ElementInfo & neighbor() const { return _neighbor; }

  KOKKOS_FUNCTION ContiguousElementID neighborID() const { return _neighbor.id; }

  KOKKOS_FUNCTION unsigned int side() const { return _side; }

  KOKKOS_FUNCTION bool isSide() const { return _side != libMesh::invalid_uint; }

  KOKKOS_FUNCTION ContiguousSubdomainID subdomain() const { return _elem.subdomain; }

  KOKKOS_FUNCTION ContiguousSubdomainID neighborSubdomain() const { return _neighbor.subdomain; }

protected:
  static KOKKOS_FUNCTION ElementInfo neighborInfo(ContiguousElementID elem,
                                                  const unsigned int side,
                                                  const Mesh & mesh);

  const Mesh & _mesh;

  const ElementInfo _elem;

  const unsigned int _side = libMesh::invalid_uint;

  const ElementInfo _neighbor;
};

KOKKOS_FUNCTION inline MeshDatum::MeshDatum(ContiguousElementID elem,
                                            const unsigned int side,
                                            const Mesh & mesh)
  : _mesh(mesh),
    _elem(elem != libMesh::DofObject::invalid_id ? _mesh.getElementInfo(elem) : ElementInfo{}),
    _side(side),
    _neighbor(neighborInfo(_elem.id, _side, _mesh))
{
}

KOKKOS_FUNCTION inline bool
MeshDatum::hasNeighbor() const
{
  return _neighbor.id != libMesh::DofObject::invalid_id;
}

KOKKOS_FUNCTION inline ElementInfo
MeshDatum::neighborInfo(ContiguousElementID elem, const unsigned int side, const Mesh & mesh)
{
  if (elem == libMesh::DofObject::invalid_id || side == libMesh::invalid_uint)
    return {};

  const auto neighbor = mesh.getNeighbor(elem, side);
  return neighbor == libMesh::DofObject::invalid_id ? ElementInfo{} : mesh.getElementInfo(neighbor);
}

class FVDatum : public MeshDatum
{
public:
  KOKKOS_FUNCTION
  FVDatum(ContiguousElementID elem, const unsigned int side, const Mesh & mesh);

  KOKKOS_FUNCTION Real3 faceCentroid() const { return _mesh.getSideCentroid(_elem.id, _side); }

  KOKKOS_FUNCTION Real3 faceNormal() const { return _mesh.getSideNormal(_elem.id, _side); }

  KOKKOS_FUNCTION Real faceDCNMag() const
  {
    return _mesh.getElementCentroidToNeighborCentroidDistance(_elem.id, _side);
  }

  KOKKOS_FUNCTION Real faceDCFMag() const
  {
    return _mesh.getElementCentroidToSideCentroidDistance(_elem.id, _side);
  }

  KOKKOS_FUNCTION Real faceArea() const { return _mesh.getSideArea(_elem.id, _side); }

  KOKKOS_FUNCTION Real3 elementCentroid() const { return _mesh.getElementCentroid(_elem.id); }

  KOKKOS_FUNCTION Real elementVolume() const { return _mesh.getElementVolume(_elem.id); }
};

KOKKOS_FUNCTION inline FVDatum::FVDatum(ContiguousElementID elem,
                                        const unsigned int side,
                                        const Mesh & mesh)
  : MeshDatum(elem, side, mesh)
{
}

class Datum : public MeshDatum
{
public:
  KOKKOS_FUNCTION
  Datum(const ContiguousElementID elem,
        const unsigned int side,
        const Assembly & assembly,
        const Array<FESystem> & systems);

  KOKKOS_FUNCTION
  Datum(const ContiguousNodeID node, const Assembly & assembly, const Array<FESystem> & systems);

  KOKKOS_FUNCTION const Assembly & assembly() const { return _assembly; }

  KOKKOS_FUNCTION const FESystem & system(unsigned int sys) const { return _systems[sys]; }

  KOKKOS_FUNCTION dof_id_type extraElemID(unsigned int index) const
  {
    return isNodal() ? libMesh::DofObject::invalid_id : _mesh.getExtraElementID(_elem.id, index);
  }

  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 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<FESystem> & _systems;

  const ContiguousNodeID _node = 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 Datum::Datum(const ContiguousElementID elem,
                                    const unsigned int side,
                                    const Assembly & assembly,
                                    const Array<FESystem> & systems)
  : MeshDatum(elem, side, assembly.kokkosMesh()),
    _assembly(assembly),
    _systems(systems),
    _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 inline Datum::Datum(const ContiguousNodeID node,
                                    const Assembly & assembly,
                                    const Array<FESystem> & systems)
  : MeshDatum(libMesh::DofObject::invalid_id, libMesh::invalid_uint, assembly.kokkosMesh()),
    _assembly(assembly),
    _systems(systems),
    _node(node)
{
}

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<FESystem> & 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<FESystem> & 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;
using FVDatum = Moose::Kokkos::FVDatum;