Reduce flex stiffness and bending memory size to the exact amount needed.

PiperOrigin-RevId: 911774147
Change-Id: I3b18eba23e75ec91ece8342945f25e0cd7c63a1a

src/engine/engine_core_constraint.c

@@ -832,9 +832,14 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
 
       // read eigenmode data from flex_stiffness
       int ndof_elem = 3 * npe;
-      const mjtNum* k_elem = m->flex_stiffness + m->flex_stiffnessadr[f]
+      int stiffnessadr = m->flex_stiffnessadr[f];
-                           + elem_idx * ndof_elem * ndof_elem;
+      int neig = 0;
-      int neig = (int)k_elem[0];
+      const mjtNum* k_elem = NULL;
+      if (stiffnessadr >= 0) {
+        k_elem = m->flex_stiffness + stiffnessadr
+               + elem_idx * ndof_elem * ndof_elem;
+        neig = (int)k_elem[0];
+      }
 
       // compute displacement in corotational frame
       mjtNum* displ_e = mjSTACKALLOC(d, ndof_elem, mjtNum);
@@ -2387,9 +2392,12 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
 
       // read eigenmode count from flex_stiffness
       int ndof_elem = 3 * npe;
-      const mjtNum* k_elem = m->flex_stiffness + m->flex_stiffnessadr[f]
+      size = 0;
-                           + elem_idx * ndof_elem * ndof_elem;
+      if (m->flex_stiffnessadr[f] >= 0) {
-      size = (int)k_elem[0];  // neig stored as first element
+        const mjtNum* k_elem = m->flex_stiffness + m->flex_stiffnessadr[f]
+                             + elem_idx * ndof_elem * ndof_elem;
+        size = (int)k_elem[0];  // neig stored as first element
+      }
 
       if (nnz) {
         // get element node body IDs

src/engine/engine_derivative.c

@@ -951,7 +951,11 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
     }
 
     // get stiffness and damping
-    mjtNum* K = m->flex_stiffness + m->flex_stiffnessadr[f];
+    int stiffnessadr = m->flex_stiffnessadr[f];
+    if (stiffnessadr < 0) {
+      continue;
+    }
+    mjtNum* K = m->flex_stiffness + stiffnessadr;
 
     // skip if rigid or no stiffness
     if (m->flex_rigid[f] || K[0] == 0) {

src/engine/engine_passive.c

@@ -62,7 +62,12 @@ static void inline GradSquaredLengths(mjtNum gradient[6][2][3],
 // passive forces for interpolated flex (stretch + bending)
 static void mj_flexPassiveInterp(const mjModel* m, mjData* d, int f,
                                  int enbl_spring, int enbl_damper) {
-  mjtNum* k = m->flex_stiffness + m->flex_stiffnessadr[f];
+  int stiffnessadr = m->flex_stiffnessadr[f];
+  if (stiffnessadr < 0) {
+    return;
+  }
+
+  mjtNum* k = m->flex_stiffness + stiffnessadr;
   int nodenum = m->flex_nodenum[f];
 
   int order = m->flex_interp[f];
@@ -230,8 +235,17 @@ static inline mjtNum mju_dphi2D(mjtNum s0, int l0, mjtNum s1, int l1,
 // accuracy for elements with varying curvature.
 static void mj_flexPassiveBendInterp(const mjModel* m, mjData* d, int f,
                                       int enbl_spring, int enbl_damper) {
+  if (m->flex_interp[f] >= 0) {
+    return;
+  }
+
+  int bendingadr = m->flex_bendingadr[f];
+  if (bendingadr < 0) {
+    return;
+  }
+
   // read bending edge data
-  const mjtNum* bdata = m->flex_bending + m->flex_bendingadr[f];
+  const mjtNum* bdata = m->flex_bending + bendingadr;
   int nedge = (int)bdata[0];
   if (nedge == 0) return;
 
@@ -398,10 +412,15 @@ static void mj_flexPassiveBend(const mjModel* m, mjData* d, int f,
     return;
   }
 
+  int bendingadr = m->flex_bendingadr[f];
+  if (bendingadr < 0) {
+    return;
+  }
+
   int edgenum = m->flex_edgenum[f];
   mjtNum* xpos = d->flexvert_xpos + 3*m->flex_vertadr[f];
   int* bodyid = m->flex_vertbodyid + m->flex_vertadr[f];
-  mjtNum* b = m->flex_bending + 17*m->flex_edgeadr[f];
+  mjtNum* b = m->flex_bending + bendingadr;
 
   for (int e = 0; e < edgenum; e++) {
     const int* edge = m->flex_edge + 2*(e+m->flex_edgeadr[f]);
@@ -469,7 +488,11 @@ static void mj_flexPassiveBend(const mjModel* m, mjData* d, int f,
 // passive forces for flex stretch
 static void mj_flexPassiveStretch(const mjModel* m, mjData* d, int f,
                                   int enbl_spring, int enbl_damper) {
-  mjtNum* k = m->flex_stiffness + m->flex_stiffnessadr[f];
+  int stiffnessadr = m->flex_stiffnessadr[f];
+  if (stiffnessadr < 0) {
+    return;
+  }
+  mjtNum* k = m->flex_stiffness + stiffnessadr;
   if (k[0] == 0) {
     return;
   }
@@ -660,9 +683,7 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
       mj_flexPassiveInterp(m, d, f, enbl_spring, enbl_damper);
 
       // interpolated shell bending forces
-      if (m->flex_interp[f] < 0) {
+      mj_flexPassiveBendInterp(m, d, f, enbl_spring, enbl_damper);
-        mj_flexPassiveBendInterp(m, d, f, enbl_spring, enbl_damper);
-      }
     } else {
       // add bending forces
       mj_flexPassiveBend(m, d, f, enbl_spring, enbl_damper);

src/user/user_mesh.cc

@@ -4918,6 +4918,14 @@ void mjCFlex::Compile(const mjVFS* vfs) {
     stiffness_cached = LoadCachedStiffness();
   }
 
+  // check if any strain equality references this flex
+  for (auto* equality : model->Equalities()) {
+    if (equality->spec.type == mjEQ_FLEXSTRAIN && *equality->spec.name1 == name) {
+      has_strain_eq = true;
+      break;
+    }
+  }
+
   if (!stiffness_cached && interpolated && (young > 0 || has_strain_eq)) {
     // use young=1 for strain constraints (eigenvectors are geometry-only)
     double K_young = has_strain_eq ? 1e1 : young;

src/user/user_model.cc

@@ -2106,21 +2106,6 @@ static size_t getpathslength(std::vector<T> list) {
   return result;
 }
 
-// compute extra stiffness/bending array size for an interpolated flex
-static int flexInterpExtraSize(int order, const int cellcount[3], bool shell) {
-  int npe, nelem;
-  int cx = cellcount[0], cy = cellcount[1], cz = cellcount[2];
-  if (shell) {
-    npe = (int)pow(order + 1, 2);
-    nelem = 2*(cy*cz + cx*cz + cx*cy);
-  } else {
-    npe = (int)pow(order + 1, 3);
-    nelem = cx * cy * cz;
-  }
-  int ndof_elem = 3 * npe;
-  return nelem * ndof_elem * ndof_elem;
-}
-
 // set array sizes
 void mjCModel::SetSizes() {
   // set from object list sizes
@@ -2194,8 +2179,6 @@ void mjCModel::SetSizes() {
   }
   nbvh = nbvhstatic + nbvhdynamic;
 
-  int extra_stiffness_size = 0;
-  int extra_bending_size = 0;
   // flex counts
   for (int i=0; i < nflex; i++) {
     nflexnode += flexes_[i]->nnode;
@@ -2207,13 +2190,8 @@ void mjCModel::SetSizes() {
     nflexshelldata += (int)flexes_[i]->shell.size();
     nflexevpair += (int)flexes_[i]->evpair.size()/2;
     nflextexcoord += (flexes_[i]->HasTexcoord() ? flexes_[i]->get_texcoord().size()/2 : 0);
-    if (flexes_[i]->spec.order != 0) {
+    nflexstiffness += flexes_[i]->stiffness.size();
-      int extra_size = flexInterpExtraSize(
+    nflexbending += flexes_[i]->bending.size();
-          flexes_[i]->spec.order, flexes_[i]->spec.cellcount,
-          flexes_[i]->elastic2d != 0);
-      extra_stiffness_size += extra_size;
-      extra_bending_size += extra_size;
-    }
     if (flexes_[i]->interpolated || flexes_[i]->rigid) {
       continue;
     }
@@ -2263,10 +2241,6 @@ void mjCModel::SetSizes() {
       }
     }
   }
-  // TODO: This can be compacted further when we update mjwarp to not rely on
-  // 21*elem_adr for non-interpolated flexes and 17*edge_adr for bending.
-  nflexstiffness = nflexelem * 21 + extra_stiffness_size;
-  nflexbending = nflexedge * 17 + extra_bending_size;
 
   // mesh counts
   for (int i=0; i < nmesh; i++) {
@@ -3351,6 +3325,7 @@ int mjCModel::CountNJten(const mjModel* m) {
 // copy objects outside kinematic tree
 void mjCModel::CopyObjects(mjModel* m) {
   mjtSize adr, bone_adr, vert_adr, node_adr, normal_adr, face_adr, texcoord_adr, oct_adr;
+  mjtSize stiffness_adr, bending_adr;
   mjtSize edge_adr, elem_adr, elemdata_adr, elemedge_adr, shelldata_adr, evpair_adr;
   mjtSize bonevert_adr, graph_adr, data_adr, bvh_adr;
   mjtSize poly_adr, polymap_adr, polyvert_adr;
@@ -3465,10 +3440,8 @@ void mjCModel::CopyObjects(mjModel* m) {
   shelldata_adr = 0;
   evpair_adr = 0;
   texcoord_adr = 0;
-  int standard_stiffness_size = 21 * m->nflexelem;
+  stiffness_adr = 0;
-  int current_extra_stiffness_adr = standard_stiffness_size;
+  bending_adr = 0;
-  int standard_bending_size = 17 * m->nflexedge;
-  int current_extra_bending_adr = standard_bending_size;
   for (int i=0; i < nflex; i++) {
     // get pointer
     mjCFlex* pfl = flexes_[i];
@@ -3491,46 +3464,19 @@ void mjCModel::CopyObjects(mjModel* m) {
     mjuu_copyvec(m->flex_rgba + 4 * i, pfl->rgba, 4);
 
     // elasticity
-    if (pfl->spec.order == 0) {
+    if (pfl->stiffness.empty()) {
-      m->flex_stiffnessadr[i] = 21 * elem_adr;
+      m->flex_stiffnessadr[i] = -1;
     } else {
-      m->flex_stiffnessadr[i] = current_extra_stiffness_adr;
+      m->flex_stiffnessadr[i] = stiffness_adr;
-      current_extra_stiffness_adr += flexInterpExtraSize(
-          pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
-    }
-
-    if (!pfl->stiffness.empty()) {
       mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i],
                    pfl->stiffness.data(), pfl->stiffness.size());
-    } else {
-      int stiff_size;
-      if (pfl->spec.order == 0) {
-        stiff_size = 21 * pfl->nelem;
-      } else {
-        stiff_size = flexInterpExtraSize(
-            pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
-      }
-      mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], stiff_size);
     }
-    if (pfl->spec.order == 0) {
+    if (pfl->bending.empty()) {
-      m->flex_bendingadr[i] = 17 * edge_adr;
+      m->flex_bendingadr[i] = -1;
     } else {
-      m->flex_bendingadr[i] = current_extra_bending_adr;
+      m->flex_bendingadr[i] = bending_adr;
-      current_extra_bending_adr += flexInterpExtraSize(
+      mjuu_copyvec(m->flex_bending + m->flex_bendingadr[i],
-          pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
+                   pfl->bending.data(), pfl->bending.size());
-    }
-
-    if (!pfl->bending.empty()) {
-      mjuu_copyvec(m->flex_bending + m->flex_bendingadr[i], pfl->bending.data(), pfl->bending.size());
-    } else {
-      int bending_size;
-      if (pfl->spec.order == 0) {
-        bending_size = 17 * pfl->nedge;
-      } else {
-        bending_size = flexInterpExtraSize(
-            pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
-      }
-      mjuu_zerovec(m->flex_bending + m->flex_bendingadr[i], bending_size);
     }
     m->flex_damping[i] = (mjtNum)pfl->damping;
 
@@ -3704,6 +3650,8 @@ void mjCModel::CopyObjects(mjModel* m) {
     evpair_adr += (int)pfl->evpair.size()/2;
     texcoord_adr += (int)pfl->texcoord_.size()/2;
     bvh_adr += pfl->tree.Nbvh();
+    stiffness_adr += pfl->stiffness.size();
+    bending_adr += pfl->bending.size();
   }
 
   // skins