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 
40 class Kernel : public KernelBase
41 {
42 public:
44 
49 
53  virtual void computeResidual() override;
57  virtual void computeJacobian() override;
58 
63 
72  KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int /* i */,
73  const unsigned int /* j */,
74  const unsigned int /* qp */,
75  ResidualDatum & /* datum */) const
76  {
77  return 0;
78  }
88  KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int /* i */,
89  const unsigned int /* j */,
90  const unsigned int /* jvar */,
91  const unsigned int /* qp */,
92  ResidualDatum & /* datum */) const
93  {
94  return 0;
95  }
100  static auto defaultJacobian() { return &Kernel::computeQpJacobian; }
107 
111  template <typename Derived>
113  KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const;
114  template <typename Derived>
115  KOKKOS_FUNCTION void operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const;
116  template <typename Derived>
117  KOKKOS_FUNCTION void
118  operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const;
120 
126 
132  template <typename Derived>
133  KOKKOS_FUNCTION void computeResidualInternal(const Derived & kernel, ResidualDatum & datum) const;
139  template <typename Derived>
140  KOKKOS_FUNCTION void computeJacobianInternal(const Derived & kernel, ResidualDatum & datum) const;
146  template <typename Derived>
147  KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived & kernel,
148  ResidualDatum & datum) const;
150 
151 protected:
176 };
177 
178 template <typename Derived>
179 KOKKOS_FUNCTION void
180 Kernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const
181 {
182  auto elem = kokkosBlockElementID(tid);
183 
184  ResidualDatum datum(elem,
186  kokkosAssembly(),
187  kokkosSystems(),
188  _kokkos_var,
189  _kokkos_var.var());
190 
191  kernel.computeResidualInternal(kernel, datum);
192 }
193 
194 template <typename Derived>
195 KOKKOS_FUNCTION void
196 Kernel::operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const
197 {
198  auto elem = kokkosBlockElementID(tid);
199 
200  ResidualDatum datum(elem,
202  kokkosAssembly(),
203  kokkosSystems(),
204  _kokkos_var,
205  _kokkos_var.var());
206 
207  kernel.computeJacobianInternal(kernel, datum);
208 }
209 
210 template <typename Derived>
211 KOKKOS_FUNCTION void
212 Kernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const
213 {
214  auto elem = kokkosBlockElementID(_thread(tid, 1));
215 
216  auto & sys = kokkosSystem(_kokkos_var.sys());
217  auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)];
218 
219  if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain))
220  return;
221 
222  ResidualDatum datum(
224 
225  kernel.computeOffDiagJacobianInternal(kernel, datum);
226 }
227 
228 template <typename Derived>
229 KOKKOS_FUNCTION void
230 Kernel::computeResidualInternal(const Derived & kernel, ResidualDatum & datum) const
231 {
233  datum,
234  [&](Real * local_re, const unsigned int ib, const unsigned int ie)
235  {
236  for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
237  {
238  datum.reinit();
239 
240  for (unsigned int i = ib; i < ie; ++i)
241  local_re[i] += datum.JxW(qp) * kernel.computeQpResidual(i, qp, datum);
242  }
243  });
244 }
245 
246 template <typename Derived>
247 KOKKOS_FUNCTION void
248 Kernel::computeJacobianInternal(const Derived & kernel, ResidualDatum & datum) const
249 {
251  datum,
252  [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
253  {
254  for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
255  {
256  datum.reinit();
257 
258  for (unsigned int ij = ijb; ij < ije; ++ij)
259  {
260  unsigned int i = ij % datum.n_jdofs();
261  unsigned int j = ij / datum.n_jdofs();
262 
263  local_ke[ij] += datum.JxW(qp) * kernel.computeQpJacobian(i, j, qp, datum);
264  }
265  }
266  });
267 }
268 
269 template <typename Derived>
270 KOKKOS_FUNCTION void
271 Kernel::computeOffDiagJacobianInternal(const Derived & kernel, ResidualDatum & datum) const
272 {
274  datum,
275  [&](Real * local_ke, const unsigned int ijb, const unsigned int ije)
276  {
277  for (unsigned int qp = 0; qp < datum.n_qps(); ++qp)
278  {
279  datum.reinit();
280 
281  for (unsigned int ij = ijb; ij < ije; ++ij)
282  {
283  unsigned int i = ij % datum.n_jdofs();
284  unsigned int j = ij / datum.n_jdofs();
285 
286  local_ke[ij] +=
287  datum.JxW(qp) * kernel.computeQpOffDiagJacobian(i, j, datum.jvar(), qp, datum);
288  }
289  }
290  });
291 }
292 
293 } // namespace Kokkos
294 } // namespace Moose
KOKKOS_FUNCTION unsigned int n_jdofs() const
Get the number of local DOFs for the coupled variable.
Definition: KokkosDatum.h:315
KOKKOS_FUNCTION unsigned int sys(unsigned int comp=0) const
Get the system number of a component.
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
KOKKOS_FUNCTION void computeJacobianInternal(ResidualDatum &datum, function body) const
The common loop structure template for computing elemental Jacobian.
static auto defaultJacobian()
Get the function pointer of the default computeQpJacobian()
Definition: KokkosKernel.h:100
const VariableValue _u
Current solution at quadrature points.
Definition: KokkosKernel.h:171
const InputParameters & parameters() const
Get the parameters of the object.
Definition: MooseBase.h:131
const VariableTestValue _test
Current test function.
Definition: KokkosKernel.h:155
The main MOOSE class responsible for handling user-defined parameters in almost every MOOSE system...
Thread _thread
Kokkos thread object.
KOKKOS_FUNCTION Real computeQpJacobian(const unsigned int, const unsigned int, const unsigned int, ResidualDatum &) const
Default methods to prevent compile errors even when these methods were not defined in the derived cla...
Definition: KokkosKernel.h:72
The Kokkos wrapper classes for MOOSE-like shape function access.
static auto defaultOffDiagJacobian()
Get the function pointer of the default computeQpOffDiagJacobian()
Definition: KokkosKernel.h:105
KOKKOS_FUNCTION void computeJacobianInternal(const Derived &kernel, ResidualDatum &datum) const
Compute diagonal Jacobian.
Definition: KokkosKernel.h:248
KOKKOS_FUNCTION Real computeQpOffDiagJacobian(const unsigned int, const unsigned int, const unsigned int, const unsigned int, ResidualDatum &) const
Compute off-diagonal Jacobian contribution on a quadrature point.
Definition: KokkosKernel.h:88
KOKKOS_FUNCTION ContiguousElementID kokkosBlockElementID(ThreadID tid) const
Get the contiguous element ID this Kokkos thread is operating on.
static InputParameters validParams()
KOKKOS_FUNCTION void computeResidualInternal(const Derived &kernel, ResidualDatum &datum) const
The parallel computation bodies that can be customized in the derived class by defining them in the d...
Definition: KokkosKernel.h:230
KOKKOS_FUNCTION void operator()(ResidualLoop, const ThreadID tid, const Derived &kernel) const
The parallel computation entry functions called by Kokkos.
Definition: KokkosKernel.h:180
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:615
The base class for Kokkos kernels.
const VariableGradient _grad_u
Gradient of the current solution at quadrature points.
Definition: KokkosKernel.h:175
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:159
Kernel(const InputParameters &parameters)
Constructor.
const VariablePhiGradient _grad_phi
Gradient of the current shape function.
Definition: KokkosKernel.h:167
KOKKOS_FUNCTION void computeResidualInternal(ResidualDatum &datum, function body) const
The common loop structure template for computing elemental residual.
KOKKOS_FUNCTION void computeOffDiagJacobianInternal(const Derived &kernel, ResidualDatum &datum) const
Compute off-diagonal Jacobian.
Definition: KokkosKernel.h:271
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:163
The base class for a user to derive their own Kokkos kernels.
Definition: KokkosKernel.h:40
KOKKOS_FUNCTION const auto & getCoupling(unsigned int var) const
Get the list of off-diagonal coupled variable numbers of a variable.
Definition: KokkosSystem.h:140
KOKKOS_FUNCTION unsigned int jvar() const
Get the coupled variable number.
Definition: KokkosDatum.h:330
The Kokkos object that holds thread-private data in the parallel operations of Kokkos residual object...
Definition: KokkosDatum.h:245
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:598