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

2-3x speedup of sparse matrix squaring.

Browse Source 
Split symbolic and numeric phases for sparse `M'*diag*M` computation. Microseconds per call for the monolithic vs the split approach for the 100_humanoids and 2humanoid100 models:

```
+-------+------+----------+------------+---------+
| Model | Arch | Col (µs) | Split (µs) | Speedup |
+-------+------+----------+------------+---------+
| 2H100 | x86  | 238.3    | 74.5       | 3.2x    |
+-------+------+----------+------------+---------+
|       | ARM  | 111.6    | 53.2       | 2.1x    |
+-------+------+----------+------------+---------+
| 100H  | x86  | 1325.3   | 656.2      | 2.0x    |
+-------+------+----------+------------+---------+
|       | ARM  | 594.8    | 306.6      | 1.9x    |
+-------+------+----------+------------+---------+
```

PiperOrigin-RevId: 900154308
Change-Id: Ia6e9b8e196e2ed37b723a0faf60e9731303a9619
This commit is contained in:
Yuval Tassa
2026-04-15 07:19:03 -07:00
committed by Copybara-Service
parent 62cffb1536
commit a2d0e33c0f
8 changed files with 1301 additions and 626 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/benchmark/chol_benchmark_test.cc
+26 -11
View File
@@ -68,13 +68,15 @@ struct HessianData {
// D diagonal
std::vector<mjtNum> D;
int nefc;
void Setup(const mjModel* m, mjData* d) {
// initialize simulation state
mj_resetDataKeyframe(m, d, 0);
mj_forward(m, d);
nv = m->nv;
int nefc = d->nefc;
nefc = d->nefc;
// compute D corresponding to quad states
D.resize(nefc);
@@ -205,14 +207,27 @@ mjModel* GetModel() {
return m;
}
template <Size S>
HessianData& GetHessianData() {
static HessianData data;
static bool initialized = false;
if (!initialized) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
data.Setup(m, d);
mj_deleteData(d);
initialized = true;
}
return data;
}
// old implementation benchmark
template <Size S>
static void BM_chol_old(benchmark::State& state) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
HessianData hd;
hd.Setup(m, d);
HessianData& hd = GetHessianData<S>();
std::vector<mjtNum> L_work(hd.nL);
std::vector<int> L_colind_work(hd.nL);
@@ -239,8 +254,7 @@ static void BM_chol_symbolic(benchmark::State& state) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
HessianData hd;
hd.Setup(m, d);
HessianData& hd = GetHessianData<S>();
std::vector<int> L_colind_work(hd.nL);
std::vector<int> LT_rownnz_work(hd.nv);
@@ -266,8 +280,7 @@ static void BM_chol_numeric(benchmark::State& state) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
HessianData hd;
hd.Setup(m, d);
HessianData& hd = GetHessianData<S>();
std::vector<mjtNum> L_work(hd.nL);
std::vector<int> L_colind_work(hd.nL);
@@ -371,9 +384,10 @@ template <Size S>
static void BM_update_old(benchmark::State& state) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
mj_resetDataKeyframe(m, d, 0);
mj_forward(m, d);
HessianData hd;
hd.Setup(m, d);
HessianData& hd = GetHessianData<S>();
int nv = hd.nv;
@@ -433,9 +447,10 @@ template <Size S>
static void BM_update_new(benchmark::State& state) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
mj_resetDataKeyframe(m, d, 0);
mj_forward(m, d);
HessianData hd;
hd.Setup(m, d);
HessianData& hd = GetHessianData<S>();
int nv = hd.nv;