Add new mju_threadpool API function, and delete old threading API.

PiperOrigin-RevId: 922838541
Change-Id: Id9f7e0fb298ffde61fcc49a802dc78971858ce51
This commit is contained in:
Kyle Bayes
2026-05-28 10:09:07 -07:00
committed by Copybara-Service
parent a22fc2423a
commit b935d4153c
47 changed files with 576 additions and 1755 deletions
+39 -44
View File
@@ -15,8 +15,8 @@
// A benchmark for parsing and compiling models from XML.
#include <cstddef>
#include <vector>
#include <string>
#include <vector>
#include <benchmark/benchmark.h>
#include <absl/base/attributes.h>
@@ -36,6 +36,9 @@ static const int kNumWarmupSteps = 1000;
// number of steps to benchmark (before resetting state)
static const int kBatchSize = 50;
// number of threads to test
static const int kNumThreads = 9;
static const char kBoxMeshPath[] =
"../test/engine/testdata/collision_convex/perf/boxmesh.xml";
static const char kBoxBoxPath[] =
@@ -47,11 +50,10 @@ static const char kMixedPath[] =
class TestHarness {
public:
TestHarness(const char* xml_path, std::string label, int disable_flags = 0) {
TestHarness(const char* xml_path, std::string label) {
// Fail test if there are any mujoco errors
MujocoErrorTestGuard guard;
model_ = LoadModelFromPath(xml_path);
model_->opt.disableflags |= disable_flags;
data_ = mj_makeData(model_);
for (int i=0; i < kNumWarmupSteps; i++) {
mj_step(model_, data_);
@@ -62,13 +64,18 @@ class TestHarness {
int size = mj_stateSize(model_, spec_);
initial_state_.resize(size);
mj_getState(model_, data_, initial_state_.data(), spec_);
label_ = label;
name_ = label;
}
void Reset() {
mj_setState(model_, data_, initial_state_.data(), spec_);
}
void SetThreads(int nthread) {
mju_threadpool(data_, nthread);
nthread_ = nthread;
}
void RunBenchmark(benchmark::State& state) {
std::size_t ncon = 0;
while (state.KeepRunningBatch(kBatchSize)) {
@@ -81,7 +88,9 @@ class TestHarness {
}
}
state.SetLabel(label_);
std::string label = name_ + " " + std::to_string(nthread_ + 1) +
" thread(s)";
state.SetLabel(label);
state.SetItemsProcessed(ncon); // report number of contacts per second
}
@@ -91,7 +100,8 @@ class TestHarness {
}
private:
std::string label_;
std::string name_;
int nthread_ = 0;
int spec_;
mjModel* model_;
mjData* data_;
@@ -103,62 +113,47 @@ class TestHarness {
// separately in CPU profiles (and don't get replaced with raw calls to
// run_parse_benchmark).
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_BoxMesh_NativeCCD(benchmark::State& state) {
static TestHarness harness(kBoxMeshPath, "boxmesh.xml (nativeccd)");
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxMesh(benchmark::State& state) {
int nthread = state.range(0);
static TestHarness harness(kBoxMeshPath, "boxmesh.xml");
harness.SetThreads(nthread);
harness.RunBenchmark(state);
}
BENCHMARK(BM_BoxMesh_NativeCCD);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_BoxMesh_LibCCD(benchmark::State& state) {
static TestHarness harness(kBoxMeshPath, "boxmesh.xml (libccd)",
mjDSBL_NATIVECCD);
harness.RunBenchmark(state);
}
BENCHMARK(BM_BoxMesh_LibCCD);
BENCHMARK(BM_BoxMesh)->Arg(0)->Arg(kNumThreads);
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBox(benchmark::State& state) {
static TestHarness harness(kBoxBoxPath, "box.xml (BoxBox)", mjDSBL_NATIVECCD);
int nthread = state.range(0);
static TestHarness harness(kBoxBoxPath, "box.xml (BoxBox)");
harness.SetThreads(nthread);
harness.RunBenchmark(state);
}
BENCHMARK(BM_BoxBox);
BENCHMARK(BM_BoxBox)->Arg(0)->Arg(kNumThreads);
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBox_NativeCCD(benchmark::State& state) {
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBoxConvex(benchmark::State& state) {
int nthread = state.range(0);
mjCOLLISIONFUNC[mjGEOM_BOX][mjGEOM_BOX] = mjc_Convex;
static TestHarness harness(kBoxBoxPath, "box.xml (NativeCCD)");
harness.SetThreads(nthread);
harness.RunBenchmark(state);
mjCOLLISIONFUNC[mjGEOM_BOX][mjGEOM_BOX] = mjc_BoxBox;
}
BENCHMARK(BM_BoxBox_NativeCCD);
BENCHMARK(BM_BoxBoxConvex)->Arg(0)->Arg(kNumThreads);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_Ellipsoid_NativeCCD(benchmark::State& state) {
static TestHarness harness(kEllipsoidPath, "ellipsoid.xml (nativeccd)");
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Ellipsoid(benchmark::State& state) {
int nthread = state.range(0);
static TestHarness harness(kEllipsoidPath, "ellipsoid.xml");
harness.SetThreads(nthread);
harness.RunBenchmark(state);
}
BENCHMARK(BM_Ellipsoid_NativeCCD);
BENCHMARK(BM_Ellipsoid)->Arg(0)->Arg(kNumThreads);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_Ellipsoid_LibCCD(benchmark::State& state) {
static TestHarness harness(kEllipsoidPath, "ellipsoid.xml (libccd)",
mjDSBL_NATIVECCD);
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Mixed(benchmark::State& state) {
int nthread = state.range(0);
static TestHarness harness(kMixedPath, "mixed.xml");
harness.SetThreads(nthread);
harness.RunBenchmark(state);
}
BENCHMARK(BM_Ellipsoid_LibCCD);
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Mixed_NativeCCD(benchmark::State& state) {
static TestHarness harness(kMixedPath, "mixed.xml (nativeccd)");
harness.RunBenchmark(state);
}
BENCHMARK(BM_Mixed_NativeCCD);
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Mixed_LibCCD(benchmark::State& state) {
static TestHarness harness(kMixedPath, "mixed.xml (libccd)",
mjDSBL_NATIVECCD);
harness.RunBenchmark(state);
}
BENCHMARK(BM_Mixed_LibCCD);
BENCHMARK(BM_Mixed)->Arg(0)->Arg(kNumThreads);
} // namespace
} // namespace mujoco
+9 -17
View File
@@ -17,7 +17,6 @@
#include <benchmark/benchmark.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjthread.h>
#include <mujoco/mujoco.h>
#include "test/fixture.h"
@@ -30,6 +29,9 @@ static const int kNumWarmupSteps = 500;
// number of steps to benchmark (before resetting state)
static const int kBatchSize = 50;
// number of threads to test
static const int kNumThreads = 6;
void BM_StepHumanoid200(benchmark::State& state) {
int nthread = state.range(0);
std::string label = std::to_string(nthread) + " thread(s)";
@@ -43,10 +45,8 @@ void BM_StepHumanoid200(benchmark::State& state) {
model->opt.disableflags &= ~mjDSBL_ISLAND; // enable islands
mjData* data = mj_makeData(model);
mjThreadPool* threadpool = nullptr;
if (nthread > 1) {
threadpool = mju_threadPoolCreate(nthread);
mju_bindThreadPool(data, threadpool);
if (nthread) {
mju_threadpool(data, nthread);
}
// warm-up rollout to get a steady state
@@ -73,9 +73,6 @@ void BM_StepHumanoid200(benchmark::State& state) {
state.SetLabel(label);
state.SetItemsProcessed(state.iterations());
mj_deleteData(data);
if (threadpool) {
mju_threadPoolDestroy(threadpool);
}
}
void BM_Step22Humanoids(benchmark::State& state) {
@@ -90,10 +87,8 @@ void BM_Step22Humanoids(benchmark::State& state) {
model->opt.disableflags &= ~mjDSBL_ISLAND; // enable islands
mjData* data = mj_makeData(model);
mjThreadPool* threadpool = nullptr;
if (nthread > 1) {
threadpool = mju_threadPoolCreate(nthread);
mju_bindThreadPool(data, threadpool);
if (nthread) {
mju_threadpool(data, nthread);
}
// warm-up rollout to get a steady state
@@ -123,12 +118,9 @@ void BM_Step22Humanoids(benchmark::State& state) {
state.SetLabel(label);
state.SetItemsProcessed(state.iterations());
mj_deleteData(data);
if (threadpool) {
mju_threadPoolDestroy(threadpool);
}
}
BENCHMARK(BM_StepHumanoid200)->Arg(1)->Arg(6);
BENCHMARK(BM_Step22Humanoids)->Arg(1)->Arg(6);
BENCHMARK(BM_StepHumanoid200)->Arg(0)->Arg(kNumThreads);
BENCHMARK(BM_Step22Humanoids)->Arg(0)->Arg(kNumThreads);
} // namespace
} // namespace mujoco