[SymForce] Add PreallocatedLinearizedDenseFactors

**Context**
Currently, if the same `LinearizedDenseFactor` is used as the output
argument to a Factor with different residual and rhs dimensions than
what it already has, Eigen will reallocate the wrong sized members, even
if the new size is smaller.

To avoid this, we ensure that we don't pass the same
`LinearizedDenseFactor` to `Factor`s with different sizes.

Simultaneously, to reduce the amount of memory used, we do not simply
have a seperate `LinearizedDenseFactor` for each dense `Factor` (as we
do for sparse factors), but instead have only one
`LinearizedDenseFactor` for each dense factor shape.

Before this commit, all of this logic was inlined in members of
`Linearizer` and in the body of `Linearizer::BuildInitialLinearization`.

**Motiviation**
The principal motivation for this commit is to allow the same logic to
be reused in a different linearizer class.

As a side benefit, I think it makes
`Linearizer::BuildInitialLinearization` easier to read.

**What I did**
I factored out the logic into `PreallocatedLinearizedDenseFactors`.

An instance of `PreallocatedLinearizedDenseFactors` is stored by the
linearizer, and contains the unique `LinearizedDenseFactors` as well as
the mapping from `Factor`s to the corresponding `LinearizedDenseFactor`.

Usage is detailed in doc-string (but also demoed in the linearizer.cc).

NOTE: requires the user pass in a very awkward temporary unordered_map
to the call to `AppendFactorSize`. I do it this way because I couldn't
think of a way to encapsulate this without either storing more data than
is necessary or adding a lot of extra boilerplate code.

Topic: extract_dense_linearized_factor_consolidation
GitOrigin-RevId: 69c672ec3228b13b4cc55a0a75f3b01dee74eccc

symforce/opt/internal/linearized_dense_factor_pool.h

@@ -0,0 +1,118 @@
+/* ----------------------------------------------------------------------------
+ * SymForce - Copyright 2022, Skydio, Inc.
+ * This source code is under the Apache 2.0 license found in the LICENSE file.
+ * ---------------------------------------------------------------------------- */
+
+#pragma once
+
+#include <vector>
+
+#include "../factor.h"
+#include "./linearizer_utils.h"
+
+namespace sym {
+
+namespace internal {
+
+/**
+ * Class to abstract out the method of getting a temporary LinearizedDenseFactor in which
+ * to evaluate factor's linearizations.
+ *
+ * Exists because adding these features directly to the Linearizer class makes that harder
+ * to read and to allow the functionality to be reused by the dense linearizer.
+ *
+ * At a more abstract level, exists because we can't pass the same MatrixXd's to the different
+ * factors because then Eigen will implicitly reallocate the memory of each MatrixXd every
+ * time the vector size changes. So we pass them MatrixXds who already have the right size.
+ *
+ * This class could be deleted if we could pass a factor a block of memory that can be made
+ * larger if needed by not performing any new allocation if not.
+ *
+ * NOTE(brad): We do not have a corresponding class for the sparse linearized factors because
+ * there we currently have a separate linearized factor per factor. This is to avoid the
+ * need to re-write the column pointers and row indices during each factor call.
+ *
+ * Usage (during creation):
+ *   LinearizedDenseFactorPool<Scalar> linearized_dense_factors;
+ *   ...
+ *   linearized_dense_factors.reserve(number_of_dense_factors);
+ *   // Don't touch size_tracker other than to pass it to AppendFactorSize
+ *   // Can be disposed of after last call to AppendFactorSize
+ *   // Same SizeTracker must be used for every call to AppendFactorSize
+ *   typename LinearizedDenseFactorPool<Scalar>::SizeTracker size_tracker;
+ *
+ *   for (Factor& factor : factors) {
+ *     int res_dim = ...;
+ *     int rhs_dim = ...;
+ *
+ *     linearized_dense_factors.AppendFactorSize(res_dim, rhs_dim, size_tracker);
+ *   }
+ */
+template <typename Scalar>
+class LinearizedDenseFactorPool {
+ public:
+  using SizeTracker = std::unordered_map<std::pair<int, int>, int, internal::StdPairHash>;
+  using LinearizedDenseFactor = typename Factor<Scalar>::LinearizedDenseFactor;
+
+  /**
+   * Reserve space for the number of dense factor to be appended with AppendFactorSize.
+   */
+  void reserve(const size_t dense_factor_count) {
+    index_per_factor_.reserve(dense_factor_count);
+  }
+
+  /**
+   * Call this method for each factor in the same order you wish to index into
+   * PreallocatedLinearizedDenseFactor with.
+   *
+   * Preconditions:
+   *   - dims_to_index was default initialized
+   *   - The same SizeTracker is used for every call to AppendFactorSize
+   *   - dims_to_index has only been mutated by AppendFactorSize with this instance of
+   *     LinearizedDenseFactorPool
+   *
+   * If not already present, adds a linearization with residual and rhs dimensions provided
+   * to unique_linearizations.
+   * Appends to factor_indices the index of the linearization in unique_linearizations with
+   * dimensions res_dim and rhs_dim.
+   */
+  void AppendFactorSize(const int res_dim, const int rhs_dim, SizeTracker& dims_to_index) {
+    // Invariant: dims_to_index contains a key (dim1, dim2) iff AppendFactorSize was called with
+    // those dims. dims_to_index[(dim1, dim2)] = n  ==> AppendFactorSize was first called with
+    // those dims on its nth call.
+    const auto index_at_and_was_inserted =
+        dims_to_index.emplace(std::make_pair(res_dim, rhs_dim), unique_linearized_factors_.size());
+
+    index_per_factor_.push_back(index_at_and_was_inserted.first->second);
+    if (index_at_and_was_inserted.second) {
+      unique_linearized_factors_.emplace_back();
+      LinearizedDenseFactor& new_linearization = unique_linearized_factors_.back();
+
+      new_linearization.residual.resize(res_dim, 1);
+      new_linearization.jacobian.resize(res_dim, rhs_dim);
+      new_linearization.hessian.resize(rhs_dim, rhs_dim);
+      new_linearization.rhs.resize(rhs_dim, 1);
+    }
+  }
+
+  /**
+   * Returns a linearized dense factor whose size is correct for the dense_index'th dense
+   * factor.
+   */
+  LinearizedDenseFactor& at(const int dense_index) {
+    return unique_linearized_factors_.at(index_per_factor_.at(dense_index));
+  }
+
+ private:
+  // One LinearizedDenseFactor for each unique pair of res_dim and rhs_dim seen
+  std::vector<LinearizedDenseFactor> unique_linearized_factors_;
+
+  // Indices into unique_linearized_factors_; one per dense factor;
+  // unique_linearized_factors_[index_per_factor_[n]] has the same res_dim and rhs_dim
+  // as the nth call to AppendFactorSize
+  std::vector<int> index_per_factor_;
+};
+
+}  // namespace internal
+
+}  // namespace sym

symforce/opt/linearizer.cc

@@ -44,7 +44,6 @@ Linearizer<ScalarType>::Linearizer(const std::string& name,
   linearized_sparse_factors_.resize(num_sparse_factors);
   sparse_factor_update_helpers_.reserve(num_sparse_factors);
 
-  linearized_dense_factor_indices_.reserve(num_dense_factors);
   dense_factor_update_helpers_.reserve(num_dense_factors);
 }
 
@@ -77,8 +76,7 @@ void Linearizer<ScalarType>::Relinearize(const Values<Scalar>& values,
         ++sparse_idx;
       } else {
         // Use temporary with the right size to avoid allocating after initialization.
-        auto& linearized_dense_factor =
-            linearized_dense_factors_.at(linearized_dense_factor_indices_.at(dense_idx));
+        auto& linearized_dense_factor = linearized_dense_factors_.at(dense_idx);
         // TODO: Only compute factor Jacobians when include_jacobians_ is true.
         factor.Linearize(values, linearized_dense_factor, &factor_indices_[i]);
 
@@ -183,10 +181,11 @@ void Linearizer<ScalarType>::BuildInitialLinearization(const Values<Scalar>& val
   // Track these to make sure that all combined keys are touched by at least one factor.
   std::unordered_set<Key> keys_touched_by_factors;
 
+  linearized_dense_factors_.reserve(factors_->size() - linearized_sparse_factors_.size());
+  typename internal::LinearizedDenseFactorPool<Scalar>::SizeTracker dense_factor_size_tracker;
+
   // Evaluate all factors, processing the dense ones in place and storing the sparse ones for
   // later
-  std::unordered_map<std::pair<int, int>, int, internal::StdPairHash>
-      linearized_dense_factor_indices_by_dims;
   LinearizedDenseFactor linearized_dense_factor{};
   size_t sparse_idx{0};
   factor_indices_.reserve(factors_->size());
@@ -214,16 +213,9 @@ void Linearizer<ScalarType>::BuildInitialLinearization(const Values<Scalar>& val
       factor.Linearize(values, linearized_dense_factor, &factor_indices_.back());
 
       // Make sure a temporary of the right dimension is kept for relinearizations
-      const std::pair<int, int> dims{linearized_dense_factor.jacobian.rows(),
-                                     linearized_dense_factor.jacobian.cols()};
-
-      const auto linearized_dense_factors_it_and_inserted =
-          linearized_dense_factor_indices_by_dims.emplace(dims, linearized_dense_factors_.size());
-      if (linearized_dense_factors_it_and_inserted.second) {
-        linearized_dense_factors_.push_back(linearized_dense_factor);
-      }
-      linearized_dense_factor_indices_.push_back(
-          linearized_dense_factors_it_and_inserted.first->second);
+      linearized_dense_factors_.AppendFactorSize(linearized_dense_factor.residual.rows(),
+                                                 linearized_dense_factor.rhs.rows(),
+                                                 dense_factor_size_tracker);
 
       // Create dense factor helper
       auto helper_and_dimension =

symforce/opt/linearizer.h

@@ -13,6 +13,7 @@
 #include <lcmtypes/sym/linearization_sparse_factor_helper_t.hpp>
 
 #include "./factor.h"
+#include "./internal/linearized_dense_factor_pool.h"
 #include "./linearization.h"
 #include "./values.h"
 
@@ -144,10 +145,8 @@ class Linearizer {
   bool include_jacobians_;
 
   // Linearized factors - stores individual factor residuals, jacobians, etc
-  std::vector<LinearizedDenseFactor> linearized_dense_factors_;  // one per Jacobian shape
-  std::vector<int> linearized_dense_factor_indices_;  // index into linearized_dense_factors_; one
-                                                      // per dense factor
-  std::vector<LinearizedSparseFactor> linearized_sparse_factors_;  // one per sparse factor
+  internal::LinearizedDenseFactorPool<Scalar> linearized_dense_factors_;  // one per Jacobian shape
+  std::vector<LinearizedSparseFactor> linearized_sparse_factors_;         // one per sparse factor
 
   // Keys that form the state vector
   std::vector<Key> keys_;