ITS: GPU: fix upc iteration

Signed-off-by: Felix Schlepper <felix.schlepper@cern.ch>

Author: f3sch

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

@@ -251,10 +251,12 @@ void TimeFrameGPU<nLayers>::loadTrackingFrameInfoDevice(const int iteration)
 template <int nLayers>
 void TimeFrameGPU<nLayers>::loadMultiplicityCutMask(const int iteration)
 {
-  if (!iteration) {
+  if (!iteration || iteration == 3) { // we need to re-load the swapped mult-mask in upc iteration
     GPUTimer timer(mGpuStreams[0], "loading multiplicity cut mask");
-    GPULog("gpu-transfer: loading multiplicity cut mask with {} elements, for {:.2f} MB.", this->mMultiplicityCutMask.size(), this->mMultiplicityCutMask.size() * sizeof(bool) / constants::MB);
-    allocMemAsync(reinterpret_cast<void**>(&mMultMaskDevice), this->mMultiplicityCutMask.size() * sizeof(uint8_t), mGpuStreams[0], this->getExtAllocator());
+    GPULog("gpu-transfer: iteration {} loading multiplicity cut mask with {} elements, for {:.2f} MB.", iteration, this->mMultiplicityCutMask.size(), this->mMultiplicityCutMask.size() * sizeof(bool) / constants::MB);
+    if (!iteration) { // only allocate on first call
+      allocMemAsync(reinterpret_cast<void**>(&mMultMaskDevice), this->mMultiplicityCutMask.size() * sizeof(uint8_t), mGpuStreams[0], this->getExtAllocator());
+    }
     GPUChkErrS(cudaMemcpyAsync(mMultMaskDevice, this->mMultiplicityCutMask.data(), this->mMultiplicityCutMask.size() * sizeof(uint8_t), cudaMemcpyHostToDevice, mGpuStreams[0].get()));
   }
 }

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

@@ -544,9 +544,10 @@ GPUg() void computeLayerTrackletsMultiROFKernel(
       const float inverseR0{1.f / currentCluster.radius};
       for (int iV{startVtx}; iV < endVtx; ++iV) {
         auto& primaryVertex{primaryVertices[iV]};
-        if (primaryVertex.isFlagSet(Vertex::Flags::UPCMode) && iteration != 3) {
+        if ((primaryVertex.isFlagSet(Vertex::Flags::UPCMode) && iteration != 3) || (iteration == 3 && !primaryVertex.isFlagSet(Vertex::Flags::UPCMode))) {
           continue;
         }
+
         const float resolution = o2::gpu::CAMath::Sqrt(math_utils::Sq(resolutionPV) / primaryVertex.getNContributors() + math_utils::Sq(positionResolution));
         const float tanLambda{(currentCluster.zCoordinate - primaryVertex.getZ()) * inverseR0};
         const float zAtRmin{tanLambda * (minR - currentCluster.radius) + currentCluster.zCoordinate};

Detectors/ITSMFT/ITS/tracking/src/TrackerTraits.cxx

@@ -107,9 +107,10 @@ void TrackerTraits<nLayers>::computeLayerTracklets(const int iteration, int iROF
 
           for (int iV = startVtx; iV < endVtx; ++iV) {
             const auto& pv = primaryVertices[iV];
-            if (pv.isFlagSet(Vertex::Flags::UPCMode) && iteration != 3) {
+            if ((pv.isFlagSet(Vertex::Flags::UPCMode) && iteration != 3) || (iteration == 3 && !pv.isFlagSet(Vertex::Flags::UPCMode))) {
               continue;
             }
+
             const float resolution = o2::gpu::CAMath::Sqrt(math_utils::Sq(mTimeFrame->getPositionResolution(iLayer)) + math_utils::Sq(mTrkParams[iteration].PVres) / float(pv.getNContributors()));
             const float tanLambda = (currentCluster.zCoordinate - pv.getZ()) * inverseR0;
             const float zAtRmin = tanLambda * (mTimeFrame->getMinR(iLayer + 1) - currentCluster.radius) + currentCluster.zCoordinate;