KokkosAssembly.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 "MooseMesh.h"

#include "libmesh/elem_range.h"
#include "libmesh/fe_base.h"
#include "libmesh/fe_type.h"

class FEProblemBase;

namespace Moose
{
namespace Kokkos
{

class Assembly : public MeshHolder
{
public:
  Assembly(FEProblemBase & problem);
  void init();

#ifdef MOOSE_KOKKOS_SCOPE

  unsigned int getFETypeID(FEType type) const { return libmesh_map_find(_fe_type_map, type); }
  KOKKOS_FUNCTION unsigned int getDimension() const { return _dimension; }
  KOKKOS_FUNCTION unsigned int getMaxQpsPerElem() const { return _max_qps_per_elem; }
  KOKKOS_FUNCTION dof_id_type getNumQps(ContiguousSubdomainID subdomain) const
  {
    return _n_subdomain_qps[subdomain];
  }
  KOKKOS_FUNCTION unsigned int getNumQps(ElementInfo info) const { return _n_qps[info.id]; }
  KOKKOS_FUNCTION dof_id_type getNumFaceQps(ContiguousSubdomainID subdomain) const
  {
    return _n_subdomain_qps_face[subdomain];
  }
  KOKKOS_FUNCTION unsigned int getNumFaceQps(ElementInfo info, unsigned int side) const
  {
    return _n_qps_face(side, info.id);
  }
  KOKKOS_FUNCTION dof_id_type getQpOffset(ElementInfo info) const { return _qp_offset[info.id]; }
  KOKKOS_FUNCTION dof_id_type getQpFaceOffset(ElementInfo info, unsigned int side) const
  {
    return _qp_offset_face(side, info.id);
  }
  KOKKOS_FUNCTION unsigned int getNumDofs(unsigned int elem_type, unsigned int fe_type) const
  {
    return _n_dofs(elem_type, fe_type);
  }
  KOKKOS_FUNCTION const auto &
  getPhi(ContiguousSubdomainID subdomain, unsigned int elem_type, unsigned int fe_type) const
  {
    return _phi(subdomain, elem_type, fe_type);
  }
  KOKKOS_FUNCTION const auto &
  getPhiFace(ContiguousSubdomainID subdomain, unsigned int elem_type, unsigned int fe_type) const
  {
    return _phi_face(subdomain, elem_type, fe_type);
  }
  KOKKOS_FUNCTION const auto &
  getGradPhi(ContiguousSubdomainID subdomain, unsigned int elem_type, unsigned int fe_type) const
  {
    return _grad_phi(subdomain, elem_type, fe_type);
  }
  KOKKOS_FUNCTION const auto & getGradPhiFace(ContiguousSubdomainID subdomain,
                                              unsigned int elem_type,
                                              unsigned int fe_type) const
  {
    return _grad_phi_face(subdomain, elem_type, fe_type);
  }
  KOKKOS_FUNCTION Real33 getJacobian(ElementInfo info, unsigned int qp) const
  {
    return _jacobian[info.subdomain][getQpOffset(info) + qp];
  }
  KOKKOS_FUNCTION Real getJxW(ElementInfo info, unsigned int qp) const
  {
    return _jxw[info.subdomain][getQpOffset(info) + qp];
  }
  KOKKOS_FUNCTION Real3 getQPoint(ElementInfo info, unsigned int qp) const
  {
    return _xyz[info.subdomain][getQpOffset(info) + qp];
  }

  KOKKOS_FUNCTION Real coordTransformFactor(const ContiguousSubdomainID subdomain,
                                            const Real3 point) const;
  KOKKOS_FUNCTION void computePhysicalMap(const ElementInfo info,
                                          const unsigned int qp,
                                          Real33 * const jacobian,
                                          Real * const JxW,
                                          Real3 * const q_points) const;
  KOKKOS_FUNCTION void computePhysicalMap(const ElementInfo info,
                                          const unsigned int side,
                                          const unsigned int qp,
                                          Real33 * const jacobian,
                                          Real * const JxW,
                                          Real3 * const q_points) const;

  KOKKOS_FUNCTION void operator()(const ThreadID tid) const;

  const auto & getMaterialBoundaries() const { return _material_boundaries; }
#endif

private:
  void initQuadrature();
  void initShape();
  void cachePhysicalMap();

  FEProblemBase & _problem;
  MooseMesh & _mesh;
  std::map<FEType, unsigned int> _fe_type_map;

  const unsigned int _dimension;
  Array<Moose::CoordinateSystemType> _coord_type;
  unsigned int _rz_radial_coord = libMesh::invalid_uint;
  Array<Pair<Real3, Real3>> _rz_axis;

  Array<dof_id_type> _qp_offset;
  Array2D<dof_id_type> _qp_offset_face;

  Array<unsigned int> _n_qps;
  Array2D<unsigned int> _n_qps_face;

  unsigned int _max_qps_per_elem = 0;

  Array<dof_id_type> _n_subdomain_qps;
  Array<dof_id_type> _n_subdomain_qps_face;

  Array2D<Array<Real3>> _q_points;
  Array2D<Array<Array<Real3>>> _q_points_face;
  Array2D<Array<Real>> _weights;
  Array2D<Array<Array<Real>>> _weights_face;

  Array3D<Array2D<Real>> _phi;
  Array3D<Array<Array2D<Real>>> _phi_face;
  Array3D<Array2D<Real3>> _grad_phi;
  Array3D<Array<Array2D<Real3>>> _grad_phi_face;
  Array2D<unsigned int> _n_dofs;

  Array2D<Array2D<Real>> _map_phi;
  Array2D<Array<Array2D<Real>>> _map_phi_face;
  Array2D<Array<Array2D<Real>>> _map_psi_face;
  Array2D<Array2D<Real3>> _map_grad_phi;
  Array2D<Array<Array2D<Real3>>> _map_grad_phi_face;
  Array2D<Array<Array2D<Real3>>> _map_grad_psi_face;

  Array<Array<Real33>> _jacobian;
  Array<Array<Real>> _jxw;
  Array<Array<Real3>> _xyz;

  std::set<BoundaryID> _material_boundaries;
};

#ifdef MOOSE_KOKKOS_SCOPE
KOKKOS_FUNCTION inline Real
Assembly::coordTransformFactor(const ContiguousSubdomainID subdomain, const Real3 point) const
{
  switch (_coord_type[subdomain])
  {
    case Moose::COORD_XYZ:
      return 1;
    case Moose::COORD_RZ:
      if (_rz_radial_coord == libMesh::invalid_uint)
        return 2 * M_PI *
               (point - _rz_axis[subdomain].first).cross_product(_rz_axis[subdomain].second).norm();
      else
        return 2 * M_PI * point(_rz_radial_coord);
    case Moose::COORD_RSPHERICAL:
      return 4 * M_PI * point(0) * point(0);
    default:
      return 0;
  }
}

KOKKOS_FUNCTION inline void
Assembly::computePhysicalMap(const ElementInfo info,
                             const unsigned int qp,
                             Real33 * const jacobian,
                             Real * const JxW,
                             Real3 * const q_points) const
{
  auto sid = info.subdomain;
  auto eid = info.id;
  auto elem_type = info.type;
  auto num_nodes = kokkosMesh().getNumNodes(elem_type);

  auto & phi = _map_phi(sid, elem_type);
  auto & grad_phi = _map_grad_phi(sid, elem_type);

  Real33 J;
  Real3 xyz;

  for (unsigned int node = 0; node < num_nodes; ++node)
  {
    auto points = kokkosMesh().getNodePoint(kokkosMesh().getContiguousNodeID(eid, node));

    if (jacobian || JxW)
      J += grad_phi(node, qp).cartesian_product(points);

    xyz += phi(node, qp) * points;
  }

  if (jacobian)
    *jacobian = J.inverse(_dimension);
  if (JxW)
    *JxW =
        J.determinant(_dimension) * _weights(sid, elem_type)[qp] * coordTransformFactor(sid, xyz);
  if (q_points)
    *q_points = xyz;
}

KOKKOS_FUNCTION inline void
Assembly::computePhysicalMap(const ElementInfo info,
                             const unsigned int side,
                             const unsigned int qp,
                             Real33 * const jacobian,
                             Real * const JxW,
                             Real3 * const q_points) const
{
  auto sid = info.subdomain;
  auto eid = info.id;
  auto elem_type = info.type;
  auto num_nodes = kokkosMesh().getNumNodes(elem_type);
  auto num_side_nodes = kokkosMesh().getNumNodes(elem_type, side);

  auto & phi = _map_phi_face(sid, elem_type)(side);
  auto & grad_phi = _map_grad_phi_face(sid, elem_type)(side);

  Real33 J;
  Real3 xyz;

  for (unsigned int node = 0; node < num_nodes; ++node)
  {
    auto points = kokkosMesh().getNodePoint(kokkosMesh().getContiguousNodeID(eid, node));

    if (jacobian)
      J += grad_phi(node, qp).cartesian_product(points);

    if (JxW || q_points)
      xyz += phi(node, qp) * points;
  }

  if (jacobian)
    *jacobian = J.inverse(_dimension);
  if (q_points)
    *q_points = xyz;

  if (JxW)
  {
    Real33 J;

    auto & grad_psi = _map_grad_psi_face(sid, elem_type)(side);

    for (unsigned int node = 0; node < num_side_nodes; ++node)
    {
      auto points = kokkosMesh().getNodePoint(kokkosMesh().getContiguousNodeID(eid, side, node));

      J += grad_psi(node, qp).cartesian_product(points);
    }

    *JxW = ::Kokkos::sqrt((J * J.transpose()).determinant(_dimension - 1)) *
           _weights_face(sid, elem_type)[side][qp] * coordTransformFactor(sid, xyz);
  }
}
#endif

class AssemblyHolder
{
public:
  AssemblyHolder(const Assembly & assembly) : _assembly_host(assembly), _assembly_device(assembly)
  {
  }
  AssemblyHolder(const AssemblyHolder & holder)
    : _assembly_host(holder._assembly_host), _assembly_device(holder._assembly_host)
  {
  }

#ifdef MOOSE_KOKKOS_SCOPE

  KOKKOS_FUNCTION const Assembly & kokkosAssembly() const
  {
    KOKKOS_IF_ON_HOST(return _assembly_host;)

    return _assembly_device;
  }
#endif

private:
  const Assembly & _assembly_host;
  const Assembly _assembly_device;
};

} // namespace Kokkos
} // namespace Moose