https://mooseframework.inl.gov
KokkosDispatcher.h
Go to the documentation of this file.
1 //* This file is part of the MOOSE framework
2 //* https://www.mooseframework.org
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 "KokkosHeader.h"
13 #include "KokkosThread.h"
14 
15 #include <typeindex>
16 
17 namespace Moose
18 {
19 namespace Kokkos
20 {
21 
22 using Policy = ::Kokkos::RangePolicy<ExecSpace, ::Kokkos::IndexType<ThreadID>>;
23 
30 {
31 public:
32  virtual ~DispatcherBase() {}
37  virtual void parallelFor(const Policy & policy) = 0;
38 };
39 
46 template <typename Operation, typename Object>
47 class Dispatcher : public DispatcherBase
48 {
49 public:
56  Dispatcher(const void * object)
57  : _functor_host(*static_cast<const Object *>(object)), _functor_device(_functor_host)
58  {
59  }
63  Dispatcher(const Dispatcher & functor)
65  {
66  }
67 
68  void parallelFor(const Policy & policy) override final
69  {
70  ::Kokkos::parallel_for(policy, *this);
71  ::Kokkos::fence();
72  }
73 
77  KOKKOS_FUNCTION void operator()(const ThreadID tid) const
78  {
79  _functor_device(Operation{}, tid, _functor_device);
80  }
81 
82 private:
86  const Object & _functor_host;
90  const Object _functor_device;
91 };
92 
99 {
100 public:
106  virtual std::unique_ptr<DispatcherBase> build(const void * object) const = 0;
107 
112  void hasUserMethod(bool flag) { _has_user_method = flag; }
117  bool hasUserMethod() const { return _has_user_method; }
118 
119 private:
123  bool _has_user_method = false;
124 };
125 
135 template <typename Operation, typename Object>
137 {
138 public:
139  std::unique_ptr<DispatcherBase> build(const void * object) const override final
140  {
141  return std::make_unique<Dispatcher<Operation, Object>>(object);
142  }
143 };
144 
149 {
150 public:
151  DispatcherRegistry() = default;
152 
153  DispatcherRegistry(DispatcherRegistry const &) = delete;
154  DispatcherRegistry & operator=(DispatcherRegistry const &) = delete;
155 
158 
165  template <typename Operation, typename Object>
166  static void add(const std::string & name)
167  {
168  auto operation = std::type_index(typeid(Operation));
169 
170  getRegistry()._dispatchers[std::make_pair(operation, name)] =
171  std::make_unique<DispatcherRegistryEntry<Operation, Object>>();
172  }
173 
180  template <typename Operation>
181  static void hasUserMethod(const std::string & name, bool flag)
182  {
183  getDispatcher<Operation>(name)->hasUserMethod(flag);
184  }
185 
192  template <typename Operation>
193  static bool hasUserMethod(const std::string & name)
194  {
195  return getDispatcher<Operation>(name)->hasUserMethod();
196  }
197 
205  template <typename Operation>
206  static std::unique_ptr<DispatcherBase> build(const void * object, const std::string & name)
207  {
208  return getDispatcher<Operation>(name)->build(object);
209  }
210 
211 private:
216  static DispatcherRegistry & getRegistry();
217 
224  template <typename Operation>
225  static auto & getDispatcher(const std::string & name)
226  {
227  auto operation = std::type_index(typeid(Operation));
228 
229  auto it = getRegistry()._dispatchers.find(std::make_pair(operation, name));
230  if (it == getRegistry()._dispatchers.end())
231  mooseError("Kokkos functor dispatcher not registered for object type '",
232  name,
233  "'. Double check that you used Kokkos-specific registration macro.");
234 
235  return it->second;
236  }
237 
242  std::map<std::pair<std::type_index, std::string>, std::unique_ptr<DispatcherRegistryEntryBase>>
244 };
245 
246 } // namespace Kokkos
247 } // namespace Moose
248 
249 #define callRegisterKokkosResidualObjectFunction(classname, objectname) \
250  static char registerKokkosResidualObject##classname() \
251  { \
252  using namespace Moose::Kokkos; \
253  \
254  DispatcherRegistry::add<classname::ResidualLoop, classname>(objectname); \
255  DispatcherRegistry::add<classname::JacobianLoop, classname>(objectname); \
256  DispatcherRegistry::add<classname::OffDiagJacobianLoop, classname>(objectname); \
257  DispatcherRegistry::hasUserMethod<classname::JacobianLoop>( \
258  objectname, &classname::computeQpJacobian != classname::defaultJacobian()); \
259  DispatcherRegistry::hasUserMethod<classname::OffDiagJacobianLoop>( \
260  objectname, &classname::computeQpOffDiagJacobian != classname::defaultOffDiagJacobian()); \
261  \
262  return 0; \
263  } \
264  \
265  static char combineNames(kokkos_dispatcher_residual_object_##classname, __COUNTER__) = \
266  registerKokkosResidualObject##classname()
267 
268 #define registerKokkosResidualObject(app, classname) \
269  registerMooseObject(app, classname); \
270  callRegisterKokkosResidualObjectFunction(classname, #classname)
271 
272 #define registerKokkosResidualObjectAliased(app, classname, alias) \
273  registerMooseObjectAliased(app, classname, alias); \
274  callRegisterKokkosResidualObjectFunction(classname, alias)
275 
276 #define callRegisterKokkosMaterialFunction(classname, objectname) \
277  static char registerKokkosMaterial##classname() \
278  { \
279  using namespace Moose::Kokkos; \
280  \
281  DispatcherRegistry::add<classname::ElementInit, classname>(objectname); \
282  DispatcherRegistry::add<classname::SideInit, classname>(objectname); \
283  DispatcherRegistry::add<classname::NeighborInit, classname>(objectname); \
284  DispatcherRegistry::add<classname::ElementCompute, classname>(objectname); \
285  DispatcherRegistry::add<classname::SideCompute, classname>(objectname); \
286  DispatcherRegistry::add<classname::NeighborCompute, classname>(objectname); \
287  DispatcherRegistry::hasUserMethod<classname::ElementInit>( \
288  objectname, &classname::initQpStatefulProperties != classname::defaultInitStateful()); \
289  DispatcherRegistry::hasUserMethod<classname::SideInit>( \
290  objectname, &classname::initQpStatefulProperties != classname::defaultInitStateful()); \
291  DispatcherRegistry::hasUserMethod<classname::NeighborInit>( \
292  objectname, &classname::initQpStatefulProperties != classname::defaultInitStateful()); \
293  \
294  return 0; \
295  } \
296  \
297  static char combineNames(kokkos_dispatcher_material_##classname, __COUNTER__) = \
298  registerKokkosMaterial##classname()
299 
300 #define registerKokkosMaterial(app, classname) \
301  registerMooseObject(app, classname); \
302  callRegisterKokkosMaterialFunction(classname, #classname)
303 
304 #define registerKokkosMaterialAliased(app, classname, alias) \
305  registerMooseObjectAliased(app, classname, alias); \
306  callRegisterKokkosMaterialFunction(classname, alias)
std::string name(const ElemQuality q)
std::unique_ptr< DispatcherBase > build(const void *object) const override final
Build a dispatcher for this operation and functor.
static void add(const std::string &name)
Register a dispatcher of an operation of a functor.
Class that dispatches an operation of a Kokkos functor.
static bool hasUserMethod(const std::string &name)
Get whether the user has overriden the hook method associated with an operation of a functor...
dof_id_type ThreadID
Definition: KokkosThread.h:18
DispatcherRegistry & operator=(DispatcherRegistry const &)=delete
void mooseError(Args &&... args)
Emit an error message with the given stringified, concatenated args and terminate the application...
Definition: MooseError.h:323
KOKKOS_FUNCTION void operator()(const ThreadID tid) const
The parallel computation entry function called by Kokkos.
const Object _functor_device
Copy of the functor on device.
static auto & getDispatcher(const std::string &name)
Get the dispatcher shell of an operation of a functor.
Class that stores the information of a dispatcher and builds it.
::Kokkos::RangePolicy< ExecSpace, ::Kokkos::IndexType< ThreadID > > Policy
Base class for dispatcher registry entry.
static void hasUserMethod(const std::string &name, bool flag)
Set whether the user has overriden the hook method associated with an operation of a functor...
virtual void parallelFor(const Policy &policy)=0
Dispatch this functor with Kokkos parallel_for() given a Kokkos execution policy. ...
static std::unique_ptr< DispatcherBase > build(const void *object, const std::string &name)
Build and get a dispatcher of an operation of a functor.
bool _has_user_method
Flag whether the user has overriden the hook method associated with this operation.
bool hasUserMethod() const
Get whether the user has overriden the hook method associated with this operation.
void hasUserMethod(bool flag)
Set whether the user has overriden the hook method associated with this operation.
Base class for Kokkos functor dispatcher.
const Object & _functor_host
Reference of the functor on host.
Dispatcher(const void *object)
Constructor.
void parallelFor(const Policy &policy) override final
Dispatch this functor with Kokkos parallel_for() given a Kokkos execution policy. ...
std::map< std::pair< std::type_index, std::string >, std::unique_ptr< DispatcherRegistryEntryBase > > _dispatchers
Map containing the dispatcher shells with the key being the pair of function tag type index and regis...
Class that registers dispatchers of all Kokkos functors.
MOOSE now contains C++17 code, so give a reasonable error message stating what the user can do to add...
static DispatcherRegistry & getRegistry()
Get the registry singleton.
virtual std::unique_ptr< DispatcherBase > build(const void *object) const =0
Build a dispatcher for this operation and functor.
Dispatcher(const Dispatcher &functor)
Copy constructor for parallel dispatch.