From 769109630f45c3296281914525ab7aa0558f6630 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Fri, 23 Jan 2026 21:59:17 -0800 Subject: [PATCH] Move flexvert_J and flexedge_J sparsity pattern to mjModel. PiperOrigin-RevId: 860392796 Change-Id: Ie1264f2df5bf04b5794a3c7d8ad1a92524a6a567 --- src/engine/engine_core_smooth.c | 77 +--------------- src/engine/engine_setconst.c | 156 ++++++++++++++++++++++++++++++++ 2 files changed, 158 insertions(+), 75 deletions(-) diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index b8d70f77..55963651 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -536,7 +536,7 @@ void mj_updateDynamicBVH(const mjModel* m, mjData* d, int bvhadr, int bvhnum) { // compute flex-related quantities void mj_flex(const mjModel* m, mjData* d) { int nv = m->nv; - int* rowadr = m->flexedge_J_rowadr, *rownnz = m->flexedge_J_rownnz; + int* rowadr = m->flexedge_J_rowadr; int* vrowadr = m->flexvert_J_rowadr, *vrownnz = m->flexvert_J_rownnz; // skip if no flexes @@ -658,10 +658,6 @@ void mj_flex(const mjModel* m, mjData* d) { int* chain = mjSTACKALLOC(d, nv, int); // clear Jacobian - mju_zeroInt(rowadr, m->nflexedge); - mju_zeroInt(rownnz, m->nflexedge); - mju_zeroInt(vrowadr, 2*m->nflexvert); - mju_zeroInt(vrownnz, 2*m->nflexvert); mju_zero(d->flexvert_J, 2*m->nJfv); // compute lengths and Jacobians of edges @@ -698,11 +694,6 @@ void mj_flex(const mjModel* m, mjData* d) { continue; } - // set rowadr - if (ebase+e > 0) { - rowadr[ebase+e] = rowadr[ebase+e-1] + rownnz[ebase+e-1]; - } - // get endpoint Jacobians, subtract int NV = mj_jacDifPair(m, d, chain, b1, b2, pos1, pos2, jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1); @@ -714,10 +705,6 @@ void mj_flex(const mjModel* m, mjData* d) { // apply chain rule to compute edge Jacobian mju_mulMatTVec(d->flexedge_J + rowadr[ebase+e], jacdif, vec, 3, NV); - - // copy sparsity info - rownnz[ebase+e] = NV; - mju_copyInt(m->flexedge_J_colind + rowadr[ebase+e], chain, NV); } // if dim=2 and constraints are active we use the vertex-based constraint defined in @@ -729,7 +716,6 @@ void mj_flex(const mjModel* m, mjData* d) { int t_adr, t0, t1, t2; mj_markStack(d); - int* buf_ind = mjSTACKALLOC(d, nv, int); // compute normal from first element t_adr = m->flex_elemdataadr[f]; @@ -858,68 +844,9 @@ void mj_flex(const mjModel* m, mjData* d) { cauchy[0][1] * cauchy[1][0] - 1; } - // 1st pass: compute vrownnz + // clear Jacobian and assemble vertex by vertex int* chain1 = mjSTACKALLOC(d, nv, int); int* chain2 = mjSTACKALLOC(d, nv, int); - - // determine start address for this flex - int v0_base = 2*vbase; - int current_adr = 0; - if (v0_base > 0) { - current_adr = vrowadr[v0_base - 1] + vrownnz[v0_base - 1]; - } - vrowadr[v0_base] = current_adr; - - for (int v=0; vflex_edge[2*(ebase+e)]; - int v2 = m->flex_edge[2*(ebase+e)+1]; - - // chains from edge e - int b1 = m->flex_vertbodyid[vbase+v1]; - int b2 = m->flex_vertbodyid[vbase+v2]; - int NV1 = mj_bodyChain(m, b1, chain1); - int NV2 = mj_bodyChain(m, b2, chain2); - - for (int j=0; jnflexvert) { - vrowadr[row1+1] = vrowadr[row1] + current_nnz; - } - - // fill colind - int count = 0; - for (int j=0; jflexvert_J_colind[vrowadr[row0]+count] = j; - m->flexvert_J_colind[vrowadr[row1]+count] = j; - count++; - } - } - } - - // 2nd pass: clear Jacobian and assemble vertex by vertex mjtNum* J0_dense = mjSTACKALLOC(d, nv, mjtNum); mjtNum* J1_dense = mjSTACKALLOC(d, nv, mjtNum); mjtNum dI1dy1[3], dI1dy2[3], FB[6]; diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 56200cc6..339190e7 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -282,9 +282,165 @@ static void setFixed(mjModel* m, mjData* d) { mj_freeStack(d); } +// compute flex sparsity: flexedge_J_{rowadr,rownnz,colind} and flexvert_J_{rowadr,rownnz} +static void makeFlexSparse(mjModel* m, mjData* d) { + int nv = m->nv; + int* rowadr = m->flexedge_J_rowadr; + int* rownnz = m->flexedge_J_rownnz; + int* colind = m->flexedge_J_colind; + int* vrowadr = m->flexvert_J_rowadr; + int* vrownnz = m->flexvert_J_rownnz; + + if (!m->nflex) { + return; + } + + mj_markStack(d); + int* chain = mjSTACKALLOC(d, nv, int); + int* chain1 = mjSTACKALLOC(d, nv, int); + int* chain2 = mjSTACKALLOC(d, nv, int); + int* buf_ind = mjSTACKALLOC(d, nv, int); + mjtNum* dummy_pos = mjSTACKALLOC(d, 3, mjtNum); + mju_zero(dummy_pos, 3); + + // clear + mju_zeroInt(rowadr, m->nflexedge); + mju_zeroInt(rownnz, m->nflexedge); + mju_zeroInt(vrowadr, 2 * m->nflexvert); + mju_zeroInt(vrownnz, 2 * m->nflexvert); + + // compute lengths and Jacobians of edges + for (int f = 0; f < m->nflex; f++) { + // skip if edges cannot generate forces + if (m->flex_rigid[f] || m->flex_interp[f]) { + continue; + } + + // skip Jacobian if no built-in passive force is needed + int skipjacobian = !m->flex_edgeequality[f] && !m->flex_edgedamping[f] && + !m->flex_edgestiffness[f] && !m->flex_damping[f]; + + // 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++) { + if (skipjacobian) { + continue; + } + + // set rowadr + if (ebase + e > 0) { + rowadr[ebase + e] = rowadr[ebase + e - 1] + rownnz[ebase + e - 1]; + } + + int v1 = m->flex_edge[2 * (ebase + e)]; + int v2 = m->flex_edge[2 * (ebase + e) + 1]; + int b1 = m->flex_vertbodyid[vbase + v1]; + int b2 = m->flex_vertbodyid[vbase + v2]; + + // get sparsity + int NV = mj_jacDifPair(m, d, chain, b1, b2, dummy_pos, dummy_pos, NULL, + NULL, NULL, NULL, NULL, NULL, /*issparse=*/1); + + // copy sparsity info + rownnz[ebase + e] = NV; + mju_copyInt(colind + rowadr[ebase + e], chain, NV); + } + + // if dim=2 and constraints are active we use the vertex-based constraint + if (m->flex_dim[f] == 2 && m->flex_edgeequality[f] == 2) { + int nvert = m->flex_vertnum[f]; + + // build vertex adjacency list local to this function + int* v_edge_cnt = mjSTACKALLOC(d, nvert, int); + 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) { + v_edge_cnt[m->flex_edge[2 * (ebase + e) + 0]]++; + v_edge_cnt[m->flex_edge[2 * (ebase + e) + 1]]++; + } + int total_adj_edges = 0; + for (int v = 0; v < nvert; ++v) { + v_edge_adr[v] = total_adj_edges; + total_adj_edges += v_edge_cnt[v]; + } + int* v_edge_fill = mjSTACKALLOC(d, nvert, int); + mju_zeroInt(v_edge_fill, nvert); + for (int e = 0; e < m->flex_edgenum[f]; ++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; + v_edge_fill[v1]++; + adj_edges[v_edge_adr[v2] + v_edge_fill[v2]] = e; + v_edge_fill[v2]++; + } + + // determine start address for this flex + int v0_base = 2 * vbase; + int current_adr = 0; + if (v0_base > 0) { + current_adr = vrowadr[v0_base - 1] + vrownnz[v0_base - 1]; + } + vrowadr[v0_base] = current_adr; + + for (int v = 0; v < nvert; ++v) { + // clear buf_ind + mju_zeroInt(buf_ind, nv); + int current_nnz = 0; + for (int i = 0; i < v_edge_cnt[v]; ++i) { + int e = adj_edges[v_edge_adr[v] + i]; + int v1 = m->flex_edge[2 * (ebase + e)]; + int v2 = m->flex_edge[2 * (ebase + e) + 1]; + + // chains from edge e + int b1 = m->flex_vertbodyid[vbase + v1]; + int b2 = m->flex_vertbodyid[vbase + v2]; + int NV1 = mj_bodyChain(m, b1, chain1); + int NV2 = mj_bodyChain(m, b2, chain2); + + for (int j = 0; j < NV1; ++j) { + if (!buf_ind[chain1[j]]) { + buf_ind[chain1[j]] = 1; + current_nnz++; + } + } + for (int j = 0; j < NV2; ++j) { + if (!buf_ind[chain2[j]]) { + buf_ind[chain2[j]] = 1; + current_nnz++; + } + } + } + int row0 = 2 * (vbase + v); + int row1 = 2 * (vbase + v) + 1; + vrownnz[row0] = vrownnz[row1] = current_nnz; + + // set rowadr for next rows + vrowadr[row1] = vrowadr[row0] + current_nnz; + if (row1 + 1 < 2 * m->nflexvert) { + vrowadr[row1 + 1] = vrowadr[row1] + current_nnz; + } + + // fill colind + int count = 0; + for (int j = 0; j < nv; j++) { + if (buf_ind[j]) { + m->flexvert_J_colind[vrowadr[row0] + count] = j; + m->flexvert_J_colind[vrowadr[row1] + count] = j; + count++; + } + } + } + } + } + + mj_freeStack(d); +} // set quantities that depend on qpos0 static void set0(mjModel* m, mjData* d) { + makeFlexSparse(m, d); int nv = m->nv; mjtNum A[36] = {0}, pos[3], quat[4]; mj_markStack(d);