KokkosDispatcher.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* 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 "KokkosHeader.h"
#include "KokkosThread.h"

#include <typeindex>

namespace Moose::Kokkos
{

using Policy = ::Kokkos::RangePolicy<ExecSpace, ::Kokkos::IndexType<ThreadID>>;

class DispatcherBase
{
public:
  virtual ~DispatcherBase() {}
  virtual void parallelFor(const Policy & /* policy */)
  {
    mooseError("parallelFor() called for an instance that is not a dispatcher.");
  }
  virtual void parallelReduce(const Policy & /* policy */,
                              ::Kokkos::View<Real *, ::Kokkos::HostSpace> & /* result */)
  {
    mooseError("parallelReduce() called for an instance that is not a reducer.");
  }
};

template <typename Operation, typename Object>
class Dispatcher : public DispatcherBase
{
public:
  Dispatcher(const void * object)
    : _functor_host(*static_cast<const Object *>(object)), _functor_device(_functor_host)
  {
  }
  Dispatcher(const Dispatcher & functor)
    : _functor_host(functor._functor_host), _functor_device(functor._functor_host)
  {
  }

  void parallelFor(const Policy & policy) override final
  {
    ::Kokkos::parallel_for(policy, *this);
    ::Kokkos::fence();
  }

  KOKKOS_FUNCTION void operator()(const ThreadID tid) const
  {
    _functor_device(Operation{}, tid, _functor_device);
  }

private:
  const Object & _functor_host;
  const Object _functor_device;
};

template <typename Operation, typename Object>
class Reducer : public DispatcherBase
{
public:
  Reducer(const void * object)
    : _functor_host(*static_cast<const Object *>(object)), _functor_device(_functor_host)
  {
  }
  Reducer(const Reducer & functor)
    : value_count(functor.value_count),
      _functor_host(functor._functor_host),
      _functor_device(functor._functor_host)
  {
  }

  void parallelReduce(const Policy & policy,
                      ::Kokkos::View<Real *, ::Kokkos::HostSpace> & result) override final
  {
    value_count = result.size();

    ::Kokkos::parallel_reduce(policy, *this, result);
    ::Kokkos::fence();
  }

  using value_type = Real[];
  using size_type = ::Kokkos::View<Real *>::size_type;

  size_type value_count;

  KOKKOS_FUNCTION void operator()(const ThreadID tid, value_type result) const
  {
    _functor_device(Operation{}, tid, _functor_device, result);
  }

  KOKKOS_FUNCTION void join(value_type result, const value_type source) const
  {
    _functor_device.template join<Object>(result, source);
  }
  KOKKOS_FUNCTION void init(value_type result) const
  {
    _functor_device.template init<Object>(result);
  }

private:
  const Object & _functor_host;
  const Object _functor_device;
};

class DispatcherRegistryEntryBase
{
public:
  virtual ~DispatcherRegistryEntryBase() {}

  virtual std::unique_ptr<DispatcherBase> build(const void * object) const = 0;

  void hasUserMethod(bool flag) { _has_user_method = flag; }
  bool hasUserMethod() const { return _has_user_method; }

private:
  bool _has_user_method = false;
};

template <typename Operation, typename Object>
class DispatcherRegistryEntry : public DispatcherRegistryEntryBase
{
public:
  std::unique_ptr<DispatcherBase> build(const void * object) const override final
  {
    return std::make_unique<Dispatcher<Operation, Object>>(object);
  }
};

template <typename Operation, typename Object>
class ReducerRegistryEntry : public DispatcherRegistryEntryBase
{
public:
  std::unique_ptr<DispatcherBase> build(const void * object) const override final
  {
    return std::make_unique<Reducer<Operation, Object>>(object);
  }
};

class DispatcherRegistry
{
public:
  DispatcherRegistry() = default;

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

  DispatcherRegistry(DispatcherRegistry &&) = delete;
  DispatcherRegistry & operator=(DispatcherRegistry &&) = delete;

  template <typename Operation, typename Object>
  static void addDispatcher(const std::string & name)
  {
    auto operation = std::type_index(typeid(Operation));

    getRegistry()._dispatchers[std::make_pair(operation, name)] =
        std::make_unique<DispatcherRegistryEntry<Operation, Object>>();
  }

  template <typename Operation, typename Object>
  static void addReducer(const std::string & name)
  {
    auto operation = std::type_index(typeid(Operation));

    getRegistry()._dispatchers[std::make_pair(operation, name)] =
        std::make_unique<ReducerRegistryEntry<Operation, Object>>();
  }

  template <typename Operation>
  static void hasUserMethod(const std::string & name, const bool flag)
  {
    getDispatcher<Operation>(name)->hasUserMethod(flag);
  }

  template <typename Operation>
  static bool hasUserMethod(const std::string & name)
  {
    return getDispatcher<Operation>(name)->hasUserMethod();
  }

  template <typename Operation>
  static std::unique_ptr<DispatcherBase> build(const void * object, const std::string & name)
  {
    return getDispatcher<Operation>(name)->build(object);
  }

private:
  static DispatcherRegistry & getRegistry();

  template <typename Operation>
  static auto & getDispatcher(const std::string & name)
  {
    auto operation = std::type_index(typeid(Operation));

    auto it = getRegistry()._dispatchers.find(std::make_pair(operation, name));
    if (it == getRegistry()._dispatchers.end())
      mooseError("Kokkos functor dispatcher not registered for object type '",
                 name,
                 "'. Double check that you used Kokkos-specific registration macro.");

    return it->second;
  }

  std::map<std::pair<std::type_index, std::string>, std::unique_ptr<DispatcherRegistryEntryBase>>
      _dispatchers;
};

} // namespace Moose::Kokkos

// Kernel, NodalKernel, BC

#define callRegisterKokkosResidualObjectFunction(classname, objectname)                            \
  static char registerKokkosResidualObject##classname()                                            \
  {                                                                                                \
    using namespace Moose::Kokkos;                                                                 \
                                                                                                   \
    DispatcherRegistry::addDispatcher<classname::ResidualLoop, classname>(objectname);             \
    DispatcherRegistry::addDispatcher<classname::JacobianLoop, classname>(objectname);             \
    DispatcherRegistry::addDispatcher<classname::OffDiagJacobianLoop, classname>(objectname);      \
    DispatcherRegistry::hasUserMethod<classname::JacobianLoop>(                                    \
        objectname,                                                                                \
        &classname::computeQpJacobian<classname> != classname::defaultJacobian<classname>());      \
    DispatcherRegistry::hasUserMethod<classname::OffDiagJacobianLoop>(                             \
        objectname,                                                                                \
        &classname::computeQpOffDiagJacobian<classname> !=                                         \
            classname::defaultOffDiagJacobian<classname>());                                       \
                                                                                                   \
    return 0;                                                                                      \
  }                                                                                                \
                                                                                                   \
  static char combineNames(kokkos_dispatcher_residual_object_##classname, __COUNTER__) =           \
      registerKokkosResidualObject##classname()

#define registerKokkosResidualObject(app, classname)                                               \
  registerMooseObject(app, classname);                                                             \
  callRegisterKokkosResidualObjectFunction(classname, #classname)

#define registerKokkosResidualObjectAliased(app, classname, alias)                                 \
  registerMooseObjectAliased(app, classname, alias);                                               \
  callRegisterKokkosResidualObjectFunction(classname, alias)

// AD Kernel, NodalKernel, BC

#define callRegisterKokkosADResidualObjectFunction(classname, objectname)                          \
  static char registerKokkosADResidualObject##classname()                                          \
  {                                                                                                \
    using namespace Moose::Kokkos;                                                                 \
                                                                                                   \
    DispatcherRegistry::addDispatcher<classname::ResidualLoop, classname>(objectname);             \
                                                                                                   \
    return 0;                                                                                      \
  }                                                                                                \
                                                                                                   \
  static char combineNames(kokkos_dispatcher_ad_residual_object_##classname, __COUNTER__) =        \
      registerKokkosADResidualObject##classname()

#define registerKokkosADResidualObject(app, classname)                                             \
  registerMooseObject(app, classname);                                                             \
  callRegisterKokkosADResidualObjectFunction(classname, #classname)

#define registerKokkosADResidualObjectAliased(app, classname, alias)                               \
  registerMooseObjectAliased(app, classname, alias);                                               \
  callRegisterKokkosADResidualObjectFunction(classname, alias)

// Material

#define callRegisterKokkosMaterialFunction(classname, objectname)                                  \
  static char registerKokkosMaterial##classname()                                                  \
  {                                                                                                \
    using namespace Moose::Kokkos;                                                                 \
                                                                                                   \
    DispatcherRegistry::addDispatcher<classname::ElementInit, classname>(objectname);              \
    DispatcherRegistry::addDispatcher<classname::SideInit, classname>(objectname);                 \
    DispatcherRegistry::addDispatcher<classname::NeighborInit, classname>(objectname);             \
    DispatcherRegistry::addDispatcher<classname::ElementCompute, classname>(objectname);           \
    DispatcherRegistry::addDispatcher<classname::SideCompute, classname>(objectname);              \
    DispatcherRegistry::addDispatcher<classname::NeighborCompute, classname>(objectname);          \
    DispatcherRegistry::hasUserMethod<classname::ElementInit>(                                     \
        objectname,                                                                                \
        &classname::initQpStatefulProperties<classname> !=                                         \
            classname::defaultInitStateful<classname>());                                          \
    DispatcherRegistry::hasUserMethod<classname::SideInit>(                                        \
        objectname,                                                                                \
        &classname::initQpStatefulProperties<classname> !=                                         \
            classname::defaultInitStateful<classname>());                                          \
    DispatcherRegistry::hasUserMethod<classname::NeighborInit>(                                    \
        objectname,                                                                                \
        &classname::initQpStatefulProperties<classname> !=                                         \
            classname::defaultInitStateful<classname>());                                          \
                                                                                                   \
    return 0;                                                                                      \
  }                                                                                                \
                                                                                                   \
  static char combineNames(kokkos_dispatcher_material_##classname, __COUNTER__) =                  \
      registerKokkosMaterial##classname()

#define registerKokkosMaterial(app, classname)                                                     \
  registerMooseObject(app, classname);                                                             \
  callRegisterKokkosMaterialFunction(classname, #classname)

#define registerKokkosMaterialAliased(app, classname, alias)                                       \
  registerMooseObjectAliased(app, classname, alias);                                               \
  callRegisterKokkosMaterialFunction(classname, alias)

// AuxKernel

#define callRegisterKokkosAuxKernelFunction(classname, objectname)                                 \
  static char registerKokkosAuxKernel##classname()                                                 \
  {                                                                                                \
    using namespace Moose::Kokkos;                                                                 \
                                                                                                   \
    DispatcherRegistry::addDispatcher<classname::ElementLoop, classname>(objectname);              \
    DispatcherRegistry::addDispatcher<classname::NodeLoop, classname>(objectname);                 \
                                                                                                   \
    return 0;                                                                                      \
  }                                                                                                \
                                                                                                   \
  static char combineNames(kokkos_dispatcher_auxkernel_##classname, __COUNTER__) =                 \
      registerKokkosAuxKernel##classname()

#define registerKokkosAuxKernel(app, classname)                                                    \
  registerMooseObject(app, classname);                                                             \
  callRegisterKokkosAuxKernelFunction(classname, #classname)

#define registerKokkosAuxKernelAliased(app, classname, alias)                                      \
  registerMooseObjectAliased(app, classname, alias);                                               \
  callRegisterKokkosAuxKernelFunction(classname, alias)

// UserObject

#define callRegisterKokkosUserObjectFunction(classname, objectname)                                \
  static char registerKokkosUserObject##classname()                                                \
  {                                                                                                \
    using namespace Moose::Kokkos;                                                                 \
                                                                                                   \
    DispatcherRegistry::addDispatcher<classname::DefaultLoop, classname>(objectname);              \
    DispatcherRegistry::addReducer<classname::ReducerLoop, classname>(objectname);                 \
    DispatcherRegistry::hasUserMethod<classname::DefaultLoop>(                                     \
        objectname, &classname::execute<classname> != classname::defaultExecute<classname>());     \
    DispatcherRegistry::hasUserMethod<classname::ReducerLoop>(                                     \
        objectname, &classname::reduce<classname> != classname::defaultReduce<classname>());       \
                                                                                                   \
    return 0;                                                                                      \
  }                                                                                                \
                                                                                                   \
  static char combineNames(kokkos_dispatcher_userobject_##classname, __COUNTER__) =                \
      registerKokkosUserObject##classname()

#define registerKokkosUserObject(app, classname)                                                   \
  registerMooseObject(app, classname);                                                             \
  callRegisterKokkosUserObjectFunction(classname, #classname)

#define registerKokkosUserObjectAliased(app, classname, alias)                                     \
  registerMooseObjectAliased(app, classname, alias);                                               \
  callRegisterKokkosUserObjectFunction(classname, alias)

// User-defined parallel operation registry

#define registerKokkosAdditionalOperation(classname, operation)                                    \
  static char registerKokkos##classname##operation()                                               \
  {                                                                                                \
    using namespace Moose::Kokkos;                                                                 \
                                                                                                   \
    DispatcherRegistry::addDispatcher<classname::operation, classname>(#classname);                \
                                                                                                   \
    return 0;                                                                                      \
  }                                                                                                \
                                                                                                   \
  static char combineNames(kokkos_##classname##_##operation, __COUNTER__) =                        \
      registerKokkos##classname##operation()