array.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_ARRAY_HPP_
#define GKO_PUBLIC_CORE_BASE_ARRAY_HPP_
 
 
#include <algorithm>
#include <iterator>
#include <memory>
#include <type_traits>
#include <utility>
 
#include <ginkgo/core/base/exception.hpp>
#include <ginkgo/core/base/exception_helpers.hpp>
#include <ginkgo/core/base/executor.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
 
 
namespace gko {
 
 
template <typename ValueType>
class array;
 
 
namespace detail {
 
 
template <typename SourceType, typename TargetType>
void convert_data(std::shared_ptr<const Executor> exec, size_type size,
                  const SourceType* src, TargetType* dst);
 
 
template <typename ValueType>
class const_array_view {
public:
    using value_type = ValueType;
 
    const_array_view(std::shared_ptr<const Executor> exec, size_type size,
                     const ValueType* data)
        : exec_{std::move(exec)}, size_{size}, data_{data}
    {}
 
    /*
     * To avoid any collisions with the value semantics of normal arrays,
     * disable assignment and copy-construction altogether.
     */
    const_array_view& operator=(const const_array_view&) = delete;
    const_array_view& operator=(const_array_view&&) = delete;
    const_array_view(const const_array_view&) = delete;
    /*
     * TODO C++17: delete this overload as well, it is no longer necessary due
     * to guaranteed RVO.
     */
    const_array_view(const_array_view&& other)
        : const_array_view{other.exec_, other.size_, other.data_}
    {
        other.size_ = 0;
        other.data_ = nullptr;
    }
 
    size_type get_size() const noexcept { return size_; }
 
    GKO_DEPRECATED("use get_size() instead")
    size_type get_num_elems() const noexcept { return get_size(); }
 
    const value_type* get_const_data() const noexcept { return data_; }
 
    std::shared_ptr<const Executor> get_executor() const noexcept
    {
        return exec_;
    }
 
    bool is_owning() const noexcept { return false; }
 
    array<ValueType> copy_to_array() const;
 
private:
    std::shared_ptr<const Executor> exec_;
    size_type size_;
    const ValueType* data_;
};
 
 
template <typename ValueType>
using ConstArrayView GKO_DEPRECATED("please use const_array_view") =
    const_array_view<ValueType>;
 
 
template <typename ValueType>
array<ValueType> array_const_cast(const_array_view<ValueType> view);
 
 
}  // namespace detail
 
 
template <typename ValueType>
class array {
public:
    using value_type = ValueType;
 
    using default_deleter = executor_deleter<value_type[]>;
 
    using view_deleter = null_deleter<value_type[]>;
 
    array() noexcept
        : size_(0), data_(nullptr, default_deleter{nullptr}), exec_(nullptr)
    {}
 
    explicit array(std::shared_ptr<const Executor> exec) noexcept
        : size_(0),
          data_(nullptr, default_deleter{exec}),
          exec_(std::move(exec))
    {}
 
    array(std::shared_ptr<const Executor> exec, size_type size)
        : size_(size),
          data_(nullptr, default_deleter{exec}),
          exec_(std::move(exec))
    {
        if (size > 0) {
            data_.reset(exec_->alloc<value_type>(size));
        }
    }
 
    template <typename DeleterType>
    array(std::shared_ptr<const Executor> exec, size_type size,
          value_type* data, DeleterType deleter)
        : size_{size}, data_(data, deleter), exec_{exec}
    {}
 
    array(std::shared_ptr<const Executor> exec, size_type size,
          value_type* data)
        : array(exec, size, data, default_deleter{exec})
    {}
 
    template <typename RandomAccessIterator>
    array(std::shared_ptr<const Executor> exec, RandomAccessIterator begin,
          RandomAccessIterator end)
        : array(exec)
    {
        array tmp(exec->get_master(), std::distance(begin, end));
        std::copy(begin, end, tmp.data_.get());
        *this = std::move(tmp);
    }
 
    template <typename T>
    array(std::shared_ptr<const Executor> exec,
          std::initializer_list<T> init_list)
        : array(exec, begin(init_list), end(init_list))
    {}
 
    array(std::shared_ptr<const Executor> exec, const array& other)
        : array(exec)
    {
        *this = other;
    }
 
    array(const array& other) : array(other.get_executor(), other) {}
 
    array(std::shared_ptr<const Executor> exec, array&& other) : array(exec)
    {
        *this = std::move(other);
    }
 
    array(array&& other) : array(other.get_executor(), std::move(other)) {}
 
    static array view(std::shared_ptr<const Executor> exec, size_type size,
                      value_type* data)
    {
        return array{exec, size, data, view_deleter{}};
    }
 
    static detail::const_array_view<ValueType> const_view(
        std::shared_ptr<const Executor> exec, size_type size,
        const value_type* data)
    {
        return {exec, size, data};
    }
 
    array<ValueType> as_view()
    {
        return view(this->get_executor(), this->get_size(), this->get_data());
    }
 
    detail::const_array_view<ValueType> as_const_view() const
    {
        return const_view(this->get_executor(), this->get_size(),
                          this->get_const_data());
    }
 
    array& operator=(const array& other)
    {
        if (&other == this) {
            return *this;
        }
        if (exec_ == nullptr) {
            exec_ = other.get_executor();
            data_ = data_manager{nullptr, other.data_.get_deleter()};
        }
        if (other.get_executor() == nullptr) {
            this->clear();
            return *this;
        }
 
        if (this->is_owning()) {
            this->resize_and_reset(other.get_size());
        } else {
            GKO_ENSURE_COMPATIBLE_BOUNDS(other.get_size(), this->get_size());
        }
        exec_->copy_from(other.get_executor(), other.get_size(),
                         other.get_const_data(), this->get_data());
        return *this;
    }
 
    array& operator=(array&& other)
    {
        if (&other == this) {
            return *this;
        }
        if (exec_ == nullptr) {
            exec_ = other.get_executor();
            data_ = data_manager{nullptr, default_deleter{exec_}};
        }
        if (other.get_executor() == nullptr) {
            this->clear();
            return *this;
        }
        if (exec_ == other.get_executor()) {
            // same device, only move the pointer
            data_ = std::exchange(
                other.data_, data_manager{nullptr, default_deleter{exec_}});
            size_ = std::exchange(other.size_, 0);
        } else {
            // different device, copy the data
            *this = other;
            other.clear();
        }
        return *this;
    }
 
    template <typename OtherValueType>
    std::enable_if_t<!std::is_same<ValueType, OtherValueType>::value, array>&
    operator=(const array<OtherValueType>& other)
    {
        if (this->exec_ == nullptr) {
            this->exec_ = other.get_executor();
            this->data_ = data_manager{nullptr, default_deleter{this->exec_}};
        }
        if (other.get_executor() == nullptr) {
            this->clear();
            return *this;
        }
 
        if (this->is_owning()) {
            this->resize_and_reset(other.get_size());
        } else {
            GKO_ENSURE_COMPATIBLE_BOUNDS(other.get_size(), this->get_size());
        }
        array<OtherValueType> tmp{this->exec_};
        const OtherValueType* source = other.get_const_data();
        // if we are on different executors: copy, then convert
        if (this->exec_ != other.get_executor()) {
            tmp = other;
            source = tmp.get_const_data();
        }
        detail::convert_data(this->exec_, other.get_size(), source,
                             this->get_data());
        return *this;
    }
 
    array& operator=(const detail::const_array_view<ValueType>& other)
    {
        if (this->exec_ == nullptr) {
            this->exec_ = other.get_executor();
            this->data_ = data_manager{nullptr, default_deleter{this->exec_}};
        }
        if (other.get_executor() == nullptr) {
            this->clear();
            return *this;
        }
 
        if (this->is_owning()) {
            this->resize_and_reset(other.get_size());
        } else {
            GKO_ENSURE_COMPATIBLE_BOUNDS(other.get_size(), this->get_size());
        }
        array tmp{this->exec_};
        const ValueType* source = other.get_const_data();
        // if we are on different executors: copy
        if (this->exec_ != other.get_executor()) {
            tmp = other.copy_to_array();
            source = tmp.get_const_data();
        }
        exec_->copy_from(other.get_executor(), other.get_size(), source,
                         this->get_data());
        return *this;
    }
 
    void clear() noexcept
    {
        size_ = 0;
        data_.reset(nullptr);
    }
 
    void resize_and_reset(size_type size)
    {
        if (size == size_) {
            return;
        }
        if (exec_ == nullptr) {
            throw gko::NotSupported(__FILE__, __LINE__, __func__,
                                    "gko::Executor (nullptr)");
        }
        if (!this->is_owning()) {
            throw gko::NotSupported(__FILE__, __LINE__, __func__,
                                    "Non owning gko::array cannot be resized.");
        }
 
        if (size > 0 && this->is_owning()) {
            size_ = size;
            data_.reset(exec_->alloc<value_type>(size));
        } else {
            this->clear();
        }
    }
 
    std::vector<value_type> copy_to_host() const
    {
        std::vector<value_type> result(this->get_size());
        auto view = make_array_view(this->get_executor()->get_master(),
                                    this->get_size(), result.data());
        view = *this;
        return result;
    }
 
    void fill(const value_type value);
 
    size_type get_size() const noexcept { return size_; }
 
    GKO_DEPRECATED("use get_size() instead")
    size_type get_num_elems() const noexcept { return get_size(); }
 
    value_type* get_data() noexcept { return data_.get(); }
 
    const value_type* get_const_data() const noexcept { return data_.get(); }
 
    std::shared_ptr<const Executor> get_executor() const noexcept
    {
        return exec_;
    }
 
    void set_executor(std::shared_ptr<const Executor> exec)
    {
        if (exec == exec_) {
            // moving to the same executor, no-op
            return;
        }
        array tmp(std::move(exec));
        tmp = *this;
        exec_ = std::move(tmp.exec_);
        data_ = std::move(tmp.data_);
    }
 
    bool is_owning()
    {
        return data_.get_deleter().target_type() == typeid(default_deleter);
    }
 
 
private:
    // Allow other array types to access private members
    template <typename OtherValueType>
    friend class array;
 
    using data_manager =
        std::unique_ptr<value_type[], std::function<void(value_type[])>>;
 
    size_type size_;
    data_manager data_;
    std::shared_ptr<const Executor> exec_;
};
 
 
template <typename ValueType>
using Array GKO_DEPRECATED("please use array") = array<ValueType>;
 
 
template <typename ValueType>
ValueType reduce_add(const array<ValueType>& input_arr,
                     const ValueType init_val = 0);
 
template <typename ValueType>
void reduce_add(const array<ValueType>& input_arr, array<ValueType>& result);
 
 
template <typename ValueType>
array<ValueType> make_array_view(std::shared_ptr<const Executor> exec,
                                 size_type size, ValueType* data)
{
    return array<ValueType>::view(exec, size, data);
}
 
 
template <typename ValueType>
detail::const_array_view<ValueType> make_const_array_view(
    std::shared_ptr<const Executor> exec, size_type size, const ValueType* data)
{
    return array<ValueType>::const_view(exec, size, data);
}
 
 
namespace detail {
 
 
template <typename T>
struct temporary_clone_helper<array<T>> {
    static std::unique_ptr<array<T>> create(
        std::shared_ptr<const Executor> exec, array<T>* ptr, bool copy_data)
    {
        if (copy_data) {
            return std::make_unique<array<T>>(std::move(exec), *ptr);
        } else {
            return std::make_unique<array<T>>(std::move(exec), ptr->get_size());
        }
    }
};
 
template <typename T>
struct temporary_clone_helper<const array<T>> {
    static std::unique_ptr<const array<T>> create(
        std::shared_ptr<const Executor> exec, const array<T>* ptr, bool)
    {
        return std::make_unique<const array<T>>(std::move(exec), *ptr);
    }
};
 
 
// specialization for non-constant arrays, copying back via assignment
template <typename T>
class copy_back_deleter<array<T>>
    : public copy_back_deleter_from_assignment<array<T>> {
public:
    using copy_back_deleter_from_assignment<
        array<T>>::copy_back_deleter_from_assignment;
};
 
 
template <typename ValueType>
array<ValueType> array_const_cast(const_array_view<ValueType> view)
{
    return array<ValueType>::view(
        view.get_executor(), view.get_size(),
        const_cast<ValueType*>(view.get_const_data()));
}
 
 
template <typename ValueType>
array<ValueType> const_array_view<ValueType>::copy_to_array() const
{
    array<ValueType> result(this->get_executor(), this->get_size());
    result.get_executor()->copy_from(this->get_executor(), this->get_size(),
                                     this->get_const_data(), result.get_data());
    return result;
}
 
 
}  // namespace detail
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_ARRAY_HPP_