ITS: track and hard limit memory allocations

Author: f3sch

Detectors/ITSMFT/ITS/tracking/CMakeLists.txt

@@ -9,6 +9,7 @@
 # granted to it by virtue of its status as an Intergovernmental Organization
 # or submit itself to any jurisdiction.
 
+#add_compile_options(-O0 -g -fPIC -fno-omit-frame-pointer)
 o2_add_library(ITStracking
                TARGETVARNAME targetName
                SOURCES src/ClusterLines.cxx
@@ -35,12 +36,8 @@ o2_add_library(ITStracking
                        O2::ITSBase
                        O2::ITSReconstruction
                        O2::ITSMFTReconstruction
-                       O2::DataFormatsITS)
-
-if (OpenMP_CXX_FOUND)
-        target_compile_definitions(${targetName} PRIVATE WITH_OPENMP)
-        target_link_libraries(${targetName} PRIVATE OpenMP::OpenMP_CXX)
-endif()
+                       O2::DataFormatsITS
+               PRIVATE_LINK_LIBRARIES TBB::tbb)
 
 o2_add_library(ITSTrackingInterface
                TARGETVARNAME targetName
@@ -50,11 +47,6 @@ o2_add_library(ITSTrackingInterface
                        O2::Framework
                        O2::GPUTracking)
 
-if (OpenMP_CXX_FOUND)
-        target_compile_definitions(${targetName} PRIVATE WITH_OPENMP)
-        target_link_libraries(${targetName} PRIVATE OpenMP::OpenMP_CXX)
-endif()
-
 o2_target_root_dictionary(ITStracking
                           HEADERS include/ITStracking/ClusterLines.h
                                   include/ITStracking/Tracklet.h

Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackerTraitsGPU.h

@@ -43,6 +43,7 @@ class TrackerTraitsGPU final : public TrackerTraits
 
   const char* getName() const noexcept final { return "GPU"; }
   bool isGPU() const noexcept final { return true; }
+  bool usesMemoryPool() const noexcept final { return false; }
 
   // TimeFrameGPU information forwarding
   int getTFNumberOfClusters() const override;
@@ -63,4 +64,4 @@ inline void TrackerTraitsGPU<nLayers>::adoptTimeFrame(TimeFrame* tf)
 } // namespace its
 } // namespace o2
 
-#endif
+#endif

Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/VertexerTraitsGPU.h

@@ -50,6 +50,7 @@ class VertexerTraitsGPU final : public VertexerTraits
 
   virtual bool isGPU() final { return true; }
   virtual const char* getName() final { return "GPU"; }
+  virtual bool usesMemoryPool() const noexcept { return false; }
 
  protected:
   IndexTableUtils* mDeviceIndexTableUtils;

Detectors/ITSMFT/ITS/tracking/include/ITStracking/BoundedAllocator.h

@@ -0,0 +1,112 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+///
+/// \file BoundedAllocator.h
+/// \brief
+///
+
+#ifndef TRACKINGITSU_INCLUDE_BOUNDEDALLOCATOR_H_
+#define TRACKINGITSU_INCLUDE_BOUNDEDALLOCATOR_H_
+
+#include <limits>
+#include <memory_resource>
+#include <atomic>
+#include <new>
+#include <vector>
+#include <string>
+#include <exception>
+
+#include "Framework/Logger.h"
+
+namespace o2::its
+{
+
+// 1. Thread-safe memory resource with global limit
+class BoundedMemoryResource final : public std::pmr::memory_resource
+{
+ public:
+  class MaxMemoryReached final : public std::bad_alloc
+  {
+   public:
+    const char* what() const noexcept final
+    {
+      return "Reached set memory limit";
+    }
+  };
+
+  BoundedMemoryResource(const std::string& name, size_t maxBytes = 0, std::pmr::memory_resource* upstream = std::pmr::get_default_resource())
+    : mMaxMemory(maxBytes), mUpstream(upstream), mName(name) {}
+
+  void* do_allocate(size_t bytes, size_t alignment) final
+  {
+    size_t new_used{0}, current_used{mUsedMemory.load(std::memory_order_relaxed)};
+    do {
+      new_used = current_used + bytes;
+      if (new_used > mMaxMemory) {
+	++mCountMaxThrow;
+	throw MaxMemoryReached();
+      }
+    } while (!mUsedMemory.compare_exchange_weak(current_used, new_used,
+						std::memory_order_acq_rel,
+						std::memory_order_relaxed));
+    ++mCountAllocate;
+    return mUpstream->allocate(bytes, alignment);
+  }
+
+  void do_deallocate(void* p, size_t bytes, size_t alignment) final
+  {
+    ++mCountDeallocate;
+    mUpstream->deallocate(p, bytes, alignment);
+    mUsedMemory.fetch_sub(bytes, std::memory_order_relaxed);
+  }
+
+  bool do_is_equal(const std::pmr::memory_resource& other) const noexcept final
+  {
+    return this == &other;
+  }
+
+  size_t getUsedMemory() const noexcept { return mUsedMemory.load(); }
+  size_t getMaxMemory() const noexcept { return mMaxMemory; }
+  void setMaxMemory(size_t max)
+  {
+    if (mUsedMemory > max) {
+      ++mCountMaxThrow;
+      throw MaxMemoryReached();
+    }
+    mMaxMemory = max;
+  }
+
+  void print() const
+  {
+    LOGP(info, "BoundedAllocator:{} alloc={} dealloc={} maxthrow={} maxmem={} used={}", mName,
+	 mCountAllocate.load(std::memory_order_relaxed),
+	 mCountDeallocate.load(std::memory_order_relaxed),
+	 mCountMaxThrow.load(std::memory_order_relaxed),
+	 mMaxMemory,
+	 mUsedMemory.load(std::memory_order_relaxed));
+  }
+
+ private:
+  size_t mMaxMemory{std::numeric_limits<size_t>::max()};
+  std::atomic<size_t> mCountMaxThrow{0};
+  std::atomic<size_t> mCountAllocate{0};
+  std::atomic<size_t> mCountDeallocate{0};
+  std::atomic<size_t> mUsedMemory{0};
+  std::pmr::memory_resource* mUpstream;
+  std::string mName;
+};
+
+template <typename T>
+using BoundedVector = std::pmr::vector<T>;
+
+} // namespace o2::its
+
+#endif

Detectors/ITSMFT/ITS/tracking/include/ITStracking/Configuration.h

@@ -90,7 +90,6 @@ struct TrackingParameters {
   float CellsPerClusterLimit = 2.f;
   /// Fitter parameters
   o2::base::PropagatorImpl<float>::MatCorrType CorrType = o2::base::PropagatorImpl<float>::MatCorrType::USEMatCorrNONE;
-  size_t MaxMemory = 12000000000UL;
   float MaxChi2ClusterAttachment = 60.f;
   float MaxChi2NDF = 30.f;
   std::vector<float> MinPt = {0.f, 0.f, 0.f, 0.f};
@@ -100,14 +99,17 @@ struct TrackingParameters {
   bool SaveTimeBenchmarks = false;
   bool DoUPCIteration = false;
   bool FataliseUponFailure = true;
-  bool DropTFUponFailure = false;
   /// Cluster attachment
   bool UseTrackFollower = false;
   bool UseTrackFollowerTop = false;
   bool UseTrackFollowerBot = false;
   bool UseTrackFollowerMix = false;
   float TrackFollowerNSigmaCutZ = 1.f;
   float TrackFollowerNSigmaCutPhi = 1.f;
+
+  bool PrintMemory = false; // print allocator usage in epilog report
+  size_t MaxMemory = 12000000000UL;
+  bool DropTFUponFailure = false;
 };
 
 struct VertexingParameters {
@@ -138,7 +140,9 @@ struct VertexingParameters {
   int zSpan = -1;
 
   int nThreads = 1;
+  bool PrintMemory = false; // print allocator usage in epilog report
   size_t MaxMemory = 12000000000UL;
+  bool DropTFUponFailure = false;
 };
 
 struct TimeFrameGPUParameters {

Detectors/ITSMFT/ITS/tracking/include/ITStracking/Constants.h

@@ -84,8 +84,8 @@ GPUhdi() constexpr std::array<float, LayersNumber> InverseZBinSize()
 {
   constexpr auto zSize = LayersZCoordinate();
   return std::array<float, LayersNumber>{0.5f * ZBins / (zSize[0]), 0.5f * ZBins / (zSize[1]), 0.5f * ZBins / (zSize[2]),
-                                                     0.5f * ZBins / (zSize[3]), 0.5f * ZBins / (zSize[4]), 0.5f * ZBins / (zSize[5]),
-                                                     0.5f * ZBins / (zSize[6])};
+                                         0.5f * ZBins / (zSize[3]), 0.5f * ZBins / (zSize[4]), 0.5f * ZBins / (zSize[5]),
+                                         0.5f * ZBins / (zSize[6])};
 }
 
 GPUhdi() constexpr float getInverseZCoordinate(const int layerIndex)

Detectors/ITSMFT/ITS/tracking/include/ITStracking/TimeFrame.h

@@ -35,6 +35,7 @@
 #include "ITStracking/Tracklet.h"
 #include "ITStracking/IndexTableUtils.h"
 #include "ITStracking/ExternalAllocator.h"
+#include "ITStracking/BoundedAllocator.h"
 
 #include "SimulationDataFormat/MCCompLabel.h"
 #include "SimulationDataFormat/MCTruthContainer.h"
@@ -59,6 +60,37 @@ class ROFRecord;
 
 namespace its
 {
+
+template <typename T>
+void deepVectorClear(std::vector<T>& vec)
+{
+  std::vector<T>().swap(vec);
+}
+
+template <typename T>
+void deepVectorClear(BoundedVector<T>& vec)
+{
+  vec = BoundedVector<T>(vec.get_allocator().resource());
+}
+
+template <typename T>
+void deepVectorClear(std::vector<BoundedVector<T>>& vec)
+{
+  for (auto& v : vec) {
+    v = BoundedVector<T>(v.get_allocator().resource());
+  }
+}
+
+template <typename T>
+void clearResizeBoundedVector(std::vector<BoundedVector<T>>& vec, size_t size, BoundedMemoryResource* bmr)
+{
+  vec.clear();
+  vec.reserve(size);
+  for (size_t i{0}; i < size; ++i) {
+    vec.emplace_back(bmr);
+  }
+}
+
 using Vertex = o2::dataformats::Vertex<o2::dataformats::TimeStamp<int>>;
 
 template <int nLayers = 7>
@@ -164,8 +196,8 @@ struct TimeFrame {
   auto& getCellsLookupTable() { return mCellsLookupTable; }
   auto& getCellsNeighbours() { return mCellsNeighbours; }
   auto& getCellsNeighboursLUT() { return mCellsNeighboursLUT; }
-  std::vector<Road<nLayers - 2>>& getRoads() { return mRoads; }
-  std::vector<TrackITSExt>& getTracks(int rofId) { return mTracks[rofId]; }
+  auto& getRoads() { return mRoads; }
+  auto& getTracks(int rofId) { return mTracks[rofId]; }
   std::vector<MCCompLabel>& getTracksLabel(const int rofId) { return mTracksLabel[rofId]; }
   std::vector<MCCompLabel>& getLinesLabel(const int rofId) { return mLinesLabels[rofId]; }
   std::vector<std::pair<MCCompLabel, float>>& getVerticesMCRecInfo() { return mVerticesMCRecInfo; }
@@ -179,8 +211,12 @@ struct TimeFrame {
   auto getNumberOfExtendedTracks() const { return mNExtendedTracks; }
   auto getNumberOfUsedExtendedClusters() const { return mNExtendedUsedClusters; }
 
+  void setMemoryPool(std::shared_ptr<BoundedMemoryResource>& pool) { mMemoryPool = pool; }
+  auto& getMemoryPool() const noexcept { return mMemoryPool; }
+  bool usesMemoryPool() const noexcept { return true; }
   bool checkMemory(unsigned long max) { return getArtefactsMemory() < max; }
-  unsigned long getArtefactsMemory();
+  unsigned long getArtefactsMemory() const;
+  void printArtefactsMemory() const;
   int getROFCutClusterMult() const { return mCutClusterMult; };
   int getROFCutVertexMult() const { return mCutVertexMult; };
   int getROFCutAllMult() const { return mCutClusterMult + mCutVertexMult; }
@@ -189,7 +225,7 @@ struct TimeFrame {
   void computeTrackletsPerROFScans();
   void computeTracletsPerClusterScans();
   int& getNTrackletsROF(int rofId, int combId) { return mNTrackletsPerROF[combId][rofId]; }
-  std::vector<Line>& getLines(int rofId) { return mLines[rofId]; }
+  auto& getLines(int rofId) { return mLines[rofId]; }
   int getNLinesTotal() const
   {
     return std::accumulate(mLines.begin(), mLines.end(), 0, [](int sum, const auto& l) { return sum + l.size(); });
@@ -234,6 +270,7 @@ struct TimeFrame {
   void addClusterExternalIndexToLayer(int layer, const int idx) { mClusterExternalIndices[layer].push_back(idx); }
 
   void resetVectors();
+  void resetTracklets();
 
   /// Debug and printing
   void checkTrackletLUTs();
@@ -279,29 +316,32 @@ struct TimeFrame {
     }
   }
 
-  std::array<std::vector<Cluster>, nLayers> mUnsortedClusters;
-  std::vector<std::vector<Tracklet>> mTracklets;
-  std::vector<std::vector<CellSeed>> mCells;
-  std::vector<std::vector<o2::track::TrackParCovF>> mCellSeeds;
-  std::vector<std::vector<float>> mCellSeedsChi2;
-  std::vector<Road<nLayers - 2>> mRoads;
-  std::vector<std::vector<TrackITSExt>> mTracks;
-  std::vector<std::vector<int>> mCellsNeighbours;
-  std::vector<std::vector<int>> mCellsLookupTable;
+  std::array<BoundedVector<Cluster>, nLayers> mUnsortedClusters;
+  std::vector<BoundedVector<Tracklet>> mTracklets;
+  std::vector<BoundedVector<CellSeed>> mCells;
+  std::vector<BoundedVector<o2::track::TrackParCovF>> mCellSeeds;
+  std::vector<BoundedVector<float>> mCellSeedsChi2;
+  BoundedVector<Road<nLayers - 2>> mRoads;
+  std::vector<BoundedVector<TrackITSExt>> mTracks;
+  std::vector<BoundedVector<int>> mCellsNeighbours;
+  std::vector<BoundedVector<int>> mCellsLookupTable;
   std::vector<uint8_t> mMultiplicityCutMask;
 
   const o2::base::PropagatorImpl<float>* mPropagatorDevice = nullptr; // Needed only for GPU
-  void dropTracks()
+
+  void wipe()
   {
-    for (auto& v : mTracks) {
+    for (auto& v : mUnsortedClusters) {
       deepVectorClear(v);
     }
-  }
-
-  template <typename T>
-  void deepVectorClear(std::vector<T>& vec)
-  {
-    std::vector<T>().swap(vec);
+    deepVectorClear(mTracks);
+    deepVectorClear(mTracklets);
+    deepVectorClear(mCells);
+    deepVectorClear(mCellSeeds);
+    deepVectorClear(mCellSeedsChi2);
+    deepVectorClear(mRoads);
+    deepVectorClear(mCellsNeighbours);
+    deepVectorClear(mCellsLookupTable);
   }
 
   virtual void prepareClusters(const TrackingParameters& trkParam, const int maxLayers);
@@ -341,11 +381,16 @@ struct TimeFrame {
   unsigned int mNoVertexROF = 0;
   std::vector<int> mTotVertPerIteration;
   // \Vertexer
+
+  std::shared_ptr<BoundedMemoryResource> mMemoryPool;
 };
 
 template <int nLayers>
 inline gsl::span<const Vertex> TimeFrame<nLayers>::getPrimaryVertices(int rofId) const
 {
+  if (mPrimaryVertices.empty()) {
+    return {};
+  }
   const int start = mROFramesPV[rofId];
   const int stop_idx = rofId >= mNrof - 1 ? mNrof : rofId + 1;
   int delta = mMultiplicityCutMask[rofId] ? mROFramesPV[stop_idx] - start : 0; // return empty span if Rof is excluded
@@ -364,6 +409,9 @@ inline gsl::span<const std::pair<MCCompLabel, float>> TimeFrame<nLayers>::getPri
 template <int nLayers>
 inline gsl::span<const Vertex> TimeFrame<nLayers>::getPrimaryVertices(int romin, int romax) const
 {
+  if (mPrimaryVertices.empty()) {
+    return {};
+  }
   return {&mPrimaryVertices[mROFramesPV[romin]], static_cast<gsl::span<const Vertex>::size_type>(mROFramesPV[romax + 1] - mROFramesPV[romin])};
 }