perf: TBB parallelization for HessianFactor.

tzvist · tzvist · commit 1e72eb19be97 · 2026-01-17T16:46:21.000+02:00
Parallelize the updateHessian operation when forming A'*A in the
HessianFactor merge constructor. For large matrices (&gt;50 rows), the
work is split across columns using TBB parallel_for, with each thread
updating a disjoint set of block columns.

Key changes:
- Add column-range overload of updateHessian() to GaussianFactor,
  HessianFactor, JacobianFactor, and RegularImplicitSchurFactor
- Add setZeroColumns() method to SymmetricBlockMatrix for efficient
  column-wise zeroing using memset

In my toy example it reducded updateHessian time from ~72s to ~41s.
(with 10 threads and GTSAM_TBB_BOUNDED_MEMORY_GROWTH=OFF)
diff --git a/gtsam/base/SymmetricBlockMatrix.h b/gtsam/base/SymmetricBlockMatrix.h
@@ -25,6 +25,7 @@
 #include <boost/serialization/nvp.hpp>
 #endif
 #include <cassert>
+#include <cstring>
 #include <stdexcept>
 #include <array>
 #include <vector>
@@ -300,6 +301,25 @@ namespace gtsam {
       matrix_.setZero();
     }
 
+    /// Set the block columns between beginCol (inclusive) and endCol (exclusive) to zero. 
+    void setZeroColumns(DenseIndex beginCol, DenseIndex endCol) {
+      assert(beginCol < endCol);
+      assert(beginCol >= 0);
+      assert(endCol >= 0);
+      assert(beginCol < nBlocks());
+      assert(endCol <= nBlocks());
+      static_assert(Matrix::IsRowMajor == 0, "setZeroColumns requires column-major storage.");
+
+      const DenseIndex denseBeginCol = offset(beginCol);
+      const DenseIndex denseEndCol = offset(endCol);
+
+      double *begin = matrix_.data() + denseBeginCol * matrix_.rows();
+      double *end = matrix_.data() + denseEndCol * matrix_.rows();
+
+      // Using memset for maximal compiler optimization.
+      memset(begin, 0, (end - begin) * sizeof(*begin));
+    }
+
     /// Negate the entire active matrix.
     void negate();
 
diff --git a/gtsam/base/tests/testSymmetricBlockMatrix.cpp b/gtsam/base/tests/testSymmetricBlockMatrix.cpp
@@ -79,6 +79,28 @@ TEST(SymmetricBlockMatrix, WriteBlocks)
   EXPECT(assert_equal(expected3, actual3));
 }
 
+/* ************************************************************************* */
+TEST(SymmetricBlockMatrix, setZeroColumns) {
+  // Expected: columns 3 and 4 are zero
+  Matrix expected = testBlockMatrix.selfadjointView().toDenseMatrix().eval();
+  expected.col(3).setZero();
+  expected.col(4).setZero();
+  expected = expected.triangularView<Eigen::Upper>().toDenseMatrix().eval();
+
+  SymmetricBlockMatrix bm = testBlockMatrix;
+
+  // Zero out the middle block (block 1, columns 3-4)
+  bm.setZeroColumns(1, 2);
+
+  Matrix result = bm.selfadjointView()
+                      .toDenseMatrix()
+                      .triangularView<Eigen::Upper>()
+                      .toDenseMatrix()
+                      .eval();
+
+  EXPECT(assert_equal(expected, result));
+}
+
 /* ************************************************************************* */
 // Verify block range access.
 TEST(SymmetricBlockMatrix, Ranges)
diff --git a/gtsam/linear/GaussianFactor.h b/gtsam/linear/GaussianFactor.h
@@ -150,6 +150,18 @@ namespace gtsam {
     virtual void updateHessian(const KeyVector& keys,
                            SymmetricBlockMatrix* info) const = 0;
 
+    /** Update an information matrix by adding the information corresponding to this factor
+     * (used internally during elimination), restricted to a range of block columns,
+     * useful for parallelization.
+     * @param keys The ordered vector of keys for the information matrix to be updated
+     * @param info The information matrix to be updated
+     * @param beginCol First block column index (inclusive) in the range to update
+     * @param endCol Last block column index (exclusive) in the range to update
+     */
+    virtual void updateHessian(const KeyVector& keys,
+                           SymmetricBlockMatrix* info,
+                           DenseIndex beginCol, DenseIndex endCol) const = 0;
+
     /// @}
     /// @name Operator interface
     /// @{
diff --git a/gtsam/linear/HessianFactor.cpp b/gtsam/linear/HessianFactor.cpp
@@ -34,6 +34,14 @@
 #include <cassert>
 #include <limits>
 
+#ifdef GTSAM_USE_TBB
+  #include <tbb/blocked_range.h>
+  #include <tbb/parallel_for.h>
+  #include <oneapi/tbb/global_control.h>
+  #include <oneapi/tbb/task_arena.h>
+  #include <algorithm>
+#endif
+
 using namespace std;
 
 namespace gtsam {
@@ -241,17 +249,65 @@ HessianFactor::HessianFactor(const GaussianFactorGraph& factors,
     const Scatter& scatter) {
   gttic(HessianFactor_MergeConstructor);
 
+  gttic(Allocate);
   Allocate(scatter);
+  gttoc(Allocate);
+
+  // Only parallelize the inner loop if GTSAM_TBB_BOUNDED_MEMORY_GROWTH_FLAG is
+  // defined. Without this flag, multiple HessianFactor constructions are
+  // already running in parallel (e.g., when constructing multiple factors
+  // concurrently), so there's no need to parallelize the inner
+  // updateHessian loop here as well.
+#if defined(GTSAM_USE_TBB) && defined(GTSAM_TBB_BOUNDED_MEMORY_GROWTH_FLAG)
+  constexpr DenseIndex kParallelThresholdHeuristic = 50;
+  if (info_.rows() > kParallelThresholdHeuristic) {
+    gttic(updateHessian_TBB);
+
+    const DenseIndex M = info_.nBlocks();
+
+    auto numThreads = std::min(
+        static_cast<int>(oneapi::tbb::global_control::active_value(
+            oneapi::tbb::global_control::max_allowed_parallelism)),
+        static_cast<int>(oneapi::tbb::this_task_arena::max_concurrency()));
+
+    if (numThreads > 1) {
+      DenseIndex grain = std::max<DenseIndex>(1, M / (2 * numThreads));
+      tbb::parallel_for(tbb::blocked_range<DenseIndex>(0, M, grain),
+                        [&, M](const tbb::blocked_range<DenseIndex>& range) {
+                          // reverse the range to start from the end because
+                          // matrix is upper triangular and therefore end is
+                          // larger than begin so we would like to start with
+                          // the last column that is the most work and go to the
+                          // first column that is least.
+                          DenseIndex beginCol = M - range.end();
+                          DenseIndex endCol = M - range.begin();
+                          info_.setZeroColumns(beginCol, endCol);
+                          for (const auto& factor : factors) {
+                            if (factor) {
+                              factor->updateHessian(keys_, &info_, beginCol,
+                                                    endCol);
+                            }
+                          }
+                        });
+      return;
+    }
+  }
+#endif
+  gttic(setAllZero);
+  info_.setAllZero();
+  gttoc(setAllZero);
+
+  gttic(updateHessian);
 
   // Form A' * A
-  gttic(update);
-  info_.setAllZero();
-  for(const auto& factor: factors)
-    if (factor)
+  for (const auto& factor : factors) {
+    if (factor) {
       factor->updateHessian(keys_, &info_);
-  gttoc(update);
+    }
+  }
 }
 
+
 /* ************************************************************************* */
 void HessianFactor::print(const std::string& s,
     const KeyFormatter& formatter) const {
@@ -352,14 +408,58 @@ void HessianFactor::updateHessian(const KeyVector& infoKeys,
   gttic(updateHessian_HessianFactor);
   const DenseIndex nrVariablesInThisFactor = size();
 
+  gttic(slots);
   vector<DenseIndex> slots(nrVariablesInThisFactor + 1);
   for (DenseIndex j = 0; j < nrVariablesInThisFactor; ++j)
     slots[j] = Slot(infoKeys, keys_[j]);
+  
   slots[nrVariablesInThisFactor] = info->nBlocks() - 1;
+  gttoc(slots);
 
   info->updateFromMappedBlocks(info_, slots);
 }
 
+/* ************************************************************************* */
+void HessianFactor::updateHessian(const KeyVector& infoKeys,
+                                  SymmetricBlockMatrix* info,
+                                  DenseIndex beginCol,
+                                  DenseIndex endCol) const {
+  assert(info);
+  const DenseIndex nrVariablesInThisFactor = size();
+
+  vector<DenseIndex> slots;
+  slots.reserve(nrVariablesInThisFactor + 1);
+
+  for (DenseIndex j = 0; j < nrVariablesInThisFactor; ++j) {
+    slots.push_back(Slot(infoKeys, keys_[j]));
+  }
+  slots.push_back(info->nBlocks() - 1);
+
+  for (DenseIndex j = 0; j <= nrVariablesInThisFactor; ++j) {
+    const DenseIndex J = slots[j];
+    // Update diagonal block if J is in range
+    if (J >= beginCol && J < endCol) {
+      info->updateDiagonalBlock(J, info_.diagonalBlock(j));
+    }
+
+    // Update off-diagonal blocks where column max(I, J) is in range
+    // Note: We process all blocks and let the maxCol check filter them,
+    // because I and J may be in different orders (slots are not necessarily sorted)
+    for (DenseIndex i = 0; i < j; ++i) {
+      const DenseIndex I = slots[i];
+      assert(i < j);
+      assert(I != J);
+
+      // The physical column index in the symmetric matrix is max(I, J)
+      const DenseIndex maxCol = std::max(I, J);
+
+      if (maxCol >= beginCol && maxCol < endCol) {
+        info->updateOffDiagonalBlock(I, J, info_.aboveDiagonalBlock(i, j));
+      }
+    }
+  }
+}
+
 /* ************************************************************************* */
 GaussianFactor::shared_ptr HessianFactor::negate() const {
   shared_ptr result = std::make_shared<This>(*this);
diff --git a/gtsam/linear/HessianFactor.h b/gtsam/linear/HessianFactor.h
@@ -326,6 +326,17 @@ namespace gtsam {
      */
     void updateHessian(const KeyVector& keys, SymmetricBlockMatrix* info) const override;
 
+    /** Update an information matrix by adding the information corresponding to this factor
+     * (used internally during elimination), restricted to a range of block columns,
+     * useful for parallelization.
+     * @param keys The ordered vector of keys for the information matrix to be updated
+     * @param info The information matrix to be updated
+     * @param beginCol First block column index (inclusive) in the range to update
+     * @param endCol Last block column index (exclusive) in the range to update
+     */
+    void updateHessian(const KeyVector& keys, SymmetricBlockMatrix* info,
+                       DenseIndex beginCol, DenseIndex endCol) const override;
+
     /** Update another Hessian factor
      * @param other the HessianFactor to be updated
      */
diff --git a/gtsam/linear/JacobianFactor.cpp b/gtsam/linear/JacobianFactor.cpp
@@ -630,6 +630,72 @@ void JacobianFactor::updateHessian(const KeyVector& infoKeys,
   }
 }
 
+/* ************************************************************************* */
+static void whitenedUpdateHessian(SymmetricBlockMatrix* info, std::vector<DenseIndex> slots, 
+  const VerticalBlockMatrix& Ab_, DenseIndex beginCol, DenseIndex endCol) {
+  const DenseIndex n = Ab_.nBlocks() - 1;
+
+  for (DenseIndex j = 0; j <= n; ++j) {
+    const DenseIndex J = slots[j];
+    Eigen::Block<const Matrix> Ab_j = Ab_(j);
+    
+    // Update diagonal block if J is in range
+    if (J >= beginCol && J < endCol) {
+      info->diagonalBlock(J).rankUpdate(Ab_j.transpose());
+    }
+    
+    // Fill off-diagonal blocks with Ai'*Aj where column max(I, J) is in range
+    for (DenseIndex i = 0; i < j; ++i) {
+      const DenseIndex I = slots[i];
+      
+      // The physical column index in the symmetric matrix is max(I, J)
+      const DenseIndex maxCol = std::max(I, J);
+      if (maxCol >= beginCol && maxCol < endCol) {
+        // Pass original indices - updateOffDiagonalBlock handles swapping internally
+        info->updateOffDiagonalBlock(I, J, Ab_(i).transpose() * Ab_j);
+      }
+    }
+  }
+}
+
+void JacobianFactor::updateHessian(const KeyVector& infoKeys,
+                                   SymmetricBlockMatrix* info,
+                                   DenseIndex beginCol, DenseIndex endCol) const {
+  if (rows() == 0) return;
+
+  // Ab_ is the augmented Jacobian matrix A, and we perform I += A'*A below.
+  DenseIndex n = Ab_.nBlocks() - 1;
+
+  // Pre-calculate slots
+  vector<DenseIndex> slots;
+  slots.reserve(n + 1);
+  bool foundCol = false;
+  for (DenseIndex j = 0; j < n; ++j) {
+    slots.push_back(Slot(infoKeys, keys_[j]));
+    if (slots[j] >= beginCol && slots[j] < endCol) {
+      foundCol = true;
+    }
+  }
+  slots.push_back(info->nBlocks() - 1);
+  if (slots[n] >= beginCol && slots[n] < endCol) {
+    foundCol = true;
+  }
+  if (!foundCol) return;
+
+  // Whiten the factor if it has a noise model
+  const SharedDiagonal& model = get_model();
+  if (model && !model->isUnit()) {
+    if (model->isConstrained())
+      throw invalid_argument(
+          "JacobianFactor::updateHessian: cannot update information with "
+          "constrained noise model");
+    JacobianFactor whitenedFactor = whiten();
+    whitenedUpdateHessian(info, std::move(slots), whitenedFactor.Ab_, beginCol, endCol);
+  } else {
+    whitenedUpdateHessian(info, std::move(slots), Ab_, beginCol, endCol);
+  }
+}
+
 /* ************************************************************************* */
 Vector JacobianFactor::operator*(const VectorValues& x) const {
   Vector Ax(Ab_.rows());
diff --git a/gtsam/linear/JacobianFactor.h b/gtsam/linear/JacobianFactor.h
@@ -331,6 +331,17 @@ namespace gtsam {
      */
     void updateHessian(const KeyVector& keys, SymmetricBlockMatrix* info) const override;
 
+    /** Update an information matrix by adding the information corresponding to this factor
+     * (used internally during elimination), restricted to a range of block columns,
+     * useful for parallelization.
+     * @param keys The ordered vector of keys for the information matrix to be updated
+     * @param info The information matrix to be updated
+     * @param beginCol First block column index (inclusive) in the range to update
+     * @param endCol Last block column index (exclusive) in the range to update
+     */
+    void updateHessian(const KeyVector& keys, SymmetricBlockMatrix* info,
+                       DenseIndex beginCol, DenseIndex endCol) const override;
+
     /** Return A*x */
     Vector operator*(const VectorValues& x) const;
 
diff --git a/gtsam/linear/tests/testHessianFactor.cpp b/gtsam/linear/tests/testHessianFactor.cpp
diff --git a/gtsam/linear/tests/testJacobianFactor.cpp b/gtsam/linear/tests/testJacobianFactor.cpp
diff --git a/gtsam/slam/RegularImplicitSchurFactor.h b/gtsam/slam/RegularImplicitSchurFactor.h
