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 // Kernel, NodalKernel, BC
250 
251 #define callRegisterKokkosResidualObjectFunction(classname, objectname) \
252  static char registerKokkosResidualObject##classname() \
253  { \
254  using namespace Moose::Kokkos; \
255  \
256  DispatcherRegistry::add<classname::ResidualLoop, classname>(objectname); \
257  DispatcherRegistry::add<classname::JacobianLoop, classname>(objectname); \
258  DispatcherRegistry::add<classname::OffDiagJacobianLoop, classname>(objectname); \
259  DispatcherRegistry::hasUserMethod<classname::JacobianLoop>( \
260  objectname, &classname::computeQpJacobian != classname::defaultJacobian()); \
261  DispatcherRegistry::hasUserMethod<classname::OffDiagJacobianLoop>( \
262  objectname, &classname::computeQpOffDiagJacobian != classname::defaultOffDiagJacobian()); \
263  \
264  return 0; \
265  } \
266  \
267  static char combineNames(kokkos_dispatcher_residual_object_##classname, __COUNTER__) = \
268  registerKokkosResidualObject##classname()
269 
270 #define registerKokkosResidualObject(app, classname) \
271  registerMooseObject(app, classname); \
272  callRegisterKokkosResidualObjectFunction(classname, #classname)
273 
274 #define registerKokkosResidualObjectAliased(app, classname, alias) \
275  registerMooseObjectAliased(app, classname, alias); \
276  callRegisterKokkosResidualObjectFunction(classname, alias)
277 
278 // Material
279 
280 #define callRegisterKokkosMaterialFunction(classname, objectname) \
281  static char registerKokkosMaterial##classname() \
282  { \
283  using namespace Moose::Kokkos; \
284  \
285  DispatcherRegistry::add<classname::ElementInit, classname>(objectname); \
286  DispatcherRegistry::add<classname::SideInit, classname>(objectname); \
287  DispatcherRegistry::add<classname::NeighborInit, classname>(objectname); \
288  DispatcherRegistry::add<classname::ElementCompute, classname>(objectname); \
289  DispatcherRegistry::add<classname::SideCompute, classname>(objectname); \
290  DispatcherRegistry::add<classname::NeighborCompute, classname>(objectname); \
291  DispatcherRegistry::hasUserMethod<classname::ElementInit>( \
292  objectname, &classname::initQpStatefulProperties != classname::defaultInitStateful()); \
293  DispatcherRegistry::hasUserMethod<classname::SideInit>( \
294  objectname, &classname::initQpStatefulProperties != classname::defaultInitStateful()); \
295  DispatcherRegistry::hasUserMethod<classname::NeighborInit>( \
296  objectname, &classname::initQpStatefulProperties != classname::defaultInitStateful()); \
297  \
298  return 0; \
299  } \
300  \
301  static char combineNames(kokkos_dispatcher_material_##classname, __COUNTER__) = \
302  registerKokkosMaterial##classname()
303 
304 #define registerKokkosMaterial(app, classname) \
305  registerMooseObject(app, classname); \
306  callRegisterKokkosMaterialFunction(classname, #classname)
307 
308 #define registerKokkosMaterialAliased(app, classname, alias) \
309  registerMooseObjectAliased(app, classname, alias); \
310  callRegisterKokkosMaterialFunction(classname, alias)
311 
312 // AuxKernel
313 
314 #define callRegisterKokkosAuxKernelFunction(classname, objectname) \
315  static char registerKokkosAuxKernel##classname() \
316  { \
317  using namespace Moose::Kokkos; \
318  \
319  DispatcherRegistry::add<classname::ElementLoop, classname>(objectname); \
320  DispatcherRegistry::add<classname::NodeLoop, classname>(objectname); \
321  \
322  return 0; \
323  } \
324  \
325  static char combineNames(kokkos_dispatcher_auxkernel_##classname, __COUNTER__) = \
326  registerKokkosAuxKernel##classname()
327 
328 #define registerKokkosAuxKernel(app, classname) \
329  registerMooseObject(app, classname); \
330  callRegisterKokkosAuxKernelFunction(classname, #classname)
331 
332 #define registerKokkosAuxKernelAliased(app, classname, alias) \
333  registerMooseObjectAliased(app, classname, alias); \
334  callRegisterKokkosAuxKernelFunction(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.