Kokkos Function System

framework/app.mk

@@ -377,8 +377,14 @@ app_KOKKOS_DEPS      := $(patsubst %.$(KOKKOS_OBJ_SUFFIX), %.$(KOKKOS_OBJ_SUFFIX
 app_KOKKOS_LIB       :=
 
 ifneq ($(app_KOKKOS_OBJECTS),)
-  app_KOKKOS_LIB     := $(APPLICATION_DIR)/lib/lib$(APPLICATION_NAME)$(KOKKOS_LIB_SUFFIX)
-  app_KOKKOS_LIBS    += $(app_KOKKOS_LIB)
+  app_KOKKOS_LIB  := $(APPLICATION_DIR)/lib/lib$(APPLICATION_NAME)$(KOKKOS_LIB_SUFFIX)
+  app_KOKKOS_LIBS += $(app_KOKKOS_LIB)
+endif
+
+ifeq ($(KOKKOS_COMPILER),CPU)
+  app_KOKKOS_LIB_COMBINED := $(MOOSE_KOKKOS_LIB) $(app_KOKKOS_LIBS)
+else
+  app_KOKKOS_LIB_COMBINED := $(APPLICATION_DIR)/lib/lib$(APPLICATION_NAME)_combined$(KOKKOS_LIB_SUFFIX)
 endif
 
 KOKKOS_OBJECTS += $(app_KOKKOS_OBJECTS)
@@ -407,6 +413,7 @@ $(app_KOKKOS_LIB): $(app_KOKKOS_OBJECTS)
 else
 
 # libtool ignores nvcc and just uses mpicxx to link, so cannot be used
+
 $(app_KOKKOS_LIB): curr_dir  := $(APPLICATION_DIR)
 $(app_KOKKOS_LIB): curr_objs := $(app_KOKKOS_OBJECTS)
 $(app_KOKKOS_LIB): $(app_KOKKOS_OBJECTS)
@@ -418,6 +425,21 @@ endif
 
 endif
 
+ifeq ($(KOKKOS_COMPILER),GPU)
+
+# Making a dummy object file for triggering device link is only required for NVCC
+
+$(app_KOKKOS_LIB_COMBINED): curr_dir := $(APPLICATION_DIR)
+$(app_KOKKOS_LIB_COMBINED): $(MOOSE_KOKKOS_LIB) $(app_KOKKOS_LIBS)
+	@mkdir -p $(curr_dir)/lib
+	@echo "Device Linking Kokkos Libraries "$@"..."
+	@echo > dlink.K
+	@$(KOKKOS_CXX) $(KOKKOS_CXXFLAGS) -c dlink.K -o dlink.o
+	@$(KOKKOS_CXX) --shared -o $@ dlink.o $(KOKKOS_LDFLAGS) $(KOKKOS_LIBS) $(MOOSE_KOKKOS_LIB) $(app_KOKKOS_LIBS)
+	@rm dlink.K dlink.o
+
+endif
+
 endif
 
 # Target-specific Variable Values (See GNU-make manual)
@@ -516,10 +538,10 @@ ifneq (,$(findstring darwin,$(libmesh_HOST)))
   endif
 endif
 
-$(app_EXEC): $(app_LIBS) $(mesh_library) $(main_object) $(app_test_LIB) $(depend_test_libs) $(app_resource) $(MOOSE_KOKKOS_LIB) $(app_KOKKOS_LIBS)
+$(app_EXEC): $(app_LIBS) $(mesh_library) $(main_object) $(app_test_LIB) $(depend_test_libs) $(app_resource) $(app_KOKKOS_LIB_COMBINED)
 	@echo "Linking Executable "$@"..."
 	@bash -c '$(libmesh_LIBTOOL) --tag=CXX $(LIBTOOLFLAGS) --mode=link --quiet \
-	  $(libmesh_CXX) $(libmesh_CXXFLAGS) -o $@ $(main_object) $(depend_test_libs_flags) $(applibs) $(ADDITIONAL_LIBS) $(LDFLAGS) $(libmesh_LDFLAGS) $(libmesh_LIBS) $(EXTERNAL_FLAGS) $(MOOSE_KOKKOS_LIB) $(app_KOKKOS_LIBS) ${SILENCE_SOME_WARNINGS}'
+	  $(libmesh_CXX) $(libmesh_CXXFLAGS) -o $@ $(main_object) $(depend_test_libs_flags) $(applibs) $(ADDITIONAL_LIBS) $(LDFLAGS) $(libmesh_LDFLAGS) $(libmesh_LIBS) $(EXTERNAL_FLAGS) $(app_KOKKOS_LIB_COMBINED) ${SILENCE_SOME_WARNINGS}'
 	@$(codesign)
 
 ###### install stuff #############

framework/doc/content/source/kokkos/actions/AddKokkosFunctionAction.md

@@ -0,0 +1,16 @@
+# AddKokkosFunctionAction
+
+!if! function=hasCapability('kokkos')
+
+!syntax description /KokkosFunctions/AddKokkosFunctionAction
+
+Kokkos `Functions` are specified as an object inside the `[KokkosFunctions]` block. This action adds them to the [Problem](syntax/Problem/index.md).
+
+More information about Kokkos `Functions` can be found on the [Kokkos Functions syntax documentation](syntax/KokkosFunctions/index.md).
+
+!syntax parameters /KokkosFunctions/AddKokkosFunctionAction
+
+!if-end!
+
+!else
+!include kokkos/kokkos_warning.md

framework/doc/content/source/kokkos/actions/AddKokkosKernelAction.md

@@ -1,4 +1,4 @@
-# AddKokkosKokkosKernelAction
+# AddKokkosKernelAction
 
 !if! function=hasCapability('kokkos')
 

framework/doc/content/source/kokkos/auxkernels/KokkosFunctionAux.md

@@ -0,0 +1,18 @@
+# KokkosFunctionAux
+
+!if! function=hasCapability('kokkos')
+
+This is the Kokkos version of [FunctionAux](FunctionAux.md). See the original document for details.
+
+## Example Syntax
+
+!listing test/tests/kokkos/functions/piecewise_constant/kokkos_piecewise_constant.i block=KokkosAuxKernels
+
+!syntax parameters /KokkosAuxKernels/KokkosFunctionAux
+
+!syntax inputs /KokkosAuxKernels/KokkosFunctionAux
+
+!if-end!
+
+!else
+!include kokkos/kokkos_warning.md

framework/doc/content/source/kokkos/functions/KokkosConstantFunction.md

@@ -0,0 +1,18 @@
+# KokkosConstantFunction
+
+!if! function=hasCapability('kokkos')
+
+This is the Kokkos version of [ConstantFunction](ConstantFunction.md). See the original document for details.
+
+## Example Syntax
+
+!listing test/tests/kokkos/functions/constant_function/kokkos_constant_function.i block=KokkosFunctions
+
+!syntax parameters /KokkosFunctions/KokkosConstantFunction
+
+!syntax inputs /KokkosFunctions/KokkosConstantFunction
+
+!if-end!
+
+!else
+!include kokkos/kokkos_warning.md

framework/doc/content/source/kokkos/functions/KokkosPiecewiseConstant.md

@@ -0,0 +1,18 @@
+# KokkosPiecewiseConstant
+
+!if! function=hasCapability('kokkos')
+
+This is the Kokkos version of [PiecewiseConstant](PiecewiseConstant.md). See the original document for details.
+
+## Example Syntax
+
+!listing test/tests/kokkos/functions/piecewise_constant/kokkos_piecewise_constant.i block=KokkosFunctions
+
+!syntax parameters /KokkosFunctions/KokkosPiecewiseConstant
+
+!syntax inputs /KokkosFunctions/KokkosPiecewiseConstant
+
+!if-end!
+
+!else
+!include kokkos/kokkos_warning.md

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

@@ -25,7 +25,7 @@ Moose `Function`s should override the following member functions:
 and may optionally override the following member functions, which is only needed
 for some particular functionality:
 
-- `Real timeIntegral(Real t1, Real t1, const Point & p)`, which computes the
+- `Real timeIntegral(Real t1, Real t2, const Point & p)`, which computes the
   time integral of the function at the spatial point `p` between the time values
   `t1` and `t2`.
 

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

@@ -433,6 +433,7 @@ The following objects are currently available in Kokkos-MOOSE:
 - [BCs](syntax/KokkosBCs/index.md)
 - [Materials](syntax/KokkosMaterials/index.md)
 - [AuxKernels](syntax/KokkosAuxKernels/index.md)
+- [Functions](syntax/KokkosFunctions/index.md)
 
 !if-end!
 

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

@@ -0,0 +1,68 @@
+# Kokkos Functions System
+
+!if! function=hasCapability('kokkos')
+
+Before reading this documentation, consider reading the following materials first for a better understanding of this documentation:
+
+- [Functions System](syntax/Functions/index.md) to understand the MOOSE function system,
+- [Getting Started with Kokkos-MOOSE](syntax/Kokkos/index.md) to understand the programming practices for Kokkos-MOOSE,
+- [Kokkos Kernels System](syntax/KokkosKernels/index.md) to understand the common design pattern of objects in Kokkos-MOOSE.
+
+!alert note
+Kokkos-MOOSE functions do not support automatic differention yet.
+
+A Kokkos-MOOSE function can be created by subclassing `Moose::Kokkos::FunctionBase` (not `Moose::Kokkos::Function`), which should be registered with `registerKokkosFunction()` instead of `registerMooseObject()`.
+The signatures of the hook methods are defined as follows:
+
+```cpp
+KOKKOS_FUNCTION Real value(Real t, Real3 p) const;
+KOKKOS_FUNCTION Real3 vectorValue(Real t, Real3 p) const;
+KOKKOS_FUNCTION Real3 gradient(Real t, Real3 p) const;
+KOKKOS_FUNCTION Real3 curl(Real t, Real3 p) const;
+KOKKOS_FUNCTION Real div(Real t, Real3 p) const;
+KOKKOS_FUNCTION Real timeDerivative(Real t, Real3 p) const;
+KOKKOS_FUNCTION Real timeIntegral(Real t1, Real t2, Real3 p) const;
+KOKKOS_FUNCTION Real integral() const;
+KOKKOS_FUNCTION Real average() const;
+```
+
+As in other Kokkos-MOOSE objects, they should be defined in the derived class as +*inlined public*+ methods instead of virtual override.
+It is not mandatory to define each hook method in the derived class, but any hook method that was not defined in the derived class should not be called.
+
+See the following source codes of `KokkosPiecewiseConstant` for an example of a function:
+
+!listing framework/include/kokkos/functions/KokkosPiecewiseConstant.h id=kokkos-piecewise-constant-header
+         caption=The `KokkosPiecewiseConstant` header file.
+
+!listing framework/src/kokkos/functions/KokkosPiecewiseConstant.K id=kokkos-piecewise-constant-source language=cpp
+         caption=The `KokkosPiecewiseConstant` source file.
+
+Functions can be acquired in your object by calling `getKokkosFunction<T>()`, where `T` should be your function type.
+Namely, the actual type of the function should be known in advance.
+This is the limitation of the current Kokkos-MOOSE function implementation which is being addressed, and more discussions can be found below.
+It is recommended to store the acquired function in a `ReferenceWrapper` instance.
+Otherwise, it should be stored as a copy, and it becomes your responsibility to maintain synchronization between the original function and the copy.
+
+## Use of Dynamic Polymorphism
+
+Unlike most of other Kokkos-MOOSE objects, functions are pluggable objects that are to be used in other objects.
+Therefore, it is desired to use functions type-agnostically in your object so that any type of function can be plugged in.
+Such design requires dynamic polymorphism, which is represented by virtual functions.
+Although using virtual functions on GPU can sacrifice performance due to the inability of code inlining and vtable lookup and is still advised to avoid if possible, it can be useful for avoiding code duplications and improving code maintainability.
+
+In fact, Kokkos-MOOSE already has the code path for acquiring functions as a type-agnostic abstract type and using virtual dispatch to evaluate actual functions.
+There is another non-template API `getKokkosFunction()` which returns an abstract type `Moose::Kokkos::Function`.
+It is not the base class of your function but a wrapper class that contains your function and provides shims to call the function methods.
+
+However, that code path is currently blocked for GPU backends, and you will hit a runtime error when you attempt to use it.
+While our target GPU backends (CUDA, HIP, and Intel SYCL) support virtual functions on GPU, it requires the relocatable device code (RDC) option to be enabled during compilation, which the current MOOSE build system is lacking.
+Kokkos imposes a restriction to have a consistent RDC configuration across the whole software stack, but PETSc, from which MOOSE is currently acquring Kokkos, requires significant changes in its build system to enable RDC.
+The non-templated `getKokkosFunction()` code path is hypothetically functional under the assumption that the MOOSE software stack can be built with the RDC option, but this hypothetical will only become a reality after a PETSc build system rework.
+The PETSc and MOOSE team are looking into this issue, and it is expected to be resolved in the foreseeable future.
+
+!syntax list /KokkosFunctions objects=True actions=False subsystems=False
+
+!if-end!
+
+!else
+!include kokkos/kokkos_warning.md

framework/doc/unregister.yml

@@ -6,6 +6,7 @@ KokkosAuxKernel: AuxKernels/*
 KokkosNodalKernel: NodalKernels/*
 KokkosBoundaryCondition: BCs/*
 KokkosMaterial: Materials/*
+KokkosFunction: Functions/*
 Kernel: KokkosKernels/*
 VectorKernel: KokkosKernels/*
 ArrayKernel: KokkosKernels/*
@@ -16,4 +17,5 @@ ArrayAuxKernel: KokkosAuxKernels/*
 NodalKernel: KokkosNodalKernels/*
 BoundaryCondition: KokkosBCs/*
 MaterialBase: KokkosMaterials/*
-FunctorMaterial: KokkosMaterials/*
+FunctorMaterial: KokkosMaterials/*
+Function: KokkosFunctions/*

framework/include/base/MooseObject.h

@@ -56,26 +56,13 @@ class MooseObject : public ParallelParamObject,
   std::shared_ptr<const MooseObject> getSharedPtr() const;
 
 #ifdef MOOSE_KOKKOS_ENABLED
-  class IsKokkosObjectKey
-  {
-    friend class BlockRestrictable;
-    friend class BoundaryRestrictable;
-    friend class MaterialPropertyInterface;
-    IsKokkosObjectKey() = default;
-    IsKokkosObjectKey(const IsKokkosObjectKey &) = delete;
-    IsKokkosObjectKey(IsKokkosObjectKey &&) = delete;
-  };
-
   /**
    * Get whether this object is a Kokkos functor
    * The parameter is set by the Kokkos base classes:
    * - Moose::Kokkos::ResidualObject in KokkosResidualObject.K
    * - Moose::Kokkos::MaterialBase in KokkosMaterialBase.K
    */
-  bool isKokkosObject(IsKokkosObjectKey &&) const
-  {
-    return parameters().isParamValid(MooseBase::kokkos_object_param);
-  }
+  bool isKokkosObject() const { return parameters().isParamValid(MooseBase::kokkos_object_param); }
 #endif
 
   // To get warnings tracked in the SolutionInvalidityOutput

framework/include/functions/Function.h

@@ -9,14 +9,7 @@
 
 #pragma once
 
-#include "MooseObject.h"
-#include "SetupInterface.h"
-#include "TransientInterface.h"
-#include "PostprocessorInterface.h"
-#include "UserObjectInterface.h"
-#include "Restartable.h"
-#include "MeshChangedInterface.h"
-#include "ScalarCoupleable.h"
+#include "FunctionBase.h"
 #include "MooseFunctor.h"
 #include "ChainedReal.h"
 
@@ -33,15 +26,7 @@ class Point;
  * Base class for function objects.  Functions override value to supply a
  * value at a point.
  */
-class Function : public MooseObject,
-                 public SetupInterface,
-                 public TransientInterface,
-                 public PostprocessorInterface,
-                 public UserObjectInterface,
-                 public Restartable,
-                 public MeshChangedInterface,
-                 public ScalarCoupleable,
-                 public Moose::FunctorBase<Real>
+class Function : public Moose::FunctionBase, public Moose::FunctorBase<Real>
 {
 public:
   /**

framework/include/functions/FunctionBase.h

@@ -0,0 +1,48 @@
+//* This file is part of the MOOSE framework
+//* https://mooseframework.inl.gov
+//*
+//* All rights reserved, see COPYRIGHT for full restrictions
+//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
+//*
+//* Licensed under LGPL 2.1, please see LICENSE for details
+//* https://www.gnu.org/licenses/lgpl-2.1.html
+
+#pragma once
+
+#include "MooseObject.h"
+#include "SetupInterface.h"
+#include "TransientInterface.h"
+#include "PostprocessorInterface.h"
+#include "UserObjectInterface.h"
+#include "Restartable.h"
+#include "MeshChangedInterface.h"
+#include "ScalarCoupleable.h"
+
+namespace Moose
+{
+
+class FunctionBase : public MooseObject,
+                     public SetupInterface,
+                     public TransientInterface,
+                     public PostprocessorInterface,
+                     public UserObjectInterface,
+                     public Restartable,
+                     public MeshChangedInterface,
+                     public ScalarCoupleable
+{
+public:
+  static InputParameters validParams();
+
+  FunctionBase(const InputParameters & parameters);
+
+#ifdef MOOSE_KOKKOS_ENABLED
+  /**
+   * Special constructor used for Kokkos functor copy during parallel dispatch
+   */
+  FunctionBase(const FunctionBase & object, const Moose::Kokkos::FunctorCopy & key);
+#endif
+
+  virtual ~FunctionBase();
+};
+
+}