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::Kokkos
      26             : {
      27             : 
      28             : /**
      29             :  * The Kokkos object that contains the information of an element
      30             :  * The IDs used in Kokkos are different from the MOOSE or libMesh IDs
      31             :  * The Kokkos IDs start from zero in each process and are contiguous
      32             :  */
      33             : struct ElementInfo
      34             : {
      35             :   /**
      36             :    * Element type ID
      37             :    */
      38             :   unsigned int type = libMesh::invalid_uint;
      39             :   /**
      40             :    * Contiguous element ID
      41             :    */
      42             :   ContiguousElementID id = libMesh::DofObject::invalid_id;
      43             :   /**
      44             :    * Contiguous subdomain ID
      45             :    */
      46             :   ContiguousSubdomainID subdomain = std::numeric_limits<ContiguousSubdomainID>::max();
      47             : };
      48             : 
      49             : /**
      50             :  * The Kokkos mesh object
      51             :  */
      52             : class Mesh
      53             : {
      54             : public:
      55             :   /**
      56             :    * Constructor
      57             :    * @param mesh The MOOSE mesh
      58             :    */
      59       51430 :   Mesh(MooseMesh & mesh) : _mesh(mesh) {}
      60             :   /**
      61             :    * Get the underyling MOOSE mesh
      62             :    * @returns The MOOSE mesh
      63             :    */
      64             :   const MooseMesh & getMesh() { return _mesh; }
      65             :   /**
      66             :    * Get whether the mesh was initialized
      67             :    */
      68     5325946 :   [[nodiscard]] bool initialized() const { return _initialized; }
      69             :   /**
      70             :    * Update the mesh
      71             :    */
      72             :   void update();
      73             : 
      74             :   /**
      75             :    * Get the number of subdomains
      76             :    * @returns The number of subdomains
      77             :    */
      78       36100 :   auto getNumSubdomains() const { return _maps->subdomain_id_mapping.size(); }
      79             :   /**
      80             :    * Get the number of local elements types
      81             :    * @returns The number of local element types
      82             :    */
      83        6795 :   auto getNumLocalElementTypes() const { return _maps->elem_type_id_mapping.size(); }
      84             :   /**
      85             :    * Get the element type ID map
      86             :    * @returns The element type ID map
      87             :    */
      88       11655 :   const auto & getElementTypeMap() const { return _maps->elem_type_id_mapping; }
      89             :   /**
      90             :    * Get the number of local elements
      91             :    * @returns The number of local elements
      92             :    */
      93      199162 :   auto getNumLocalElements() const { return _num_local_elems; }
      94             :   /**
      95             :    * Get the number of local elements in a MOOSE subdomain
      96             :    * @param subdomain The MOOSE subdomain ID
      97             :    * @returns The local number of elements in the subdomain
      98             :    */
      99         321 :   auto getNumSubdomainLocalElements(const SubdomainID subdomain) const
     100             :   {
     101         321 :     auto range = libmesh_map_find(_maps->subdomain_elem_id_ranges, subdomain);
     102         321 :     return range.second - range.first;
     103             :   }
     104             :   /**
     105             :    * Get the number of local nodes including semi-local nodes
     106             :    * @returns The number of local nodes including semi-local nodes
     107             :    */
     108        6795 :   auto getNumLocalNodes() const { return _num_local_nodes; }
     109             :   /**
     110             :    * Get the list of local nodes including semi-local nodes
     111             :    * @returns The list of local nodes including semi-local nodes
     112             :    */
     113        6284 :   const auto & getLocalNodes() const { return _maps->local_nodes; }
     114             :   /**
     115             :    * Get the contiguous subdomain ID of a MOOSE subdomain
     116             :    * @param subdomain The MOOSE subdomain ID
     117             :    * @returns The contiguous subdomain ID
     118             :    */
     119             :   ContiguousSubdomainID getContiguousSubdomainID(const SubdomainID subdomain) const;
     120             :   /**
     121             :    * Get the contiguous boundary ID of a boundary
     122             :    * @param boundary The MOOSE boundary ID
     123             :    * @returns The contiguous boundary ID
     124             :    */
     125             :   ContiguousBoundaryID getContiguousBoundaryID(const BoundaryID boundary) const;
     126             :   /**
     127             :    * Get the element type ID of an element
     128             :    * @param elem The libMesh element
     129             :    * @returns The element type ID
     130             :    */
     131             :   unsigned int getElementTypeID(const Elem * elem) const;
     132             :   /**
     133             :    * Get the contiguous element ID of an element
     134             :    * @param elem The libMesh element
     135             :    * @returns The contiguous element ID
     136             :    */
     137             :   ContiguousElementID getContiguousElementID(const Elem * elem) const;
     138             :   /**
     139             :    * Get the range of contiguous element IDs for a subdomain
     140             :    * @param subdomain The MOOSE subdomain ID
     141             :    * @returns The range of contiguous element IDs in the subdomain
     142             :    */
     143       17980 :   auto getSubdomainContiguousElementIDRange(const SubdomainID subdomain) const
     144             :   {
     145       17980 :     const auto & range = libmesh_map_find(_maps->subdomain_elem_id_ranges, subdomain);
     146       17980 :     return libMesh::make_range(range.first, range.second);
     147             :   }
     148             :   /**
     149             :    * Get the contiguous node ID of a node
     150             :    * @param node The libMesh node
     151             :    * @returns The contiguous node ID that starts from zero in each process
     152             :    */
     153             :   ContiguousNodeID getContiguousNodeID(const Node * node) const;
     154             :   /**
     155             :    * Get the list of local contiguous node IDs for a subdomain
     156             :    * NOTE: This list excludes semi-local nodes
     157             :    * @param subdomain The MOOSE subdomain ID
     158             :    * @returns The list of local contiguous node IDs in the subdomain
     159             :    */
     160       11994 :   const auto & getSubdomainContiguousNodeIDs(const SubdomainID subdomain) const
     161             :   {
     162       11994 :     return libmesh_map_find(_maps->subdomain_node_ids, subdomain);
     163             :   }
     164             :   /**
     165             :    * Get the list of local contiguous node IDs for a boundary
     166             :    * NOTE: This list excludes semi-local nodes
     167             :    * @param boundary The MOOSE boundary ID
     168             :    * @returns The list of local contiguous node IDs on the boundary
     169             :    */
     170       13668 :   const auto & getBoundaryContiguousNodeIDs(const BoundaryID boundary) const
     171             :   {
     172       13668 :     return libmesh_map_find(_maps->boundary_node_ids, boundary);
     173             :   }
     174             : #ifdef MOOSE_KOKKOS_SCOPE
     175             :   /**
     176             :    * Get the element information object
     177             :    * @param elem The contiguous element ID
     178             :    * @returns The element information object
     179             :    */
     180    91022215 :   KOKKOS_FUNCTION const auto & getElementInfo(ContiguousElementID elem) const
     181             :   {
     182    91022215 :     return _elem_info[elem];
     183             :   }
     184             :   /**
     185             :    * Get the neighbor contiguous element ID
     186             :    * @param elem The contiguous element ID
     187             :    * @param side The side index
     188             :    * @returns The neighbor contiguous element ID
     189             :    */
     190      578628 :   KOKKOS_FUNCTION ContiguousElementID getNeighbor(ContiguousElementID elem, unsigned int side) const
     191             :   {
     192      578628 :     return _elem_neighbor(side, elem);
     193             :   }
     194             :   /**
     195             :    * Get the number of sides of an element type
     196             :    * @param elem_type The element type ID
     197             :    * @returns The number of sides of the element type
     198             :    */
     199           0 :   KOKKOS_FUNCTION unsigned int getNumSides(unsigned int elem_type) const
     200             :   {
     201           0 :     return _num_sides[elem_type];
     202             :   }
     203             :   /**
     204             :    * Get the number of nodes of an element type
     205             :    * @param elem_type The element type ID
     206             :    * @returns The number of nodes of the element type
     207             :    */
     208     1682796 :   KOKKOS_FUNCTION unsigned int getNumNodes(unsigned int elem_type) const
     209             :   {
     210     1682796 :     return _num_nodes[elem_type];
     211             :   }
     212             :   /**
     213             :    * Get the number of nodes on a side of an element type
     214             :    * @param elem_type The element type ID
     215             :    * @param side The side index
     216             :    * @returns The number of nodes on the side of the element type
     217             :    */
     218     1063280 :   KOKKOS_FUNCTION unsigned int getNumNodes(unsigned int elem_type, unsigned int side) const
     219             :   {
     220     1063280 :     return _num_side_nodes[elem_type][side];
     221             :   }
     222             :   /**
     223             :    * Get the extra element ID of an element
     224             :    * @param elem The contiguous element ID
     225             :    * @param index The extra element ID index
     226             :    * @returns The extra element ID
     227             :    */
     228         500 :   KOKKOS_FUNCTION dof_id_type getExtraElementID(ContiguousElementID elem, unsigned int index) const
     229             :   {
     230             :     KOKKOS_ASSERT(index < _extra_elem_ids.n(1));
     231             : 
     232         500 :     return _extra_elem_ids(elem, index);
     233             :   }
     234             :   /**
     235             :    * Get the starting contiguous element ID of a subdomain
     236             :    * @param subdomain The contiguous subdomain ID
     237             :    * @returns The starting contiguous element ID
     238             :    */
     239    22175078 :   KOKKOS_FUNCTION dof_id_type getStartingContiguousElementID(ContiguousSubdomainID subdomain) const
     240             :   {
     241    22175078 :     return _starting_elem_id[subdomain];
     242             :   }
     243             :   /**
     244             :    * Get the contiguous node ID for an element
     245             :    * @param elem The contiguous element ID
     246             :    * @param node The node index
     247             :    * @returns The contiguous node ID
     248             :    */
     249     6956460 :   KOKKOS_FUNCTION ContiguousNodeID getContiguousNodeID(ContiguousElementID elem,
     250             :                                                        unsigned int node) const
     251             :   {
     252     6956460 :     return _nodes(node, elem);
     253             :   }
     254             :   /**
     255             :    * Get the contiguous node ID for a side
     256             :    * @param elem The contiguous element ID
     257             :    * @param side The side index
     258             :    * @param node The node index
     259             :    * @returns The contiguous node ID
     260             :    */
     261     2137244 :   KOKKOS_FUNCTION ContiguousNodeID getContiguousNodeID(ElementInfo info,
     262             :                                                        unsigned int side,
     263             :                                                        unsigned int node) const
     264             :   {
     265     2137244 :     return _nodes(_local_side_node[info.type](node, side), info.id);
     266             :   }
     267             :   /**
     268             :    * Get the coordinate of a node
     269             :    * @param node The contiguous node ID
     270             :    * @returns The node coordinate
     271             :    */
     272     9534070 :   KOKKOS_FUNCTION Real3 getNodePoint(ContiguousNodeID node) const { return _points[node]; }
     273             :   /**
     274             :    * Get whether a node is on a boundary
     275             :    * @param node The contiguous node ID
     276             :    * @param boundary The contiguous boundary ID
     277             :    * @returns Whether the node is on the boundary
     278             :    */
     279             :   KOKKOS_FUNCTION bool isBoundaryNode(ContiguousNodeID node, ContiguousBoundaryID boundary) const;
     280             : #endif
     281             : 
     282             : private:
     283             :   /**
     284             :    * Initialize host maps
     285             :    */
     286             :   void initMap();
     287             :   /**
     288             :    * Initialize device element data
     289             :    */
     290             :   void initElement();
     291             : 
     292             :   /**
     293             :    * Reference of the MOOSE mesh
     294             :    */
     295             :   MooseMesh & _mesh;
     296             :   /**
     297             :    * Flag whether the mesh was initialized
     298             :    */
     299             :   bool _initialized = false;
     300             : 
     301             :   /**
     302             :    * The wrapper of host maps
     303             :    */
     304             :   struct MeshMap
     305             :   {
     306             :     /**
     307             :      * Map from the MOOSE subdomain ID to the contiguous subdomain ID
     308             :      */
     309             :     std::unordered_map<SubdomainID, ContiguousSubdomainID> subdomain_id_mapping;
     310             :     /**
     311             :      * Map from the MOOSE boundary ID to the contiguous boundary ID
     312             :      */
     313             :     std::unordered_map<BoundaryID, ContiguousBoundaryID> boundary_id_mapping;
     314             :     /**
     315             :      * Map from the MOOSE element type to the element type ID
     316             :      */
     317             :     std::unordered_map<ElemType, unsigned int> elem_type_id_mapping;
     318             :     /**
     319             :      * List of local nodes including semi-local nodes
     320             :      */
     321             :     std::vector<Node *> local_nodes;
     322             :     /**
     323             :      * Map from the libMesh semi-local node to the contiguous node ID
     324             :      */
     325             :     std::unordered_map<const Node *, ContiguousNodeID> semi_local_node_id_mapping;
     326             :     /**
     327             :      * Range of the contiguous element IDs in each subdomain
     328             :      */
     329             :     std::unordered_map<SubdomainID, std::pair<ContiguousElementID, ContiguousElementID>>
     330             :         subdomain_elem_id_ranges;
     331             :     /**
     332             :      * List of the contiguous node IDs in each subdomain
     333             :      * NOTE: This list excludes semi-local nodes
     334             :      */
     335             :     std::unordered_map<SubdomainID, std::vector<ContiguousNodeID>> subdomain_node_ids;
     336             :     /**
     337             :      * List of the contiguous node IDs on each boundary
     338             :      * NOTE: This list excludes semi-local nodes
     339             :      */
     340             :     std::unordered_map<BoundaryID, std::set<ContiguousNodeID>> boundary_node_ids;
     341             :   };
     342             :   /**
     343             :    * A shared pointer holding all the host maps to avoid deep copy
     344             :    */
     345             :   std::shared_ptr<MeshMap> _maps;
     346             : 
     347             :   /**
     348             :    * Element integer for Kokkos contiguous element ID
     349             :    */
     350             :   unsigned int _elem_id_integer = libMesh::invalid_uint;
     351             :   /**
     352             :    * Node integer for Kokkos contiguous node ID
     353             :    */
     354             :   unsigned int _node_id_integer = libMesh::invalid_uint;
     355             :   /**
     356             :    * Number of local elements
     357             :    */
     358             :   dof_id_type _num_local_elems = 0;
     359             :   /**
     360             :    * Number of local nodes including semi-local nodes
     361             :    */
     362             :   dof_id_type _num_local_nodes = 0;
     363             :   /**
     364             :    * Element information
     365             :    */
     366             :   Array<ElementInfo> _elem_info;
     367             :   /**
     368             :    * Neighbor contiguous element IDs of each element
     369             :    */
     370             :   Array2D<ContiguousElementID> _elem_neighbor;
     371             :   /**
     372             :    * Starting contiguous element ID of each subdomain
     373             :    */
     374             :   Array<ContiguousElementID> _starting_elem_id;
     375             :   /**
     376             :    * Extra element IDs
     377             :    */
     378             :   Array2D<dof_id_type> _extra_elem_ids;
     379             :   /**
     380             :    * Number of sides of each element type
     381             :    */
     382             :   Array<unsigned int> _num_sides;
     383             :   /**
     384             :    * Number of nodes of each element type
     385             :    */
     386             :   Array<unsigned int> _num_nodes;
     387             :   /**
     388             :    * Number of nodes per side of each element side
     389             :    */
     390             :   Array<Array<unsigned int>> _num_side_nodes;
     391             :   /**
     392             :    * Map from local side node index to local element node index
     393             :    */
     394             :   Array<Array2D<unsigned int>> _local_side_node;
     395             :   /**
     396             :    * Node coordinates
     397             :    */
     398             :   Array<Real3> _points;
     399             :   /**
     400             :    * Contiguous node IDs of each element
     401             :    */
     402             :   Array2D<ContiguousNodeID> _nodes;
     403             :   /**
     404             :    * Contiguous node IDs on each boundary
     405             :    */
     406             :   Array<Array<ContiguousNodeID>> _boundary_nodes;
     407             : };
     408             : 
     409             : #ifdef MOOSE_KOKKOS_SCOPE
     410             : KOKKOS_FUNCTION inline bool
     411      664470 : Mesh::isBoundaryNode(ContiguousNodeID node, ContiguousBoundaryID boundary) const
     412             : {
     413      664470 :   if (!_boundary_nodes[boundary].size())
     414      106880 :     return false;
     415             : 
     416      557590 :   auto begin = &_boundary_nodes[boundary].begin();
     417      557590 :   auto end = &_boundary_nodes[boundary].end();
     418      557590 :   auto target = Utils::find(node, begin, end);
     419             : 
     420      557590 :   return target != end;
     421             : }
     422             : #endif
     423             : 
     424             : /**
     425             :  * The Kokkos interface that holds the host reference of the Kokkos mesh and copies it to device
     426             :  * during parallel dispatch
     427             :  * Maintains synchronization between host and device Kokkos mesh and provides access to the
     428             :  * appropriate Kokkos mesh depending on the architecture
     429             :  */
     430             : class MeshHolder
     431             : {
     432             : public:
     433             :   /**
     434             :    * Constructor
     435             :    * @param mesh The Kokkos mesh
     436             :    */
     437       63245 :   MeshHolder(const Mesh & mesh) : _mesh_host(mesh), _mesh_device(mesh) {}
     438             :   /**
     439             :    * Copy constructor
     440             :    */
     441     1416947 :   MeshHolder(const MeshHolder & holder)
     442      816625 :     : _mesh_host(holder._mesh_host), _mesh_device(holder._mesh_host)
     443             :   {
     444     1416947 :   }
     445             : 
     446             : #ifdef MOOSE_KOKKOS_SCOPE
     447             :   /**
     448             :    * Get the const reference of the Kokkos mesh
     449             :    * @returns The const reference of the Kokkos mesh depending on the architecture this function
     450             :    * is being called on
     451             :    */
     452   121828825 :   KOKKOS_FUNCTION const Mesh & kokkosMesh() const
     453             :   {
     454   121828825 :     KOKKOS_IF_ON_HOST(
     455             :         if (!_mesh_host.initialized()) mooseError(
     456             :             "kokkosMesh() was called too early. Kokkos mesh is available after problem "
     457             :             "initialization. Override initialSetup() if you need to setup your object data "
     458             :             "using the Kokkos mesh.");
     459             : 
     460             :         return _mesh_host;)
     461             : 
     462   116502879 :     return _mesh_device;
     463             :   }
     464             : #endif
     465             : 
     466             : private:
     467             :   /**
     468             :    * Host reference of the Kokkos mesh
     469             :    */
     470             :   const Mesh & _mesh_host;
     471             :   /**
     472             :    * Device copy of the Kokkos mesh
     473             :    */
     474             :   const Mesh _mesh_device;
     475             : };
     476             : 
     477             : } // namespace Moose::Kokkos