ITS: GPU report found neighbours (#14438)

Author: f3sch

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

@@ -130,7 +130,9 @@ class TimeFrameGPU : public TimeFrame<nLayers>
   // Host-specific getters
   gsl::span<int, nLayers - 1> getNTracklets() { return mNTracklets; }
   gsl::span<int, nLayers - 2> getNCells() { return mNCells; }
-  std::array<int, nLayers - 2>& getArrayNCells() { return mNCells; }
+  auto& getArrayNCells() { return mNCells; }
+  gsl::span<int, nLayers - 3> getNNeighbours() { return mNNeighbours; }
+  auto& getArrayNNeighbours() { return mNNeighbours; }
 
   // Host-available device getters
   gsl::span<int*> getDeviceTrackletsLUTs() { return mTrackletsLUTDevice; }
@@ -139,7 +141,9 @@ class TimeFrameGPU : public TimeFrame<nLayers>
   gsl::span<CellSeed*> getDeviceCells() { return mCellsDevice; }
 
   // Overridden getters
-  int getNumberOfCells() const;
+  int getNumberOfTracklets() const final;
+  int getNumberOfCells() const final;
+  int getNumberOfNeighbours() const final;
 
  private:
   void allocMemAsync(void**, size_t, Stream*, bool); // Abstract owned and unowned memory allocations
@@ -149,6 +153,7 @@ class TimeFrameGPU : public TimeFrame<nLayers>
   // Host-available device buffer sizes
   std::array<int, nLayers - 1> mNTracklets;
   std::array<int, nLayers - 2> mNCells;
+  std::array<int, nLayers - 3> mNNeighbours;
 
   // Device pointers
   IndexTableUtils* mIndexTableUtilsDevice;
@@ -218,12 +223,24 @@ inline std::vector<unsigned int> TimeFrameGPU<nLayers>::getClusterSizes()
   return sizes;
 }
 
+template <int nLayers>
+inline int TimeFrameGPU<nLayers>::getNumberOfTracklets() const
+{
+  return std::accumulate(mNTracklets.begin(), mNTracklets.end(), 0);
+}
+
 template <int nLayers>
 inline int TimeFrameGPU<nLayers>::getNumberOfCells() const
 {
   return std::accumulate(mNCells.begin(), mNCells.end(), 0);
 }
 
+template <int nLayers>
+inline int TimeFrameGPU<nLayers>::getNumberOfNeighbours() const
+{
+  return std::accumulate(mNNeighbours.begin(), mNNeighbours.end(), 0);
+}
+
 } // namespace o2::its::gpu
 
 #endif

Detectors/ITSMFT/ITS/tracking/GPU/cuda/TimeFrameGPU.cu

@@ -342,6 +342,9 @@ void TimeFrameGPU<nLayers>::createNeighboursIndexTablesDevice()
     LOGP(debug, "gpu-transfer: loading neighbours LUT for {} elements on layer {}, for {} MB.", mNCells[iLayer], iLayer, mNCells[iLayer] * sizeof(CellSeed) / MB);
     allocMemAsync(reinterpret_cast<void**>(&mNeighboursIndexTablesDevice[iLayer]), (mNCells[iLayer] + 1) * sizeof(int), nullptr, this->getExtAllocator());
     GPUChkErrS(cudaMemsetAsync(mNeighboursIndexTablesDevice[iLayer], 0, (mNCells[iLayer] + 1) * sizeof(int), mGpuStreams[0]->get()));
+    if (iLayer < nLayers - 3) {
+      mNNeighbours[iLayer] = 0;
+    }
   }
   STOP_GPU_STREAM_TIMER(mGpuStreams[0]->get());
 }

Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackerTraitsGPU.cxx

@@ -169,7 +169,7 @@ template <int nLayers>
 void TrackerTraitsGPU<nLayers>::findCellsNeighbours(const int iteration)
 {
   mTimeFrameGPU->createNeighboursIndexTablesDevice();
-  auto& conf = o2::its::ITSGpuTrackingParamConfig::Instance();
+  const auto& conf = o2::its::ITSGpuTrackingParamConfig::Instance();
   for (int iLayer{0}; iLayer < this->mTrkParams[iteration].CellsPerRoad() - 1; ++iLayer) {
     const int nextLayerCellsNum{static_cast<int>(mTimeFrameGPU->getNCells()[iLayer + 1])};
 
@@ -208,10 +208,11 @@ void TrackerTraitsGPU<nLayers>::findCellsNeighbours(const int iteration)
                                  conf.nBlocks,
                                  conf.nThreads);
 
-    filterCellNeighboursHandler(mTimeFrameGPU->getDeviceNeighbourPairs(iLayer),
-                                mTimeFrameGPU->getDeviceNeighbours(iLayer),
-                                nNeigh,
-                                mTimeFrameGPU->getExternalAllocator());
+    nNeigh = filterCellNeighboursHandler(mTimeFrameGPU->getDeviceNeighbourPairs(iLayer),
+                                         mTimeFrameGPU->getDeviceNeighbours(iLayer),
+                                         nNeigh,
+                                         mTimeFrameGPU->getExternalAllocator());
+    mTimeFrameGPU->getArrayNNeighbours()[iLayer] = nNeigh;
   }
   mTimeFrameGPU->createNeighboursDeviceArray();
   mTimeFrameGPU->unregisterRest();

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

@@ -175,9 +175,9 @@ struct TimeFrame {
   auto& getVerticesMCRecInfo() { return mVerticesMCRecInfo; }
 
   int getNumberOfClusters() const;
-  int getNumberOfCells() const;
-  int getNumberOfTracklets() const;
-  int getNumberOfNeighbours() const;
+  virtual int getNumberOfCells() const;
+  virtual int getNumberOfTracklets() const;
+  virtual int getNumberOfNeighbours() const;
   size_t getNumberOfTracks() const;
   size_t getNumberOfUsedClusters() const;
   auto getNumberOfExtendedTracks() const { return mNExtendedTracks; }