dyn Hessian

Author: Ahdhn

include/rxmesh/context.h

@@ -70,6 +70,25 @@ class Context
         return m_num_vertices;
     }
 
+    /**
+     * @brief Total number of vertices in mesh
+     */
+    template <typename HandleT>
+    __device__ __forceinline__ uint32_t get_num() const
+    {
+        if constexpr (std::is_same_v<HandleT, VertexHandle>) {
+            return m_num_vertices[0];
+        }
+
+        if constexpr (std::is_same_v<HandleT, EdgeHandle>) {
+            return m_num_edges[0];
+        }
+
+        if constexpr (std::is_same_v<HandleT, FaceHandle>) {
+            return m_num_faces[0];
+        }
+    }
+
     /**
      * @brief Total number of patches in mesh
      */

include/rxmesh/diff/dyn_hessian_sparse_matrix.h

@@ -0,0 +1,74 @@
+#pragma once
+
+#include "rxmesh/diff/hessian_sparse_matrix.h"
+
+#include "rxmesh/diff/scalar.h"
+
+namespace rxmesh {
+
+/**
+ * @brief Construct the 'dynamic' sparse Hessian of type T with K variables per
+ * vertex. Dynamic here means the sparsity of the Hessian is going to change
+ during runtime. The matrix size is (V * K) X (V * K) where V is the number of
+ * vertices in the mesh.
+ *
+ */
+template <typename T, int K>
+struct DynamicHessianSparseMatrix : public HessianSparseMatrix<T, k>
+{
+    using Type = T;
+
+    static constexpr int K_ = K;
+
+    using ScalarT = Scalar<T, K, true>;
+
+    using IndexT = typename SparseMatrix<T>::IndexT;
+
+
+    DynamicHessianSparseMatrix() : HessianSparseMatrix<T, k>()
+    {
+    }
+
+    DynamicHessianSparseMatrix(const RXMeshStatic& rx, Op op = Op::VV)
+        : HessianSparseMatrix<T>(rx, op)
+    {
+    }
+
+
+    /**
+     * @brief insert more entries to the hessian matrix. The input here is the
+     * list of vertices that will be interacting. Thus, we need to extend them
+     * by the number of their replicate, i.e., k
+     */
+    __host__ void insert(uint32_t size, IndexT* d_rows, IndexT* d_cols)
+    {
+        // Here, we assume the number of rows and cols is the same and only
+        // the sparsity is changing
+
+        constexpr uint32_t blockThreads = 256;
+
+        // fill the new row_ptr with the data from the mesh connectivity
+        rx.run_kernel<blockThreads>(
+            {Op::VV},
+            detail::sparse_mat_prescan<Op::VV, blockThreads>,
+            m_d_row_ptr,
+            k);
+
+        // fill the new row_ptr with the data from the new entries
+        detail::
+            sparse_mat_prescan<<<DIVIDE_UP(size, blockThreads), blockThreads>>>(
+                m_d_row_ptr, size, d_rows, d_cols, k);
+    }
+
+    // delete the functions that access the matrix using only the VertexHandle
+    // since with the Hessian, we should also have the local index (the index
+    // within the kxk matrix)
+    __device__ __host__ const T& operator()(const VertexHandle& row_v,
+                                            const VertexHandle& col_v) const =
+        delete;
+
+    __device__ __host__ T& operator()(const VertexHandle& row_v,
+                                      const VertexHandle& col_v) = delete;
+};
+
+}  // namespace rxmesh

include/rxmesh/matrix/sparse_matrix_kernels.cuh

@@ -15,6 +15,8 @@ __global__ static void sparse_mat_prescan(const rxmesh::Context context,
                                           IndexT*               row_ptr,
                                           IndexT                replicate)
 {
+    bool is_aos = true;
+
     using namespace rxmesh;
 
     using HandleT = typename InputHandle<op>::type;
@@ -27,10 +29,17 @@ __global__ static void sparse_mat_prescan(const rxmesh::Context context,
         IndexT   size     = iter.size() + 1;
         size *= replicate;
         IndexT offset = context.prefix<HandleT>()[patch_id] + local_id;
-        offset *= replicate;
 
-        for (IndexT i = 0; i < replicate; ++i) {
-            row_ptr[offset + i] = size;
+        if (is_aos) {
+            offset *= replicate;
+            for (IndexT i = 0; i < replicate; ++i) {
+                row_ptr[offset + i] = size;
+            }
+        } else {
+            const uint32_t num_elements = context.get_num<HandleT>();
+            for (IndexT i = 0; i < replicate; ++i) {
+                row_ptr[num_elements * i + offset] = size;
+            }
         }
     };
 
@@ -40,6 +49,23 @@ __global__ static void sparse_mat_prescan(const rxmesh::Context context,
     query.dispatch<op>(block, shrd_alloc, init_lambda);
 }
 
+template <typename IndexT = int>
+__global__ static void sparse_mat_prescan(IndexT*       row_ptr,
+                                          const IndexT  size,
+                                          const IndexT* rows,
+                                          const IndexT* cols,
+                                          const IndexT  replicate)
+{
+    const uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;
+    if (tid < size) {
+        const int row = rows[tid] * replicate;
+
+        for (IndexT i = 0; i < replicate; ++i) {
+            ::atomicAdd(row_ptr + (row + i), replicate);
+        }
+    }
+}
+
 template <Op op, uint32_t blockThreads, typename IndexT = int>
 __global__ static void sparse_mat_col_fill(const rxmesh::Context context,
                                            IndexT*               row_ptr,