Skip redundant kinematics in inner-term forward sweep

The per-column forward sweep in ComputeABASecondOrderDerivatives was
recomputing q-dependent quantities (jdata via jmodel.calc, liMi, oMi,
J_cols, and the leaf body inertia in world frame) on every one of the
3*nv iterations, even though they only depend on q and are already
populated by the preceding computeABADerivatives and
ComputeRNEASecondOrderDerivatives calls.

Restrict the per-column forward visitor to the acceleration-dependent
work only: propagate data.oa[i] with zero gravity at the root, write
ddJ_cols = oa x J_cols, reset data.oYcrb[i] to the cached leaf value
(which the backward sweep then composites in place), and compute
data.of[i] = oYcrb[i] * oa. The leaf body inertia is cached once in
oYcrb_leaf before the loop, since the backward sweep mutates
data.oYcrb on each per-column iteration.

For humanoidRandom (nv=32) this drops ComputeABASecondOrderDerivatives
median from ~2765us to ~2524us (~9 percent faster).

Refs: stack-of-tasks#2891

Author: shubhamsingh91

include/pinocchio/src/algorithm/aba-second-order-derivatives.hxx

@@ -16,32 +16,41 @@
 namespace pinocchio
 {
 
-  // Forward sweep used by ComputeABASecondOrderDerivatives to set up the
-  // articulated-body-style inner-product recursion: forward kinematics, the
-  // body spatial inertia oYcrb[i] (composited into a CRBA-style composite
-  // inertia by the backward sweep), and the spatial force oY * oa
-  // corresponding to the input acceleration `a` propagated through the chain.
-  // Gravity is deliberately zeroed at the base so the result of the
-  // subsequent backward sweep is (dM/dq)*a only -- no gravity contribution.
+  // Forward sweep used by ComputeABASecondOrderDerivatives to propagate the
+  // input acceleration `a` through the chain. Only acceleration-dependent
+  // quantities are recomputed per call:
+  //   - data.oa[i]: spatial acceleration with zero gravity at the root, so
+  //     the backward sweep yields (dM/dq)*a only (no gravity contribution).
+  //   - ddJ_cols = oa x J_cols.
+  //   - data.of[i] = oYcrb_leaf[i] * oa, the spatial force at body i.
+  //   - data.oYcrb[i] is reset to oYcrb_leaf[i], since the backward sweep
+  //     composites oYcrb in place and we need the leaf values at the start
+  //     of each per-column iteration.
+  // q-dependent quantities (data.oMi, data.liMi, data.J, data.joints[i] with
+  // a valid jdata.S(), and the world-frame leaf inertia oYcrb_leaf cached by
+  // the caller) are assumed to already be set by the prior computeABADerivatives
+  // and ComputeRNEASecondOrderDerivatives calls, so we deliberately skip
+  // jmodel.calc and the liMi/oMi/J_cols updates that a standard ABA forward
+  // sweep would perform.
   template<
     typename Scalar,
     int Options,
     template<typename, int> class JointCollectionTpl,
-    typename ConfigVectorType,
     typename InputAccType>
   struct ComputeABASecondOrderDerivativesInnerForwardStep
   : public fusion::JointUnaryVisitorBase<ComputeABASecondOrderDerivativesInnerForwardStep<
       Scalar,
       Options,
       JointCollectionTpl,
-      ConfigVectorType,
       InputAccType>>
   {
     typedef ModelTpl<Scalar, Options, JointCollectionTpl> Model;
     typedef DataTpl<Scalar, Options, JointCollectionTpl> Data;
+    typedef typename Data::Inertia Inertia;
+    typedef std::vector<Inertia> InertiaVector;
 
     typedef boost::fusion::
-      vector<const Model &, Data &, const ConfigVectorType &, const InputAccType &>
+      vector<const Model &, Data &, const InputAccType &, const InertiaVector &>
         ArgsType;
 
     template<typename JointModel>
@@ -50,44 +59,32 @@ namespace pinocchio
       JointDataBase<typename JointModel::JointDataDerived> & jdata,
       const Model & model,
       Data & data,
-      const Eigen::MatrixBase<ConfigVectorType> & q,
-      const Eigen::MatrixBase<InputAccType> & a)
+      const Eigen::MatrixBase<InputAccType> & a,
+      const InertiaVector & oYcrb_leaf)
     {
       typedef typename Model::JointIndex JointIndex;
       typedef typename Data::Motion Motion;
-      typedef typename Data::Inertia Inertia;
 
       const JointIndex i = jmodel.id();
       const JointIndex parent = model.parents[i];
       Motion & oa = data.oa[i];
 
-      jmodel.calc(jdata.derived(), q.derived());
-      data.liMi[i] = model.jointPlacements[i] * jdata.M();
-
       if (parent > 0)
-      {
-        data.oMi[i] = data.oMi[parent] * data.liMi[i];
         oa = data.oa[parent];
-      }
       else
-      {
-        data.oMi[i] = data.liMi[i];
         oa.setZero(); // zero gravity: result is (dM/dq)*a, not (dM/dq)*a + dh/dq
-      }
 
       typedef
         typename SizeDepType<JointModel::NV>::template ColsReturn<typename Data::Matrix6x>::Type
           ColsBlock;
       ColsBlock J_cols = jmodel.jointCols(data.J);
       ColsBlock ddJ_cols = jmodel.jointCols(data.ddJ);
 
-      J_cols.noalias() = data.oMi[i].act(jdata.S());
       motionSet::motionAction(oa, J_cols, ddJ_cols);
       oa += data.oMi[i].act(jdata.S() * jmodel.jointVelocitySelector(a));
 
-      Inertia & oY = data.oYcrb[i];
-      oY = data.oMi[i].act(model.inertias[i]);
-      data.of[i] = oY * oa;
+      data.oYcrb[i] = oYcrb_leaf[i];
+      data.of[i] = data.oYcrb[i] * oa;
     }
   };
 
@@ -216,9 +213,7 @@ namespace pinocchio
     // 1) First-order ABA derivatives: fills data.ddq, data.ddq_dq, data.ddq_dv
     //    and (upper triangle of) data.Minv.
     computeABADerivatives(model, data, q.derived(), v.derived(), tau.derived());
-    // Symmetrize Minv: computeABADerivatives only fills the upper triangle, but
-    // the dense GEMM in step 4 and the inner-term sweep for the dtaudq block
-    // (step 3, fill_inner(data.Minv, prodqdd)) both require the full matrix.
+    // Symmetrize Minv: computeABADerivatives only fills the upper triangle
     data.Minv.template triangularView<Eigen::StrictlyLower>() =
       data.Minv.transpose().template triangularView<Eigen::StrictlyLower>();
 
@@ -230,13 +225,26 @@ namespace pinocchio
 
     // 3) Inner-Term: for each column w of (ddq_dq, ddq_dv, Minv), compute
     //    (dM/dq) * column and store as page w of the corresponding tensor.
+    //
+    // The per-column forward sweep below is the acceleration-only pass; it
+    // reuses the q-dependent kinematics (data.oMi, data.J, data.joints[i])
+    // already populated by computeABADerivatives and SO RNEA, and the leaf
+    // body inertias cached in oYcrb_leaf right below.
     typedef ComputeABASecondOrderDerivativesInnerForwardStep<
-      Scalar, Options, JointCollectionTpl, ConfigVectorType, VectorXs>
+      Scalar, Options, JointCollectionTpl, VectorXs>
       InnerPass1;
     typedef ComputeABASecondOrderDerivativesInnerBackwardStep<
       Scalar, Options, JointCollectionTpl, Matrix6x, MatrixXs>
       InnerPass2;
 
+    // Cache the world-frame leaf body inertia once. SO RNEA leaves
+    // data.oYcrb in its composited state, and the inner-term backward sweep
+    // mutates data.oYcrb again per call, so we restore the leaf values at
+    // the start of each per-column forward pass.
+    typename InnerPass1::InertiaVector oYcrb_leaf((std::size_t)model.njoints);
+    for (JointIndex i = 1; i < (JointIndex)model.njoints; ++i)
+      oYcrb_leaf[i] = data.oMi[i].act(model.inertias[i]);
+
     // TODO: preallocate these scratch buffers (three nv-cubed tensors, two
     // 6-by-nv matrices, one nv-by-nv matrix, plus the two nv-by-4*nv^2
     // matrices in step 4) as Data fields to avoid the per-call allocations
@@ -254,7 +262,7 @@ namespace pinocchio
         for (JointIndex i = 1; i < (JointIndex)model.njoints; ++i)
           InnerPass1::run(
             model.joints[i], data.joints[i],
-            typename InnerPass1::ArgsType(model, data, q.derived(), a_input));
+            typename InnerPass1::ArgsType(model, data, a_input, oYcrb_leaf));
 
         Ftmp1.setZero();
         Ftmp3.setZero();