generalize the blocks in sparse matrix to non-square blocks

Author: Ahdhn

include/rxmesh/diff/hessian_sparse_matrix.h

@@ -32,7 +32,12 @@ struct HessianSparseMatrix : public SparseMatrix<T>
     HessianSparseMatrix(const RXMeshStatic& rx,
                         const int           extra_nnz_entries,
                         Op                  op = Op::VV)
-        : SparseMatrix<T>(rx, 1.0f, extra_nnz_entries, op, K)
+        : SparseMatrix<T>(rx,
+                          1.0f,
+                          extra_nnz_entries,
+                          op,
+                          detail::BlockDim(K, K),
+                          true)
     {
     }
 
@@ -49,9 +54,9 @@ struct HessianSparseMatrix : public SparseMatrix<T>
         assert(is_non_zero(row_v, col_v));
 
         const IndexT r_id =
-            this->get_row_id(row_v) * this->m_replicate + local_i;
+            this->get_row_id(row_v) * this->m_block_dim.x + local_i;
         const IndexT c_id =
-            this->get_row_id(col_v) * this->m_replicate + local_j;
+            this->get_row_id(col_v) * this->m_block_dim.y + local_j;
 
         return SparseMatrix<T>::operator()(r_id, c_id);
     }
@@ -68,9 +73,9 @@ struct HessianSparseMatrix : public SparseMatrix<T>
         assert(is_non_zero(row_v, col_v));
 
         const IndexT r_id =
-            this->get_row_id(row_v) * this->m_replicate + local_i;
+            this->get_row_id(row_v) * this->m_block_dim.x + local_i;
         const IndexT c_id =
-            this->get_row_id(col_v) * this->m_replicate + local_j;
+            this->get_row_id(col_v) * this->m_block_dim.y + local_j;
 
         return SparseMatrix<T>::operator()(r_id, c_id);
     }
@@ -88,9 +93,9 @@ struct HessianSparseMatrix : public SparseMatrix<T>
         assert(is_non_zero(row_v, col_v));
 
         const IndexT r_id =
-            this->get_row_id(row_v) * this->m_replicate + local_i;
+            this->get_row_id(row_v) * this->m_block_dim.x + local_i;
         const IndexT c_id =
-            this->get_row_id(col_v) * this->m_replicate + local_j;
+            this->get_row_id(col_v) * this->m_block_dim.y + local_j;
 
         return {r_id, c_id};
     }
@@ -102,8 +107,8 @@ struct HessianSparseMatrix : public SparseMatrix<T>
     __device__ __host__ const bool is_non_zero(const VertexHandle& row_v,
                                                const VertexHandle& col_v) const
     {
-        const IndexT r_id = this->get_row_id(row_v) * this->m_replicate;
-        const IndexT c_id = this->get_row_id(col_v) * this->m_replicate;
+        const IndexT r_id = this->get_row_id(row_v) * this->m_block_dim.x;
+        const IndexT c_id = this->get_row_id(col_v) * this->m_block_dim.y;
 
         return SparseMatrix<T>::is_non_zero(r_id, c_id);
     }

include/rxmesh/matrix/block_dim.h

@@ -0,0 +1,22 @@
+#pragma once
+
+namespace rxmesh {
+namespace detail {
+// The dimension of a (dense) block in the sparse matrix. In most cases
+// the block size is just 1x1 which represent a single non-zero value.
+// But in, e.g., Hessians, the blocks could be kxk.
+struct BlockDim
+{
+    int x, y;
+
+    BlockDim() : x(1), y(1)
+    {
+    }
+
+    BlockDim(int x_, int y_) : x(x_), y(y_)
+    {
+    }
+};
+
+}  // namespace detail
+}  // namespace rxmesh