Cache OBJ files in MuJoCo compiler.

PiperOrigin-RevId: 616211896
Change-Id: I3b765f4529957914e6154f0e62ef36fd9df93c96

src/user/CMakeLists.txt

@@ -15,6 +15,8 @@
 set(MUJOCO_USER_SRCS
     user_api.cc
     user_api.h
+    user_cache.cc
+    user_cache.h
     user_composite.cc
     user_composite.h
     user_flexcomp.cc

src/user/user_api.cc

@@ -13,8 +13,11 @@
 // limitations under the License.
 
 #include "user/user_api.h"
+
+#include <cstddef>
 #include <functional>
 #include <map>
+#include <new>
 #include <string>
 #include <utility>
 #include <vector>
@@ -22,10 +25,9 @@
 #include <mujoco/mujoco.h>
 #include "user/user_model.h"
 #include "user/user_objects.h"
+#include "user/user_cache.h"
 #include "xml/xml_util.h"
 
-
-
 // create model
 mjSpec* mjm_createSpec() {
   mjCModel* modelC = new mjCModel;
@@ -600,4 +602,17 @@ const char* mjm_setFullInertia(mjmBody* bodyspec, double quat[4], double inertia
   return body->FullInertia(quat, inertia);
 }
 
+// -------------------------- GLOBAL ASSET CACHE -------------------------------
 
+void mj_setCacheSize(mjCache cache, std::size_t size) {
+  mjCCache* ccache = reinterpret_cast<mjCCache*>(cache);
+  if (ccache) {
+    ccache->SetMaxSize(size);
+  }
+}
+
+
+
+mjCache mj_globalCache() {
+  return NULL;  // currently disabled
+}

src/user/user_api.h

@@ -983,6 +983,16 @@ MJAPI void mjm_defaultKey(mjmKey& key);
 // Default plugin attributes.
 MJAPI void mjm_defaultPlugin(mjmPlugin& plugin);
 
+//------------------------- Cache functions ------------------------------------
+
+typedef struct _mjCache* mjCache;
+
+// Set the size of the cache in bytes.
+MJAPI void mj_setCacheSize(mjCache cache, size_t size);
+
+// Get internal global cache context.
+MJAPI mjCache mj_globalCache();
+
 #ifdef __cplusplus
 }
 #endif

src/user/user_cache.cc

@@ -39,7 +39,8 @@ template<typename T> std::size_t mjCAsset::Add(const std::string& name,
   const uint8_t* ptr = reinterpret_cast<const uint8_t*>(data);
   mjCAssetData& block = it->second;
 
-  block.bytes = std::make_shared<uint8_t[]>(nbytes);
+  block.bytes = std::shared_ptr<uint8_t>(new uint8_t[nbytes],
+                                        [](uint8_t *p) {delete [] p;});
   std::copy(ptr, ptr + nbytes, block.bytes.get());
   block.nbytes = nbytes;
   nbytes_ += nbytes;

src/user/user_cache.h

@@ -30,7 +30,7 @@
 
 // data associated with an asset
 struct mjCAssetData {
-  std::shared_ptr<uint8_t[]> bytes;  // raw serialized bytes of cached data
+  std::shared_ptr<uint8_t> bytes;    // raw serialized bytes of cached data
   std::size_t nbytes;                // number of bytes stored
 };
 
@@ -80,6 +80,11 @@ class mjCAsset {
     return blocks_.find(name) != blocks_.end();
   }
 
+  const std::string& Timestamp() const { return timestamp_; }
+  const std::string& Id() const { return id_; }
+  std::size_t InsertNum() const { return insert_num_; }
+  std::size_t AccessCount() const { return access_count_; }
+
  private:
   mjCAsset() = default;
 
@@ -105,10 +110,6 @@ class mjCAsset {
   void SetTimestamp(std::string timestamp) { timestamp_ = timestamp; }
 
   // accessors
-  const std::string& Id() const { return id_; }
-  const std::string& Timestamp() const { return timestamp_; }
-  std::size_t InsertNum() const { return insert_num_; }
-  std::size_t AccessCount() const { return access_count_; }
   std::size_t BytesCount() const { return nbytes_; }
   const std::unordered_map<std::string, mjCAssetData>& Blocks() const {
     return blocks_;
@@ -128,14 +129,23 @@ class mjCAsset {
   std::set<std::string> references_;
 };
 
+struct mjCAssetCompare {
+  bool operator()(const mjCAsset* e1, const mjCAsset* e2) const {
+    if (e1->AccessCount() != e2->AccessCount()) {
+      return e1->AccessCount() < e2->AccessCount();
+    }
+    return e1->InsertNum() < e2->InsertNum();
+  }
+};
+
 // the class container for a thread-safe asset cache
 class mjCCache {
  public:
   explicit mjCCache(std::size_t size)  :  max_size_(size) {}
 
   // move only
-  mjCCache(mjCCache&& other) = default;
+  mjCCache(mjCCache&& other) = delete;
-  mjCCache& operator=(mjCCache&& other) = default;
+  mjCCache& operator=(mjCCache&& other) = delete;
   mjCCache(const mjCCache& other) = delete;
   mjCCache& operator=(const mjCCache& other) = delete;
 
@@ -186,20 +196,11 @@ class mjCCache {
   std::size_t size_ = 0;        // current size of the cache in bytes
   std::size_t max_size_ = 0;    // max size of the cache in bytes
 
-  // compare function for the priority queue
-  static constexpr auto compare_ = [](const mjCAsset* e1,
-                                      const mjCAsset* e2) {
-    if (e1->AccessCount() != e2->AccessCount()) {
-      return e1->AccessCount() < e2->AccessCount();
-    }
-    return e1->InsertNum() < e2->InsertNum();
-  };
-
   // internal constant look up table for assets
   std::unordered_map<std::string, mjCAsset> lookup_;
 
   // internal priority queue for the cache
-  std::set<mjCAsset*, decltype(compare_)> entries_;
+  std::set<mjCAsset*, mjCAssetCompare> entries_;
 
   // models using the cache along with the assets they reference
   std::unordered_map<std::string, std::unordered_set<mjCAsset*>> models_;

src/user/user_mesh.cc

@@ -60,6 +60,7 @@
 #include "engine/engine_util_misc.h"
 #include "engine/engine_util_solve.h"
 #include "engine/engine_util_spatial.h"
+#include "user/user_cache.h"
 #include "user/user_model.h"
 #include "user/user_objects.h"
 #include "user/user_util.h"
@@ -967,6 +968,19 @@ void mjCMesh::RemoveRepeated() {
 void mjCMesh::LoadOBJ(mjResource* resource) {
   tinyobj::ObjReader objReader;
   const void* bytes = nullptr;
+
+  // try loading from cache
+  mjCCache *cache = reinterpret_cast<mjCCache*>(mj_globalCache());
+  if (cache) {
+    auto asset = cache->Get(resource->name);
+    if (asset.has_value() &&
+        !mju_isModifiedResource(resource, asset->Timestamp().c_str())) {
+      if (LoadCachedOBJ(asset.value())) {
+        return;
+      }
+    }
+  }
+
   int buffer_sz = mju_readResource(resource, &bytes);
   if (buffer_sz < 0) {
     throw mjCError(this, "could not read OBJ file '%s'", resource->name);
@@ -1030,9 +1044,117 @@ void mjCMesh::LoadOBJ(mjResource* resource) {
   for (int i=0; i<usertexcoord_.size()/2; i++) {
     usertexcoord_[2*i+1] = 1-usertexcoord_[2*i+1];
   }
+
+  // try caching asset
+  if (cache) {
+    mjCAsset asset("", resource->name, resource->timestamp);
+
+    asset.AddVector("uservert_", uservert_);
+    asset.AddVector("usernormal_", usernormal_);
+    asset.AddVector("usertexcoord_", usertexcoord_);
+
+    if (!objReader.GetShapes().empty()) {
+      const auto& mesh = objReader.GetShapes()[0].mesh;
+      std::vector<int> vertex_index;
+      vertex_index.reserve(mesh.indices.size());
+
+      std::vector<int> normal_index;
+      normal_index.reserve(mesh.indices.size());
+
+      std::vector<int> texcoord_index;
+      texcoord_index.reserve(mesh.indices.size());
+
+      for (tinyobj::index_t index : mesh.indices) {
+        vertex_index.push_back(index.vertex_index);
+        normal_index.push_back(index.normal_index);
+        texcoord_index.push_back(index.texcoord_index);
+      }
+
+      asset.AddVector("num_face_vertices", mesh.num_face_vertices);
+      asset.AddVector("vertex_index", vertex_index);
+      asset.AddVector("normal_index", normal_index);
+      asset.AddVector("texcoord_index", texcoord_index);
+    } else {
+      asset.AddVector("num_face_vertices", std::vector<unsigned char>());
+      asset.AddVector("vertex_index", std::vector<int>());
+      asset.AddVector("normal_index", std::vector<int>());
+      asset.AddVector("texcoord_index", std::vector<int>());
+    }
+    cache->Insert(std::move(asset));
+  }
 }
 
 
+
+// load OBJ from cached asset, return true on success
+bool mjCMesh::LoadCachedOBJ(const mjCAsset& asset) {
+  // check that asset has all data
+  if (!asset.HasData("uservert_") || !asset.HasData("usernormal_")
+      || !asset.HasData("usertexcoord_") || !asset.HasData("num_face_vertices")
+      || !asset.HasData("vertex_index") || !asset.HasData("normal_index")
+      || !asset.HasData("texcoord_index")) {
+    return false;
+  }
+  uservert_ = asset.GetVector<float>("uservert_").value();
+  usernormal_ = asset.GetVector<float>("usernormal_").value();
+  usertexcoord_ = asset.GetVector<float>("usertexcoord_").value();
+
+  vector<int> vertex_index = asset.GetVector<int>("vertex_index").value();
+  vector<int> normal_index = asset.GetVector<int>("normal_index").value();
+  vector<int> texcoord_index = asset.GetVector<int>("texcoord_index").value();
+  vector<unsigned char> num_face_vertices =
+      asset.GetVector<unsigned char>("num_face_vertices").value();
+
+  bool righthand = (scale[0] * scale[1] * scale[2]) > 0;
+
+
+  for (int face = 0, i = 0; i < vertex_index.size();) {
+    int nfacevert = num_face_vertices[face];
+    if (nfacevert < 3 || nfacevert > 4) {
+      throw mjCError(
+          this, "only tri or quad meshes are supported for OBJ (file '%s')",
+          asset.Id().c_str());
+    }
+
+    userface_.push_back(vertex_index[i]);
+    userface_.push_back(vertex_index[i + (righthand == 1 ? 1 : 2)]);
+    userface_.push_back(vertex_index[i + (righthand == 1 ? 2 : 1)]);
+
+    if (!usernormal_.empty()) {
+      userfacenormal_.push_back(normal_index[i]);
+      userfacenormal_.push_back(normal_index[i + (righthand == 1 ? 1 : 2)]);
+      userfacenormal_.push_back(normal_index[i + (righthand == 1 ? 2 : 1)]);
+    }
+
+    if (!usertexcoord_.empty()) {
+      userfacetexcoord_.push_back(texcoord_index[i]);
+      userfacetexcoord_.push_back(texcoord_index[i + (righthand == 1 ? 1 : 2)]);
+      userfacetexcoord_.push_back(texcoord_index[i + (righthand == 1 ? 2 : 1)]);
+    }
+
+    if (nfacevert == 4) {
+      userface_.push_back(vertex_index[i]);
+      userface_.push_back(vertex_index[i + (righthand == 1 ? 2 : 3)]);
+      userface_.push_back(vertex_index[i + (righthand == 1 ? 3 : 2)]);
+
+      if (!usernormal_.empty()) {
+        userfacenormal_.push_back(normal_index[i]);
+        userfacenormal_.push_back(normal_index[i + (righthand == 1 ? 1 : 2)]);
+        userfacenormal_.push_back(normal_index[i + (righthand == 1 ? 2 : 1)]);
+      }
+
+      if (!usertexcoord_.empty()) {
+        userfacetexcoord_.push_back(texcoord_index[i]);
+        userfacetexcoord_.push_back(texcoord_index[i + (righthand == 1 ? 1 : 2)]);
+        userfacetexcoord_.push_back(texcoord_index[i + (righthand == 1 ? 2 : 1)]);
+      }
+    }
+    i += nfacevert;
+    ++face;
+  }
+  return true;
+}
+
 // load STL binary mesh
 void mjCMesh::LoadSTL(mjResource* resource) {
   bool righthand = (scale[0]*scale[1]*scale[2]>0);

src/user/user_objects.h

@@ -26,6 +26,7 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include "user/user_api.h"
+#include "user/user_cache.h"
 
 // forward declarations of all mjC/X classes
 class mjCError;
@@ -827,6 +828,7 @@ class mjCMesh: public mjCMesh_, private mjmMesh {
 
  private:
   void LoadOBJ(mjResource* resource);         // load mesh in wavefront OBJ format
+  bool LoadCachedOBJ(const mjCAsset& asset);  // load OBJ from cache asset, return true on success
   void LoadSTL(mjResource* resource);         // load mesh in STL BIN format
   void LoadMSH(mjResource* resource);         // load mesh in MSH BIN format
   void LoadSDF();                             // generate mesh using marching cubes

test/user/user_api_test.cc

@@ -16,6 +16,7 @@
 
 #include <array>
 #include <filesystem>
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -116,5 +117,58 @@ TEST_F(PluginTest, RecompileCompare) {
   }
 }
 
+// ------------------- test cache with modified assets -------------------------
+TEST_F(PluginTest, RecompileCompareCache) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <asset>
+      <mesh file="cube.obj"/>
+    </asset>
+    <worldbody>
+      <geom type="mesh" mesh="cube"/>
+    </worldbody>
+  </mujoco>)";
+
+  static constexpr char cube1[] = R"(
+  v -0.500000 -0.500000  0.500000
+  v  0.500000 -0.500000  0.500000
+  v -0.500000  0.500000  0.500000
+  v  0.500000  0.500000  0.500000
+  v -0.500000  0.500000 -0.500000
+  v  0.500000  0.500000 -0.500000
+  v -0.500000 -0.500000 -0.500000
+  v  0.500000 -0.500000 -0.500000)";
+
+  static constexpr char cube2[] = R"(
+  v -1 -1  1
+  v  1 -1  1
+  v -1  1  1
+  v  1  1  1
+  v -1  1 -1
+  v  1  1 -1
+  v -1 -1 -1
+  v  1 -1 -1)";
+
+  auto vfs = std::make_unique<mjVFS>();
+  mj_defaultVFS(vfs.get());
+  mj_addBufferVFS(vfs.get(), "cube.obj", cube1, sizeof(cube1));
+
+  std::array<char, 1024> error;
+
+  // load model once
+  mjModel* m = LoadModelFromString(xml, error.data(), error.size(), vfs.get());
+  EXPECT_EQ(m->mesh_vert[0], -0.5);
+  mj_deleteModel(m);
+
+  // update cube.obj, load again
+  mj_deleteFileVFS(vfs.get(), "cube.obj");
+  mj_addBufferVFS(vfs.get(), "cube.obj", cube2, sizeof(cube2));
+  m = LoadModelFromString(xml, error.data(), error.size(), vfs.get());
+  EXPECT_EQ(m->mesh_vert[0], -1);
+  mj_deleteModel(m);
+
+  mj_deleteVFS(vfs.get());
+}
+
 }  // namespace
 }  // namespace mujoco