Use banded solver for implicit flex integration.
The flex interpolation stiffness matrix within the implicit/implicitfast solvers is now built and factorized in a banded format instead of a dense one. This involves: - Calculating the bandwidth based on the sparsity of the mass/damping matrix and the connectivity within flex cells. - Allocating and populating a banded matrix `H`. - Using `mju_cholFactorBand` and `mju_cholSolveBand` for factorization and solving. This change improves performance for flexes with many DOFs but local coupling. PiperOrigin-RevId: 901297952 Change-Id: I3efe06353d1903ea65ab30dc49685cede228bb68
This commit is contained in:
Alessio Quaglino
committed by
Copybara-Service
Copybara-Service
parent
6c7ed66781
commit
b16383dfaf
@@ -872,12 +872,12 @@ typedef enum {
|
||||
|
||||
// shared kernel for flex interpolation derivatives, scale = s1 + s2*damping
|
||||
// op: operation type (VEC, or ADDH)
|
||||
// res: output vector (VEC) or dense H matrix (ADDH)
|
||||
// res: output vector (VEC) or banded H matrix (ADDH)
|
||||
// vec: input vector for VEC operation, NULL otherwise
|
||||
// dof_indices, ndof: DOF mapping for ADDH, ignored 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,
|
||||
mjtNum* res, const mjtNum* vec, mjtNum s1, mjtNum s2,
|
||||
const int* dof_indices, int ndof) {
|
||||
const int* dof_indices, int ndof, int nband) {
|
||||
int nv = m->nv;
|
||||
|
||||
// build global2local map for ADDH
|
||||
@@ -1021,7 +1021,7 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
|
||||
J_val, K_rot_cell, dim_c);
|
||||
} else if (op == mjFLEXOP_ADDH) {
|
||||
mj_markStack(d);
|
||||
// H -= J_cell^T * K_rot_cell * J_cell
|
||||
// H -= J_cell^T * K_rot_cell * J_cell (banded format)
|
||||
mjtNum* J_reduced = mjSTACKALLOC(d, dim_c*ndof, mjtNum);
|
||||
mju_zero(J_reduced, dim_c*ndof);
|
||||
|
||||
@@ -1041,14 +1041,14 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
|
||||
mjtNum* KJ = mjSTACKALLOC(d, dim_c*ndof, mjtNum);
|
||||
mju_mulMatMat(KJ, K_rot_cell, J_reduced, dim_c, dim_c, ndof);
|
||||
|
||||
// H[i,j] -= J_reduced[k,i] * KJ[k,j]
|
||||
// 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 = 0; j < ndof; j++) {
|
||||
for (int j = mjMAX(0, i-nband+1); j <= i; j++) {
|
||||
mjtNum val = 0;
|
||||
for (int dim_idx = 0; dim_idx < dim_c; dim_idx++) {
|
||||
val += J_reduced[dim_idx*ndof + i] * KJ[dim_idx*ndof + j];
|
||||
}
|
||||
res[i*ndof + j] -= val;
|
||||
res[i*nband + nband-1-(i-j)] -= val;
|
||||
}
|
||||
}
|
||||
mj_freeStack(d);
|
||||
@@ -1072,15 +1072,16 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
|
||||
// compute res += (h^2 + h*damping) * J'*K*J * vec, for all interpolated flexes
|
||||
void mjd_flexInterp_mulKD(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec, mjtNum h) {
|
||||
// s1=h*h, s2=h => scale = h*h + h*damping
|
||||
mjd_flexInterp_kernel(m, d, mjFLEXOP_VEC, res, vec, h * h, h, NULL, 0);
|
||||
mjd_flexInterp_kernel(m, d, mjFLEXOP_VEC, res, vec, h * h, h, NULL, 0, 0);
|
||||
}
|
||||
|
||||
|
||||
// add (h^2 + h*damping) * J'*K*J to dense matrix H, for all interpolated flexes
|
||||
// H: dense ndof x ndof matrix
|
||||
// add (h^2 + h*damping) * J'*K*J to banded matrix H, for all interpolated flexes
|
||||
// H: banded ndof x nband matrix (lower triangle, band storage)
|
||||
// dof_indices: maps local indices to global DOFs
|
||||
void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices, int ndof, mjtNum h) {
|
||||
mjd_flexInterp_kernel(m, d, mjFLEXOP_ADDH, H, NULL, h * h, h, dof_indices, ndof);
|
||||
void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices,
|
||||
int ndof, int nband, mjtNum h) {
|
||||
mjd_flexInterp_kernel(m, d, mjFLEXOP_ADDH, H, NULL, h * h, h, dof_indices, ndof, nband);
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -49,7 +49,7 @@ MJAPI void mjd_flexInterp_mulKD(const mjModel* m, mjData* d, mjtNum* res, const
|
||||
|
||||
// assemble flex stiffness matrix H_flex: H += h*h*K + h*D
|
||||
// H is a dense matrix of size ndof x ndof, dof_indices maps local rows/cols to global DOFs
|
||||
MJAPI void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices, int ndof, mjtNum h);
|
||||
MJAPI void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices, int ndof, int nband, mjtNum h);
|
||||
|
||||
|
||||
#ifdef __cplusplus
|
||||
|
||||
+40
-15
@@ -1350,11 +1350,12 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
|
||||
}
|
||||
|
||||
|
||||
// context for flex interp reduced dense factorization/solve
|
||||
// context for flex interp reduced banded factorization/solve
|
||||
typedef struct {
|
||||
mjtNum* H; // dense Cholesky-factored matrix (ndof x ndof)
|
||||
mjtNum* H; // banded Cholesky-factored matrix (ndof x nband)
|
||||
int* dof_indices; // global DOF index for each local flex DOF
|
||||
int ndof; // number of flex DOFs
|
||||
int nband; // half-bandwidth + 1 (number of band columns)
|
||||
int ncoupling; // number of off-diagonal coupling terms
|
||||
mjtNum* coupling_val; // coupling coefficient values
|
||||
int* coupling_row; // local flex row index for each coupling term
|
||||
@@ -1391,7 +1392,7 @@ static void flexInterp_collect(const mjModel* m, int f,
|
||||
}
|
||||
|
||||
|
||||
// build and factor the reduced dense matrix for flex interp DOFs
|
||||
// build and factor the reduced banded matrix for flex interp DOFs
|
||||
// mark/free stack handled by caller
|
||||
static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv) {
|
||||
FlexInterpContext ctx = {0};
|
||||
@@ -1456,19 +1457,36 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv)
|
||||
const int* colind = implicit ? m->D_colind : m->M_colind;
|
||||
const mjtNum* source = implicit ? d->qLU : d->qH;
|
||||
|
||||
// count coupling terms (off-diagonal: flex row, non-flex col)
|
||||
// get precomputed bandwidth
|
||||
int bandwidth = 0;
|
||||
for (int f=0; f < m->nflex; f++) {
|
||||
if (m->flex_interp[f]) {
|
||||
if (m->flex_bandwidth[f] > bandwidth) {
|
||||
bandwidth = m->flex_bandwidth[f];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// compute ncoupling from sparse matrix entries
|
||||
int ncoupling = 0;
|
||||
for (int i=0; i < ndof; i++) {
|
||||
int row = dof_indices[i];
|
||||
int start = rowadr[row];
|
||||
int end = start + rownnz[row];
|
||||
for (int k=start; k < end; k++) {
|
||||
if (global2local[colind[k]] < 0) {
|
||||
int local_j = global2local[colind[k]];
|
||||
if (local_j < 0) {
|
||||
ncoupling++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// nband = bandwidth + 1 (includes diagonal)
|
||||
int nband = bandwidth + 1;
|
||||
|
||||
// cap nband at ndof (dense fallback for small systems)
|
||||
if (nband > ndof) nband = ndof;
|
||||
|
||||
// allocate coupling storage
|
||||
mjtNum* coupling_val = NULL;
|
||||
int* coupling_row = NULL;
|
||||
@@ -1479,9 +1497,9 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv)
|
||||
coupling_col = mjSTACKALLOC(d, ncoupling, int);
|
||||
}
|
||||
|
||||
// build H_flex (dense) from qLU (implicit) or qH (implicitfast)
|
||||
mjtNum* H = mjSTACKALLOC(d, ndof*ndof, mjtNum);
|
||||
mju_zero(H, ndof*ndof);
|
||||
// build H_flex (banded) from qLU (implicit) or qH (implicitfast)
|
||||
mjtNum* H = mjSTACKALLOC(d, ndof*nband, mjtNum);
|
||||
mju_zero(H, ndof*nband);
|
||||
|
||||
int coup_cnt = 0;
|
||||
for (int i=0; i < ndof; i++) {
|
||||
@@ -1492,7 +1510,13 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv)
|
||||
int col = colind[k];
|
||||
int local_j = global2local[col];
|
||||
if (local_j >= 0) {
|
||||
H[i*ndof+local_j] = source[k];
|
||||
// store lower triangle only: row i, col local_j, where i >= local_j
|
||||
if (i >= local_j) {
|
||||
H[i*nband + nband-1-(i-local_j)] = source[k];
|
||||
} else {
|
||||
// upper triangle entry: store symmetrically in lower triangle
|
||||
H[local_j*nband + nband-1-(local_j-i)] = source[k];
|
||||
}
|
||||
} else if (coup_cnt < ncoupling) {
|
||||
coupling_val[coup_cnt] = source[k];
|
||||
coupling_row[coup_cnt] = i;
|
||||
@@ -1502,14 +1526,15 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv)
|
||||
}
|
||||
}
|
||||
|
||||
// add flex stiffness and factorize
|
||||
mjd_flexInterp_addH(m, d, H, dof_indices, ndof, m->opt.timestep);
|
||||
mju_cholFactor(H, ndof, mjMINVAL);
|
||||
// add flex stiffness in banded format and factorize
|
||||
mjd_flexInterp_addH(m, d, H, dof_indices, ndof, nband, m->opt.timestep);
|
||||
mju_cholFactorBand(H, ndof, nband, 0, 0, 0);
|
||||
|
||||
// store results in context
|
||||
ctx.H = H;
|
||||
ctx.dof_indices = dof_indices;
|
||||
ctx.ndof = ndof;
|
||||
ctx.nband = nband;
|
||||
ctx.ncoupling = ncoupling;
|
||||
ctx.coupling_val = coupling_val;
|
||||
ctx.coupling_row = coupling_row;
|
||||
@@ -1518,7 +1543,7 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv)
|
||||
}
|
||||
|
||||
|
||||
// solve the reduced dense system for flex interp DOFs, overwrite qacc
|
||||
// solve the reduced banded system for flex interp DOFs, overwrite qacc
|
||||
static void flexInterp_solve(const mjModel* m, mjData* d, const FlexInterpContext* ctx,
|
||||
mjtNum* qacc, const mjtNum* qfrc, int nv) {
|
||||
int ndof = ctx->ndof;
|
||||
@@ -1544,8 +1569,8 @@ static void flexInterp_solve(const mjModel* m, mjData* d, const FlexInterpContex
|
||||
qfrc_flex[ctx->coupling_row[k]] -= ctx->coupling_val[k] * qacc[ctx->coupling_col[k]];
|
||||
}
|
||||
|
||||
// solve and scatter back
|
||||
mju_cholSolve(qfrc_flex, ctx->H, qfrc_flex, ndof);
|
||||
// solve with banded Cholesky and scatter back
|
||||
mju_cholSolveBand(qfrc_flex, ctx->H, qfrc_flex, ndof, ctx->nband, 0);
|
||||
mju_scatter(qacc, qfrc_flex, ctx->dof_indices, ndof);
|
||||
}
|
||||
|
||||
|
||||
@@ -636,6 +636,135 @@ static void makeFlexSparse(mjModel* m, mjData* d) {
|
||||
mj_freeStack(d);
|
||||
}
|
||||
|
||||
// compute flex bandwidth for trilinear interpolation
|
||||
static void makeFlexBandwidth(mjModel* m, mjData* d) {
|
||||
if (!m->nflex) {
|
||||
return;
|
||||
}
|
||||
|
||||
mj_markStack(d);
|
||||
int* chain_dofs = mjSTACKALLOC(d, m->nv, int);
|
||||
int* seen_dof = mjSTACKALLOC(d, m->nv, int);
|
||||
int* dof_indices = mjSTACKALLOC(d, m->nv, int);
|
||||
int* global2local = mjSTACKALLOC(d, m->nv, int);
|
||||
|
||||
mju_zeroInt(seen_dof, m->nv);
|
||||
for (int i = 0; i < m->nv; i++) {
|
||||
global2local[i] = -1;
|
||||
}
|
||||
|
||||
int ndof = 0;
|
||||
for (int f = 0; f < m->nflex; f++) {
|
||||
if (m->flex_interp[f]) {
|
||||
int nodenum = m->flex_nodenum[f];
|
||||
int nodeadr = m->flex_nodeadr[f];
|
||||
for (int n = 0; n < nodenum; n++) {
|
||||
int b = m->flex_nodebodyid[nodeadr + n];
|
||||
// only the body's own DOFs enter the reduced banded flex system;
|
||||
// ancestor DOFs are solved by the global factorization and coupled
|
||||
// via off-diagonal correction (see flexInterp_solve in engine_forward)
|
||||
int chain_nnz;
|
||||
if (m->body_dofnum[b] == 0) {
|
||||
chain_nnz = mj_bodyChain(m, b, chain_dofs);
|
||||
} else {
|
||||
chain_nnz = m->body_dofnum[b];
|
||||
for (int j = 0; j < chain_nnz; j++) {
|
||||
chain_dofs[j] = m->body_dofadr[b] + j;
|
||||
}
|
||||
}
|
||||
for (int i = 0; i < chain_nnz; i++) {
|
||||
int dof = chain_dofs[i];
|
||||
if (!seen_dof[dof]) {
|
||||
seen_dof[dof] = 1;
|
||||
dof_indices[ndof] = dof;
|
||||
global2local[dof] = ndof++;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int bandwidth = 0;
|
||||
if (ndof > 0) {
|
||||
// check sparse matrix coupling (both D and M)
|
||||
for (int integrator = 0; integrator < 2; integrator++) {
|
||||
const int* rownnz = (integrator == 0) ? m->D_rownnz : m->M_rownnz;
|
||||
const int* rowadr = (integrator == 0) ? m->D_rowadr : m->M_rowadr;
|
||||
const int* colind = (integrator == 0) ? m->D_colind : m->M_colind;
|
||||
|
||||
// D arrays are only allocated for implicit integrators
|
||||
if (!rownnz) continue;
|
||||
|
||||
for (int i = 0; i < ndof; i++) {
|
||||
int row = dof_indices[i];
|
||||
int start = rowadr[row];
|
||||
int end = start + rownnz[row];
|
||||
for (int k = start; k < end; k++) {
|
||||
int local_j = global2local[colind[k]];
|
||||
if (local_j >= 0) {
|
||||
int diff = i - local_j;
|
||||
if (diff < 0) diff = -diff;
|
||||
if (diff > bandwidth) bandwidth = diff;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// check stiffness coupling
|
||||
for (int f = 0; f < m->nflex; f++) {
|
||||
if (!m->flex_interp[f]) continue;
|
||||
int order = m->flex_interp[f];
|
||||
int nodeadr = m->flex_nodeadr[f];
|
||||
int nodenum = m->flex_nodenum[f];
|
||||
int cx = m->flex_cellnum[3*f+0];
|
||||
int cy = m->flex_cellnum[3*f+1];
|
||||
int cz = m->flex_cellnum[3*f+2];
|
||||
int ny = cy * order + 1;
|
||||
int nz = cz * order + 1;
|
||||
|
||||
for (int icx = 0; icx < cx; icx++) {
|
||||
for (int icy = 0; icy < cy; icy++) {
|
||||
for (int icz = 0; icz < cz; icz++) {
|
||||
int min_local = ndof, max_local = -1;
|
||||
for (int lx = 0; lx <= order; lx++) {
|
||||
for (int ly = 0; ly <= order; ly++) {
|
||||
for (int lz = 0; lz <= order; lz++) {
|
||||
int gx = icx * order + lx;
|
||||
int gy = icy * order + ly;
|
||||
int gz = icz * order + lz;
|
||||
int node_idx = gx * ny * nz + gy * nz + gz; // non-negative by construction
|
||||
if (node_idx < nodenum) {
|
||||
int b = m->flex_nodebodyid[nodeadr + node_idx];
|
||||
int chain_nnz = mj_bodyChain(m, b, chain_dofs);
|
||||
for (int i = 0; i < chain_nnz; i++) {
|
||||
int dof = chain_dofs[i];
|
||||
int local = global2local[dof];
|
||||
if (local >= 0) {
|
||||
if (local < min_local) min_local = local;
|
||||
if (local > max_local) max_local = local;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (max_local >= 0 && max_local - min_local > bandwidth) {
|
||||
bandwidth = max_local - min_local;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// store bandwidth for all flexes (global max)
|
||||
for (int f = 0; f < m->nflex; f++) {
|
||||
m->flex_bandwidth[f] = bandwidth;
|
||||
}
|
||||
|
||||
mj_freeStack(d);
|
||||
}
|
||||
|
||||
// align 2D flexes to the XY plane
|
||||
static void mj_alignFlex(mjModel* m, mjData* d) {
|
||||
for (int f = 0; f < m->nflex; f++) {
|
||||
@@ -687,6 +816,7 @@ static void mj_alignFlex(mjModel* m, mjData* d) {
|
||||
static void set0(mjModel* m, mjData* d) {
|
||||
makeTendonSparse(m);
|
||||
makeFlexSparse(m, d);
|
||||
makeFlexBandwidth(m, d);
|
||||
mj_alignFlex(m, d);
|
||||
int nv = m->nv;
|
||||
mjtNum A[36] = {0}, pos[3], quat[4];
|
||||
|
||||
Reference in New Issue
Block a user