ITS: GPU: add stream abstraction + use for trackleting

Author: f3sch

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

@@ -22,8 +22,6 @@
 namespace o2::its::gpu
 {
 
-class Stream;
-
 class DefaultGPUAllocator : public ExternalAllocator
 {
   void* allocate(size_t size) override;
@@ -80,10 +78,11 @@ class TimeFrameGPU : public TimeFrame
   void downloadCellsLUTDevice();
   void unregisterRest();
   template <Task task>
-  Stream& getStream(const size_t stream)
+  auto getStream(const size_t stream)
   {
-    return *mGpuStreams[stream];
+    return mGpuStreams[stream];
   }
+  auto& getStreams() { return mGpuStreams; }
   void wipe(const int);
 
   /// interface
@@ -141,7 +140,8 @@ class TimeFrameGPU : public TimeFrame
   int getNumberOfCells() const;
 
  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
+  void allocMemAsync(void**, size_t, Stream&, bool, int); // Abstract owned and unowned memory allocations and set
   bool mHostRegistered = false;
   TimeFrameGPUParameters mGpuParams;
 
@@ -194,7 +194,7 @@ class TimeFrameGPU : public TimeFrame
   const TrackingFrameInfo** mTrackingFrameInfoDeviceArray;
 
   // State
-  std::vector<Stream*> mGpuStreams;
+  Streams mGpuStreams;
   size_t mAvailMemGB;
   bool mFirstInit = true;
 

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

@@ -79,7 +79,8 @@ void countTrackletsInROFsHandler(const IndexTableUtils* utils,
                                  std::vector<float>& radii,
                                  std::vector<float>& mulScatAng,
                                  const int nBlocks,
-                                 const int nThreads);
+                                 const int nThreads,
+                                 gpu::Streams& streams);
 
 template <int nLayers = 7>
 void computeTrackletsInROFsHandler(const IndexTableUtils* utils,
@@ -112,7 +113,8 @@ void computeTrackletsInROFsHandler(const IndexTableUtils* utils,
                                    std::vector<float>& radii,
                                    std::vector<float>& mulScatAng,
                                    const int nBlocks,
-                                   const int nThreads);
+                                   const int nThreads,
+                                   gpu::Streams& streams);
 
 void countCellsHandler(const Cluster** sortedClusters,
                        const Cluster** unsortedClusters,

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

@@ -16,7 +16,10 @@
 #ifndef ITSTRACKINGGPU_UTILS_H_
 #define ITSTRACKINGGPU_UTILS_H_
 
+#include <vector>
+
 #include "GPUCommonDef.h"
+#include "GPUCommonHelpers.h"
 
 namespace o2
 {
@@ -31,11 +34,6 @@ struct gpuPair {
 namespace gpu
 {
 
-template <typename T>
-void discardResult(const T&)
-{
-}
-
 // Poor man implementation of a span-like struct. It is very limited.
 template <typename T>
 struct gpuSpan {
@@ -96,8 +94,76 @@ GPUhd() const T* getPtrFromRuler(int index, const T* src, const int* ruler, cons
 {
   return src + ruler[index] * stride;
 }
+
+// Abstract stream class
+class Stream
+{
+ public:
+#if defined(__HIPCC__)
+  using Handle = hipStream_t;
+  static constexpr Handle Default = 0;
+#elif defined(__CUDACC__)
+  using Handle = cudaStream_t;
+  static constexpr Handle Default = 0;
+#else
+  using Handle = void*;
+  static constexpr Handle Default = nullptr;
+#endif
+
+  Stream(unsigned int flags = 0)
+  {
+#if defined(__HIPCC__)
+    GPUChkErrS(hipStreamCreateWithFlags(&mHandle, flags));
+#elif defined(__CUDACC__)
+    GPUChkErrS(cudaStreamCreateWithFlags(&mHandle, flags));
+#endif
+  }
+
+  Stream(Handle h) : mHandle(h) {}
+  ~Stream()
+  {
+    if (mHandle != Default) {
+#if defined(__HIPCC__)
+      GPUChkErrS(hipStreamDestroy(mHandle));
+#elif defined(__CUDACC__)
+      GPUChkErrS(cudaStreamDestroy(mHandle));
+#endif
+    }
+  }
+
+  operator bool() const { return mHandle != Default; }
+  const Handle& get() { return mHandle; }
+  void sync() const
+  {
+#if defined(__HIPCC__)
+    GPUChkErrS(hipStreamSynchronize(mHandle));
+#elif defined(__CUDACC__)
+    GPUChkErrS(cudaStreamSynchronize(mHandle));
+#endif
+  }
+
+ private:
+  Handle mHandle{Default};
+};
+static_assert(sizeof(Stream) == sizeof(void*), "Stream type must match pointer type!");
+
+// Abstract vector for streams.
+// Handles specifically wrap around.
+class Streams
+{
+ public:
+  size_t size() const noexcept { return mStreams.size(); }
+  void resize(size_t n) { mStreams.resize(n); }
+  void clear() { mStreams.clear(); }
+  auto& operator[](size_t i) { return mStreams[i % mStreams.size()]; }
+  void push_back(const Stream& stream) { mStreams.push_back(stream); }
+
+ private:
+  std::vector<Stream> mStreams;
+};
+
 } // namespace gpu
 } // namespace its
 } // namespace o2
 
-#endif
+#endif