chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Cache OBJ files in MuJoCo compiler.

Browse Source 
PiperOrigin-RevId: 616211896
Change-Id: I3b765f4529957914e6154f0e62ef36fd9df93c96
This commit is contained in:
Kyle Bayes
2024-03-15 12:30:51 -07:00
committed by Copybara-Service
parent bf66801bc9
commit 74d3d8cb4e
8 changed files with 227 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/user/CMakeLists.txt
+2
View File
@@ -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
+17 -2
View File
@@ -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
+10
View File
@@ -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
+2 -1
View File
@@ -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
+18 -17
View File
@@ -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& operator=(mjCCache&& other) = default;
mjCCache(mjCCache&& other) = delete;
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
+122
View File
@@ -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
+2
View File
@@ -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
+54
View File
@@ -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