From 2f28473bc1de1efd4bc67bc13d37a74987d3e0e9 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Thu, 21 Aug 2025 07:20:29 -0700 Subject: [PATCH] Remove transposed Jacobian fields from `mjData` No longer computed unconditionally, only ad-hoc where required. PiperOrigin-RevId: 797753297 Change-Id: I2030fc342c98ff33575b0526dd72e2110c4fcb74 --- doc/includes/references.h | 10 --- include/mujoco/mjdata.h | 10 --- include/mujoco/mjxmacro.h | 10 --- python/mujoco/introspect/structs.py | 80 ------------------- src/engine/engine_core_constraint.c | 14 +--- src/engine/engine_island.c | 14 ---- src/engine/engine_print.c | 2 - .../engine_util_sparse_benchmark_test.cc | 40 +++++++--- test/engine/engine_island_test.cc | 21 ----- unity/Runtime/Bindings/MjBindings.cs | 10 --- 10 files changed, 33 insertions(+), 178 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index b769c729..90d939ff 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -349,12 +349,7 @@ struct mjData_ { int* efc_J_rowadr; // row start address in colind array (nefc x 1) int* efc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) int* efc_J_colind; // column indices in constraint Jacobian (nJ x 1) - int* efc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nv x 1) - int* efc_JT_rowadr; // row start address in colind array T (nv x 1) - int* efc_JT_rowsuper; // number of subsequent rows in supernode T (nv x 1) - int* efc_JT_colind; // column indices in constraint Jacobian T (nJ x 1) mjtNum* efc_J; // constraint Jacobian (nJ x 1) - mjtNum* efc_JT; // constraint Jacobian transposed (nJ x 1) mjtNum* efc_pos; // constraint position (equality, contact) (nefc x 1) mjtNum* efc_margin; // inclusion margin (contact) (nefc x 1) mjtNum* efc_frictionloss; // frictionloss (friction) (nefc x 1) @@ -399,12 +394,7 @@ struct mjData_ { int* iefc_J_rowadr; // row start address in colind array (nefc x 1) int* iefc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) int* iefc_J_colind; // column indices in constraint Jacobian (nJ x 1) - int* iefc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nidof x 1) - int* iefc_JT_rowadr; // row start address in colind array T (nidof x 1) - int* iefc_JT_rowsuper; // number of subsequent rows in supernode T (nidof x 1) - int* iefc_JT_colind; // column indices in constraint Jacobian T (nJ x 1) mjtNum* iefc_J; // constraint Jacobian (nJ x 1) - mjtNum* iefc_JT; // constraint Jacobian transposed (nJ x 1) mjtNum* iefc_frictionloss; // frictionloss (friction) (nefc x 1) mjtNum* iefc_D; // constraint mass (nefc x 1) mjtNum* iefc_R; // inverse constraint mass (nefc x 1) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 6e12cb98..f1ad77c7 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -377,12 +377,7 @@ struct mjData_ { int* efc_J_rowadr; // row start address in colind array (nefc x 1) int* efc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) int* efc_J_colind; // column indices in constraint Jacobian (nJ x 1) - int* efc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nv x 1) - int* efc_JT_rowadr; // row start address in colind array T (nv x 1) - int* efc_JT_rowsuper; // number of subsequent rows in supernode T (nv x 1) - int* efc_JT_colind; // column indices in constraint Jacobian T (nJ x 1) mjtNum* efc_J; // constraint Jacobian (nJ x 1) - mjtNum* efc_JT; // constraint Jacobian transposed (nJ x 1) mjtNum* efc_pos; // constraint position (equality, contact) (nefc x 1) mjtNum* efc_margin; // inclusion margin (contact) (nefc x 1) mjtNum* efc_frictionloss; // frictionloss (friction) (nefc x 1) @@ -427,12 +422,7 @@ struct mjData_ { int* iefc_J_rowadr; // row start address in colind array (nefc x 1) int* iefc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) int* iefc_J_colind; // column indices in constraint Jacobian (nJ x 1) - int* iefc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nidof x 1) - int* iefc_JT_rowadr; // row start address in colind array T (nidof x 1) - int* iefc_JT_rowsuper; // number of subsequent rows in supernode T (nidof x 1) - int* iefc_JT_colind; // column indices in constraint Jacobian T (nJ x 1) mjtNum* iefc_J; // constraint Jacobian (nJ x 1) - mjtNum* iefc_JT; // constraint Jacobian transposed (nJ x 1) mjtNum* iefc_frictionloss; // frictionloss (friction) (nefc x 1) mjtNum* iefc_D; // constraint mass (nefc x 1) mjtNum* iefc_R; // inverse constraint mass (nefc x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index b5a1cdb8..606b4752 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -729,12 +729,7 @@ XNV( int, efc_J_rowadr, MJ_D(nefc), 1 ) \ XNV( int, efc_J_rowsuper, MJ_D(nefc), 1 ) \ XNV( int, efc_J_colind, MJ_D(nJ), 1 ) \ - XNV( int, efc_JT_rownnz, MJ_M(nv), 1 ) \ - XNV( int, efc_JT_rowadr, MJ_M(nv), 1 ) \ - XNV( int, efc_JT_rowsuper, MJ_M(nv), 1 ) \ - XNV( int, efc_JT_colind, MJ_D(nJ), 1 ) \ XNV( mjtNum, efc_J, MJ_D(nJ), 1 ) \ - XNV( mjtNum, efc_JT, MJ_D(nJ), 1 ) \ X ( mjtNum, efc_pos, MJ_D(nefc), 1 ) \ X ( mjtNum, efc_margin, MJ_D(nefc), 1 ) \ X ( mjtNum, efc_frictionloss, MJ_D(nefc), 1 ) \ @@ -786,12 +781,7 @@ XNV( int, iefc_J_rowadr, MJ_D(nefc), 1 ) \ XNV( int, iefc_J_rowsuper, MJ_D(nefc), 1 ) \ XNV( int, iefc_J_colind, MJ_D(nJ), 1 ) \ - XNV( int, iefc_JT_rownnz, MJ_D(nidof), 1 ) \ - XNV( int, iefc_JT_rowadr, MJ_D(nidof), 1 ) \ - XNV( int, iefc_JT_rowsuper, MJ_D(nidof), 1 ) \ - XNV( int, iefc_JT_colind, MJ_D(nJ), 1 ) \ XNV( mjtNum, iefc_J, MJ_D(nJ), 1 ) \ - XNV( mjtNum, iefc_JT, MJ_D(nJ), 1 ) \ X ( mjtNum, iefc_frictionloss, MJ_D(nefc), 1 ) \ X ( mjtNum, iefc_D, MJ_D(nefc), 1 ) \ X ( mjtNum, iefc_R, MJ_D(nefc), 1 ) \ diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index 4407bd8d..537b7394 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -5914,38 +5914,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='column indices in constraint Jacobian', array_extent=('nJ',), ), - StructFieldDecl( - name='efc_JT_rownnz', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='number of non-zeros in constraint Jacobian row T', - array_extent=('nv',), - ), - StructFieldDecl( - name='efc_JT_rowadr', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='row start address in colind array T', - array_extent=('nv',), - ), - StructFieldDecl( - name='efc_JT_rowsuper', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='number of subsequent rows in supernode T', - array_extent=('nv',), - ), - StructFieldDecl( - name='efc_JT_colind', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='column indices in constraint Jacobian T', - array_extent=('nJ',), - ), StructFieldDecl( name='efc_J', type=PointerType( @@ -5954,14 +5922,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='constraint Jacobian', array_extent=('nJ',), ), - StructFieldDecl( - name='efc_JT', - type=PointerType( - inner_type=ValueType(name='mjtNum'), - ), - doc='constraint Jacobian transposed', - array_extent=('nJ',), - ), StructFieldDecl( name='efc_pos', type=PointerType( @@ -6250,38 +6210,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='column indices in constraint Jacobian', array_extent=('nJ',), ), - StructFieldDecl( - name='iefc_JT_rownnz', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='number of non-zeros in constraint Jacobian row T', - array_extent=('nidof',), - ), - StructFieldDecl( - name='iefc_JT_rowadr', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='row start address in colind array T', - array_extent=('nidof',), - ), - StructFieldDecl( - name='iefc_JT_rowsuper', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='number of subsequent rows in supernode T', - array_extent=('nidof',), - ), - StructFieldDecl( - name='iefc_JT_colind', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='column indices in constraint Jacobian T', - array_extent=('nJ',), - ), StructFieldDecl( name='iefc_J', type=PointerType( @@ -6290,14 +6218,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='constraint Jacobian', array_extent=('nJ',), ), - StructFieldDecl( - name='iefc_JT', - type=PointerType( - inner_type=ValueType(name='mjtNum'), - ), - doc='constraint Jacobian transposed', - array_extent=('nJ',), - ), StructFieldDecl( name='iefc_frictionloss', type=PointerType( diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 63c17db2..b30f1b1a 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -386,10 +386,10 @@ void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* } // sparse Jacobian - if (mj_isSparse(m)) - mju_mulMatVecSparse(res, d->efc_JT, vec, m->nv, - d->efc_JT_rownnz, d->efc_JT_rowadr, - d->efc_JT_colind, d->efc_JT_rowsuper); + if (mj_isSparse(m)) { + mju_mulMatTVecSparse(res, d->efc_J, vec, d->nefc, m->nv, + d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind); + } // dense Jacobian else { @@ -1996,12 +1996,6 @@ void mj_makeConstraint(const mjModel* m, mjData* d) { // transpose sparse Jacobian, make row supernodes if (mj_isSparse(m)) { - // transpose - mju_transposeSparse(d->efc_JT, d->efc_J, d->nefc, m->nv, - d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper, - d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind); - - #ifdef mjUSEAVX // compute supernodes of J; used by mju_mulMatVecSparse_avx mju_superSparse(d->nefc, d->efc_J_rowsuper, diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 8382776d..06d4f7d7 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -597,20 +597,6 @@ void mj_island(const mjModel* m, mjData* d) { mju_superSparse(d->island_nefc[island], d->iefc_J_rowsuper + adr, d->iefc_J_rownnz + adr, d->iefc_J_rowadr + adr, d->iefc_J_colind); } - - // block-diagonalize Jacobian-transpose - mju_blockDiagSparse(d->iefc_JT, d->iefc_JT_rownnz, d->iefc_JT_rowadr, d->iefc_JT_colind, - d->efc_JT, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, - nidof, nisland, - d->map_idof2dof, d->map_efc2iefc, - d->island_idofadr, d->island_iefcadr, NULL, NULL); - - // recompute rowsuper per island - for (int island=0; island < nisland; island++) { - int adr = d->island_idofadr[island]; - mju_superSparse(d->island_nv[island], d->iefc_JT_rowsuper + adr, - d->iefc_JT_rownnz + adr, d->iefc_JT_rowadr + adr, d->iefc_JT_colind); - } } // copy position-dependent efc vectors required by solver diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 20cea7f5..c31b021c 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1236,8 +1236,6 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena printArrayInt("EFC_J_ROWADR", d->nefc, 1, d->efc_J_rowadr, fp); printSparse("EFC_J", d->efc_J, d->nefc, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, fp, float_format); - mj_printSparsity("JT: constraint Jacobian transposed", m->nv, d->nefc, d->efc_JT_rowadr, NULL, - d->efc_JT_rownnz, d->efc_JT_rowsuper, d->efc_JT_colind, fp); if (mj_isDual(m)) { printArrayInt("EFC_AR_ROWNNZ", d->nefc, 1, d->efc_AR_rownnz, fp); printArrayInt("EFC_AR_ROWADR", d->nefc, 1, d->efc_AR_rowadr, fp); diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc index fdeadb31..02bfab1e 100644 --- a/test/benchmark/engine_util_sparse_benchmark_test.cc +++ b/test/benchmark/engine_util_sparse_benchmark_test.cc @@ -448,14 +448,23 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) { } } + int* JT_rownnz = mj_stackAllocInt(d, m->nv); + int* JT_rowadr = mj_stackAllocInt(d, m->nv); + int* JT_rowsuper = mj_stackAllocInt(d, m->nv); + int* JT_colind = mj_stackAllocInt(d, d->nJ); + mjtNum* JT = mj_stackAllocNum(d, d->nJ); + mju_transposeSparse(JT, d->efc_J, d->nefc, m->nv, + JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper, + d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind); + // compute H = J'*D*J, uncompressed layout mju_sqrMatTDUncompressedInit(rowadr, m->nv); - mju_sqrMatTDSparse(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, + mju_sqrMatTDSparse(H, d->efc_J, JT, D, d->nefc, m->nv, rownnz, rowadr, colind, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, d->efc_J_rowsuper, - d->efc_JT_rownnz, d->efc_JT_rowadr, - d->efc_JT_colind, d->efc_JT_rowsuper, d, + JT_rownnz, JT_rowadr, + JT_colind, JT_rowsuper, d, diagind); // compute H = M + J'*D*J @@ -605,19 +614,28 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) { } } + int* JT_rownnz = mj_stackAllocInt(d, m->nv); + int* JT_rowadr = mj_stackAllocInt(d, m->nv); + int* JT_rowsuper = mj_stackAllocInt(d, m->nv); + int* JT_colind = mj_stackAllocInt(d, d->nJ); + mjtNum* JT = mj_stackAllocNum(d, d->nJ); + mju_transposeSparse(JT, d->efc_J, d->nefc, m->nv, + JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper, + d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind); + // time benchmark if (func) { mju_sqrMatTDSparseCount(rownnz, rowadr, m->nv, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, - d->efc_JT_rownnz, d->efc_JT_rowadr, - d->efc_JT_colind, nullptr, d, 1); + JT_rownnz, JT_rowadr, + JT_colind, nullptr, d, 1); for (auto s : state) { // compute H = J'*D*J, compressed layout - func(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind, + func(H, d->efc_J, JT, D, d->nefc, m->nv, rownnz, rowadr, colind, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL, - d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, - d->efc_JT_rowsuper, d, diagind); + JT_rownnz, JT_rowadr, JT_colind, + JT_rowsuper, d, diagind); } } else { for (auto s : state) { @@ -627,10 +645,10 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) { // compute H = J'*D*J, uncompressed layout mju_sqrMatTDSparse_baseline( - H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind, + H, d->efc_J, JT, D, d->nefc, m->nv, rownnz, rowadr, colind, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, d->efc_J_rowsuper, - d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, - d->efc_JT_rowsuper, d, /*unused=*/nullptr); + JT_rownnz, JT_rowadr, JT_colind, + JT_rowsuper, d, /*unused=*/nullptr); } } diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 0329aca5..4a6b1865 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -467,27 +467,6 @@ TEST_F(IslandTest, IslandJacobian) { EXPECT_EQ(J_island[i * nv_island + j], J[efc * nv + dof]); } } - - // === test JT (if sparse) - - // get pointer to J_island, dense (nefc_island x nv_island) submatrix - if (jac == mjJAC_SPARSE) { - // dense copy of island in iJ (here used as scratch) - mju_sparse2dense(iJ, d->iefc_JT, nv_island, nefc_island, - d->iefc_JT_rownnz + idof, - d->iefc_JT_rowadr + idof, - d->iefc_JT_colind); - J_island = iJ; - - // sequential memory in J_island equals random access memory in J - for (int i=0; i < nv_island; i++) { - for (int j=0; j < nefc_island; j++) { - int dof = d->map_idof2dof[idof + i]; - int efc = d->map_iefc2efc[iefc + j]; - EXPECT_EQ(J_island[i * nefc_island + j], J[efc * nv + dof]); - } - } - } } mju_free(iJ); diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index cdc0c118..8bb01d64 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -5001,12 +5001,7 @@ public unsafe struct mjData_ { public int* efc_J_rowadr; public int* efc_J_rowsuper; public int* efc_J_colind; - public int* efc_JT_rownnz; - public int* efc_JT_rowadr; - public int* efc_JT_rowsuper; - public int* efc_JT_colind; public double* efc_J; - public double* efc_JT; public double* efc_pos; public double* efc_margin; public double* efc_frictionloss; @@ -5043,12 +5038,7 @@ public unsafe struct mjData_ { public int* iefc_J_rowadr; public int* iefc_J_rowsuper; public int* iefc_J_colind; - public int* iefc_JT_rownnz; - public int* iefc_JT_rowadr; - public int* iefc_JT_rowsuper; - public int* iefc_JT_colind; public double* iefc_J; - public double* iefc_JT; public double* iefc_frictionloss; public double* iefc_D; public double* iefc_R;