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

@@ -63,6 +63,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 occupancyTotal;                 // Total occupancy in the TPC (nCl / nHbf)
+  uint32_t occupancyMapSize;               // Size of occupancy map
 
   GPUParamSector SectorParam[GPUCA_NSECTORS];
 

GPU/GPUTracking/Base/GPUParam.inc

@@ -213,7 +213,7 @@ GPUdi() float GPUParam::GetUnscaledMult(float time) const
   if (!occupancyMap) {
     return 0.f;
   }
-  const uint32_t bin = CAMath::Max(0.f, time / rec.tpc.occupancyMapTimeBins);
+  const uint32_t bin = CAMath::Min<uint32_t>(occupancyMapSize - 1, CAMath::Max(0.f, time / rec.tpc.occupancyMapTimeBins));
   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,9 +356,13 @@ 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().occupancyMapSize = mapSize;
   param().occupancyTotal = occupancyTotal;
   if (IsGPU()) {
     if (!((size_t)&param().occupancyTotal - (size_t)&param().occupancyMap == sizeof(param().occupancyMap) && sizeof(param().occupancyMap) == sizeof(size_t) && sizeof(param().occupancyTotal) < sizeof(size_t))) {

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

@@ -26,6 +26,7 @@
 #include "TPCBase/CRU.h"
 #include "TPCBase/RDHUtils.h"
 #include "DataFormatsTPC/ZeroSuppression.h"
+#include "GPUTPCClusterOccupancyMap.h"
 #include <gsl/span>
 
 using namespace o2::gpu;
@@ -132,6 +133,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);
     }
   }
 }

GPU/GPUTracking/Standalone/Benchmark/standalone.cxx

@@ -453,7 +453,7 @@ int32_t SetupReconstruction()
   if (configStandalone.proc.doublePipeline) {
     recPipeline->SetSettings(&grp, &recSet, &procSet, &steps);
   }
-  if (configStandalone.testSyncAsync) {
+  if (configStandalone.testSyncAsync) { // TODO: Add --async mode / flag
     // Set settings for asynchronous
     steps.steps.setBits(GPUDataTypes::RecoStep::TPCDecompression, true);
     steps.steps.setBits(GPUDataTypes::RecoStep::TPCdEdx, true);