From 55f6469cf75fe063cba0bd26e094b6696c66ceb1 Mon Sep 17 00:00:00 2001 From: Namjae Choi Date: Fri, 27 Feb 2026 08:12:26 -0700 Subject: [PATCH 1/3] Support multi-variable Kokkos kernels/BCs #30655 --- framework/include/kokkos/base/KokkosDatum.h | 15 ++++- .../kokkos/base/KokkosResidualObject.h | 29 +++++--- .../include/kokkos/base/KokkosVariable.h | 67 ++++++++++++++++--- .../include/kokkos/base/KokkosVariableValue.h | 64 +++++++++++++++--- .../kokkos/bcs/KokkosDirichletBCBase.h | 15 +++-- .../include/kokkos/bcs/KokkosIntegratedBC.h | 19 +++--- framework/include/kokkos/bcs/KokkosNodalBC.h | 37 +++++----- .../include/kokkos/kernels/KokkosKernel.h | 21 +++--- .../kokkos/kernels/KokkosTimeDerivative.h | 4 +- .../include/kokkos/kernels/KokkosTimeKernel.h | 7 +- .../nodalkernels/KokkosBoundNodalKernel.h | 5 +- .../kokkos/nodalkernels/KokkosNodalKernel.h | 41 +++++++----- .../KokkosTimeDerivativeNodalKernel.h | 4 +- .../nodalkernels/KokkosTimeNodalKernel.h | 7 +- .../src/kokkos/base/KokkosResidualObject.K | 42 ++++++++++-- framework/src/kokkos/base/KokkosVariable.K | 28 ++++++-- framework/src/kokkos/bcs/KokkosIntegratedBC.K | 24 +++++-- framework/src/kokkos/bcs/KokkosNodalBC.K | 23 +++++-- .../src/kokkos/interfaces/KokkosCoupleable.K | 14 ++-- .../src/kokkos/kernels/KokkosCoupledForce.K | 7 +- framework/src/kokkos/kernels/KokkosKernel.K | 24 +++++-- .../kokkos/kernels/KokkosMatCoupledForce.K | 9 +-- .../src/kokkos/kernels/KokkosTimeKernel.K | 16 ++++- .../KokkosCoupledForceNodalKernel.K | 9 +-- .../kokkos/nodalkernels/KokkosNodalKernel.K | 24 +++++-- .../nodalkernels/KokkosTimeNodalKernel.K | 16 ++++- 26 files changed, 422 insertions(+), 149 deletions(-) diff --git a/framework/include/kokkos/base/KokkosDatum.h b/framework/include/kokkos/base/KokkosDatum.h index fb2afedc61b2..a697cf7a519b 100644 --- a/framework/include/kokkos/base/KokkosDatum.h +++ b/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 diff --git a/framework/include/kokkos/base/KokkosResidualObject.h b/framework/include/kokkos/base/KokkosResidualObject.h index 608a8d63822e..1c9ab89e790b 100644 --- a/framework/include/kokkos/base/KokkosResidualObject.h +++ b/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 _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()); } } } diff --git a/framework/include/kokkos/base/KokkosVariable.h b/framework/include/kokkos/base/KokkosVariable.h index eddf6b23a278..0f017d57180c 100644 --- a/framework/include/kokkos/base/KokkosVariable.h +++ b/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 & variables, const TagID tag) + { + init(variables, tag); + } + Variable(const std::vector & 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 & variables, + const TagName & tag_name = Moose::SOLUTION_TAG) + { + init(variables, tag_name); + } + Variable(const std::vector & 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 & variables, const TagID tag, CoupleableKey); + ///@{ + void init(const std::vector & variables, const TagID tag); + void init(const std::vector & 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 & variables, + const TagName & tag_name = Moose::SOLUTION_TAG); + void init(const std::vector & variables, + const TagName & tag_name = Moose::SOLUTION_TAG); + ///@} /** * Initialize the variable with coupled default values * @param values The default coupled values diff --git a/framework/include/kokkos/base/KokkosVariableValue.h b/framework/include/kokkos/base/KokkosVariableValue.h index 0cadb3e76c0d..06d891a191ff 100644 --- a/framework/include/kokkos/base/KokkosVariableValue.h +++ b/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 & vars, + const TagName & tag = Moose::SOLUTION_TAG, + bool dof = false) + : _var(vars, tag), _dof(dof) + { + } + VariableValue(const std::vector & 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 & vars, + const TagName & tag = Moose::SOLUTION_TAG) + : _var(vars, tag) + { + } + VariableGradient(const std::vector 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 diff --git a/framework/include/kokkos/bcs/KokkosDirichletBCBase.h b/framework/include/kokkos/bcs/KokkosDirichletBCBase.h index e907e5bf73d7..144ea4f38c59 100644 --- a/framework/include/kokkos/bcs/KokkosDirichletBCBase.h +++ b/framework/include/kokkos/bcs/KokkosDirichletBCBase.h @@ -84,10 +84,12 @@ template void DirichletBCBase::presetSolution(TagID tag) { + _thread.resize({variables().size(), numKokkosBoundaryNodes()}); + _solution_tag = tag; ::Kokkos::parallel_for( - ::Kokkos::RangePolicy>(0, numKokkosBoundaryNodes()), + ::Kokkos::RangePolicy>(0, _thread.size()), *static_cast(this)); } @@ -96,14 +98,17 @@ KOKKOS_FUNCTION void DirichletBCBase::operator()(const ThreadID tid) const { auto bc = static_cast(this); - auto node = kokkosBoundaryNodeID(tid); - auto & sys = kokkosSystem(_kokkos_var.sys()); - auto dof = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var()); + + auto comp = _thread(tid, 0); + auto node = kokkosBoundaryNodeID(_thread(tid, 1)); + auto & sys = kokkosSystem(_kokkos_var.sys(comp)); + auto dof = sys.getNodeLocalDofIndex(node, 0, _kokkos_var.var(comp)); if (dof == libMesh::DofObject::invalid_id) return; - AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var()); + AssemblyDatum datum( + node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var(comp), comp); sys.getVectorDofValue(dof, _solution_tag) = bc->computeValue(0, datum); } diff --git a/framework/include/kokkos/bcs/KokkosIntegratedBC.h b/framework/include/kokkos/bcs/KokkosIntegratedBC.h index bced2478eaac..6ef55528602d 100644 --- a/framework/include/kokkos/bcs/KokkosIntegratedBC.h +++ b/framework/include/kokkos/bcs/KokkosIntegratedBC.h @@ -209,10 +209,11 @@ template KOKKOS_FUNCTION void IntegratedBC::operator()(ResidualLoop, const ThreadID tid, const Derived & bc) const { - auto [elem, side] = kokkosBoundaryElementSideID(tid); + auto comp = _thread(tid, 0); + auto [elem, side] = kokkosBoundaryElementSideID(_thread(tid, 1)); AssemblyDatum datum( - elem, side, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var()); + elem, side, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var(comp), comp); bc.computeResidualInternal(bc, datum); } @@ -221,10 +222,11 @@ template KOKKOS_FUNCTION void IntegratedBC::operator()(JacobianLoop, const ThreadID tid, const Derived & bc) const { - auto [elem, side] = kokkosBoundaryElementSideID(tid); + auto comp = _thread(tid, 0); + auto [elem, side] = kokkosBoundaryElementSideID(_thread(tid, 1)); AssemblyDatum datum( - elem, side, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var()); + elem, side, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var(comp), comp); bc.computeJacobianInternal(bc, datum); } @@ -233,15 +235,16 @@ template KOKKOS_FUNCTION void IntegratedBC::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & bc) const { - auto [elem, side] = kokkosBoundaryElementSideID(_thread(tid, 1)); + auto comp = _thread(tid, 0); + auto [elem, side] = kokkosBoundaryElementSideID(_thread(tid, 2)); - auto & sys = kokkosSystem(_kokkos_var.sys()); - auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)]; + auto & sys = kokkosSystem(_kokkos_var.sys(comp)); + auto jvar = sys.getCoupling(_kokkos_var.var(comp))[_thread(tid, 1)]; if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain)) return; - AssemblyDatum datum(elem, side, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar); + AssemblyDatum datum(elem, side, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar, comp); bc.computeOffDiagJacobianInternal(bc, datum); } diff --git a/framework/include/kokkos/bcs/KokkosNodalBC.h b/framework/include/kokkos/bcs/KokkosNodalBC.h index 58e4b0eb9eac..aa6fd8b16683 100644 --- a/framework/include/kokkos/bcs/KokkosNodalBC.h +++ b/framework/include/kokkos/bcs/KokkosNodalBC.h @@ -143,53 +143,58 @@ template KOKKOS_FUNCTION void NodalBC::operator()(ResidualLoop, const ThreadID tid, const Derived & bc) const { - auto node = kokkosBoundaryNodeID(tid); - auto & sys = kokkosSystem(_kokkos_var.sys()); + auto comp = _thread(tid, 0); + auto node = kokkosBoundaryNodeID(_thread(tid, 1)); + auto & sys = kokkosSystem(_kokkos_var.sys(comp)); - if (!sys.isNodalDefined(node, _kokkos_var.var())) + if (!sys.isNodalDefined(node, _kokkos_var.var(comp))) return; - AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var()); + AssemblyDatum datum( + node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var(comp), comp); Real local_re = bc.computeQpResidualShim(bc, 0, datum); - accumulateTaggedNodalResidual(false, local_re, node); + accumulateTaggedNodalResidual(false, local_re, node, comp); } template KOKKOS_FUNCTION void NodalBC::operator()(JacobianLoop, const ThreadID tid, const Derived & bc) const { - auto node = kokkosBoundaryNodeID(tid); - auto & sys = kokkosSystem(_kokkos_var.sys()); + auto comp = _thread(tid, 0); + auto node = kokkosBoundaryNodeID(_thread(tid, 1)); + auto & sys = kokkosSystem(_kokkos_var.sys(comp)); - if (!sys.isNodalDefined(node, _kokkos_var.var())) + if (!sys.isNodalDefined(node, _kokkos_var.var(comp))) return; - AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var()); + AssemblyDatum datum( + node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var(comp), comp); Real local_ke = bc.computeQpJacobianShim(bc, 0, datum); // This initializes the row to zero except the diagonal - accumulateTaggedNodalMatrix(false, local_ke, node, _kokkos_var.var()); + accumulateTaggedNodalMatrix(false, local_ke, node, _kokkos_var.var(comp), comp); } template KOKKOS_FUNCTION void NodalBC::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & bc) const { - auto node = kokkosBoundaryNodeID(_thread(tid, 1)); - auto & sys = kokkosSystem(_kokkos_var.sys()); - auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)]; + auto comp = _thread(tid, 0); + auto node = kokkosBoundaryNodeID(_thread(tid, 2)); + auto & sys = kokkosSystem(_kokkos_var.sys(comp)); + auto jvar = sys.getCoupling(_kokkos_var.var(comp))[_thread(tid, 1)]; - if (!sys.isNodalDefined(node, _kokkos_var.var())) + if (!sys.isNodalDefined(node, _kokkos_var.var(comp))) return; - AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar); + AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar, comp); Real local_ke = bc.computeQpOffDiagJacobianShim(bc, jvar, 0, datum); - accumulateTaggedNodalMatrix(true, local_ke, node, jvar); + accumulateTaggedNodalMatrix(true, local_ke, node, jvar, comp); } } // namespace Moose::Kokkos diff --git a/framework/include/kokkos/kernels/KokkosKernel.h b/framework/include/kokkos/kernels/KokkosKernel.h index e58cdc71c2f5..48813462fc5d 100644 --- a/framework/include/kokkos/kernels/KokkosKernel.h +++ b/framework/include/kokkos/kernels/KokkosKernel.h @@ -209,14 +209,16 @@ template KOKKOS_FUNCTION void Kernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const { - auto elem = kokkosBlockElementID(tid); + auto comp = _thread(tid, 0); + auto elem = kokkosBlockElementID(_thread(tid, 1)); AssemblyDatum datum(elem, libMesh::invalid_uint, kokkosAssembly(), kokkosSystems(), _kokkos_var, - _kokkos_var.var()); + _kokkos_var.var(comp), + comp); kernel.computeResidualInternal(kernel, datum); } @@ -225,14 +227,16 @@ template KOKKOS_FUNCTION void Kernel::operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const { - auto elem = kokkosBlockElementID(tid); + auto comp = _thread(tid, 0); + auto elem = kokkosBlockElementID(_thread(tid, 1)); AssemblyDatum datum(elem, libMesh::invalid_uint, kokkosAssembly(), kokkosSystems(), _kokkos_var, - _kokkos_var.var()); + _kokkos_var.var(comp), + comp); kernel.computeJacobianInternal(kernel, datum); } @@ -241,16 +245,17 @@ template KOKKOS_FUNCTION void Kernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const { - auto elem = kokkosBlockElementID(_thread(tid, 1)); + auto comp = _thread(tid, 0); + auto elem = kokkosBlockElementID(_thread(tid, 2)); - auto & sys = kokkosSystem(_kokkos_var.sys()); - auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)]; + auto & sys = kokkosSystem(_kokkos_var.sys(comp)); + auto jvar = sys.getCoupling(_kokkos_var.var(comp))[_thread(tid, 1)]; if (!sys.isVariableActive(jvar, kokkosMesh().getElementInfo(elem).subdomain)) return; AssemblyDatum datum( - elem, libMesh::invalid_uint, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar); + elem, libMesh::invalid_uint, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar, comp); kernel.computeOffDiagJacobianInternal(kernel, datum); } diff --git a/framework/include/kokkos/kernels/KokkosTimeDerivative.h b/framework/include/kokkos/kernels/KokkosTimeDerivative.h index 99b6ae172baa..9f18227a6907 100644 --- a/framework/include/kokkos/kernels/KokkosTimeDerivative.h +++ b/framework/include/kokkos/kernels/KokkosTimeDerivative.h @@ -73,7 +73,7 @@ KokkosTimeDerivative::computeJacobianInternal(const Derived & kernel, AssemblyDa unsigned int j = ij / datum.n_jdofs(); accumulateTaggedElementalMatrix( - local_ke[ij - ijb], datum.elem().id, i, _lumping ? i : j, datum.jvar()); + local_ke[ij - ijb], datum.elem().id, i, _lumping ? i : j, datum.jvar(), datum.comp()); } } } @@ -92,5 +92,5 @@ KokkosTimeDerivative::computeQpJacobian(const unsigned int i, const unsigned int qp, AssemblyDatum & datum) const { - 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()]; } diff --git a/framework/include/kokkos/kernels/KokkosTimeKernel.h b/framework/include/kokkos/kernels/KokkosTimeKernel.h index b7db90b6328a..bd20115ab9c5 100644 --- a/framework/include/kokkos/kernels/KokkosTimeKernel.h +++ b/framework/include/kokkos/kernels/KokkosTimeKernel.h @@ -27,6 +27,11 @@ class TimeKernel : public Kernel */ TimeKernel(const InputParameters & parameters); + /** + * Copy constructor for parallel dispatch + */ + TimeKernel(const TimeKernel & object); + /** * Hook for additional computation for residual after the standard calls * @param ib The beginning element-local DOF index @@ -56,7 +61,7 @@ class TimeKernel : public Kernel /** * Derivative of u_dot with respect to u */ - const Scalar _du_dot_du; + Array _du_dot_du; }; template diff --git a/framework/include/kokkos/nodalkernels/KokkosBoundNodalKernel.h b/framework/include/kokkos/nodalkernels/KokkosBoundNodalKernel.h index ead3de94f520..fb850349e926 100644 --- a/framework/include/kokkos/nodalkernels/KokkosBoundNodalKernel.h +++ b/framework/include/kokkos/nodalkernels/KokkosBoundNodalKernel.h @@ -61,8 +61,9 @@ template KokkosBoundNodalKernel::KokkosBoundNodalKernel(const InputParameters & parameters) : NodalKernel(parameters), _v_var(coupled("v")), _v(kokkosCoupledNodalValue("v")) { - if (_var.number() == _v_var) - paramError("v", "Coupled variable needs to be different from 'variable'"); + for (const auto var : variables()) + if (var->number() == _v_var) + paramError("v", "Coupled variable needs to be different from variable"); } template diff --git a/framework/include/kokkos/nodalkernels/KokkosNodalKernel.h b/framework/include/kokkos/nodalkernels/KokkosNodalKernel.h index 5ad5f65563da..20b08ea4d8c4 100644 --- a/framework/include/kokkos/nodalkernels/KokkosNodalKernel.h +++ b/framework/include/kokkos/nodalkernels/KokkosNodalKernel.h @@ -149,53 +149,60 @@ template KOKKOS_FUNCTION void NodalKernel::operator()(ResidualLoop, const ThreadID tid, const Derived & kernel) const { - auto node = _boundary_restricted ? kokkosBoundaryNodeID(tid) : kokkosBlockNodeID(tid); - auto & sys = kokkosSystem(_kokkos_var.sys()); + auto comp = _thread(tid, 0); + auto node = _boundary_restricted ? kokkosBoundaryNodeID(_thread(tid, 1)) + : kokkosBlockNodeID(_thread(tid, 1)); + auto & sys = kokkosSystem(_kokkos_var.sys(comp)); - if (!sys.isNodalDefined(node, _kokkos_var.var())) + if (!sys.isNodalDefined(node, _kokkos_var.var(comp))) return; - AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var()); + AssemblyDatum datum( + node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var(comp), comp); Real local_re = kernel.computeQpResidualShim(kernel, 0, datum); - accumulateTaggedNodalResidual(true, local_re, node); + accumulateTaggedNodalResidual(true, local_re, node, comp); } template KOKKOS_FUNCTION void NodalKernel::operator()(JacobianLoop, const ThreadID tid, const Derived & kernel) const { - auto node = _boundary_restricted ? kokkosBoundaryNodeID(tid) : kokkosBlockNodeID(tid); - auto & sys = kokkosSystem(_kokkos_var.sys()); + auto comp = _thread(tid, 0); + auto node = _boundary_restricted ? kokkosBoundaryNodeID(_thread(tid, 1)) + : kokkosBlockNodeID(_thread(tid, 1)); + auto & sys = kokkosSystem(_kokkos_var.sys(comp)); - if (!sys.isNodalDefined(node, _kokkos_var.var())) + if (!sys.isNodalDefined(node, _kokkos_var.var(comp))) return; - AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var()); + AssemblyDatum datum( + node, kokkosAssembly(), kokkosSystems(), _kokkos_var, _kokkos_var.var(comp), comp); Real local_ke = kernel.computeQpJacobianShim(kernel, 0, datum); - accumulateTaggedNodalMatrix(true, local_ke, node, _kokkos_var.var()); + accumulateTaggedNodalMatrix(true, local_ke, node, _kokkos_var.var(comp), comp); } template KOKKOS_FUNCTION void NodalKernel::operator()(OffDiagJacobianLoop, const ThreadID tid, const Derived & kernel) const { - auto node = _boundary_restricted ? kokkosBoundaryNodeID(_thread(tid, 1)) - : kokkosBlockNodeID(_thread(tid, 1)); - auto & sys = kokkosSystem(_kokkos_var.sys()); - auto jvar = sys.getCoupling(_kokkos_var.var())[_thread(tid, 0)]; + auto comp = _thread(tid, 0); + auto node = _boundary_restricted ? kokkosBoundaryNodeID(_thread(tid, 2)) + : kokkosBlockNodeID(_thread(tid, 2)); + auto & sys = kokkosSystem(_kokkos_var.sys(comp)); + auto jvar = sys.getCoupling(_kokkos_var.var(comp))[_thread(tid, 1)]; - if (!sys.isNodalDefined(node, _kokkos_var.var())) + if (!sys.isNodalDefined(node, _kokkos_var.var(comp))) return; - AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar); + AssemblyDatum datum(node, kokkosAssembly(), kokkosSystems(), _kokkos_var, jvar, comp); Real local_ke = kernel.computeQpOffDiagJacobianShim(kernel, jvar, 0, datum); - accumulateTaggedNodalMatrix(true, local_ke, node, jvar); + accumulateTaggedNodalMatrix(true, local_ke, node, jvar, comp); } } // namespace Moose::Kokkos diff --git a/framework/include/kokkos/nodalkernels/KokkosTimeDerivativeNodalKernel.h b/framework/include/kokkos/nodalkernels/KokkosTimeDerivativeNodalKernel.h index 880d72ea2355..220462ec2eb7 100644 --- a/framework/include/kokkos/nodalkernels/KokkosTimeDerivativeNodalKernel.h +++ b/framework/include/kokkos/nodalkernels/KokkosTimeDerivativeNodalKernel.h @@ -31,7 +31,7 @@ KokkosTimeDerivativeNodalKernel::computeQpResidual(const unsigned int qp, KOKKOS_FUNCTION inline Real KokkosTimeDerivativeNodalKernel::computeQpJacobian(const unsigned int /* qp */, - AssemblyDatum & /* datum */) const + AssemblyDatum & datum) const { - return _du_dot_du; + return _du_dot_du[datum.comp()]; } diff --git a/framework/include/kokkos/nodalkernels/KokkosTimeNodalKernel.h b/framework/include/kokkos/nodalkernels/KokkosTimeNodalKernel.h index 75ae4bd1f0ed..472ed22d0613 100644 --- a/framework/include/kokkos/nodalkernels/KokkosTimeNodalKernel.h +++ b/framework/include/kokkos/nodalkernels/KokkosTimeNodalKernel.h @@ -27,6 +27,11 @@ class TimeNodalKernel : public NodalKernel */ TimeNodalKernel(const InputParameters & parameters); + /** + * Copy constructor for parallel dispatch + */ + TimeNodalKernel(const TimeNodalKernel & object); + protected: /** * Time derivative of the current solution at nodes @@ -35,7 +40,7 @@ class TimeNodalKernel : public NodalKernel /** * Derivative of u_dot with respect to u */ - const Scalar _du_dot_du; + Array _du_dot_du; }; } // namespace Moose::Kokkos diff --git a/framework/src/kokkos/base/KokkosResidualObject.K b/framework/src/kokkos/base/KokkosResidualObject.K index 20fbc513a1f2..9d3c45a39f59 100644 --- a/framework/src/kokkos/base/KokkosResidualObject.K +++ b/framework/src/kokkos/base/KokkosResidualObject.K @@ -19,6 +19,10 @@ ResidualObject::validParams() { auto params = ::ResidualObject::validParams(); + params.makeParamNotRequired("variable"); + params.addParam>( + "variables", "The names of the variables that this residual object operates on"); + params.addParam("use_displaced_mesh", false, "Whether or not this object should use the " @@ -42,17 +46,41 @@ ResidualObject::ResidualObject(const InputParameters & parameters, MeshHolder(*_fe_problem.mesh().getKokkosMesh()), AssemblyHolder(_fe_problem.kokkosAssembly()), SystemHolder(_fe_problem.getKokkosSystems()), - _var(_subproblem.getVariable(_tid, - getParam("variable"), - Moose::VarKindType::VAR_SOLVER, - field_type)), _t(TransientInterface::_t), _t_old(TransientInterface::_t_old), _t_step(TransientInterface::_t_step), _dt(TransientInterface::_dt), _dt_old(TransientInterface::_dt_old) { - _kokkos_var.init(_var); + if (isParamValid("variable") && isParamValid("variables")) + mooseError("Cannot specify both 'variable' and 'variables' at the same time."); + else if (!isParamValid("variable") && !isParamValid("variables")) + mooseError("Either 'variable' or 'variables' should be specified."); + + std::vector variables; + + if (isParamValid("variable")) + variables.push_back(getParam("variable")); + else + variables = getParam>("variables"); + + for (const auto & variable : variables) + _vars.push_back( + &_subproblem.getVariable(_tid, variable, Moose::VarKindType::VAR_SOLVER, field_type)); + + if (_vars.size() > 1) + for (unsigned int i = 1; i < _vars.size(); ++i) + { + if (_vars[0]->sys().number() != _vars[i]->sys().number()) + mooseError( + "For multi-variable kernels/BCs, all variables should reside in the same system."); + + if (_vars[0]->feType() != _vars[i]->feType()) + mooseError("For multi-variable kernels/BCs, all variables should have the same FE family " + "and order."); + } + + _kokkos_var.init(_vars); } ResidualObject::ResidualObject(const ResidualObject & object) @@ -60,14 +88,14 @@ ResidualObject::ResidualObject(const ResidualObject & object) MeshHolder(object), AssemblyHolder(object), SystemHolder(object), - _var(object._var), _kokkos_var(object._kokkos_var), _thread(object._thread), _t(object._t), _t_old(object._t_old), _t_step(object._t_step), _dt(object._dt), - _dt_old(object._dt_old) + _dt_old(object._dt_old), + _vars(object._vars) { _vector_tags = object.getVectorTags({}); _matrix_tags = object.getMatrixTags({}); diff --git a/framework/src/kokkos/base/KokkosVariable.K b/framework/src/kokkos/base/KokkosVariable.K index b00b73837611..0920873eed16 100644 --- a/framework/src/kokkos/base/KokkosVariable.K +++ b/framework/src/kokkos/base/KokkosVariable.K @@ -16,13 +16,13 @@ namespace Moose::Kokkos { void -Variable::init(const MooseVariableBase & variable, const TagName & tag_name) +Variable::init(const MooseVariableFieldBase & variable, const TagName & tag_name) { init(variable, variable.sys().feProblem().getVectorTagID(tag_name)); } void -Variable::init(const MooseVariableBase & variable, const TagID tag) +Variable::init(const MooseVariableFieldBase & variable, const TagID tag) { _initialized = true; _coupled = true; @@ -44,9 +44,27 @@ Variable::init(const MooseVariableBase & variable, const TagID tag) } void -Variable::init(const std::vector & variables, - const TagID tag, - CoupleableKey) +Variable::init(const std::vector & variables, const TagName & tag_name) +{ + init(std::vector(variables.begin(), variables.end()), + variables[0]->sys().feProblem().getVectorTagID(tag_name)); +} + +void +Variable::init(const std::vector & variables, + const TagName & tag_name) +{ + init(variables, variables[0]->sys().feProblem().getVectorTagID(tag_name)); +} + +void +Variable::init(const std::vector & variables, const TagID tag) +{ + init(std::vector(variables.begin(), variables.end()), tag); +} + +void +Variable::init(const std::vector & variables, const TagID tag) { _initialized = true; _coupled = true; diff --git a/framework/src/kokkos/bcs/KokkosIntegratedBC.K b/framework/src/kokkos/bcs/KokkosIntegratedBC.K index 8d047c4751b4..35bf193faaeb 100644 --- a/framework/src/kokkos/bcs/KokkosIntegratedBC.K +++ b/framework/src/kokkos/bcs/KokkosIntegratedBC.K @@ -25,16 +25,18 @@ IntegratedBC::IntegratedBC(const InputParameters & parameters) _grad_test(), _phi(), _grad_phi(), - _u(_var), - _grad_u(_var) + _u(_kokkos_var), + _grad_u(_kokkos_var) { - addMooseVariableDependency(&_var); + addMooseVariableDependency(variables()); } void IntegratedBC::computeResidual() { - Policy policy(0, numKokkosBoundarySides()); + _thread.resize({variables().size(), numKokkosBoundarySides()}); + + Policy policy(0, _thread.size()); if (!_residual_dispatcher) _residual_dispatcher = DispatcherRegistry::build(this, type()); @@ -47,7 +49,9 @@ IntegratedBC::computeJacobian() { if (DispatcherRegistry::hasUserMethod(type())) { - Policy policy(0, numKokkosBoundarySides()); + _thread.resize({variables().size(), numKokkosBoundarySides()}); + + Policy policy(0, _thread.size()); if (!_jacobian_dispatcher) _jacobian_dispatcher = DispatcherRegistry::build(this, type()); @@ -59,7 +63,15 @@ IntegratedBC::computeJacobian() { auto & sys = kokkosSystem(_kokkos_var.sys()); - _thread.resize({sys.getCoupling(_kokkos_var.var()).size(), numKokkosBoundarySides()}); + if (variables().size() > 1) + for (unsigned int i = 1; i < variables().size(); ++i) + if (sys.getCoupling(_kokkos_var.var(0)).size() != + sys.getCoupling(_kokkos_var.var(i)).size()) + mooseError("For multi-variable kernels/BCs, all variables should have the same number of " + "off-diagonal coupling variables."); + + _thread.resize( + {variables().size(), sys.getCoupling(_kokkos_var.var()).size(), numKokkosBoundarySides()}); Policy policy(0, _thread.size()); diff --git a/framework/src/kokkos/bcs/KokkosNodalBC.K b/framework/src/kokkos/bcs/KokkosNodalBC.K index 3bcefc81ffd4..9d63eb5e8f19 100644 --- a/framework/src/kokkos/bcs/KokkosNodalBC.K +++ b/framework/src/kokkos/bcs/KokkosNodalBC.K @@ -20,16 +20,17 @@ NodalBC::validParams() } NodalBC::NodalBC(const InputParameters & parameters) - : NodalBCBase(parameters, Moose::VarFieldType::VAR_FIELD_STANDARD), - _u(_var, Moose::SOLUTION_TAG, true) + : NodalBCBase(parameters, Moose::VarFieldType::VAR_FIELD_STANDARD), _u(_kokkos_var, true) { - addMooseVariableDependency(&_var); + addMooseVariableDependency(variables()); } void NodalBC::computeResidual() { - Policy policy(0, numKokkosBoundaryNodes()); + _thread.resize({variables().size(), numKokkosBoundaryNodes()}); + + Policy policy(0, _thread.size()); if (!_residual_dispatcher) _residual_dispatcher = DispatcherRegistry::build(this, type()); @@ -40,7 +41,9 @@ NodalBC::computeResidual() void NodalBC::computeJacobian() { - Policy policy(0, numKokkosBoundaryNodes()); + _thread.resize({variables().size(), numKokkosBoundaryNodes()}); + + Policy policy(0, _thread.size()); if (!_jacobian_dispatcher) _jacobian_dispatcher = DispatcherRegistry::build(this, type()); @@ -51,7 +54,15 @@ NodalBC::computeJacobian() { auto & sys = kokkosSystem(_kokkos_var.sys()); - _thread.resize({sys.getCoupling(_kokkos_var.var()).size(), numKokkosBoundaryNodes()}); + if (variables().size() > 1) + for (unsigned int i = 1; i < variables().size(); ++i) + if (sys.getCoupling(_kokkos_var.var(0)).size() != + sys.getCoupling(_kokkos_var.var(i)).size()) + mooseError("For multi-variable kernels/BCs, all variables should have the same number of " + "off-diagonal coupling variables."); + + _thread.resize( + {variables().size(), sys.getCoupling(_kokkos_var.var()).size(), numKokkosBoundaryNodes()}); Policy policy(0, _thread.size()); diff --git a/framework/src/kokkos/interfaces/KokkosCoupleable.K b/framework/src/kokkos/interfaces/KokkosCoupleable.K index 737b2d185740..3ca8e8da37f3 100644 --- a/framework/src/kokkos/interfaces/KokkosCoupleable.K +++ b/framework/src/kokkos/interfaces/KokkosCoupleable.K @@ -13,6 +13,8 @@ #include "SystemBase.h" #include "FEProblemBase.h" +using CoupleableKey = Moose::Kokkos::Variable::CoupleableKey; + Moose::Kokkos::Variable Coupleable::kokkosCoupledVectorTagVariable(const std::string & var_name, const std::string & tag_name, @@ -34,10 +36,10 @@ Coupleable::kokkosCoupledVectorTagVariable(const std::string & var_name, const_cast(this)->addFEVariableCoupleableVectorTag(tag); - variable.init({var}, tag, {}); + variable.init(*var, tag); } else - variable.init({_c_parameters.defaultCoupledValue(var_name, comp)}, {}); + variable.init({_c_parameters.defaultCoupledValue(var_name, comp)}, CoupleableKey{}); return variable; } @@ -52,7 +54,7 @@ Coupleable::kokkosCoupledVectorTagVariables(const std::string & var_name, if (isCoupled(var_name)) { - std::vector vars; + std::vector vars; for (unsigned int comp = 0; comp < components; ++comp) { @@ -71,7 +73,7 @@ Coupleable::kokkosCoupledVectorTagVariables(const std::string & var_name, const_cast(this)->addFEVariableCoupleableVectorTag(tag); - variable.init({vars}, tag, {}); + variable.init(vars, tag); } else { @@ -80,7 +82,7 @@ Coupleable::kokkosCoupledVectorTagVariables(const std::string & var_name, for (unsigned int comp = 0; comp < components; ++comp) default_values[comp] = _c_parameters.defaultCoupledValue(var_name, comp); - variable.init(default_values, {}); + variable.init(default_values, CoupleableKey{}); } return variable; @@ -91,7 +93,7 @@ Coupleable::kokkosZeroVariable() const { Moose::Kokkos::Variable variable; - variable.init({0}, {}); + variable.init({0}, CoupleableKey{}); return variable; } diff --git a/framework/src/kokkos/kernels/KokkosCoupledForce.K b/framework/src/kokkos/kernels/KokkosCoupledForce.K index 60e4382123b7..601f0f2b23cc 100644 --- a/framework/src/kokkos/kernels/KokkosCoupledForce.K +++ b/framework/src/kokkos/kernels/KokkosCoupledForce.K @@ -31,7 +31,8 @@ KokkosCoupledForce::KokkosCoupledForce(const InputParameters & parameters) _v(kokkosCoupledValue("v")), _coef(getParam("coef")) { - if (_var.number() == _v_var) - paramError( - "v", "Coupled variable 'v' needs to be different from 'variable' with KokkosCoupledForce."); + for (const auto var : variables()) + if (var->number() == _v_var) + paramError( + "v", "Coupled variable 'v' needs to be different from variable with KokkosCoupledForce."); } diff --git a/framework/src/kokkos/kernels/KokkosKernel.K b/framework/src/kokkos/kernels/KokkosKernel.K index 9fcc264c5fa7..6d52083d0ad6 100644 --- a/framework/src/kokkos/kernels/KokkosKernel.K +++ b/framework/src/kokkos/kernels/KokkosKernel.K @@ -25,16 +25,18 @@ Kernel::Kernel(const InputParameters & parameters) _grad_test(), _phi(), _grad_phi(), - _u(_var), - _grad_u(_var) + _u(_kokkos_var), + _grad_u(_kokkos_var) { - addMooseVariableDependency(&_var); + addMooseVariableDependency(variables()); } void Kernel::computeResidual() { - Policy policy(0, numKokkosBlockElements()); + _thread.resize({variables().size(), numKokkosBlockElements()}); + + Policy policy(0, _thread.size()); if (!_residual_dispatcher) _residual_dispatcher = DispatcherRegistry::build(this, type()); @@ -47,7 +49,9 @@ Kernel::computeJacobian() { if (DispatcherRegistry::hasUserMethod(type())) { - Policy policy(0, numKokkosBlockElements()); + _thread.resize({variables().size(), numKokkosBlockElements()}); + + Policy policy(0, _thread.size()); if (!_jacobian_dispatcher) _jacobian_dispatcher = DispatcherRegistry::build(this, type()); @@ -59,7 +63,15 @@ Kernel::computeJacobian() { auto & sys = kokkosSystem(_kokkos_var.sys()); - _thread.resize({sys.getCoupling(_kokkos_var.var()).size(), numKokkosBlockElements()}); + if (variables().size() > 1) + for (unsigned int i = 1; i < variables().size(); ++i) + if (sys.getCoupling(_kokkos_var.var(0)).size() != + sys.getCoupling(_kokkos_var.var(i)).size()) + mooseError("For multi-variable kernels/BCs, all variables should have the same number of " + "off-diagonal coupling variables."); + + _thread.resize( + {variables().size(), sys.getCoupling(_kokkos_var.var()).size(), numKokkosBlockElements()}); Policy policy(0, _thread.size()); diff --git a/framework/src/kokkos/kernels/KokkosMatCoupledForce.K b/framework/src/kokkos/kernels/KokkosMatCoupledForce.K index 47c0a44a8217..707620a82d09 100644 --- a/framework/src/kokkos/kernels/KokkosMatCoupledForce.K +++ b/framework/src/kokkos/kernels/KokkosMatCoupledForce.K @@ -43,10 +43,11 @@ KokkosMatCoupledForce::KokkosMatCoupledForce(const InputParameters & parameters) { _v_var_to_index[_v_var[j]] = j; - if (_var.number() == _v_var[j]) - paramError("v", - "Coupled variable 'v' needs to be different from 'variable', consider using " - "KokkosReaction or something similar"); + for (const auto var : variables()) + if (var->number() == _v_var[j]) + paramError("v", + "Coupled variable 'v' needs to be different from variable, consider using " + "KokkosReaction or something similar"); } _v_var_to_index.copyToDevice(); diff --git a/framework/src/kokkos/kernels/KokkosTimeKernel.K b/framework/src/kokkos/kernels/KokkosTimeKernel.K index cf2a89a0fe67..8e6ffd8c6ead 100644 --- a/framework/src/kokkos/kernels/KokkosTimeKernel.K +++ b/framework/src/kokkos/kernels/KokkosTimeKernel.K @@ -24,10 +24,20 @@ TimeKernel::validParams() } TimeKernel::TimeKernel(const InputParameters & parameters) - : Kernel(parameters), - _u_dot(_var, Moose::SOLUTION_DOT_TAG), - _du_dot_du(_var.sys().duDotDu(_var.number())) + : Kernel(parameters), _u_dot(variables(), Moose::SOLUTION_DOT_TAG), _du_dot_du(variables().size()) { } +TimeKernel::TimeKernel(const TimeKernel & object) + : Kernel(object), _u_dot(object._u_dot), _du_dot_du(object._du_dot_du) +{ + // TODO: We want to use Array> for _du_dot_du so that it is automatically + // synchronized, but currently there is an NVCC bug preventing it + + for (unsigned int i = 0; i < variables().size(); ++i) + _du_dot_du[i] = variables()[i]->sys().duDotDu(variables()[i]->number()); + + _du_dot_du.copyToDevice(); +} + } // namespace Moose::Kokkos diff --git a/framework/src/kokkos/nodalkernels/KokkosCoupledForceNodalKernel.K b/framework/src/kokkos/nodalkernels/KokkosCoupledForceNodalKernel.K index 000c1b9f0901..96c43f02a0a5 100644 --- a/framework/src/kokkos/nodalkernels/KokkosCoupledForceNodalKernel.K +++ b/framework/src/kokkos/nodalkernels/KokkosCoupledForceNodalKernel.K @@ -29,8 +29,9 @@ KokkosCoupledForceNodalKernel::KokkosCoupledForceNodalKernel(const InputParamete _v(kokkosCoupledNodalValue("v")), _coef(getParam("coef")) { - if (_var.number() == _v_var) - mooseError( - "Coupled variable 'v' needs to be different from 'variable' with CoupledForceNodalKernel, " - "consider using Reaction or somethig similar"); + for (const auto var : variables()) + if (var->number() == _v_var) + mooseError( + "Coupled variable 'v' needs to be different from variable with CoupledForceNodalKernel, " + "consider using Reaction or somethig similar"); } diff --git a/framework/src/kokkos/nodalkernels/KokkosNodalKernel.K b/framework/src/kokkos/nodalkernels/KokkosNodalKernel.K index 53d92add61dd..9a395cd05bde 100644 --- a/framework/src/kokkos/nodalkernels/KokkosNodalKernel.K +++ b/framework/src/kokkos/nodalkernels/KokkosNodalKernel.K @@ -21,16 +21,19 @@ NodalKernel::validParams() NodalKernel::NodalKernel(const InputParameters & parameters) : NodalKernelBase(parameters, Moose::VarFieldType::VAR_FIELD_STANDARD), - _u(_var, Moose::SOLUTION_TAG, true), + _u(_kokkos_var, true), _boundary_restricted(boundaryRestricted()) { - addMooseVariableDependency(&_var); + addMooseVariableDependency(variables()); } void NodalKernel::computeResidual() { - Policy policy(0, _boundary_restricted ? numKokkosBoundaryNodes() : numKokkosBlockNodes()); + _thread.resize({variables().size(), + _boundary_restricted ? numKokkosBoundaryNodes() : numKokkosBlockNodes()}); + + Policy policy(0, _thread.size()); if (!_residual_dispatcher) _residual_dispatcher = DispatcherRegistry::build(this, type()); @@ -43,7 +46,10 @@ NodalKernel::computeJacobian() { if (DispatcherRegistry::hasUserMethod(type())) { - Policy policy(0, _boundary_restricted ? numKokkosBoundaryNodes() : numKokkosBlockNodes()); + _thread.resize({variables().size(), + _boundary_restricted ? numKokkosBoundaryNodes() : numKokkosBlockNodes()}); + + Policy policy(0, _thread.size()); if (!_jacobian_dispatcher) _jacobian_dispatcher = DispatcherRegistry::build(this, type()); @@ -55,7 +61,15 @@ NodalKernel::computeJacobian() { auto & sys = kokkosSystem(_kokkos_var.sys()); - _thread.resize({sys.getCoupling(_kokkos_var.var()).size(), + if (variables().size() > 1) + for (unsigned int i = 1; i < variables().size(); ++i) + if (sys.getCoupling(_kokkos_var.var(0)).size() != + sys.getCoupling(_kokkos_var.var(i)).size()) + mooseError("For multi-variable kernels/BCs, all variables should have the same number of " + "off-diagonal coupling variables."); + + _thread.resize({variables().size(), + sys.getCoupling(_kokkos_var.var()).size(), _boundary_restricted ? numKokkosBoundaryNodes() : numKokkosBlockNodes()}); Policy policy(0, _thread.size()); diff --git a/framework/src/kokkos/nodalkernels/KokkosTimeNodalKernel.K b/framework/src/kokkos/nodalkernels/KokkosTimeNodalKernel.K index 15c9d7be3ec6..c262b8818584 100644 --- a/framework/src/kokkos/nodalkernels/KokkosTimeNodalKernel.K +++ b/framework/src/kokkos/nodalkernels/KokkosTimeNodalKernel.K @@ -25,9 +25,21 @@ TimeNodalKernel::validParams() TimeNodalKernel::TimeNodalKernel(const InputParameters & parameters) : NodalKernel(parameters), - _u_dot(_var, Moose::SOLUTION_DOT_TAG, true), - _du_dot_du(_var.sys().duDotDu(_var.number())) + _u_dot(variables(), Moose::SOLUTION_DOT_TAG, true), + _du_dot_du(variables().size()) { } +TimeNodalKernel::TimeNodalKernel(const TimeNodalKernel & object) + : NodalKernel(object), _u_dot(object._u_dot), _du_dot_du(object._du_dot_du) +{ + // TODO: We want to use Array> for _du_dot_du so that it is automatically + // synchronized, but currently there is an NVCC bug preventing it + + for (unsigned int i = 0; i < variables().size(); ++i) + _du_dot_du[i] = variables()[i]->sys().duDotDu(variables()[i]->number()); + + _du_dot_du.copyToDevice(); +} + } // namespace Moose::Kokkos From f8de5559894f4b4640f8200afbf910f438fae87b Mon Sep 17 00:00:00 2001 From: Namjae Choi Date: Mon, 16 Feb 2026 18:44:24 -0700 Subject: [PATCH 2/3] Update documentation #30655 --- .../doc/content/syntax/KokkosKernels/index.md | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/framework/doc/content/syntax/KokkosKernels/index.md b/framework/doc/content/syntax/KokkosKernels/index.md index 38f12f857f95..fc4226077f8e 100644 --- a/framework/doc/content/syntax/KokkosKernels/index.md +++ b/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: From 3ecda2338bb4b9429ee267edb7b04a78dab7fcc2 Mon Sep 17 00:00:00 2001 From: Namjae Choi Date: Wed, 18 Feb 2026 09:36:53 -0700 Subject: [PATCH 3/3] Add tests #30655 --- .../kokkos_multi_variable_kernel_test.i | 92 +++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 test/tests/kokkos/kernels/multi_variable/kokkos_multi_variable_kernel_test.i diff --git a/test/tests/kokkos/kernels/multi_variable/kokkos_multi_variable_kernel_test.i b/test/tests/kokkos/kernels/multi_variable/kokkos_multi_variable_kernel_test.i new file mode 100644 index 000000000000..516d581a03bb --- /dev/null +++ b/test/tests/kokkos/kernels/multi_variable/kokkos_multi_variable_kernel_test.i @@ -0,0 +1,92 @@ +[Mesh] + [square] + type = GeneratedMeshGenerator + nx = 2 + ny = 2 + dim = 2 + [] +[] + +[Variables] + [u] + order = FIRST + family = LAGRANGE + [] + [v] + order = FIRST + family = LAGRANGE + [] + [w] + order = FIRST + family = LAGRANGE + [] +[] + +[Kernels] + [diff] + type = KokkosDiffusion + variables = 'u v w' + [] + [bf] + type = KokkosBodyForce + variables = 'u v' + postprocessor = ramp + [] +[] + +[Functions] + [ramp] + type = ParsedFunction + expression = 't' + [] +[] + +[Postprocessors] + [ramp] + type = FunctionValuePostprocessor + function = ramp + execute_on = linear + [] +[] + +[BCs] + [left] + type = KokkosDirichletBC + variables = 'v w' + boundary = 3 + value = 0 + [] + + [right] + type = KokkosDirichletBC + variables = 'v w' + boundary = 1 + value = 0 + [] + + [top] + type = KokkosDirichletBC + variable = 'u' + boundary = 2 + value = 0 + [] + + [bottom] + type = KokkosDirichletBC + variable = 'u' + boundary = 0 + value = 0 + [] +[] + +[Executioner] + type = Transient + dt = 1.0 + end_time = 1.0 + + solve_type = 'NEWTON' +[] + +[Outputs] + exodus = true +[]