CSRMatrix.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2023 NeoN authors
#pragma once
 
#include "NeoN/core/vector.hpp"
 
#include <type_traits>
 
namespace NeoN::la
{
 
template<typename ValueType, typename IndexType>
struct CSRMatrixView
{
    CSRMatrixView(
        const View<ValueType>& valueView,
        const View<IndexType>& colIdxsView,
        const View<IndexType>& rowOffsView
    )
        : values(valueView), colIdxs(colIdxsView), rowOffs(rowOffsView) {};
 
    ~CSRMatrixView() = default;
 
    KOKKOS_INLINE_FUNCTION
    ValueType& entry(const IndexType i, const IndexType j) const
    {
        const IndexType rowSize = rowOffs[i + 1] - rowOffs[i];
        for (std::remove_const_t<IndexType> ic = 0; ic < rowSize; ++ic)
        {
            const IndexType localCol = rowOffs[i] + ic;
            if (colIdxs[localCol] == j)
            {
                return values[localCol];
            }
            if (colIdxs[localCol] > j) break;
        }
        Kokkos::abort("Memory not allocated for CSR matrix component.");
        return values[values.size()]; // compiler warning suppression.
    }
 
    KOKKOS_INLINE_FUNCTION
    ValueType& entry(const IndexType offset) const { return values[offset]; }
 
    View<ValueType> values;  
    View<IndexType> colIdxs; 
    View<IndexType> rowOffs; 
};
 
template<typename ValueType, typename IndexType>
class CSRMatrix
{
 
public:
 
    CSRMatrix(
        const Vector<ValueType>& values,
        const Vector<IndexType>& colIdxs,
        const Vector<IndexType>& rowOffs
    )
        : values_(values), colIdxs_(colIdxs), rowOffs_(rowOffs)
    {
        NF_ASSERT(values.exec() == colIdxs_.exec(), "Executors are not the same");
        NF_ASSERT(values.exec() == rowOffs_.exec(), "Executors are not the same");
    }
 
    CSRMatrix(const Executor exec) : values_(exec, 0), colIdxs_(exec, 0), rowOffs_(exec, 0) {}
 
    ~CSRMatrix() = default;
 
    [[nodiscard]] const Executor& exec() const { return values_.exec(); }
 
    [[nodiscard]] IndexType nRows() const
    {
        return static_cast<IndexType>(rowOffs_.size())
             - static_cast<IndexType>(static_cast<bool>(rowOffs_.size()));
    }
 
    [[nodiscard]] IndexType nNonZeros() const { return static_cast<IndexType>(values_.size()); }
 
    [[nodiscard]] Vector<ValueType>& values() { return values_; }
 
    [[nodiscard]] Vector<IndexType>& colIdxs() { return colIdxs_; }
 
    [[nodiscard]] Vector<IndexType>& rowPtrs() { return rowOffs_; }
 
    [[nodiscard]] const Vector<ValueType>& values() const { return values_; }
 
    [[nodiscard]] const Vector<IndexType>& colIdxs() const { return colIdxs_; }
 
    [[nodiscard]] const Vector<IndexType>& rowPtrs() const { return rowOffs_; }
 
    [[nodiscard]] CSRMatrix<ValueType, IndexType> copyToExecutor(Executor dstExec) const
    {
        if (dstExec == values_.exec())
        {
            return *this;
        }
        CSRMatrix<ValueType, IndexType> other(
            values_.copyToHost(), colIdxs_.copyToHost(), rowOffs_.copyToHost()
        );
        return other;
    }
 
    [[nodiscard]] CSRMatrix<ValueType, IndexType> copyToHost() const
    {
        return copyToExecutor(SerialExecutor());
    }
 
    [[nodiscard]] CSRMatrixView<ValueType, IndexType> view()
    {
        return CSRMatrixView(values_.view(), colIdxs_.view(), rowOffs_.view());
    }
 
    [[nodiscard]] const CSRMatrixView<const ValueType, const IndexType> view() const
    {
        return CSRMatrixView(values_.view(), colIdxs_.view(), rowOffs_.view());
    }
 
private:
 
    Vector<ValueType> values_;  
    Vector<IndexType> colIdxs_; 
    Vector<IndexType> rowOffs_; 
};
 
// /* @brief given a csr matrix this function copies the matrix and converts to requested target
// types
//  *
//  */
// template<typename ValueTypeIn, typename IndexTypeIn, typename ValueTypeOut, typename
// IndexTypeOut> la::CSRMatrix<ValueTypeOut, IndexTypeOut> convert(const Executor exec, const
// la::CSRMatrixView<const ValueTypeIn, const IndexTypeIn> in)
// {
//     Vector<IndexTypeOut> colIdxsTmp(exec, in.colIdxs.size());
//     Vector<IndexTypeOut> rowPtrsTmp(exec, in.rowOffs.size());
//     Vector<ValueTypeOut> valuesTmp(exec, in.values.data(), in.values.size());
 
//     parallelFor(
//         colIdxsTmp, KOKKOS_LAMBDA(const localIdx i) { return IndexTypeOut {in.colIdxs[i]}; }
//     );
//     parallelFor(
//         rowPtrsTmp, KOKKOS_LAMBDA(const localIdx i) { return IndexTypeOut {in.rowOffs[i]}; }
//     );
//     parallelFor(
//         valuesTmp, KOKKOS_LAMBDA(const localIdx i) { return ValueTypeOut {in.values[i]}; }
//     );
 
//     return la::CSRMatrix<ValueTypeOut, IndexTypeOut> {valuesTmp, colIdxsTmp, rowPtrsTmp};
// }
 
 
} // namespace NeoN