From e933081ff17a32b98519098cd59685c152ae3a80 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Sun, 25 Jan 2026 05:22:06 -0800 Subject: [PATCH] Optimize flexvert Jacobian computation. The dense Jacobian arrays `J0_dense` and `J1_dense` are now zeroed only once before the vertex loop. Inside the loop, only the entries that were actually used are zeroed out after being added to the sparse Jacobian, allowing for efficient reuse of the dense arrays across all vertices. PiperOrigin-RevId: 860799056 Change-Id: Ia61daae5eeb2f4dbc0c5d8c0064622a881e92b7a --- src/engine/engine_core_smooth.c | 74 ++++++++++++++++++++------------- 1 file changed, 45 insertions(+), 29 deletions(-) diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index ab6ab889..caf971fc 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -533,6 +533,17 @@ void mj_updateDynamicBVH(const mjModel* m, mjData* d, int bvhadr, int bvhnum) { } +// C(3x2) = A(3x2) * B(2x2) +static inline void mju_mulMatMat322(mjtNum* C, const mjtNum* A, const mjtNum* B) { + C[0] = A[0]*B[0] + A[1]*B[2]; + C[1] = A[0]*B[1] + A[1]*B[3]; + C[2] = A[2]*B[0] + A[3]*B[2]; + C[3] = A[2]*B[1] + A[3]*B[3]; + C[4] = A[4]*B[0] + A[5]*B[2]; + C[5] = A[4]*B[1] + A[5]*B[3]; +} + + // compute flex-related quantities void mj_flex(const mjModel* m, mjData* d) { int nv = m->nv; @@ -659,6 +670,7 @@ void mj_flex(const mjModel* m, mjData* d) { // clear Jacobian mju_zero(d->flexvert_J, 2*m->nJfv); + mju_zero(d->flexedge_J, m->nJfe); // compute lengths and Jacobians of edges for (int f=0; f < m->nflex; f++) { @@ -667,8 +679,8 @@ void mj_flex(const mjModel* m, mjData* d) { continue; } - // skip Jacobian if no built-in passive force is needed - int skipjacobian = !m->flex_edgeequality[f] && + // skip edge Jacobian if no built-in passive force is needed + int skipjacobian = m->flex_edgeequality[f] != 1 && !m->flex_edgedamping[f] && !m->flex_edgestiffness[f] && !m->flex_damping[f]; @@ -676,7 +688,8 @@ void mj_flex(const mjModel* m, mjData* d) { // process edges of this flex int vbase = m->flex_vertadr[f]; int ebase = m->flex_edgeadr[f]; - for (int e=0; e < m->flex_edgenum[f]; e++) { + int edgenum = m->flex_edgenum[f]; + for (int e=0; e < edgenum; e++) { int v1 = m->flex_edge[2*(ebase+e)]; int v2 = m->flex_edge[2*(ebase+e)+1]; int b1 = m->flex_vertbodyid[vbase+v1]; @@ -715,9 +728,9 @@ void mj_flex(const mjModel* m, mjData* d) { mj_markStack(d); // compute edge vectors - mjtNum* edge_dx = mjSTACKALLOC(d, 3*m->flex_edgenum[f], mjtNum); - mjtNum* edge_dy = mjSTACKALLOC(d, 3*m->flex_edgenum[f], mjtNum); - for (int e=0; e < m->flex_edgenum[f]; e++) { + mjtNum* edge_dx = mjSTACKALLOC(d, 3*edgenum, mjtNum); + mjtNum* edge_dy = mjSTACKALLOC(d, 3*edgenum, mjtNum); + for (int e=0; e < edgenum; e++) { int v1 = m->flex_edge[2*(ebase+e)]; int v2 = m->flex_edge[2*(ebase+e)+1]; mju_sub3(edge_dx + 3 * e, m->flex_vert0 + 3 * (vbase + v2), @@ -739,7 +752,7 @@ void mj_flex(const mjModel* m, mjData* d) { int* v_edge_adr = mjSTACKALLOC(d, nvert, int); int* adj_edges = mjSTACKALLOC(d, 2*m->flex_edgenum[f], int); mju_zeroInt(v_edge_cnt, nvert); - for (int e = 0; e < m->flex_edgenum[f]; ++e) { + for (int e = 0; e < edgenum; ++e) { v_edge_cnt[m->flex_edge[2*(ebase+e)+0]]++; v_edge_cnt[m->flex_edge[2*(ebase+e)+1]]++; } @@ -750,7 +763,7 @@ void mj_flex(const mjModel* m, mjData* d) { } int* v_edge_fill = mjSTACKALLOC(d, nvert, int); mju_zeroInt(v_edge_fill, nvert); - for (int e = 0; e < m->flex_edgenum[f]; ++e) { + for (int e = 0; e < edgenum; ++e) { int v1 = m->flex_edge[2*(ebase+e)+0]; int v2 = m->flex_edge[2*(ebase+e)+1]; adj_edges[v_edge_adr[v1] + v_edge_fill[v1]] = e; @@ -784,16 +797,16 @@ void mj_flex(const mjModel* m, mjData* d) { } // accumulate A += w * dy * dx', B += w * dx * dx' - for (int row=0; row < 3; row++) { - for (int col=0; col < 2; col++) { - A[2 * row + col] += weight * dy[row] * dx[col]; - } - } - for (int row=0; row < 2; row++) { - for (int col=0; col < 2; col++) { - B[2 * row + col] += weight * dx[row] * dx[col]; - } - } + A[0] += weight * dy[0] * dx[0]; + A[1] += weight * dy[0] * dx[1]; + A[2] += weight * dy[1] * dx[0]; + A[3] += weight * dy[1] * dx[1]; + A[4] += weight * dy[2] * dx[0]; + A[5] += weight * dy[2] * dx[1]; + B[0] += weight * dx[0] * dx[0]; + B[1] += weight * dx[0] * dx[1]; + B[2] += weight * dx[1] * dx[0]; + B[3] += weight * dx[1] * dx[1]; } int vadr = vbase+v; @@ -814,7 +827,7 @@ void mj_flex(const mjModel* m, mjData* d) { } // compute deformation gradient F = A * Binv - mju_mulMatMat(F, A, Binv, 3, 2, 2); + mju_mulMatMat322(F, A, Binv); // compute Cauchy strain tensor F^T F cauchy[0][0] = F[0]*F[0] + F[2]*F[2] + F[4]*F[4]; @@ -837,15 +850,15 @@ void mj_flex(const mjModel* m, mjData* d) { mjtNum dI2dy1[3], dI2dy2[3]; mjtNum cauchy[4], adj[4], Fadj[6], FadjBinv[6], dI2dy[3]; - for (int v=0; vflexvert_J[vrowadr[row0]+j] += J0_dense[m->flexvert_J_colind[vrowadr[row0]+j]]; + for (int j = 0; j < nnz0; j++) { + int col = m->flexvert_J_colind[vrowadr[row0] + j]; + d->flexvert_J[vrowadr[row0] + j] += J0_dense[col]; + J0_dense[col] = 0; } int row1 = 2*(vbase+v)+1; int nnz1 = vrownnz[row1]; for (int j = 0; j < nnz1; j++) { - d->flexvert_J[vrowadr[row1] + j] += - J1_dense[m->flexvert_J_colind[vrowadr[row1] + j]]; + int col = m->flexvert_J_colind[vrowadr[row1] + j]; + d->flexvert_J[vrowadr[row1] + j] += J1_dense[col]; + J1_dense[col] = 0; } // mass scaling: scale constraint by sqrt(mass) to improve condition