segmented_array.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#pragma once
#include <numeric>
 
#include <ginkgo/config.hpp>
#include <ginkgo/core/base/array.hpp>
 
 
namespace gko {
 
template <typename T>
struct segmented_array {
    explicit segmented_array(std::shared_ptr<const Executor> exec);
 
    static segmented_array create_from_sizes(const gko::array<int64>& sizes);
 
    static segmented_array create_from_sizes(gko::array<T> buffer,
                                             const gko::array<int64>& sizes);
 
    static segmented_array create_from_offsets(gko::array<int64> offsets);
 
    static segmented_array create_from_offsets(gko::array<T> buffer,
                                               gko::array<int64> offsets);
 
    segmented_array(std::shared_ptr<const Executor> exec,
                    const segmented_array& other);
 
    segmented_array(std::shared_ptr<const Executor> exec,
                    segmented_array&& other);
 
    segmented_array(const segmented_array& other);
 
    segmented_array(segmented_array&& other) noexcept(false);
 
    segmented_array& operator=(const segmented_array& other);
 
    segmented_array& operator=(segmented_array&&) noexcept(false);
 
    size_type get_size() const;
 
    size_type get_segment_count() const;
 
    T* get_flat_data();
 
    const T* get_const_flat_data() const;
 
    const gko::array<int64>& get_offsets() const;
 
    std::shared_ptr<const Executor> get_executor() const;
 
private:
    gko::array<T> buffer_;
    gko::array<int64> offsets_;
};
 
 
namespace detail {
 
 
template <typename T>
struct temporary_clone_helper<segmented_array<T>> {
    static std::unique_ptr<segmented_array<T>> create(
        std::shared_ptr<const Executor> exec, segmented_array<T>* ptr,
        bool copy_data)
    {
        if (copy_data) {
            return std::make_unique<segmented_array<T>>(
                make_array_view(exec, ptr->get_size(), ptr->get_flat_data()),
                ptr->get_offsets());
        } else {
            return std::make_unique<segmented_array<T>>(std::move(exec),
                                                        ptr->get_offsets());
        }
    }
};
 
template <typename T>
struct temporary_clone_helper<const segmented_array<T>> {
    static std::unique_ptr<const segmented_array<T>> create(
        std::shared_ptr<const Executor> exec, const segmented_array<T>* ptr,
        bool)
    {
        return std::make_unique<segmented_array<T>>(
            make_array_view(exec, ptr->get_size(), ptr->get_const_flat_data()),
            ptr->get_offsets());
    }
};
 
 
template <typename T>
class copy_back_deleter<segmented_array<T>>
    : public copy_back_deleter_from_assignment<segmented_array<T>> {
public:
    using copy_back_deleter_from_assignment<
        segmented_array<T>>::copy_back_deleter_from_assignment;
};
 
 
}  // namespace detail
}  // namespace gko