segmentedVector.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2023 - 2026 NeoN authors
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include "NeoN/core/view.hpp"
#include "NeoN/core/copyTo.hpp"
#include "NeoN/core/parallelAlgorithms.hpp"
#include "NeoN/core/primitives/label.hpp"
#include "NeoN/core/vector/vector.hpp"
 
namespace NeoN
{
 
namespace detail
{
 
/* @brief Prefix-sum kernel of segmentsFromIntervals.
 *
 * A named functor rather than a NEON_LAMBDA: under nvcc a NEON_LAMBDA is an
 * extended lambda whose host-side callable is registered per translation unit,
 * and for an enclosing *function template* the registration belongs to a comdat
 * group the linker may drop while keeping the body. The surviving copy then
 * calls a null pointer — every segmentsFromIntervals<localIdx> call inside
 * libNeoN segfaults (e.g. CellToFaceStencil::computeInternalStencil, the
 * stencil CellLimitedGrad builds). A functor carries no such registration.
 */
template<typename IndexType>
struct IntervalOffsetScan
{
    using value_type = IndexType;
 
    View<const IndexType> intervals;
    View<IndexType> offsets;
 
    KOKKOS_INLINE_FUNCTION
    void operator()(const localIdx i, IndexType& update, const bool final) const
    {
        update += intervals[i];
        if (final)
        {
            // offsets is a view, thus [] takes unsigned idx
            offsets[i] = update;
        }
    }
};
 
} // namespace detail
 
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()},
        detail::IntervalOffsetScan<IndexType> {inView, offsView},
        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 SupportsCopyTo<SegmentedVector<ValueType, IndexType>>
{
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; }
 
    [[nodiscard]] SegmentedVector<ValueType, IndexType> copyToExecutor(Executor exec) const override
    {
        return SegmentedVector<ValueType, IndexType>(
            values_.copyToExecutor(exec), segments_.copyToExecutor(exec)
        );
    }
 
 
    [[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