Program Listing for File rnea.hpp

↰ Return to documentation for file (include/pinocchio/algorithm/parallel/rnea.hpp)

//
// Copyright (c) 2021 INRIA
//

#ifndef __pinocchio_algorithm_parallel_rnea_hpp__
#define __pinocchio_algorithm_parallel_rnea_hpp__

#include <omp.h>

#include "pinocchio/multibody/pool/model.hpp"
#include "pinocchio/algorithm/rnea.hpp"

namespace pinocchio
{
  template<typename Scalar, int Options, template<typename,int> class JointCollectionTpl, typename ConfigVectorPool, typename TangentVectorPool1, typename TangentVectorPool2, typename TangentVectorPool3>
  void rnea(const int num_threads,
            ModelPoolTpl<Scalar,Options,JointCollectionTpl> & pool,
            const Eigen::MatrixBase<ConfigVectorPool> & q,
            const Eigen::MatrixBase<TangentVectorPool1> & v,
            const Eigen::MatrixBase<TangentVectorPool2> & a,
            const Eigen::MatrixBase<TangentVectorPool3> & tau)
  {
    typedef ModelPoolTpl<Scalar,Options,JointCollectionTpl> Pool;
    typedef typename Pool::Model Model;
    typedef typename Pool::Data Data;
    typedef typename Pool::DataVector DataVector;

    const Model & model = pool.model();
    DataVector & datas = pool.datas();
    TangentVectorPool3 & res = tau.const_cast_derived();

    PINOCCHIO_CHECK_INPUT_ARGUMENT(num_threads <= pool.size(), "The pool is too small");
    PINOCCHIO_CHECK_ARGUMENT_SIZE(q.rows(), model.nq);
    PINOCCHIO_CHECK_ARGUMENT_SIZE(v.rows(), model.nv);
    PINOCCHIO_CHECK_ARGUMENT_SIZE(a.rows(), model.nv);
    PINOCCHIO_CHECK_ARGUMENT_SIZE(res.rows(), model.nv);

    PINOCCHIO_CHECK_ARGUMENT_SIZE(q.cols(), v.cols());
    PINOCCHIO_CHECK_ARGUMENT_SIZE(q.cols(), a.cols());
    PINOCCHIO_CHECK_ARGUMENT_SIZE(q.cols(), res.cols());

    omp_set_num_threads(num_threads);
    const Eigen::DenseIndex batch_size = res.cols();
    Eigen::DenseIndex i = 0;

#pragma omp parallel for
    for(i = 0; i < batch_size; i++)
    {
      const int thread_id = omp_get_thread_num();
      Data & data = datas[(size_t)thread_id];
      res.col(i) = rnea(model,data,q.col(i),v.col(i),a.col(i));
    }
  }
}

#endif // ifndef __pinocchio_algorithm_parallel_rnea_hpp__