2-3x speedup of sparse matrix squaring.
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
committed by
Copybara-Service
Copybara-Service
parent
62cffb1536
commit
a2d0e33c0f
@@ -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;
|
||||
|
||||
|
||||
@@ -38,116 +38,7 @@ static const int kNumWarmupSteps = 500;
|
||||
|
||||
// ----------------------------- old functions --------------------------------
|
||||
|
||||
void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline(
|
||||
mjtNum* res, const mjtNum* mat, const mjtNum* matT, const mjtNum* diag,
|
||||
int nr, int nc, int* res_rownnz, int* res_rowadr, int* res_colind,
|
||||
const int* rownnz, const int* rowadr, const int* colind,
|
||||
const int* rowsuper, const int* rownnzT, const int* rowadrT,
|
||||
const int* colindT, const int* rowsuperT, mjData* d, int* unused) {
|
||||
mj_markStack(d);
|
||||
int* chain = mj_stackAllocInt(d, 2 * nc);
|
||||
mjtNum* buffer = mj_stackAllocNum(d, nc);
|
||||
|
||||
for (int r = 0; r < nc; r++) {
|
||||
res_rowadr[r] = r * nc;
|
||||
}
|
||||
|
||||
for (int r = 0; r < nc; r++) {
|
||||
if (rowsuperT && r > 0 && rowsuperT[r - 1] > 0) {
|
||||
res_rownnz[r] = res_rownnz[r - 1];
|
||||
memcpy(res_colind + res_rowadr[r], res_colind + res_rowadr[r - 1],
|
||||
res_rownnz[r] * sizeof(int));
|
||||
|
||||
if (rownnzT[r]) {
|
||||
res_colind[res_rowadr[r] + res_rownnz[r]] = r;
|
||||
res_rownnz[r]++;
|
||||
}
|
||||
} else {
|
||||
int nchain = 0;
|
||||
int inew = 0, iold = nc;
|
||||
int lastadded = -1;
|
||||
for (int i = 0; i < rownnzT[r]; i++) {
|
||||
int c = colindT[rowadrT[r] + i];
|
||||
if (rowsuper && lastadded >= 0 &&
|
||||
(c - lastadded) <= rowsuper[lastadded]) {
|
||||
continue;
|
||||
} else {
|
||||
lastadded = c;
|
||||
}
|
||||
|
||||
int adr = inew;
|
||||
inew = iold;
|
||||
iold = adr;
|
||||
|
||||
int nnewchain = 0;
|
||||
adr = 0;
|
||||
int end = rowadr[c] + rownnz[c];
|
||||
for (int adr1 = rowadr[c]; adr1 < end; adr1++) {
|
||||
int col_mat = colind[adr1];
|
||||
while (adr < nchain && chain[iold + adr] < col_mat &&
|
||||
chain[iold + adr] <= r) {
|
||||
chain[inew + nnewchain++] = chain[iold + adr++];
|
||||
}
|
||||
|
||||
if (col_mat > r) {
|
||||
break;
|
||||
}
|
||||
|
||||
if (adr < nchain && chain[iold + adr] == col_mat) {
|
||||
adr++;
|
||||
}
|
||||
chain[inew + nnewchain++] = col_mat;
|
||||
}
|
||||
|
||||
while (adr < nchain && chain[iold + adr] <= r) {
|
||||
chain[inew + nnewchain++] = chain[iold + adr++];
|
||||
}
|
||||
nchain = nnewchain;
|
||||
}
|
||||
res_rownnz[r] = nchain;
|
||||
if (nchain) {
|
||||
memcpy(res_colind + res_rowadr[r], chain + inew, nchain * sizeof(int));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (int r = 0; r < nc; r++) {
|
||||
int adr = res_rowadr[r];
|
||||
for (int i = 0; i < res_rownnz[r]; i++) {
|
||||
buffer[res_colind[adr + i]] = 0;
|
||||
}
|
||||
for (int i = 0; i < rownnzT[r]; i++) {
|
||||
int c = colindT[rowadrT[r] + i];
|
||||
mjtNum matTrc = matT[rowadrT[r] + i];
|
||||
if (diag) {
|
||||
matTrc *= diag[c];
|
||||
}
|
||||
|
||||
int end = rowadr[c] + rownnz[c];
|
||||
for (int adr = rowadr[c]; adr < end; adr++) {
|
||||
int adr1;
|
||||
if ((adr1 = colind[adr]) > r) {
|
||||
break;
|
||||
}
|
||||
buffer[adr1] += matTrc * mat[adr];
|
||||
}
|
||||
}
|
||||
adr = res_rowadr[r];
|
||||
for (int i = 0; i < res_rownnz[r]; i++) {
|
||||
res[adr + i] = buffer[res_colind[adr + i]];
|
||||
}
|
||||
}
|
||||
for (int r = 1; r < nc; r++) {
|
||||
int end = res_rowadr[r] + res_rownnz[r] - 1;
|
||||
for (int adr = res_rowadr[r]; adr < end; adr++) {
|
||||
int adr1 = res_rowadr[res_colind[adr]] + res_rownnz[res_colind[adr]]++;
|
||||
res[adr1] = res[adr];
|
||||
res_colind[adr1] = r;
|
||||
}
|
||||
}
|
||||
|
||||
mj_freeStack(d);
|
||||
}
|
||||
|
||||
// transpose sparse matrix (uncompressed)
|
||||
void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline(
|
||||
@@ -506,15 +397,27 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old(
|
||||
}
|
||||
BENCHMARK(BM_combineSparse_old);
|
||||
|
||||
enum class Size { H2_100, H100 };
|
||||
|
||||
template <Size S>
|
||||
const char* ModelPath() {
|
||||
if constexpr (S == Size::H2_100) {
|
||||
return "../test/benchmark/testdata/2humanoid100_chol.xml";
|
||||
} else {
|
||||
return "../test/benchmark/testdata/100_humanoids_chol.xml";
|
||||
}
|
||||
}
|
||||
|
||||
enum class Supernode {
|
||||
None,
|
||||
PostProcess,
|
||||
Inline
|
||||
};
|
||||
|
||||
template <Size S>
|
||||
static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func,
|
||||
Supernode super) {
|
||||
static mjModel* m = LoadModelFromPath("humanoid/humanoid100.xml");
|
||||
static mjModel* m = LoadModelFromPath(ModelPath<S>());
|
||||
|
||||
// force use of sparse matrices
|
||||
m->opt.jacobian = mjJAC_SPARSE;
|
||||
@@ -553,131 +456,67 @@ static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func,
|
||||
}
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_transposeSparse_old(benchmark::State& state) {
|
||||
BM_transposeSparse_2H100_old(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_transposeSparse(state, &transposeSparse_baseline, Supernode::None);
|
||||
BM_transposeSparse<Size::H2_100>(state, &transposeSparse_baseline,
|
||||
Supernode::None);
|
||||
}
|
||||
BENCHMARK(BM_transposeSparse_old);
|
||||
BENCHMARK(BM_transposeSparse_2H100_old);
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_transposeSparse_new(benchmark::State& state) {
|
||||
BM_transposeSparse_2H100_new(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_transposeSparse(state, &mju_transposeSparse, Supernode::None);
|
||||
BM_transposeSparse<Size::H2_100>(state, &mju_transposeSparse,
|
||||
Supernode::None);
|
||||
}
|
||||
BENCHMARK(BM_transposeSparse_new);
|
||||
BENCHMARK(BM_transposeSparse_2H100_new);
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_transposeSparse_superpost(benchmark::State& state) {
|
||||
BM_transposeSparse_2H100_superpost(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_transposeSparse(state, &mju_transposeSparse, Supernode::PostProcess);
|
||||
BM_transposeSparse<Size::H2_100>(state, &mju_transposeSparse,
|
||||
Supernode::PostProcess);
|
||||
}
|
||||
BENCHMARK(BM_transposeSparse_superpost);
|
||||
BENCHMARK(BM_transposeSparse_2H100_superpost);
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_transposeSparse_superinline(benchmark::State& state) {
|
||||
BM_transposeSparse_2H100_superinline(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_transposeSparse(state, &mju_transposeSparse, Supernode::Inline);
|
||||
}
|
||||
BENCHMARK(BM_transposeSparse_superinline);
|
||||
|
||||
static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
|
||||
static mjModel* m =
|
||||
LoadModelFromPath("../test/benchmark/testdata/2humanoid100.xml");
|
||||
|
||||
// force use of sparse matrices, Newton solver, no islands
|
||||
m->opt.jacobian = mjJAC_SPARSE;
|
||||
m->opt.solver = mjSOL_NEWTON;
|
||||
m->opt.disableflags |= mjDSBL_ISLAND;
|
||||
|
||||
mjData* d = mj_makeData(m);
|
||||
|
||||
// warm-up rollout to get a typical state
|
||||
while (d->time < 2) {
|
||||
mj_step(m, d);
|
||||
}
|
||||
|
||||
// allocate
|
||||
mj_markStack(d);
|
||||
mjtNum* H = mj_stackAllocNum(d, m->nv * m->nv);
|
||||
int* rownnz = mj_stackAllocInt(d, m->nv);
|
||||
int* rowadr = mj_stackAllocInt(d, m->nv);
|
||||
int* colind = mj_stackAllocInt(d, m->nv * m->nv);
|
||||
int* diagind = mj_stackAllocInt(d, m->nv);
|
||||
|
||||
// compute D corresponding to quad states
|
||||
mjtNum* D = mj_stackAllocNum(d, d->nefc);
|
||||
for (int i = 0; i < d->nefc; i++) {
|
||||
if (d->efc_state[i] == mjCNSTRSTATE_QUADRATIC) {
|
||||
D[i] = d->efc_D[i];
|
||||
} else {
|
||||
D[i] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
int* JT_rownnz = mj_stackAllocInt(d, m->nv);
|
||||
int* JT_rowadr = mj_stackAllocInt(d, m->nv);
|
||||
int* JT_rowsuper = mj_stackAllocInt(d, m->nv);
|
||||
int* JT_colind = mj_stackAllocInt(d, d->nJ);
|
||||
mjtNum* JT = mj_stackAllocNum(d, d->nJ);
|
||||
mju_transposeSparse(JT, d->efc_J, d->nefc, m->nv,
|
||||
JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
|
||||
|
||||
// time benchmark
|
||||
if (func) {
|
||||
mju_sqrMatTDSparseCount(rownnz, rowadr, m->nv,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
|
||||
JT_rownnz, JT_rowadr,
|
||||
JT_colind, nullptr, d, 1);
|
||||
|
||||
for (auto s : state) {
|
||||
// compute H = J'*D*J, compressed layout
|
||||
func(H, d->efc_J, JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
|
||||
JT_rownnz, JT_rowadr, JT_colind,
|
||||
JT_rowsuper, d, diagind);
|
||||
}
|
||||
} else {
|
||||
for (auto s : state) {
|
||||
// baseline depends on efc_J_rowsuper
|
||||
mju_superSparse(d->nefc, d->efc_J_rowsuper,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
|
||||
|
||||
// compute H = J'*D*J, uncompressed layout
|
||||
mju_sqrMatTDSparse_baseline(
|
||||
H, d->efc_J, JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, d->efc_J_rowsuper,
|
||||
JT_rownnz, JT_rowadr, JT_colind,
|
||||
JT_rowsuper, d, /*unused=*/nullptr);
|
||||
}
|
||||
}
|
||||
|
||||
// finalize
|
||||
mj_freeStack(d);
|
||||
mj_deleteData(d);
|
||||
state.SetItemsProcessed(state.iterations());
|
||||
BM_transposeSparse<Size::H2_100>(state, &mju_transposeSparse,
|
||||
Supernode::Inline);
|
||||
}
|
||||
BENCHMARK(BM_transposeSparse_2H100_superinline);
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_sqrMatTDSparse_col(benchmark::State& state) {
|
||||
BM_transposeSparse_100H_old(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTDSparse(state, &mju_sqrMatTDSparse);
|
||||
BM_transposeSparse<Size::H100>(state, &transposeSparse_baseline,
|
||||
Supernode::None);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTDSparse_col);
|
||||
BENCHMARK(BM_transposeSparse_100H_old);
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_sqrMatTDSparse_row(benchmark::State& state) {
|
||||
BM_transposeSparse_100H_new(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTDSparse(state, &mju_sqrMatTDSparse_row);
|
||||
BM_transposeSparse<Size::H100>(state, &mju_transposeSparse, Supernode::None);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTDSparse_row);
|
||||
BENCHMARK(BM_transposeSparse_100H_new);
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_sqrMatTDSparse_uncompressed(benchmark::State& state) {
|
||||
BM_transposeSparse_100H_superpost(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTDSparse(state, nullptr);
|
||||
BM_transposeSparse<Size::H100>(state, &mju_transposeSparse,
|
||||
Supernode::PostProcess);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTDSparse_uncompressed);
|
||||
BENCHMARK(BM_transposeSparse_100H_superpost);
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_transposeSparse_100H_superinline(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_transposeSparse<Size::H100>(state, &mju_transposeSparse,
|
||||
Supernode::Inline);
|
||||
}
|
||||
BENCHMARK(BM_transposeSparse_100H_superinline);
|
||||
|
||||
} // namespace
|
||||
} // namespace mujoco
|
||||
|
||||
@@ -0,0 +1,423 @@
|
||||
// Copyright 2026 DeepMind Technologies Limited
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
// Benchmarks for sparse matrix operations.
|
||||
|
||||
#include <cstring>
|
||||
#include <vector>
|
||||
|
||||
#include "benchmark/benchmark.h"
|
||||
#include <absl/base/attributes.h>
|
||||
#include <mujoco/mjdata.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "src/engine/engine_util_sparse.h"
|
||||
#include "test/fixture.h"
|
||||
|
||||
namespace mujoco {
|
||||
namespace {
|
||||
|
||||
// ================================ Test Data ==================================
|
||||
// Stores pre-computed sparse matrix inputs extracted from MuJoCo simulations.
|
||||
// Each benchmark computes its own outputs (H, L, etc.) from these inputs.
|
||||
|
||||
struct SparseTestData {
|
||||
// Dimensions
|
||||
int nv; // number of DoFs
|
||||
int nefc; // number of constraint rows
|
||||
int nJ; // nnz in J
|
||||
|
||||
// J (Jacobian) - nefc x nv sparse
|
||||
std::vector<mjtNum> J;
|
||||
std::vector<int> J_rownnz, J_rowadr, J_colind, J_rowsuper;
|
||||
|
||||
// J' (transpose)
|
||||
std::vector<mjtNum> JT;
|
||||
std::vector<int> JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper;
|
||||
|
||||
// D (diagonal weights for constraints)
|
||||
std::vector<mjtNum> D;
|
||||
|
||||
// M structure (mass matrix, lower triangle)
|
||||
std::vector<int> M_rownnz, M_rowadr, M_colind;
|
||||
|
||||
void Setup(const mjModel* m, mjData* d) {
|
||||
// initialize simulation state
|
||||
mj_resetDataKeyframe(m, d, 0);
|
||||
mj_step(m, d);
|
||||
mj_forward(m, d);
|
||||
nv = m->nv;
|
||||
nefc = d->nefc;
|
||||
nJ = d->nJ;
|
||||
|
||||
// copy J
|
||||
J.assign(d->efc_J, d->efc_J + nJ);
|
||||
J_rownnz.assign(d->efc_J_rownnz, d->efc_J_rownnz + nefc);
|
||||
J_rowadr.assign(d->efc_J_rowadr, d->efc_J_rowadr + nefc);
|
||||
J_colind.assign(d->efc_J_colind, d->efc_J_colind + nJ);
|
||||
J_rowsuper.assign(d->efc_J_rowsuper, d->efc_J_rowsuper + nefc);
|
||||
|
||||
// transpose J
|
||||
JT.assign(nJ, 0);
|
||||
JT_rownnz.assign(nv, 0);
|
||||
JT_rowadr.assign(nv, 0);
|
||||
JT_colind.assign(nJ, 0);
|
||||
JT_rowsuper.assign(nv, 0);
|
||||
mju_transposeSparse(JT.data(), J.data(), nefc, nv, JT_rownnz.data(),
|
||||
JT_rowadr.data(), JT_colind.data(), JT_rowsuper.data(),
|
||||
J_rownnz.data(), J_rowadr.data(), J_colind.data());
|
||||
|
||||
// compute D corresponding to quadratic constraint states
|
||||
D.resize(nefc);
|
||||
for (int i = 0; i < nefc; i++) {
|
||||
if (d->efc_state[i] == mjCNSTRSTATE_QUADRATIC) {
|
||||
D[i] = d->efc_D[i];
|
||||
} else {
|
||||
D[i] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
// copy M structure
|
||||
M_rownnz.assign(m->M_rownnz, m->M_rownnz + nv);
|
||||
M_rowadr.assign(m->M_rowadr, m->M_rowadr + nv);
|
||||
int nM = M_rowadr[nv - 1] + M_rownnz[nv - 1];
|
||||
M_colind.assign(m->M_colind, m->M_colind + nM);
|
||||
}
|
||||
};
|
||||
|
||||
// ================================ Model Sizes ================================
|
||||
|
||||
enum class Size { H2_100, H100 };
|
||||
|
||||
template <Size S>
|
||||
const char* ModelPath() {
|
||||
if constexpr (S == Size::H2_100) {
|
||||
return "../test/benchmark/testdata/2humanoid100_chol.xml";
|
||||
} else {
|
||||
return "../test/benchmark/testdata/100_humanoids_chol.xml";
|
||||
}
|
||||
}
|
||||
|
||||
template <Size S>
|
||||
mjModel* GetModel() {
|
||||
static mjModel* m = LoadModelFromPath(ModelPath<S>());
|
||||
m->opt.jacobian = mjJAC_SPARSE;
|
||||
m->opt.solver = mjSOL_NEWTON;
|
||||
m->opt.disableflags |= mjDSBL_ISLAND;
|
||||
return m;
|
||||
}
|
||||
|
||||
template <Size S>
|
||||
SparseTestData& GetData() {
|
||||
static SparseTestData 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;
|
||||
}
|
||||
|
||||
// ========================== Baseline Implementations =========================
|
||||
|
||||
|
||||
|
||||
// Baseline sqrMatTD (uncompressed layout, from old implementation)
|
||||
void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline(
|
||||
mjtNum* res, const mjtNum* mat, const mjtNum* matT, const mjtNum* diag,
|
||||
int nr, int nc, int* res_rownnz, int* res_rowadr, int* res_colind,
|
||||
const int* rownnz, const int* rowadr, const int* colind,
|
||||
const int* rowsuper, const int* rownnzT, const int* rowadrT,
|
||||
const int* colindT, const int* rowsuperT, mjData* d) {
|
||||
mj_markStack(d);
|
||||
int* chain = mj_stackAllocInt(d, 2 * nc);
|
||||
mjtNum* buffer = mj_stackAllocNum(d, nc);
|
||||
|
||||
for (int r = 0; r < nc; r++) {
|
||||
res_rowadr[r] = r * nc;
|
||||
}
|
||||
|
||||
for (int r = 0; r < nc; r++) {
|
||||
if (rowsuperT && r > 0 && rowsuperT[r - 1] > 0) {
|
||||
res_rownnz[r] = res_rownnz[r - 1];
|
||||
memcpy(res_colind + res_rowadr[r], res_colind + res_rowadr[r - 1],
|
||||
res_rownnz[r] * sizeof(int));
|
||||
if (rownnzT[r]) {
|
||||
res_colind[res_rowadr[r] + res_rownnz[r]] = r;
|
||||
res_rownnz[r]++;
|
||||
}
|
||||
} else {
|
||||
int nchain = 0;
|
||||
int inew = 0, iold = nc;
|
||||
int lastadded = -1;
|
||||
for (int i = 0; i < rownnzT[r]; i++) {
|
||||
int c = colindT[rowadrT[r] + i];
|
||||
if (rowsuper && lastadded >= 0 &&
|
||||
(c - lastadded) <= rowsuper[lastadded]) {
|
||||
continue;
|
||||
} else {
|
||||
lastadded = c;
|
||||
}
|
||||
|
||||
int adr = inew;
|
||||
inew = iold;
|
||||
iold = adr;
|
||||
|
||||
int nnewchain = 0;
|
||||
adr = 0;
|
||||
int end = rowadr[c] + rownnz[c];
|
||||
for (int adr1 = rowadr[c]; adr1 < end; adr1++) {
|
||||
int col_mat = colind[adr1];
|
||||
while (adr < nchain && chain[iold + adr] < col_mat &&
|
||||
chain[iold + adr] <= r) {
|
||||
chain[inew + nnewchain++] = chain[iold + adr++];
|
||||
}
|
||||
if (col_mat > r) {
|
||||
break;
|
||||
}
|
||||
if (adr < nchain && chain[iold + adr] == col_mat) {
|
||||
adr++;
|
||||
}
|
||||
chain[inew + nnewchain++] = col_mat;
|
||||
}
|
||||
|
||||
while (adr < nchain && chain[iold + adr] <= r) {
|
||||
chain[inew + nnewchain++] = chain[iold + adr++];
|
||||
}
|
||||
nchain = nnewchain;
|
||||
}
|
||||
res_rownnz[r] = nchain;
|
||||
if (nchain) {
|
||||
memcpy(res_colind + res_rowadr[r], chain + inew, nchain * sizeof(int));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (int r = 0; r < nc; r++) {
|
||||
int adr = res_rowadr[r];
|
||||
for (int i = 0; i < res_rownnz[r]; i++) {
|
||||
buffer[res_colind[adr + i]] = 0;
|
||||
}
|
||||
for (int i = 0; i < rownnzT[r]; i++) {
|
||||
int c = colindT[rowadrT[r] + i];
|
||||
mjtNum matTrc = matT[rowadrT[r] + i];
|
||||
if (diag) {
|
||||
matTrc *= diag[c];
|
||||
}
|
||||
|
||||
int end = rowadr[c] + rownnz[c];
|
||||
for (int adr2 = rowadr[c]; adr2 < end; adr2++) {
|
||||
int adr1;
|
||||
if ((adr1 = colind[adr2]) > r) {
|
||||
break;
|
||||
}
|
||||
buffer[adr1] += matTrc * mat[adr2];
|
||||
}
|
||||
}
|
||||
adr = res_rowadr[r];
|
||||
for (int i = 0; i < res_rownnz[r]; i++) {
|
||||
res[adr + i] = buffer[res_colind[adr + i]];
|
||||
}
|
||||
}
|
||||
for (int r = 1; r < nc; r++) {
|
||||
int end = res_rowadr[r] + res_rownnz[r] - 1;
|
||||
for (int adr = res_rowadr[r]; adr < end; adr++) {
|
||||
int adr1 = res_rowadr[res_colind[adr]] + res_rownnz[res_colind[adr]]++;
|
||||
res[adr1] = res[adr];
|
||||
res_colind[adr1] = r;
|
||||
}
|
||||
}
|
||||
|
||||
mj_freeStack(d);
|
||||
}
|
||||
|
||||
|
||||
|
||||
// ========================== SqrMatTD Benchmarks ==============================
|
||||
|
||||
enum class SqrMatTDVariant {
|
||||
kBaseline,
|
||||
kRow,
|
||||
kCol,
|
||||
kSplitCol
|
||||
};
|
||||
|
||||
template <Size S>
|
||||
static void BM_sqrMatTD_impl(benchmark::State& state, SqrMatTDVariant variant) {
|
||||
SparseTestData& data = GetData<S>();
|
||||
mjModel* m = GetModel<S>();
|
||||
mjData* d = mj_makeData(m);
|
||||
|
||||
int nv = data.nv;
|
||||
|
||||
// nothing to benchmark if no constraints
|
||||
if (data.nefc == 0) {
|
||||
for (auto s : state) {}
|
||||
mj_deleteData(d);
|
||||
return;
|
||||
}
|
||||
|
||||
// allocate H output (uncompressed for baseline, compressed for others)
|
||||
int max_nnz = (variant == SqrMatTDVariant::kBaseline) ? nv * nv : 0;
|
||||
std::vector<mjtNum> H;
|
||||
std::vector<int> H_rownnz(nv);
|
||||
std::vector<int> H_rowadr(nv);
|
||||
std::vector<int> H_colind;
|
||||
std::vector<int> diagind(nv);
|
||||
|
||||
if (variant == SqrMatTDVariant::kBaseline) {
|
||||
H.resize(max_nnz);
|
||||
H_colind.resize(max_nnz);
|
||||
} else if (variant == SqrMatTDVariant::kSplitCol ||
|
||||
variant == SqrMatTDVariant::kCol) {
|
||||
// use symbolic to count nnz
|
||||
int nH = mju_sqrMatTDSparseSymbolic(
|
||||
H_rownnz.data(), H_rowadr.data(), nullptr, nullptr,
|
||||
data.nefc, nv, data.J_rownnz.data(), data.J_rowadr.data(),
|
||||
data.J_colind.data(), data.JT_rownnz.data(), data.JT_rowadr.data(),
|
||||
data.JT_colind.data(), data.JT_rowsuper.data(), d);
|
||||
H.resize(nH);
|
||||
H_colind.resize(nH);
|
||||
} else {
|
||||
// row: use Count (lower triangle only)
|
||||
mju_sqrMatTDSparseCount(
|
||||
H_rownnz.data(), H_rowadr.data(), nv, data.J_rownnz.data(),
|
||||
data.J_rowadr.data(), data.J_colind.data(), data.JT_rownnz.data(),
|
||||
data.JT_rowadr.data(), data.JT_colind.data(), nullptr, d, 0);
|
||||
int nH = H_rowadr[nv - 1] + H_rownnz[nv - 1];
|
||||
H.resize(nH);
|
||||
H_colind.resize(nH);
|
||||
}
|
||||
|
||||
for (auto s : state) {
|
||||
switch (variant) {
|
||||
case SqrMatTDVariant::kBaseline:
|
||||
mju_superSparse(data.nefc, data.J_rowsuper.data(), data.J_rownnz.data(),
|
||||
data.J_rowadr.data(), data.J_colind.data());
|
||||
mju_sqrMatTDSparse_baseline(
|
||||
H.data(), data.J.data(), data.JT.data(), data.D.data(), data.nefc,
|
||||
nv, H_rownnz.data(), H_rowadr.data(), H_colind.data(),
|
||||
data.J_rownnz.data(), data.J_rowadr.data(), data.J_colind.data(),
|
||||
data.J_rowsuper.data(), data.JT_rownnz.data(),
|
||||
data.JT_rowadr.data(), data.JT_colind.data(),
|
||||
data.JT_rowsuper.data(), d);
|
||||
break;
|
||||
case SqrMatTDVariant::kRow:
|
||||
mju_sqrMatTDSparseCount(
|
||||
H_rownnz.data(), H_rowadr.data(), nv, data.J_rownnz.data(),
|
||||
data.J_rowadr.data(), data.J_colind.data(), data.JT_rownnz.data(),
|
||||
data.JT_rowadr.data(), data.JT_colind.data(), nullptr, d, 0);
|
||||
mju_sqrMatTDSparse_row(
|
||||
H.data(), data.J.data(), data.JT.data(), data.D.data(), data.nefc,
|
||||
nv, H_rownnz.data(), H_rowadr.data(), H_colind.data(),
|
||||
data.J_rownnz.data(), data.J_rowadr.data(), data.J_colind.data(),
|
||||
nullptr, data.JT_rownnz.data(), data.JT_rowadr.data(),
|
||||
data.JT_colind.data(), data.JT_rowsuper.data(), d, nullptr);
|
||||
break;
|
||||
case SqrMatTDVariant::kCol:
|
||||
mju_sqrMatTDSparseCount(
|
||||
H_rownnz.data(), H_rowadr.data(), nv, data.J_rownnz.data(),
|
||||
data.J_rowadr.data(), data.J_colind.data(), data.JT_rownnz.data(),
|
||||
data.JT_rowadr.data(), data.JT_colind.data(), nullptr, d, 0);
|
||||
mju_sqrMatTDSparse(
|
||||
H.data(), data.J.data(), data.JT.data(), data.D.data(), data.nefc,
|
||||
nv, H_rownnz.data(), H_rowadr.data(), H_colind.data(),
|
||||
data.J_rownnz.data(), data.J_rowadr.data(), data.J_colind.data(),
|
||||
nullptr, data.JT_rownnz.data(), data.JT_rowadr.data(),
|
||||
data.JT_colind.data(), data.JT_rowsuper.data(), d, nullptr);
|
||||
break;
|
||||
|
||||
case SqrMatTDVariant::kSplitCol:
|
||||
mju_sqrMatTDSparseSymbolic(
|
||||
H_rownnz.data(), H_rowadr.data(), nullptr, nullptr,
|
||||
data.nefc, nv, data.J_rownnz.data(), data.J_rowadr.data(),
|
||||
data.J_colind.data(), data.JT_rownnz.data(), data.JT_rowadr.data(),
|
||||
data.JT_colind.data(), data.JT_rowsuper.data(), d);
|
||||
mju_sqrMatTDSparseSymbolic(
|
||||
H_rownnz.data(), H_rowadr.data(), H_colind.data(), nullptr,
|
||||
data.nefc, nv, data.J_rownnz.data(), data.J_rowadr.data(),
|
||||
data.J_colind.data(), data.JT_rownnz.data(), data.JT_rowadr.data(),
|
||||
data.JT_colind.data(), data.JT_rowsuper.data(), d);
|
||||
mju_sqrMatTDSparseNumeric(
|
||||
H.data(), nv, H_rownnz.data(), H_rowadr.data(),
|
||||
H_colind.data(), nullptr, data.J.data(), data.J_rownnz.data(),
|
||||
data.J_rowadr.data(), data.J_colind.data(), data.JT.data(),
|
||||
data.JT_rownnz.data(), data.JT_rowadr.data(), data.JT_colind.data(),
|
||||
data.JT_rowsuper.data(), data.D.data(), d);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
mj_deleteData(d);
|
||||
state.SetItemsProcessed(state.iterations());
|
||||
}
|
||||
|
||||
void BM_sqrMatTD_2H100_baseline(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTD_impl<Size::H2_100>(state, SqrMatTDVariant::kBaseline);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTD_2H100_baseline);
|
||||
|
||||
void BM_sqrMatTD_2H100_row(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTD_impl<Size::H2_100>(state, SqrMatTDVariant::kRow);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTD_2H100_row);
|
||||
|
||||
void BM_sqrMatTD_2H100_col(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTD_impl<Size::H2_100>(state, SqrMatTDVariant::kCol);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTD_2H100_col);
|
||||
|
||||
void BM_sqrMatTD_2H100_splitCol(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTD_impl<Size::H2_100>(state, SqrMatTDVariant::kSplitCol);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTD_2H100_splitCol);
|
||||
|
||||
void BM_sqrMatTD_100H_baseline(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTD_impl<Size::H100>(state, SqrMatTDVariant::kBaseline);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTD_100H_baseline);
|
||||
|
||||
void BM_sqrMatTD_100H_row(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTD_impl<Size::H100>(state, SqrMatTDVariant::kRow);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTD_100H_row);
|
||||
|
||||
void BM_sqrMatTD_100H_col(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTD_impl<Size::H100>(state, SqrMatTDVariant::kCol);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTD_100H_col);
|
||||
|
||||
void BM_sqrMatTD_100H_splitCol(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_sqrMatTD_impl<Size::H100>(state, SqrMatTDVariant::kSplitCol);
|
||||
}
|
||||
BENCHMARK(BM_sqrMatTD_100H_splitCol);
|
||||
|
||||
} // namespace
|
||||
} // namespace mujoco
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
benchmark::Initialize(&argc, argv);
|
||||
benchmark::RunSpecifiedBenchmarks();
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user