field.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2023 NeoFOAM authors
#pragma once
 
#include <Kokkos_Core.hpp>
 
#include <iostream>
#include <span>
 
#include "NeoFOAM/core/error.hpp"
#include "NeoFOAM/core/executor/executor.hpp"
#include "NeoFOAM/core/primitives/scalar.hpp"
#include "NeoFOAM/fields/fieldFreeFunctions.hpp"
#include "NeoFOAM/fields/fieldTypeDefs.hpp"
 
namespace NeoFOAM
{
 
namespace detail
{
template<typename ValueType>
auto deepCopyVisitor(size_t size, const ValueType* srcPtr, ValueType* dstPtr)
{
    return [size, srcPtr, dstPtr](const auto& srcExec, const auto& dstExec)
    {
        Kokkos::deep_copy(
            dstExec.createKokkosView(dstPtr, size), srcExec.createKokkosView(srcPtr, size)
        );
    };
}
}
 
template<typename ValueType>
class Field
{
 
public:
 
    Field(const Executor& exec, size_t size) : size_(size), data_(nullptr), exec_(exec)
    {
        void* ptr = nullptr;
        std::visit(
            [&ptr, size](const auto& concreteExec)
            { ptr = concreteExec.alloc(size * sizeof(ValueType)); },
            exec_
        );
        data_ = static_cast<ValueType*>(ptr);
    }
 
    Field(
        const Executor& exec, const ValueType* in, size_t size, Executor hostExec = SerialExecutor()
    )
        : size_(size), data_(nullptr), exec_(exec)
    {
        void* ptr = nullptr;
        std::visit(
            [&ptr, size](const auto& concreteExec)
            { ptr = concreteExec.alloc(size * sizeof(ValueType)); },
            exec_
        );
        data_ = static_cast<ValueType*>(ptr);
        std::visit(detail::deepCopyVisitor<ValueType>(size_, in, data_), hostExec, exec_);
    }
 
    Field(const Executor& exec, size_t size, ValueType value)
        : size_(size), data_(nullptr), exec_(exec)
    {
        void* ptr = nullptr;
        std::visit(
            [&ptr, size](const auto& execu) { ptr = execu.alloc(size * sizeof(ValueType)); }, exec_
        );
        data_ = static_cast<ValueType*>(ptr);
        NeoFOAM::fill(*this, value);
    }
 
    Field(const Executor& exec, std::vector<ValueType> in) : Field(exec, in.data(), in.size()) {}
 
    Field(const Executor& exec, const Field<ValueType>& in)
        : Field(exec, in.data(), in.size(), in.exec())
    {}
 
    Field(const Field<ValueType>& rhs) : data_(nullptr), exec_(rhs.exec_)
    {
        resize(rhs.size_);
        setField(*this, rhs.span());
    }
 
    Field(Field<ValueType>&& rhs) noexcept : size_(rhs.size_), data_(rhs.data_), exec_(rhs.exec_)
    {
        rhs.data_ = nullptr;
        rhs.size_ = 0;
    }
 
    ~Field()
    {
        std::visit([this](const auto& exec) { exec.free(data_); }, exec_);
        data_ = nullptr;
    }
 
    template<typename func>
    void apply(func f)
    {
        map(*this, f);
    }
 
    [[nodiscard]] Field<ValueType> copyToExecutor(Executor dstExec) const
    {
        if (dstExec == exec_) return *this;
 
        Field<ValueType> result(dstExec, size_);
        std::visit(detail::deepCopyVisitor(size_, data_, result.data()), exec_, dstExec);
 
        return result;
    }
 
    [[nodiscard]] Field<ValueType> copyToHost() const { return copyToExecutor(SerialExecutor()); }
 
    void copyToHost(Field<ValueType>& result)
    {
        NF_DEBUG_ASSERT(
            result.size() == size_, "Parsed Field size not the same as current field size"
        );
        result = copyToExecutor(SerialExecutor());
    }
 
    KOKKOS_INLINE_FUNCTION
    ValueType& operator[](const size_t i) { return data_[i]; }
 
    KOKKOS_INLINE_FUNCTION
    const ValueType& operator[](const size_t i) const { return data_[i]; }
 
    KOKKOS_INLINE_FUNCTION
    ValueType& operator()(const size_t i) { return data_[i]; }
 
    KOKKOS_INLINE_FUNCTION
    const ValueType& operator()(const size_t i) const { return data_[i]; }
 
    void operator=(const ValueType& rhs) { fill(*this, rhs); }
 
    void operator=(const Field<ValueType>& rhs)
    {
        NF_ASSERT(exec_ == rhs.exec_, "Executors are not the same");
        if (this->size() != rhs.size())
        {
            this->resize(rhs.size());
        }
        setField(*this, rhs.span());
    }
 
    Field<ValueType>& operator+=(const Field<ValueType>& rhs)
    {
        validateOtherField(rhs);
        add(*this, rhs);
        return *this;
    }
 
    Field<ValueType>& operator-=(const Field<ValueType>& rhs)
    {
        validateOtherField(rhs);
        sub(*this, rhs);
        return *this;
    }
 
    [[nodiscard]] Field<ValueType> operator*(const Field<scalar>& rhs)
    {
        validateOtherField(rhs);
        Field<ValueType> result(exec_, size_);
        result = *this;
        mul(result, rhs);
        return result;
    }
 
    [[nodiscard]] Field<ValueType> operator*(const scalar rhs)
    {
        Field<ValueType> result(exec_, size_);
        result = *this;
        scalarMul(result, rhs);
        return result;
    }
 
    void resize(const size_t size)
    {
        void* ptr = nullptr;
        if (!empty())
        {
            std::visit(
                [this, &ptr, size](const auto& exec)
                { ptr = exec.realloc(this->data_, size * sizeof(ValueType)); },
                exec_
            );
        }
        else
        {
            std::visit(
                [&ptr, size](const auto& exec) { ptr = exec.alloc(size * sizeof(ValueType)); },
                exec_
            );
        }
        data_ = static_cast<ValueType*>(ptr);
        size_ = size;
    }
 
    [[nodiscard]] ValueType* data() { return data_; }
 
    [[nodiscard]] const ValueType* data() const { return data_; }
 
    [[nodiscard]] const Executor& exec() const { return exec_; }
 
    [[nodiscard]] size_t size() const { return size_; }
 
    [[nodiscard]] bool empty() const { return size() == 0; }
 
    // ensures no return a span of a temporary object --> invalid memory access
    std::span<ValueType> span() && = delete;
 
    // ensures no return a span of a temporary object --> invalid memory access
    std::span<const ValueType> span() const&& = delete;
 
    [[nodiscard]] std::span<ValueType> span() & { return std::span<ValueType>(data_, size_); }
 
    [[nodiscard]] std::span<const ValueType> span() const&
    {
        return std::span<const ValueType>(data_, size_);
    }
 
    // ensures no return a span of a temporary object --> invalid memory access
    [[nodiscard]] std::span<ValueType> span(std::pair<size_t, size_t> range) && = delete;
 
    // ensures no return a span of a temporary object --> invalid memory access
    [[nodiscard]] std::span<const ValueType> span(std::pair<size_t, size_t> range) const&& = delete;
 
    [[nodiscard]] std::span<ValueType> span(std::pair<size_t, size_t> range) &
    {
        return std::span<ValueType>(data_ + range.first, range.second - range.first);
    }
 
    [[nodiscard]] std::span<const ValueType> span(std::pair<size_t, size_t> range) const&
    {
        return std::span<const ValueType>(data_ + range.first, range.second - range.first);
    }
 
    [[nodiscard]] std::pair<size_t, size_t> range() const { return {0, size()}; }
 
private:
 
    size_t size_ {0};           
    ValueType* data_ {nullptr}; 
    const Executor exec_;       
 
    void validateOtherField(const Field<ValueType>& rhs) const
    {
        NF_DEBUG_ASSERT(size() == rhs.size(), "Fields are not the same size.");
        NF_DEBUG_ASSERT(exec() == rhs.exec(), "Executors are not the same.");
    }
};
 
template<typename T>
[[nodiscard]] Field<T> operator+(Field<T> lhs, const Field<T>& rhs)
{
    lhs += rhs;
    return lhs;
}
 
template<typename T>
[[nodiscard]] Field<T> operator-(Field<T> lhs, const Field<T>& rhs)
{
    lhs -= rhs;
    return lhs;
}
 
} // namespace NeoFOAM