Move cellFinder on a pure GPU-based implementation

mconcas · Nov 5, 2024 · a7f52e5 · a7f52e5
1 parent 688eac2
commit a7f52e5
Show file tree

Hide file tree

Showing 10 changed files with 585 additions and 314 deletions.
diff --git a/DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrizationWithError.h b/DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrizationWithError.h
@@ -42,7 +42,7 @@ class TrackParametrizationWithError : public TrackParametrization<value_T>
   using MatrixDSym5 = o2::math_utils::SMatrix<double, kNParams, kNParams, o2::math_utils::MatRepSym<double, kNParams>>;
   using MatrixD5 = o2::math_utils::SMatrix<double, kNParams, kNParams, o2::math_utils::MatRepStd<double, kNParams, kNParams>>;
 
-  GPUd() TrackParametrizationWithError();
+  GPUhd() TrackParametrizationWithError();
   GPUd() TrackParametrizationWithError(value_t x, value_t alpha, const params_t& par, const covMat_t& cov, int charge = 1, const PID pid = PID::Pion);
   GPUd() TrackParametrizationWithError(const dim3_t& xyz, const dim3_t& pxpypz,
                                        const gpu::gpustd::array<value_t, kLabCovMatSize>& cv, int sign, bool sectorAlpha = true, const PID pid = PID::Pion);
@@ -145,7 +145,7 @@ class TrackParametrizationWithError : public TrackParametrization<value_T>
 
 //__________________________________________________________________________
 template <typename value_T>
-GPUdi() TrackParametrizationWithError<value_T>::TrackParametrizationWithError() : TrackParametrization<value_T>{}
+GPUhdi() TrackParametrizationWithError<value_T>::TrackParametrizationWithError() : TrackParametrization<value_T>{}
 {
 }
 

diff --git a/Detectors/ITSMFT/ITS/tracking/CMakeLists.txt b/Detectors/ITSMFT/ITS/tracking/CMakeLists.txt
@@ -37,7 +37,6 @@ o2_add_library(ITStracking
                        O2::ITSMFTReconstruction
                        O2::DataFormatsITS)
 
-
 if (OpenMP_CXX_FOUND)
         target_compile_definitions(${targetName} PRIVATE WITH_OPENMP)
         target_link_libraries(${targetName} PRIVATE OpenMP::OpenMP_CXX)

diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TimeFrameGPU.h b/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TimeFrameGPU.h
@@ -52,22 +52,28 @@ class TimeFrameGPU : public TimeFrame
   void initDevice(IndexTableUtils*, const TrackingParameters& trkParam, const TimeFrameGPUParameters&, const int, const int);
   void initDeviceSAFitting();
   void loadTrackingFrameInfoDevice(const int);
-  void loadUnsortedClustersDevice();
-  void loadClustersDevice();
+  void loadUnsortedClustersDevice(const int);
+  void loadClustersDevice(const int);
   void loadTrackletsDevice();
+  void loadTrackletsLUTDevice();
   void loadCellsDevice();
-  void loadCellsLUT();
+  void loadCellsLUTDevice();
   void loadTrackSeedsDevice();
   void loadTrackSeedsChi2Device();
   void loadRoadsDevice();
   void loadTrackSeedsDevice(std::vector<CellSeed>&);
+  void createCellsBuffers(const int);
+  void createCellsDevice();
+  void createCellsLUTDevice();
+  void createNeighboursDevice();
   void createNeighboursDevice(const unsigned int& layer, std::vector<std::pair<int, int>>& neighbours);
   void createNeighboursLUTDevice(const int, const unsigned int);
   void createTrackITSExtDevice(std::vector<CellSeed>&);
   void downloadTrackITSExtDevice(std::vector<CellSeed>&);
-  void downloadCellsNeighbours(std::vector<std::vector<std::pair<int, int>>>&, const int);
-  void downloadNeighboursLUT(std::vector<int>&, const int);
-  void downloadCellsDevice(const int);
+  void downloadCellsNeighboursDevice(std::vector<std::vector<std::pair<int, int>>>&, const int);
+  void downloadNeighboursLUTDevice(std::vector<int>&, const int);
+  void downloadCellsDevice();
+  void downloadCellsLUTDevice();
   void unregisterRest();
   void initDeviceChunks(const int, const int);
   template <Task task>
@@ -98,11 +104,11 @@ class TimeFrameGPU : public TimeFrame
   int* getDeviceNeighboursLUT(const int layer) { return mNeighboursLUTDevice[layer]; }
   gpuPair<int, int>* getDeviceNeighbours(const int layer) { return mNeighboursDevice[layer]; }
   TrackingFrameInfo* getDeviceTrackingFrameInfo(const int);
-  // TrackingFrameInfo** getDeviceArrayTrackingFrameInfo() { return mTrackingFrameInfoDeviceArray; }
   const TrackingFrameInfo** getDeviceArrayTrackingFrameInfo() const { return mTrackingFrameInfoDeviceArray; }
-  Cluster** getDeviceArrayClusters() const { return mClustersDeviceArray; }
-  Cluster** getDeviceArrayUnsortedClusters() const { return mUnsortedClustersDeviceArray; }
-  Tracklet** getDeviceArrayTracklets() const { return mTrackletsDeviceArray; }
+  const Cluster** getDeviceArrayClusters() const { return mClustersDeviceArray; }
+  const Cluster** getDeviceArrayUnsortedClusters() const { return mUnsortedClustersDeviceArray; }
+  const Tracklet** getDeviceArrayTracklets() const { return mTrackletsDeviceArray; }
+  const int** getDeviceArrayTrackletsLUT() const { return mTrackletsLUTDeviceArray; }
   int** getDeviceArrayCellsLUT() const { return mCellsLUTDeviceArray; }
   int** getDeviceArrayNeighboursCellLUT() const { return mNeighboursCellLUTDeviceArray; }
   CellSeed** getDeviceArrayCells() const { return mCellsDeviceArray; }
@@ -117,13 +123,21 @@ class TimeFrameGPU : public TimeFrame
   gsl::span<int> getHostNTracklets(const int chunkId);
   gsl::span<int> getHostNCells(const int chunkId);
 
+  // Host-available device getters
+  gsl::span<int*> getDeviceCellLUTs() { return mCellsLUTDevice; }
+  gsl::span<CellSeed*> getDeviceCells() { return mCellsDevice; }
+  gsl::span<int, nLayers - 2> getNCellsDevice() { return mNCells; }
+
  private:
   void allocMemAsync(void**, size_t, Stream*, bool); // Abstract owned and unowned memory allocations
   bool mHostRegistered = false;
   std::vector<GpuTimeFrameChunk<nLayers>> mMemChunks;
   TimeFrameGPUParameters mGpuParams;
   StaticTrackingParameters<nLayers> mStaticTrackingParams;
 
+  // Host-available device buffer sizes
+  std::array<int, nLayers - 2> mNCells;
+
   // Device pointers
   StaticTrackingParameters<nLayers>* mTrackingParamsDevice;
   IndexTableUtils* mIndexTableUtilsDevice;
@@ -135,12 +149,15 @@ class TimeFrameGPU : public TimeFrame
   // Hybrid pref
   std::array<Cluster*, nLayers> mClustersDevice;
   std::array<Cluster*, nLayers> mUnsortedClustersDevice;
-  Cluster** mClustersDeviceArray;
-  Cluster** mUnsortedClustersDeviceArray;
+  const Cluster** mClustersDeviceArray;
+  const Cluster** mUnsortedClustersDeviceArray;
   std::array<Tracklet*, nLayers - 1> mTrackletsDevice;
-  Tracklet** mTrackletsDeviceArray;
+  const Tracklet** mTrackletsDeviceArray;
+  const int** mTrackletsLUTDeviceArray;
+  std::array<int*, nLayers - 2> mTrackletsLUTDevice;
   std::array<int*, nLayers - 2> mCellsLUTDevice;
   std::array<int*, nLayers - 3> mNeighboursLUTDevice;
+
   int** mCellsLUTDeviceArray;
   int** mNeighboursCellDeviceArray;
   int** mNeighboursCellLUTDeviceArray;

diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackingKernels.h b/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackingKernels.h
@@ -49,6 +49,41 @@ GPUg() void fitTrackSeedsKernel(
   const o2::base::PropagatorF::MatCorrType matCorrType = o2::base::PropagatorF::MatCorrType::USEMatCorrLUT);
 #endif
 } // namespace gpu
+
+void countCellsHandler(const Cluster** sortedClusters,
+                       const Cluster** unsortedClusters,
+                       const TrackingFrameInfo** tfInfo,
+                       const Tracklet** tracklets,
+                       const int** trackletsLUT,
+                       const int nTracklets,
+                       const int layer,
+                       CellSeed* cells,
+                       int** cellsLUTsDeviceArray,
+                       int* cellsLUTsHost,
+                       const float bz,
+                       const float maxChi2ClusterAttachment,
+                       const float cellDeltaTanLambdaSigma,
+                       const float nSigmaCut,
+                       const int nBlocks,
+                       const int nThreads);
+
+void computeCellsHandler(const Cluster** sortedClusters,
+                         const Cluster** unsortedClusters,
+                         const TrackingFrameInfo** tfInfo,
+                         const Tracklet** tracklets,
+                         const int** trackletsLUT,
+                         const int nTracklets,
+                         const int layer,
+                         CellSeed* cells,
+                         int** cellsLUTsDeviceArray,
+                         int* cellsLUTsHost,
+                         const float bz,
+                         const float maxChi2ClusterAttachment,
+                         const float cellDeltaTanLambdaSigma,
+                         const float nSigmaCut,
+                         const int nBlocks,
+                         const int nThreads);
+
 void countCellNeighboursHandler(CellSeed** cellsLayersDevice,
                                 int* neighboursLUTs,
                                 int** cellsLUTs,

diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/CMakeLists.txt b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/CMakeLists.txt
@@ -14,6 +14,7 @@ if(CUDA_ENABLED)
 find_package(CUDAToolkit)
 message(STATUS "Building ITS CUDA tracker")
 # add_compile_options(-O0 -g -lineinfo -fPIC)
+# add_compile_definitions(ITS_MEASURE_GPU_TIME)
 o2_add_library(ITStrackingCUDA
                SOURCES ClusterLinesGPU.cu
                        Context.cu