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             : #include "KokkosMaterialPropertyStorage.h"
      11             : 
      12             : #include "FEProblemBase.h"
      13             : #include "MaterialBase.h"
      14             : 
      15             : namespace Moose
      16             : {
      17             : namespace Kokkos
      18             : {
      19             : 
      20             : std::unordered_map<std::type_index, PropertyStore> MaterialPropertyStorage::_store_functions;
      21             : std::unordered_map<std::type_index, PropertyLoad> MaterialPropertyStorage::_load_functions;
      22             : 
      23      128010 : MaterialPropertyStorage::MaterialPropertyStorage(MaterialPropertyRegistry & registry,
      24      130800 :                                                  FEProblemBase & problem)
      25      508320 :   : ::MaterialPropertyStorage(registry, problem)
      26             : {
      27      128010 : }
      28             : 
      29             : MaterialPropertyStorage &
      30        6033 : MaterialPropertyStorage::cast(::MaterialPropertyStorage & storage)
      31             : {
      32        6033 :   return static_cast<MaterialPropertyStorage &>(storage);
      33             : }
      34             : 
      35             : MaterialPropertyBase &
      36        3140 : MaterialPropertyStorage::addKokkosProperty(const std::string & prop_name,
      37             :                                            const std::type_info & type,
      38             :                                            const unsigned int state,
      39             :                                            const ::MaterialBase * declarer,
      40             :                                            std::shared_ptr<MaterialPropertyBase> shell)
      41             : {
      42        3140 :   auto id = addProperty(prop_name, type, state, declarer);
      43             : 
      44        3140 :   if (state)
      45        1760 :     return addKokkosPropertyState(prop_name, state, shell);
      46             : 
      47        1380 :   auto & record = _kokkos_prop_records[prop_name];
      48             : 
      49        1380 :   if (record.id == libMesh::invalid_uint)
      50             :   {
      51         840 :     record.type = MooseUtils::prettyCppType(libMesh::demangle(type.name()));
      52         840 :     record.id = id;
      53             :   }
      54             : 
      55        1380 :   if (declarer)
      56        1023 :     record.declarers.insert(declarer);
      57             : 
      58        1380 :   if (!_kokkos_props[0].count(prop_name))
      59             :   {
      60         840 :     _kokkos_props[0][prop_name] = shell;
      61         840 :     _kokkos_props[0][prop_name]->init(record, {});
      62             :   }
      63             : 
      64        1380 :   return *_kokkos_props[0][prop_name];
      65             : }
      66             : 
      67             : MaterialPropertyBase &
      68        2532 : MaterialPropertyStorage::addKokkosPropertyState(const std::string & prop_name,
      69             :                                                 const unsigned int state,
      70             :                                                 std::shared_ptr<MaterialPropertyBase> shell)
      71             : {
      72        2532 :   auto & record = _kokkos_prop_records[prop_name];
      73             : 
      74        2532 :   if (!_kokkos_props[state].count(prop_name))
      75             :   {
      76         642 :     _kokkos_props[state][prop_name] = shell;
      77         642 :     _kokkos_props[state][prop_name]->init(record, {});
      78             :   }
      79             : 
      80        2532 :   return *_kokkos_props[state][prop_name];
      81             : }
      82             : 
      83             : MaterialPropertyBase &
      84        1027 : MaterialPropertyStorage::declareKokkosProperty(const std::string & prop_name,
      85             :                                                const std::type_info & type,
      86             :                                                const ::MaterialBase * declarer,
      87             :                                                const std::vector<unsigned int> & dims,
      88             :                                                const bool bnd,
      89             :                                                std::shared_ptr<MaterialPropertyBase> shell)
      90             : {
      91        1027 :   auto & record = _kokkos_prop_records[prop_name];
      92             : 
      93        1027 :   if (record.declarers.size() && record.dims != dims)
      94             :   {
      95           4 :     std::stringstream ss;
      96             : 
      97           4 :     ss << "The declared " << dims.size() << "D Kokkos material property '" << prop_name << "'";
      98           4 :     if (dims.size())
      99          20 :       ss << " with dimensions (" << Moose::stringify(dims) << ")";
     100           4 :     ss << " was already declared as a " << record.dims.size() << "D property";
     101           4 :     if (record.dims.size())
     102          20 :       ss << " with dimensions (" << Moose::stringify(record.dims) << ")";
     103           4 :     ss << ".";
     104             : 
     105           4 :     mooseError(ss.str());
     106           0 :   }
     107             : 
     108        1023 :   if (!record.declarers.size())
     109             :   {
     110         834 :     record.dims = dims;
     111         834 :     record.bnd = bnd;
     112             :   }
     113             : 
     114        1023 :   return addKokkosProperty(prop_name, type, 0, declarer, shell);
     115             : }
     116             : 
     117             : MaterialPropertyBase &
     118         357 : MaterialPropertyStorage::getKokkosProperty(std::string prop_name, unsigned int state)
     119             : {
     120         357 :   return *_kokkos_props[state][prop_name];
     121             : }
     122             : 
     123             : void
     124        2079 : MaterialPropertyStorage::allocateKokkosProperties()
     125             : {
     126        2079 :   auto & mesh = _problem.mesh();
     127        2079 :   auto & assembly = _problem.kokkosAssembly();
     128             : 
     129        4800 :   for (unsigned int state = 0; state < numStates(); ++state)
     130        4185 :     for (auto & [name, prop] : _kokkos_props[state])
     131             :     {
     132        1464 :       auto & record = _kokkos_prop_records[name];
     133             : 
     134        3171 :       for (auto declarer : record.declarers)
     135             :       {
     136         348 :         auto & blocks =
     137        1359 :             declarer->boundaryRestricted() ? declarer->meshBlockIDs() : declarer->blockIDs();
     138             : 
     139        1707 :         prop->allocate(mesh, assembly, blocks, record.bnd, {});
     140             :       }
     141             :     }
     142        2079 : }
     143             : 
     144             : bool
     145           0 : MaterialPropertyStorage::haveKokkosProperty(std::string prop_name, unsigned int state)
     146             : {
     147           0 :   return _kokkos_props[state].count(prop_name);
     148             : }
     149             : 
     150             : void
     151        1760 : MaterialPropertyStorage::registerLoadStore(std::type_index type,
     152             :                                            PropertyStore store,
     153             :                                            PropertyLoad load)
     154             : {
     155        1760 :   _store_functions[type] = store;
     156        1760 :   _load_functions[type] = load;
     157        1760 : }
     158             : 
     159             : void
     160        1178 : MaterialPropertyStorage::shift()
     161             : {
     162        3046 :   for (unsigned int state = maxState(); state != 0; --state)
     163        3736 :     for (auto & [name, prop] : _kokkos_props[state])
     164        1868 :       prop->swap(*_kokkos_props[state - 1][name], {});
     165        1178 : }
     166             : 
     167             : void
     168         600 : MaterialPropertyStorage::copy()
     169             : {
     170        1212 :   for (unsigned int state : statefulIndexRange())
     171        1224 :     for (auto & [name, prop] : _kokkos_props[state])
     172         612 :       if (!isRestoredProperty(name))
     173         484 :         prop->copy(*_kokkos_props[0][name], {});
     174         600 : }
     175             : 
     176             : } // namespace Kokkos
     177             : } // namespace Moose
     178             : 
     179             : void
     180      100470 : dataStore(std::ostream & stream, Moose::Kokkos::MaterialPropertyStorage & storage, void * context)
     181             : {
     182      100470 :   auto & base = static_cast<MaterialPropertyStorage &>(storage);
     183             : 
     184      100470 :   dataStore(stream, base, context);
     185             : 
     186      100470 :   auto & props = storage._kokkos_props;
     187      100470 :   auto & records = storage._kokkos_prop_records;
     188             : 
     189      100470 :   unsigned int num_stateful_properties = 0;
     190             : 
     191      100728 :   for (auto & prop : props[0])
     192             :   {
     193         258 :     auto name = prop.first;
     194         258 :     auto id = records[name].id;
     195             : 
     196         258 :     if (storage.getPropRecord(id).stateful())
     197         190 :       num_stateful_properties++;
     198         258 :   }
     199             : 
     200      100470 :   dataStore(stream, num_stateful_properties, nullptr);
     201             : 
     202      100728 :   for (auto & prop : props[0])
     203             :   {
     204         258 :     auto name = prop.first;
     205         258 :     auto id = records[name].id;
     206             : 
     207         258 :     if (storage.getPropRecord(id).stateful())
     208             :     {
     209         190 :       dataStore(stream, name, nullptr);
     210             : 
     211         190 :       unsigned int num_states = storage.getPropRecord(id).state + 1;
     212         190 :       dataStore(stream, num_states, nullptr);
     213             : 
     214             :       // Store data in a temporary stream to allow skipping loading
     215         190 :       std::stringstream ss;
     216             : 
     217         630 :       for (unsigned int state = 0; state < num_states; ++state)
     218             :       {
     219         440 :         auto it = storage._store_functions.find(props[state][name]->propertyType());
     220             :         mooseAssert(it != storage._store_functions.end(),
     221             :                     "Store function was not properly registered.");
     222             : 
     223         440 :         it->second(ss, props[state][name].get());
     224             :       }
     225             : 
     226             :       // Write the temporary stream to the output stream
     227         190 :       dataStore(stream, ss, nullptr);
     228         190 :     }
     229         258 :   }
     230      100470 : }
     231             : 
     232             : void
     233       32697 : dataLoad(std::istream & stream, Moose::Kokkos::MaterialPropertyStorage & storage, void * context)
     234             : {
     235       32697 :   auto & base = static_cast<MaterialPropertyStorage &>(storage);
     236             : 
     237             :   // All exceptions are expected to be handled here
     238       32697 :   dataLoad(stream, base, context);
     239             : 
     240       32691 :   auto & props = storage._kokkos_props;
     241       32691 :   auto & records = storage._kokkos_prop_records;
     242             : 
     243             :   unsigned int num_stateful_properties;
     244       32691 :   dataLoad(stream, num_stateful_properties, nullptr);
     245             : 
     246       32813 :   for (unsigned int prop = 0; prop < num_stateful_properties; ++prop)
     247             :   {
     248         122 :     std::string name;
     249         122 :     dataLoad(stream, name, nullptr);
     250             : 
     251             :     unsigned int num_states;
     252         122 :     dataLoad(stream, num_states, nullptr);
     253             : 
     254         122 :     std::stringstream ss;
     255         122 :     dataLoad(stream, ss, nullptr);
     256             : 
     257         122 :     if (records.count(name) && storage.getPropRecord(records[name].id).stateful())
     258         324 :       for (unsigned int state = 0; state < num_states; ++state)
     259             :       {
     260         226 :         auto it = storage._load_functions.find(props[state][name]->propertyType());
     261             :         mooseAssert(it != storage._load_functions.end(),
     262             :                     "Load function was not properly registered.");
     263             : 
     264         226 :         it->second(ss, props[state][name].get());
     265             :       }
     266         122 :   }
     267       32691 : }