refactor: Split GainMatrixUpdater compilation (#3486)

Once again, our compilation memory usage is out of control, among others due to `GainMatrixUpdater`. This currently takes about 1.2G peak memory. In this PR, I'm moving the function doing the heavy instantiation back into a header, but I'm suppressing the instantiation of the template via `extern template`. Then, I have CMake generate `.cpp` files where each file instantiates one dimension. The linker ultimately resolves these separate template insantiations. Now, this is a bit ugly, but it does reduce **peak** memory consumption to under 600M, at the cost of now having 6 instead of one compilation unit: ``` [ 16.73M, max: 295.39M] [ 1.63s] - Core/src/TrackFitting/GainMatrixUpdater.cpp [ 16.58M, max: 431.10M] [ 2.95s] - build/Core/src/TrackFitting/GainMatrixUpdaterImpl1.cpp [ 16.58M, max: 456.64M] [ 3.13s] - build/Core/src/TrackFitting/GainMatrixUpdaterImpl2.cpp [ 17.06M, max: 485.18M] [ 3.40s] - build/Core/src/TrackFitting/GainMatrixUpdaterImpl3.cpp [ 16.86M, max: 456.70M] [ 3.24s] - build/Core/src/TrackFitting/GainMatrixUpdaterImpl4.cpp [ 16.58M, max: 568.34M] [ 4.37s] - build/Core/src/TrackFitting/GainMatrixUpdaterImpl5.cpp [ 16.73M, max: 537.15M] [ 4.11s] - build/Core/src/TrackFitting/GainMatrixUpdaterImpl6.cpp ``` Not sure this is the way to go, but I wanted to suggest it.
acts-project · Aug 7, 2024 · ea794f2 · ea794f2
1 parent 95da6aa
commit ea794f2
Show file tree

Hide file tree

Showing 6 changed files with 145 additions and 60 deletions.
diff --git a/Core/include/Acts/EventData/MeasurementHelpers.hpp b/Core/include/Acts/EventData/MeasurementHelpers.hpp
@@ -63,10 +63,12 @@ auto visit_measurement(A&& param, B&& cov, std::size_t dim, L&& lambda) {
 /// @tparam L The generic lambda type to call
 /// @param dim The runtime dimension of the measurement
 /// @param lambda The generic lambda instance to call
+/// @param args Additional arguments passed to @p lambda
 /// @return Returns the lambda return value
-template <typename L>
-auto visit_measurement(std::size_t dim, L&& lambda) {
-  return template_switch_lambda<1, eBoundSize>(dim, lambda);
+template <typename L, typename... Args>
+auto visit_measurement(std::size_t dim, L&& lambda, Args&&... args) {
+  return template_switch_lambda<1, eBoundSize>(dim, lambda,
+                                               std::forward<Args>(args)...);
 }
 
 }  // namespace Acts
diff --git a/Core/include/Acts/TrackFitting/GainMatrixUpdater.hpp b/Core/include/Acts/TrackFitting/GainMatrixUpdater.hpp
@@ -96,6 +96,10 @@ class GainMatrixUpdater {
  private:
   std::tuple<double, std::error_code> visitMeasurement(
       InternalTrackState trackState, const Logger& logger) const;
+
+  template <std::size_t N>
+  std::tuple<double, std::error_code> visitMeasurementImpl(
+      InternalTrackState trackState, const Logger& logger) const;
 };
 
 }  // namespace Acts
diff --git a/Core/include/Acts/TrackFitting/detail/GainMatrixUpdaterImpl.hpp b/Core/include/Acts/TrackFitting/detail/GainMatrixUpdaterImpl.hpp
@@ -0,0 +1,93 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2024 CERN for the benefit of the Acts project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#pragma once
+
+#include "Acts/TrackFitting/GainMatrixUpdater.hpp"
+#include "Acts/Utilities/Logger.hpp"
+
+#include <cstddef>
+#include <tuple>
+
+namespace Acts {
+
+template <std::size_t N>
+std::tuple<double, std::error_code> GainMatrixUpdater::visitMeasurementImpl(
+    InternalTrackState trackState, const Logger& logger) const {
+  double chi2 = 0;
+
+  constexpr std::size_t kMeasurementSize = N;
+  using ParametersVector = ActsVector<kMeasurementSize>;
+  using CovarianceMatrix = ActsSquareMatrix<kMeasurementSize>;
+
+  typename TrackStateTraits<kMeasurementSize, true>::Calibrated calibrated{
+      trackState.calibrated};
+  typename TrackStateTraits<kMeasurementSize, true>::CalibratedCovariance
+      calibratedCovariance{trackState.calibratedCovariance};
+
+  ACTS_VERBOSE("Measurement dimension: " << kMeasurementSize);
+  ACTS_VERBOSE("Calibrated measurement: " << calibrated.transpose());
+  ACTS_VERBOSE("Calibrated measurement covariance:\n" << calibratedCovariance);
+
+  const auto H = trackState.projector
+                     .template topLeftCorner<kMeasurementSize, eBoundSize>()
+                     .eval();
+
+  ACTS_VERBOSE("Measurement projector H:\n" << H);
+
+  const auto K = (trackState.predictedCovariance * H.transpose() *
+                  (H * trackState.predictedCovariance * H.transpose() +
+                   calibratedCovariance)
+                      .inverse())
+                     .eval();
+
+  ACTS_VERBOSE("Gain Matrix K:\n" << K);
+
+  if (K.hasNaN()) {
+    // set to error abort execution
+    return {0, KalmanFitterError::UpdateFailed};
+  }
+
+  trackState.filtered =
+      trackState.predicted + K * (calibrated - H * trackState.predicted);
+  trackState.filteredCovariance =
+      (BoundSquareMatrix::Identity() - K * H) * trackState.predictedCovariance;
+  ACTS_VERBOSE("Filtered parameters: " << trackState.filtered.transpose());
+  ACTS_VERBOSE("Filtered covariance:\n" << trackState.filteredCovariance);
+
+  ParametersVector residual;
+  residual = calibrated - H * trackState.filtered;
+  ACTS_VERBOSE("Residual: " << residual.transpose());
+
+  CovarianceMatrix m =
+      ((CovarianceMatrix::Identity() - H * K) * calibratedCovariance).eval();
+
+  chi2 = (residual.transpose() * m.inverse() * residual).value();
+
+  ACTS_VERBOSE("Chi2: " << chi2);
+
+  return {chi2, {}};
+}
+
+// Ensure thet the compiler does not implicitly instantiate the template
+
+#define _EXTERN(N)                                                          \
+  extern template std::tuple<double, std::error_code>                       \
+  GainMatrixUpdater::visitMeasurementImpl<N>(InternalTrackState trackState, \
+                                             const Logger& logger) const
+
+_EXTERN(1);
+_EXTERN(2);
+_EXTERN(3);
+_EXTERN(4);
+_EXTERN(5);
+_EXTERN(6);
+
+#undef _EXTERN
+
+}  // namespace Acts
diff --git a/Core/src/TrackFitting/CMakeLists.txt b/Core/src/TrackFitting/CMakeLists.txt
@@ -12,3 +12,21 @@ target_sources(
     GlobalChiSquareFitter.cpp
     MbfSmoother.cpp
 )
+
+foreach(DIM RANGE 1 6)
+
+  set(dim_file ${CMAKE_CURRENT_BINARY_DIR}/GainMatrixUpdaterImpl${DIM}.cpp)
+  configure_file(
+    ${CMAKE_CURRENT_SOURCE_DIR}/GainMatrixUpdaterImpl.cpp.in
+    ${dim_file}
+    @ONLY
+  )
+
+  set_source_files_properties(
+  ${dim_file}
+  PROPERTIES COMPILE_DEFINITIONS ACTS_GAIN_MATRIX_UPDATER_INSTANTIATE=${DIM})
+  target_sources(ActsCore
+  PRIVATE
+  ${dim_file})
+
+endforeach()
diff --git a/Core/src/TrackFitting/GainMatrixUpdater.cpp b/Core/src/TrackFitting/GainMatrixUpdater.cpp
@@ -15,6 +15,7 @@
 
 #include <algorithm>
 #include <cstddef>
+#include <type_traits>
 #include <utility>
 
 #include <Eigen/src/Core/MatrixBase.h>
@@ -24,64 +25,13 @@ namespace Acts {
 std::tuple<double, std::error_code> GainMatrixUpdater::visitMeasurement(
     InternalTrackState trackState, const Logger& logger) const {
   // default-constructed error represents success, i.e. an invalid error code
-  std::error_code error;
-  double chi2 = 0;
 
-  visit_measurement(trackState.calibratedSize, [&](auto N) -> void {
-    constexpr std::size_t kMeasurementSize = decltype(N)::value;
-    using ParametersVector = ActsVector<kMeasurementSize>;
-    using CovarianceMatrix = ActsSquareMatrix<kMeasurementSize>;
-
-    typename TrackStateTraits<kMeasurementSize, true>::Calibrated calibrated{
-        trackState.calibrated};
-    typename TrackStateTraits<kMeasurementSize, true>::CalibratedCovariance
-        calibratedCovariance{trackState.calibratedCovariance};
-
-    ACTS_VERBOSE("Measurement dimension: " << kMeasurementSize);
-    ACTS_VERBOSE("Calibrated measurement: " << calibrated.transpose());
-    ACTS_VERBOSE("Calibrated measurement covariance:\n"
-                 << calibratedCovariance);
-
-    const auto H = trackState.projector
-                       .template topLeftCorner<kMeasurementSize, eBoundSize>()
-                       .eval();
-
-    ACTS_VERBOSE("Measurement projector H:\n" << H);
-
-    const auto K = (trackState.predictedCovariance * H.transpose() *
-                    (H * trackState.predictedCovariance * H.transpose() +
-                     calibratedCovariance)
-                        .inverse())
-                       .eval();
-
-    ACTS_VERBOSE("Gain Matrix K:\n" << K);
-
-    if (K.hasNaN()) {
-      // set to error abort execution
-      error = KalmanFitterError::UpdateFailed;
-      return;
-    }
-
-    trackState.filtered =
-        trackState.predicted + K * (calibrated - H * trackState.predicted);
-    trackState.filteredCovariance = (BoundSquareMatrix::Identity() - K * H) *
-                                    trackState.predictedCovariance;
-    ACTS_VERBOSE("Filtered parameters: " << trackState.filtered.transpose());
-    ACTS_VERBOSE("Filtered covariance:\n" << trackState.filteredCovariance);
-
-    ParametersVector residual;
-    residual = calibrated - H * trackState.filtered;
-    ACTS_VERBOSE("Residual: " << residual.transpose());
-
-    CovarianceMatrix m =
-        ((CovarianceMatrix::Identity() - H * K) * calibratedCovariance).eval();
-
-    chi2 = (residual.transpose() * m.inverse() * residual).value();
-
-    ACTS_VERBOSE("Chi2: " << chi2);
-  });
-
-  return {chi2, error};
+  return visit_measurement(
+      trackState.calibratedSize,
+      [&, this]<std::size_t N>(std::integral_constant<std::size_t, N>)
+          -> std::tuple<double, std::error_code> {
+        return visitMeasurementImpl<N>(trackState, logger);
+      });
 }
 
 }  // namespace Acts
diff --git a/Core/src/TrackFitting/GainMatrixUpdaterImpl.cpp.in b/Core/src/TrackFitting/GainMatrixUpdaterImpl.cpp.in
@@ -0,0 +1,18 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2024 CERN for the benefit of the Acts project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#include "Acts/TrackFitting/detail/GainMatrixUpdaterImpl.hpp"
+
+// clang-format off
+namespace Acts {
+
+template std::tuple<double, std::error_code>
+GainMatrixUpdater::visitMeasurementImpl<@DIM@>(InternalTrackState trackState,
+                                                const Logger& logger) const;
+
+}  // namespace Acts