- added CSRMatrixSpan as a helper, removed entry from CSRMatrix.

- working read/write approach.
exasim-project · Feb 26, 2025 · 8652683 · 8652683
1 parent 458d892
commit 8652683
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 30 deletions.
diff --git a/include/NeoFOAM/linearAlgebra/CSRMatrix.hpp b/include/NeoFOAM/linearAlgebra/CSRMatrix.hpp
@@ -14,6 +14,50 @@ enum BlockStructType
     component
 };
 
+template<typename ValueType, typename IndexType>
+class CSRMatrixSpan
+{
+public:
+
+    // Constructor for non-const spans
+    CSRMatrixSpan(
+        const std::span<ValueType>& values,
+        const std::span<IndexType>& colIdxs,
+        const std::span<IndexType>& rowPtrs
+    )
+        : values_(values), colIdxs_(colIdxs), rowPtrs_(rowPtrs) {};
+
+    // Constructor for const spans - needed for const matrix objects
+    CSRMatrixSpan(
+        const std::span<const ValueType>& values,
+        const std::span<const IndexType>& colIdxs,
+        const std::span<const IndexType>& rowPtrs
+    )
+        : values_(values), colIdxs_(colIdxs), rowPtrs_(rowPtrs) {};
+
+    ~CSRMatrixSpan() = default;
+
+    KOKKOS_INLINE_FUNCTION
+    ValueType& entry(const IndexType i, const IndexType j) const
+    {
+        for (IndexType ic = 0; ic < (rowPtrs_[i + 1] - rowPtrs_[i]); ++ic)
+        {
+            if (colIdxs_[rowPtrs_[i] + ic] == j)
+            {
+                return values_[rowPtrs_[i] + ic];
+            }
+            if (colIdxs_[rowPtrs_[i] + ic] > j) break;
+        }
+        Kokkos::abort("Memory not allocated for CSR matrix component.");
+    }
+
+private:
+
+    std::span<ValueType> values_;
+    std::span<IndexType> colIdxs_;
+    std::span<IndexType> rowPtrs_;
+};
+
 template<typename ValueType, typename IndexType>
 class CSRMatrix
 {
@@ -49,7 +93,6 @@ class CSRMatrix
     [[nodiscard]] const std::span<const IndexType> colIdxs() const { return colIdxs_.span(); }
     [[nodiscard]] const std::span<const IndexType> rowPtrs() const { return rowPtrs_.span(); }
 
-
     [[nodiscard]] CSRMatrix<ValueType, IndexType> copyToExecutor(Executor dstExec) const
     {
         if (dstExec == values_.exec()) return *this;
@@ -65,35 +108,14 @@ class CSRMatrix
         return copyToExecutor(SerialExecutor());
     }
 
-    KOKKOS_INLINE_FUNCTION
-    const ValueType& entry(const IndexType i, const IndexType j) const
+    [[nodiscard]] CSRMatrixSpan<ValueType, IndexType> span()
     {
-        for (IndexType ic = 0; ic < (rowPtrs_[i + 1] - rowPtrs_[i]); ++ic)
-        {
-            if (colIdxs_[rowPtrs_[i] + ic] == j)
-            {
-                return values_[rowPtrs_[i] + ic];
-            }
-            if (colIdxs_[rowPtrs_[i] + ic] > j) break;
-        }
-        Kokkos::abort("Memory not allocated for CSR matrix component.");
+        return CSRMatrixSpan(values_.span(), colIdxs_.span(), rowPtrs_.span());
     }
 
-    KOKKOS_INLINE_FUNCTION
-    ValueType& entry(IndexType i, IndexType j)
+    [[nodiscard]] const CSRMatrixSpan<const ValueType, const IndexType> span() const
     {
-        for (IndexType ic = 0; ic < (rowPtrs_[i + 1] - rowPtrs_[i]); ++ic)
-        {
-            if (colIdxs_[rowPtrs_[i] + ic] == j)
-            {
-                return values_[rowPtrs_[i] + ic];
-            }
-            if (colIdxs_[rowPtrs_[i] + ic] > j)
-            {
-                Kokkos::abort("WIP.");
-            };
-        }
-        Kokkos::abort("Memory not allocated for CSR matrix component.");
+        return CSRMatrixSpan(values_.span(), colIdxs_.span(), rowPtrs_.span());
     }
 
 private:

diff --git a/test/linearAlgebra/CSRMatrix.cpp b/test/linearAlgebra/CSRMatrix.cpp
@@ -114,17 +114,21 @@ TEST_CASE("CSRMatrix")
     SECTION("Update existing entry on " + execName)
     {
         // Sparse
+        auto csrSpan = denseMatrix.span();
         parallelFor(
             exec,
             {0, 1},
             KOKKOS_LAMBDA(const size_t i) {
-                sparseMatrix.entry(0, 0) = -1.0;
-                sparseMatrix.entry(1, 1) = -5.0;
+                csrSpan.entry(0, 0) = -1.0;
+                // sparseMatrix.entry(1, 1) = -5.0;
             }
         );
 
-        // auto checkHost = checkSparse.copyToHost();
-        // REQUIRE(checkHost.span()[0] == 1.0);
+
+        REQUIRE(denseMatrix.copyToHost().values()[0] == -1.0);
+
+        // auto checkHost = denseMatrixConst.values.copyToHost();
+        // REQUIRE(denseMatrixConst.values()[0] == -1.0);
         // REQUIRE(checkHost.span()[1] == 5.0);
         // REQUIRE(checkHost.span()[2] == 6.0);
         // REQUIRE(checkHost.span()[3] == 8.0);