diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index c2fddb81..800aefa3 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -198,6 +198,11 @@ typedef struct mjData_ { // sleep state 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 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) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 8f36c166..71f8ed39 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -929,7 +929,10 @@ X ( mjtNum, qfrc_inverse, nv, 1 ) \ X ( mjtNum, cacc, 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 @@ -1080,5 +1083,6 @@ // to obtain only X macros for fields that are relevant for mjvScene creation, // redefine XNV to expand to nothing #define XNV X +#define XIC(type, name, nr, nc) #endif // MUJOCO_MJXMACRO_H_ diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 482f5915..285561fb 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -998,7 +998,11 @@ static void mj_setPtrData(const mjModel* m, mjData* d) { ASAN_POISON_MEMORY_REGION(ptr, PTRDIFF(d->name, ptr)); \ ptr += SKIP((intptr_t)ptr) + sizeof(type)*(m->nr)*(nc); + #undef XIC + #define XIC X MJDATA_POINTERS + #undef XIC + #define XIC(type, name, nr, nc) #undef X // check size @@ -1085,7 +1089,11 @@ void mj_makeRawData(mjData** dest, const mjModel* m) { return; \ } + #undef XIC + #define XIC X MJDATA_POINTERS + #undef XIC + #define XIC(type, name, nr, nc) #undef X // copy stack size from model @@ -1108,6 +1116,9 @@ void mj_makeRawData(mjData** dest, const mjModel* m) { // set pointers into buffer mj_setPtrData(m, d); + if (m->ntree) { + d->island_cache_tree[0] = 0; + } // clear threadpool d->threadpool = 0; @@ -1173,6 +1184,9 @@ mjData* mj_copyDataVisual(mjData* dest, const mjModel* m, const mjData* src, int dest->arena = save_arena; dest->threadpool = 0; 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 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) #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)); MJDATA_POINTERS #undef X + #undef XIC + #define XIC(type, name, nr, nc) } #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 #ifdef MEMORY_SANITIZER // under MSAN, mark the entire buffer as uninitialized __msan_allocated_memory(d->buffer, d->nbuffer); + if (m->ntree) { + d->island_cache_tree[0] = 0; + } #endif + if (debug_value && m->ntree) { + d->island_cache_tree[0] = 0; + } + // zero out user-settable state and input arrays (MSAN: mark as initialized) mju_zero(d->qpos, m->nq); mju_zero(d->qvel, m->nv); diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 34403b11..207ccad8 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -14,8 +14,10 @@ #include "engine/engine_island.h" +#include #include #include +#include #include #include @@ -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 ----------------------------------------------------- // discover islands: @@ -440,6 +647,16 @@ void mj_island(const mjModel* m, mjData* d) { 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); // union direct tree incidence and assign deterministic components @@ -605,5 +822,6 @@ void mj_island(const mjModel* m, mjData* d) { // SHOULD NOT OCCUR if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount"); + saveIslandCache(m, d); mj_freeStack(d); } diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 3a4462ab..00762da2 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -295,6 +295,69 @@ TEST_F(IslandTest, ProductionFlexEqualityRescansRows) { EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), ElementsAre(0, 1)); } +TEST_F(IslandTest, EqualityTopologyCacheInvalidatesExactly) { + static constexpr char xml[] = R"( + + + + + + + + + + + + +)"; + 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) { constexpr int kTreeCount = 1024; constexpr size_t kArenaBytes = 2 * 1024 * 1024;