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

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_sensor.c
+17 -28
View File
@@ -35,8 +35,7 @@
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
#include "engine/engine_util_spatial.h"
#include "thread/thread_pool.h"
#include "thread/thread_task.h"
#include "engine/engine_thread.h"
@@ -64,8 +63,6 @@ mjPARTIAL_SORT(ContactSelect, ContactInfo, ContactInfoCompare);
// arguments for parallel tactile sensor computation
typedef struct mjTactileTaskArgs_ {
const mjModel* m;
mjData* d;
int sensor_id;
int mesh_id;
int geom_id;
@@ -80,10 +77,8 @@ typedef struct mjTactileTaskArgs_ {
// worker function for parallel tactile computation over taxel batches
static void* tactile_taxel_batch(void* args) {
static void* tactile_taxel_batch(const mjModel* m, mjData* d, void* args) {
mjTactileTaskArgs* t = (mjTactileTaskArgs*)args;
const mjModel* m = t->m;
mjData* d = t->d;
int mesh_id = t->mesh_id;
int geom_id = t->geom_id;
int parent_weld = t->parent_weld;
@@ -193,6 +188,12 @@ static void* tactile_taxel_batch(void* args) {
}
static void tactileTask(const mjModel* m, mjData* d, void* arg, int thread_id, int task_id) {
mjTactileTaskArgs* args_array = (mjTactileTaskArgs*)arg;
tactile_taxel_batch(m, d, &args_array[task_id]);
}
// apply cutoff to sensor i, clamping values in data buffer
static void apply_cutoff(const mjModel* m, int i, mjtNum* data) {
mjtNum cutoff = m->sensor_cutoff[i];
@@ -1261,18 +1262,15 @@ static void mj_computeSensorAcc(const mjModel* m, mjData* d, int i, mjtNum* sens
// threshold for parallelization (taxel count below which sequential is faster)
const int kTactileParallelThreshold = 1000;
// parallel path: use threadpool to process taxel batches
if (d->threadpool && ncon >= kTactileParallelThreshold) {
int nthreads = mju_threadPoolNumberOfThreads((mjThreadPool*)d->threadpool);
int batch_size = (ncon + nthreads - 1) / nthreads;
int ntasks = (ncon + batch_size - 1) / batch_size;
// parallel path: use mj_batch to process taxel batches
int nthread = mju_numThread(d);
if (nthread > 0 && ncon >= kTactileParallelThreshold) {
int batch_size = (ncon + nthread - 1) / nthread;
int ntask = (ncon + batch_size - 1) / batch_size;
mjTask* tasks = mjSTACKALLOC(d, ntasks, mjTask);
mjTactileTaskArgs* task_args = mjSTACKALLOC(d, ntasks, mjTactileTaskArgs);
mjTactileTaskArgs* task_args = mjSTACKALLOC(d, ntask, mjTactileTaskArgs);
for (int t = 0; t < ntasks; t++) {
task_args[t].m = m;
task_args[t].d = d;
for (int t = 0; t < ntask; t++) {
task_args[t].sensor_id = i;
task_args[t].mesh_id = mesh_id;
task_args[t].geom_id = geom_id;
@@ -1283,22 +1281,13 @@ static void mj_computeSensorAcc(const mjModel* m, mjData* d, int i, mjtNum* sens
task_args[t].start_taxel = t * batch_size;
task_args[t].end_taxel = mju_min((t+1) * batch_size, ncon);
task_args[t].forcesT = forcesT;
mju_defaultTask(&tasks[t]);
tasks[t].func = tactile_taxel_batch;
tasks[t].args = &task_args[t];
mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[t]);
}
for (int t = 0; t < ntasks; t++) {
mju_taskJoin(&tasks[t]);
}
mju_dispatch(m, d, tactileTask, task_args, ntask);
}
// sequential path: call tactile_taxel_batch with full range
else {
mjTactileTaskArgs args;
args.m = m;
args.d = d;
args.sensor_id = i;
args.mesh_id = mesh_id;
args.geom_id = geom_id;
@@ -1309,7 +1298,7 @@ static void mj_computeSensorAcc(const mjModel* m, mjData* d, int i, mjtNum* sens
args.start_taxel = 0;
args.end_taxel = ncon;
args.forcesT = forcesT;
tactile_taxel_batch(&args);
tactile_taxel_batch(m, d, &args);
}
// compute sensor output