Kokkos Multivariable Kernels

framework/doc/content/syntax/KokkosKernels/index.md

@@ -102,6 +102,23 @@ See the following source codes of `KokkosCoupledForce` for another example of a
 !alert note
 [Every GPU function needs to be inlineable](syntax/Kokkos/index.md#kokkos_execution_space) and thus should be defined in headers.
 
+## Multi-variable Kernels
+
+Unlike the original MOOSE kernels which can only operate on a single variable, Kokkos-MOOSE kernels can operate on multiple variables in parallel.
+Instead of the `variable` parameter, you can specify the `variables` parameter with arbitrarily many variables you want the kernel to operate on.
+In the code, you can obtain which component of the variables you are currently operating on using `datum.comp()`.
+This provides additional parallelism over the variables and potentially enhances performance by providing higher chance of memory coalescing.
+The restrictions are that all variables should reside in the same system and have the same finite element shape function family and order.
+Also, all variables should have the same number of off-diagonal coupling.
+
+Multi-variable kernels can be useful in certain applications where the same kernel applies to multiple variables but using array variables is not desired.
+One example is neutronics, where there are multiple flux variables representing different energy groups that use the same kernel.
+It is desired to keep the flux variable of each group as separate standard variables, so that they can be individually coupled to other objects.
+Using an array variable here makes it cumbersome to access the flux solution of each group.
+
+!alert note
+All other Kokkos-MOOSE residual objects such as [BCs](syntax/KokkosBCs/index.md) and [NodalKernels](syntax/KokkosNodalKernels/index.md) also support operating on multiple variables with the same restrictions.
+
 ## Optimized Kernel Objects
 
 [Similarly to the original MOOSE](syntax/Kernels/index.md#optimized), Kokkos-MOOSE provides `Moose::Kokkos::KernelValue` and `Moose::Kokkos::KernelGrad` for creating an optimized kernel by factoring out test functions in residual and Jacobian calculations.
@@ -140,6 +157,7 @@ See the following source codes of `KokkosConvectionPrecompute` and `KokkosDiffus
 
 [Like the original MOOSE](syntax/Kernels/index.md#time-derivative), you can create a time-derivative kernel by subclassing `Moose::Kokkos::TimeKernel`.
 In Kokkos-MOOSE, the dummy `_qp` indexing of the `du_dot_du` term was lifted.
+Instead, the current variable component index `datum.comp()` should be passed as an argument.
 The following shows the conversion of the example presented in the original page into the Kokkos version:
 
 - For `computeQpResidual()` whose original code is:
@@ -163,7 +181,7 @@ return _test[_i][_qp] * _phi[_j][_qp] * _du_dot_du[_qp];
 the Kokkos version will look like:
 
 ```cpp
-return _test(datum, i, qp) * _phi(datum, j, qp) * _du_dot_du;
+return _test(datum, i, qp) * _phi(datum, j, qp) * _du_dot_du[datum.comp()];
 ```
 
 See the following source codes of `KokkosCoupledTimeDerivative` for an example of a time-derivative kernel:

framework/include/kokkos/base/KokkosDatum.h

@@ -343,7 +343,8 @@ class AssemblyDatum : public Datum
       _ife(systems[ivar.sys(comp)].getFETypeID(_ivar)),
       _jfe(systems[ivar.sys(comp)].getFETypeID(_jvar)),
       _n_idofs(assembly.getNumDofs(_elem.type, _ife)),
-      _n_jdofs(assembly.getNumDofs(_elem.type, _jfe))
+      _n_jdofs(assembly.getNumDofs(_elem.type, _jfe)),
+      _comp(comp)
   {
   }
   /**
@@ -367,7 +368,8 @@ class AssemblyDatum : public Datum
       _ivar(ivar.var(comp)),
       _jvar(jvar),
       _ife(systems[ivar.sys(comp)].getFETypeID(_ivar)),
-      _jfe(systems[ivar.sys(comp)].getFETypeID(_jvar))
+      _jfe(systems[ivar.sys(comp)].getFETypeID(_jvar)),
+      _comp(comp)
   {
   }
 
@@ -416,6 +418,11 @@ class AssemblyDatum : public Datum
    * @returns The variable FE type ID
    */
   KOKKOS_FUNCTION unsigned int jfe() const { return _jfe; }
+  /**
+   * Get the variable component
+   * @returns The variable component
+   */
+  KOKKOS_FUNCTION unsigned int comp() const { return _comp; }
 
 protected:
   /**
@@ -434,6 +441,10 @@ class AssemblyDatum : public Datum
    * Number of local DOFs
    */
   const unsigned int _n_idofs = 1, _n_jdofs = 1;
+  /**
+   * Variable component
+   */
+  const unsigned int _comp;
 };
 
 } // namespace Moose::Kokkos

framework/include/kokkos/base/KokkosResidualObject.h

@@ -58,7 +58,13 @@ class ResidualObject : public ::ResidualObject,
   };
   ///@}
 
-  virtual const MooseVariableBase & variable() const override { return _var; }
+  virtual const MooseVariableBase & variable() const override { return *_vars[0]; }
+
+  /**
+   * Get MOOSE variables
+   * @returns The MOOSE variables
+   */
+  const auto & variables() { return _vars; }
 
   virtual void computeOffDiagJacobian(unsigned int) override final
   {
@@ -71,10 +77,6 @@ class ResidualObject : public ::ResidualObject,
   }
 
 protected:
-  /**
-   * Reference of the MOOSE variable
-   */
-  MooseVariableFieldBase & _var;
   /**
    * Kokkos variable
    */
@@ -128,7 +130,7 @@ class ResidualObject : public ::ResidualObject,
   KOKKOS_FUNCTION void accumulateTaggedElementalResidual(const Real local_re,
                                                          const ContiguousElementID elem,
                                                          const unsigned int i,
-                                                         const unsigned int comp = 0) const;
+                                                         const unsigned int comp) const;
   /**
    * Accumulate or set local nodal residual contribution to tagged vectors
    * @param add The flag whether to add or set the local residual
@@ -139,7 +141,7 @@ class ResidualObject : public ::ResidualObject,
   KOKKOS_FUNCTION void accumulateTaggedNodalResidual(const bool add,
                                                      const Real local_re,
                                                      const ContiguousNodeID node,
-                                                     const unsigned int comp = 0) const;
+                                                     const unsigned int comp) const;
   /**
    * Accumulate local elemental Jacobian contribution to tagged matrices
    * @param local_ke The local elemental Jacobian contribution
@@ -154,7 +156,7 @@ class ResidualObject : public ::ResidualObject,
                                                        const unsigned int i,
                                                        const unsigned int j,
                                                        const unsigned int jvar,
-                                                       const unsigned int comp = 0) const;
+                                                       const unsigned int comp) const;
   /**
    * Accumulate or set local nodal Jacobian contribution to tagged matrices
    * @param add The flag whether to add or set the local residual
@@ -167,7 +169,7 @@ class ResidualObject : public ::ResidualObject,
                                                    const Real local_ke,
                                                    const ContiguousNodeID node,
                                                    const unsigned int jvar,
-                                                   const unsigned int comp = 0) const;
+                                                   const unsigned int comp) const;
 
   /**
    * The common loop structure template for computing elemental residual
@@ -185,6 +187,10 @@ class ResidualObject : public ::ResidualObject,
   KOKKOS_FUNCTION void computeJacobianInternal(AssemblyDatum & datum, function body) const;
 
 private:
+  /**
+   * MOOSE variables
+   */
+  std::vector<MooseVariableFieldBase *> _vars;
   /**
    * Tags this object operates on
    */
@@ -320,7 +326,7 @@ ResidualObject::computeResidualInternal(AssemblyDatum & datum, function body) co
     body(local_re - ib, ib, ie);
 
     for (unsigned int i = ib; i < ie; ++i)
-      accumulateTaggedElementalResidual(local_re[i - ib], datum.elem().id, i);
+      accumulateTaggedElementalResidual(local_re[i - ib], datum.elem().id, i, datum.comp());
   }
 }
 
@@ -350,7 +356,8 @@ ResidualObject::computeJacobianInternal(AssemblyDatum & datum, function body) co
       unsigned int i = ij % datum.n_jdofs();
       unsigned int j = ij / datum.n_jdofs();
 
-      accumulateTaggedElementalMatrix(local_ke[ij - ijb], datum.elem().id, i, j, datum.jvar());
+      accumulateTaggedElementalMatrix(
+          local_ke[ij - ijb], datum.elem().id, i, j, datum.jvar(), datum.comp());
     }
   }
 }

framework/include/kokkos/base/KokkosVariable.h

@@ -12,7 +12,7 @@
 #include "KokkosTypes.h"
 
 #include "MooseTypes.h"
-#include "MooseVariableBase.h"
+#include "MooseVariableFieldBase.h"
 #include "MoosePassKey.h"
 
 class Coupleable;
@@ -38,36 +38,85 @@ class Variable
    * @param variable The MOOSE variable
    * @param tag The vector tag ID
    */
-  Variable(const MooseVariableBase & variable, const TagID tag) { init(variable, tag); }
+  Variable(const MooseVariableFieldBase & variable, const TagID tag) { init(variable, tag); }
   /**
    * Constructor
    * Initialize the variable with a MOOSE variable and vector tag name
    * @param variable The MOOSE variable
    * @param tag_name The vector tag name
    */
-  Variable(const MooseVariableBase & variable, const TagName & tag_name = Moose::SOLUTION_TAG)
+  Variable(const MooseVariableFieldBase & variable, const TagName & tag_name = Moose::SOLUTION_TAG)
   {
     init(variable, tag_name);
   }
+  /**
+   * Constructor
+   * Initialize the variable with multiple MOOSE variables and vector tag ID
+   * @param variables The MOOSE variables
+   * @param tag The vector tag ID
+   */
+  ///@{
+  Variable(const std::vector<const MooseVariableFieldBase *> & variables, const TagID tag)
+  {
+    init(variables, tag);
+  }
+  Variable(const std::vector<MooseVariableFieldBase *> & variables, const TagID tag)
+  {
+    init(variables, tag);
+  }
+  ///@}
+  /**
+   * Constructor
+   * Initialize the variable with multiple MOOSE variables and vector tag name
+   * @param variables The MOOSE variables
+   * @param tag The vector tag ID
+   */
+  ///@{
+  Variable(const std::vector<const MooseVariableFieldBase *> & variables,
+           const TagName & tag_name = Moose::SOLUTION_TAG)
+  {
+    init(variables, tag_name);
+  }
+  Variable(const std::vector<MooseVariableFieldBase *> & variables,
+           const TagName & tag_name = Moose::SOLUTION_TAG)
+  {
+    init(variables, tag_name);
+  }
+  ///@}
+
   /**
    * Initialize the variable with a MOOSE variable and vector tag ID
    * @param variable The MOOSE variable
    * @param tag The vector tag ID
    */
-  void init(const MooseVariableBase & variable, const TagID tag);
+  void init(const MooseVariableFieldBase & variable, const TagID tag);
   /**
    * Initialize the variable with a MOOSE variable and vector tag name
    * @param variable The MOOSE variable
    * @param tag_name The vector tag name
    */
-  void init(const MooseVariableBase & variable, const TagName & tag_name = Moose::SOLUTION_TAG);
+  void init(const MooseVariableFieldBase & variable,
+            const TagName & tag_name = Moose::SOLUTION_TAG);
   /**
-   * Initialize the variable with coupled MOOSE variables
-   * @param variables The coupled MOOSE variables
+   * Initialize the variable with multiple MOOSE variables and vector tag ID
+   * @param variables The MOOSE variables
    * @param tag The vector tag ID
    */
-  void
-  init(const std::vector<const MooseVariableBase *> & variables, const TagID tag, CoupleableKey);
+  ///@{
+  void init(const std::vector<const MooseVariableFieldBase *> & variables, const TagID tag);
+  void init(const std::vector<MooseVariableFieldBase *> & variables, const TagID tag);
+  ///@}
+  /**
+   * Initialize the variable with multiple MOOSE variables and vector tag name
+   * @param variables The MOOSE variables
+   * @param tag_name The vector tag name
+   */
+  ///@{
+  void init(const std::vector<const MooseVariableFieldBase *> & variables,
+            const TagName & tag_name = Moose::SOLUTION_TAG);
+  void init(const std::vector<MooseVariableFieldBase *> & variables,
+            const TagName & tag_name = Moose::SOLUTION_TAG);
+  ///@}
   /**
    * Initialize the variable with coupled default values
    * @param values The default coupled values

framework/include/kokkos/base/KokkosVariableValue.h

@@ -11,7 +11,7 @@
 
 #include "KokkosDatum.h"
 
-#include "MooseVariableBase.h"
+#include "MooseVariableFieldBase.h"
 
 namespace Moose::Kokkos
 {
@@ -134,12 +134,32 @@ class VariableValue
    * @param tag The vector tag name
    * @param dof Whether to get DOF values
    */
-  VariableValue(const MooseVariableBase & var,
+  VariableValue(const MooseVariableFieldBase & var,
                 const TagName & tag = Moose::SOLUTION_TAG,
                 bool dof = false)
     : _var(var, tag), _dof(dof)
   {
   }
+  /**
+   * Constructor
+   * @param vars The MOOSE variables
+   * @param tag The vector tag name
+   * @param dof Whether to get DOF values
+   */
+  ///@{
+  VariableValue(const std::vector<const MooseVariableFieldBase *> & vars,
+                const TagName & tag = Moose::SOLUTION_TAG,
+                bool dof = false)
+    : _var(vars, tag), _dof(dof)
+  {
+  }
+  VariableValue(const std::vector<MooseVariableFieldBase *> & vars,
+                const TagName & tag = Moose::SOLUTION_TAG,
+                bool dof = false)
+    : _var(vars, tag), _dof(dof)
+  {
+  }
+  ///@}
 
   /**
    * Get whether the variable was coupled
@@ -150,11 +170,17 @@ class VariableValue
   /**
    * Get the current variable value
    * @param datum The Datum object of the current thread
-   * @param qp The local quadrature point index
+   * @param idx The local quadrature point index or DOF index
    * @param comp The variable component
    * @returns The variable value
    */
-  KOKKOS_FUNCTION Real operator()(Datum & datum, unsigned int qp, unsigned int comp = 0) const;
+  ///@{
+  KOKKOS_FUNCTION Real operator()(Datum & datum, unsigned int idx, unsigned int comp = 0) const;
+  KOKKOS_FUNCTION Real operator()(AssemblyDatum & datum, unsigned int idx) const
+  {
+    return this->operator()(datum, idx, datum.comp());
+  }
+  ///@}
 
   /**
    * Get the Kokkos variable
@@ -190,10 +216,27 @@ class VariableGradient
    * @param var The MOOSE variable
    * @param tag The vector tag name
    */
-  VariableGradient(const MooseVariableBase & var, const TagName & tag = Moose::SOLUTION_TAG)
+  VariableGradient(const MooseVariableFieldBase & var, const TagName & tag = Moose::SOLUTION_TAG)
     : _var(var, tag)
   {
   }
+  /**
+   * Constructor
+   * @param vars The MOOSE variables
+   * @param tag The vector tag name
+   */
+  ///@{
+  VariableGradient(const std::vector<const MooseVariableFieldBase *> & vars,
+                   const TagName & tag = Moose::SOLUTION_TAG)
+    : _var(vars, tag)
+  {
+  }
+  VariableGradient(const std::vector<MooseVariableFieldBase *> vars,
+                   const TagName & tag = Moose::SOLUTION_TAG)
+    : _var(vars, tag)
+  {
+  }
+  ///@}
 
   /**
    * Get whether the variable was coupled
@@ -204,12 +247,17 @@ class VariableGradient
   /**
    * Get the current variable gradient
    * @param datum The Datum object of the current thread
-   * @param idx The local quadrature point or DOF index
+   * @param qp The local quadrature point index
    * @param comp The variable component
    * @returns The variable gradient
    */
-  KOKKOS_FUNCTION Real3 operator()(Datum & datum, unsigned int idx, unsigned int comp = 0) const;
-
+  ///@{
+  KOKKOS_FUNCTION Real3 operator()(Datum & datum, unsigned int qp, unsigned int comp = 0) const;
+  KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int qp) const
+  {
+    return this->operator()(datum, qp, datum.comp());
+  }
+  ///@}
   /**
    * Get the Kokkos variable
    * @returns The Kokkos variable