Represent only the lower triangle in Newton solver's reduced dof-dof matrix.

PiperOrigin-RevId: 712488529
Change-Id: Iad91c72654376539791d7856765a0d0ac9088251

Author: yuvaltassa

src/engine/engine_core_constraint.c

@@ -2176,12 +2176,12 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
 
     // AR = JM2 * JM2'
     mju_sqrMatTDSparseInit(d->efc_AR_rownnz, d->efc_AR_rowadr, nefc, rownnzT,
-                           rowadrT, colindT, rownnz, rowadr, colind, rowsuper, d);
+                           rowadrT, colindT, rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1);
 
     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);
+                       rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1);
 
     // add R to diagonal of AR
     for (int i=0; i < nefc; i++) {

src/engine/engine_io.c

@@ -943,8 +943,11 @@ static void makeDofDofSparse(const mjModel* m, mjData* d,
     // process below diagonal unless reduced and dof is simple
     if (!reduced || !m->dof_simplenum[i]) {
       while ((j = m->dof_parentid[j]) >= 0) {
+        // both reduced and non-reduced have lower triangle
         rownnz[i]++;
-        rownnz[j]++;
+
+        // only non-reduced has upper triangle
+        if (!reduced) rownnz[j]++;
       }
     }
   }
@@ -969,8 +972,11 @@ static void makeDofDofSparse(const mjModel* m, mjData* d,
         remaining[i]--;
         colind[rowadr[i] + remaining[i]] = j;
 
-        remaining[j]--;
-        colind[rowadr[j] + remaining[j]] = i;
+        // only non-reduced has upper triangle
+        if (!reduced) {
+          remaining[j]--;
+          colind[rowadr[j] + remaining[j]] = i;
+        }
       }
     }
   }
@@ -1152,8 +1158,11 @@ static void copyM2Sparse(const mjModel* m, mjData* d, int* dst, const int* src,
         remaining[i]--;
         dst[rowadr[i] + remaining[i]] = src[adr];
 
-        remaining[j]--;
-        dst[rowadr[j] + remaining[j]] = src[adr];
+        // only non-reduced has upper triangle
+        if (!reduced) {
+          remaining[j]--;
+          dst[rowadr[j] + remaining[j]] = src[adr];
+        }
 
         adr++;
       }
@@ -1172,7 +1181,7 @@ static void copyM2Sparse(const mjModel* m, mjData* d, int* dst, const int* src,
 
 
 
-// integer valued dst[M] = src[D lower], handle different sparsity representations
+// integer valued dst[M] = src[D lower]
 static void copyD2MSparse(const mjModel* m, const mjData* d, int* dst, const int* src) {
   int nv = m->nv;
 
@@ -1197,7 +1206,7 @@ static void copyD2MSparse(const mjModel* m, const mjData* d, int* dst, const int
 
 
 // construct index mappings between M <-> D and M -> C
-static void makeDmap(const mjModel* m, mjData* d) {
+static void makeDofDofmap(const mjModel* m, mjData* d) {
   int nM = m->nM, nC = m->nC, nD = m->nD;
   mj_markStack(d);
 
@@ -1955,7 +1964,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
 
     // make C
     makeDofDofSparse(m, d, d->C_rownnz, d->C_rowadr, d->C_diag, d->C_colind, /*reduced=*/1);
-    makeDmap(m, d);
+    makeDofDofmap(m, d);
   }
 
   // restore pluginstate and plugindata

src/engine/engine_solver.c

@@ -1403,7 +1403,7 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
     mju_sqrMatTDSparseInit(ctx->H_rownnz, ctx->H_rowadr, nv,
                            d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
                            d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
-                           d);
+                           d, /*flg_upper=*/0);
 
     // add nC to Hessian total nonzeros (unavoidable overcounting since H_colind is still unknown)
     ctx->nH = m->nC + ctx->H_rowadr[nv - 1] + ctx->H_rownnz[nv - 1];
@@ -1424,14 +1424,24 @@ 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);
+                       d, /*flg_upper=*/0);
 
     // add mass matrix: H = J'*D*J + C
     mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
                   ctx->C, d->C_rownnz, d->C_rowadr, d->C_colind);
 
+    // transiently compute H'; mju_cholFactorNNZ is memory-contiguous in upper triangle layout
+    mj_markStack(d);
+    int* HT_rownnz = mjSTACKALLOC(d, nv, int);
+    int* HT_rowadr = mjSTACKALLOC(d, nv, int);
+    int* HT_colind = mjSTACKALLOC(d, ctx->nH, int);
+    mju_transposeSparse(NULL, NULL, nv, nv,
+                        HT_rownnz, HT_rowadr, HT_colind,
+                        ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind);
+
     // count total and row non-zeros of reverse-Cholesky factor L
-    ctx->nL = mju_cholFactorNNZ(ctx->L_rownnz, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv, d);
+    ctx->nL = mju_cholFactorNNZ(ctx->L_rownnz, HT_rownnz, HT_rowadr, HT_colind, nv, d);
+    mj_freeStack(d);
 
     // compute L row adresses: rowadr = cumsum(rownnz)
     ctx->L_rowadr[0] = 0;
@@ -1508,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);
+                        d, /*flg_upper=*/0);
 
       // 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

@@ -628,7 +628,7 @@ void mju_superSparse(int nr, int* rowsuper,
 void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
                             const int* rownnz, const int* rowadr, const int* colind,
                             const int* rownnzT, const int* rowadrT, const int* colindT,
-                            const int* rowsuperT, mjData* d) {
+                            const int* rowsuperT, mjData* d, int flg_upper) {
   mj_markStack(d);
   int* chain = mjSTACKALLOC(d, 2*nr, int);
   int nchain = 0;
@@ -640,8 +640,10 @@ void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
       res_rownnz[r] = res_rownnz[r - 1];
 
       // fill in upper triangle
-      for (int j=0; j < nchain; j++) {
-        res_rownnz[res_colind[j]]++;
+      if (flg_upper) {
+        for (int j=0; j < nchain; j++) {
+          res_rownnz[res_colind[j]]++;
+        }
       }
 
       // update chain with diagonal
@@ -691,15 +693,17 @@ void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
       res_colind = chain + inew;
 
       // update upper triangle
-      int nchain_end = nchain;
+      if (flg_upper) {
+        int nchain_end = nchain;
 
-      // avoid double counting.
-      if (nchain > 0 && res_colind[nchain-1] == r) {
-        nchain_end = nchain - 1;
-      }
+        // 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]]++;
+        for (int j=0; j < nchain_end; j++) {
+          res_rownnz[res_colind[j]]++;
+        }
       }
     }
   }
@@ -731,7 +735,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) {
+                        mjData* d, int flg_upper) {
   // allocate space for accumulation buffer and matT
   mj_markStack(d);
 
@@ -846,13 +850,15 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
 
 
   // fill upper triangle
-  for (int i=0; i < nc; i++) {
-    int start = res_rowadr[i];
-    int end = start + res_rownnz[i] - 1;
-    for (int j=start; j < end; j++) {
-      int adr = res_rowadr[res_colind[j]] + res_rownnz[res_colind[j]]++;
-      res[adr] = res[j];
-      res_colind[adr] = i;
+  if (flg_upper) {
+    for (int i=0; i < nc; i++) {
+      int start = res_rowadr[i];
+      int end = start + res_rownnz[i] - 1;
+      for (int j=start; j < end; j++) {
+        int adr = res_rowadr[res_colind[j]] + res_rownnz[res_colind[j]]++;
+        res[adr] = res[j];
+        res_colind[adr] = i;
+      }
     }
   }
 

src/engine/engine_util_sparse.h

@@ -96,13 +96,13 @@ 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);
+                              mjData* d, int flg_upper);
 
 // precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse
 MJAPI void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
                                   const int* rownnz, const int* rowadr, const int* colind,
                                   const int* rownnzT, const int* rowadrT, const int* colindT,
-                                  const int* rowsuperT, mjData* d);
+                                  const int* rowsuperT, mjData* d, int flg_upper);
 
 // precompute res_rowadr for mju_sqrMatTDSparse using uncompressed memory
 MJAPI void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc);

src/user/user_model.cc

@@ -2495,7 +2495,7 @@ void mjCModel::CopyTree(mjModel* m) {
       }
     }
   }
-  m->nC = nC = 2 * nOD + nv;
+  m->nC = nC = nOD + nv;
 }
 
 // copy plugin data

test/benchmark/engine_util_sparse_benchmark_test.cc

@@ -43,7 +43,7 @@ void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline(
     int nr, int nc, int* res_rownnz, int* res_rowadr, int* res_colind,
     const int* rownnz, const int* rowadr, const int* colind,
     const int* rowsuper, const int* rownnzT, const int* rowadrT,
-    const int* colindT, const int* rowsuperT, mjData* d) {
+    const int* colindT, const int* rowsuperT, mjData* d, int unused) {
   mj_markStack(d);
   int* chain = mj_stackAllocInt(d, 2 * nc);
   mjtNum* buffer = mj_stackAllocNum(d, nc);
@@ -453,7 +453,8 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
                      d->efc_J_rownnz, d->efc_J_rowadr,
                      d->efc_J_colind, d->efc_J_rowsuper,
                      d->efc_JT_rownnz, d->efc_JT_rowadr,
-                     d->efc_JT_colind, d->efc_JT_rowsuper, d);
+                     d->efc_JT_colind, d->efc_JT_rowsuper, d,
+                     /*flg_upper=*/1);
 
   // compute H = M + J'*D*J
   mj_addM(m, d, H, rownnz, rowadr, colind);
@@ -578,7 +579,7 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
       func(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, 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);
+         d->efc_JT_rowsuper, d, /*flg_upper=*/1);
     }
   } else {
     for (auto s : state) {
@@ -587,10 +588,11 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
                       d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
 
       // compute H = J'*D*J, uncompressed layout
-      mju_sqrMatTDSparse_baseline(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
-                                  d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, d->efc_J_rowsuper,
-                                  d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
-                                  d->efc_JT_rowsuper, d);
+      mju_sqrMatTDSparse_baseline(
+          H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
+          d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, d->efc_J_rowsuper,
+          d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
+          d->efc_JT_rowsuper, d, /*unused=*/0);
     }
   }
 

test/engine/engine_core_smooth_test.cc

@@ -481,8 +481,12 @@ TEST_F(CoreSmoothTest, SolveLDs) {
   vector<mjtNum> LDdense2(nv*nv);
   mj_fullM(m, LDdense2.data(), d->qLD);
 
-  // expect dense matrices to match exactly
-  for (int i=0; i < nv*nv; i++) EXPECT_EQ(LDdense[i], LDdense2[i]);
+  // expect lower triangles to match exactly
+  for (int i=0; i < nv; i++) {
+    for (int j=0; j < i; j++) {
+      EXPECT_EQ(LDdense[i*nv+j], LDdense2[i*nv+j]);
+    }
+  }
 
   // compare LD and LDs vector solve
   vector<mjtNum> vec(nv);

test/engine/engine_solver_test.cc

@@ -61,8 +61,9 @@ static const char* const kIlslandEfcPath =
 // compare accelerations produced by CG solver with and without islands
 TEST_F(SolverTest, IslandsEquivalent) {
   const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
-  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
-  ASSERT_THAT(model, NotNull());
+  char error[1024];
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
   model->opt.solver = mjSOL_CG;                 // use CG solver
   model->opt.tolerance = 0;                     // set tolerance to 0
   model->opt.enableflags &= ~mjENBL_ISLAND;     // disable islands

test/engine/engine_support_test.cc

@@ -725,8 +725,12 @@ TEST_F(InertiaTest, DenseSameAsSparse) {
   // dense addM
   mj_addM(m, d, dst_dense.data(), nullptr, nullptr, nullptr);
 
-  // dense comparison, should be same matrix
-  EXPECT_THAT(dst_dense, ElementsAreArray(dst_sparse));
+  // dense comparison, lower triangle should match
+  for (int i=0; i < nv; i++) {
+    for (int j=0; j < i; j++) {
+      EXPECT_EQ(dst_dense[i*nv+j], dst_sparse[i*nv+j]);
+    }
+  }
 
   // clean up
   mj_deleteData(d);