sundials.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2023 NeoN authors
 
#pragma once
 
#if NN_WITH_SUNDIALS
 
#include <concepts>
#include <functional>
#include <memory>
 
#include <sundials/sundials_nvector.h>
#include <sundials/sundials_core.hpp>
#include <nvector/nvector_serial.h>
#include <nvector/nvector_kokkos.hpp>
#include <arkode/arkode_arkstep.h>
#include <arkode/arkode_erkstep.h>
 
#include "NeoN/core/error.hpp"
#include "NeoN/core/parallelAlgorithms.hpp"
#include "NeoN/fields/field.hpp"
 
namespace NeoN::sundials
{
 
inline auto SUN_CONTEXT_DELETER = [](SUNContext* ctx)
{
    if (ctx != nullptr)
    {
        SUNContext_Free(ctx);
    }
};
 
inline auto SUN_ARK_DELETER = [](char* ark)
{
    if (ark != nullptr)
    {
        void* arkodMem = reinterpret_cast<void*>(ark);
        ARKodeFree(&arkodMem);
    }
};
 
inline ARKODE_ERKTableID stringToERKTable(const std::string& key)
{
    if (key == "Forward-Euler") return ARKODE_FORWARD_EULER_1_1;
    if (key == "Heun")
    {
        NF_ERROR_EXIT("Currently unsupported until field time step-stage indexing resolved.");
        return ARKODE_HEUN_EULER_2_1_2;
    }
    if (key == "Midpoint")
    {
        NF_ERROR_EXIT("Currently unsupported until field time step-stage indexing resolved.");
        return ARKODE_EXPLICIT_MIDPOINT_EULER_2_1_2;
    }
    NF_ERROR_EXIT(
        "Unsupported Runge-Kutta time integration method selectied: " + key + ".\n"
        + "Supported methods are: Forward-Euler, Heun, Midpoint."
    );
    return ARKODE_ERK_NONE; // avoids compiler warnings.
}
 
template<typename SKVectorType, typename ValueType>
void fieldToSunNVectorImpl(const NeoN::Vector<ValueType>& field, N_Vector& vector)
{
    auto view = ::sundials::kokkos::GetVec<SKVectorType>(vector)->View();
    auto fieldView = field.view();
    NeoN::parallelFor(
        field.exec(), field.range(), KOKKOS_LAMBDA(const localIdx i) { view(i) = fieldView[i]; }
    );
};
 
template<typename ValueType>
void fieldToSunNVector(const NeoN::Vector<ValueType>& field, N_Vector& vector)
{
    // CHECK FOR N_Vector on correct space in DEBUG
    if (std::holds_alternative<NeoN::GPUExecutor>(field.exec()))
    {
        fieldToSunNVectorImpl<::sundials::kokkos::Vector<Kokkos::DefaultExecutionSpace>>(
            field, vector
        );
        return;
    }
    if (std::holds_alternative<NeoN::CPUExecutor>(field.exec()))
    {
        fieldToSunNVectorImpl<::sundials::kokkos::Vector<Kokkos::DefaultHostExecutionSpace>>(
            field, vector
        );
        return;
    }
    if (std::holds_alternative<NeoN::SerialExecutor>(field.exec()))
    {
        fieldToSunNVectorImpl<::sundials::kokkos::Vector<Kokkos::Serial>>(field, vector);
        return;
    }
    NF_ERROR_EXIT("Unsupported NeoN executor for field.");
};
 
template<typename SKVectorType, typename ValueType>
void sunNVectorToVectorImpl(const N_Vector& vector, NeoN::Vector<ValueType>& field)
{
    auto view = ::sundials::kokkos::GetVec<SKVectorType>(vector)->View();
    ValueType* fieldData = field.data();
    NeoN::parallelFor(
        field.exec(), field.range(), KOKKOS_LAMBDA(const localIdx i) { fieldData[i] = view(i); }
    );
};
 
template<typename ValueType>
void sunNVectorToVector(const N_Vector& vector, NeoN::Vector<ValueType>& field)
{
    if (std::holds_alternative<NeoN::GPUExecutor>(field.exec()))
    {
        sunNVectorToVectorImpl<::sundials::kokkos::Vector<Kokkos::DefaultExecutionSpace>>(
            vector, field
        );
        return;
    }
    if (std::holds_alternative<NeoN::CPUExecutor>(field.exec()))
    {
        sunNVectorToVectorImpl<::sundials::kokkos::Vector<Kokkos::DefaultHostExecutionSpace>>(
            vector, field
        );
        return;
    }
    if (std::holds_alternative<NeoN::SerialExecutor>(field.exec()))
    {
        sunNVectorToVectorImpl<::sundials::kokkos::Vector<Kokkos::Serial>>(vector, field);
        return;
    }
    NF_ERROR_EXIT("Unsupported NeoN executor for field.");
};
 
template<typename SolutionVectorType>
int explicitRKSolve([[maybe_unused]] sunrealtype t, N_Vector y, N_Vector ydot, void* userData)
{
    // Pointer wrangling
    using ValueType = typename SolutionVectorType::VectorValueType;
    NeoN::dsl::Expression<ValueType>* pdeExpre =
        reinterpret_cast<NeoN::dsl::Expression<ValueType>*>(userData);
    sunrealtype* yDotArray = N_VGetArrayPointer(ydot);
    sunrealtype* yArray = N_VGetArrayPointer(y);
 
    NF_ASSERT(
        yDotArray != nullptr && yArray != nullptr && pdeExpre != nullptr,
        "Failed to dereference pointers in sundails."
    );
 
    auto size = static_cast<localIdx>(N_VGetLength(y));
    // Copy initial value from y to source.
    NeoN::Vector<NeoN::scalar> source = pdeExpre->explicitOperation(size) * -1.0; // compute spatial
    if (std::holds_alternative<NeoN::GPUExecutor>(pdeExpre->exec()))
    {
        Kokkos::fence();
    }
    NeoN::sundials::fieldToSunNVector(source, ydot); // assign rhs to ydot.
    return 0;
}
 
namespace detail
{
 
template<typename Vector>
void initNVector(size_t size, std::shared_ptr<SUNContext> context, Vector& vec)
{
    vec.initNVector(size, context);
}
 
template<typename Vector>
const N_Vector& sunNVector(const Vector& vec)
{
    return vec.sunNVector();
}
 
template<typename Vector>
N_Vector& sunNVector(Vector& vec)
{
    return vec.sunNVector();
}
}
 
template<typename ValueType>
class SKVectorSerial
{
public:
 
    SKVectorSerial() {};
    ~SKVectorSerial() = default;
    SKVectorSerial(const SKVectorSerial& other)
        : kvector_(other.kvector_), svector_(other.kvector_) {};
    SKVectorSerial(SKVectorSerial&& other) noexcept
        : kvector_(std::move(other.kvector_)), svector_(std::move(other.svector_)) {};
    SKVectorSerial& operator=(const SKVectorSerial& other) = delete;
    SKVectorSerial& operator=(SKVectorSerial&& other) = delete;
 
 
    using KVector = ::sundials::kokkos::Vector<Kokkos::Serial>;
    void initNVector(size_t size, std::shared_ptr<SUNContext> context)
    {
        kvector_ = KVector(size, *context);
        svector_ = kvector_;
    };
    const N_Vector& sunNVector() const { return svector_; };
    N_Vector& sunNVector() { return svector_; };
 
private:
 
    KVector kvector_ {}; 
    N_Vector svector_ {nullptr};
};
 
template<typename ValueType>
class SKVectorHostDefault
{
public:
 
    using KVector = ::sundials::kokkos::Vector<Kokkos::DefaultHostExecutionSpace>;
 
    SKVectorHostDefault() = default;
    ~SKVectorHostDefault() = default;
    SKVectorHostDefault(const SKVectorHostDefault& other)
        : kvector_(other.kvector_), svector_(other.kvector_) {};
    SKVectorHostDefault(SKVectorHostDefault&& other) noexcept
        : kvector_(std::move(other.kvector_)), svector_(std::move(other.svector_)) {};
    SKVectorHostDefault& operator=(const SKVectorHostDefault& other) = delete;
    SKVectorHostDefault& operator=(SKVectorHostDefault&& other) = delete;
 
    void initNVector(size_t size, std::shared_ptr<SUNContext> context)
    {
        kvector_ = KVector(size, *context);
        svector_ = kvector_;
    };
    const N_Vector& sunNVector() const { return svector_; };
    N_Vector& sunNVector() { return svector_; };
 
private:
 
    KVector kvector_ {}; 
    N_Vector svector_ {nullptr};
};
 
template<typename ValueType>
class SKVectorDefault
{
public:
 
    using KVector = ::sundials::kokkos::Vector<Kokkos::DefaultExecutionSpace>;
 
    SKVectorDefault() = default;
    ~SKVectorDefault() = default;
    SKVectorDefault(const SKVectorDefault& other)
        : kvector_(other.kvector_), svector_(other.kvector_) {};
    SKVectorDefault(SKVectorDefault&& other) noexcept
        : kvector_(std::move(other.kvector_)), svector_(std::move(other.svector_)) {};
    SKVectorDefault& operator=(const SKVectorDefault& other) = delete;
    SKVectorDefault& operator=(SKVectorDefault&& other) = delete;
 
    void initNVector(size_t size, std::shared_ptr<SUNContext> context)
    {
        kvector_ = KVector(size, *context);
        svector_ = kvector_;
    };
 
    const N_Vector& sunNVector() const { return svector_; };
 
    N_Vector& sunNVector() { return svector_; };
 
private:
 
    KVector kvector_ {}; 
    N_Vector svector_ {nullptr};
};
 
template<typename ValueType>
class SKVector
{
public:
 
    using SKVectorSerialV = SKVectorSerial<ValueType>;
    using SKVectorHostDefaultV = SKVectorHostDefault<ValueType>;
    using SKDefaultVectorV = SKVectorDefault<ValueType>;
    using SKVectorVariant = std::variant<SKVectorSerialV, SKVectorHostDefaultV, SKDefaultVectorV>;
 
    SKVector() { vector_.template emplace<SKVectorHostDefaultV>(); };
 
    ~SKVector() = default;
 
    SKVector(const SKVector&) = default;
 
    SKVector& operator=(const SKVector&) = delete;
 
    SKVector(SKVector&&) noexcept = default;
 
    SKVector& operator=(SKVector&&) noexcept = delete;
 
    void setExecutor(const NeoN::Executor& exec)
    {
        if (std::holds_alternative<NeoN::GPUExecutor>(exec))
        {
            vector_.template emplace<SKDefaultVectorV>();
            return;
        }
        if (std::holds_alternative<NeoN::CPUExecutor>(exec))
        {
            vector_.template emplace<SKVectorHostDefaultV>();
            return;
        }
        if (std::holds_alternative<NeoN::SerialExecutor>(exec))
        {
            vector_.template emplace<SKVectorSerialV>();
            return;
        }
 
        NF_ERROR_EXIT(
            "Unsupported NeoN executor: "
            << std::visit([](const auto& e) { return e.name(); }, exec) << "."
        );
    }
 
    void initNVector(size_t size, std::shared_ptr<SUNContext> context)
    {
        std::visit(
            [size, &context](auto& vec) { detail::initNVector(size, context, vec); }, vector_
        );
    }
 
    const N_Vector& sunNVector() const
    {
        return std::visit(
            [](const auto& vec) -> const N_Vector& { return detail::sunNVector(vec); }, vector_
        );
    }
 
    N_Vector& sunNVector()
    {
        return std::visit([](auto& vec) -> N_Vector& { return detail::sunNVector(vec); }, vector_);
    }
 
    const SKVectorVariant& variant() const { return vector_; }
 
    SKVectorVariant& variant() { return vector_; }
 
private:
 
    SKVectorVariant vector_; 
};
}
 
#endif