Share GPU framework allocation with ITSMFT

Initialize GPU memory pool sizes and reconstruction in the legacy ITS workflow. Move ExternalAllocator and its PMR adaptor into common ITSMFT tracking, retain legacy aliases, and pass allocation flags directly to one shared framework adapter. Expose that adapter through GPUChainITS and GPUO2Interface.

Author: mpuccio

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

@@ -26,6 +26,11 @@
 #include "ITStracking/LineProjection.h"
 #include "DetectorsBase/Propagator.h"
 
+namespace o2::itsmft::tracking
+{
+class ExternalAllocator;
+}
+
 namespace o2::its
 {
 using o2::itsmft::tracking::bounded_vector;
@@ -41,7 +46,7 @@ template <int>
 class IndexTableUtils;
 class Cluster;
 class TrackITSExt;
-class ExternalAllocator;
+using ExternalAllocator = o2::itsmft::tracking::ExternalAllocator;
 
 template <int NLayers>
 struct TrackingKernels {

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

@@ -16,72 +16,13 @@
 #ifndef TRACKINGITSU_INCLUDE_EXTERNALALLOCATOR_H_
 #define TRACKINGITSU_INCLUDE_EXTERNALALLOCATOR_H_
 
-#include <memory_resource>
 #include "GPUO2ExternalUser.h"
-#include "Base/GPUMemoryResource.h"
+#include "ITSMFTTracking/ExternalAllocator.h"
 
 namespace o2::its
 {
-
-class ExternalAllocator
-{
-  using Type = std::underlying_type_t<o2::gpu::GPUMemoryResource::MemoryType>;
-
- public:
-  virtual void deallocate(char*, size_t) = 0;
-  virtual void* allocate(size_t) = 0;
-  void* allocate(size_t s, Type type)
-  {
-    auto old = mType;
-    mType = type;
-    void* p = allocate(s);
-    mType = old;
-    return p;
-  }
-
-  void* allocateStack(size_t s)
-  {
-    return allocate(s, (o2::gpu::GPUMemoryResource::MEMORY_GPU | o2::gpu::GPUMemoryResource::MEMORY_STACK));
-  }
-  virtual void pushTagOnStack(uint64_t) = 0;
-  virtual void popTagOffStack(uint64_t) = 0;
-
-  void setType(Type t) noexcept { mType = t; }
-  Type getType() const noexcept { return mType; }
-
- protected:
-  Type mType;
-};
-
-class ExternalAllocatorAdaptor final : public std::pmr::memory_resource
-{
- public:
-  explicit ExternalAllocatorAdaptor(ExternalAllocator* alloc) : mAlloc(alloc) {}
-
- protected:
-  void* do_allocate(size_t bytes, size_t alignment) override
-  {
-    void* p = mAlloc->allocate(bytes, o2::gpu::GPUMemoryResource::MemoryType::MEMORY_HOST);
-    if (!p) {
-      throw std::bad_alloc();
-    }
-    return p;
-  }
-
-  void do_deallocate(void* p, size_t bytes, size_t) override
-  {
-    mAlloc->deallocate(static_cast<char*>(p), bytes);
-  }
-
-  bool do_is_equal(const std::pmr::memory_resource& other) const noexcept override
-  {
-    return this == &other;
-  }
-
- private:
-  ExternalAllocator* mAlloc;
-};
-
+using ExternalAllocator = o2::itsmft::tracking::ExternalAllocator;
+using ExternalAllocatorAdaptor = o2::itsmft::tracking::ExternalAllocatorAdaptor;
 } // namespace o2::its
 
 #endif

Detectors/ITSMFT/ITS/workflow/src/TrackerSpec.cxx

@@ -45,6 +45,9 @@ void TrackerDPL::init(InitContext& ic)
   mTimer.Reset();
   o2::base::GRPGeomHelper::instance().setRequest(mGGCCDBRequest);
   mChainITS.reset(mRecChain->AddChain<o2::gpu::GPUChainITS>());
+  if (mRecChain->Init()) {
+    throw std::runtime_error("GPU reconstruction initialization failed");
+  }
   mITSTrackingInterface.setTraitsFromProvider(mChainITS->GetITSVertexerTraits(),
                                               mChainITS->GetITSTrackerTraits(),
                                               mChainITS->GetITSTimeframe());

Detectors/ITSMFT/common/tracking/include/ITSMFTTracking/ExternalAllocator.h

@@ -0,0 +1,84 @@
+// 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 ExternalAllocator.h
+/// \brief
+///
+
+#ifndef ALICEO2_ITSMFT_TRACKING_EXTERNALALLOCATOR_H_
+#define ALICEO2_ITSMFT_TRACKING_EXTERNALALLOCATOR_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <memory_resource>
+#include <type_traits>
+#include "Base/GPUMemoryResource.h"
+
+namespace o2::itsmft::tracking
+{
+
+class ExternalAllocator
+{
+ public:
+  using Type = std::underlying_type_t<o2::gpu::GPUMemoryResource::MemoryType>;
+  virtual ~ExternalAllocator() = default;
+  virtual void deallocate(char*, size_t) = 0;
+  virtual void* allocate(size_t, Type) = 0;
+  void* allocate(size_t size) { return allocate(size, mType); }
+
+  void* allocateDevice(size_t size, bool stack)
+  {
+    return allocate(size, o2::gpu::GPUMemoryResource::MEMORY_GPU | (stack ? o2::gpu::GPUMemoryResource::MEMORY_STACK : 0));
+  }
+
+  void* allocateStack(size_t size) { return allocateDevice(size, true); }
+  virtual void pushTagOnStack(uint64_t) = 0;
+  virtual void popTagOffStack(uint64_t) = 0;
+
+  void setType(Type t) noexcept { mType = t; }
+  Type getType() const noexcept { return mType; }
+
+ protected:
+  Type mType{o2::gpu::GPUMemoryResource::MEMORY_HOST};
+};
+
+class ExternalAllocatorAdaptor final : public std::pmr::memory_resource
+{
+ public:
+  explicit ExternalAllocatorAdaptor(ExternalAllocator* alloc) : mAlloc(alloc) {}
+
+ protected:
+  void* do_allocate(size_t bytes, size_t alignment) override
+  {
+    void* p = mAlloc->allocate(bytes, o2::gpu::GPUMemoryResource::MemoryType::MEMORY_HOST);
+    if (!p) {
+      throw std::bad_alloc();
+    }
+    return p;
+  }
+
+  void do_deallocate(void* p, size_t bytes, size_t) override
+  {
+    mAlloc->deallocate(static_cast<char*>(p), bytes);
+  }
+
+  bool do_is_equal(const std::pmr::memory_resource& other) const noexcept override
+  {
+    return this == &other;
+  }
+
+ private:
+  ExternalAllocator* mAlloc;
+};
+
+} // namespace o2::itsmft::tracking
+
+#endif

Detectors/Upgrades/ALICE3/GlobalReconstruction/reconstruction/include/ALICE3GlobalReconstruction/GPUExternalAllocator.h

@@ -30,7 +30,7 @@ class GPUExternalAllocator final : public o2::its::ExternalAllocator
   GPUExternalAllocator() = default;
   ~GPUExternalAllocator();
 
-  void* allocate(size_t size) override;
+  void* allocate(size_t size, Type type) override;
   void deallocate(char* ptr, size_t size) override;
   void pushTagOnStack(uint64_t tag) override;
   void popTagOffStack(uint64_t tag) override;

Detectors/Upgrades/ALICE3/GlobalReconstruction/reconstruction/src/GPUExternalAllocator.cu

@@ -37,9 +37,8 @@ GPUExternalAllocator::~GPUExternalAllocator()
   releaseAll();
 }
 
-void* GPUExternalAllocator::allocate(size_t size)
+void* GPUExternalAllocator::allocate(size_t size, Type type)
 {
-  const auto type = static_cast<MemoryType>(getType());
   const bool useHost = (type & static_cast<MemoryType>(o2::gpu::GPUMemoryResource::MEMORY_HOST)) != 0;
   const bool useStack = (type & static_cast<MemoryType>(o2::gpu::GPUMemoryResource::MEMORY_STACK)) != 0;
 

GPU/GPUTracking/CMakeLists.txt

@@ -379,6 +379,7 @@ if(ALIGPU_BUILD_TYPE STREQUAL "O2")
                                        O2::TOFBase
                                        O2::ITSBase
                                        O2::ITStracking
+                                       O2::ITSMFTTracking
                                        O2::TPCFastTransformation
                                        O2::DetectorsRaw
                                        O2::Steer

GPU/GPUTracking/Global/GPUChainITS.cxx

@@ -14,43 +14,44 @@
 
 #include "GPUChainITS.h"
 #include "GPUConstantMem.h"
+#include "GPUDefParametersConstants.h"
 #include "DataFormatsITS/TrackITS.h"
-#include "ITStracking/ExternalAllocator.h"
+#include "ITSMFTTracking/ExternalAllocator.h"
 #include "GPUReconstructionIncludesITS.h"
 
 using namespace o2::gpu;
 
-namespace o2::its
+namespace
 {
-class GPUFrameworkExternalAllocator final : public o2::its::ExternalAllocator
+class GPUFrameworkExternalAllocator final : public o2::itsmft::tracking::ExternalAllocator
 {
  public:
-  void* allocate(size_t size) final
+  explicit GPUFrameworkExternalAllocator(GPUReconstruction* fwr) : mFWReco(fwr) {}
+  void* allocate(size_t size, Type type) final
   {
-    return mFWReco->AllocateDirectMemory(size, mType);
+    return mFWReco->AllocateDirectMemory(size, type);
   }
-  void deallocate(char* ptr, size_t size) final {} // this is a simple no-op
-  void pushTagOnStack(uint64_t tag) final
-  {
-    mFWReco->PushNonPersistentMemory(tag);
-  }
-  void popTagOffStack(uint64_t tag) final
-  {
-    mFWReco->PopNonPersistentMemory(gpudatatypes::RecoStep::ITSTracking, tag);
-  }
-  void setReconstructionFramework(o2::gpu::GPUReconstruction* fwr) { mFWReco = fwr; }
+  void deallocate(char*, size_t) final {}
+  void pushTagOnStack(uint64_t tag) final { mFWReco->PushNonPersistentMemory(tag); }
+  void popTagOffStack(uint64_t tag) final { mFWReco->PopNonPersistentMemory(gpudatatypes::RecoStep::ITSTracking, tag); }
 
  private:
-  o2::gpu::GPUReconstruction* mFWReco;
+  GPUReconstruction* mFWReco;
 };
-} // namespace o2::its
+} // namespace
 
 GPUChainITS::~GPUChainITS() = default;
 
 GPUChainITS::GPUChainITS(GPUReconstruction* rec) : GPUChain(rec) {}
 
 int32_t GPUChainITS::Init() { return 0; }
 
+void GPUChainITS::MemorySize(size_t& gpuMem, size_t& pageLockedHostMem)
+{
+  gpuMem = constants::GPU_DEFAULT_MEMORY_SIZE;
+  pageLockedHostMem = constants::GPU_DEFAULT_HOST_MEMORY_SIZE;
+}
+
 o2::its::TrackerTraits<7>* GPUChainITS::GetITSTrackerTraits()
 {
   if (mITSTrackerTraits == nullptr) {
@@ -74,14 +75,20 @@ o2::its::TimeFrame<7>* GPUChainITS::GetITSTimeframe()
   }
 #if !defined(GPUCA_STANDALONE)
   if (mITSTimeFrame->isGPU()) {
-    mFrameworkAllocator.reset(new o2::its::GPUFrameworkExternalAllocator());
-    mFrameworkAllocator->setReconstructionFramework(rec());
-    mITSTimeFrame->setFrameworkAllocator(mFrameworkAllocator.get());
+    mITSTimeFrame->setFrameworkAllocator(GetITSMFTFrameworkAllocator());
   }
 #endif
   return mITSTimeFrame.get();
 }
 
+o2::itsmft::tracking::ExternalAllocator* GPUChainITS::GetITSMFTFrameworkAllocator()
+{
+  if (mFrameworkAllocator == nullptr && mRec->IsGPU()) {
+    mFrameworkAllocator = std::make_unique<GPUFrameworkExternalAllocator>(rec());
+  }
+  return mFrameworkAllocator.get();
+}
+
 int32_t GPUChainITS::PrepareEvent() { return 0; }
 
 int32_t GPUChainITS::Finalize() { return 0; }

GPU/GPUTracking/Global/GPUChainITS.h

@@ -21,8 +21,11 @@ namespace o2::its
 struct Cluster;
 struct TrackingFrameInfo;
 class TrackITSExt;
-class GPUFrameworkExternalAllocator;
 } // namespace o2::its
+namespace o2::itsmft::tracking
+{
+class ExternalAllocator;
+} // namespace o2::itsmft::tracking
 
 namespace o2::gpu
 {
@@ -39,15 +42,16 @@ class GPUChainITS final : public GPUChain
 
   void RegisterPermanentMemoryAndProcessors() final {};
   void RegisterGPUProcessors() final {};
-  void MemorySize(size_t&, size_t&) final {};
+  void MemorySize(size_t& gpuMem, size_t& pageLockedHostMem) final;
 
   o2::its::TrackerTraits<7>* GetITSTrackerTraits();
   o2::its::VertexerTraits<7>* GetITSVertexerTraits();
   o2::its::TimeFrame<7>* GetITSTimeframe();
+  o2::itsmft::tracking::ExternalAllocator* GetITSMFTFrameworkAllocator();
 
  protected:
   GPUChainITS(GPUReconstruction* rec);
-  std::unique_ptr<o2::its::GPUFrameworkExternalAllocator> mFrameworkAllocator;
+  std::unique_ptr<o2::itsmft::tracking::ExternalAllocator> mFrameworkAllocator;
   std::unique_ptr<o2::its::TimeFrame<7>> mITSTimeFrame;
   std::unique_ptr<o2::its::TrackerTraits<7>> mITSTrackerTraits;
   std::unique_ptr<o2::its::VertexerTraits<7>> mITSVertexerTraits;

GPU/GPUTracking/Interface/GPUO2Interface.cxx

@@ -261,6 +261,11 @@ void GPUO2Interface::setErrorCodeOutput(std::vector<std::array<uint32_t, 4>>* v)
   }
 }
 
+o2::itsmft::tracking::ExternalAllocator* GPUO2Interface::GetITSMFTFrameworkAllocator()
+{
+  return mChainITS->GetITSMFTFrameworkAllocator();
+}
+
 void GPUO2Interface::GetITSTraits(o2::its::TrackerTraits<7>*& trackerTraits, o2::its::VertexerTraits<7>*& vertexerTraits, o2::its::TimeFrame<7>*& timeFrame)
 {
   trackerTraits = mChainITS->GetITSTrackerTraits();