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>
#include <type_traits>

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;
};

template <typename T, typename = void>
struct has_split_linear_fv_flux_dispatchers : std::false_type
{
};

template <typename T>
struct has_split_linear_fv_flux_dispatchers<T,
                                            std::void_t<typename T::InternalRightHandSideLoop,
                                                        typename T::BoundaryRightHandSideLoop,
                                                        typename T::InternalMatrixLoop,
                                                        typename T::BoundaryMatrixLoop>>
  : std::true_type
{
};

template <typename Object>
bool
hasLinearFVMatrixContribution()
{
  return &Object::template computeMatrixContribution<Object> !=
             Object::template defaultMatrixContribution<Object>() ||
         &Object::template computeNeighborMatrixContribution<Object> !=
             Object::template defaultNeighborMatrixContribution<Object>();
}

template <typename Object>
bool
hasInternalLinearFVFluxMatrixContribution()
{
  return &Object::template computeInternalMatrixContribution<Object> !=
             Object::template defaultInternalMatrixContribution<Object>() ||
         &Object::template computeInternalNeighborMatrixContribution<Object> !=
             Object::template defaultInternalNeighborMatrixContribution<Object>();
}

template <typename Object>
bool
hasBoundaryLinearFVFluxMatrixContribution()
{
  return &Object::template computeBoundaryMatrixContribution<Object> !=
         Object::template defaultBoundaryMatrixContribution<Object>();
}

template <typename Object>
void
registerLinearFVKernelDispatchers(const std::string & objectname)
{
  if constexpr (has_split_linear_fv_flux_dispatchers<Object>::value)
  {
    DispatcherRegistry::addDispatcher<typename Object::InternalRightHandSideLoop, Object>(
        objectname);
    DispatcherRegistry::addDispatcher<typename Object::BoundaryRightHandSideLoop, Object>(
        objectname);
    DispatcherRegistry::addDispatcher<typename Object::InternalMatrixLoop, Object>(objectname);
    DispatcherRegistry::addDispatcher<typename Object::BoundaryMatrixLoop, Object>(objectname);

    DispatcherRegistry::hasUserMethod<typename Object::InternalMatrixLoop>(
        objectname, hasInternalLinearFVFluxMatrixContribution<Object>());
    DispatcherRegistry::hasUserMethod<typename Object::BoundaryMatrixLoop>(
        objectname, hasBoundaryLinearFVFluxMatrixContribution<Object>());
  }
  else
  {
    DispatcherRegistry::addDispatcher<typename Object::RightHandSideLoop, Object>(objectname);
    DispatcherRegistry::addDispatcher<typename Object::MatrixLoop, Object>(objectname);
    DispatcherRegistry::hasUserMethod<typename Object::MatrixLoop>(
        objectname, hasLinearFVMatrixContribution<Object>());
  }
}

} // namespace Moose::Kokkos

// Kernel, NodalKernel, BC

namespace Moose::Kokkos
{

template <typename Object>
bool
hasUserJacobianHook()
{
  if constexpr (Object::use_precompute_hooks)
    return &Object::template precomputeQpJacobian<Object> !=
           Object::template defaultJacobian<Object>();
  else
    return &Object::template computeQpJacobian<Object> !=
           Object::template defaultJacobian<Object>();
}

template <typename Object>
bool
hasUserOffDiagJacobianHook()
{
  if constexpr (Object::use_precompute_hooks)
    return &Object::template precomputeQpOffDiagJacobian<Object> !=
           Object::template defaultOffDiagJacobian<Object>();
  else
    return &Object::template computeQpOffDiagJacobian<Object> !=
           Object::template defaultOffDiagJacobian<Object>();
}

} // namespace Moose::Kokkos

#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,                         \
                                                               hasUserJacobianHook<classname>());  \
    DispatcherRegistry::hasUserMethod<classname::OffDiagJacobianLoop>(                             \
        objectname, hasUserOffDiagJacobianHook<classname>());                                      \
                                                                                                   \
    return 0;                                                                                      \
  }                                                                                                \
                                                                                                   \
  [[maybe_unused]] 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;                                                                                      \
  }                                                                                                \
                                                                                                   \
  [[maybe_unused]] 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)

#define callRegisterKokkosLinearFVKernelFunction(classname, objectname)                            \
  static char registerKokkosLinearFVKernel##classname()                                            \
  {                                                                                                \
    using namespace Moose::Kokkos;                                                                 \
                                                                                                   \
    registerLinearFVKernelDispatchers<classname>(objectname);                                      \
                                                                                                   \
    return 0;                                                                                      \
  }                                                                                                \
                                                                                                   \
  static char combineNames(kokkos_dispatcher_linear_fv_kernel_##classname, __COUNTER__) =          \
      registerKokkosLinearFVKernel##classname()

#define registerKokkosLinearFVKernel(app, classname)                                               \
  registerMooseObject(app, classname);                                                             \
  callRegisterKokkosLinearFVKernelFunction(classname, #classname)

#define registerKokkosLinearFVKernelAliased(app, classname, alias)                                 \
  registerMooseObjectAliased(app, classname, alias);                                               \
  callRegisterKokkosLinearFVKernelFunction(classname, alias)

#define callRegisterKokkosLinearFVBoundaryConditionFunction(classname, objectname)                 \
  static char registerKokkosLinearFVBoundaryCondition##classname()                                 \
  {                                                                                                \
    using namespace Moose::Kokkos;                                                                 \
                                                                                                   \
    DispatcherRegistry::addDispatcher<classname::BoundaryValueLoop, classname>(objectname);        \
    DispatcherRegistry::addDispatcher<classname::BoundaryNormalGradientLoop, classname>(           \
        objectname);                                                                               \
    DispatcherRegistry::hasUserMethod<classname::BoundaryValueLoop>(                               \
        objectname,                                                                                \
        &classname::computeBoundaryValue<classname> !=                                             \
            classname::defaultBoundaryValue<classname>());                                         \
    DispatcherRegistry::hasUserMethod<classname::BoundaryNormalGradientLoop>(                      \
        objectname,                                                                                \
        &classname::computeBoundaryNormalGradient<classname> !=                                    \
            classname::defaultBoundaryNormalGradient<classname>());                                \
                                                                                                   \
    return 0;                                                                                      \
  }                                                                                                \
                                                                                                   \
  static char combineNames(kokkos_dispatcher_linear_fv_boundary_condition_##classname,             \
                           __COUNTER__) = registerKokkosLinearFVBoundaryCondition##classname()

#define registerKokkosLinearFVBoundaryCondition(app, classname)                                    \
  registerMooseObject(app, classname);                                                             \
  callRegisterKokkosLinearFVBoundaryConditionFunction(classname, #classname)

#define registerKokkosLinearFVBoundaryConditionAliased(app, classname, alias)                      \
  registerMooseObjectAliased(app, classname, alias);                                               \
  callRegisterKokkosLinearFVBoundaryConditionFunction(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;                                                                                      \
  }                                                                                                \
                                                                                                   \
  [[maybe_unused]] 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;                                                                                      \
  }                                                                                                \
                                                                                                   \
  [[maybe_unused]] 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;                                                                                      \
  }                                                                                                \
                                                                                                   \
  [[maybe_unused]] 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;                                                                                      \
  }                                                                                                \
                                                                                                   \
  [[maybe_unused]] static char combineNames(kokkos_##classname##_##operation, __COUNTER__) =       \
      registerKokkos##classname##operation()