GPU TPC: occuapncy Map needs bounds check

Author: davidrohr

GPU/GPUTracking/Base/GPUParam.cxx

@@ -35,6 +35,7 @@ void GPUParam::SetDefaults(float solenoidBz, bool assumeConstantBz)
   memset((void*)this, 0, sizeof(*this));
   new (&rec) GPUSettingsRec;
   occupancyMap = nullptr;
+  occupancyMap = 0;
   occupancyTotal = 0;
 
 #ifdef GPUCA_TPC_GEOMETRY_O2

GPU/GPUTracking/Base/GPUParam.h

@@ -62,6 +62,7 @@ struct GPUParam_t {
 
   GPUTPCGMPolynomialField polynomialField; // Polynomial approx. of magnetic field for TPC GM
   const uint32_t* occupancyMap;            // Ptr to TPC occupancy map
+  uint32_t occupancyMapSize;               // Size of occupancy map
   uint32_t occupancyTotal;                 // Total occupancy in the TPC (nCl / nHbf)
 
   GPUParamSector SectorParam[GPUCA_NSECTORS];

GPU/GPUTracking/Base/GPUParam.inc

@@ -214,6 +214,9 @@ GPUdi() float GPUParam::GetUnscaledMult(float time) const
     return 0.f;
   }
   const uint32_t bin = CAMath::Max(0.f, time / rec.tpc.occupancyMapTimeBins);
+  if (bin >= occupancyMapSize) {
+    return 0.f;
+  }
   return occupancyMap[bin];
 }
 

GPU/GPUTracking/Base/GPUReconstructionCPU.cxx

@@ -31,6 +31,7 @@
 #include "GPULogging.h"
 #include "GPUMemorySizeScalers.h"
 #include "GPUReconstructionProcessingKernels.inc"
+#include "GPUTPCClusterOccupancyMap.h"
 
 #include <atomic>
 #include <ctime>
@@ -355,8 +356,11 @@ void GPUReconstructionCPU::ResetDeviceProcessorTypes()
   }
 }
 
-void GPUReconstructionCPU::UpdateParamOccupancyMap(const uint32_t* mapHost, const uint32_t* mapGPU, uint32_t occupancyTotal, int32_t stream)
+void GPUReconstructionCPU::UpdateParamOccupancyMap(const uint32_t* mapHost, const uint32_t* mapGPU, uint32_t occupancyTotal, uint32_t mapSize, int32_t stream)
 {
+  if (mapHost && mapSize != GPUTPCClusterOccupancyMapBin::getNBins(param())) {
+    throw std::runtime_error("Updating occupancy map with object of invalid size");
+  }
   param().occupancyMap = mapHost;
   param().occupancyTotal = occupancyTotal;
   if (IsGPU()) {

GPU/GPUTracking/Base/GPUReconstructionCPU.h

@@ -48,7 +48,7 @@ class GPUReconstructionCPU : public GPUReconstructionProcessing::KernelInterface
 
   int32_t RunChains() override;
 
-  void UpdateParamOccupancyMap(const uint32_t* mapHost, const uint32_t* mapGPU, uint32_t occupancyTotal, int32_t stream = -1);
+  void UpdateParamOccupancyMap(const uint32_t* mapHost, const uint32_t* mapGPU, uint32_t occupancyTotal, uint32_t mapSize, int32_t stream = -1);
 
  protected:
   struct GPUProcessorProcessors : public GPUProcessor {

GPU/GPUTracking/Global/GPUChainTrackingSectorTracker.cxx

@@ -126,19 +126,19 @@ int32_t GPUChainTracking::RunTPCTrackingSectors_internal()
     auto* ptrTmp = (GPUTPCClusterOccupancyMapBin*)mRec->AllocateVolatileMemory(GPUTPCClusterOccupancyMapBin::getTotalSize(param()), doGPU);
     runKernel<GPUMemClean16>(GetGridAutoStep(streamInitAndOccMap, RecoStep::TPCSectorTracking), ptrTmp, GPUTPCClusterOccupancyMapBin::getTotalSize(param()));
     runKernel<GPUTPCCreateOccupancyMap, GPUTPCCreateOccupancyMap::fill>(GetGridBlk(GPUCA_NSECTORS * GPUCA_ROW_COUNT, streamInitAndOccMap), ptrTmp);
-    runKernel<GPUTPCCreateOccupancyMap, GPUTPCCreateOccupancyMap::fold>(GetGridBlk(GPUTPCClusterOccupancyMapBin::getNBins(param()), streamInitAndOccMap), ptrTmp, ptr + 2);
+    runKernel<GPUTPCCreateOccupancyMap, GPUTPCCreateOccupancyMap::fold>(GetGridBlk(mInputsHost->mTPCClusterOccupancyMapSize, streamInitAndOccMap), ptrTmp, ptr + 2);
     mRec->ReturnVolatileMemory();
     mInputsHost->mTPCClusterOccupancyMap[1] = param().rec.tpc.occupancyMapTimeBins * 0x10000 + param().rec.tpc.occupancyMapTimeBinsAverage;
     if (doGPU) {
-      GPUMemCpy(RecoStep::TPCSectorTracking, mInputsHost->mTPCClusterOccupancyMap + 2, mInputsShadow->mTPCClusterOccupancyMap + 2, sizeof(*ptr) * GPUTPCClusterOccupancyMapBin::getNBins(mRec->GetParam()), streamInitAndOccMap, false, &mEvents->init);
+      GPUMemCpy(RecoStep::TPCSectorTracking, mInputsHost->mTPCClusterOccupancyMap + 2, mInputsShadow->mTPCClusterOccupancyMap + 2, sizeof(*ptr) * mInputsHost->mTPCClusterOccupancyMapSize, streamInitAndOccMap, false, &mEvents->init);
     } else {
       TransferMemoryResourceLinkToGPU(RecoStep::TPCSectorTracking, mInputsHost->mResourceOccupancyMap, streamInitAndOccMap, &mEvents->init);
     }
   }
   if (param().rec.tpc.occupancyMapTimeBins || param().rec.tpc.sysClusErrorC12Norm) {
     uint32_t& occupancyTotal = *mInputsHost->mTPCClusterOccupancyMap;
     occupancyTotal = CAMath::Float2UIntRn(mRec->MemoryScalers()->nTPCHits / (mIOPtrs.settingsTF && mIOPtrs.settingsTF->hasNHBFPerTF ? mIOPtrs.settingsTF->nHBFPerTF : 128));
-    mRec->UpdateParamOccupancyMap(param().rec.tpc.occupancyMapTimeBins ? mInputsHost->mTPCClusterOccupancyMap + 2 : nullptr, doGPU && param().rec.tpc.occupancyMapTimeBins ? mInputsShadow->mTPCClusterOccupancyMap + 2 : nullptr, occupancyTotal, streamInitAndOccMap);
+    mRec->UpdateParamOccupancyMap(param().rec.tpc.occupancyMapTimeBins ? mInputsHost->mTPCClusterOccupancyMap + 2 : nullptr, doGPU && param().rec.tpc.occupancyMapTimeBins ? mInputsShadow->mTPCClusterOccupancyMap + 2 : nullptr, occupancyTotal, mInputsHost->mTPCClusterOccupancyMapSize, streamInitAndOccMap);
   }
 
   int32_t streamMap[NSECTORS];

GPU/GPUTracking/Global/GPUTrackingInputProvider.cxx

@@ -82,7 +82,8 @@ void* GPUTrackingInputProvider::SetPointersInputTRD(void* mem)
 void* GPUTrackingInputProvider::SetPointersTPCOccupancyMap(void* mem)
 {
   if (mHoldTPCOccupancyMap) {
-    computePointerWithAlignment(mem, mTPCClusterOccupancyMap, (mRec->GetParam().rec.tpc.occupancyMapTimeBins ? GPUTPCClusterOccupancyMapBin::getNBins(mRec->GetParam()) + 1 : 0) + 1); // +1 for total occupancy estimator, +1 for sanity check information
+    mTPCClusterOccupancyMapSize = mRec->GetParam().rec.tpc.occupancyMapTimeBins ? GPUTPCClusterOccupancyMapBin::getNBins(mRec->GetParam()) : 0;
+    computePointerWithAlignment(mem, mTPCClusterOccupancyMap, (mRec->GetParam().rec.tpc.occupancyMapTimeBins ? mTPCClusterOccupancyMapSize + 1 : 0) + 1); // +1 for total occupancy estimator, +1 for sanity check information
   }
   return mem;
 }

GPU/GPUTracking/Global/GPUTrackingInputProvider.h

@@ -81,6 +81,7 @@ class GPUTrackingInputProvider : public GPUProcessor
   o2::tpc::ClusterNative* mPclusterNativeOutput = nullptr;
 
   uint32_t* mTPCClusterOccupancyMap = nullptr;
+  uint32_t mTPCClusterOccupancyMapSize = 0;
 
   uint32_t* mErrorCodes = nullptr;
 };

GPU/GPUTracking/Interface/GPUO2InterfaceUtils.cxx

@@ -132,6 +132,7 @@ void GPUO2InterfaceUtils::paramUseExternalOccupancyMap(GPUParam* param, uint32_t
     param->occupancyTotal = *occupancymap;
     if (param->rec.tpc.occupancyMapTimeBins) {
       param->occupancyMap = occupancymap + 2;
+      param->occupancyMapSize = GPUTPCClusterOccupancyMapBin::getNBins(*param);
     }
   }
 }