Add nnz per row precount method for mju_sqrMatTDSparse.

PiperOrigin-RevId: 521744893
Change-Id: Ie0fbfab552a680127f4acf9041b07916f8a2a490

src/engine/engine_core_constraint.c

@@ -1801,16 +1801,16 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
     // construct supernodes
     mju_superSparse(nefc, rowsuper, rownnz, rowadr, colind);
 
-    // AR = JM2 * JM2', uncompressed layout
+    // AR = JM2 * JM2'
+    mju_sqrMatTDSparseInit(d->efc_AR_rownnz, d->efc_AR_rowadr, JM2T, JM2,
+                           nv, nefc, rownnzT, rowadrT, colindT, rownnz,
+                           rowadr, colind, rowsuper, d);
+
     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);
 
-    // compress layout of AR
-    mju_compressSparse(d->efc_AR, nefc, nefc,
-                       d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind);
-
     // add R to diagonal of AR
     for (int i=0; i<nefc; i++) {
       for (int j=0; j<d->efc_AR_rownnz[i]; j++) {

src/engine/engine_solver.c

@@ -1362,7 +1362,9 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
 
   // sparse
   if (mj_isSparse(m)) {
-    // compute H = J'*D*J, uncompressed layout
+    // compute H = J'*D*J
+    // TODO(b/266802572): remove uncompressed layout
+    mju_sqrMatTDUncompressedInit(ctx->rowadr, nv);
     mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, D, nefc, nv,
                        ctx->rownnz, ctx->rowadr, ctx->colind,
                        d->efc_J_rownnz, d->efc_J_rowadr,

src/engine/engine_util_sparse.c

@@ -411,8 +411,111 @@ void mju_superSparse(int nr, int* rowsuper,
 }
 
 
+// precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse
+void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr,
+                            const mjtNum* mat, const mjtNum* matT,
+                            int nr, int nc,  const int* rownnz,
+                            const int* rowadr, const int* colind,
+                            const int* rownnzT, const int* rowadrT,
+                            const int* colindT, const int* rowsuperT,
+                            mjData* d) {
+  mjMARKSTACK;
+  int* chain = mj_stackAllocInt(d, 2*nc);
+  int nchain = 0;
+  int* res_colind = NULL;
+
+  for (int r=0; r<nc; r++) {
+
+    // supernode; copy everything to next row
+    if (rowsuperT && r>0 && rowsuperT[r-1]>0) {
+      res_rownnz[r] = res_rownnz[r - 1];
+
+      // fill in upper triangle
+      for (int j=0; j <nchain; j++) {
+        res_rownnz[res_colind[j]]++;
+      }
+
+      // update chain with diagonal
+      if (rownnzT[r]) {
+        res_colind[nchain++] = r;
+        res_rownnz[r]++;
+      }
+    } else {
+      int inew = 0, iold = nc;
+      nchain = 0;
+
+      for (int i=0; i<rownnzT[r]; i++) {
+        int c = colindT[rowadrT[r] + i];
+
+        int adr = inew;
+        inew = iold;
+        iold = adr;
+
+        int nnewchain = 0;
+        adr = 0;
+        int end = rowadr[c] + rownnz[c];
+        for (int adr1=rowadr[c]; adr1<end; adr1++) {
+          int col_mat = colind[adr1];
+          while (adr<nchain && chain[iold + adr] < col_mat &&
+                 chain[iold + adr]<=r) {
+            chain[inew + nnewchain++] = chain[iold + adr++];
+          }
+
+          // skip upper triangle
+          if (col_mat>r) {
+            break;
+          }
+
+          if (adr < nchain && chain[iold + adr] == col_mat) {
+            adr++;
+          }
+          chain[inew + nnewchain++] = col_mat;
+        }
+
+        while (adr<nchain && chain[iold + adr]<=r) {
+          chain[inew + nnewchain++] = chain[iold + adr++];
+        }
+        nchain = nnewchain;
+      }
+
+      // only computed for lower triangle
+      res_rownnz[r] = nchain;
+      res_colind = chain + inew;
+
+      // update upper triangle
+      int nchain_end = nchain;
+
+      // avoid double counting.
+      if (nchain>0 && res_colind[nchain-1]==r) {
+        nchain_end = nchain - 1;
+      }
+
+      for (int j=0; j<nchain_end; j++) {
+        res_rownnz[res_colind[j]]++;
+      }
+    }
+  }
+
+  res_rowadr[0] = 0;
+  for (int r = 1; r < nc; r++) {
+    res_rowadr[r] = res_rowadr[r-1] + res_rownnz[r-1];
+  }
+
+  mjFREESTACK;
+}
+
+
+// precompute res_rowadr for mju_sqrMatTDSparse using uncompressed memory
+void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc) {
+  for (int r=0; r<nc; r++) {
+    res_rowadr[r] = r*nc;
+  }
+}
+
+
 
 // compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout
+// res_rowadr is required to be precomputed
 void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
                         const mjtNum* diag, int nr, int nc,
                         int* res_rownnz, int* res_rowadr, int* res_colind,
@@ -424,11 +527,6 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
   // allocate space for accumulation buffer and matT
   mjMARKSTACK;
 
-  // set uncompressed layout (the following doesn't depend on this layout)
-  for (int r=0; r<nc; r++) {
-    res_rowadr[r] = r*nc;
-  }
-
   // a dense row buffer that stores the current row in the resulting matrix
   mjtNum* buffer = mj_stackAlloc(d, nc);
 

src/engine/engine_util_sparse.h

@@ -70,6 +70,7 @@ 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
 MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
                               const mjtNum* diag, int nr, int nc,
                               int* res_rownnz, int* res_rowadr, int* res_colind,
@@ -79,6 +80,18 @@ MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT
                               const int* colindT, const int* rowsuperT,
                               mjData* d);
 
+// precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse
+MJAPI void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr,
+                                  const mjtNum* mat, const mjtNum* matT,
+                                  int nr, int nc,  const int* rownnz,
+                                  const int* rowadr, const int* colind,
+                                  const int* rownnzT, const int* rowadrT,
+                                  const int* colindT, const int* rowsuperT,
+                                  mjData* d);
+
+// precompute res_rowadr for mju_sqrMatTDSparse using uncompressed memory
+MJAPI void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc);
+
 
 #ifdef __cplusplus
 }