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

Cache island topology data

Browse Source 
Store island topology in `mjData` so repeated island solves can reuse stable connect/weld equality partitions. Add cache invalidation checks for active equality changes and cover the fast path with an island regression test.

Signed-off-by: teerthsharma <teerths57@gmail.com>
This commit is contained in:
teerthsharma
2026-07-11 19:08:15 +05:30
committed by teerthsharma
parent 03203e2f32
commit bc32db7f25
5 changed files with 319 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/mujoco/mjdata.h
+5
View File
@@ -198,6 +198,11 @@ typedef struct mjData_ {
// sleep state // sleep state
int* tree_asleep; // <0: awake; >=0: index cycle of sleeping trees (ntree x 1) int* tree_asleep; // <0: awake; >=0: index cycle of sleeping trees (ntree x 1)
// opaque internal island topology cache; users must not modify
int* island_cache_tree; // tree and island topology (ntree x 28)
int* island_cache_dof; // DOF topology (nv x 8)
int* island_cache_eq; // equality topology (neq x 40)
// plugins // plugins
int* plugin; // copy of m->plugin, required for deletion (nplugin x 1) int* plugin; // copy of m->plugin, required for deletion (nplugin x 1)
uintptr_t* plugin_data; // pointer to plugin-managed data structure (nplugin x 1) uintptr_t* plugin_data; // pointer to plugin-managed data structure (nplugin x 1)
include/mujoco/mjxmacro.h
+5 -1
View File
@@ -929,7 +929,10 @@
X ( mjtNum, qfrc_inverse, nv, 1 ) \ X ( mjtNum, qfrc_inverse, nv, 1 ) \
X ( mjtNum, cacc, nbody, 6 ) \ X ( mjtNum, cacc, nbody, 6 ) \
X ( mjtNum, cfrc_int, nbody, 6 ) \ X ( mjtNum, cfrc_int, nbody, 6 ) \
X ( mjtNum, cfrc_ext, nbody, 6 ) X ( mjtNum, cfrc_ext, nbody, 6 ) \
XIC ( int, island_cache_tree, ntree, 28 ) \
XIC ( int, island_cache_dof, nv, 8 ) \
XIC ( int, island_cache_eq, neq, 40 )
// macro for annotating that an array size in an X macro is a member of mjData // macro for annotating that an array size in an X macro is a member of mjData
@@ -1080,5 +1083,6 @@
// to obtain only X macros for fields that are relevant for mjvScene creation, // to obtain only X macros for fields that are relevant for mjvScene creation,
// redefine XNV to expand to nothing // redefine XNV to expand to nothing
#define XNV X #define XNV X
#define XIC(type, name, nr, nc)
#endif // MUJOCO_MJXMACRO_H_ #endif // MUJOCO_MJXMACRO_H_
src/engine/engine_io.c
+28 -1
View File
@@ -998,7 +998,11 @@ static void mj_setPtrData(const mjModel* m, mjData* d) {
ASAN_POISON_MEMORY_REGION(ptr, PTRDIFF(d->name, ptr)); \ ASAN_POISON_MEMORY_REGION(ptr, PTRDIFF(d->name, ptr)); \
ptr += SKIP((intptr_t)ptr) + sizeof(type)*(m->nr)*(nc); ptr += SKIP((intptr_t)ptr) + sizeof(type)*(m->nr)*(nc);
#undef XIC
#define XIC X
MJDATA_POINTERS MJDATA_POINTERS
#undef XIC
#define XIC(type, name, nr, nc)
#undef X #undef X
// check size // check size
@@ -1085,7 +1089,11 @@ void mj_makeRawData(mjData** dest, const mjModel* m) {
return; \ return; \
} }
#undef XIC
#define XIC X
MJDATA_POINTERS MJDATA_POINTERS
#undef XIC
#define XIC(type, name, nr, nc)
#undef X #undef X
// copy stack size from model // copy stack size from model
@@ -1108,6 +1116,9 @@ void mj_makeRawData(mjData** dest, const mjModel* m) {
// set pointers into buffer // set pointers into buffer
mj_setPtrData(m, d); mj_setPtrData(m, d);
if (m->ntree) {
d->island_cache_tree[0] = 0;
}
// clear threadpool // clear threadpool
d->threadpool = 0; d->threadpool = 0;
@@ -1173,6 +1184,9 @@ mjData* mj_copyDataVisual(mjData* dest, const mjModel* m, const mjData* src, int
dest->arena = save_arena; dest->arena = save_arena;
dest->threadpool = 0; dest->threadpool = 0;
mj_setPtrData(m, dest); mj_setPtrData(m, dest);
if (m->ntree) {
dest->island_cache_tree[0] = 0;
}
// save plugin_data, since the X macro copying block below will override it // save plugin_data, since the X macro copying block below will override it
const size_t plugin_data_size = sizeof(*dest->plugin_data) * dest->nplugin; const size_t plugin_data_size = sizeof(*dest->plugin_data) * dest->nplugin;
@@ -1373,19 +1387,32 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
// fill buffer with debug_value (normally 0) // fill buffer with debug_value (normally 0)
#ifdef ADDRESS_SANITIZER #ifdef ADDRESS_SANITIZER
{ {
#undef XIC
#define XIC(type, name, nr, nc)
#define X(type, name, nr, nc) memset(d->name, (int)debug_value, sizeof(type)*(m->nr)*(nc)); #define X(type, name, nr, nc) memset(d->name, (int)debug_value, sizeof(type)*(m->nr)*(nc));
MJDATA_POINTERS MJDATA_POINTERS
#undef X #undef X
#undef XIC
#define XIC(type, name, nr, nc)
} }
#else #else
memset(d->buffer, (int)debug_value, d->nbuffer); size_t cache_offset = d->island_cache_tree ?
(size_t)PTRDIFF(d->island_cache_tree, d->buffer) : d->nbuffer;
memset(d->buffer, (int)debug_value, cache_offset);
#endif #endif
#ifdef MEMORY_SANITIZER #ifdef MEMORY_SANITIZER
// under MSAN, mark the entire buffer as uninitialized // under MSAN, mark the entire buffer as uninitialized
__msan_allocated_memory(d->buffer, d->nbuffer); __msan_allocated_memory(d->buffer, d->nbuffer);
if (m->ntree) {
d->island_cache_tree[0] = 0;
}
#endif #endif
if (debug_value && m->ntree) {
d->island_cache_tree[0] = 0;
}
// zero out user-settable state and input arrays (MSAN: mark as initialized) // zero out user-settable state and input arrays (MSAN: mark as initialized)
mju_zero(d->qpos, m->nq); mju_zero(d->qpos, m->nq);
mju_zero(d->qvel, m->nv); mju_zero(d->qvel, m->nv);
src/engine/engine_island.c
+218
View File
@@ -14,8 +14,10 @@
#include "engine/engine_island.h" #include "engine/engine_island.h"
#include <stdint.h>
#include <stdio.h> #include <stdio.h>
#include <stddef.h> #include <stddef.h>
#include <string.h>
#include <mujoco/mjdata.h> #include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h> #include <mujoco/mjmodel.h>
@@ -427,6 +429,211 @@ static void assignConstraintIslands(const mjModel* m, mjData* d, const int* tree
} }
enum {
kCacheMagic = 0x49534C44,
kCacheHeader = 16,
};
typedef struct mjIslandCacheView_ {
int* tree_island;
int* island_ntree;
int* island_itreeadr;
int* map_itree2tree;
int* island_nv;
int* island_idofadr;
int* island_dofadr;
int* island_ne;
int* island_nf;
int* island_nefc;
int* island_iefcadr;
int* tree_dofnum;
int* dof_treeid;
int* dof_island;
int* map_dof2idof;
int* map_idof2dof;
int* efc_island;
int* efc_type;
int* efc_id;
int* map_efc2iefc;
int* map_iefc2efc;
int* eq_active;
int* eq_type;
int* eq_tree1;
int* eq_tree2;
} mjIslandCacheView;
static mjIslandCacheView islandCacheView(const mjModel* m, const mjData* d) {
int* tree = d->island_cache_tree + kCacheHeader;
int* dof = d->island_cache_dof;
int* eq = d->island_cache_eq;
int nisland = d->island_cache_tree[4];
int nefc = d->island_cache_tree[3];
mjIslandCacheView view;
#define TAKE(base, name, count) view.name = base; base += (count)
TAKE(tree, tree_island, m->ntree);
TAKE(tree, island_ntree, nisland);
TAKE(tree, island_itreeadr, nisland);
TAKE(tree, map_itree2tree, m->ntree);
TAKE(tree, island_nv, nisland);
TAKE(tree, island_idofadr, nisland);
TAKE(tree, island_dofadr, nisland);
TAKE(tree, island_ne, nisland);
TAKE(tree, island_nf, nisland);
TAKE(tree, island_nefc, nisland);
TAKE(tree, island_iefcadr, nisland);
TAKE(tree, tree_dofnum, m->ntree);
TAKE(dof, dof_treeid, m->nv);
TAKE(dof, dof_island, m->nv);
TAKE(dof, map_dof2idof, m->nv);
TAKE(dof, map_idof2dof, m->nv);
TAKE(eq, efc_island, nefc);
TAKE(eq, efc_type, nefc);
TAKE(eq, efc_id, nefc);
TAKE(eq, map_efc2iefc, nefc);
TAKE(eq, map_iefc2efc, nefc);
TAKE(eq, eq_active, m->neq);
TAKE(eq, eq_type, m->neq);
TAKE(eq, eq_tree1, m->neq);
TAKE(eq, eq_tree2, m->neq);
#undef TAKE
return view;
}
static int islandCacheMatches(const mjModel* m, const mjData* d) {
if (!m->ntree || !m->neq || d->island_cache_tree[0] != kCacheMagic ||
d->island_cache_tree[1] != m->ntree || d->island_cache_tree[2] != m->nv ||
d->island_cache_tree[3] != d->nefc || d->island_cache_tree[6] != m->neq ||
d->nefc != d->ne || d->nefc < 0 || (int64_t)d->nefc > 6*(int64_t)m->neq ||
d->island_cache_tree[4] < 0 || d->island_cache_tree[4] > m->ntree ||
d->island_cache_tree[5] < 0 || d->island_cache_tree[5] > m->nv) {
return 0;
}
mjIslandCacheView view = islandCacheView(m, d);
if (memcmp(view.efc_type, d->efc_type, d->nefc*sizeof(int)) ||
memcmp(view.efc_id, d->efc_id, d->nefc*sizeof(int)) ||
memcmp(view.dof_treeid, m->dof_treeid, m->nv*sizeof(int)) ||
memcmp(view.tree_dofnum, m->tree_dofnum, m->ntree*sizeof(int))) {
return 0;
}
for (int i=0; i < m->neq; i++) {
if (view.eq_active[i] != d->eq_active[i] || view.eq_type[i] != m->eq_type[i] ||
(m->eq_type[i] != mjEQ_CONNECT && m->eq_type[i] != mjEQ_WELD)) {
return 0;
}
int obj1 = m->eq_obj1id[i];
int obj2 = m->eq_obj2id[i];
if (m->eq_objtype[i] == mjOBJ_SITE) {
obj1 = m->site_bodyid[obj1];
obj2 = m->site_bodyid[obj2];
}
if (view.eq_tree1[i] != m->body_treeid[obj1] ||
view.eq_tree2[i] != m->body_treeid[obj2]) {
return 0;
}
}
return 1;
}
static void restoreIslandCache(const mjModel* m, mjData* d) {
mjIslandCacheView view = islandCacheView(m, d);
#define RESTORE(name, count) mju_copyInt(d->name, view.name, (count))
RESTORE(tree_island, m->ntree);
RESTORE(island_ntree, d->nisland);
RESTORE(island_itreeadr, d->nisland);
RESTORE(map_itree2tree, m->ntree);
RESTORE(dof_island, m->nv);
RESTORE(island_nv, d->nisland);
RESTORE(island_idofadr, d->nisland);
RESTORE(island_dofadr, d->nisland);
RESTORE(map_dof2idof, m->nv);
RESTORE(map_idof2dof, m->nv);
RESTORE(efc_island, d->nefc);
RESTORE(island_ne, d->nisland);
RESTORE(island_nf, d->nisland);
RESTORE(island_nefc, d->nisland);
RESTORE(island_iefcadr, d->nisland);
RESTORE(map_efc2iefc, d->nefc);
RESTORE(map_iefc2efc, d->nefc);
#undef RESTORE
}
static void saveIslandCache(const mjModel* m, mjData* d) {
if (!m->ntree || !m->neq || d->nefc != d->ne ||
(int64_t)d->nefc > 6*(int64_t)m->neq) {
if (m->ntree) d->island_cache_tree[0] = 0;
return;
}
d->island_cache_tree[0] = 0;
d->island_cache_tree[1] = m->ntree;
d->island_cache_tree[2] = m->nv;
d->island_cache_tree[3] = d->nefc;
d->island_cache_tree[4] = d->nisland;
d->island_cache_tree[5] = d->nidof;
d->island_cache_tree[6] = m->neq;
d->island_cache_tree[7] = 1;
mjIslandCacheView view = islandCacheView(m, d);
for (int i=0; i < m->neq; i++) {
view.eq_active[i] = d->eq_active[i];
view.eq_type[i] = m->eq_type[i];
if (m->eq_type[i] != mjEQ_CONNECT && m->eq_type[i] != mjEQ_WELD) return;
int obj1 = m->eq_obj1id[i];
int obj2 = m->eq_obj2id[i];
if (m->eq_objtype[i] == mjOBJ_SITE) {
obj1 = m->site_bodyid[obj1];
obj2 = m->site_bodyid[obj2];
}
view.eq_tree1[i] = m->body_treeid[obj1];
view.eq_tree2[i] = m->body_treeid[obj2];
}
#define SAVE(name, count) mju_copyInt(view.name, d->name, (count))
SAVE(tree_island, m->ntree);
SAVE(island_ntree, d->nisland);
SAVE(island_itreeadr, d->nisland);
SAVE(map_itree2tree, m->ntree);
SAVE(dof_island, m->nv);
SAVE(island_nv, d->nisland);
SAVE(island_idofadr, d->nisland);
SAVE(island_dofadr, d->nisland);
SAVE(map_dof2idof, m->nv);
SAVE(map_idof2dof, m->nv);
SAVE(efc_island, d->nefc);
SAVE(island_ne, d->nisland);
SAVE(island_nf, d->nisland);
SAVE(island_nefc, d->nisland);
SAVE(island_iefcadr, d->nisland);
SAVE(map_efc2iefc, d->nefc);
SAVE(map_iefc2efc, d->nefc);
mju_copyInt(view.efc_type, d->efc_type, d->nefc);
mju_copyInt(view.efc_id, d->efc_id, d->nefc);
mju_copyInt(view.dof_treeid, m->dof_treeid, m->nv);
mju_copyInt(view.tree_dofnum, m->tree_dofnum, m->ntree);
#undef SAVE
for (int i=0; i < d->nefc; i++) {
if (d->map_iefc2efc[i] != i) {
d->island_cache_tree[7] = 0;
break;
}
}
d->island_cache_tree[0] = kCacheMagic;
}
static void copyIslandEfcVectors(mjData* d) {
if (d->island_cache_tree[0] == kCacheMagic && d->island_cache_tree[7]) {
d->iefc_type = d->efc_type;
d->iefc_id = d->efc_id;
d->iefc_frictionloss = d->efc_frictionloss;
d->iefc_D = d->efc_D;
d->iefc_R = d->efc_R;
return;
}
mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, d->nefc);
mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, d->nefc);
mju_gather(d->iefc_frictionloss, d->efc_frictionloss, d->map_iefc2efc, d->nefc);
mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, d->nefc);
mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, d->nefc);
}
//-------------------------- main entry-point ----------------------------------------------------- //-------------------------- main entry-point -----------------------------------------------------
// discover islands: // discover islands:
@@ -440,6 +647,16 @@ void mj_island(const mjModel* m, mjData* d) {
return; return;
} }
// exact fast path for topology-stable connect/weld equality constraints
if (islandCacheMatches(m, d)) {
d->nisland = d->island_cache_tree[4];
d->nidof = d->island_cache_tree[5];
if (!arenaAllocIsland(m, d)) return;
restoreIslandCache(m, d);
copyIslandEfcVectors(d);
return;
}
mj_markStack(d); mj_markStack(d);
// union direct tree incidence and assign deterministic components // union direct tree incidence and assign deterministic components
@@ -605,5 +822,6 @@ void mj_island(const mjModel* m, mjData* d) {
// SHOULD NOT OCCUR // SHOULD NOT OCCUR
if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount"); if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount");
saveIslandCache(m, d);
mj_freeStack(d); mj_freeStack(d);
} }
test/engine/engine_island_test.cc
+63
View File
@@ -295,6 +295,69 @@ TEST_F(IslandTest, ProductionFlexEqualityRescansRows) {
EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), ElementsAre(0, 1)); EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), ElementsAre(0, 1));
} }
TEST_F(IslandTest, EqualityTopologyCacheInvalidatesExactly) {
static constexpr char xml[] = R"(
<mujoco>
<option jacobian="sparse"><flag contact="disable" gravity="disable"/></option>
<worldbody>
<body name="b0"><joint type="slide"/><geom size=".1"/></body>
<body name="b1"><joint type="slide"/><geom size=".1"/></body>
<body name="b2"><joint type="slide"/><geom size=".1"/></body>
</worldbody>
<equality>
<connect body1="b0" body2="b1" anchor="0 0 0"/>
<connect body1="b1" body2="b2" anchor="0 0 0"/>
</equality>
</mujoco>
)";
char error[1024] = {};
MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(model.get(), NotNull()) << error;
MjDataPtr data = MakeData(model);
auto expect_all_connected = [&] {
EXPECT_EQ(data->nisland, 1);
EXPECT_EQ(data->nidof, 3);
EXPECT_EQ(data->nefc, 6);
EXPECT_THAT(AsVector(data->tree_island, model->ntree),
ElementsAre(0, 0, 0));
EXPECT_THAT(AsVector(data->dof_island, model->nv), ElementsAre(0, 0, 0));
EXPECT_THAT(AsVector(data->efc_island, data->nefc),
ElementsAre(0, 0, 0, 0, 0, 0));
};
// Identical topology is stable across repeated evaluation and ordinary reset.
mj_fwdPosition(model.get(), data.get());
expect_all_connected();
mj_fwdPosition(model.get(), data.get());
expect_all_connected();
// Copying a warm cache produces independent data with the same result.
MjDataPtr copy(mj_copyData(nullptr, model.get(), data.get()));
ASSERT_THAT(copy.get(), NotNull());
mj_fwdPosition(model.get(), copy.get());
EXPECT_THAT(AsVector(copy->tree_island, model->ntree),
ElementsAre(0, 0, 0));
mj_resetData(model.get(), data.get());
mj_fwdPosition(model.get(), data.get());
expect_all_connected();
// Equality activation changes invalidate the cached partition.
data->eq_active[1] = 0;
mj_fwdPosition(model.get(), data.get());
EXPECT_EQ(data->nisland, 1);
EXPECT_EQ(data->nidof, 2);
EXPECT_EQ(data->nefc, 3);
EXPECT_THAT(AsVector(data->tree_island, model->ntree), ElementsAre(0, 0, -1));
EXPECT_THAT(AsVector(data->dof_island, model->nv), ElementsAre(0, 0, -1));
EXPECT_THAT(AsVector(data->efc_island, data->nefc), ElementsAre(0, 0, 0));
// The earlier copy remains independent after the source topology changes.
mj_fwdPosition(model.get(), copy.get());
EXPECT_THAT(AsVector(copy->tree_island, model->ntree), ElementsAre(0, 0, 0));
}
TEST_F(IslandTest, BoundedArenaSupports1024Trees) { TEST_F(IslandTest, BoundedArenaSupports1024Trees) {
constexpr int kTreeCount = 1024; constexpr int kTreeCount = 1024;
constexpr size_t kArenaBytes = 2 * 1024 * 1024; constexpr size_t kArenaBytes = 2 * 1024 * 1024;