From 2937efd5f208cfa4e601bc59be058d6add6f4479 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Thu, 3 Jul 2025 10:51:35 -0700
Subject: [PATCH] Use sparse math in `mjd_actuator_vel`.

PiperOrigin-RevId: 778924297
Change-Id: I221fb338c894f278e9f8401f25056539c62ccc65
---
 src/engine/engine_derivative.c | 30 ++++++++++++------------------
 1 file changed, 12 insertions(+), 18 deletions(-)

diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c
index deb2c0c9..d8b4c8bc 100644
--- a/src/engine/engine_derivative.c
+++ b/src/engine/engine_derivative.c
@@ -734,20 +734,23 @@ static void addJTBJSparse(
   // compute qDeriv(k,p) += sum_{i,j} ( J(i,k)*B(i,j)*J(j,p) )
   for (int i = 0; i < n; i++) {
     for (int j = 0; j < n; j++) {
-      int offset_i = offset+i, offset_j = offset+j;
       if (!B[i*n+j]) {
         continue;
       }
 
       // loop over non-zero elements of J(i,:)
-      for (int k = 0; k < J_rownnz[offset_i]; k++) {
-        int ik = J_rowadr[offset_i] + k;
+      int nnz_i = J_rownnz[offset+i];
+      int adr_i = J_rowadr[offset+i];
+      int nnz_j = J_rownnz[offset+j];
+      int adr_j = J_rowadr[offset+j];
+      for (int k = 0; k < nnz_i; k++) {
+        int ik = adr_i + k;
         int colik = J_colind[ik];
 
         // qDeriv(k,:) += J(j,:) * J(i,k)*B(i,j)
-        mju_addToSclSparseInc(d->qDeriv + d->D_rowadr[colik], J + J_rowadr[offset_j],
+        mju_addToSclSparseInc(d->qDeriv + d->D_rowadr[colik], J + adr_j,
                               d->D_rownnz[colik], d->D_colind + d->D_rowadr[colik],
-                              J_rownnz[offset_j], J_colind + J_rowadr[offset_j],
+                              nnz_j, J_colind + adr_j,
                               J[ik]*B[i*n+j]);
       }
     }
@@ -810,17 +813,13 @@ static mjtNum mjd_muscleGain_vel(mjtNum len, mjtNum vel, const mjtNum lengthrang
 
 // add (d qfrc_actuator / d qvel) to qDeriv
 void mjd_actuator_vel(const mjModel* m, mjData* d) {
-  int nv = m->nv, nu = m->nu;
+  int nu = m->nu;
 
   // disabled: nothing to add
   if (mjDISABLED(mjDSBL_ACTUATION)) {
     return;
   }
 
-  // allocate dense actuator_moment row
-  mj_markStack(d);
-  mjtNum* moment = mjSTACKALLOC(d, nv, mjtNum);
-
   // process actuators
   for (int i=0; i < nu; i++) {
     // skip if disabled
@@ -864,14 +863,10 @@ void mjd_actuator_vel(const mjModel* m, mjData* d) {
 
     // add
     if (bias_vel != 0) {
-      mju_sparse2dense(moment, d->actuator_moment, 1, nv, d->moment_rownnz + i,
-                       d->moment_rowadr + i, d->moment_colind);
-      addJTBJ(m, d, moment, &bias_vel, 1);
+      addJTBJSparse(m, d, d->actuator_moment, &bias_vel, 1, i,
+                    d->moment_rownnz, d->moment_rowadr, d->moment_colind);
     }
   }
-
-  // free space
-  mj_freeStack(d);
 }
 
 
@@ -1260,8 +1255,7 @@ void mjd_ellipsoidFluid(const mjModel* m, mjData* d, int bodyid) {
 
     if (mj_isSparse(m)) {
       addJTBJSparse(m, d, J, B, 6, 0, rownnz, rowadr, colind_compressed);
-    }
-    else {
+    } else {
       addJTBJ(m, d, J, B, 6);
     }
   }