Add a member function to BsrMatrix to convert to Crs (#2809)

jgfouca · web-flow · commit ac991afa713c · 2025-10-08T14:00:17.000-06:00
* Add a member function to BsrMatrix to conver to Crs

Signed-off-by: James Foucar &lt;jgfouca@sandia.gov&gt;

* Update docs

Signed-off-by: James Foucar &lt;jgfouca@sandia.gov&gt;

* Fix doc

Signed-off-by: James Foucar &lt;jgfouca@sandia.gov&gt;

* Fix mistakes in bsr test

Signed-off-by: James Foucar &lt;jgfouca@sandia.gov&gt;

* Fix veclen max

Signed-off-by: James Foucar &lt;jgfouca@sandia.gov&gt;

---------

Signed-off-by: James Foucar &lt;jgfouca@sandia.gov&gt;
diff --git a/docs/source/API/sparse/bsr_matrix.rst b/docs/source/API/sparse/bsr_matrix.rst
@@ -321,10 +321,25 @@ unmanaged_block_const
 Return a const view of the i-th block in the matrix.
 
 
+convertToCrs
+^^^^^^^^^^^^
+
+.. code:: cppkokkos
+
+  template <typename CrsMatrixType = KokkosSparse::CrsMatrix<ScalarType, OrdinalType, Device, MemoryTraits, SizeType>>
+  CrsMatrixType convertToCrs() const;
+
+Convert the Bsr into a CrsMatrix
+
+The default return type will be a CrsMatrix with all the same template arguments
+as this Bsr, but you can provide your own type if needed. The only requirement
+is that the execution spaces match.
+
+This is a host function.
+
 Example
 =======
 
 .. literalinclude:: ../../../../example/wiki/sparse/KokkosSparse_wiki_bsrmatrix_2.cpp
   :language: c++
   :lines: 16-
-
diff --git a/sparse/src/KokkosSparse_BsrMatrix.hpp b/sparse/src/KokkosSparse_BsrMatrix.hpp
@@ -912,6 +912,96 @@ class BsrMatrix {
     return const_block_type(&values(i * blockDim_ * blockDim_), blockDim_, blockDim_);
   }
 
+  /// \brief Convert the Bsr into a CrsMatrix
+  ///
+  /// The default return type will be a CrsMatrix with all the same template arguments
+  /// as this Bsr, but you can provide your own type if needed. The only requirement
+  /// is that the execution spaces match.
+  ///
+  /// This is a host function.
+  ///
+  template <typename CrsMatrixType = KokkosSparse::CrsMatrix<ScalarType, OrdinalType, Device, MemoryTraits, SizeType>>
+  CrsMatrixType convertToCrs() const {
+    using crs_size_t      = typename CrsMatrixType::size_type;
+    using crs_rowmap_t    = typename CrsMatrixType::row_map_type::non_const_type;
+    using crs_entries_t   = typename CrsMatrixType::index_type::non_const_type;
+    using crs_values_t    = typename CrsMatrixType::values_type::non_const_type;
+    using crs_exe_space_t = typename CrsMatrixType::execution_space;
+    using policy_t        = typename Kokkos::TeamPolicy<execution_space>;
+    using member_t        = typename policy_t::member_type;
+
+    static_assert(std::is_same_v<crs_exe_space_t, execution_space>,
+                  "We do not currently support converting a BsrMatrix to a CsrMatrix with a different execution space");
+
+    // Get size/dimension info from this Bsr. We will use crs_size_t for all int types
+    const crs_size_t blockDim        = this->blockDim();
+    const crs_size_t blockSize       = blockDim * blockDim;
+    const crs_size_t numBlockRows    = numRows();
+    const crs_size_t numBlockCols    = numCols();
+    const crs_size_t numBlockEntries = nnz();
+
+    // Get graph and values from this Bsr
+    const auto blockRowMap  = graph.row_map;
+    const auto blockEntries = graph.entries;
+    const auto blockValues  = values;
+
+    // Compute Csr row/col/entry sizes by multiplying Bsr sizes by block dimension
+    const crs_size_t numCrsRows    = numBlockRows * blockDim;
+    const crs_size_t numCrsCols    = numBlockCols * blockDim;
+    const crs_size_t numCrsEntries = numBlockEntries * blockSize;
+
+    // Allocate CrsMatrix views
+    crs_rowmap_t crsRowMap("crsRowMap", numCrsRows + 1);
+    crs_entries_t crsEntries("crsEntries", numCrsEntries);
+    crs_values_t crsValues("crsValues", numCrsEntries);
+
+    // Create the policy, we have 3 levels of parallelism available in the algorithm
+    const crs_size_t maxvec = policy_t::vector_length_max();
+    const crs_size_t veclen = (blockDim <= maxvec) ? blockDim : maxvec;
+    policy_t policy(numBlockRows, Kokkos::AUTO(), veclen);
+
+    // Fill CrsMatrix row map, entries, and values
+    Kokkos::parallel_for(
+        "ConvertBsrToCrs", policy, KOKKOS_LAMBDA(const member_t& team) {
+          const crs_size_t blockRow      = team.league_rank();
+          const crs_size_t blockRowStart = blockRowMap(blockRow);
+          const crs_size_t blockRowEnd   = blockRowMap(blockRow + 1);
+          const crs_size_t blockRowCount = blockRowEnd - blockRowStart;
+
+          // Iterate over block entries in this row.
+          Kokkos::parallel_for(
+              Kokkos::TeamThreadRange(team, blockRowStart, blockRowEnd), [&](const crs_size_t& blockNnz) {
+                const crs_size_t blockCol = blockEntries(blockNnz);
+                const crs_size_t blockNum = blockNnz - blockRowStart;
+
+                // Iterate over block dim to get the unblocked rows
+                Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, blockDim), [&](const crs_size_t& blockRowOffset) {
+                  const crs_size_t crsRow = blockRow * blockDim + blockRowOffset;
+                  // Each unblocked row has blockRowCount * blockDim items
+                  const crs_size_t crsRowStart = blockRowStart * blockSize + blockRowCount * blockDim * blockRowOffset;
+                  crsRowMap(crsRow)            = crsRowStart;
+
+                  // Iterate over block dim to get the unblocked cols
+                  for (crs_size_t blockColOffset = 0; blockColOffset < blockDim; ++blockColOffset) {
+                    const crs_size_t crsCol = blockCol * blockDim + blockColOffset;
+                    const crs_size_t crsNnz = crsRowStart + blockNum * blockDim + blockColOffset;
+                    crsEntries(crsNnz)      = crsCol;
+                    crsValues(crsNnz) = blockValues(blockNnz * blockSize + blockRowOffset * blockDim + blockColOffset);
+                  }
+                });
+              });
+
+          // Finalize CrsMatrix row map
+          if (blockRow == numBlockRows - 1) {
+            crsRowMap(numCrsRows) = blockRowMap(numBlockRows) * blockSize;
+          }
+        });
+
+    // Construct CrsMatrix
+    return CrsMatrixType("convertedFromBsrMatrix", numCrsRows, numCrsCols, crsEntries.extent(0), crsValues, crsRowMap,
+                         crsEntries);
+  }
+
  protected:
   enum class valueOperation { ADD, ASSIGN };
 
diff --git a/sparse/unit_test/Test_Sparse_BsrMatrix.hpp b/sparse/unit_test/Test_Sparse_BsrMatrix.hpp
@@ -189,16 +189,14 @@ struct TestFunctor {
             auto row_ptr = iblockrow.local_row_in_block(blk, lrow);
             for (auto lcol = 0; lcol < A.blockDim(); ++lcol) {
               auto entry = iblockrow.local_block_value(blk, lrow, lcol);
-              // std::cout << "check0: " << ( entry == row_ptr[lcol] );
-              // std::cout << "check1: " << ( entry == view_blk(lrow,lcol) );
-              check0 = check0 && (entry == row_ptr[lcol]);
-              check1 = check1 && (entry == view_blk(lrow, lcol));
+              check0 &= (entry == row_ptr[lcol]);
+              check1 &= (entry == view_blk(lrow, lcol));
             }  // end local col in row
           }    // end local row in blk
         }      // end blk
       }
-      d_results(0) = check0;
-      d_results(1) = check1;
+      d_results(0) &= check0;
+      d_results(1) &= check1;
 
       // Test BsrRowViewConst
       {
@@ -210,16 +208,14 @@ struct TestFunctor {
             auto row_ptr = iblockrow.local_row_in_block(blk, lrow);
             for (auto lcol = 0; lcol < A.blockDim(); ++lcol) {
               auto entry = iblockrow.local_block_value(blk, lrow, lcol);
-              check2     = check2 && (entry == row_ptr[lcol]);
-              check3     = check3 && (entry == view_blk(lrow, lcol));
+              check2 &= (entry == row_ptr[lcol]);
+              check3 &= (entry == view_blk(lrow, lcol));
             }  // end local col in row
           }    // end local row in blk
         }      // end blk
       }
-      d_results(0) = check0;
-      d_results(1) = check1;
-      d_results(2) = check2;
-      d_results(3) = check3;
+      d_results(2) &= check2;
+      d_results(3) &= check3;
     }  // end for blk rows
 
     // Test sumIntoValues
@@ -243,14 +239,14 @@ struct TestFunctor {
         auto row_ptr = iblockrow.local_row_in_block(relBlk, lrow);
         for (auto lcol = 0; lcol < A.blockDim(); ++lcol) {
           auto entry = iblockrow.local_block_value(relBlk, lrow, lcol);
-          check0     = check0 && (entry == row_ptr[lcol]);
-          check1     = check1 && (entry == view_blk(lrow, lcol));
-          check2     = check2 && (entry == result[lrow * A.blockDim() + lcol]);
+          check0 &= (entry == row_ptr[lcol]);
+          check1 &= (entry == view_blk(lrow, lcol));
+          check2 &= (entry == result[lrow * A.blockDim() + lcol]);
         }  // end local col in row
       }    // end local row in blk
-      d_results(4) = check0;
-      d_results(5) = check1;
-      d_results(6) = check2;
+      d_results(4) &= check0;
+      d_results(5) &= check1;
+      d_results(6) &= check2;
     }
 
     // Test replaceValues
@@ -273,19 +269,79 @@ struct TestFunctor {
         auto row_ptr = iblockrow.local_row_in_block(relBlk, lrow);
         for (auto lcol = 0; lcol < A.blockDim(); ++lcol) {
           auto entry = iblockrow.local_block_value(relBlk, lrow, lcol);
-          check0     = check0 && (entry == row_ptr[lcol]);
-          check1     = check1 && (entry == view_blk(lrow, lcol));
-          check2     = check2 && (entry == valsreplace[lrow * A.blockDim() + lcol]);
+          check0 &= (entry == row_ptr[lcol]);
+          check1 &= (entry == view_blk(lrow, lcol));
+          check2 &= (entry == valsreplace[lrow * A.blockDim() + lcol]);
         }  // end local col in row
       }    // end local row in blk
-      d_results(7) = check0;
-      d_results(8) = check1;
-      d_results(9) = check2;
+      d_results(7) &= check0;
+      d_results(8) &= check1;
+      d_results(9) &= check2;
     }
-
   }  // end operator()(i)
 };   // end TestFunctor
 
+template <class BsrMatrixType, class ResultsType>
+struct TestConvFunctor {
+  typedef typename BsrMatrixType::value_type scalar_t;
+  typedef typename BsrMatrixType::ordinal_type lno_t;
+  typedef typename BsrMatrixType::size_type size_type;
+
+  // Members
+  BsrMatrixType A_;
+  BsrMatrixType A_bsr_conv_;
+  ResultsType d_results_;
+
+  // Constructor
+  TestConvFunctor(BsrMatrixType &A, ResultsType &d_results) : A_(A), d_results_(d_results) {
+    auto A_crs  = A_.convertToCrs();
+    A_bsr_conv_ = BsrMatrixType(A_crs, A_.blockDim());
+  }
+
+  KOKKOS_INLINE_FUNCTION
+  void operator()(const int /*rid*/) const {
+    // Test 1: Converting to Crs and then back to Bsr should give us an indentical bsr
+    bool check0    = true;
+    bool check1    = true;
+    bool check2    = true;
+    bool check3    = true;
+    bool check4    = true;
+    bool check5    = true;
+    bool check6    = true;
+    int result_idx = 0;
+
+    check0 = A_bsr_conv_.numRows() == A_.numRows();
+    check1 = A_bsr_conv_.numCols() == A_.numCols();
+    check2 = A_bsr_conv_.blockDim() == A_.blockDim();
+    check3 = A_bsr_conv_.nnz() == A_.nnz();
+
+    const auto blockRowMap1  = A_.graph.row_map;
+    const auto blockEntries1 = A_.graph.entries;
+    const auto blockValues1  = A_.values;
+
+    const auto blockRowMap2  = A_.graph.row_map;
+    const auto blockEntries2 = A_.graph.entries;
+    const auto blockValues2  = A_.values;
+
+    for (lno_t i = 0; i < A_.numRows() + 1; ++i) {
+      check4 &= (blockRowMap1(i) == blockRowMap2(i));
+    }
+
+    for (size_type i = 0; i < A_.nnz(); ++i) {
+      check5 &= (blockEntries1(i) == blockEntries2(i));
+      check6 &= (blockValues1(i) == blockValues2(i));
+    }
+
+    d_results_(result_idx++) = check0;
+    d_results_(result_idx++) = check1;
+    d_results_(result_idx++) = check2;
+    d_results_(result_idx++) = check3;
+    d_results_(result_idx++) = check4;
+    d_results_(result_idx++) = check5;
+    d_results_(result_idx++) = check6;
+  }  // end operator()(i)
+};   // end TestConvFunctor
+
 }  // namespace Test_Bsr
 
 // Create a CrsMatrix and BsrMatrix and test member functions.
@@ -299,13 +355,26 @@ void testBsrMatrix() {
   crs_matrix_type crsA = makeCrsMatrix_BlockStructure<crs_matrix_type>();
   bsr_matrix_type A    = makeBsrMatrix<bsr_matrix_type>();
 
-  const int num_entries = 10;
+  static constexpr int num_entries = 10;
   typedef Kokkos::View<bool[num_entries], device> result_view_type;
   result_view_type d_results("d_results");
+  Kokkos::deep_copy(d_results, true);
   auto h_results = Kokkos::create_mirror_view(d_results);
+  Test_Bsr::TestFunctor<bsr_matrix_type, result_view_type> functor(A, d_results);
+
+  Kokkos::parallel_for("KokkosSparse::Test_Bsr::BsrMatrix", Kokkos::RangePolicy<typename device::execution_space>(0, 1),
+                       functor);
+
+  Kokkos::deep_copy(h_results, d_results);
+
+  for (decltype(h_results.extent(0)) i = 0; i < h_results.extent(0); ++i) {
+    EXPECT_EQ(h_results[i], true);
+  }
 
+  Kokkos::deep_copy(d_results, true);
+  Test_Bsr::TestConvFunctor<bsr_matrix_type, result_view_type> conv_functor(A, d_results);
   Kokkos::parallel_for("KokkosSparse::Test_Bsr::BsrMatrix", Kokkos::RangePolicy<typename device::execution_space>(0, 1),
-                       Test_Bsr::TestFunctor<bsr_matrix_type, result_view_type>(A, d_results));
+                       conv_functor);
 
   Kokkos::deep_copy(h_results, d_results);
 
