array.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2025 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/view.hpp"
#include "NeoN/core/containerFreeFunctions.hpp"
 
#include <variant>
#include <vector>
 
 
namespace NeoN
{
 
template<typename ValueType>
class Array
{
 
public:
 
    using ArrayValueType = ValueType;
 
    Array(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);
    }
 
    Array(
        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_);
    }
 
 
    Array(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);
    }
 
    Array(const Executor& exec, std::vector<ValueType> in)
        : Array(exec, in.data(), static_cast<localIdx>(in.size()))
    {}
 
 
    Array(const Executor& exec, const Array<ValueType>& in)
        : Array(exec, in.data(), in.size(), in.exec())
    {}
 
    Array(const Array<ValueType>& rhs) : Array(rhs.exec(), rhs.data(), rhs.size(), rhs.exec()) {}
 
 
    Array(Array<ValueType>&& rhs) noexcept : size_(rhs.size_), data_(rhs.data_), exec_(rhs.exec_)
    {
        rhs.data_ = nullptr;
        rhs.size_ = 0;
    };
 
    ~Array()
    {
        std::visit([this](const auto& exec) { exec.free(data_); }, exec_);
        data_ = nullptr;
    }
 
    template<typename func>
    void apply(func f)
    {
        map(*this, f);
    }
 
    [[nodiscard]] Array<ValueType> copyToExecutor(Executor dstExec) const
    {
        if (dstExec == exec_) return *this;
 
        Array<ValueType> result(dstExec, size_);
        std::visit(detail::deepCopyVisitor(size_, data_, result.data()), exec_, dstExec);
 
        return result;
    }
 
    [[nodiscard]] Array<ValueType> copyToHost() const { return copyToExecutor(SerialExecutor()); }
 
    void copyToHost(Array<ValueType>& result)
    {
        NF_DEBUG_ASSERT(
            result.size() == size_, "Parsed Array size not the same as current field size"
        );
        result = copyToExecutor(SerialExecutor());
    }
 
    // ensures no return of device address on host --> invalid memory access
    Array& operator[](const localIdx i) = delete;
 
    // ensures no return of device address on host --> invalid memory access
    const Array& operator[](const localIdx i) const = delete;
 
    void operator=(const ValueType& rhs)
    {
        NF_ERROR_EXIT("Not implemented");
        fill(*this, rhs);
    }
 
    void operator=(const Array<ValueType>& rhs)
    {
        NF_ASSERT(exec_ == rhs.exec_, "Executors are not the same");
        if (this->size() != rhs.size())
        {
            this->resize(rhs.size());
        }
        setContainer(*this, rhs.view());
    }
 
    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]] inline ValueType* data() { return data_; }
 
    [[nodiscard]] inline const ValueType* data() const { return data_; }
 
    [[nodiscard]] inline const Executor& exec() const { return exec_; }
 
    [[nodiscard]] inline localIdx size() const { return size_; }
 
    [[nodiscard]] inline label ssize() const { return static_cast<label>(size_); }
 
    [[nodiscard]] inline 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]] inline View<ValueType> view() &
    {
        return View<ValueType>(data_, static_cast<size_t>(size_));
    }
 
    [[nodiscard]] inline 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]] inline View<ValueType> view(std::pair<localIdx, localIdx> range) &
    {
        return View<ValueType>(
            data_ + range.first, static_cast<size_t>(range.second - range.first)
        );
    }
 
    [[nodiscard]] inline 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]] inline std::pair<localIdx, localIdx> range() const { return {0, size()}; }
 
private:
 
    localIdx size_ {0};         
    ValueType* data_ {nullptr}; 
    const Executor exec_;       
 
    void validateOtherArray(const Array<ValueType>& rhs) const
    {
        NF_DEBUG_ASSERT(size() == rhs.size(), "Arrays are not the same size.");
        NF_DEBUG_ASSERT(exec() == rhs.exec(), "Executors are not the same.");
    }
};
 
} // namespace NeoN