enable changing the heap size on device via input parameter (apt-sim#355)

This PR allows to adjust the HeapSize on device via input parameter.
This is needed for the ATLAS geometry. Note that this might become
redundant when the heap usage on device is fixed but for now it is
needed.

Author: SeverinDiederichs

include/AdePT/core/AdePTConfiguration.hh

@@ -37,6 +37,7 @@ public:
   void SetMillionsOfHitSlots(double millionSlots) { fMillionsOfHitSlots = millionSlots; }
   void SetHitBufferFlushThreshold(float threshold) { fHitBufferFlushThreshold = threshold; }
   void SetCUDAStackLimit(int limit) { fCUDAStackLimit = limit; }
+  void SetCUDAHeapLimit(int limit) { fCUDAHeapLimit = limit; }
 
   // We temporarily load VecGeom geometry from GDML
   void SetVecGeomGDML(std::string filename) { fVecGeomGDML = filename; }
@@ -47,6 +48,7 @@ public:
   int GetVerbosity() { return fVerbosity; };
   int GetTransportBufferThreshold() { return fTransportBufferThreshold; }
   int GetCUDAStackLimit() { return fCUDAStackLimit; }
+  int GetCUDAHeapLimit() { return fCUDAHeapLimit; }
   float GetHitBufferFlushThreshold() { return fHitBufferFlushThreshold; }
   double GetMillionsOfTrackSlots() { return fMillionsOfTrackSlots; }
   double GetMillionsOfHitSlots() { return fMillionsOfHitSlots; }
@@ -62,6 +64,7 @@ private:
   int fVerbosity{0};
   int fTransportBufferThreshold{200};
   int fCUDAStackLimit{0};
+  int fCUDAHeapLimit{0};
   float fHitBufferFlushThreshold{0.8};
   double fMillionsOfTrackSlots{1};
   double fMillionsOfHitSlots{1};

include/AdePT/core/AdePTTransport.h

@@ -67,6 +67,7 @@ class AdePTTransport : public AdePTTransportInterface {
   void SetGPURegionNames(std::vector<std::string> const *regionNames) { fGPURegionNames = regionNames; }
   /// @brief Set CUDA device stack limit
   void SetCUDAStackLimit(int limit) { fCUDAStackLimit = limit; }
+  void SetCUDAHeapLimit(int limit) { fCUDAHeapLimit = limit; }
   std::vector<std::string> const *GetGPURegionNames() { return fGPURegionNames; }
   /// @brief Create material-cut couple index array
   /// @brief Initialize service and copy geometry & physics data on device
@@ -89,6 +90,7 @@ class AdePTTransport : public AdePTTransportInterface {
   size_t fBufferThreshold{20};                         ///< Buffer threshold for flushing AdePT transport buffer
   int fDebugLevel{1};                                  ///< Debug level
   int fCUDAStackLimit{0};                              ///< CUDA device stack limit
+  int fCUDAHeapLimit{0};                               ///< CUDA device heap limit
   GPUstate *fGPUstate{nullptr};                        ///< CUDA state placeholder
   AdeptScoring *fScoring{nullptr};                     ///< User scoring object
   AdeptScoring *fScoring_dev{nullptr};                 ///< Device ptr for scoring data

include/AdePT/core/AdePTTransport.icc

@@ -48,6 +48,7 @@ AdePTTransport<IntegrationLayer>::AdePTTransport(AdePTConfiguration &configurati
   fTrackInAllRegions = configuration.GetTrackInAllRegions();
   fGPURegionNames    = configuration.GetGPURegionNames();
   fCUDAStackLimit    = configuration.GetCUDAStackLimit();
+  fCUDAHeapLimit     = configuration.GetCUDAHeapLimit();
   fCapacity          = 1024 * 1024 * configuration.GetMillionsOfTrackSlots() / configuration.GetNumThreads();
   fHitBufferCapacity = 1024 * 1024 * configuration.GetMillionsOfHitSlots() / configuration.GetNumThreads();
 
@@ -99,6 +100,10 @@ bool AdePTTransport<IntegrationLayer>::InitializeGeometry(const vecgeom::cxx::VP
     std::cout << "CUDA Device stack limit: " << fCUDAStackLimit << "\n";
     cudaDeviceSetLimit(cudaLimitStackSize, fCUDAStackLimit);
   }
+  if (fCUDAHeapLimit > 0) {
+    std::cout << "CUDA Device heap limit: " << fCUDAHeapLimit << "\n";
+    cudaDeviceSetLimit(cudaLimitMallocHeapSize, fCUDAHeapLimit);
+  }
   bool success = true;
 #ifdef ADEPT_USE_SURF
 #ifdef ADEPT_USE_SURF_SINGLE

include/AdePT/core/AdePTTransportInterface.hh

@@ -41,6 +41,7 @@ public:
   virtual void SetGPURegionNames(std::vector<std::string> const *regionNames) = 0;
   virtual std::vector<std::string> const *GetGPURegionNames()                 = 0;
   virtual void SetCUDAStackLimit(int limit)                                   = 0;
+  virtual void SetCUDAHeapLimit(int limit)                                    = 0;
   /// @brief Initialize service and copy geometry & physics data on device
   virtual void Initialize(bool common_data = false) = 0;
   /// @brief Initialize the ApplyCuts flag on device

include/AdePT/core/AsyncAdePTTransport.hh

@@ -46,6 +46,7 @@ private:
   unsigned int fScoringCapacity{0}; ///< Number of hit slots to allocate on device
   int fDebugLevel{0};               ///< Debug level
   int fCUDAStackLimit{0};           ///< CUDA device stack limit
+  int fCUDAHeapLimit{0};            ///< CUDA device stack limit
   std::vector<IntegrationLayer> fIntegrationLayerObjects;
   std::unique_ptr<GPUstate, GPUstateDeleter> fGPUstate{nullptr}; ///< CUDA state placeholder
   std::vector<AdePTScoring> fScoring;                            ///< User scoring objects per G4 worker
@@ -90,6 +91,7 @@ public:
   bool GetTrackInAllRegions() const override { return fTrackInAllRegions; }
   void SetGPURegionNames(std::vector<std::string> const *regionNames) override { fGPURegionNames = regionNames; }
   void SetCUDAStackLimit(int limit) override {};
+  void SetCUDAHeapLimit(int limit) override {};
   std::vector<std::string> const *GetGPURegionNames() override { return fGPURegionNames; }
   /// No effect
   void Initialize(bool) override {}

include/AdePT/core/AsyncAdePTTransport.icc

@@ -84,7 +84,8 @@ AsyncAdePTTransport<IntegrationLayer>::AsyncAdePTTransport(AdePTConfiguration &c
       fScoringCapacity{(uint)(1024 * 1024 * configuration.GetMillionsOfHitSlots())},
       fDebugLevel{configuration.GetVerbosity()}, fIntegrationLayerObjects(fNThread), fEventStates(fNThread),
       fGPUNetEnergy(fNThread, 0.0), fTrackInAllRegions{configuration.GetTrackInAllRegions()},
-      fGPURegionNames{configuration.GetGPURegionNames()}, fCUDAStackLimit{configuration.GetCUDAStackLimit()}
+      fGPURegionNames{configuration.GetGPURegionNames()}, fCUDAStackLimit{configuration.GetCUDAStackLimit()},
+      fCUDAHeapLimit{configuration.GetCUDAHeapLimit()}
 {
   if (fNThread > kMaxThreads)
     throw std::invalid_argument("AsyncAdePTTransport limited to " + std::to_string(kMaxThreads) + " threads");
@@ -176,6 +177,10 @@ bool AsyncAdePTTransport<IntegrationLayer>::InitializeGeometry(const vecgeom::cx
     std::cout << "CUDA Device stack limit: " << fCUDAStackLimit << "\n";
     cudaDeviceSetLimit(cudaLimitStackSize, fCUDAStackLimit);
   }
+  if (fCUDAHeapLimit > 0) {
+    std::cout << "CUDA Device heap limit: " << fCUDAHeapLimit << "\n";
+    cudaDeviceSetLimit(cudaLimitMallocHeapSize, fCUDAHeapLimit);
+  }
   cudaManager.LoadGeometry(world);
   auto world_dev = cudaManager.Synchronize();
   // Initialize BVH

include/AdePT/integration/AdePTConfigurationMessenger.hh

@@ -34,6 +34,7 @@ private:
 
   G4UIdirectory *fDir;
   G4UIcmdWithAnInteger *fSetCUDAStackLimitCmd;
+  G4UIcmdWithAnInteger *fSetCUDAHeapLimitCmd;
   G4UIcmdWithABool *fSetTrackInAllRegionsCmd;
   G4UIcmdWithAString *fAddRegionCmd;
   G4UIcmdWithABool *fActivateAdePTCmd;

src/AdePTConfigurationMessenger.cc

@@ -58,6 +58,8 @@ AdePTConfigurationMessenger::AdePTConfigurationMessenger(AdePTConfiguration *ade
 
   fSetCUDAStackLimitCmd = new G4UIcmdWithAnInteger("/adept/setCUDAStackLimit", this);
   fSetCUDAStackLimitCmd->SetGuidance("Set the CUDA device stack limit");
+  fSetCUDAHeapLimitCmd = new G4UIcmdWithAnInteger("/adept/setCUDAHeapLimit", this);
+  fSetCUDAHeapLimitCmd->SetGuidance("Set the CUDA device heap limit");
 }
 
 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
@@ -66,6 +68,7 @@ AdePTConfigurationMessenger::~AdePTConfigurationMessenger()
 {
   delete fDir;
   delete fSetCUDAStackLimitCmd;
+  delete fSetCUDAHeapLimitCmd;
   delete fSetTrackInAllRegionsCmd;
   delete fAddRegionCmd;
   delete fActivateAdePTCmd;
@@ -101,6 +104,8 @@ void AdePTConfigurationMessenger::SetNewValue(G4UIcommand *command, G4String new
     fAdePTConfiguration->SetVecGeomGDML(newValue);
   } else if (command == fSetCUDAStackLimitCmd) {
     fAdePTConfiguration->SetCUDAStackLimit(fSetCUDAStackLimitCmd->GetNewIntValue(newValue));
+  } else if (command == fSetCUDAHeapLimitCmd) {
+    fAdePTConfiguration->SetCUDAHeapLimit(fSetCUDAHeapLimitCmd->GetNewIntValue(newValue));
   }
 }