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