https://mooseframework.inl.gov
KokkosResidualObject.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 "KokkosDispatcher.h"
13 #include "KokkosVariableValue.h"
15 
16 #include "MooseVariableBase.h"
17 #include "ResidualObject.h"
18 
19 namespace Moose
20 {
21 namespace Kokkos
22 {
23 
28  public MeshHolder,
29  public AssemblyHolder,
30  public SystemHolder
31 {
32 public:
34 
41  Moose::VarFieldType field_type,
42  bool nodal = false);
46  ResidualObject(const ResidualObject & object);
47 
51  struct ResidualLoop
53  {
54  };
55  struct JacobianLoop
56  {
57  };
59  {
60  };
62 
63  virtual const MooseVariableBase & variable() const override { return _var; }
64 
65  virtual void computeOffDiagJacobian(unsigned int) override final
66  {
67  mooseError("computeOffDiagJacobian() is not used for Kokkos residual objects.");
68  }
69  virtual void computeResidualAndJacobian() override final
70  {
73  }
74 
75 protected:
91  std::unique_ptr<DispatcherBase> _residual_dispatcher;
93  std::unique_ptr<DispatcherBase> _jacobian_dispatcher;
94  std::unique_ptr<DispatcherBase> _offdiag_jacobian_dispatcher;
96 
100 
122 
130  KOKKOS_FUNCTION void accumulateTaggedElementalResidual(const Real local_re,
131  const ContiguousElementID elem,
132  const unsigned int i,
133  const unsigned int comp = 0) const;
141  KOKKOS_FUNCTION void accumulateTaggedNodalResidual(const bool add,
142  const Real local_re,
143  const ContiguousNodeID node,
144  const unsigned int comp = 0) const;
154  KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const Real local_ke,
155  const ContiguousElementID elem,
156  const unsigned int i,
157  const unsigned int j,
158  const unsigned int jvar,
159  const unsigned int comp = 0) const;
168  KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add,
169  const Real local_ke,
170  const ContiguousNodeID node,
171  const unsigned int jvar,
172  const unsigned int comp = 0) const;
173 
179  template <typename function>
180  KOKKOS_FUNCTION void computeResidualInternal(ResidualDatum & datum, function body) const;
186  template <typename function>
187  KOKKOS_FUNCTION void computeJacobianInternal(ResidualDatum & datum, function body) const;
188 
189 private:
197 
204 };
205 
206 KOKKOS_FUNCTION inline void
208  const ContiguousElementID elem,
209  const unsigned int i,
210  const unsigned int comp) const
211 {
212  if (!local_re)
213  return;
214 
215  auto & sys = kokkosSystem(_kokkos_var.sys());
216  auto dof = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
217 
218  for (dof_id_type t = 0; t < _vector_tags.size(); ++t)
219  {
220  auto tag = _vector_tags[t];
221 
222  if (sys.isResidualTagActive(tag))
223  {
224  bool has_nodal_bc =
225  _is_extra_vector_tag[t] ? sys.hasNodalBCResidualTag(dof, tag) : sys.hasNodalBC(dof);
226 
227  if (!has_nodal_bc)
228  ::Kokkos::atomic_add(&sys.getVectorDofValue(dof, tag), local_re);
229  }
230  }
231 }
232 
233 KOKKOS_FUNCTION inline void
235  const Real local_re,
236  const ContiguousNodeID node,
237  const unsigned int comp) const
238 {
239  if (!local_re)
240  return;
241 
242  auto & sys = kokkosSystem(_kokkos_var.sys());
243  auto dof = sys.getNodeLocalDofIndex(node, _kokkos_var.var(comp));
244 
245  for (dof_id_type t = 0; t < _vector_tags.size(); ++t)
246  {
247  auto tag = _vector_tags[t];
248 
249  if (sys.isResidualTagActive(tag))
250  {
251  if (add)
252  sys.getVectorDofValue(dof, tag) += local_re;
253  else
254  sys.getVectorDofValue(dof, tag) = local_re;
255  }
256  }
257 }
258 
259 KOKKOS_FUNCTION inline void
261  const ContiguousElementID elem,
262  const unsigned int i,
263  const unsigned int j,
264  const unsigned int jvar,
265  const unsigned int comp) const
266 {
267  if (!local_ke)
268  return;
269 
270  auto & sys = kokkosSystem(_kokkos_var.sys());
271  auto row = sys.getElemLocalDofIndex(elem, i, _kokkos_var.var(comp));
272  auto col = sys.getElemGlobalDofIndex(elem, j, jvar);
273 
274  for (dof_id_type t = 0; t < _matrix_tags.size(); ++t)
275  {
276  auto tag = _matrix_tags[t];
277 
278  if (sys.isMatrixTagActive(tag))
279  {
280  bool has_nodal_bc =
281  _is_extra_matrix_tag[t] ? sys.hasNodalBCMatrixTag(row, tag) : sys.hasNodalBC(row);
282 
283  if (!has_nodal_bc)
284  ::Kokkos::atomic_add(&sys.getMatrixValue(row, col, tag), local_ke);
285  }
286  }
287 }
288 
289 KOKKOS_FUNCTION inline void
291  const Real local_ke,
292  const ContiguousNodeID node,
293  const unsigned int jvar,
294  const unsigned int comp) const
295 {
296  if (!local_ke)
297  return;
298 
299  auto & sys = kokkosSystem(_kokkos_var.sys());
300  auto row = sys.getNodeLocalDofIndex(node, _kokkos_var.var(comp));
301  auto col = sys.getNodeGlobalDofIndex(node, jvar);
302 
303  for (dof_id_type t = 0; t < _matrix_tags.size(); ++t)
304  {
305  auto tag = _matrix_tags[t];
306 
307  if (sys.isMatrixTagActive(tag))
308  {
309  auto & matrix = sys.getMatrix(tag);
310 
311  if (add)
312  matrix(row, col) += local_ke;
313  else
314  {
315  matrix.zero(row);
316  matrix(row, col) = local_ke;
317  }
318  }
319  }
320 }
321 
322 template <typename function>
323 KOKKOS_FUNCTION void
325 {
326  Real local_re[MAX_CACHED_DOF];
327 
328  unsigned int num_batches = datum.n_dofs() / MAX_CACHED_DOF;
329 
330  if (datum.n_dofs() % MAX_CACHED_DOF)
331  ++num_batches;
332 
333  for (unsigned int batch = 0; batch < num_batches; ++batch)
334  {
335  unsigned int ib = batch * MAX_CACHED_DOF;
336  unsigned int ie = ::Kokkos::min(ib + MAX_CACHED_DOF, datum.n_dofs());
337 
338  for (unsigned int i = ib; i < ie; ++i)
339  local_re[i - ib] = 0;
340 
341  body(local_re - ib, ib, ie);
342 
343  for (unsigned int i = ib; i < ie; ++i)
344  accumulateTaggedElementalResidual(local_re[i - ib], datum.elem().id, i);
345  }
346 }
347 
348 template <typename function>
349 KOKKOS_FUNCTION void
351 {
352  Real local_ke[MAX_CACHED_DOF];
353 
354  unsigned int num_batches = datum.n_idofs() * datum.n_jdofs() / MAX_CACHED_DOF;
355 
356  if ((datum.n_idofs() * datum.n_jdofs()) % MAX_CACHED_DOF)
357  ++num_batches;
358 
359  for (unsigned int batch = 0; batch < num_batches; ++batch)
360  {
361  unsigned int ijb = batch * MAX_CACHED_DOF;
362  unsigned int ije = ::Kokkos::min(ijb + MAX_CACHED_DOF, datum.n_idofs() * datum.n_jdofs());
363 
364  for (unsigned int ij = ijb; ij < ije; ++ij)
365  local_ke[ij - ijb] = 0;
366 
367  body(local_ke - ijb, ijb, ije);
368 
369  for (unsigned int ij = ijb; ij < ije; ++ij)
370  {
371  unsigned int i = ij % datum.n_jdofs();
372  unsigned int j = ij / datum.n_jdofs();
373 
374  accumulateTaggedElementalMatrix(local_ke[ij - ijb], datum.elem().id, i, j, datum.jvar());
375  }
376  }
377 }
378 
379 } // namespace Kokkos
380 } // namespace Moose
VarFieldType
Definition: MooseTypes.h:722
KOKKOS_FUNCTION unsigned int n_jdofs() const
Get the number of local DOFs for the coupled variable.
Definition: KokkosDatum.h:315
KOKKOS_FUNCTION dof_id_type getElemLocalDofIndex(ContiguousElementID elem, unsigned int i, unsigned int var) const
Get the local DOF index of a variable for an element.
Definition: KokkosSystem.h:225
Scalar< Real > _dt_old
Size of the old time step.
std::unique_ptr< DispatcherBase > _jacobian_dispatcher
KOKKOS_FUNCTION unsigned int sys(unsigned int comp=0) const
Get the system number of a component.
Scalar< Real > _t
TODO: Move to TransientInterface.
dof_id_type ContiguousElementID
Definition: KokkosMesh.h:20
KOKKOS_FUNCTION void computeJacobianInternal(ResidualDatum &datum, function body) const
The common loop structure template for computing elemental Jacobian.
Array< bool > _is_extra_vector_tag
Flag whether each tag.
const InputParameters & parameters() const
Get the parameters of the object.
Definition: MooseBase.h:131
ResidualObject(const InputParameters &parameters, Moose::VarFieldType field_type, bool nodal=false)
Constructor.
KOKKOS_FUNCTION void accumulateTaggedNodalResidual(const bool add, const Real local_re, const ContiguousNodeID node, const unsigned int comp=0) const
Accumulate or set local nodal residual contribution to tagged vectors.
KOKKOS_FUNCTION dof_id_type getNodeLocalDofIndex(ContiguousNodeID node, unsigned int var) const
Get the local DOF index of a variable for a node.
Definition: KokkosSystem.h:237
virtual void computeOffDiagJacobian(unsigned int) override final
Computes this object&#39;s contribution to off-diagonal blocks of the system Jacobian matrix...
The main MOOSE class responsible for handling user-defined parameters in almost every MOOSE system...
The base class for Kokkos residual objects.
Thread _thread
Kokkos thread object.
MooseVariableFieldBase & _var
Reference of the MOOSE variable.
This class provides an interface for common operations on field variables of both FE and FV types wit...
virtual void computeResidualAndJacobian() override final
Compute this object&#39;s contribution to the residual and Jacobian simultaneously.
The Kokkos interface that holds the host reference of the Kokkos systems and copies it to device duri...
Definition: KokkosSystem.h:576
The Kokkos interface that holds the host reference of the Kokkos mesh and copies it to device during ...
Definition: KokkosMesh.h:338
KOKKOS_FUNCTION unsigned int n_dofs() const
Get the number of local DOFs.
Definition: KokkosDatum.h:305
KOKKOS_FUNCTION void accumulateTaggedElementalMatrix(const Real local_ke, const ContiguousElementID elem, const unsigned int i, const unsigned int j, const unsigned int jvar, const unsigned int comp=0) const
Accumulate local elemental Jacobian contribution to tagged matrices.
dof_id_type ContiguousNodeID
Definition: KokkosMesh.h:21
virtual void computeResidual()=0
Compute this object&#39;s contribution to the residual.
KOKKOS_FUNCTION const System & kokkosSystem(unsigned int sys) const
Get the const reference of a Kokkos system.
Definition: KokkosSystem.h:615
The Kokkos interface that holds the host reference of the Kokkos assembly and copies it to device dur...
virtual void computeJacobian()=0
Compute this object&#39;s contribution to the diagonal Jacobian entries.
Scalar< const Real > _t_old
Old time.
Scalar< int > _t_step
The number of the time step.
std::unique_ptr< DispatcherBase > _offdiag_jacobian_dispatcher
constexpr unsigned int MAX_CACHED_DOF
Maximum number of DOFs to cache during residual and Jacobian computation.
Definition: KokkosHeader.h:79
DIE A HORRIBLE DEATH HERE typedef LIBMESH_DEFAULT_SCALAR_TYPE Real
The Kokkos variable object that carries the coupled variable and tag information. ...
KOKKOS_FUNCTION void accumulateTaggedElementalResidual(const Real local_re, const ContiguousElementID elem, const unsigned int i, const unsigned int comp=0) const
Accumulate local elemental residual contribution to tagged vectors.
KOKKOS_FUNCTION const ElementInfo & elem() const
Get the element information object.
Definition: KokkosDatum.h:80
void mooseError(Args &&... args) const
Emits an error prefixed with object name and type and optionally a file path to the top-level block p...
Definition: MooseBase.h:271
KOKKOS_FUNCTION void computeResidualInternal(ResidualDatum &datum, function body) const
The common loop structure template for computing elemental residual.
Scalar< Real > _dt
Time step size.
Variable _kokkos_var
Kokkos variable.
MOOSE now contains C++17 code, so give a reasonable error message stating what the user can do to add...
KOKKOS_FUNCTION unsigned int var(unsigned int comp=0) const
Get the variable number of a component.
std::unique_ptr< DispatcherBase > _residual_dispatcher
Kokkos functor dispatchers.
auto min(const L &left, const R &right)
static InputParameters validParams()
KOKKOS_FUNCTION unsigned int jvar() const
Get the coupled variable number.
Definition: KokkosDatum.h:330
Array< TagID > _vector_tags
Tags this object operates on.
The Kokkos object that holds thread-private data in the parallel operations of Kokkos residual object...
Definition: KokkosDatum.h:245
KOKKOS_FUNCTION void accumulateTaggedNodalMatrix(const bool add, const Real local_ke, const ContiguousNodeID node, const unsigned int jvar, const unsigned int comp=0) const
Accumulate or set local nodal Jacobian contribution to tagged matrices.
The Kokkos thread object that aids in converting the one-dimensional thread index into multi-dimensio...
Definition: KokkosThread.h:29
KOKKOS_FUNCTION unsigned int n_idofs() const
Get the number of local DOFs.
Definition: KokkosDatum.h:310
virtual const MooseVariableBase & variable() const override
Returns the variable that this object operates on.
ContiguousElementID id
Contiguous element ID.
Definition: KokkosMesh.h:44
Base variable class.
uint8_t dof_id_type