Update GPU tracking to follow new development

AliceO2Group · Nov 4, 2023 · c53ccd7 · c53ccd7
1 parent 029e62e
commit c53ccd7
Show file tree

Hide file tree

Showing 7 changed files with 232 additions and 63 deletions.
diff --git a/DataFormats/Detectors/ITSMFT/ITS/include/DataFormatsITS/TrackITS.h b/DataFormats/Detectors/ITSMFT/ITS/include/DataFormatsITS/TrackITS.h
@@ -92,8 +92,8 @@ class TrackITS : public o2::track::TrackParCov
 
   bool isBetter(const TrackITS& best, float maxChi2) const;
 
-  o2::track::TrackParCov& getParamIn() { return *this; }
-  const o2::track::TrackParCov& getParamIn() const { return *this; }
+  GPUhdi() o2::track::TrackParCov& getParamIn() { return *this; }
+  GPUhdi() const o2::track::TrackParCov& getParamIn() const { return *this; }
 
   GPUhdi() o2::track::TrackParCov& getParamOut() { return mParamOut; }
   GPUhdi() const o2::track::TrackParCov& getParamOut() const { return mParamOut; }

diff --git a/DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrizationWithError.h b/DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrizationWithError.h
@@ -46,7 +46,7 @@ class TrackParametrizationWithError : public TrackParametrization<value_T>
   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);
 
-  GPUdDefault() TrackParametrizationWithError(const TrackParametrizationWithError& src) = default;
+  GPUhdDefault() TrackParametrizationWithError(const TrackParametrizationWithError& src) = default;
   GPUdDefault() TrackParametrizationWithError(TrackParametrizationWithError&& src) = default;
   GPUhdDefault() TrackParametrizationWithError& operator=(const TrackParametrizationWithError& src) = default;
   GPUhdDefault() TrackParametrizationWithError& operator=(TrackParametrizationWithError&& src) = default;

diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TimeFrameGPU.h b/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TimeFrameGPU.h
@@ -198,7 +198,9 @@ class TimeFrameGPU : public TimeFrame
   void loadTrackSeedsDevice();
   void loadTrackSeedsChi2Device();
   void loadRoadsDevice();
-  void createTrackITSExtDevice();
+  void loadTrackSeedsDevice(std::vector<CellSeed>&);
+  void createTrackITSExtDevice(const std::vector<CellSeed>& seeds);
+  void createTrackITSExtDevice(); // deprecated
   void downloadTrackITSExtDevice();
   void initDeviceChunks(const int, const int);
   template <Task task>
@@ -232,6 +234,7 @@ class TimeFrameGPU : public TimeFrame
   Cluster** getDeviceArrayUnsortedClusters() const { return mUnsortedClustersDeviceArray; }
   Tracklet** getDeviceArrayTracklets() const { return mTrackletsDeviceArray; }
   CellSeed** getDeviceArrayCells() const { return mCellsDeviceArray; }
+  CellSeed* getDeviceTrackSeeds() { return mTrackSeedsDevice; }
   o2::track::TrackParCovF** getDeviceArrayTrackSeeds() { return mCellSeedsDeviceArray; }
   float** getDeviceArrayTrackSeedsChi2() { return mCellSeedsChi2DeviceArray; }
   void setDevicePropagator(const o2::base::PropagatorImpl<float>*) override;
@@ -263,6 +266,7 @@ class TimeFrameGPU : public TimeFrame
   std::array<Tracklet*, nLayers - 1> mTrackletsDevice;
   Tracklet** mTrackletsDeviceArray;
   std::array<CellSeed*, nLayers - 2> mCellsDevice;
+  CellSeed* mTrackSeedsDevice;
   CellSeed** mCellsDeviceArray;
   std::array<o2::track::TrackParCovF*, nLayers - 2> mCellSeedsDevice;
   o2::track::TrackParCovF** mCellSeedsDeviceArray;

diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TimeFrameGPU.cu b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TimeFrameGPU.cu
@@ -575,6 +575,26 @@ void TimeFrameGPU<nLayers>::loadRoadsDevice()
   checkGPUError(cudaMemcpyAsync(mRoadsDevice, mRoads.data(), mRoads.size() * sizeof(Road<nLayers - 2>), cudaMemcpyHostToDevice, mGpuStreams[0].get()));
 }
 
+template <int nLayers>
+void TimeFrameGPU<nLayers>::loadTrackSeedsDevice(std::vector<CellSeed>& seeds)
+{
+  LOGP(debug, "gpu-transfer: loading {} track seeds, for {} MB.", seeds.size(), seeds.size() * sizeof(CellSeed) / MB);
+  allocMemAsync(reinterpret_cast<void**>(&mTrackSeedsDevice), seeds.size() * sizeof(CellSeed), &(mGpuStreams[0]), false);
+  checkGPUError(cudaHostRegister(seeds.data(), seeds.size() * sizeof(CellSeed), cudaHostRegisterPortable));
+  checkGPUError(cudaMemcpyAsync(mTrackSeedsDevice, seeds.data(), seeds.size() * sizeof(CellSeed), cudaMemcpyHostToDevice, mGpuStreams[0].get()));
+}
+
+template <int nLayers>
+void TimeFrameGPU<nLayers>::createTrackITSExtDevice(const std::vector<CellSeed>& seeds)
+{
+  mTrackITSExt.clear();
+  mTrackITSExt.resize(seeds.size());
+  LOGP(debug, "gpu-allocation: reserving {} tracks, for {} MB.", seeds.size(), seeds.size() * sizeof(o2::its::TrackITSExt) / MB);
+  allocMemAsync(reinterpret_cast<void**>(&mTrackITSExtDevice), seeds.size() * sizeof(o2::its::TrackITSExt), &(mGpuStreams[0]), false);
+  checkGPUError(cudaMemsetAsync(mTrackITSExtDevice, 0, seeds.size() * sizeof(o2::its::TrackITSExt), mGpuStreams[0].get()));
+  checkGPUError(cudaHostRegister(mTrackITSExt.data(), seeds.size() * sizeof(o2::its::TrackITSExt), cudaHostRegisterPortable));
+}
+
 template <int nLayers>
 void TimeFrameGPU<nLayers>::createTrackITSExtDevice()
 {
@@ -589,8 +609,10 @@ void TimeFrameGPU<nLayers>::createTrackITSExtDevice()
 template <int nLayers>
 void TimeFrameGPU<nLayers>::downloadTrackITSExtDevice()
 {
-  LOGP(debug, "gpu-transfer: downloading {} tracks, for {} MB.", mRoads.size(), mRoads.size() * sizeof(o2::its::TrackITSExt) / MB);
-  checkGPUError(cudaMemcpyAsync(mTrackITSExt.data(), mTrackITSExtDevice, mRoads.size() * sizeof(o2::its::TrackITSExt), cudaMemcpyDeviceToHost, mGpuStreams[0].get()));
+  LOGP(debug, "gpu-transfer: downloading {} tracks, for {} MB.", mTrackITSExt.size(), mTrackITSExt.size() * sizeof(o2::its::TrackITSExt) / MB);
+  checkGPUError(cudaMemcpyAsync(mTrackITSExt.data(), mTrackITSExtDevice, mTrackITSExt.size() * sizeof(o2::its::TrackITSExt), cudaMemcpyDeviceToHost, mGpuStreams[0].get()));
+  checkGPUError(cudaHostUnregister(mTrackITSExt.data()));
+  discardResult(cudaDeviceSynchronize());
 }
 
 template <int nLayers>