KokkosFunctorRegistry.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://mooseframework.inl.gov
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#pragma once

#include "KokkosFunctorWrapper.h"
#include "KokkosFunctionWrapper.h"

namespace Moose::Kokkos
{

class FunctorRegistryEntryBase
{
public:
  virtual ~FunctorRegistryEntryBase() {}
  virtual std::unique_ptr<FunctorWrapperHostBase> build(const void * object) const = 0;
};

template <typename Object>
class FunctorRegistryEntry : public FunctorRegistryEntryBase
{
public:
  std::unique_ptr<FunctorWrapperHostBase> build(const void * object) const override final
  {
    return std::make_unique<FunctorWrapperHost<Object>>(object);
  }
};

class FunctionRegistryEntryBase
{
public:
  virtual ~FunctionRegistryEntryBase() {}
  virtual std::unique_ptr<FunctionWrapperHostBase> build(const void * object) const = 0;
};

template <typename Object>
class FunctionRegistryEntry : public FunctionRegistryEntryBase
{
public:
  std::unique_ptr<FunctionWrapperHostBase> build(const void * object) const override final
  {
    return std::make_unique<FunctionWrapperHost<Object>>(object);
  }
};

class FunctorRegistry
{
public:
  FunctorRegistry() = default;

  FunctorRegistry(FunctorRegistry const &) = delete;
  FunctorRegistry & operator=(FunctorRegistry const &) = delete;

  FunctorRegistry(FunctorRegistry &&) = delete;
  FunctorRegistry & operator=(FunctorRegistry &&) = delete;

  template <typename Object>
  static char addFunctor(const std::string & name)
  {
    getRegistry()._functors[name] = std::make_unique<FunctorRegistryEntry<Object>>();

    return 0;
  }

  template <typename Object>
  static char addFunction(const std::string & name)
  {
    getRegistry()._functions[name] = std::make_unique<FunctionRegistryEntry<Object>>();

    return 0;
  }

  static std::unique_ptr<FunctorWrapperHostBase> buildFunctor(const void * object,
                                                              const std::string & name)
  {
    auto it = getRegistry()._functors.find(name);
    if (it == getRegistry()._functors.end())
      mooseError("Kokkos functor not registered for type '",
                 name,
                 "'. Double check that you used Kokkos-specific registration macro.");

    return it->second->build(object);
  }

  static std::unique_ptr<FunctionWrapperHostBase> buildFunction(const void * object,
                                                                const std::string & name)
  {
    auto it = getRegistry()._functions.find(name);
    if (it == getRegistry()._functions.end())
      mooseError("Kokkos function not registered for type '",
                 name,
                 "'. Double check that you used Kokkos-specific registration macro.");

    return it->second->build(object);
  }

private:
  static FunctorRegistry & getRegistry();

  std::map<std::string, std::unique_ptr<FunctorRegistryEntryBase>> _functors;
  std::map<std::string, std::unique_ptr<FunctionRegistryEntryBase>> _functions;
};

} // namespace Moose::Kokkos

#define registerKokkosFunction(app, classname)                                                     \
  registerMooseObject(app, classname);                                                             \
  static char combineNames(kokkos_functor_##classname, __COUNTER__) =                              \
      Moose::Kokkos::FunctorRegistry::addFunctor<classname>(#classname);                           \
  static char combineNames(kokkos_function_##classname, __COUNTER__) =                             \
      Moose::Kokkos::FunctorRegistry::addFunction<classname>(#classname)

#define registerKokkosFunctionAliased(app, classname, alias)                                       \
  registerMooseObjectAliased(app, classname, alias);                                               \
  static char combineNames(kokkos_functor_##classname, __COUNTER__) =                              \
      Moose::Kokkos::FunctorRegistry::addFunctor<classname>(alias);                                \
  static char combineNames(kokkos_function_##classname, __COUNTER__) =                             \
      Moose::Kokkos::FunctorRegistry::addFunction<classname>(alias)