adding optional access to the contact during line search

Author: Ahdhn

apps/NeoHookean/barrier_energy.h

@@ -81,8 +81,6 @@ void add_contact(RXMeshStatic&      rx,
         T d = (xi - x_dbc).dot(normal);
 
         if (d < dhat) {
-            // if (vh.local_id() == 20 || vh.local_id() == 21 || vh.local_id()
-            // == 22 || vh.local_id() == 23 || vh.local_id() == 24) {
             bool inserted = contact_pairs.insert(vh, dbc_vertex);
             assert(inserted);
 
@@ -91,19 +89,8 @@ void add_contact(RXMeshStatic&      rx,
 
             inserted = contact_pairs.insert(vh, dbc_vertex2);
             assert(inserted);
-
-            printf("\n addec contact pair between %d, (%d, %d, %d)",
-                   vh.local_id(),
-                   dbc_vertex.local_id(),
-                   dbc_vertex1.local_id(),
-                   dbc_vertex2.local_id());
         }
     });
-
-    {
-        CUDA_ERROR(cudaGetLastError());
-        CUDA_ERROR(cudaDeviceSynchronize());
-    }
 }
 
 template <typename ProblemT,

apps/NeoHookean/neo_hookean.cu

@@ -198,11 +198,6 @@ void neo_hookean(RXMeshStatic& rx, T dx)
     timer.add("LinearSolver");
     timer.add("Diff");
 
-    //add_contact(
-    //    rx, problem.vv_pairs, v_dbc[0], v_dbc[1], v_dbc[2], is_dbc, x, dhat);
-    //problem.update_hessian();
-    //printf("\n");
-
     auto step_forward = [&]() {
         // x_tilde = x + v*h
         timer.start("Step");
@@ -241,10 +236,6 @@ void neo_hookean(RXMeshStatic& rx, T dx)
                     dhat);
         problem.update_hessian();
         problem.eval_terms();
-        {
-            CUDA_ERROR(cudaGetLastError());
-            CUDA_ERROR(cudaDeviceSynchronize());
-        }
 
         grad.copy_from(problem.grad, DEVICE, DEVICE);
         // DBC satisfied
@@ -297,8 +288,17 @@ void neo_hookean(RXMeshStatic& rx, T dx)
             line_search_init_step = std::min(nh_step, bar_step);
 
             // TODO: line search should pass the step to the friction energy
-            bool ls_success =
-                newton_solver.line_search(line_search_init_step, 0.5, 64, 0.0);
+            bool ls_success = newton_solver.line_search(
+                line_search_init_step, 0.5, 64, 0.0, [&](auto temp_x) {
+                    add_contact(rx,
+                                problem.vv_pairs,
+                                v_dbc[0],
+                                v_dbc[1],
+                                v_dbc[2],
+                                is_dbc,
+                                temp_x,
+                                dhat);
+                });
 
             if (!ls_success) {
                 RXMESH_WARN("Line search failed!");
@@ -315,10 +315,6 @@ void neo_hookean(RXMeshStatic& rx, T dx)
                         dhat);
             problem.update_hessian();
             problem.eval_terms();
-            {
-                CUDA_ERROR(cudaGetLastError());
-                CUDA_ERROR(cudaDeviceSynchronize());
-            }
 
             T f = problem.get_current_loss();
 

include/rxmesh/diff/candidate_pairs.h

@@ -73,18 +73,12 @@ struct CandidatePairs
         /// if we have a space to insert new pairs
 
         auto add_candidate_with_indices = [&](int id) {
-            // add the handles
-            m_pairs_handle(id).first  = c0;
-            m_pairs_handle(id).second = c1;
-
             // add the indices
-            id *= m_variable_dim * m_variable_dim * 2;
-
             for (int i = 0; i < m_variable_dim; ++i) {
                 for (int j = 0; j < m_variable_dim; ++j) {
 
-                    int r_id = m_context.linear_id(c0) * m_variable_dim + i;
-                    int c_id = m_context.linear_id(c1) * m_variable_dim + j;
+                    int r_id = m_hess.get_row_id(c0) * m_variable_dim + i;
+                    int c_id = m_hess.get_row_id(c1) * m_variable_dim + j;
 
                     m_pairs_id(id, 0) = r_id;
                     m_pairs_id(id, 1) = c_id;
@@ -107,13 +101,13 @@ struct CandidatePairs
 
 #ifdef __CUDA_ARCH__
         int id = ::atomicAdd(m_current_num_pairs.data(DEVICE), 1);
-        if (id < m_pairs_handle.rows()) {            
-            if (m_hess.is_non_zero(c0, c1)) {
+        if (id < m_pairs_handle.rows()) {
+            add_candidate(id);
+            if (!m_hess.is_non_zero(c0, c1)) {
                 add_candidate(id);
-            } else {
-                ::atomicAdd(m_current_num_index.data(DEVICE),
-                            m_variable_dim * m_variable_dim * 2);
-                add_candidate_with_indices(id);
+                int idd = ::atomicAdd(m_current_num_index.data(DEVICE),
+                                      m_variable_dim * m_variable_dim * 2);
+                add_candidate_with_indices(idd);
             }
             return true;
         } else {
@@ -123,12 +117,12 @@ struct CandidatePairs
 #else
         if (m_current_num_pairs(0) < m_pairs_handle.rows()) {
             int id = m_current_num_pairs(0);
-            m_current_num_pairs(0)++;            
-            if (m_hess.is_non_zero(c0, c1)) {
-                add_candidate(id);
-            } else {
+            m_current_num_pairs(0)++;
+            add_candidate(id);
+            if (!m_hess.is_non_zero(c0, c1)) {
+                int idd = m_current_num_index(0);
                 m_current_num_index(0) += m_variable_dim * m_variable_dim * 2;
-                add_candidate_with_indices(id);
+                add_candidate_with_indices(idd);
             }
             return true;
         } else {

include/rxmesh/diff/diff_query_kernel.cuh

@@ -399,6 +399,8 @@ __global__ static void diff_kernel_active_pair(
 
         std::pair<HandleT0, HandleT1> pair = pairs.get_pair(id);
 
+        assert(hess.is_non_zero(pair.first, pair.second));
+
         PairIterator<HandleT0> iter(pair.first, pair.second);
 
         ScalarT res = user_func(diff_handle, iter, objective);

include/rxmesh/diff/hessian_sparse_matrix.h

@@ -46,6 +46,8 @@ struct HessianSparseMatrix : public SparseMatrix<T>
                                             const IndexT        local_i,
                                             const IndexT        local_j) const
     {
+        assert(is_non_zero(row_v, col_v));
+
         const IndexT r_id =
             this->get_row_id(row_v) * this->m_replicate + local_i;
         const IndexT c_id =
@@ -63,6 +65,8 @@ struct HessianSparseMatrix : public SparseMatrix<T>
                                       const IndexT        local_i,
                                       const IndexT        local_j)
     {
+        assert(is_non_zero(row_v, col_v));
+
         const IndexT r_id =
             this->get_row_id(row_v) * this->m_replicate + local_i;
         const IndexT c_id =
@@ -81,6 +85,8 @@ struct HessianSparseMatrix : public SparseMatrix<T>
         const IndexT        local_i,
         const IndexT        local_j) const
     {
+        assert(is_non_zero(row_v, col_v));
+
         const IndexT r_id =
             this->get_row_id(row_v) * this->m_replicate + local_i;
         const IndexT c_id =

include/rxmesh/diff/lbfgs_solver.h

@@ -49,12 +49,7 @@ struct LBFGSSolver
             y_list[i].reset(0, LOCATION_ALL);
         }
     }
-
-    inline void solve(cudaStream_t stream = NULL)
-    {
-        compute_direction(stream);
-        line_search(stream);
-    }
+     
 
     inline void compute_direction(cudaStream_t stream = NULL)
     {

include/rxmesh/diff/newton_solver.h

@@ -50,15 +50,7 @@ struct NetwtonSolver
                 };
         }
     }
-
-    /**
-     * @brief
-     */
-    inline void solve(cudaStream_t stream = NULL)
-    {
-        compute_direction(stream);
-        line_search(stream);
-    }
+ 
 
     /**
      * @brief solve to get Newton direction
@@ -140,10 +132,13 @@ struct NetwtonSolver
     /**
      * @brief line search
      */
+    using Callback = std::function<void(Attribute<T, ObjHandleT>)>;
+
     inline bool line_search(const T      s_max        = 1.0,
                             const T      shrink       = 0.8,
                             const int    max_iters    = 64,
                             const T      armijo_const = 1e-4,
+                            Callback     cb           = {},
                             cudaStream_t stream       = NULL)
     {
         // we are going to keep trying to update temp_objective until we reach
@@ -183,6 +178,9 @@ struct NetwtonSolver
 
 
             // eval new obj func
+            if (cb) {
+                cb(*temp_objective);
+            }
             problem.eval_terms_passive(temp_objective.get(), stream);
 
             // get the new value of the objective function

include/rxmesh/matrix/sparse_matrix.h

@@ -602,7 +602,7 @@ struct SparseMatrix
      * @brief check if the input entry (x,y) is a non-zero (i.e., if it
      * allocated in the CSR representation)
      */
-    __device__ __host__ bool is_non_zero(const IndexT x, IndexT y) const
+    __device__ __host__ bool is_non_zero(const IndexT x, const IndexT y) const
     {
         const IndexT start = row_ptr()[x];
         const IndexT end   = row_ptr()[x + 1];