segmentedVector.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2023 - 2025 NeoN authors
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include "NeoN/core/view.hpp"
#include "NeoN/core/parallelAlgorithms.hpp"
#include "NeoN/core/primitives/label.hpp"
#include "NeoN/core/vector/vector.hpp"
 
namespace NeoN
{
 
template<typename IndexType>
IndexType segmentsFromIntervals(const Vector<IndexType>& intervals, Vector<IndexType>& offsets)
{
    IndexType finalValue = 0;
    const auto inView = intervals.view();
    // skip the first element of the offsets
    // assumed to be zero
    auto offsView = offsets.view().subview(1);
    // NOTE avoid compiler warning by static_casting to localIdx since offsView
    // is a View
    NF_ASSERT_EQUAL(inView.size(), offsView.size());
    NeoN::parallelScan(
        intervals.exec(),
        {0, offsView.size()},
        KOKKOS_LAMBDA(const localIdx i, IndexType& update, const bool final) {
            update += inView[i];
            if (final)
            {
                // offView is a view, thus [] takes unsigned idx
                offsView[i] = update;
            }
        },
        finalValue
    );
    return finalValue;
}
 
template<typename ValueType, typename IndexType = NeoN::localIdx>
class SegmentedVectorView
{
public:
 
    View<ValueType> values;
 
    View<IndexType> segments;
 
    KOKKOS_INLINE_FUNCTION
    Kokkos::pair<IndexType, IndexType> bounds(localIdx segI) const
    {
        return Kokkos::pair<IndexType, IndexType> {segments[segI], segments[segI + 1]};
    }
 
    KOKKOS_INLINE_FUNCTION
    Kokkos::pair<IndexType, IndexType> range(localIdx segI) const
    {
        return Kokkos::pair<IndexType, IndexType> {
            segments[segI], segments[segI + 1] - segments[segI]
        };
    }
 
    KOKKOS_INLINE_FUNCTION View<ValueType> view(localIdx segI) const
    {
        auto [start, length] = range(segI);
        return values.subview(start, length);
    }
 
    KOKKOS_INLINE_FUNCTION
    IndexType operator[](localIdx i) const { return segments[i]; }
};
 
template<typename ValueType, typename IndexType>
class SegmentedVector
{
public:
 
 
    SegmentedVector(const Executor& exec, localIdx size, localIdx numSegments)
        : values_(exec, size), segments_(exec, numSegments + 1)
    {}
 
    /*
     * @brief Create a segmented vector from intervals.
     * @param intervals The intervals to create the segmented vector from.
     * @note The intervals are the lengths of each segment
     */
    SegmentedVector(const Vector<IndexType>& intervals)
        : values_(intervals.exec(), 0),
          segments_(intervals.exec(), intervals.size() + 1, IndexType(0))
    {
        IndexType valueSize = segmentsFromIntervals(intervals, segments_);
        values_ = Vector<ValueType>(intervals.exec(), valueSize, ValueType(0));
    }
 
 
    SegmentedVector(const Vector<ValueType>& values, const Vector<IndexType>& segments)
        : values_(values), segments_(segments)
    {
        NF_ASSERT(values.exec() == segments.exec(), "Executors are not the same.");
    }
 
 
    const Executor& exec() const { return values_.exec(); }
 
    localIdx size() const { return values_.size(); }
 
    localIdx numSegments() const { return segments_.size() - 1; }
 
    SegmentedVector<ValueType, IndexType> copyToHost() const
    {
        SegmentedVector<ValueType, IndexType> result(values_.copyToHost(), segments_.copyToHost());
        return result;
    }
 
 
    [[nodiscard]] SegmentedVectorView<ValueType, IndexType> view() &
    {
        return SegmentedVectorView<ValueType, IndexType> {values_.view(), segments_.view()};
    }
 
    // ensures no return a view of a temporary object --> invalid memory access
    [[nodiscard]] SegmentedVectorView<ValueType, IndexType> view() && = delete;
 
    [[nodiscard]] std::pair<View<ValueType>, View<IndexType>> views() &
    {
        return {values_.view(), segments_.view()};
    }
 
 
    // ensures not to return a view of a temporary object --> invalid memory access
    [[nodiscard]] std::pair<View<ValueType>, View<IndexType>> views() && = delete;
 
    const Vector<ValueType>& values() const { return values_; }
 
    const Vector<IndexType>& segments() const { return segments_; }
 
private:
 
    Vector<ValueType> values_;
    Vector<IndexType> segments_; 
};
 
} // namespace NeoN