Line data    Source code

       1             : //* This file is part of the MOOSE framework
       2             : //* https://www.mooseframework.org
       3             : //*
       4             : //* All rights reserved, see COPYRIGHT for full restrictions
       5             : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
       6             : //*
       7             : //* Licensed under LGPL 2.1, please see LICENSE for details
       8             : //* https://www.gnu.org/licenses/lgpl-2.1.html
       9             : 
      10             : #pragma once
      11             : 
      12             : #include "KokkosTypes.h"
      13             : 
      14             : #ifdef MOOSE_KOKKOS_SCOPE
      15             : #include "KokkosUtils.h"
      16             : #endif
      17             : 
      18             : using ContiguousSubdomainID = SubdomainID;
      19             : using ContiguousBoundaryID = BoundaryID;
      20             : using ContiguousElementID = dof_id_type;
      21             : using ContiguousNodeID = dof_id_type;
      22             : 
      23             : class MooseMesh;
      24             : 
      25             : namespace Moose
      26             : {
      27             : namespace Kokkos
      28             : {
      29             : 
      30             : /**
      31             :  * The Kokkos object that contains the information of an element
      32             :  * The IDs used in Kokkos are different from the MOOSE or libMesh IDs
      33             :  * The Kokkos IDs start from zero in each process and are contiguous
      34             :  */
      35             : struct ElementInfo
      36             : {
      37             :   /**
      38             :    * Element type ID
      39             :    */
      40             :   unsigned int type;
      41             :   /**
      42             :    * Contiguous element ID
      43             :    */
      44             :   ContiguousElementID id;
      45             :   /**
      46             :    * Contiguous subdomain ID
      47             :    */
      48             :   ContiguousSubdomainID subdomain;
      49             : };
      50             : 
      51             : /**
      52             :  * The Kokkos mesh object
      53             :  */
      54             : class Mesh
      55             : {
      56             : public:
      57             :   /**
      58             :    * Constructor
      59             :    * @param mesh The MOOSE mesh
      60             :    */
      61       46979 :   Mesh(const MooseMesh & mesh) : _mesh(mesh) {}
      62             :   /**
      63             :    * Get the underyling MOOSE mesh
      64             :    * @returns The MOOSE mesh
      65             :    */
      66             :   const MooseMesh & getMesh() { return _mesh; }
      67             :   /**
      68             :    * Update the mesh
      69             :    */
      70             :   void update();
      71             : 
      72             :   /**
      73             :    * Get the contiguous subdomain ID of a MOOSE subdomain
      74             :    * @param subdomain The MOOSE subdomain ID
      75             :    * @returns The contiguous subdomain ID
      76             :    */
      77             :   ContiguousSubdomainID getContiguousSubdomainID(const SubdomainID subdomain) const;
      78             :   /**
      79             :    * Get the contiguous boundary ID of a boundary
      80             :    * @param boundary The MOOSE boundary ID
      81             :    * @returns The contiguous boundary ID
      82             :    */
      83             :   ContiguousBoundaryID getContiguousBoundaryID(const BoundaryID boundary) const;
      84             :   /**
      85             :    * Get the element type ID of an element
      86             :    * @param elem The libMesh element
      87             :    * @returns The element type ID
      88             :    */
      89             :   unsigned int getElementTypeID(const Elem * elem) const;
      90             :   /**
      91             :    * Get the contiguous element ID of an element
      92             :    * @param elem The libMesh element
      93             :    * @returns The contiguous element ID
      94             :    */
      95             :   ContiguousElementID getContiguousElementID(const Elem * elem) const;
      96             :   /**
      97             :    * Get the element type ID map
      98             :    * @returns The element type ID map
      99             :    */
     100        6420 :   const auto & getElementTypeMap() const { return _maps->_elem_type_id_mapping; }
     101             :   /**
     102             :    * Get the contiguous element ID map
     103             :    * @returns The contiguous element ID map
     104             :    */
     105       71500 :   const auto & getContiguousElementMap() const { return _maps->_local_elem_id_mapping; }
     106             :   /**
     107             :    * Get the list of contiguous element IDs for a subdomain
     108             :    * @param subdomain The MOOSE subdomain ID
     109             :    * @returns The list of contiguous element IDs in the subdomain
     110             :    */
     111        3836 :   const auto & getSubdomainContiguousElementIDs(const SubdomainID subdomain) const
     112             :   {
     113        3836 :     return _maps->_subdomain_elem_ids.at(subdomain);
     114             :   }
     115             :   /**
     116             :    * Get the contiguous node ID of a node
     117             :    * @param node The libMesh node
     118             :    * @returns The contiguous node ID that starts from zero in each process
     119             :    */
     120             :   ContiguousNodeID getContiguousNodeID(const Node * node) const;
     121             :   /**
     122             :    * Get the contiguous node ID map
     123             :    * This list contains the nodes of local elements, so some nodes may belong to other processes
     124             :    * @returns The contiguous node ID map
     125             :    */
     126        5352 :   const auto & getContiguousNodeMap() const { return _maps->_local_node_id_mapping; }
     127             :   /**
     128             :    * Get the list of contiguous node IDs for a subdomain
     129             :    * This list strictly contains the nodes local to the current process
     130             :    * @param subdomain The MOOSE subdomain ID
     131             :    * @returns The list of contiguous node IDs in the subdomain
     132             :    */
     133        3286 :   const auto & getSubdomainContiguousNodeIDs(const SubdomainID subdomain) const
     134             :   {
     135        3286 :     return _maps->_subdomain_node_ids.at(subdomain);
     136             :   }
     137             : 
     138             : #ifdef MOOSE_KOKKOS_SCOPE
     139             :   /**
     140             :    * Get the element information object
     141             :    * @param elem The contiguous element ID
     142             :    * @returns The element information object
     143             :    */
     144    21884527 :   KOKKOS_FUNCTION const auto & getElementInfo(ContiguousElementID elem) const
     145             :   {
     146    21884527 :     return _elem_info[elem];
     147             :   }
     148             :   /**
     149             :    * Get the neighbor contiguous element ID
     150             :    * @param elem The contiguous element ID
     151             :    * @param side The side index
     152             :    * @returns The neighbor contiguous element ID
     153             :    */
     154       94030 :   KOKKOS_FUNCTION ContiguousElementID getNeighbor(ContiguousElementID elem, unsigned int side) const
     155             :   {
     156       94030 :     return _elem_neighbor(side, elem);
     157             :   }
     158             :   /**
     159             :    * Get the number of sides of an element type
     160             :    * @param elem_type The element type ID
     161             :    * @returns The number of sides of the element type
     162             :    */
     163           0 :   KOKKOS_FUNCTION unsigned int getNumSides(unsigned int elem_type) const
     164             :   {
     165           0 :     return _num_sides[elem_type];
     166             :   }
     167             :   /**
     168             :    * Get the number of nodes of an element type
     169             :    * @param elem_type The element type ID
     170             :    * @returns The number of nodes of the element type
     171             :    */
     172      183742 :   KOKKOS_FUNCTION unsigned int getNumNodes(unsigned int elem_type) const
     173             :   {
     174      183742 :     return _num_nodes[elem_type];
     175             :   }
     176             :   /**
     177             :    * Get the number of nodes on a side of an element type
     178             :    * @param elem_type The element type ID
     179             :    * @param side The side index
     180             :    * @returns The number of nodes on the side of the element type
     181             :    */
     182       89454 :   KOKKOS_FUNCTION unsigned int getNumNodes(unsigned int elem_type, unsigned int side) const
     183             :   {
     184       89454 :     return _num_side_nodes[elem_type][side];
     185             :   }
     186             :   /**
     187             :    * Get the contiguous node ID for an element
     188             :    * @param elem The contiguous element ID
     189             :    * @param node The node index
     190             :    * @returns The contiguous node ID
     191             :    */
     192      765188 :   KOKKOS_FUNCTION ContiguousNodeID getContiguousNodeID(ContiguousElementID elem,
     193             :                                                        unsigned int node) const
     194             :   {
     195      765188 :     return _nodes(node, elem);
     196             :   }
     197             :   /**
     198             :    * Get the contiguous node ID for a side
     199             :    * @param elem The contiguous element ID
     200             :    * @param side The side index
     201             :    * @param node The node index
     202             :    * @returns The contiguous node ID
     203             :    */
     204      178382 :   KOKKOS_FUNCTION ContiguousNodeID getContiguousNodeID(ContiguousElementID elem,
     205             :                                                        unsigned int side,
     206             :                                                        unsigned int node) const
     207             :   {
     208      178382 :     return _nodes_face(node, side, elem);
     209             :   }
     210             :   /**
     211             :    * Get the coordinate of a node
     212             :    * @param node The contiguous node ID
     213             :    * @returns The node coordinate
     214             :    */
     215      943570 :   KOKKOS_FUNCTION Real3 getNodePoint(ContiguousNodeID node) const { return _points[node]; }
     216             :   /**
     217             :    * Get whether a node is on a boundary
     218             :    * @param node The contiguous node ID
     219             :    * @param boundary The contiguous boundary ID
     220             :    * @returns Whether the node is on the boundary
     221             :    */
     222             :   KOKKOS_FUNCTION bool isBoundaryNode(ContiguousNodeID node, ContiguousBoundaryID boundary) const;
     223             : #endif
     224             : 
     225             : private:
     226             :   /**
     227             :    * Initialize host maps
     228             :    */
     229             :   void initMap();
     230             :   /**
     231             :    * Initialize device element data
     232             :    */
     233             :   void initElement();
     234             : 
     235             :   /**
     236             :    * Reference of the MOOSE mesh
     237             :    */
     238             :   const MooseMesh & _mesh;
     239             :   /**
     240             :    * The wrapper of host maps
     241             :    */
     242             :   struct MeshMap
     243             :   {
     244             :     /**
     245             :      * Map from the MOOSE subdomain ID to the contiguous subdomain ID
     246             :      */
     247             :     std::unordered_map<SubdomainID, ContiguousSubdomainID> _subdomain_id_mapping;
     248             :     /**
     249             :      * Map from the MOOSE boundary ID to the contiguous boundary ID
     250             :      */
     251             :     std::unordered_map<BoundaryID, ContiguousBoundaryID> _boundary_id_mapping;
     252             :     /**
     253             :      * Map from the MOOSE element type to the element type ID
     254             :      */
     255             :     std::unordered_map<ElemType, unsigned int> _elem_type_id_mapping;
     256             :     /**
     257             :      * Map from the libMesh element to the contiguous element ID
     258             :      */
     259             :     std::unordered_map<const Elem *, ContiguousElementID> _local_elem_id_mapping;
     260             :     /**
     261             :      * Map from the libMesh node to the contiguous node ID
     262             :      * This list contains the nodes of local elements, so some nodes may belong to other processes
     263             :      */
     264             :     std::unordered_map<const Node *, ContiguousNodeID> _local_node_id_mapping;
     265             :     /**
     266             :      * List of the contiguous element IDs in each subdomain
     267             :      */
     268             :     std::unordered_map<SubdomainID, std::unordered_set<ContiguousElementID>> _subdomain_elem_ids;
     269             :     /**
     270             :      * List of the contiguous node IDs in each subdomain
     271             :      * This list strictly contains the nodes local to the current process
     272             :      */
     273             :     std::unordered_map<SubdomainID, std::unordered_set<ContiguousNodeID>> _subdomain_node_ids;
     274             :   };
     275             :   /**
     276             :    * A shared pointer holding all the host maps to avoid deep copy
     277             :    */
     278             :   std::shared_ptr<MeshMap> _maps;
     279             : 
     280             :   /**
     281             :    * Element information
     282             :    */
     283             :   Array<ElementInfo> _elem_info;
     284             :   /**
     285             :    * Neighbor contiguous element IDs of each element
     286             :    */
     287             :   Array2D<ContiguousElementID> _elem_neighbor;
     288             :   /**
     289             :    * Number of sides of each element type
     290             :    */
     291             :   Array<unsigned int> _num_sides;
     292             :   /**
     293             :    * Number of nodes of each element type
     294             :    */
     295             :   Array<unsigned int> _num_nodes;
     296             :   /**
     297             :    * number of nodes per side of each element side
     298             :    */
     299             :   Array<Array<unsigned int>> _num_side_nodes;
     300             :   /**
     301             :    * Node coordinates
     302             :    */
     303             :   Array<Real3> _points;
     304             :   /**
     305             :    * Contiguous node IDs of each element and side
     306             :    */
     307             :   ///@{
     308             :   Array2D<ContiguousNodeID> _nodes;
     309             :   Array3D<ContiguousNodeID> _nodes_face;
     310             :   ///@}
     311             :   /**
     312             :    * Contiguous node IDs on each boundary
     313             :    */
     314             :   Array<Array<ContiguousNodeID>> _boundary_nodes;
     315             : };
     316             : 
     317             : #ifdef MOOSE_KOKKOS_SCOPE
     318             : KOKKOS_FUNCTION inline bool
     319      664470 : Mesh::isBoundaryNode(ContiguousNodeID node, ContiguousBoundaryID boundary) const
     320             : {
     321      664470 :   if (!_boundary_nodes[boundary].size())
     322      106880 :     return false;
     323             : 
     324      557590 :   auto begin = &_boundary_nodes[boundary].begin();
     325      557590 :   auto end = &_boundary_nodes[boundary].end();
     326      557590 :   auto target = Utils::find(node, begin, end);
     327             : 
     328      557590 :   return target != end;
     329             : }
     330             : #endif
     331             : 
     332             : /**
     333             :  * The Kokkos interface that holds the host reference of the Kokkos mesh and copies it to device
     334             :  * during parallel dispatch
     335             :  * Maintains synchronization between host and device Kokkos mesh and provides access to the
     336             :  * appropriate Kokkos mesh depending on the architecture
     337             :  */
     338             : class MeshHolder
     339             : {
     340             : public:
     341             :   /**
     342             :    * Constructor
     343             :    * @param assembly The Kokkos mesh
     344             :    */
     345       47981 :   MeshHolder(const Mesh & mesh) : _mesh_host(mesh), _mesh_device(mesh) {}
     346             :   /**
     347             :    * Copy constructor
     348             :    */
     349      520902 :   MeshHolder(const MeshHolder & holder)
     350      423619 :     : _mesh_host(holder._mesh_host), _mesh_device(holder._mesh_host)
     351             :   {
     352      520902 :   }
     353             : 
     354             : #ifdef MOOSE_KOKKOS_SCOPE
     355             :   /**
     356             :    * Get the const reference of the Kokkos mesh
     357             :    * @returns The const reference of the Kokkos mesh depending on the architecture this function
     358             :    * is being called on
     359             :    */
     360    25044966 :   KOKKOS_FUNCTION const Mesh & kokkosMesh() const
     361             :   {
     362    25044966 :     KOKKOS_IF_ON_HOST(return _mesh_host;)
     363             : 
     364    24803363 :     return _mesh_device;
     365             :   }
     366             : #endif
     367             : 
     368             : private:
     369             :   /**
     370             :    * Host reference of the Kokkos mesh
     371             :    */
     372             :   const Mesh & _mesh_host;
     373             :   /**
     374             :    * Device copy of the Kokkos mesh
     375             :    */
     376             :   const Mesh _mesh_device;
     377             : };
     378             : 
     379             : } // namespace Kokkos
     380             : } // namespace Moose