Kokkos Thermodynamics

core/src/CMakeLists.txt

@@ -30,6 +30,7 @@ set(BaseSources
     "StructureFactory.cpp"
     "ModelArray.cpp"
     "ModelArraySlice.cpp"
+    "ModelArrayStore.cpp"
     "ModelComponent.cpp"
     "ModelMetadata.cpp"
     "NetcdfMetadataConfiguration.cpp"
@@ -61,6 +62,10 @@ if(ENABLE_MPI)
     list(APPEND NextsimSourcesTemp "${ParallelNetCDFSources}")
 endif()
 
+if(WITH_KOKKOS)
+    list(APPEND NextsimSourcesTemp "${CMAKE_CURRENT_SOURCE_DIR}/kokkos/KokkosUtils.cpp")
+endif()
+
 set(NextsimSources "${NextsimSourcesTemp}" PARENT_SCOPE)
 
 set(NextsimIncludeDirs

core/src/FieldAdvection.cpp

@@ -20,4 +20,13 @@ ModelArray& FieldAdvection::advectField(
     return field;
 }
 
+#ifdef USE_KOKKOS
+void FieldAdvection::advectField(
+    const DeviceViewMA& field, const TimestepTime& tst, double lowerLimit, double upperLimit)
+{
+    Module::getImplementation<IDynamics>().advectField(
+        tst.step.seconds(), field, lowerLimit, upperLimit);
+}
+#endif
+
 } /* namespace Nextsim */

core/src/Model.cpp

@@ -169,6 +169,8 @@ Model::HelpMap& Model::getHelpRecursive(HelpMap& map, bool getAll)
 
 void Model::run()
 {
+    assert(Nextsim::ModelComponent::getStore().checkAllRegistered());
+
     try {
         iterator.run();
     } catch (const std::exception& e) {

core/src/ModelArrayStore.cpp

@@ -0,0 +1,103 @@
+/*!
+ * @author  Robert Jendersie <[email protected]>
+ */
+
+#include "include/ModelArrayStore.hpp"
+#include "include/Logged.hpp"
+#include <iostream>
+namespace Nextsim {
+
+#ifdef USE_KOKKOS
+
+//double ConstDeviceModelArray::operator[](DeviceIndex i) const { return m_deviceView(i, 0); }
+ConstDeviceModelArray::operator ConstDeviceViewMA() const { return m_deviceView; }
+
+//double& DeviceModelArray::operator[](DeviceIndex i) const { return m_deviceView(i, 0); }
+DeviceModelArray::operator const DeviceViewMA&() const { return m_deviceView; }
+
+/*************************************************************/
+HostViewMA ModelArrayStore::ExtModelArray::hostView()
+{
+    assert(modelArray.trueSize() > 0 && "ModelArray is allocated");
+    return makeKokkosHostView(modelArray.getDataRef());
+}
+
+const DeviceViewMA& ModelArrayStore::ExtModelArray::deviceView()
+{
+    assert(modelArray.trueSize() > 0 && "ModelArray is allocated");
+    if (!m_deviceModelArray.m_deviceView.is_allocated()) {
+        m_deviceModelArray.m_deviceView = makeKokkosDeviceView(name, modelArray.getDataRef());
+    }
+
+    return m_deviceModelArray.m_deviceView;
+}
+
+DeviceModelArray& ModelArrayStore::ExtModelArray::deviceModelArray()
+{
+    // handles initialization
+    deviceView();
+    return m_deviceModelArray;
+}
+#endif
+
+/*************************************************************/
+ModelArrayStore::ExtModelArrayFlagged::ExtModelArrayFlagged(
+    const std::string& name, bool _isReadWrite, bool _isRegistered)
+    : isReadWrite(_isReadWrite)
+    , isRegistered(_isRegistered)
+{
+    extModelArray.name = name;
+}
+
+/*************************************************************/
+std::unordered_map<std::string, const ModelArray*> ModelArrayStore::getAllData() const
+{
+    std::unordered_map<std::string, const ModelArray*> dataMap;
+
+    for (const auto& [name, extArrFlagged] : store) {
+        dataMap.emplace(name, &extArrFlagged.extModelArray.modelArray);
+    }
+
+    return dataMap;
+}
+
+bool ModelArrayStore::checkAllRegistered() const
+{
+    bool b = true;
+    for (const auto& [_, extArrFlagged] : store) {
+        if (!extArrFlagged.isRegistered) {
+            Logged::warning("Field \"" + extArrFlagged.extModelArray.name
+                + "\" was not properly initialized by a registering ModelArrayAccessor.\n");
+            b = false;
+        }
+    }
+
+    return b;
+}
+
+ModelArrayStore::ExtModelArray& ModelArrayStore::getRW(const std::string& field)
+{
+    if (auto it = store.find(field); it != store.end()) {
+        ExtModelArrayFlagged& extArrFlagged = it->second;
+        if (!extArrFlagged.isReadWrite) {
+            if (extArrFlagged.isRegistered) {
+                throw std::logic_error(
+                    "Trying to access the read-only ModelArray \"" + field + "\" as read-write.");
+            } else {
+                // promote to read-write because we don't now the true restriction yet
+                extArrFlagged.isReadWrite = true;
+            }
+        }
+        return extArrFlagged.extModelArray;
+    }
+
+    // Regular emplace would be fine here since we know that it does not exist but try_emplace
+    // has a more ergonomic signature.
+    return store.try_emplace(field, field, RW, false).first->second.extModelArray;
+}
+
+ModelArrayStore::ExtModelArray& ModelArrayStore::getRO(const std::string& field)
+{
+    return store.try_emplace(field, field, RO, false).first->second.extModelArray;
+}
+}

core/src/ModelComponent.cpp

@@ -8,6 +8,11 @@ namespace Nextsim {
 
 size_t ModelComponent::nOcean = 0;
 std::vector<size_t> ModelComponent::oceanIndex;
+bool ModelComponent::columnPhysicsStoreIsDestroyed; // initialized to 0 because of static lifetime
+
+#if USE_KOKKOS
+KokkosDeviceMapView<DeviceIndex> ModelComponent::oceanIndexDevice;
+#endif
 
 ModelComponent::ModelComponent()
 {
@@ -38,6 +43,10 @@ void ModelComponent::setOceanMask(const ModelArray& mask)
             oceanIndex[iOceanIndex++] = i;
         }
     }
+
+#if USE_KOKKOS
+    makeOceanIndexDevice();
+#endif
 }
 
 // Fills the nOcean and OceanIndex variables for the zero land case
@@ -53,6 +62,10 @@ void ModelComponent::noLandMask()
     for (size_t i = 0; i < ModelArray::size(ModelArray::Type::H); ++i) {
         oceanIndex[i] = i;
     }
+
+#if USE_KOKKOS
+    makeOceanIndexDevice();
+#endif
 }
 
 ModelArray ModelComponent::mask(const ModelArray& data, const double missingValue)
@@ -67,4 +80,20 @@ ModelArray ModelComponent::mask(const ModelArray& data, const double missingValu
 
 const ModelArray& ModelComponent::oceanMask() { return oceanMaskSingleton(); }
 
+#if USE_KOKKOS
+void ModelComponent::makeOceanIndexDevice()
+{
+    // some tests don't need Kokkos
+    if (!Kokkos::is_initialized()) {
+        return;
+    }
+
+    std::vector<DeviceIndex> buf(oceanIndex.begin(), oceanIndex.end());
+    oceanIndexDevice = makeKokkosDeviceViewMap("oceanIndex", buf, MakeViewOptions::ALWAYS_COPY);
+    Finalizer::registerUnique(destroyOceanIndex);
+}
+
+void ModelComponent::destroyOceanIndex() { oceanIndexDevice.assign_data(nullptr); }
+#endif
+
 } /* namespace Nextsim */

core/src/PrognosticData.cpp

@@ -7,7 +7,6 @@
 
 #include "include/FieldAdvection.hpp"
 #include "include/Finalizer.hpp"
-#include "include/ModelArrayRef.hpp"
 #include "include/NextsimModule.hpp"
 #include "include/gridNames.hpp"
 
@@ -23,19 +22,15 @@ static constexpr bool checkFieldsFastDefault = true;
 
 PrognosticData::PrognosticData()
     : m_dt(1)
-    , hice(ModelArray::AdvectionType, { 0, 50 })
-    , cice(ModelArray::AdvectionType, { 0, 1 })
-    , damage(ModelArray::AdvectionType, { 0, 1 })
-    , hsnow(ModelArray::AdvectionType, { 0, 10 })
+    , hiceAccessor(getStore(), RW, ModelArray::AdvectionType, std::pair(0.0, 50.0))
+    , ciceAccessor(getStore(), RW, ModelArray::AdvectionType, std::pair(0.0, 1.0))
+    , damageAccessor(getStore(), RW, ModelArray::AdvectionType, std::pair(0.0, 1.0))
+    , hsnowAccessor(getStore(), RW, ModelArray::AdvectionType, std::pair(0.0, 10.0))
     , pAtmBdy(nullptr)
     , pOcnBdy(nullptr)
     , pDynamics(nullptr)
 {
-    getStore().registerArray(Shared::DAMAGE, &damage, RW);
-    getStore().registerArray(Shared::H_ICE_DG, &hice, RW);
-    getStore().registerArray(Shared::C_ICE_DG, &cice, RW);
-    getStore().registerArray(Shared::H_SNOW_DG, &hsnow, RW);
-}
+                }
 
 void PrognosticData::configure()
 {
@@ -64,8 +59,8 @@ void PrognosticData::configure()
         }
     } else if (checkFast) {
         addChecks({
-            { "thickness", &hice },
-            { "concentration", &cice },
+            { "thickness", &hiceAccessor.getHostRO() },
+            { "concentration", &ciceAccessor.getHostRO() },
         });
     }
 }
@@ -82,7 +77,11 @@ void copyMeanComponent(const ModelArray& source, ModelArray& sink)
 
 void PrognosticData::setData(const ModelState::DataMap& ms)
 {
-
+    AdvectedField& hice = hiceAccessor.getHostRW();
+    AdvectedField& cice = ciceAccessor.getHostRW();
+    AdvectedField& hsnow = hsnowAccessor.getHostRW();
+    AdvectedField& damage = damageAccessor.getHostRW();
+
     if (ms.count(maskName)) {
         setOceanMask(ms.at(maskName));
     } else {
@@ -154,9 +153,9 @@ ModelState PrognosticData::getStatePrognostic() const
 {
     ModelState state = { {
                              { maskName, ModelArray(oceanMask()) }, // make a copy
-                             { hiceName, hice },
-                             { ciceName, cice },
-                             { hsnowName, hsnow },
+                             { hiceName, hiceAccessor.getHostRO() },
+                             { ciceName, ciceAccessor.getHostRO() },
+                             { hsnowName, hsnowAccessor.getHostRO() },
                          },
         ModelComponent::getConfiguration() };
 

core/src/discontinuousgalerkin/include/ModelArrayDetails.hpp

@@ -25,7 +25,7 @@ enum class Type {
 
 enum class Base { Cell, Vertex };
 
-static const Type AdvectionType = Type::DG;
+static constexpr Type AdvectionType = Type::DG;
 
 static ModelArray HField() { return ModelArray(Type::H); }
 static ModelArray VertexField() { return ModelArray(Type::VERTEX); }

core/src/include/FieldAdvection.hpp

@@ -11,6 +11,10 @@
 #include "include/ModelArray.hpp"
 #include "include/Time.hpp"
 
+#ifdef USE_KOKKOS
+#include "include/ModelArrayStore.hpp" // needed for DeviceViewMA
+#endif
+
 namespace Nextsim {
 
 class FieldAdvection {
@@ -26,6 +30,18 @@ class FieldAdvection {
     static ModelArray& advectField(ModelArray& field, const TimestepTime& tst,
         double lowerLimit = -std::numeric_limits<double>::infinity(),
         double upperLimit = std::numeric_limits<double>::infinity());
+
+#ifdef USE_KOKKOS
+    static void advectField(const DeviceViewMA& field, const TimestepTime& tst,
+        double lowerLimit = -std::numeric_limits<double>::infinity(),
+        double upperLimit = std::numeric_limits<double>::infinity());
+
+/*    template<typename ExecSpace>
+    static void advectField(ExecSpace execSpace, const DeviceViewMA& field, const TimestepTime& tst,
+        double lowerLimit = -std::numeric_limits<double>::infinity(),
+        double upperLimit = std::numeric_limits<double>::infinity())
+    {}*/
+#endif
 };
 
 } /* namespace Nextsim */

core/src/include/KernelAlternatives.hpp

@@ -0,0 +1,51 @@
+/*!
+ * @author Robert Jendersie <[email protected]>
+ */
+
+#ifndef KERNEL_ALTERNATIVES_HPP
+#define KERNEL_ALTERNATIVES_HPP
+
+#ifdef USE_KOKKOS
+#include <Kokkos_Core.hpp>
+#include "../kokkos/include/KokkosUtils.hpp"
+#else
+#include <cmath>
+#include <algorithm>
+#endif
+
+namespace Nextsim {
+
+// namespace Utils contains standard math functions (sin, sqrt, min, ...)
+#ifdef USE_KOKKOS
+namespace Utils = Kokkos;
+#else
+namespace Utils = std;
+#endif
+
+// KERNEL_IMPL_FUNCTION annotates functions that need to be callable inside a kernel
+#ifdef USE_KOKKOS
+#define KERNEL_IMPL_FUNCTION KOKKOS_IMPL_FUNCTION
+#else
+// no special annotation needed
+#define KERNEL_IMPL_FUNCTION
+#endif
+
+// KERNEL_LAMBDA capture list of a lambda used in a overElements call
+#ifdef USE_KOKKOS
+#define OVER_ELEMENTS_LAMBDA KOKKOS_LAMBDA
+using ElementIndex = DeviceIndex;
+#else
+#define OVER_ELEMENTS_LAMBDA [&]
+using ElementIndex = size_t;
+#endif
+
+// Execution space that is used to signal async operations.
+#ifdef USE_KOKKOS
+using DefaultExecutionSpace = Kokkos::DefaultExecutionSpace;
+#else
+struct ExecutionSpaceDummy {};
+using DefaultExecutionSpace = ExecutionSpaceDummy;
+#endif
+}
+
+#endif /* KERNEL_ALTERNATIVES_HPP */