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

Precompute rotated stiffness matrix before CG in implicit flex integrator.

Browse Source 
This gives a 3x speedup in implicitfast.

Also cleanup old code that was used in the dense factorization of the stiffness matrix before we switched to CG.

PiperOrigin-RevId: 915900315
Change-Id: Id6973c4bfd7d371a43ec6db982703969b23a3550
This commit is contained in:
Alessio Quaglino
2026-05-15 02:57:18 -07:00
committed by Copybara-Service
parent 85bf2b9660
commit 7ca1bc6a56
5 changed files with 104 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_derivative.c
+88 -101
View File
@@ -864,20 +864,14 @@ static void addJTBJ_mulSparse(const mjModel* m, mjData* d, mjtNum* res, const mj
}
// operation type for flex interpolation derivative kernel
typedef enum {
mjFLEXOP_VEC, // res += J'*K*J*vec
mjFLEXOP_ADDH // H -= J'*K*J to H (dense)
} mjtFlexOp;
// shared kernel for flex interpolation derivatives, scale = s1 + s2*damping
// op: operation type (VEC, or ADDH)
// res: output vector (VEC) or banded H matrix (ADDH)
// vec: input vector for VEC operation, NULL otherwise
// dof_indices, ndof, nband: DOF mapping and band width for ADDH, ignored otherwise
static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
// res: output vector (res += J'*K*J*vec), or NULL for cache-only mode
// vec: input vector, or NULL for cache-only mode
// K_rot_cache: if non-NULL, use pre-cached K_rot values instead of computing them
// K_rot_out: if non-NULL and K_rot_cache is NULL, store computed K_rot values here
static void mjd_flexInterp_kernel(const mjModel* m, mjData* d,
mjtNum* res, const mjtNum* vec, mjtNum s1, mjtNum s2,
const int* dof_indices, int ndof, int nband) {
const mjtNum* K_rot_cache, mjtNum* K_rot_out) {
int nv = m->nv;
// compute upper bounds across all interpolated flexes
@@ -909,15 +903,6 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
// single unconditional markStack
mj_markStack(d);
// global2local map for ADDH
int* global2local = mjSTACKALLOC(d, nv, int);
if (op == mjFLEXOP_ADDH) {
mju_fillInt(global2local, -1, nv);
for (int i=0; i<ndof; i++) {
global2local[dof_indices[i]] = i;
}
}
// per-flex node positions (upper bound)
mjtNum* xpos = mjSTACKALLOC(d, 3*max_nodenum, mjtNum);
@@ -935,14 +920,6 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
int* chain_colind = mjSTACKALLOC(d, nv, int);
mjtNum* blk_jac = mjSTACKALLOC(d, 3*nv, mjtNum);
// ADDH-specific allocations (upper bound)
mjtNum* J_reduced = NULL;
mjtNum* KJ = NULL;
if (op == mjFLEXOP_ADDH) {
J_reduced = mjSTACKALLOC(d, max_dim_c*ndof, mjtNum);
KJ = mjSTACKALLOC(d, max_dim_c*ndof, mjtNum);
}
// loop over flexes
for (int f=0; f < m->nflex; f++) {
// only process flex_interp
@@ -1012,52 +989,85 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
continue;
}
// gather element-local node positions
int gindices[125]; // max npe = 125 for quadratic 3D
mjtNum quat[4];
if (shell_mode) {
mju_flexGatherFaceState(order, cx, cy, cz, fe, xpos, NULL, NULL,
xpos_c, NULL, NULL, gindices, quat);
// use cached K_rot or compute from scratch
int gindices[125]; // element node indices (max npe = 125 for quadratic 3D)
int krot_adr = stiffnessadr + fe * dim_e * dim_e;
if (K_rot_cache) {
// read K_rot from cache and apply scale
for (int i = 0; i < dim_e*dim_e; i++) {
K_rot_cell[i] = scale * K_rot_cache[krot_adr + i];
}
// recompute gindices (cheap index-only call, no rotation)
if (shell_mode) {
mju_flexGatherFaceState(order, cx, cy, cz, fe, NULL, NULL, NULL,
NULL, NULL, NULL, gindices, NULL);
} else {
int ci = fe / (cy * cz);
int cj = (fe / cz) % cy;
int ck = fe % cz;
mju_flexGatherCellState(order, cy, cz, ci, cj, ck, NULL, NULL, NULL,
NULL, NULL, NULL, gindices, NULL);
}
} else {
int ci = fe / (cy * cz);
int cj = (fe / cz) % cy;
int ck = fe % cz;
mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos, NULL, NULL,
xpos_c, NULL, NULL, gindices, quat);
}
// gather element-local node positions and rotation
mjtNum quat[4];
if (shell_mode) {
mju_flexGatherFaceState(order, cx, cy, cz, fe, xpos, NULL, NULL,
xpos_c, NULL, NULL, gindices, quat);
} else {
int ci = fe / (cy * cz);
int cj = (fe / cz) % cy;
int ck = fe % cz;
mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos, NULL, NULL,
xpos_c, NULL, NULL, gindices, quat);
}
// R = R_global2local, RT = R_local2global
mjtNum R[9], RT[9];
mju_quat2Mat(R, quat);
mju_transpose(RT, R, 3, 3);
// R = R_global2local, RT = R_local2global
mjtNum R[9], RT[9];
mju_quat2Mat(R, quat);
mju_transpose(RT, R, 3, 3);
// compute K_rot = RT * K_elem * R (block-wise)
mju_zero(K_rot_cell, dim_e*dim_e);
for (int a = 0; a < npe; a++) {
for (int b = 0; b < npe; b++) {
mjtNum blk[9], tmp[9];
// compute K_rot = RT * K_elem * R (block-wise)
mju_zero(K_rot_cell, dim_e*dim_e);
for (int a = 0; a < npe; a++) {
for (int b = 0; b < npe; b++) {
mjtNum blk[9], tmp[9];
// get K_elem(a,b) 3x3 block
int adr_cell = (3*a)*(3*npe) + 3*b;
for (int r = 0; r < 3; r++) {
for (int c = 0; c < 3; c++) {
blk[3*r+c] = k_elem[adr_cell + r*(3*npe) + c];
// get K_elem(a,b) 3x3 block
int adr_cell = (3*a)*(3*npe) + 3*b;
for (int r = 0; r < 3; r++) {
for (int c = 0; c < 3; c++) {
blk[3*r+c] = k_elem[adr_cell + r*(3*npe) + c];
}
}
}
// tmp = K * R
mju_mulMatMat3(tmp, blk, R);
// blk = RT * tmp = RT * K * R
mju_mulMatMat3(blk, RT, tmp);
// tmp = K * R
mju_mulMatMat3(tmp, blk, R);
// blk = RT * tmp = RT * K * R
mju_mulMatMat3(blk, RT, tmp);
// store in K_rot_cell at (a, b)
int adr_out = (3*a)*dim_e + 3*b;
for (int r = 0; r < 3; r++) {
for (int c = 0; c < 3; c++) {
K_rot_cell[adr_out + r*dim_e + c] = scale * blk[3*r+c];
// store in K_rot_cell at (a, b)
int adr_out = (3*a)*dim_e + 3*b;
for (int r = 0; r < 3; r++) {
for (int c = 0; c < 3; c++) {
K_rot_cell[adr_out + r*dim_e + c] = scale * blk[3*r+c];
}
}
}
}
// optionally store unscaled K_rot to output cache
if (K_rot_out) {
for (int i = 0; i < dim_e*dim_e; i++) {
K_rot_out[krot_adr + i] = K_rot_cell[i] / scale;
}
}
}
// skip Jacobian construction and op when only caching (res == NULL)
if (!res) {
continue;
}
// construct sparse Jacobian for this element's nodes
@@ -1081,40 +1091,9 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
}
}
// apply operation with element's K_rot and J
if (op == mjFLEXOP_VEC) {
addJTBJ_mulSparse(m, d, res, vec, J_rownnz, J_rowadr, J_colind,
J_val, K_rot_cell, dim_e);
} else if (op == mjFLEXOP_ADDH) {
// H -= J_elem^T * K_rot * J_elem (banded format)
mju_zero(J_reduced, dim_e*ndof);
for (int i = 0; i < dim_e; i++) {
int nnz = J_rownnz[i];
int adr = J_rowadr[i];
for (int idx = 0; idx < nnz; idx++) {
int global_col = J_colind[adr + idx];
int local_idx = global2local[global_col];
if (local_idx >= 0) {
J_reduced[i*ndof + local_idx] = J_val[adr + idx];
}
}
}
// KJ = K_rot * J_reduced (dim_e x ndof)
mju_mulMatMat(KJ, K_rot_cell, J_reduced, dim_e, dim_e, ndof);
// H[i,j] -= J_reduced[k,i] * KJ[k,j], store lower triangle in banded format
for (int i = 0; i < ndof; i++) {
for (int j = mjMAX(0, i-nband+1); j <= i; j++) {
mjtNum val = 0;
for (int dim_idx = 0; dim_idx < dim_e; dim_idx++) {
val += J_reduced[dim_idx*ndof + i] * KJ[dim_idx*ndof + j];
}
res[i*nband + nband-1-(i-j)] -= val;
}
}
}
// res += J'*K_rot*J*vec
addJTBJ_mulSparse(m, d, res, vec, J_rownnz, J_rowadr, J_colind,
J_val, K_rot_cell, dim_e);
}
}
@@ -1124,9 +1103,17 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
// compute res += (s1 + s2*damping) * J'*K*J * vec, for all interpolated flexes
// K_rot_cache: if non-NULL, use pre-cached K_rot (same layout as m->flex_stiffness)
void mjd_flexInterp_mul(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec,
mjtNum s1, mjtNum s2) {
mjd_flexInterp_kernel(m, d, mjFLEXOP_VEC, res, vec, s1, s2, NULL, 0, 0);
mjtNum s1, mjtNum s2, const mjtNum* K_rot_cache) {
mjd_flexInterp_kernel(m, d, res, vec, s1, s2, K_rot_cache, NULL);
}
// precompute unscaled K_rot for all elements into cache (same layout as m->flex_stiffness)
void mjd_flexInterp_cacheKrot(const mjModel* m, mjData* d, mjtNum* K_rot_out) {
// use s1=1, s2=0 so scale=1 and K_rot_out gets unscaled values
mjd_flexInterp_kernel(m, d, NULL, NULL, 1, 0, NULL, K_rot_out);
}
src/engine/engine_derivative.h
+5 -1
View File
@@ -44,8 +44,12 @@ MJAPI void mjd_passive_vel(const mjModel* m, mjData* d);
MJAPI void mjd_rne_vel_dense(const mjModel* m, mjData* d);
// compute res += (s1 + s2*damping) * J'*K*J * vec, for all interpolated flexes
// K_rot_cache: if non-NULL, use pre-cached K_rot (same layout as m->flex_stiffness)
MJAPI void mjd_flexInterp_mul(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec,
mjtNum s1, mjtNum s2);
mjtNum s1, mjtNum s2, const mjtNum* K_rot_cache);
// precompute unscaled K_rot for all elements into cache (same layout as m->flex_stiffness)
MJAPI void mjd_flexInterp_cacheKrot(const mjModel* m, mjData* d, mjtNum* K_rot_out);
// compute res += scale * K_bend * vec for standard (non-interp) flex bending
// scale = s1 + s2 * flex_damping[f] per flex
src/engine/engine_forward.c
+9 -3
View File
@@ -1414,16 +1414,22 @@ static void flexInterp_cgsolve(const mjModel* m, mjData* d,
mjtNum* Ap = mjSTACKALLOC(d, nv, mjtNum);
mjtNum* temp = mjSTACKALLOC(d, nv, mjtNum);
// precompute K_rot cache: same layout as m->flex_stiffness
int krot_size = m->nflexstiffness;
mjtNum* K_rot_cache = mjSTACKALLOC(d, krot_size, mjtNum);
mju_zero(K_rot_cache, krot_size);
mjd_flexInterp_cacheKrot(m, d, K_rot_cache);
// build RHS: rhs = qfrc
mju_copy(rhs, qfrc, nv);
// flex_interp velocity correction: rhs += h*K_interp*qvel (K_interp is NSD)
mjd_flexInterp_mul(m, d, rhs, d->qvel, h, 0); // rhs += h*K_interp*v
mjd_flexInterp_mul(m, d, rhs, d->qvel, h, 0, K_rot_cache);
// standard flex bending velocity correction: rhs -= h*K_bend*qvel
mjd_flexBend_mul(m, d, rhs, d->qvel, -h, 0); // rhs -= h*K_bend*v
// --- helper lambda-style inline: compute Ap = A*x ---
// --- helper: compute Ap = A*x ---
// A*x = (M - h*qDeriv)*x - (h^2+h*d)*K_interp*x + (h^2+h*d)*K_bend*x
#define FLEX_CG_MATVEC(Ap_out, x_in) \
mju_mulMatVecSparse(Ap_out, d->qDeriv, x_in, nv, m->D_rownnz, m->D_rowadr, \
@@ -1431,7 +1437,7 @@ static void flexInterp_cgsolve(const mjModel* m, mjData* d,
mju_mulSymVecSparse(temp, d->M, x_in, nv, m->M_rownnz, m->M_rowadr, \
m->M_colind); \
mju_addScl(Ap_out, temp, Ap_out, -h, nv); \
mjd_flexInterp_mul(m, d, Ap_out, x_in, -(h*h), -h); \
mjd_flexInterp_mul(m, d, Ap_out, x_in, -(h*h), -h, K_rot_cache); \
mjd_flexBend_mul(m, d, Ap_out, x_in, h*h, h)
// --- helper: preconditioner solve z = (M - h*qDeriv)^{-1} * r ---
test/engine/engine_derivative_test.cc
+1 -1
View File
@@ -1485,7 +1485,7 @@ static void mulKD_dense(mjModel* m, mjData* d, mjtNum* H_dense,
mju_zero(e_i.data(), nv);
mju_zero(col.data(), nv);
e_i[i] = 1.0;
mjd_flexInterp_mul(m, d, col.data(), e_i.data(), h * h, h);
mjd_flexInterp_mul(m, d, col.data(), e_i.data(), h * h, h, NULL);
// col = +(h^2 + h*damp)*K*e_i, negate to match addH convention (H -= K)
for (int j = 0; j < nv; j++) {
H_dense[j * nv + i] = -col[j];
test/engine/engine_forward_test.cc
+1 -1
View File
@@ -3147,7 +3147,7 @@ TEST_F(ForwardTest, FlexTrilinearInstability) {
// using mulKD for legacy check consistency, but we know it applies h^2+h*d
// scaling; actually, let's stick to the high-level property checks from
// FlexStiffnessSign which used mulKD
mjd_flexInterp_mul(model, data, flex_Kv.data(), v.data(), h * h, h);
mjd_flexInterp_mul(model, data, flex_Kv.data(), v.data(), h * h, h, NULL);
// compute v^T*M*v and v^T*scale*K*v
mjtNum vMv = mju_dot(v.data(), Mv.data(), nv);