Refactor sparse Cholesky factorization into symbolic and numeric phases.
The new symbolic function is a generalization of the function it replaces. In this CL it takes two unused temp arrays. The actual change in behavior happens in the followup.
New benchmark test output below ("L" is 2 humanoids and 100 free objects, "XL" is 100 humanoids). Note that `symbolic` is only ever called once per Newton iteration, while `numeric` is sometimes called multiple times (when the rank-1 update fails), hence timing them separately is valuable.
```
Benchmark Time(ns) CPU(ns) Iterations
--------------------------------------------------------------
BM_old_L_mean 84382 84703 19547 11.807k items/s
BM_symbolic_L_mean 16345 16381 88414 61.055k items/s
BM_numeric_L_mean 10986 10994 120000 90.999k items/s
BM_old_XL_mean 1241208 1244212 1200 803.924 items/s
BM_symbolic_XL_mean 130917 131042 12720 7.631k items/s
BM_numeric_XL_mean 77004 76767 21116 13.029k items/s
```
PiperOrigin-RevId: 846704054
Change-Id: Ib0c365724d63bf2b81606ca5353756a6496c3a26
This commit is contained in:
Yuval Tassa
committed by
Copybara-Service
Copybara-Service
parent
d1fd11bccd
commit
45b0153067
@@ -1558,16 +1558,14 @@ static void MakeHessian(mjData* d, mjCGContext* ctx) {
|
||||
HT_rownnz, HT_rowadr, HT_colind, NULL,
|
||||
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind);
|
||||
|
||||
// count total and row non-zeros of reverse-Cholesky factor L
|
||||
ctx->nL = mju_cholFactorCount(ctx->L_rownnz, HT_rownnz, HT_rowadr, HT_colind, nv, d);
|
||||
// count total and row non-zeros of reverse-Cholesky factors L and LT
|
||||
int* LT_rownnz_temp = mjSTACKALLOC(d, nv, int);
|
||||
int* LT_rowadr_temp = mjSTACKALLOC(d, nv, int);
|
||||
ctx->nL = mju_cholFactorSymbolic(NULL, ctx->L_rownnz, ctx->L_rowadr, NULL,
|
||||
LT_rownnz_temp, LT_rowadr_temp, NULL,
|
||||
HT_rownnz, HT_rowadr, HT_colind, nv, d);
|
||||
mj_freeStack(d);
|
||||
|
||||
// compute L row addresses: rowadr = cumsum(rownnz)
|
||||
ctx->L_rowadr[0] = 0;
|
||||
for (int r=1; r < nv; r++) {
|
||||
ctx->L_rowadr[r] = ctx->L_rowadr[r-1] + ctx->L_rownnz[r-1];
|
||||
}
|
||||
|
||||
// allocate L_colind, L, Lcone
|
||||
ctx->L_colind = mjSTACKALLOC(d, ctx->nL, int);
|
||||
ctx->L = mjSTACKALLOC(d, ctx->nL, mjtNum);
|
||||
|
||||
+152
-23
@@ -187,21 +187,52 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
|
||||
return rank;
|
||||
}
|
||||
|
||||
|
||||
// precount row non-zeros of reverse-Cholesky factor L, return total non-zeros
|
||||
// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package'
|
||||
// reads pattern from upper triangle
|
||||
int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind,
|
||||
int n, mjData* d) {
|
||||
// symbolic reverse-Cholesky: compute both L (CSR) and LT (CSC) structures
|
||||
// if L_colind is NULL, perform counting logic (fill rownnz/rowadr arrays and return total nnz)
|
||||
// if L_colind is not NULL, assume rownnz/rowadr are precomputed and fill colind/map arrays
|
||||
// reads pattern from upper triangle
|
||||
// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package'
|
||||
int mju_cholFactorSymbolic(int* restrict L_colind, int* restrict L_rownnz, int* restrict L_rowadr,
|
||||
int* restrict LT_colind, int* restrict LT_rownnz,
|
||||
int* restrict LT_rowadr, int* restrict LT_map,
|
||||
const int* rownnz, const int* rowadr, const int* colind, int n,
|
||||
mjData* d) {
|
||||
mj_markStack(d);
|
||||
int* parent = mjSTACKALLOC(d, n, int);
|
||||
int* flag = mjSTACKALLOC(d, n, int);
|
||||
int* restrict parent = mjSTACKALLOC(d, n, int);
|
||||
int* restrict flag = mjSTACKALLOC(d, n, int);
|
||||
int* restrict cursor = NULL;
|
||||
int* LT_write = NULL;
|
||||
|
||||
// filling phase: initialize write positions
|
||||
if (L_colind) {
|
||||
cursor = mjSTACKALLOC(d, n, int);
|
||||
LT_write = mjSTACKALLOC(d, n, int);
|
||||
for (int r = 0; r < n; r++) {
|
||||
cursor[r] = L_rowadr[r] + L_rownnz[r] - 2; // end of row r (before diagonal)
|
||||
LT_write[r] = LT_rowadr[r]; // start of LT row r
|
||||
}
|
||||
}
|
||||
|
||||
// loop over rows in reverse order
|
||||
for (int r = n - 1; r >= 0; r--) {
|
||||
parent[r] = -1;
|
||||
flag[r] = r;
|
||||
L_rownnz[r] = 1; // start with 1 for diagonal
|
||||
|
||||
// counting phase: start with 1 for diagonal
|
||||
if (!L_colind) {
|
||||
L_rownnz[r] = 1;
|
||||
LT_rownnz[r] = 1;
|
||||
}
|
||||
|
||||
// filling phase: write diagonals
|
||||
else {
|
||||
int diag_idx = L_rowadr[r] + L_rownnz[r] - 1;
|
||||
L_colind[diag_idx] = r;
|
||||
int write_idx = LT_write[r];
|
||||
LT_colind[write_idx] = r;
|
||||
LT_map[write_idx] = diag_idx;
|
||||
LT_write[r]++;
|
||||
}
|
||||
|
||||
// loop over non-zero columns of upper triangle
|
||||
int start = rowadr[r];
|
||||
@@ -221,8 +252,23 @@ int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr, con
|
||||
parent[i] = r;
|
||||
}
|
||||
|
||||
// increment non-zeros, flag row i, advance to parent
|
||||
L_rownnz[i]++;
|
||||
// counting phase: increment non-zeros
|
||||
if (!L_colind) {
|
||||
L_rownnz[i]++;
|
||||
LT_rownnz[r]++;
|
||||
}
|
||||
|
||||
// filling phase: write L[i, r] and LT[r, i]
|
||||
else {
|
||||
int L_idx = cursor[i];
|
||||
cursor[i]--;
|
||||
L_colind[L_idx] = r;
|
||||
LT_colind[LT_write[r]] = i;
|
||||
LT_map[LT_write[r]] = L_idx;
|
||||
LT_write[r]++;
|
||||
}
|
||||
|
||||
// flag row i, advance to parent
|
||||
flag[i] = r;
|
||||
i = parent[i];
|
||||
}
|
||||
@@ -231,15 +277,98 @@ int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr, con
|
||||
|
||||
mj_freeStack(d);
|
||||
|
||||
// sum up all row non-zeros
|
||||
// counting phase: compute row addresses, add up total non-zeros
|
||||
int nnz = 0;
|
||||
for (int r = 0; r < n; r++) {
|
||||
nnz += L_rownnz[r];
|
||||
if (!L_colind) {
|
||||
nnz = L_rownnz[0];
|
||||
L_rowadr[0] = 0;
|
||||
LT_rowadr[0] = 0;
|
||||
for (int r = 1; r < n; r++) {
|
||||
L_rowadr[r] = L_rowadr[r - 1] + L_rownnz[r - 1];
|
||||
LT_rowadr[r] = LT_rowadr[r - 1] + LT_rownnz[r - 1];
|
||||
nnz += L_rownnz[r];
|
||||
}
|
||||
}
|
||||
|
||||
return nnz;
|
||||
}
|
||||
|
||||
// numeric reverse-Cholesky: compute L values given fixed sparsity pattern, returns rank
|
||||
// L_colind must already contain the correct sparsity pattern (from mju_cholFactorSymbolic)
|
||||
// LT_map[k] gives index in L for LT_colind[k]
|
||||
int mju_cholFactorNumeric(mjtNum* restrict L, int n, mjtNum mindiag,
|
||||
const int* L_rownnz, const int* L_rowadr, const int* L_colind,
|
||||
const int* LT_rownnz, const int* LT_rowadr, const int* LT_colind,
|
||||
const int* LT_map, const mjtNum* H,
|
||||
const int* H_rownnz, const int* H_rowadr, const int* H_colind,
|
||||
mjData* d) {
|
||||
int rank = n;
|
||||
|
||||
// single-row dense accumulator
|
||||
mj_markStack(d);
|
||||
mjtNum* restrict dense = mjSTACKALLOC(d, n, mjtNum);
|
||||
mju_zero(dense, n);
|
||||
|
||||
// backpass over rows
|
||||
for (int r = n - 1; r >= 0; r--) {
|
||||
// scatter H[r, 0:r] into dense
|
||||
mju_scatter(dense, H + H_rowadr[r], H_colind + H_rowadr[r], H_rownnz[r]);
|
||||
|
||||
// accumulate updates from rows c > r where L[c,r] != 0
|
||||
// use CSC transpose: LT column r contains rows that have column r
|
||||
// start from k=1 to skip the diagonal entry (LT_colind[LT_adr] = r)
|
||||
int LT_adr = LT_rowadr[r];
|
||||
int LT_nnz = LT_rownnz[r];
|
||||
for (int k = 1; k < LT_nnz; k++) {
|
||||
int c = LT_colind[LT_adr + k]; // row c has L[c,r] != 0, c > r guaranteed
|
||||
|
||||
// get L[c,r] index directly from LT_map
|
||||
int L_cr_idx = LT_map[LT_adr + k];
|
||||
mjtNum L_cr = L[L_cr_idx];
|
||||
|
||||
// get row c info
|
||||
int c_adr = L_rowadr[c];
|
||||
|
||||
// dense[j] -= L[c,r] * L[c,j] for all j <= r in L[c]
|
||||
// L_cr_idx - c_adr gives the position of r in row c
|
||||
int num_cols = L_cr_idx - c_adr + 1;
|
||||
const int* colptr = L_colind + c_adr;
|
||||
const mjtNum* Lptr = L + c_adr;
|
||||
for (int i = 0; i < num_cols; i++) {
|
||||
dense[colptr[i]] -= L_cr * Lptr[i];
|
||||
}
|
||||
}
|
||||
|
||||
// factor row r diagonal, handle rank-deficient case
|
||||
mjtNum diag = dense[r];
|
||||
if (diag < mindiag) {
|
||||
diag = mindiag;
|
||||
rank--;
|
||||
}
|
||||
|
||||
// scale off-diagonals
|
||||
mjtNum L_rr = mju_sqrt(diag);
|
||||
mjtNum L_rr_inv = 1.0 / L_rr;
|
||||
int L_adr = L_rowadr[r];
|
||||
int L_nnz = L_rownnz[r];
|
||||
const int* colptr = L_colind + L_adr;
|
||||
mjtNum* Lptr = L + L_adr;
|
||||
for (int i = 0; i < L_nnz - 1; i++) {
|
||||
Lptr[i] = dense[colptr[i]] * L_rr_inv;
|
||||
}
|
||||
|
||||
// store diagonal
|
||||
L[L_adr + L_nnz - 1] = L_rr;
|
||||
|
||||
// clear dense workspace
|
||||
for (int i = 0; i < L_nnz; i++) {
|
||||
dense[colptr[i]] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
mj_freeStack(d);
|
||||
return rank;
|
||||
}
|
||||
|
||||
// sparse reverse-order Cholesky solve
|
||||
void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n,
|
||||
@@ -528,7 +657,7 @@ void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int n
|
||||
mju_zero(res, ntotal*ntotal);
|
||||
|
||||
// sparse part
|
||||
for(int i=0; i < nsparse; i++) {
|
||||
for (int i=0; i < nsparse; i++) {
|
||||
// number of non-zeros left of (i,i)
|
||||
int width = mjMIN(i, nband-1);
|
||||
|
||||
@@ -537,13 +666,13 @@ void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int n
|
||||
}
|
||||
|
||||
// dense part
|
||||
for(int i=nsparse; i < ntotal; i++) {
|
||||
for (int i=nsparse; i < ntotal; i++) {
|
||||
mju_copy(res + i*ntotal, mat + nsparse*nband + (i-nsparse)*ntotal, i+1);
|
||||
}
|
||||
|
||||
// make symmetric
|
||||
if (flg_sym) {
|
||||
for(int i=0; i < ntotal; i++) {
|
||||
for (int i=0; i < ntotal; i++) {
|
||||
for (int j=i+1; j < ntotal; j++) {
|
||||
res[i*ntotal + j] = res[j*ntotal + i];
|
||||
}
|
||||
@@ -557,7 +686,7 @@ void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int n
|
||||
int nsparse = ntotal-ndense;
|
||||
|
||||
// sparse part
|
||||
for(int i=0; i < nsparse; i++) {
|
||||
for (int i=0; i < nsparse; i++) {
|
||||
// number of non-zeros left of (i,i)
|
||||
int width = mjMIN(i, nband-1);
|
||||
|
||||
@@ -566,7 +695,7 @@ void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int n
|
||||
}
|
||||
|
||||
// dense part
|
||||
for(int i=nsparse; i < ntotal; i++) {
|
||||
for (int i=nsparse; i < ntotal; i++) {
|
||||
mju_copy(res + nsparse*nband + (i-nsparse)*ntotal, mat + i*ntotal, i+1);
|
||||
}
|
||||
}
|
||||
@@ -578,13 +707,13 @@ void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
|
||||
int nsparse = ntotal-ndense;
|
||||
|
||||
// handle multiple vectors
|
||||
for(int j=0; j < nvec; j++ ) {
|
||||
for (int j=0; j < nvec; j++) {
|
||||
// precompute pointer to corresponding vector in vec and res
|
||||
const mjtNum* vec_j = vec + ntotal*j;
|
||||
mjtNum* res_j = res + ntotal*j;
|
||||
|
||||
// sparse part
|
||||
for(int i=0; i < nsparse; i++) {
|
||||
for (int i=0; i < nsparse; i++) {
|
||||
int width = mjMIN(i+1, nband);
|
||||
int adr = i*nband + nband - width;
|
||||
int offset = mjMAX(0, i-nband+1);
|
||||
@@ -596,7 +725,7 @@ void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
|
||||
}
|
||||
|
||||
// dense part
|
||||
for(int i=nsparse; i < ntotal; i++) {
|
||||
for (int i=nsparse; i < ntotal; i++) {
|
||||
int adr = nsparse*nband + (i-nsparse)*ntotal;
|
||||
res_j[i] = mju_dot(mat+adr, vec_j, i+1);
|
||||
if (flg_sym) {
|
||||
@@ -808,7 +937,7 @@ int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum ma
|
||||
mju_normalize4(quat);
|
||||
}
|
||||
|
||||
// sort eigenvalues in decreasing order (bubblesort: 0, 1, 0)
|
||||
// sort eigenvalues in decreasing order (bubble sort: 0, 1, 0)
|
||||
for (int j=0; j < 3; j++) {
|
||||
int j1 = j%2; // lead index
|
||||
|
||||
|
||||
@@ -36,9 +36,25 @@ MJAPI int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus);
|
||||
MJAPI int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
|
||||
int* rownnz, const int* rowadr, int* colind, mjData* d);
|
||||
|
||||
// precount row non-zeros of reverse-Cholesky factor L, return total
|
||||
MJAPI int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr,
|
||||
const int* colind, int n, mjData* d);
|
||||
// symbolic reverse-Cholesky: compute both L (CSR) and LT (CSC) structures
|
||||
// if L_colind is NULL, perform counting logic (fill rownnz/rowadr arrays and return total nnz)
|
||||
// if L_colind is not NULL, assume rownnz/rowadr are precomputed and fill colind/map arrays
|
||||
// reads pattern from upper triangle
|
||||
// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package'
|
||||
MJAPI int mju_cholFactorSymbolic(int* L_colind, int* L_rownnz, int* L_rowadr,
|
||||
int* LT_colind, int* LT_rownnz, int* LT_rowadr, int* LT_map,
|
||||
const int* rownnz, const int* rowadr, const int* colind,
|
||||
int n, mjData* d);
|
||||
|
||||
// numeric reverse-Cholesky: compute L values given fixed sparsity pattern, returns rank
|
||||
// L_colind must already contain the correct sparsity pattern (from mju_cholFactorSymbolic)
|
||||
// LT_map[k] gives index in L for LT_colind[k]
|
||||
MJAPI int mju_cholFactorNumeric(mjtNum* L, int n, mjtNum mindiag,
|
||||
const int* L_rownnz, const int* L_rowadr, const int* L_colind,
|
||||
const int* LT_rownnz, const int* LT_rowadr, const int* LT_colind,
|
||||
const int* LT_map, const mjtNum* H,
|
||||
const int* H_rownnz, const int* H_rowadr, const int* H_colind,
|
||||
mjData* d);
|
||||
|
||||
// sparse reverse-order Cholesky solve
|
||||
void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n,
|
||||
|
||||
@@ -0,0 +1,333 @@
|
||||
// Copyright 2025 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.
|
||||
|
||||
// A benchmark for comparing old vs new Cholesky factorization implementations.
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstring>
|
||||
#include <vector>
|
||||
|
||||
#include <benchmark/benchmark.h>
|
||||
#include <mujoco/mjdata.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "src/engine/engine_support.h"
|
||||
#include "src/engine/engine_util_solve.h"
|
||||
#include "src/engine/engine_util_sparse.h"
|
||||
#include "test/fixture.h"
|
||||
|
||||
namespace mujoco {
|
||||
namespace {
|
||||
|
||||
// Helper to compute H = M + J'*D*J using sparse matrices
|
||||
struct HessianData {
|
||||
int nv;
|
||||
int nL;
|
||||
int nH;
|
||||
|
||||
// H sparse structure
|
||||
std::vector<mjtNum> H;
|
||||
std::vector<int> H_rownnz;
|
||||
std::vector<int> H_rowadr;
|
||||
std::vector<int> H_colind;
|
||||
|
||||
// H transpose for symbolics
|
||||
std::vector<int> HT_rownnz;
|
||||
std::vector<int> HT_rowadr;
|
||||
std::vector<int> HT_colind;
|
||||
|
||||
// L factor structure
|
||||
std::vector<int> L_rownnz;
|
||||
std::vector<int> L_rowadr;
|
||||
|
||||
// L initial values for BM_chol_old (lower triangle of H, zero-filled for
|
||||
// fill-in)
|
||||
std::vector<mjtNum> L_init;
|
||||
std::vector<int> L_rownnz_init;
|
||||
std::vector<int> L_colind_init;
|
||||
|
||||
// J transpose
|
||||
std::vector<mjtNum> JT;
|
||||
std::vector<int> JT_rownnz;
|
||||
std::vector<int> JT_rowadr;
|
||||
std::vector<int> JT_colind;
|
||||
std::vector<int> JT_rowsuper;
|
||||
|
||||
// D diagonal
|
||||
std::vector<mjtNum> D;
|
||||
|
||||
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;
|
||||
|
||||
// compute D corresponding to quad 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;
|
||||
}
|
||||
}
|
||||
|
||||
// transpose J
|
||||
JT.resize(d->nJ);
|
||||
JT_rownnz.resize(nv);
|
||||
JT_rowadr.resize(nv);
|
||||
JT_colind.resize(d->nJ);
|
||||
JT_rowsuper.resize(nv);
|
||||
mju_transposeSparse(JT.data(), d->efc_J, nefc, nv, JT_rownnz.data(),
|
||||
JT_rowadr.data(), JT_colind.data(), JT_rowsuper.data(),
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
|
||||
|
||||
// count H sparsity: nH from J'*D*J
|
||||
H_rownnz.resize(nv);
|
||||
H_rowadr.resize(nv);
|
||||
mju_sqrMatTDSparseCount(H_rownnz.data(), H_rowadr.data(), nv,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
|
||||
JT_rownnz.data(), JT_rowadr.data(),
|
||||
JT_colind.data(), nullptr, d, 1);
|
||||
|
||||
// add M elements to the H row counts and addresses
|
||||
for (int r = 0; r < nv; r++) {
|
||||
H_rownnz[r] += m->M_rownnz[r];
|
||||
}
|
||||
H_rowadr[0] = 0;
|
||||
for (int r = 1; r < nv; r++) {
|
||||
H_rowadr[r] = H_rowadr[r - 1] + H_rownnz[r - 1];
|
||||
}
|
||||
nH = H_rowadr[nv - 1] + H_rownnz[nv - 1];
|
||||
|
||||
// allocate H and colind with proper sparse size, zero-initialize H
|
||||
H.assign(nH, 0);
|
||||
H_colind.assign(nH, 0);
|
||||
|
||||
// reset rownnz for filling (sqrMatTDSparse will fill it again)
|
||||
std::fill(H_rownnz.begin(), H_rownnz.end(), 0);
|
||||
|
||||
// recount just J'*D*J (without M shift)
|
||||
mju_sqrMatTDSparseCount(H_rownnz.data(), H_rowadr.data(), nv,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
|
||||
JT_rownnz.data(), JT_rowadr.data(),
|
||||
JT_colind.data(), nullptr, d, 1);
|
||||
|
||||
// add shift for M to rowadr
|
||||
int shift = 0;
|
||||
for (int r = 0; r < nv - 1; r++) {
|
||||
shift += m->M_rownnz[r];
|
||||
H_rowadr[r + 1] += shift;
|
||||
}
|
||||
|
||||
// compute H = J'*D*J
|
||||
mju_sqrMatTDSparse(H.data(), d->efc_J, JT.data(), D.data(), nefc, nv,
|
||||
H_rownnz.data(), H_rowadr.data(), H_colind.data(),
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
|
||||
nullptr, JT_rownnz.data(), JT_rowadr.data(),
|
||||
JT_colind.data(), JT_rowsuper.data(), d, nullptr);
|
||||
|
||||
// add M to H using mj_addM
|
||||
mj_addM(m, d, H.data(), H_rownnz.data(), H_rowadr.data(), H_colind.data());
|
||||
|
||||
// transpose H for symbolic
|
||||
HT_rownnz.resize(nv);
|
||||
HT_rowadr.resize(nv);
|
||||
HT_colind.resize(nH);
|
||||
mju_transposeSparse(nullptr, nullptr, nv, nv, HT_rownnz.data(),
|
||||
HT_rowadr.data(), HT_colind.data(), nullptr,
|
||||
H_rownnz.data(), H_rowadr.data(), H_colind.data());
|
||||
|
||||
// count L fill-in (also counts LT structure)
|
||||
L_rownnz.resize(nv);
|
||||
L_rowadr.resize(nv);
|
||||
std::vector<int> LT_rownnz_temp(nv);
|
||||
std::vector<int> LT_rowadr_temp(nv);
|
||||
nL = mju_cholFactorSymbolic(
|
||||
nullptr, L_rownnz.data(), L_rowadr.data(), nullptr,
|
||||
LT_rownnz_temp.data(), LT_rowadr_temp.data(), nullptr, HT_rownnz.data(),
|
||||
HT_rowadr.data(), HT_colind.data(), nv, d);
|
||||
|
||||
// precompute initial L state for BM_chol_old
|
||||
// extract lower triangle of H into L format, zero-fill for fill-in
|
||||
L_init.assign(nL, 0);
|
||||
L_colind_init.assign(nL, 0);
|
||||
L_rownnz_init.resize(nv);
|
||||
for (int r = 0; r < nv; r++) {
|
||||
int l_adr = L_rowadr[r];
|
||||
int h_adr = H_rowadr[r];
|
||||
int lower_nnz = 0;
|
||||
for (int i = 0; i < H_rownnz[r]; i++) {
|
||||
int col = H_colind[h_adr + i];
|
||||
if (col <= r) {
|
||||
L_init[l_adr + lower_nnz] = H[h_adr + i];
|
||||
L_colind_init[l_adr + lower_nnz] = col;
|
||||
lower_nnz++;
|
||||
}
|
||||
}
|
||||
L_rownnz_init[r] = lower_nnz;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// ----------------------------- benchmark ------------------------------------
|
||||
|
||||
enum class Size { L, XL };
|
||||
|
||||
template <Size S>
|
||||
const char* ModelPath() {
|
||||
if constexpr (S == Size::L) {
|
||||
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;
|
||||
return m;
|
||||
}
|
||||
|
||||
// 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);
|
||||
|
||||
std::vector<mjtNum> L_work(hd.nL);
|
||||
std::vector<int> L_colind_work(hd.nL);
|
||||
std::vector<int> L_rownnz_work(hd.nv);
|
||||
|
||||
for (auto s : state) {
|
||||
// fast reset using memcpy from precomputed initial state
|
||||
std::memcpy(L_work.data(), hd.L_init.data(), hd.nL * sizeof(mjtNum));
|
||||
std::memcpy(L_colind_work.data(), hd.L_colind_init.data(),
|
||||
hd.nL * sizeof(int));
|
||||
std::memcpy(L_rownnz_work.data(), hd.L_rownnz_init.data(),
|
||||
hd.nv * sizeof(int));
|
||||
mju_cholFactorSparse(L_work.data(), hd.nv, mjMINVAL, L_rownnz_work.data(),
|
||||
hd.L_rowadr.data(), L_colind_work.data(), d);
|
||||
}
|
||||
|
||||
mj_deleteData(d);
|
||||
state.SetItemsProcessed(state.iterations());
|
||||
}
|
||||
|
||||
// new symbolic implementation benchmark
|
||||
template <Size S>
|
||||
static void BM_chol_symbolic(benchmark::State& state) {
|
||||
mjModel* m = GetModel<S>();
|
||||
mjData* d = mj_makeData(m);
|
||||
|
||||
HessianData hd;
|
||||
hd.Setup(m, d);
|
||||
|
||||
std::vector<int> L_colind_work(hd.nL);
|
||||
std::vector<int> LT_rownnz_work(hd.nv);
|
||||
std::vector<int> LT_rowadr_work(hd.nv);
|
||||
std::vector<int> LT_colind_work(hd.nL);
|
||||
std::vector<int> LT_pos_work(hd.nL);
|
||||
|
||||
for (auto s : state) {
|
||||
mju_cholFactorSymbolic(L_colind_work.data(), hd.L_rownnz.data(),
|
||||
hd.L_rowadr.data(), LT_colind_work.data(),
|
||||
LT_rownnz_work.data(), LT_rowadr_work.data(),
|
||||
LT_pos_work.data(), hd.HT_rownnz.data(),
|
||||
hd.HT_rowadr.data(), hd.HT_colind.data(), hd.nv, d);
|
||||
}
|
||||
|
||||
mj_deleteData(d);
|
||||
state.SetItemsProcessed(state.iterations());
|
||||
}
|
||||
|
||||
// new numeric implementation benchmark
|
||||
template <Size S>
|
||||
static void BM_chol_numeric(benchmark::State& state) {
|
||||
mjModel* m = GetModel<S>();
|
||||
mjData* d = mj_makeData(m);
|
||||
|
||||
HessianData hd;
|
||||
hd.Setup(m, d);
|
||||
|
||||
std::vector<mjtNum> L_work(hd.nL);
|
||||
std::vector<int> L_colind_work(hd.nL);
|
||||
std::vector<int> LT_rownnz_work(hd.nv);
|
||||
std::vector<int> LT_rowadr_work(hd.nv);
|
||||
std::vector<int> LT_colind_work(hd.nL);
|
||||
std::vector<int> LT_pos_work(hd.nL);
|
||||
|
||||
// symbolic setup (not benchmarked)
|
||||
mju_cholFactorSymbolic(L_colind_work.data(), hd.L_rownnz.data(),
|
||||
hd.L_rowadr.data(), LT_colind_work.data(),
|
||||
LT_rownnz_work.data(), LT_rowadr_work.data(),
|
||||
LT_pos_work.data(), hd.HT_rownnz.data(),
|
||||
hd.HT_rowadr.data(), hd.HT_colind.data(), hd.nv, d);
|
||||
|
||||
for (auto s : state) {
|
||||
mju_cholFactorNumeric(
|
||||
L_work.data(), hd.nv, mjMINVAL, hd.L_rownnz.data(), hd.L_rowadr.data(),
|
||||
L_colind_work.data(), LT_rownnz_work.data(), LT_rowadr_work.data(),
|
||||
LT_colind_work.data(), LT_pos_work.data(), hd.H.data(),
|
||||
hd.H_rownnz.data(), hd.H_rowadr.data(), hd.H_colind.data(), d);
|
||||
}
|
||||
|
||||
mj_deleteData(d);
|
||||
state.SetItemsProcessed(state.iterations());
|
||||
}
|
||||
|
||||
void BM_old_L(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_chol_old<Size::L>(state);
|
||||
}
|
||||
BENCHMARK(BM_old_L);
|
||||
|
||||
void BM_symbolic_L(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_chol_symbolic<Size::L>(state);
|
||||
}
|
||||
BENCHMARK(BM_symbolic_L);
|
||||
|
||||
void BM_numeric_L(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_chol_numeric<Size::L>(state);
|
||||
}
|
||||
BENCHMARK(BM_numeric_L);
|
||||
|
||||
void BM_old_XL(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_chol_old<Size::XL>(state);
|
||||
}
|
||||
BENCHMARK(BM_old_XL);
|
||||
|
||||
void BM_symbolic_XL(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_chol_symbolic<Size::XL>(state);
|
||||
}
|
||||
BENCHMARK(BM_symbolic_XL);
|
||||
|
||||
void BM_numeric_XL(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_chol_numeric<Size::XL>(state);
|
||||
}
|
||||
BENCHMARK(BM_numeric_XL);
|
||||
|
||||
} // namespace
|
||||
} // namespace mujoco
|
||||
+10812
File diff suppressed because one or more lines are too long
+829
File diff suppressed because one or more lines are too long
@@ -28,6 +28,7 @@
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "src/engine/engine_util_blas.h"
|
||||
#include "src/engine/engine_util_misc.h"
|
||||
#include "src/engine/engine_util_sparse.h"
|
||||
#include "test/fixture.h"
|
||||
|
||||
namespace mujoco {
|
||||
@@ -256,7 +257,7 @@ TEST_F(BoxQPTest, BoundedQP) {
|
||||
maxiter, mingrad, backtrack,
|
||||
minstep, armijo, log, logsz);
|
||||
|
||||
// EXPECT_TRUE(false) << log; // uncomment to print `log` to error log
|
||||
// ADD_FAILURE() << log; // uncomment to print `log` to error log
|
||||
|
||||
// check solution
|
||||
EXPECT_GT(nfree, -1);
|
||||
@@ -653,73 +654,72 @@ TEST_F(BandMatrixTest, Solve) {
|
||||
|
||||
using EngineUtilSolveTest = MujocoTest;
|
||||
|
||||
TEST_F(EngineUtilSolveTest, MjuCholFactorNNZ) {
|
||||
TEST_F(EngineUtilSolveTest, MjuCholFactorSymbolic) {
|
||||
mjModel* model = LoadModelFromString("<mujoco/>");
|
||||
mjData* d = mj_makeData(model);
|
||||
|
||||
int nA = 2;
|
||||
mjtNum matA[4] = {1, 0,
|
||||
0, 1};
|
||||
mjtNum sparseA[4];
|
||||
int rownnzA[2];
|
||||
int rowadrA[2];
|
||||
int colindA[4];
|
||||
int rownnzA_factor[2];
|
||||
mju_dense2sparse(sparseA, matA, nA, nA, rownnzA, rowadrA, colindA, 4);
|
||||
int nnzA = mju_cholFactorCount(rownnzA_factor,
|
||||
rownnzA, rowadrA, colindA, nA, d);
|
||||
// Test matrix (upper triangular, representing symmetric matrix):
|
||||
// [10 1 2 3]
|
||||
// [ 0 10 0 0]
|
||||
// [ 0 0 10 1]
|
||||
// [ 0 0 0 10]
|
||||
//
|
||||
// mju_cholFactorSymbolic reads entries where col >= row
|
||||
|
||||
EXPECT_EQ(nnzA, 2);
|
||||
EXPECT_THAT(AsVector(rownnzA_factor, 2), ElementsAre(1, 1));
|
||||
int n = 4;
|
||||
mjtNum H[16] = {10, 1, 2, 3, 0, 10, 0, 0, 0, 0, 10, 1, 0, 0, 0, 10};
|
||||
|
||||
int nB = 3;
|
||||
mjtNum matB[9] = {10, 1, 0,
|
||||
0, 10, 1,
|
||||
0, 0, 10};
|
||||
mjtNum sparseB[9];
|
||||
int rownnzB[3];
|
||||
int rowadrB[3];
|
||||
int colindB[9];
|
||||
int rownnzB_factor[3];
|
||||
mju_dense2sparse(sparseB, matB, nB, nB, rownnzB, rowadrB, colindB, 9);
|
||||
int nnzB = mju_cholFactorCount(rownnzB_factor,
|
||||
rownnzB, rowadrB, colindB, nB, d);
|
||||
// convert to sparse
|
||||
mjtNum sparseH[16];
|
||||
int H_rownnz[4], H_rowadr[4], H_colind[16];
|
||||
mju_dense2sparse(sparseH, H, n, n, H_rownnz, H_rowadr, H_colind, 16);
|
||||
|
||||
EXPECT_EQ(nnzB, 5);
|
||||
EXPECT_THAT(AsVector(rownnzB_factor, 3), ElementsAre(1, 2, 2));
|
||||
// phase 1: counting
|
||||
int L_rownnz[4], L_rowadr[4];
|
||||
int LT_rownnz[4], LT_rowadr[4];
|
||||
int nnz = mju_cholFactorSymbolic(nullptr, L_rownnz, L_rowadr, nullptr,
|
||||
LT_rownnz, LT_rowadr, nullptr, H_rownnz,
|
||||
H_rowadr, H_colind, n, d);
|
||||
|
||||
int nC = 3;
|
||||
mjtNum matC[9] = {10, 1, 0,
|
||||
0, 10, 0,
|
||||
0, 0, 10};
|
||||
mjtNum sparseC[9];
|
||||
int rownnzC[3];
|
||||
int rowadrC[3];
|
||||
int colindC[9];
|
||||
int rownnzC_factor[3];
|
||||
mju_dense2sparse(sparseC, matC, nC, nC, rownnzC, rowadrC, colindC, 9);
|
||||
int nnzC = mju_cholFactorCount(rownnzC_factor,
|
||||
rownnzC, rowadrC, colindC, nC, d);
|
||||
// verify counting phase outputs
|
||||
EXPECT_EQ(nnz, 8);
|
||||
// L (CSR) structure for reverse Cholesky (rows filled in reverse order)
|
||||
EXPECT_THAT(AsVector(L_rownnz, 4), ElementsAre(1, 2, 2, 3));
|
||||
EXPECT_THAT(AsVector(L_rowadr, 4), ElementsAre(0, 1, 3, 5));
|
||||
// LT (CSC) structure: transpose of L
|
||||
EXPECT_THAT(AsVector(LT_rownnz, 4), ElementsAre(4, 1, 2, 1));
|
||||
EXPECT_THAT(AsVector(LT_rowadr, 4), ElementsAre(0, 4, 5, 7));
|
||||
|
||||
EXPECT_EQ(nnzC, 4);
|
||||
EXPECT_THAT(AsVector(rownnzC_factor, 3), ElementsAre(1, 2, 1));
|
||||
// phase 2: filling
|
||||
int L_colind[8], LT_colind[8], LT_pos[8];
|
||||
mju_cholFactorSymbolic(L_colind, L_rownnz, L_rowadr, LT_colind, LT_rownnz,
|
||||
LT_rowadr, LT_pos, H_rownnz, H_rowadr, H_colind, n, d);
|
||||
|
||||
int nD = 4;
|
||||
mjtNum matD[16] = {10, 1, 2, 3,
|
||||
0, 10, 0, 0,
|
||||
0, 0, 10, 1,
|
||||
0, 0, 0, 10};
|
||||
mjtNum sparseD[16];
|
||||
int rownnzD[4];
|
||||
int rowadrD[4];
|
||||
int colindD[16];
|
||||
int rownnzD_factor[4];
|
||||
mju_dense2sparse(sparseD, matD, nD, nD, rownnzD, rowadrD, colindD, 16);
|
||||
int nnzD = mju_cholFactorCount(rownnzD_factor,
|
||||
rownnzD, rowadrD, colindD, nD, d);
|
||||
// verify L_colind
|
||||
EXPECT_THAT(AsVector(L_colind, 8), ElementsAre(0, 0, 1, 0, 2, 0, 2, 3));
|
||||
// Explanation (reverse Cholesky builds from bottom to top):
|
||||
// Row 0 (1 entry): diagonal 0
|
||||
// Row 1 (2 entries): col 0, then diagonal 1
|
||||
// Row 2 (2 entries): col 0, then diagonal 2
|
||||
// Row 3 (3 entries): col 0, col 2, then diagonal 3
|
||||
|
||||
EXPECT_EQ(nnzD, 8);
|
||||
EXPECT_THAT(AsVector(rownnzD_factor, 4), ElementsAre(1, 2, 2, 3));
|
||||
// verify LT_colind: transpose of L
|
||||
EXPECT_THAT(AsVector(LT_colind, 8), ElementsAre(0, 1, 2, 3, 1, 2, 3, 3));
|
||||
// Explanation:
|
||||
// Column 0 (4 entries): rows 0, 1, 2, 3 (all have L[row, 0] != 0)
|
||||
// Column 1 (1 entry): row 1
|
||||
// Column 2 (2 entries): rows 2, 3
|
||||
// Column 3 (1 entry): row 3
|
||||
|
||||
// verify LT_pos: for each entry in LT, should point to correct position in L
|
||||
for (int c = 0; c < n; c++) {
|
||||
int adr = LT_rowadr[c];
|
||||
for (int k = 0; k < LT_rownnz[c]; k++) {
|
||||
int L_idx = LT_pos[adr + k];
|
||||
EXPECT_EQ(L_colind[L_idx], c)
|
||||
<< "LT_pos mismatch at column " << c << ", entry " << k;
|
||||
}
|
||||
}
|
||||
|
||||
mj_deleteData(d);
|
||||
mj_deleteModel(model);
|
||||
@@ -938,5 +938,101 @@ TEST_F(EngineUtilSolveTest, MjuCholUpdateSparse) {
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
// Test that mju_cholFactorSymbolic + mju_cholFactorNumeric produces identical
|
||||
// results to the reference implementation mju_cholFactorSparse
|
||||
TEST_F(EngineUtilSolveTest, CholFactorSymbolicNumeric) {
|
||||
mjModel* model = LoadModelFromString("<mujoco/>");
|
||||
mjData* d = mj_makeData(model);
|
||||
|
||||
// test matrix with fill-in: upper triangle structure
|
||||
int n = 4;
|
||||
mjtNum H[16] = {10, 1, 2, 3, 1, 10, 0, 0, 2, 0, 10, 1, 3, 0, 1, 10};
|
||||
|
||||
// convert to sparse (lower triangle only)
|
||||
mjtNum sparseH[16];
|
||||
int H_rownnz[4], H_rowadr[4], H_colind[16];
|
||||
mju_dense2sparse(sparseH, H, n, n, H_rownnz, H_rowadr, H_colind, 16);
|
||||
|
||||
// transpose for upper triangle (needed by cholFactorSymbolic)
|
||||
int HT_rownnz[4], HT_rowadr[4], HT_colind[16];
|
||||
mju_transposeSparse(nullptr, nullptr, n, n, HT_rownnz, HT_rowadr, HT_colind,
|
||||
nullptr, H_rownnz, H_rowadr, H_colind);
|
||||
|
||||
// count fill-in (also computes LT structure)
|
||||
int L_rownnz[4], L_rowadr[4];
|
||||
int LT_rownnz[4], LT_rowadr[4];
|
||||
int nnz = mju_cholFactorSymbolic(nullptr, L_rownnz, L_rowadr, nullptr,
|
||||
LT_rownnz, LT_rowadr, nullptr, HT_rownnz,
|
||||
HT_rowadr, HT_colind, n, d);
|
||||
|
||||
// filling phase: compute L_colind, LT_colind, and LT_pos
|
||||
int L_colind[16], LT_colind[16], LT_pos[16];
|
||||
mju_cholFactorSymbolic(L_colind, L_rownnz, L_rowadr, LT_colind, LT_rownnz,
|
||||
LT_rowadr, LT_pos, HT_rownnz, HT_rowadr, HT_colind, n,
|
||||
d);
|
||||
|
||||
// verify LT structure matches what we'd get from a separate transpose
|
||||
int LT_rownnz_ref[4], LT_rowadr_ref[4], LT_colind_ref[16];
|
||||
mju_transposeSparse(nullptr, nullptr, n, n, LT_rownnz_ref, LT_rowadr_ref,
|
||||
LT_colind_ref, nullptr, L_rownnz, L_rowadr, L_colind);
|
||||
|
||||
// LT structure should match
|
||||
EXPECT_THAT(AsVector(LT_rownnz, 4),
|
||||
ElementsAre(LT_rownnz_ref[0], LT_rownnz_ref[1], LT_rownnz_ref[2],
|
||||
LT_rownnz_ref[3]));
|
||||
EXPECT_THAT(AsVector(LT_rowadr, 4),
|
||||
ElementsAre(LT_rowadr_ref[0], LT_rowadr_ref[1], LT_rowadr_ref[2],
|
||||
LT_rowadr_ref[3]));
|
||||
|
||||
// verify LT_colind and LT_pos match
|
||||
for (int r = 0; r < n; r++) {
|
||||
int adr = LT_rowadr[r];
|
||||
for (int k = 0; k < LT_rownnz[r]; k++) {
|
||||
int L_idx = LT_pos[adr + k]; // index in L array
|
||||
// verify L_colind at this position is indeed r
|
||||
EXPECT_EQ(L_colind[L_idx], r)
|
||||
<< "LT_pos mismatch at LT[" << r << ", " << k << "]";
|
||||
}
|
||||
}
|
||||
|
||||
// numeric factorization using new function
|
||||
mjtNum L_new[16];
|
||||
int rank_new = mju_cholFactorNumeric(
|
||||
L_new, n, 1e-10, L_rownnz, L_rowadr, L_colind, LT_rownnz, LT_rowadr,
|
||||
LT_colind, LT_pos, sparseH, H_rownnz, H_rowadr, H_colind, d);
|
||||
|
||||
// reference implementation: copy sparse H into L_ref, then factor in-place
|
||||
mjtNum L_ref[16];
|
||||
int L_ref_rownnz[4], L_ref_colind[16];
|
||||
for (int r = 0; r < n; r++) {
|
||||
int nnz_r = H_rownnz[r];
|
||||
// count lower triangle elements for this row
|
||||
int lower_nnz = 0;
|
||||
for (int i = 0; i < nnz_r; i++) {
|
||||
if (H_colind[H_rowadr[r] + i] <= r) {
|
||||
L_ref[L_rowadr[r] + lower_nnz] = sparseH[H_rowadr[r] + i];
|
||||
L_ref_colind[L_rowadr[r] + lower_nnz] = H_colind[H_rowadr[r] + i];
|
||||
lower_nnz++;
|
||||
}
|
||||
}
|
||||
L_ref_rownnz[r] = lower_nnz;
|
||||
}
|
||||
int rank_ref = mju_cholFactorSparse(L_ref, n, 1e-10, L_ref_rownnz, L_rowadr,
|
||||
L_ref_colind, d);
|
||||
|
||||
// compare results
|
||||
EXPECT_EQ(rank_new, rank_ref);
|
||||
EXPECT_EQ(rank_new, n);
|
||||
|
||||
// compare L values
|
||||
mjtNum eps = 1e-10;
|
||||
for (int i = 0; i < nnz; i++) {
|
||||
EXPECT_NEAR(L_new[i], L_ref[i], eps) << "mismatch at index " << i;
|
||||
}
|
||||
|
||||
mj_deleteData(d);
|
||||
mj_deleteModel(model);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mujoco
|
||||
|
||||
Reference in New Issue
Block a user