ITS: GPU: put trackleting properly on different streams

Signed-off-by: Felix Schlepper <[email protected]>

Author: f3sch

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

@@ -79,6 +79,7 @@ class TimeFrameGPU : public TimeFrame<nLayers>
     return mGpuStreams[stream];
   }
   auto& getStreams() { return mGpuStreams; }
+  void syncStreams();
   virtual void wipe() final;
 
   /// interface
@@ -108,7 +109,7 @@ class TimeFrameGPU : public TimeFrame<nLayers>
   std::vector<unsigned int> getClusterSizes();
   const unsigned char** getDeviceArrayUsedClusters() const { return mUsedClustersDeviceArray; }
   const int** getDeviceROframeClusters() const { return mROFrameClustersDeviceArray; }
-  Tracklet** getDeviceArrayTracklets() { return mTrackletsDeviceArray; }
+  Tracklet** getDeviceArrayTracklets() { return mTrackletsDevice.data(); }
   int** getDeviceArrayTrackletsLUT() const { return mTrackletsLUTDeviceArray; }
   int** getDeviceArrayCellsLUT() const { return mCellsLUTDeviceArray; }
   int** getDeviceArrayNeighboursCellLUT() const { return mNeighboursCellLUTDeviceArray; }
@@ -140,7 +141,8 @@ class TimeFrameGPU : public TimeFrame<nLayers>
   int getNumberOfNeighbours() const final;
 
  private:
-  void allocMemAsync(void**, size_t, Stream&, bool); // Abstract owned and unowned memory allocations
+  void allocMemAsync(void**, size_t, Stream&, bool); // Abstract owned and unowned memory allocations on specific stream
+  void allocMem(void**, size_t, bool);               // Abstract owned and unowned memory allocations on default stream
   bool mHostRegistered = false;
   TimeFrameGPUParameters mGpuParams;
 
@@ -167,7 +169,6 @@ class TimeFrameGPU : public TimeFrame<nLayers>
   const unsigned char** mUsedClustersDeviceArray;
   const int** mROFrameClustersDeviceArray;
   std::array<Tracklet*, nLayers - 1> mTrackletsDevice;
-  Tracklet** mTrackletsDeviceArray;
   std::array<int*, nLayers - 1> mTrackletsLUTDevice;
   std::array<int*, nLayers - 2> mCellsLUTDevice;
   std::array<int*, nLayers - 3> mNeighboursLUTDevice;
@@ -195,8 +196,6 @@ class TimeFrameGPU : public TimeFrame<nLayers>
 
   // State
   Streams mGpuStreams;
-  size_t mAvailMemGB;
-  bool mFirstInit = true;
 
   // Temporary buffer for storing output tracks from GPU tracking
   bounded_vector<TrackITSExt> mTrackITSExt;

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

@@ -94,16 +94,21 @@ class Stream
  public:
 #if defined(__HIPCC__)
   using Handle = hipStream_t;
-  static constexpr Handle Default = 0;
+  static constexpr Handle DefaultStream = 0;
+  // static constexpr unsigned int DefaultFlag = hipStreamNonBlocking; TODO replace once ready
+  static constexpr unsigned int DefaultFlag = 0;
 #elif defined(__CUDACC__)
   using Handle = cudaStream_t;
-  static constexpr Handle Default = 0;
+  static constexpr Handle DefaultStream = 0;
+  // static constexpr unsigned int DefaultFlag = cudaStreamNonBlocking; TODO replace once ready
+  static constexpr unsigned int DefaultFlag = 0;
 #else
   using Handle = void*;
-  static constexpr Handle Default = nullptr;
+  static constexpr Handle DefaultStream = nullptr;
+  static constexpr unsigned int DefaultFlag = 0;
 #endif
 
-  Stream(unsigned int flags = 0)
+  Stream(unsigned int flags = DefaultFlag)
   {
 #if defined(__HIPCC__)
     GPUChkErrS(hipStreamCreateWithFlags(&mHandle, flags));
@@ -115,7 +120,7 @@ class Stream
   Stream(Handle h) : mHandle(h) {}
   ~Stream()
   {
-    if (mHandle != Default) {
+    if (mHandle != DefaultStream) {
 #if defined(__HIPCC__)
       GPUChkErrS(hipStreamDestroy(mHandle));
 #elif defined(__CUDACC__)
@@ -124,7 +129,7 @@ class Stream
     }
   }
 
-  operator bool() const { return mHandle != Default; }
+  operator bool() const { return mHandle != DefaultStream; }
   const Handle& get() { return mHandle; }
   void sync() const
   {
@@ -136,7 +141,7 @@ class Stream
   }
 
  private:
-  Handle mHandle{Default};
+  Handle mHandle{DefaultStream};
 };
 static_assert(sizeof(Stream) == sizeof(void*), "Stream type must match pointer type!");
 
@@ -150,6 +155,12 @@ class Streams
   void clear() { mStreams.clear(); }
   auto& operator[](size_t i) { return mStreams[i % mStreams.size()]; }
   void push_back(const Stream& stream) { mStreams.push_back(stream); }
+  void sync()
+  {
+    for (auto& s : mStreams) {
+      s.sync();
+    }
+  }
 
  private:
   std::vector<Stream> mStreams;