vector.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2023 NeoN authors
#pragma once
 
#include <Kokkos_Core.hpp>
 
#include "NeoN/core/error.hpp"
#include "NeoN/core/executor/executor.hpp"
#include "NeoN/core/primitives/label.hpp"
#include "NeoN/core/primitives/scalar.hpp"
#include "NeoN/core/view.hpp"
#include "NeoN/fields/fieldFreeFunctions.hpp"
#include "NeoN/fields/fieldTypeDefs.hpp"
 
#include <variant>
#include <vector>
 
 
namespace NeoN
{
 
namespace detail
{
template<typename ValueType>
auto deepCopyVisitor(localIdx ssize, const ValueType* srcPtr, ValueType* dstPtr)
{
    size_t size = static_cast<size_t>(ssize);
    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 Vector
{
 
public:
 
    using VectorValueType = ValueType;
 
    Vector(const Executor& exec, localIdx size) : size_(size), data_(nullptr), exec_(exec)
    {
        void* ptr = nullptr;
        std::visit(
            [&ptr, size](const auto& concreteExec)
            { ptr = concreteExec.alloc(static_cast<size_t>(size) * sizeof(ValueType)); },
            exec_
        );
        data_ = static_cast<ValueType*>(ptr);
    }
 
    Vector(
        const Executor& exec,
        const ValueType* in,
        localIdx size,
        Executor hostExec = SerialExecutor()
    )
        : size_(size), data_(nullptr), exec_(exec)
    {
        void* ptr = nullptr;
        std::visit(
            [&ptr, size](const auto& concreteExec)
            { ptr = concreteExec.alloc(static_cast<size_t>(size) * sizeof(ValueType)); },
            exec_
        );
        data_ = static_cast<ValueType*>(ptr);
        std::visit(detail::deepCopyVisitor<ValueType>(size_, in, data_), hostExec, exec_);
    }
 
    Vector(const Executor& exec, localIdx size, ValueType value)
        : size_(size), data_(nullptr), exec_(exec)
    {
        void* ptr = nullptr;
        std::visit(
            [&ptr, size](const auto& execu)
            { ptr = execu.alloc(static_cast<size_t>(size) * sizeof(ValueType)); },
            exec_
        );
        data_ = static_cast<ValueType*>(ptr);
        NeoN::fill(*this, value);
    }
 
    Vector(const Executor& exec, std::vector<ValueType> in)
        : Vector(exec, in.data(), static_cast<localIdx>(in.size()))
    {}
 
    Vector(const Executor& exec, const Vector<ValueType>& in)
        : Vector(exec, in.data(), in.size(), in.exec())
    {}
 
    Vector(const Vector<ValueType>& rhs) : Vector(rhs.exec(), rhs.data(), rhs.size(), rhs.exec()) {}
 
    Vector(Vector<ValueType>&& rhs) noexcept : size_(rhs.size_), data_(rhs.data_), exec_(rhs.exec_)
    {
        rhs.data_ = nullptr;
        rhs.size_ = 0;
    }
 
    ~Vector()
    {
        std::visit([this](const auto& exec) { exec.free(data_); }, exec_);
        data_ = nullptr;
    }
 
    template<typename func>
    void apply(func f)
    {
        map(*this, f);
    }
 
    [[nodiscard]] Vector<ValueType> copyToExecutor(Executor dstExec) const
    {
        if (dstExec == exec_) return *this;
 
        Vector<ValueType> result(dstExec, size_);
        std::visit(detail::deepCopyVisitor(size_, data_, result.data()), exec_, dstExec);
 
        return result;
    }
 
    [[nodiscard]] Vector<ValueType> copyToHost() const { return copyToExecutor(SerialExecutor()); }
 
    void copyToHost(Vector<ValueType>& result)
    {
        NF_DEBUG_ASSERT(
            result.size() == size_, "Parsed Vector size not the same as current field size"
        );
        result = copyToExecutor(SerialExecutor());
    }
 
    // ensures no return of device address on host --> invalid memory access
    ValueType& operator[](const localIdx i) = delete;
 
    // ensures no return of device address on host --> invalid memory access
    const ValueType& operator[](const localIdx i) const = delete;
 
    void operator=(const ValueType& rhs) { fill(*this, rhs); }
 
    void operator=(const Vector<ValueType>& rhs)
    {
        NF_ASSERT(exec_ == rhs.exec_, "Executors are not the same");
        if (this->size() != rhs.size())
        {
            this->resize(rhs.size());
        }
        setVector(*this, rhs.view());
    }
 
    Vector<ValueType>& operator+=(const Vector<ValueType>& rhs)
    {
        validateOtherVector(rhs);
        add(*this, rhs);
        return *this;
    }
 
    Vector<ValueType>& operator-=(const Vector<ValueType>& rhs)
    {
        validateOtherVector(rhs);
        sub(*this, rhs);
        return *this;
    }
 
    [[nodiscard]] Vector<ValueType> operator*(const Vector<scalar>& rhs)
    {
        validateOtherVector(rhs);
        Vector<ValueType> result(exec_, size_);
        result = *this;
        mul(result, rhs);
        return result;
    }
 
    [[nodiscard]] Vector<ValueType> operator*(const scalar rhs)
    {
        Vector<ValueType> result(exec_, size_);
        result = *this;
        scalarMul(result, rhs);
        return result;
    }
 
    Vector<ValueType>& operator*=(const Vector<scalar>& rhs)
    {
        validateOtherVector(rhs);
        Vector<ValueType>& result = *this;
        mul(result, rhs);
        return result;
    }
 
    Vector<ValueType>& operator*=(const scalar rhs)
    {
        Vector<ValueType>& result = *this;
        scalarMul(result, rhs);
        return result;
    }
 
    void resize(const localIdx size)
    {
        void* ptr = nullptr;
        if (!empty())
        {
            std::visit(
                [this, &ptr, size](const auto& exec)
                { ptr = exec.realloc(this->data_, static_cast<size_t>(size) * sizeof(ValueType)); },
                exec_
            );
        }
        else
        {
            std::visit(
                [&ptr, size](const auto& exec)
                { ptr = exec.alloc(static_cast<size_t>(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]] localIdx size() const { return size_; }
 
    [[nodiscard]] label ssize() const { return static_cast<label>(size_); }
 
    [[nodiscard]] bool empty() const { return size() == 0; }
 
    // return of a temporary --> invalid memory access
    View<ValueType> view() && = delete;
 
    // return of a temporary --> invalid memory access
    View<const ValueType> view() const&& = delete;
 
    [[nodiscard]] View<ValueType> view() &
    {
        return View<ValueType>(data_, static_cast<size_t>(size_));
    }
 
    [[nodiscard]] View<const ValueType> view() const&
    {
        return View<const ValueType>(data_, static_cast<size_t>(size_));
    }
 
    // return of a temporary --> invalid memory access
    [[nodiscard]] View<ValueType> view(std::pair<localIdx, localIdx> range) && = delete;
 
    // return of a temporary --> invalid memory access
    [[nodiscard]] View<const ValueType> view(std::pair<localIdx, localIdx> range) const&& = delete;
 
    [[nodiscard]] View<ValueType> view(std::pair<localIdx, localIdx> range) &
    {
        return View<ValueType>(
            data_ + range.first, static_cast<size_t>(range.second - range.first)
        );
    }
 
    [[nodiscard]] View<const ValueType> view(std::pair<localIdx, localIdx> range) const&
    {
        return View<const ValueType>(
            data_ + range.first, static_cast<size_t>(range.second - range.first)
        );
    }
 
    [[nodiscard]] std::pair<localIdx, localIdx> range() const { return {0, size()}; }
 
private:
 
    localIdx size_ {0};         
    ValueType* data_ {nullptr}; 
    const Executor exec_;       
 
    void validateOtherVector(const Vector<ValueType>& rhs) const
    {
        NF_DEBUG_ASSERT(size() == rhs.size(), "Vectors are not the same size.");
        NF_DEBUG_ASSERT(exec() == rhs.exec(), "Executors are not the same.");
    }
};
 
template<typename T>
[[nodiscard]] Vector<T> operator+(Vector<T> lhs, const Vector<T>& rhs)
{
    lhs += rhs;
    return lhs;
}
 
template<typename T>
[[nodiscard]] Vector<T> operator-(Vector<T> lhs, const Vector<T>& rhs)
{
    lhs -= rhs;
    return lhs;
}
 
} // namespace NeoN