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

_realloc_island

Browse Source 
PiperOrigin-RevId: 941618362
Change-Id: Idef1deb0ff4976de42aa585993988ae3363cd1fc
This commit is contained in:
Taylor Howell
2026-07-02 04:19:19 -07:00
committed by Copybara-Service
parent 0e58c48504
commit c3adb50898
2 changed files with 112 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/functions.cc
+69
View File
@@ -1696,6 +1696,7 @@ PYBIND11_MODULE(_functions, pymodule, pybind11::mod_gil_not_used()) {
#define X(type, name, nr, nc) data->name = nullptr;
MJDATA_ARENA_POINTERS_SOLVER
MJDATA_ARENA_POINTERS_DUAL
MJDATA_ARENA_POINTERS_ISLAND
#undef X
};
@@ -1744,6 +1745,74 @@ PYBIND11_MODULE(_functions, pymodule, pybind11::mod_gil_not_used()) {
},
py::arg("d"), py::arg("ncon"), py::arg("nefc"), py::arg("nJ") = -1,
py::call_guard<py::gil_scoped_release>());
pymodule.def(
"_realloc_island",
[](MjDataWrapper& d, int nisland, int nidof) {
raw::MjData* data = d.get();
size_t parena_start = data->parena;
// Find island block start in arena to reclaim memory on re-allocation.
char* min_ptr = nullptr;
#define X(type, name, nr, nc) \
if (data->name && \
(!min_ptr || reinterpret_cast<char*>(data->name) < min_ptr)) { \
min_ptr = reinterpret_cast<char*>(data->name); \
}
MJDATA_ARENA_POINTERS_ISLAND
#undef X
if (min_ptr && data->arena) {
parena_start = min_ptr - static_cast<char*>(data->arena);
}
auto cleanup = [](raw::MjData* data, size_t target_parena) {
#define X(type, name, nr, nc) data->name = nullptr;
MJDATA_ARENA_POINTERS_ISLAND
#undef X
data->nisland = 0;
data->nidof = 0;
data->parena = target_parena;
#ifdef ADDRESS_SANITIZER
ASAN_POISON_MEMORY_REGION(
static_cast<char*>(data->arena) + target_parena,
data->narena - data->pstack - target_parena);
#endif
};
cleanup(data, parena_start);
char error_msg[128];
error_msg[0] = '\0';
const char* error_msg_fmt =
"Insufficient arena memory, currently allocated memory=\"%s\". "
"Increase using <size memory=\"X\"/>.";
data->nisland = nisland;
data->nidof = nidof;
#undef MJ_M
#define MJ_M(x) d.model().get()->x
#undef MJ_D
#define MJ_D(x) data->x
#define X(type, name, nr, nc) \
data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAllocByte)( \
data, sizeof(type) * (nr) * (nc), alignof(type))); \
if (!data->name) { \
cleanup(data, parena_start); \
std::snprintf(error_msg, sizeof(error_msg), error_msg_fmt, \
mju_writeNumBytes(data->narena)); \
throw FatalError(error_msg); \
}
MJDATA_ARENA_POINTERS_ISLAND
#undef X
#undef MJ_D
#define MJ_D(x) x
#undef MJ_M
#define MJ_M(x) x
},
py::arg("d"), py::arg("nisland"), py::arg("nidof"),
py::call_guard<py::gil_scoped_release>());
} // PYBIND11_MODULE NOLINT(readability/fn_size)
} // namespace
} // namespace mujoco::python