Compute diagonal indices in mj_sqrMatTDSparse

PiperOrigin-RevId: 713714087
Change-Id: Icc8eae74e6e47ba1d12a6e9aa0d774ab006cfe38

src/engine/engine_core_constraint.c

@@ -2174,23 +2174,20 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
     // construct supernodes
     mju_superSparse(nefc, rowsuper, rownnz, rowadr, colind);
 
-    // AR = JM2 * JM2'
+    // pre-count efc_AR_rownnz, efc_AR_rowadr
     mju_sqrMatTDSparseCount(d->efc_AR_rownnz, d->efc_AR_rowadr, nefc, rownnzT,
                             rowadrT, colindT, rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1);
 
+    // AR = JM2 * JM2'
+    int* diagind = mjSTACKALLOC(d, nefc, int);
     mju_sqrMatTDSparse(d->efc_AR, JM2T, JM2, NULL, nv, nefc,
                        d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind,
                        rownnzT, rowadrT, colindT, NULL,
-                       rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1);
+                       rownnz, rowadr, colind, rowsuper, d, diagind);
 
     // add R to diagonal of AR
     for (int i=0; i < nefc; i++) {
-      for (int j=0; j < d->efc_AR_rownnz[i]; j++) {
-        if (i == d->efc_AR_colind[d->efc_AR_rowadr[i]+j]) {
-          d->efc_AR[d->efc_AR_rowadr[i]+j] += d->efc_R[i];
-          break;
-        }
-      }
+      d->efc_AR[diagind[i]] += d->efc_R[i];
     }
   }
 

src/engine/engine_solver.c

@@ -1424,7 +1424,7 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
                        ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
                        d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
                        d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
-                       d, /*flg_upper=*/0);
+                       d, /*diagind=*/NULL);
 
     // add mass matrix: H = J'*D*J + C
     mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
@@ -1518,7 +1518,7 @@ static void FactorizeHessian(const mjModel* m, mjData* d, mjCGContext* ctx,
                         ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
                         d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
                         d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
-                        d, /*flg_upper=*/0);
+                        d, /*diagind=*/NULL);
 
       // add mass matrix: H = J'*D*J + C
       mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,

src/engine/engine_util_sparse.c

@@ -724,7 +724,7 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
                         const int* colind, const int* rowsuper,
                         const int* rownnzT, const int* rowadrT,
                         const int* colindT, const int* rowsuperT,
-                        mjData* d, int flg_upper) {
+                        mjData* d, int* diagind) {
   // allocate space for accumulation buffer and matT
   mj_markStack(d);
 
@@ -838,8 +838,14 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
   }
 
 
-  // fill upper triangle
-  if (flg_upper) {
+  // diagonal indices requested: fill upper triangle
+  if (diagind) {
+    // save diagonal indices
+    for (int i=0; i < nc; i++) {
+      diagind[i] = res_rowadr[i] + res_rownnz[i] - 1;
+    }
+
+    // fill upper triangle
     for (int i=0; i < nc; i++) {
       int start = res_rowadr[i];
       int end = start + res_rownnz[i] - 1;

src/engine/engine_util_sparse.h

@@ -89,8 +89,7 @@ MJAPI void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
 MJAPI void mju_superSparse(int nr, int* rowsuper,
                            const int* rownnz, const int* rowadr, const int* colind);
 
-// compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout
-// res_rowadr is required to be precomputed
+// compute sparse M'*diag*M (diag=NULL: compute M'*M), res_rowadr must be precomputed
 MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
                               const mjtNum* diag, int nr, int nc,
                               int* res_rownnz, const int* res_rowadr, int* res_colind,
@@ -98,7 +97,7 @@ MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT
                               const int* colind, const int* rowsuper,
                               const int* rownnzT, const int* rowadrT,
                               const int* colindT, const int* rowsuperT,
-                              mjData* d, int flg_upper);
+                              mjData* d, int* diagind);
 
 // precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse
 MJAPI void mju_sqrMatTDSparseCount(int* res_rownnz, int* res_rowadr, int nr,