From 9924cce4b75badc7bb3e22d2f2dfd5ab2e349570 Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Tue, 4 Apr 2023 06:20:32 -0700
Subject: [PATCH] Add nnz per row precount method for mju_sqrMatTDSparse.

PiperOrigin-RevId: 521744893
Change-Id: Ie0fbfab552a680127f4acf9041b07916f8a2a490
---
 src/engine/engine_core_constraint.c           |  10 +-
 src/engine/engine_solver.c                    |   4 +-
 src/engine/engine_util_sparse.c               | 108 ++++++++++++-
 src/engine/engine_util_sparse.h               |  13 ++
 .../engine_util_sparse_benchmark_test.cc      |   3 +
 test/engine/engine_util_sparse_test.cc        | 152 ++++++++++++++++--
 6 files changed, 267 insertions(+), 23 deletions(-)

diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c
index d9dc9c2a..e1b2c3ae 100644
--- a/src/engine/engine_core_constraint.c
+++ b/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++) {
diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c
index 6ae66944..eec03aeb 100644
--- a/src/engine/engine_solver.c
+++ b/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,
diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c
index 770f75c3..1ac75006 100644
--- a/src/engine/engine_util_sparse.c
+++ b/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);
 
diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h
index 103ceccb..3da92b8c 100644
--- a/src/engine/engine_util_sparse.h
+++ b/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
 }
diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc
index 048371f1..10628a93 100644
--- a/test/benchmark/engine_util_sparse_benchmark_test.cc
+++ b/test/benchmark/engine_util_sparse_benchmark_test.cc
@@ -448,6 +448,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
   }
 
   // compute H = J'*D*J, uncompressed layout
+  mju_sqrMatTDUncompressedInit(rowadr, m->nv);
   mju_sqrMatTDSparse(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv,
                       rownnz, rowadr, colind,
                       d->efc_J_rownnz, d->efc_J_rowadr,
@@ -572,6 +573,8 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
   // time benchmark
   if (func) {
     for (auto s : state) {
+      mju_sqrMatTDUncompressedInit(rowadr, m->nv);
+
       // compute H = J'*D*J, uncompressed layout
       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,
diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc
index 6c5b1617..0f9594b8 100644
--- a/test/engine/engine_util_sparse_test.cc
+++ b/test/engine/engine_util_sparse_test.cc
@@ -206,6 +206,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) {
   int rownnzH[] = {0, 0, 0};
   int rowadrH[] = {0, 0, 0};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 3, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, NULL, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 3);
   mju_sqrMatTDSparse(matH, mat, matT, NULL, 3, 3, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
                      NULL, data);
@@ -242,6 +251,16 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse2) {
   int rownnzH[] = {0, 0, 0};
   int rowadrH[] = {0, 0, 0};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 3, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, NULL, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
+
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 3);
   mju_sqrMatTDSparse(matH, mat, matT, NULL, 3, 3, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
                      NULL, data);
@@ -280,6 +299,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse3) {
 
   mjtNum diag[] = {2, 3, 4};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 3, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, NULL, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(2, 2, 0));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 2, 4));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 3);
   mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
                      NULL, data);
@@ -318,6 +346,16 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse4) {
 
   mjtNum diag[] = {2, 3, 4};
 
+
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 3, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, NULL, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(2, 0, 2));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 2, 2));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 3);
   mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
                      NULL, data);
@@ -354,6 +392,16 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse5) {
   int rownnzH[] = {0, 0, 0};
   int rowadrH[] = {0, 0, 0};
 
+
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 3, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, NULL, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 3, 5));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 3);
   mju_sqrMatTDSparse(matH, mat, matT, NULL, 3, 3, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
                      NULL, data);
@@ -390,6 +438,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse6) {
   int rownnzH[] = {0, 0, 0};
   int rowadrH[] = {0, 0, 0};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 3, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, NULL, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(2, 1, 2));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 2, 3));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 3);
   mju_sqrMatTDSparse(matH, mat, matT, NULL, 3, 3, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
                      NULL, data);
@@ -428,6 +485,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse7) {
 
   mjtNum diag[] = {2, 3, 4};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 2, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, NULL, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(2, 2));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 2));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 2);
   mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 2, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
                      NULL, data);
@@ -465,6 +531,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse8) {
 
   mjtNum diag[] = {2, 3};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 2, 3, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, NULL, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 3, 5));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 3);
   mju_sqrMatTDSparse(matH, mat, matT, diag, 2, 3, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
                      NULL, data);
@@ -503,6 +578,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse9) {
 
   mjtNum diag[] = {2, 3, 4};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 3, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, NULL, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 3);
   mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
                      NULL, data);
@@ -533,7 +617,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse10) {
   int colindT[] = {0, 1, 2, 0, 1, 2, 0, 1, 2};
   int rownnzT[] = {3, 3, 3};
   int rowadrT[] = {0, 3, 6};
-  int superowT[] = {2, 1, 0};
+  int rowsuperT[] = {2, 1, 0};
 
   mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
   int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
@@ -542,9 +626,18 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse10) {
 
   mjtNum diag[] = {1, 1, 1};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 3, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, rowsuperT, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 3);
   mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
-                     superowT, data);
+                     rowsuperT, data);
 
   EXPECT_THAT(matH, ElementsAre(14, 14, 14, 14, 14, 14, 14, 14, 14));
   EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
@@ -572,7 +665,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse11) {
   int colindT[] = {0, 0, 2, 0, 2};
   int rownnzT[] = {1, 2, 2};
   int rowadrT[] = {0, 1, 3};
-  int superowT[] = {0, 1, 0};
+  int rowsuperT[] = {0, 1, 0};
 
   mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
   int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
@@ -581,9 +674,18 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse11) {
 
   mjtNum diag[] = {1, 1, 1};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 3, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, rowsuperT, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 3);
   mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
-                     superowT, data);
+                     rowsuperT, data);
 
   EXPECT_THAT(matH, ElementsAre(1, 1, 1, 1, 10, 10, 1, 10, 10));
   EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
@@ -611,7 +713,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse12) {
   int colindT[] = {0, 0, 0, 2, 0, 2};
   int rownnzT[] = {1, 1, 2, 2};
   int rowadrT[] = {0, 1, 2, 4};
-  int superowT[] = {1, 0, 1, 0};
+  int rowsuperT[] = {1, 0, 1, 0};
 
   mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
   int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
@@ -620,9 +722,18 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse12) {
 
   mjtNum diag[] = {1, 1, 1};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 4, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, rowsuperT, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(4, 4, 4, 4));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 4, 8, 12));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 4);
   mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 4, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
-                     superowT, data);
+                     rowsuperT, data);
 
   EXPECT_THAT(matH,
               ElementsAre(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 10, 10, 1, 1, 10, 10));
@@ -652,7 +763,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse13) {
   int colindT[] = {0, 1, 2, 0, 1, 2};
   int rownnzT[] = {3, 3, 0, 0, 0};
   int rowadrT[] = {0, 3, 6, 6, 6};
-  int superowT[] = {1, 0, 2, 1, 0};
+  int rowsuperT[] = {1, 0, 2, 1, 0};
 
   mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
@@ -663,9 +774,18 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse13) {
 
   mjtNum diag[] = {1, 1, 1};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 3, 5, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, rowsuperT, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(2, 2, 0, 0, 0));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 2, 4, 4, 4));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 5);
   mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 5, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
-                     superowT, data);
+                     rowsuperT, data);
 
   EXPECT_THAT(matH, ElementsAre(3, 3, 0, 0, 0, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                                 0, 0, 0, 0, 0, 0, 0, 0, 0));
@@ -693,7 +813,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
   int colindT[] = {0, 0, 0, 0, 0, 0, 0};
   int rownnzT[] = {1, 1, 1, 1, 1, 1, 1};
   int rowadrT[] = {0, 1, 2, 3, 4, 5, 6};
-  int superowT[] = {3, 2, 1, 0, 2, 1, 0};
+  int rowsuperT[] = {3, 2, 1, 0, 2, 1, 0};
 
   mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
@@ -704,9 +824,18 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
   int rownnzH[] = {0, 0, 0, 0, 0, 0, 0};
   int rowadrH[] = {0, 0, 0, 0, 0, 0, 0};
 
+  // test precount
+  mju_sqrMatTDSparseInit(rownnzH, rowadrH, mat, matT, 1, 7, rownnz, rowadr,
+                         colind, rownnzT, rowadrT, colindT, rowsuperT, data);
+
+  EXPECT_THAT(rownnzH, ElementsAre(7, 7, 7, 7, 7, 7, 7));
+  EXPECT_THAT(rowadrH, ElementsAre(0, 7, 14, 21, 28, 35, 42));
+
+  // test computation
+  mju_sqrMatTDUncompressedInit(rowadrH, 7);
   mju_sqrMatTDSparse(matH, mat, matT, NULL, 1, 7, rownnzH, rowadrH, colindH,
                      rownnz, rowadr, colind, NULL, rownnzT, rowadrT, colindT,
-                     superowT, data);
+                     rowsuperT, data);
 
   EXPECT_THAT(
       matH, ElementsAre(1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2,
@@ -715,9 +844,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
   EXPECT_THAT(colindH,
               ElementsAre(0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3,
                           4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 0,
-                          1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6
+                          1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6));
 
-                          ));
   EXPECT_THAT(rownnzH, ElementsAre(7, 7, 7, 7, 7, 7, 7));
   EXPECT_THAT(rowadrH, ElementsAre(0, 7, 14, 21, 28, 35, 42));