Skip to content

Commit

Permalink
[gar] Change LQRKnotTpl class to work with single block of allocate…
Browse files Browse the repository at this point in the history
…d memory per block

+ using `MemReq`

Give no_init_t tag type a constexpr explicit ctor
  • Loading branch information
ManifoldFR committed Nov 1, 2024
1 parent 8c0fc17 commit 274a13b
Show file tree
Hide file tree
Showing 4 changed files with 214 additions and 45 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- [gar] Rework `RiccatiSolverDense` to not use inner struct `FactorData`
- Various changes to `gar` tests and `test_util`, add `LQRKnot::isApprox()`
- Add `MemReq` struct to handle requests for single blocks of memory for multiple arrays
- Change `LQRKnotTpl` to work with a single block of allocated memory using `MemReq`

### Removed

Expand Down
5 changes: 3 additions & 2 deletions bindings/python/src/gar/expose-gar.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,9 @@ void exposeGAR() {
exposeBlockMatrices();

bp::class_<knot_t>("LQRKnot", bp::no_init)
.def(bp::init<uint, uint, uint>(("nx"_a, "nu", "nc")))
.def(bp::init<uint, uint, uint, uint>(("nx"_a, "nu", "nc", "nx2")))
.def(bp::init<uint, uint, uint>(("self"_a, "nx", "nu", "nc")))
.def(bp::init<uint, uint, uint, uint, uint>(
("self"_a, "nx"_a, "nu", "nc", "nx2", "nth"_a = 0)))
.def_readonly("nx", &knot_t::nx)
.def_readonly("nu", &knot_t::nu)
.def_readonly("nc", &knot_t::nc)
Expand Down
71 changes: 59 additions & 12 deletions gar/include/aligator/gar/lqr-problem.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
#pragma once

#include "aligator/math.hpp"
#include "mem-req.hpp"
#include <fmt/format.h>

#include <optional>
Expand All @@ -27,35 +28,81 @@ namespace gar {
///
template <typename Scalar> struct LQRKnotTpl {
ALIGATOR_DYNAMIC_TYPEDEFS(Scalar);
enum { Alignment = Eigen::AlignedMax };
using VectorMap = Eigen::Map<VectorXs, Alignment>;
using MatrixMap = Eigen::Map<MatrixXs, Alignment>;

// tag type
struct no_alloc_t {
explicit constexpr no_alloc_t() {}
};
static constexpr no_alloc_t no_alloc{};

uint nx, nu, nc, nx2, nth;
MatrixXs Q, S, R;
VectorXs q, r;
MatrixXs A, B, E;
VectorXs f;
MatrixXs C, D;
VectorXs d;

MatrixXs Gth;
MatrixXs Gx;
MatrixXs Gu;
MatrixXs Gv;
VectorXs gamma;

MatrixMap Q, S, R;
VectorMap q, r;
MatrixMap A, B, E;
VectorMap f;
MatrixMap C, D;
VectorMap d;

MatrixMap Gth;
MatrixMap Gx;
MatrixMap Gu;
MatrixMap Gv;
VectorMap gamma;

LQRKnotTpl(uint nx, uint nu, uint nc, uint nx2, uint nth = 0);

LQRKnotTpl(uint nx, uint nu, uint nc) : LQRKnotTpl(nx, nu, nc, nx) {}

void allocate();
// initialize the matrices.
void initialize();
// reallocates entire buffer for contigousness
void addParameterization(uint nth);

LQRKnotTpl(const LQRKnotTpl &other);
LQRKnotTpl(LQRKnotTpl &&other);
LQRKnotTpl &operator=(const LQRKnotTpl &other);
LQRKnotTpl &operator=(LQRKnotTpl &&other);

~LQRKnotTpl();

friend void swap(LQRKnotTpl &lhs, LQRKnotTpl &rhs) {
using std::swap;
swap(lhs.nx, rhs.nx);
swap(lhs.nu, rhs.nu);
swap(lhs.nc, rhs.nc);
swap(lhs.nx2, rhs.nx2);
swap(lhs.nth, rhs.nth);
// only swap the memory ptr, do not swap the Eigen::Map objects.
swap(lhs.memory, rhs.memory);
swap(lhs.req, rhs.req);

lhs.initialize();
rhs.initialize();
}

bool isApprox(const LQRKnotTpl &other,
Scalar prec = std::numeric_limits<Scalar>::epsilon()) const;

friend bool operator==(const LQRKnotTpl &lhs, const LQRKnotTpl &rhs) {
return lhs.isApprox(rhs);
}

private:
LQRKnotTpl(no_alloc_t, uint nx, uint nu, uint nc, uint nx2, uint nth);
Scalar *memory;
MemReq req;
};

template <typename Scalar> LQRKnotTpl<Scalar>::~LQRKnotTpl() {
if (memory)
std::free(memory);
}

template <typename Scalar> struct LQRProblemTpl {
ALIGATOR_DYNAMIC_TYPEDEFS(Scalar);
using KnotType = LQRKnotTpl<Scalar>;
Expand Down
182 changes: 151 additions & 31 deletions gar/include/aligator/gar/lqr-problem.hxx
Original file line number Diff line number Diff line change
Expand Up @@ -4,43 +4,163 @@

namespace aligator::gar {

namespace detail {
template <typename EigenType, int MapOptions,
typename Scalar = typename EigenType::Scalar>
void emplaceMap(Eigen::Map<EigenType, MapOptions> &map, long size,
Scalar *ptr) {
using MapType = Eigen::Map<EigenType, MapOptions>;
new (&map) MapType{ptr, size};
}

/// \brief Placement-new a map type using the provided memory pointer.
template <typename EigenType, int MapOptions,
typename Scalar = typename EigenType::Scalar>
void emplaceMap(Eigen::Map<EigenType, MapOptions> &map, long rows, long cols,
Scalar *ptr) {
using MapType = Eigen::Map<EigenType, MapOptions>;
new (&map) MapType{ptr, rows, cols};
}
} // namespace detail

template <typename Scalar>
LQRKnotTpl<Scalar>::LQRKnotTpl(no_alloc_t, uint nx, uint nu, uint nc, uint nx2,
uint nth)
: nx(nx), nu(nu), nc(nc), nx2(nx2), nth(nth), //
Q(NULL, 0, 0), S(NULL, 0, 0), R(NULL, 0, 0), q(NULL, 0), r(NULL, 0), //
A(NULL, 0, 0), B(NULL, 0, 0), E(NULL, 0, 0), f(NULL, 0), //
C(NULL, 0, 0), D(NULL, 0, 0), d(NULL, 0), //
Gth(NULL, 0, 0), Gx(NULL, 0, 0), Gu(NULL, 0, 0), Gv(NULL, 0, 0),
gamma(NULL, 0), //
memory(NULL), req(Alignment) {}

template <typename Scalar>
LQRKnotTpl<Scalar>::LQRKnotTpl(uint nx, uint nu, uint nc, uint nx2, uint nth)
: nx(nx), nu(nu), nc(nc), nx2(nx2), nth(nth), //
Q(nx, nx), S(nx, nu), R(nu, nu), q(nx), r(nu), //
A(nx2, nx), B(nx2, nu), E(nx2, nx), f(nx2), //
C(nc, nx), D(nc, nu), d(nc), Gth(nth, nth), Gx(nx, nth), Gu(nu, nth),
Gv(nc, nth), gamma(nth) {
Q.setZero();
S.setZero();
R.setZero();
q.setZero();
r.setZero();

A.setZero();
B.setZero();
E.setZero();
f.setZero();

C.setZero();
D.setZero();
d.setZero();

Gth.setZero();
Gx.setZero();
Gu.setZero();
Gv.setZero();
gamma.setZero();
: LQRKnotTpl(no_alloc, nx, nu, nc, nx2, nth) {

this->allocate();
this->initialize();
}

template <typename Scalar> void LQRKnotTpl<Scalar>::allocate() {
req.addArray<double>(nx, nx) // Q
.addArray<double>(nx, nu) // S
.addArray<double>(nu, nu) // R
.addArray<double>(nx) // q
.addArray<double>(nu) // r
.addArray<double>(nx2, nx) // A
.addArray<double>(nx2, nu) // B
.addArray<double>(nx2, nx2) // E
.addArray<double>(nx2) // f
.addArray<double>(nc, nx) // C
.addArray<double>(nc, nu) // D
.addArray<double>(nc) // d
.addArray<double>(nth, nth) // Gth
.addArray<double>(nx, nth) // Gx
.addArray<double>(nu, nth) // Gu
.addArray<double>(nc, nth) // Gv
.addArray<double>(nth); // gamma

this->memory = static_cast<Scalar *>(req.allocate());
std::memset(memory, 0, req.totalBytes());
}

template <typename Scalar> void LQRKnotTpl<Scalar>::initialize() {
Scalar *ptr = memory;
detail::emplaceMap(Q, nx, nx, ptr);
req.advance(ptr);
detail::emplaceMap(S, nx, nu, ptr);
req.advance(ptr);
detail::emplaceMap(R, nu, nu, ptr);
req.advance(ptr);
detail::emplaceMap(q, nx, ptr);
req.advance(ptr);
detail::emplaceMap(r, nu, ptr);
req.advance(ptr);

detail::emplaceMap(A, nx2, nx, ptr);
req.advance(ptr);
detail::emplaceMap(B, nx2, nu, ptr);
req.advance(ptr);
detail::emplaceMap(E, nx2, nx2, ptr);
req.advance(ptr);
detail::emplaceMap(f, nx2, ptr);
req.advance(ptr);

detail::emplaceMap(C, nc, nx, ptr);
req.advance(ptr);
detail::emplaceMap(D, nc, nu, ptr);
req.advance(ptr);
detail::emplaceMap(d, nc, ptr);
req.advance(ptr);

detail::emplaceMap(Gth, nth, nth, ptr);
req.advance(ptr);
detail::emplaceMap(Gx, nx, nth, ptr);
req.advance(ptr);
detail::emplaceMap(Gu, nu, nth, ptr);
req.advance(ptr);
detail::emplaceMap(Gv, nc, nth, ptr);
req.advance(ptr);
detail::emplaceMap(gamma, nth, ptr);
req.advance(ptr);

req.reset();
}

template <typename Scalar>
void LQRKnotTpl<Scalar>::addParameterization(uint nth) {
this->nth = nth;
Gth.setZero(nth, nth);
Gx.setZero(nx, nth);
Gu.setZero(nu, nth);
Gv.setZero(nc, nth);
gamma.setZero(nth);
LQRKnotTpl copy(nx, nu, nc, nx2, nth);
copy.Q = Q;
copy.S = S;
copy.R = R;
copy.q = q;
copy.r = r;

copy.A = A;
copy.B = B;
copy.E = E;
copy.f = f;

copy.C = C;
copy.D = D;
copy.d = d;

*this = LQRKnotTpl{copy};
}

template <typename Scalar>
LQRKnotTpl<Scalar>::LQRKnotTpl(const LQRKnotTpl &other)
: LQRKnotTpl(no_alloc, other.nx, other.nu, other.nc, other.nx2, other.nth) {
this->allocate();
assert(req.totalBytes() == other.req.totalBytes());
// copy memory over from other
std::memcpy(memory, other.memory, other.req.totalBytes());
this->initialize();
}

template <typename Scalar>
LQRKnotTpl<Scalar>::LQRKnotTpl(LQRKnotTpl &&other)
: LQRKnotTpl(no_alloc, other.nx, other.nu, other.nc, other.nx2, other.nth) {
// no need to allocate, just bring in the other
// memory buffer
memory = other.memory;
other.memory = NULL;
req = other.req;
this->initialize();
}

template <typename Scalar>
LQRKnotTpl<Scalar> &LQRKnotTpl<Scalar>::operator=(const LQRKnotTpl &other) {
this->~LQRKnotTpl();
new (this) LQRKnotTpl{other};
return *this;
}

template <typename Scalar>
LQRKnotTpl<Scalar> &LQRKnotTpl<Scalar>::operator=(LQRKnotTpl &&other) {
swap(*this, other);
return *this;
}

template <typename Scalar>
Expand Down

0 comments on commit 274a13b

Please sign in to comment.