device_matrix_data.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_DEVICE_MATRIX_DATA_HPP_
#define GKO_PUBLIC_CORE_BASE_DEVICE_MATRIX_DATA_HPP_
 
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/dim.hpp>
#include <ginkgo/core/base/exception_helpers.hpp>
#include <ginkgo/core/base/executor.hpp>
#include <ginkgo/core/base/matrix_data.hpp>
#include <ginkgo/core/base/temporary_clone.hpp>
 
 
namespace gko {
 

template <typename ValueType, typename IndexType>
class device_matrix_data {
    GKO_ASSERT_SUPPORTED_VALUE_AND_INDEX_TYPE;
 
public:
    using value_type = ValueType;
    using index_type = IndexType;
    using nonzero_type = matrix_data_entry<value_type, index_type>;
    using host_type = matrix_data<value_type, index_type>;

    explicit device_matrix_data(std::shared_ptr<const Executor> exec,
                                dim<2> size = {}, size_type num_entries = 0);

    device_matrix_data(std::shared_ptr<const Executor> exec,
                       const device_matrix_data& data);

    device_matrix_data(std::shared_ptr<const Executor> exec, dim<2> size,
                       array<index_type> row_idxs, array<index_type> col_idxs,
                       array<value_type> values);

    template <typename InputValueType, typename RowIndexType,
              typename ColIndexType>
    GKO_DEPRECATED(
        "explicitly construct the gko::array arguments instead of passing "
        "initializer lists")
    device_matrix_data(std::shared_ptr<const Executor> exec, dim<2> size,
                       std::initializer_list<RowIndexType> row_idxs,
                       std::initializer_list<ColIndexType> col_idxs,
                       std::initializer_list<InputValueType> values)
        : device_matrix_data{exec, size,
                             array<index_type>{exec, std::move(row_idxs)},
                             array<index_type>{exec, std::move(col_idxs)},
                             array<value_type>{exec, std::move(values)}}
    {}

    host_type copy_to_host() const;

    static device_matrix_data create_from_host(
        std::shared_ptr<const Executor> exec, const host_type& data);

    void fill_zero();

    void sort_row_major();

    void remove_zeros();

    void sum_duplicates();

    std::shared_ptr<const Executor> get_executor() const
    {
        return values_.get_executor();
    }

    dim<2> get_size() const { return size_; }

    GKO_DEPRECATED("use get_num_stored_elements()")
    size_type get_num_elems() const { return get_num_stored_elements(); }

    size_type get_num_stored_elements() const { return values_.get_size(); }

    index_type* get_row_idxs() { return row_idxs_.get_data(); }

    const index_type* get_const_row_idxs() const
    {
        return row_idxs_.get_const_data();
    }

    index_type* get_col_idxs() { return col_idxs_.get_data(); }

    const index_type* get_const_col_idxs() const
    {
        return col_idxs_.get_const_data();
    }

    value_type* get_values() { return values_.get_data(); }

    const value_type* get_const_values() const
    {
        return values_.get_const_data();
    }

    void resize_and_reset(size_type new_num_entries);

    void resize_and_reset(dim<2> new_size, size_type new_num_entries);

    struct arrays {
        array<index_type> row_idxs;
        array<index_type> col_idxs;
        array<value_type> values;
    };

    arrays empty_out();
 
private:
    dim<2> size_;
    array<index_type> row_idxs_;
    array<index_type> col_idxs_;
    array<value_type> values_;
};
 
 
namespace detail {
 
 
template <typename ValueType, typename IndexType>
struct temporary_clone_helper<device_matrix_data<ValueType, IndexType>> {
    static std::unique_ptr<device_matrix_data<ValueType, IndexType>> create(
        std::shared_ptr<const Executor> exec,
        device_matrix_data<ValueType, IndexType>* ptr, bool copy_data)
    {
        if (copy_data) {
            return std::make_unique<device_matrix_data<ValueType, IndexType>>(
                std::move(exec), *ptr);
        } else {
            return std::make_unique<device_matrix_data<ValueType, IndexType>>(
                std::move(exec), ptr->get_size(),
                ptr->get_num_stored_elements());
        }
    }
};
 
template <typename ValueType, typename IndexType>
struct temporary_clone_helper<const device_matrix_data<ValueType, IndexType>> {
    static std::unique_ptr<const device_matrix_data<ValueType, IndexType>>
    create(std::shared_ptr<const Executor> exec,
           const device_matrix_data<ValueType, IndexType>* ptr, bool)
    {
        return std::make_unique<const device_matrix_data<ValueType, IndexType>>(
            std::move(exec), *ptr);
    }
};
 
 
template <typename ValueType, typename IndexType>
class copy_back_deleter<device_matrix_data<ValueType, IndexType>>
    : public copy_back_deleter_from_assignment<
          device_matrix_data<ValueType, IndexType>> {
public:
    using copy_back_deleter_from_assignment<device_matrix_data<
        ValueType, IndexType>>::copy_back_deleter_from_assignment;
};
 
 
}  // namespace detail
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_DEVICE_MATRIX_DATA_HPP_