Line data    Source code

       1             : //* This file is part of the MOOSE framework
       2             : //* https://mooseframework.inl.gov
       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 "KokkosFunctorWrapper.h"
      13             : #include "KokkosFunctionWrapper.h"
      14             : 
      15             : namespace Moose::Kokkos
      16             : {
      17             : 
      18             : class FunctorRegistryEntryBase
      19             : {
      20             : public:
      21           0 :   virtual ~FunctorRegistryEntryBase() {}
      22             :   /**
      23             :    * Build a host wrapper for this functor
      24             :    * @param object The pointer to the functor
      25             :    */
      26             :   virtual std::unique_ptr<FunctorWrapperHostBase> build(const void * object) const = 0;
      27             : };
      28             : 
      29             : template <typename Object>
      30             : class FunctorRegistryEntry : public FunctorRegistryEntryBase
      31             : {
      32             : public:
      33           0 :   std::unique_ptr<FunctorWrapperHostBase> build(const void * object) const override final
      34             :   {
      35           0 :     return std::make_unique<FunctorWrapperHost<Object>>(object);
      36             :   }
      37             : };
      38             : 
      39             : class FunctionRegistryEntryBase
      40             : {
      41             : public:
      42           0 :   virtual ~FunctionRegistryEntryBase() {}
      43             :   /**
      44             :    * Build a host wrapper for this function
      45             :    * @param object The pointer to the function
      46             :    */
      47             :   virtual std::unique_ptr<FunctionWrapperHostBase> build(const void * object) const = 0;
      48             : };
      49             : 
      50             : template <typename Object>
      51             : class FunctionRegistryEntry : public FunctionRegistryEntryBase
      52             : {
      53             : public:
      54         119 :   std::unique_ptr<FunctionWrapperHostBase> build(const void * object) const override final
      55             :   {
      56         119 :     return std::make_unique<FunctionWrapperHost<Object>>(object);
      57             :   }
      58             : };
      59             : 
      60             : class FunctorRegistry
      61             : {
      62             : public:
      63       40355 :   FunctorRegistry() = default;
      64             : 
      65             :   FunctorRegistry(FunctorRegistry const &) = delete;
      66             :   FunctorRegistry & operator=(FunctorRegistry const &) = delete;
      67             : 
      68             :   FunctorRegistry(FunctorRegistry &&) = delete;
      69             :   FunctorRegistry & operator=(FunctorRegistry &&) = delete;
      70             : 
      71             :   /**
      72             :    * Register a functor
      73             :    * @tparam Object The functor class type
      74             :    * @param name The registered functor type name
      75             :    */
      76             :   template <typename Object>
      77      121065 :   static char addFunctor(const std::string & name)
      78             :   {
      79      121065 :     getRegistry()._functors[name] = std::make_unique<FunctorRegistryEntry<Object>>();
      80             : 
      81      121065 :     return 0;
      82             :   }
      83             : 
      84             :   /**
      85             :    * Register a function
      86             :    * @tparam Object The function class type
      87             :    * @param name The registered function type name
      88             :    */
      89             :   template <typename Object>
      90      121067 :   static char addFunction(const std::string & name)
      91             :   {
      92      121067 :     getRegistry()._functions[name] = std::make_unique<FunctionRegistryEntry<Object>>();
      93             : 
      94      121067 :     return 0;
      95             :   }
      96             : 
      97             :   /**
      98             :    * Build and get a host wrapper of a functor
      99             :    * @param object The pointer to the functor
     100             :    * @param name The registered functor type name
     101             :    * @returns The host wrapper
     102             :    */
     103             :   static std::unique_ptr<FunctorWrapperHostBase> buildFunctor(const void * object,
     104             :                                                               const std::string & name)
     105             :   {
     106             :     auto it = getRegistry()._functors.find(name);
     107             :     if (it == getRegistry()._functors.end())
     108             :       mooseError("Kokkos functor not registered for type '",
     109             :                  name,
     110             :                  "'. Double check that you used Kokkos-specific registration macro.");
     111             : 
     112             :     return it->second->build(object);
     113             :   }
     114             : 
     115             :   /**
     116             :    * Build and get a host wrapper of a function
     117             :    * @param object The pointer to the function
     118             :    * @param name The registered function type name
     119             :    * @returns The host wrapper
     120             :    */
     121         119 :   static std::unique_ptr<FunctionWrapperHostBase> buildFunction(const void * object,
     122             :                                                                 const std::string & name)
     123             :   {
     124         119 :     auto it = getRegistry()._functions.find(name);
     125         119 :     if (it == getRegistry()._functions.end())
     126           0 :       mooseError("Kokkos function not registered for type '",
     127             :                  name,
     128             :                  "'. Double check that you used Kokkos-specific registration macro.");
     129             : 
     130         238 :     return it->second->build(object);
     131             :   }
     132             : 
     133             : private:
     134             :   /**
     135             :    * Get the registry singleton
     136             :    * @returns The registry singleton
     137             :    */
     138             :   static FunctorRegistry & getRegistry();
     139             : 
     140             :   /**
     141             :    * Map containing the host wrapper shells of functors with the key being the registered object
     142             :    * type name
     143             :    */
     144             :   std::map<std::string, std::unique_ptr<FunctorRegistryEntryBase>> _functors;
     145             :   /**
     146             :    * Map containing the host wrapper shells of functions with the key being the registered object
     147             :    * type name
     148             :    */
     149             :   std::map<std::string, std::unique_ptr<FunctionRegistryEntryBase>> _functions;
     150             : };
     151             : 
     152             : } // namespace Moose::Kokkos
     153             : 
     154             : #define registerKokkosFunction(app, classname)                                                     \
     155             :   registerMooseObject(app, classname);                                                             \
     156             :   static char combineNames(kokkos_functor_##classname, __COUNTER__) =                              \
     157             :       Moose::Kokkos::FunctorRegistry::addFunctor<classname>(#classname);                           \
     158             :   static char combineNames(kokkos_function_##classname, __COUNTER__) =                             \
     159             :       Moose::Kokkos::FunctorRegistry::addFunction<classname>(#classname)
     160             : 
     161             : #define registerKokkosFunctionAliased(app, classname, alias)                                       \
     162             :   registerMooseObjectAliased(app, classname, alias);                                               \
     163             :   static char combineNames(kokkos_functor_##classname, __COUNTER__) =                              \
     164             :       Moose::Kokkos::FunctorRegistry::addFunctor<classname>(alias);                                \
     165             :   static char combineNames(kokkos_function_##classname, __COUNTER__) =                             \
     166             :       Moose::Kokkos::FunctorRegistry::addFunction<classname>(alias)