Add caching for flex stiffness computation.

This change introduces methods to cache and load the computed stiffness matrix and nodal masses for flex elements. The cache key is generated based on material properties (Young's modulus, Poisson's ratio), interpolation order, and the bounding box of the flex. This avoids redundant expensive computations when compiling models with identical flex definitions.

PiperOrigin-RevId: 874689771
Change-Id: I6d1e299baf294a984b507e2a2602b2c068610536

src/user/user_mesh.cc

@@ -4053,6 +4053,88 @@ void mjCFlex::ResolveReferences(const mjCModel* m) {
 }
 
 
+std::string mjCFlex::ComputeStiffnessCacheKey() const {
+  std::size_t hash = 0;
+  auto combine = [&hash](std::size_t v) {
+    hash ^= v + 0x9e3779b9 + (hash << 6) + (hash >> 2);
+  };
+
+  combine(std::hash<double>{}(young));
+  combine(std::hash<double>{}(poisson));
+  combine(std::hash<int>{}(order_));
+
+  // compute bounding box from vertex positions
+  if (!vert_.empty()) {
+    double minx = vert_[0], maxx = vert_[0];
+    double miny = vert_[1], maxy = vert_[1];
+    double minz = vert_[2], maxz = vert_[2];
+    for (std::size_t i = 3; i < vert_.size(); i += 3) {
+      minx = std::min(minx, vert_[i]);
+      maxx = std::max(maxx, vert_[i]);
+      miny = std::min(miny, vert_[i + 1]);
+      maxy = std::max(maxy, vert_[i + 1]);
+      minz = std::min(minz, vert_[i + 2]);
+      maxz = std::max(maxz, vert_[i + 2]);
+    }
+    combine(std::hash<double>{}(maxx - minx));
+    combine(std::hash<double>{}(maxy - miny));
+    combine(std::hash<double>{}(maxz - minz));
+  }
+
+  for (std::size_t i = 0; i < vert_.size(); i += std::max(1, (int)vert_.size()/100)) {
+    combine(std::hash<double>{}(vert_[i]));
+  }
+
+  for (std::size_t i = 0; i < shell.size(); i += std::max(1, (int)shell.size()/50)) {
+    combine(std::hash<int>{}(shell[i]));
+  }
+
+  return "flex_stiffness:" + std::to_string(hash);
+}
+
+
+bool mjCFlex::LoadCachedStiffness() {
+  mjCCache* cache = reinterpret_cast<mjCCache*>(mj_getCache()->impl_);
+  if (!cache) return false;
+
+  std::string key = ComputeStiffnessCacheKey();
+
+  auto load_fn = [this](const void* data) {
+    const auto* cached = static_cast<const std::vector<double>*>(data);
+    stiffness = *cached;
+    return true;
+  };
+
+  mjResource dummy_resource{};
+  dummy_resource.name = const_cast<char*>(key.c_str());
+  dummy_resource.timestamp[0] = '\0';
+
+  return cache->PopulateData(key, &dummy_resource, load_fn);
+}
+
+
+void mjCFlex::CacheStiffness() {
+  mjCCache* cache = reinterpret_cast<mjCCache*>(mj_getCache()->impl_);
+  if (!cache || stiffness.empty()) return;
+
+  std::string key = ComputeStiffnessCacheKey();
+
+  auto* cached = new std::vector<double>(stiffness);
+
+  std::size_t size = sizeof(*cached) + sizeof(double) * stiffness.size();
+
+  std::shared_ptr<const void> cached_data(cached, [](const void* data) {
+    delete static_cast<const std::vector<double>*>(data);
+  });
+
+  mjResource dummy_resource{};
+  dummy_resource.name = const_cast<char*>(key.c_str());
+  dummy_resource.timestamp[0] = '\0';
+
+  cache->Insert("", key, &dummy_resource, cached_data, size);
+}
+
+
 // compiler
 void mjCFlex::Compile(const mjVFS* vfs) {
   CopyFromSpec();
@@ -4295,7 +4377,6 @@ void mjCFlex::Compile(const mjVFS* vfs) {
       if (min_size > nelem) {
         throw mjCError(this, "Trilinear dofs are require at least %d elements", "", min_size);
       }
-      ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson, order_);
     }
 
     // geometrically nonlinear elasticity
@@ -4345,6 +4426,16 @@ void mjCFlex::Compile(const mjVFS* vfs) {
   // create shell fragments and element-vertex collision pairs
   CreateShellPair();
 
+  // compute linear stiffness for interpolated elements (cached)
+  bool stiffness_cached = false;
+  if (young > 0 && interpolated) {
+    stiffness_cached = LoadCachedStiffness();
+  }
+
+  if (!stiffness_cached && young > 0 && interpolated) {
+    ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson, order_);
+  }
+
   // create bounding volume hierarchy
   CreateBVH();
 

src/user/user_objects.h

@@ -992,6 +992,9 @@ class mjCFlex_ : public mjCBase {
   std::vector<int> spec_elem_;
   std::vector<float> spec_texcoord_;
   std::vector<int> spec_elemtexcoord_;
+
+  // caching
+  std::vector<double> cached_stiffness_;   // cached stiffness matrix
 };
 
 class mjCFlex: public mjCFlex_, private mjsFlex {
@@ -1040,6 +1043,11 @@ class mjCFlex: public mjCFlex_, private mjsFlex {
   std::vector<double> node0_;             // node Cartesian positions
 
   int order_ = 0;                         // interpolation order
+
+  // stiffness caching
+  std::string ComputeStiffnessCacheKey() const;
+  bool LoadCachedStiffness();
+  void CacheStiffness();
 };
 
 

test/user/user_flex_test.cc

@@ -456,6 +456,46 @@ TEST_F(UserFlexTest, StiffnessMatrix) {
   mj_deleteModel(m);
 }
 
+TEST_F(UserFlexTest, StiffnessCacheDiffersByGeometry) {
+  std::array<char, 1024> error;
+
+  // Create two flexes with same material but different bounding boxes
+  static constexpr char xml_small[] = R"(
+  <mujoco>
+  <worldbody>
+    <flexcomp name="test" type="grid" count="3 3 3" spacing="1 1 1" dim="3" dof="trilinear">
+      <contact selfcollide="none" internal="false"/>
+      <elasticity young="1" poisson="0.3"/>
+    </flexcomp>
+  </worldbody>
+  </mujoco>
+  )";
+
+  static constexpr char xml_large[] = R"(
+  <mujoco>
+  <worldbody>
+    <flexcomp name="test" type="grid" count="3 3 3" spacing="2 2 2" dim="3" dof="trilinear">
+      <contact selfcollide="none" internal="false"/>
+      <elasticity young="1" poisson="0.3"/>
+    </flexcomp>
+  </worldbody>
+  </mujoco>
+  )";
+
+  mjModel* m_small = LoadModelFromString(xml_small, error.data(), error.size());
+  ASSERT_THAT(m_small, NotNull()) << error.data();
+
+  mjModel* m_large = LoadModelFromString(xml_large, error.data(), error.size());
+  ASSERT_THAT(m_large, NotNull()) << error.data();
+
+  // Same number of nodes but different stiffness due to different geometry
+  EXPECT_EQ(m_small->nflexnode, m_large->nflexnode);
+  EXPECT_NE(m_small->flex_stiffness[0], m_large->flex_stiffness[0]);
+
+  mj_deleteModel(m_small);
+  mj_deleteModel(m_large);
+}
+
 TEST_F(UserFlexTest, LoadTexture) {
   const std::string xml_path =
       GetTestDataFilePath("user/testdata/textured_torus_flex.xml");