range_accessors.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_RANGE_ACCESSORS_HPP_
#define GKO_PUBLIC_CORE_BASE_RANGE_ACCESSORS_HPP_
 
 
#include <array>
 
#include <ginkgo/core/base/range.hpp>
#include <ginkgo/core/base/types.hpp>
 
 
namespace gko {
namespace accessor {
 
 
template <typename ValueType, size_type Dimensionality>
class row_major {
public:
    friend class range<row_major>;
 
    static_assert(Dimensionality == 2,
                  "This accessor is only implemented for matrices");
 
    using value_type = ValueType;
 
    using data_type = value_type*;
 
    static constexpr size_type dimensionality = 2;
 
protected:
    constexpr GKO_ATTRIBUTES explicit row_major(data_type data,
                                                size_type num_rows,
                                                size_type num_cols,
                                                size_type stride)
        : data{data}, lengths{num_rows, num_cols}, stride{stride}
    {}
 
public:
    constexpr GKO_ATTRIBUTES value_type& operator()(size_type row,
                                                    size_type col) const
    {
        return GKO_ASSERT(row < lengths[0]), GKO_ASSERT(col < lengths[1]),
               data[row * stride + col];
    }
 
    constexpr GKO_ATTRIBUTES range<row_major> operator()(const span& rows,
                                                         const span& cols) const
    {
        return GKO_ASSERT(rows.is_valid()), GKO_ASSERT(cols.is_valid()),
               GKO_ASSERT(rows <= span{lengths[0]}),
               GKO_ASSERT(cols <= span{lengths[1]}),
               range<row_major>(data + rows.begin * stride + cols.begin,
                                rows.end - rows.begin, cols.end - cols.begin,
                                stride);
    }
 
    constexpr GKO_ATTRIBUTES size_type length(size_type dimension) const
    {
        return dimension < 2 ? lengths[dimension] : 1;
    }
 
    template <typename OtherAccessor>
    GKO_ATTRIBUTES void copy_from(const OtherAccessor& other) const
    {
        for (size_type i = 0; i < lengths[0]; ++i) {
            for (size_type j = 0; j < lengths[1]; ++j) {
                (*this)(i, j) = other(i, j);
            }
        }
    }
 
    const data_type data;
 
    const std::array<const size_type, dimensionality> lengths;
 
    const size_type stride;
};
 
 
}  // namespace accessor
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_RANGE_ACCESSORS_HPP_