https://mooseframework.inl.gov
KokkosKernel.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 "KokkosKernelBase.h"
13 
14 namespace Moose
15 {
16 namespace Kokkos
17 {
18 
39 class Kernel : public KernelBase
40 {
41 public:
43 
48 
52  virtual void computeResidual() override;
56  virtual void computeJacobian() override;
57 
62 
71  KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
72  const unsigned int /* j */,
73  const unsigned int /* qp */,
74  AssemblyDatum & /* datum */) const
75  {
76  return 0;
77  }
87  KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
88  const unsigned int /* j */,
89  const unsigned int /* jvar */,
90  const unsigned int /* qp */,
91  AssemblyDatum & /* datum */) const
92  {
93  return 0;
94  }
99  static auto defaultJacobian() { return &Kernel::computeQpJacobian; }
106 
110  template <typename Derived>
112  KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const;
113  template <typename Derived>
114  KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const;
115  template <typename Derived>
116  KOKKOS_FUNCTION void
117  operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const;
119 
125 
131  template <typename Derived>
132  KOKKOS_FUNCTION void computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const;
138  template <typename Derived>
139  KOKKOS_FUNCTION void computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const;
145  template <typename Derived>
146  KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived & kernel,
147  AssemblyDatum & datum) const;
149 
150 protected:
175 };
176 
177 template <typename Derived>
178 KOKKOS_FUNCTION void
179 Kernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const
180 {
181  auto elem = kokkosBlockElementID(tid);
182 
183  AssemblyDatum datum(elem,
185  kokkosAssembly(),
186  kokkosSystems(),
187  _kokkos_var,
188  _kokkos_var.var());
189 
190  kernel.computeResidualInternal(kernel, datum);
191 }
192 
193 template <typename Derived>
194 KOKKOS_FUNCTION void
195 Kernel::operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const
196 {
197  auto elem = kokkosBlockElementID(tid);
198 
199  AssemblyDatum datum(elem,
201  kokkosAssembly(),
202  kokkosSystems(),
203  _kokkos_var,
204  _kokkos_var.var());
205 
206  kernel.computeJacobianInternal(kernel, datum);
207 }
208 
209 template <typename Derived>
210 KOKKOS_FUNCTION void
211 Kernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const
212 {
213  auto elem = kokkosBlockElementID(_thread(tid, 1));
214 
215  auto & sys = kokkosSystem(_kokkos_var.sys());
216  auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];
217 
218  if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
219  return;
220 
221  AssemblyDatum datum(
223 
224  kernel.computeOffDiagJacobianInternal(kernel, datum);
225 }
226 
227 template <typename Derived>
228 KOKKOS_FUNCTION void
229 Kernel::computeResidualInternal(const Derived & kernel, AssemblyDatum & datum) const
230 {
232  datum,
233  [&](Real * local_re, const unsigned int ib, const unsigned int ie)
234  {
235  for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
236  {
237  datum.reinit();
238 
239  for (unsigned int i = ib; i < ie; ++i)
240  local_re[i] += datum.JxW(qp) * kernel.computeQpResidual(i, qp, datum);
241  }
242  });
243 }
244 
245 template <typename Derived>
246 KOKKOS_FUNCTION void
247 Kernel::computeJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
248 {
250  datum,
251  [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
252  {
253  for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
254  {
255  datum.reinit();
256 
257  for (unsigned int ij = ijb; ij < ije; ++ij)
258  {
259  unsigned int i = ij % datum.n_jdofs();
260  unsigned int j = ij / datum.n_jdofs();
261 
262  local_ke[ij] += datum.JxW(qp) * kernel.computeQpJacobian(i, j, qp, datum);
263  }
264  }
265  });
266 }
267 
268 template <typename Derived>
269 KOKKOS_FUNCTION void
270 Kernel::computeOffDiagJacobianInternal(const Derived & kernel, AssemblyDatum & datum) const
271 {
273  datum,
274  [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
275  {
276  for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
277  {
278  datum.reinit();
279 
280  for (unsigned int ij = ijb; ij < ije; ++ij)
281  {
282  unsigned int i = ij % datum.n_jdofs();
283  unsigned int j = ij / datum.n_jdofs();
284 
285  local_ke[ij] +=
286  datum.JxW(qp) * kernel.computeQpOffDiagJacobian(i, j, datum.jvar(), qp, datum);
287  }
288  }
289  });
290 }
291 
292 } // namespace Kokkos
293 } // namespace Moose
KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int, const unsigned int, const unsigned int, const unsigned int, AssemblyDatum &) const
Compute off-diagonal Jacobian contribution on a quadrature point.
Definition: KokkosKernel.h:87
KOKKOS_FUNCTION unsigned int sys(unsigned int comp=0) const
Get the system number of a component.
KOKKOS_FUNCTION void computeJacobianInternal(AssemblyDatum &datum, function body) const
The common loop structure template for computing elemental Jacobian.
virtual void computeResidual() override
Dispatch residual calculation.
const unsigned int invalid_uint
KOKKOS_FUNCTION const Assembly & kokkosAssembly() const
Get the const reference of the Kokkos assembly.
dof_id_type ThreadID
Definition: KokkosThread.h:18
static auto defaultJacobian()
Get the function pointer of the default computeQpJacobian()
Definition: KokkosKernel.h:99
const VariableValue _u
Current solution at quadrature points.
Definition: KokkosKernel.h:170
const InputParameters & parameters() const
Get the parameters of the object.
Definition: MooseBase.h:131
const VariableTestValue _test
Current test function.
Definition: KokkosKernel.h:154
The main MOOSE class responsible for handling user-defined parameters in almost every MOOSE system...
Thread _thread
Kokkos thread object.
KOKKOS_FUNCTION unsigned int jvar() const
Get the coupled variable number.
Definition: KokkosDatum.h:329
KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int, const unsigned int, const unsigned int, AssemblyDatum &) const
Default methods to prevent compile errors even when these methods were not defined in the derived cla...
Definition: KokkosKernel.h:71
KOKKOS_FUNCTION unsigned int n_jdofs() const
Get the number of local DOFs for the coupled variable.
Definition: KokkosDatum.h:314
The Kokkos wrapper classes for MOOSE-like shape function access.
KOKKOS_FUNCTION void computeResidualInternal(const Derived &kernel, AssemblyDatum &datum) const
The parallel computation bodies that can be customized in the derived class by defining them in the d...
Definition: KokkosKernel.h:229
static auto defaultOffDiagJacobian()
Get the function pointer of the default computeQpOffDiagJacobian()
Definition: KokkosKernel.h:104
KOKKOS_FUNCTION void computeJacobianInternal(const Derived &kernel, AssemblyDatum &datum) const
Compute diagonal Jacobian.
Definition: KokkosKernel.h:247
KOKKOS_FUNCTION ContiguousElementID kokkosBlockElementID(ThreadID tid) const
Get the contiguous element ID this Kokkos thread is operating on.
static InputParameters validParams()
KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived &kernel) const
The parallel computation entry functions called by Kokkos.
Definition: KokkosKernel.h:179
KOKKOS_FUNCTION unsigned int n_qps() const
Get the number of local quadrature points.
Definition: KokkosDatum.h:100
KOKKOS_FUNCTION const System & kokkosSystem(unsigned int sys) const
Get the const reference of a Kokkos system.
Definition: KokkosSystem.h:625
The base class for Kokkos kernels.
const VariableGradient _grad_u
Gradient of the current solution at quadrature points.
Definition: KokkosKernel.h:174
virtual void computeJacobian() override
Dispatch diagonal and off-diagonal Jacobian calculation.
KOKKOS_FUNCTION const Mesh & kokkosMesh() const
Get the const reference of the Kokkos mesh.
Definition: KokkosMesh.h:360
DIE A HORRIBLE DEATH HERE typedef LIBMESH_DEFAULT_SCALAR_TYPE Real
const VariableTestGradient _grad_test
Gradient of the current test function.
Definition: KokkosKernel.h:158
Kernel(const InputParameters &parameters)
Constructor.
KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived &kernel, AssemblyDatum &datum) const
Compute off-diagonal Jacobian.
Definition: KokkosKernel.h:270
const VariablePhiGradient _grad_phi
Gradient of the current shape function.
Definition: KokkosKernel.h:166
The Kokkos object that holds thread-private data in the parallel operations of Kokkos kernels...
Definition: KokkosDatum.h:244
KOKKOS_FUNCTION Real JxW(const unsigned int qp)
Get the transformed Jacobian weight.
Definition: KokkosDatum.h:139
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...
The Kokkos wrapper classes for MOOSE-like variable value access.
KOKKOS_FUNCTION unsigned int var(unsigned int comp=0) const
Get the variable number of a component.
const VariablePhiValue _phi
Current shape function.
Definition: KokkosKernel.h:162
The base class for a user to derive their own Kokkos kernels.
Definition: KokkosKernel.h:39
KOKKOS_FUNCTION const auto & getCoupling(unsigned int var) const
Get the list of off-diagonal coupled variable numbers of a variable.
Definition: KokkosSystem.h:147
KOKKOS_FUNCTION void computeResidualInternal(AssemblyDatum &datum, function body) const
The common loop structure template for computing elemental residual.
KOKKOS_FUNCTION void reinit()
Reset the reinit flag.
Definition: KokkosDatum.h:166
KOKKOS_FUNCTION const Array< System > & kokkosSystems() const
Get the const reference of the Kokkos systems.
Definition: KokkosSystem.h:608