segmentedVector.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2023 NeoN authors
#pragma once
 
#include "NeoN/core/view.hpp"
#include "NeoN/core/primitives/label.hpp"
#include "NeoN/core/vector.hpp"
 
namespace NeoN
{
 
template<typename IndexType>
IndexType segmentsFromIntervals(const Vector<IndexType>& intervals, Vector<IndexType>& offsets)
{
    IndexType finalValue = 0;
    const auto inSpan = intervals.view();
    // skip the first element of the offsets
    // assumed to be zero
    auto offsSpan = offsets.view().subspan(1);
    // NOTE avoid compiler warning by static_casting to localIdx since offsSpan
    // is a std::span
    NF_ASSERT_EQUAL(inSpan.size(), offsSpan.size());
    NeoN::parallelScan(
        intervals.exec(),
        {0, offsSpan.size()},
        KOKKOS_LAMBDA(const localIdx i, IndexType& update, const bool final) {
            update += inSpan[i];
            if (final)
            {
                // offsSpan is a std::span, thus [] takes unsigned idx
                offsSpan[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> span(localIdx segI) const
    {
        auto [start, length] = range(segI);
        return values.subspan(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 field from intervals.
     * @param intervals The intervals to create the segmented field 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);
    }
 
 
    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; }
 
 
    [[nodiscard]] SegmentedVectorView<ValueType, IndexType> view() &
    {
        return SegmentedVectorView<ValueType, IndexType> {values_.view(), segments_.view()};
    }
 
    // ensures no return a span of a temporary object --> invalid memory access
    [[nodiscard]] SegmentedVectorView<ValueType, IndexType> view() && = delete;
 
    [[nodiscard]] std::pair<View<ValueType>, View<IndexType>> spans() &
    {
        return {values_.view(), segments_.view()};
    }
 
    // ensures not to return a span of a temporary object --> invalid memory access
    [[nodiscard]] std::pair<View<ValueType>, View<IndexType>> spans() && = delete;
 
    const Vector<ValueType>& values() const { return values_; }
 
    const Vector<IndexType>& segments() const { return segments_; }
 
private:
 
    Vector<ValueType> values_;
    Vector<IndexType> segments_; 
};
 
} // namespace NeoN