Merge pull request #367 from Simple-Robotics/alloc-integration

Integrate `mimalloc` memory resource as default resource for `SolverProxDDP`, move `blk-matrix.hpp` header to `aligator/core`

Author: ManifoldFR

CHANGELOG.md

@@ -19,6 +19,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- move header `<aligator/gar/blk-matrix.hpp>` to `<aligator/core/blk-matrix>`
+- solver-proxddp : make new `mimalloc` memory resource the allocator resource for ProxDDP solver (https://github.com/Simple-Robotics/aligator/pull/367)
 - tests/gar/block-matrix.cpp : test against BunchKaufman factorization
 - move headers `allocator.hpp` and `arena-matrix.hpp` to `aligator/core` dir (https://github.com/Simple-Robotics/aligator/pull/362)
 - python: aligator now requires eigenpy version 3.10.3 at least

bench/gar-riccati.cpp

@@ -4,6 +4,7 @@
 
 #include <benchmark/benchmark.h>
 
+#include "aligator/core/mimalloc-resource.hpp"
 #include "aligator/gar/proximal-riccati.hpp"
 #include "aligator/gar/parallel-solver.hpp"
 #include "aligator/gar/dense-riccati.hpp"
@@ -22,11 +23,24 @@ static constexpr double mueq = 1e-11;
 static std::mt19937 rng;
 static normal_unary_op normal_op{rng};
 
+static aligator::mimalloc_resource mim_resource;
+
+auto get_allocator(int64_t ID) -> aligator::polymorphic_allocator {
+  static std::pmr::memory_resource *RESOURCES[2] = {
+      std::pmr::get_default_resource(), &mim_resource};
+  assert(ID < 2);
+  return aligator::polymorphic_allocator{RESOURCES[ID]};
+}
+
+#define GET_PROBLEM(state)                                                     \
+  auto allocator = get_allocator(state.range(1));                              \
+  uint horz = (uint)state.range(0);                                            \
+  VectorXs x0 = VectorXs::NullaryExpr(nx, normal_op);                          \
+  LqrProblemTpl<double> problem =                                              \
+      generateLqProblem(rng, x0, horz, nx, nu, 0, nc, true, allocator)
+
 static void BM_serial(benchmark::State &state) {
-  uint horz = (uint)state.range(0);
-  VectorXs x0 = VectorXs::NullaryExpr(nx, normal_op);
-  const LqrProblemTpl<double> problem =
-      generateLqProblem(rng, x0, horz, nx, nu, 0, nc);
+  GET_PROBLEM(state);
   ProximalRiccatiSolver<double> solver(problem);
   auto [xs, us, vs, lbdas] = lqrInitializeSolution(problem);
   for (auto _ : state) {
@@ -37,10 +51,7 @@ static void BM_serial(benchmark::State &state) {
 
 #ifdef ALIGATOR_MULTITHREADING
 template <uint NPROC> static void BM_parallel(benchmark::State &state) {
-  uint horz = (uint)state.range(0);
-  VectorXs x0 = VectorXs::NullaryExpr(nx, normal_op);
-  LqrProblemTpl<double> problem =
-      generateLqProblem(rng, x0, horz, nx, nu, 0, nc);
+  GET_PROBLEM(state);
   ParallelRiccatiSolver<double> solver(problem, NPROC);
   auto [xs, us, vs, lbdas] = lqrInitializeSolution(problem);
   for (auto _ : state) {
@@ -51,10 +62,7 @@ template <uint NPROC> static void BM_parallel(benchmark::State &state) {
 #endif
 
 static void BM_stagedense(benchmark::State &state) {
-  uint horz = (uint)state.range(0);
-  VectorXs x0 = VectorXs::NullaryExpr(nx, normal_op);
-  LqrProblemTpl<double> problem =
-      generateLqProblem(rng, x0, horz, nx, nu, 0, nc);
+  GET_PROBLEM(state);
   RiccatiSolverDense<double> solver(problem);
   auto [xs, us, vs, lbdas] = lqrInitializeSolution(problem);
   for (auto _ : state) {
@@ -65,8 +73,10 @@ static void BM_stagedense(benchmark::State &state) {
 
 static void customArgs(benchmark::internal::Benchmark *b) {
   for (uint e = 4; e <= 10; e++) {
-    b->Arg(1 << e);
+    b->Args({1 << e, 0});
+    b->Args({1 << e, 1});
   }
+  b->ArgNames({"N", "alloc"});
   b->Unit(benchmark::kMillisecond);
   b->UseRealTime();
 }

bindings/python/include/aligator/python/blk-matrix.hpp

@@ -1,24 +1,31 @@
+/// @copyright Copyright (C) 2023-2024 LAAS-CNRS, 2023-2025 INRIA
+/// @author Wilson Jallet
 #include "fwd.hpp"
-#include "aligator/gar/blk-matrix.hpp"
+#include "aligator/core/blk-matrix.hpp"
 
 namespace aligator {
 namespace python {
 namespace bp = boost::python;
 
+template <typename T> struct PrintableVisitor;
+
 template <typename BlockMatrixType> struct BlkMatrixPythonVisitor;
 
 template <typename MatrixType, int N, int M>
 struct BlkMatrixPythonVisitor<BlkMatrix<MatrixType, N, M>>
     : bp::def_visitor<BlkMatrixPythonVisitor<BlkMatrix<MatrixType, N, M>>> {
   using BlockMatrixType = BlkMatrix<MatrixType, N, M>;
   using RefType = Eigen::Ref<MatrixType>;
+  static constexpr bool IsVector = BlockMatrixType::IsVectorAtCompileTime;
 
   using Self = BlkMatrixPythonVisitor<BlockMatrixType>;
 
   static RefType get_block(BlockMatrixType &bmt, size_t i, size_t j) {
     return bmt(i, j);
   }
 
+  static RefType get_block2(BlockMatrixType &bmt, size_t i) { return bmt[i]; }
+
   static RefType blockRow(BlockMatrixType &mat, size_t i) {
     if (i >= mat.rowDims().size()) {
       PyErr_SetString(PyExc_IndexError, "Index out of range.");
@@ -27,19 +34,34 @@ struct BlkMatrixPythonVisitor<BlkMatrix<MatrixType, N, M>>
     return mat.blockRow(i);
   }
 
+  static RefType blockCol(BlockMatrixType &mat, size_t i) {
+    if (i >= mat.rowDims().size()) {
+      PyErr_SetString(PyExc_IndexError, "Index out of range.");
+      bp::throw_error_already_set();
+    }
+    return mat.blockCol(i);
+  }
+
   template <class... Args> void visit(bp::class_<Args...> &obj) const {
     obj.add_property(
            "matrix", +[](BlockMatrixType &m) -> RefType { return m.matrix(); })
-        .def_readonly("rows", &BlockMatrixType::rows)
-        .def_readonly("cols", &BlockMatrixType::cols)
+        .add_property("rows", &BlockMatrixType::rows)
+        .add_property("cols", &BlockMatrixType::cols)
         .add_property("rowDims",
                       bp::make_function(&BlockMatrixType::rowDims,
                                         bp::return_internal_reference<>()))
         .add_property("colDims",
                       bp::make_function(&BlockMatrixType::colDims,
                                         bp::return_internal_reference<>()))
         .def("blockRow", blockRow, "Get a block row by index.")
-        .def("__call__", get_block, ("self"_a, "i", "j"));
+        .def("blockCol", blockCol, "Get a block row by index.")
+        .def("setZero", &BlockMatrixType::setZero, ("self"_a),
+             "Set all coefficients to zero.")
+        .def("__call__", get_block, ("self"_a, "i", "j"))
+        .def(PrintableVisitor<BlockMatrixType>{});
+    if constexpr (BlockMatrixType::IsVectorAtCompileTime) {
+      obj.def("__call__", get_block2, ("self"_a, "i"));
+    }
   }
 
   static void expose(const char *name) {

examples/clqr.cpp

@@ -7,7 +7,7 @@
 #include "aligator/modelling/constraints.hpp"
 #include <aligator/fmt-eigen.hpp>
 #include <iostream>
-#include <random>
+#include "../tests/test_util.hpp"
 
 using namespace aligator;
 
@@ -23,20 +23,15 @@ using context::StageModel;
 using context::TrajOptProblem;
 using context::VectorXs;
 
-static std::mt19937_64 urng{42};
-struct NormalGen {
-  double operator()() const { return norm(urng); }
-  mutable std::normal_distribution<double> norm;
-};
-
 int main() {
+  std::mt19937 rng{42};
   std::srand(42);
   const size_t nsteps = 100;
   const auto nx = 4;
   const auto nu = 2;
   const auto space = Space(nx);
 
-  NormalGen norm_gen;
+  normal_unary_op norm_gen{rng};
   MatrixXs A;
   // clang-format off
   A.setIdentity(nx, nx);

examples/talos_arms.py

@@ -83,7 +83,7 @@ class Args(ArgsBase):
 
 
 TOL = 1e-5
-mu_init = 1e-3
+mu_init = 1e-2
 max_iters = 200
 verbose = aligator.VerboseLevel.VERBOSE
 solver = aligator.SolverProxDDP(TOL, mu_init, verbose=verbose)

include/aligator/gar/blk-matrix.hpp renamed to include/aligator/core/blk-matrix.hpp

@@ -1,4 +1,4 @@
-/// @copyright Copyright (C) 2023-2024 LAAS-CNRS, INRIA
+/// @copyright Copyright (C) 2023-2024 LAAS-CNRS, 2023-2025 INRIA
 /// @author Wilson Jallet
 #pragma once
 
@@ -22,7 +22,9 @@ template <typename _MatrixType, int _N, int _M> class BlkMatrix {
   using PlainObject = typename MatrixType::PlainObject;
   using Scalar = typename MatrixType::Scalar;
   using Index = Eigen::Index;
-  enum { N = _N, M = _M, Options = PlainObject::Options };
+  static constexpr int N = _N;
+  static constexpr int M = _M;
+  static constexpr int Options = PlainObject::Options;
   static constexpr bool IsVectorAtCompileTime =
       MatrixType::IsVectorAtCompileTime;
 

include/aligator/core/lagrangian.hpp

@@ -5,7 +5,7 @@
 #include "aligator/core/stage-data.hpp"
 #include "aligator/core/explicit-dynamics.hpp"
 #include "aligator/core/cost-abstract.hpp"
-#include "aligator/gar/blk-matrix.hpp"
+#include "aligator/core/blk-matrix.hpp"
 #include "aligator/tracy.hpp"
 
 namespace aligator {

include/aligator/gar/block-tridiagonal.hpp

@@ -2,7 +2,7 @@
 /// @author Wilson Jallet
 #pragma once
 
-#include "aligator/gar/blk-matrix.hpp"
+#include "aligator/core/blk-matrix.hpp"
 #include "aligator/utils/exceptions.hpp"
 #include "aligator/tracy.hpp"
 

include/aligator/gar/dense-kernel.hpp

@@ -3,9 +3,9 @@
 #pragma once
 
 #include "aligator/core/bunchkaufman.hpp"
-
-#include "blk-matrix.hpp"
+#include "aligator/core/blk-matrix.hpp"
 #include "lqr-problem.hpp"
+
 #include <boost/core/span.hpp>
 
 namespace aligator::gar {

include/aligator/gar/riccati-kernel.hpp

@@ -2,13 +2,11 @@
 /// @author Wilson Jallet
 #pragma once
 
-#include "aligator/context.hpp"
-#include "lqr-problem.hpp"
-#include "blk-matrix.hpp"
 #include "aligator/core/bunchkaufman.hpp"
+#include "aligator/core/blk-matrix.hpp"
+#include "lqr-problem.hpp"
 
 #include <boost/core/make_span.hpp>
-
 #include <optional>
 
 namespace aligator {