Keep some Track members in registers

bernhardmgruber · bernhardmgruber · commit c08531a75ec2 · 2022-10-31T19:11:40.000+01:00
The track variables are loaded at the beginning of the kernel and
written back when the track survives.
The signature of InitAsSecondary() had to be changed to not load from
the track again.
diff --git a/examples/Example19/electrons.cu b/examples/Example19/electrons.cu
@@ -47,13 +47,21 @@ static __device__ __forceinline__ void TransportElectrons(Track *electrons, cons
   for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < activeSize; i += blockDim.x * gridDim.x) {
     const int globalSlot = (*active)[i];
     Track &currentTrack  = electrons[globalSlot];
-    const auto volume    = currentTrack.navState.Top();
+    auto energy          = currentTrack.energy;
+    auto pos             = currentTrack.pos;
+    auto dir             = currentTrack.dir;
+    auto navState        = currentTrack.navState;
+    const auto volume    = navState.Top();
     const int volumeID   = volume->id();
     // the MCC vector is indexed by the logical volume id
     const int lvolID     = volume->GetLogicalVolume()->id();
     const int theMCIndex = MCIndex[lvolID];
 
     auto survive = [&](bool push = true) {
+      currentTrack.energy   = energy;
+      currentTrack.pos      = pos;
+      currentTrack.dir      = dir;
+      currentTrack.navState = navState;
       if (push) activeQueue->push_back(globalSlot);
     };
 
@@ -63,9 +71,9 @@ static __device__ __forceinline__ void TransportElectrons(Track *electrons, cons
     // Init a track with the needed data to call into G4HepEm.
     G4HepEmElectronTrack elTrack;
     G4HepEmTrack *theTrack = elTrack.GetTrack();
-    theTrack->SetEKin(currentTrack.energy);
+    theTrack->SetEKin(energy);
     theTrack->SetMCIndex(theMCIndex);
-    theTrack->SetOnBoundary(currentTrack.navState.IsOnBoundary());
+    theTrack->SetOnBoundary(navState.IsOnBoundary());
     theTrack->SetCharge(Charge);
     G4HepEmMSCTrackData *mscData = elTrack.GetMSCTrackData();
     mscData->fIsFirstStep        = currentTrack.initialRange < 0;
@@ -80,8 +88,8 @@ static __device__ __forceinline__ void TransportElectrons(Track *electrons, cons
 
     // Compute safety, needed for MSC step limit.
     double safety = 0;
-    if (!currentTrack.navState.IsOnBoundary()) {
-      safety = BVHNavigator::ComputeSafety(currentTrack.pos, currentTrack.navState);
+    if (!navState.IsOnBoundary()) {
+      safety = BVHNavigator::ComputeSafety(pos, navState);
     }
     theTrack->SetSafety(safety);
 
@@ -102,9 +110,9 @@ static __device__ __forceinline__ void TransportElectrons(Track *electrons, cons
 
     bool restrictedPhysicalStepLength = false;
     if (BzFieldValue != 0) {
-      const double momentumMag = sqrt(currentTrack.energy * (currentTrack.energy + 2.0 * Mass));
+      const double momentumMag = sqrt(energy * (energy + 2.0 * Mass));
       // Distance along the track direction to reach the maximum allowed error
-      const double safeLength = fieldPropagatorBz.ComputeSafeLength(momentumMag, Charge, currentTrack.dir);
+      const double safeLength = fieldPropagatorBz.ComputeSafeLength(momentumMag, Charge, dir);
 
       constexpr int MaxSafeLength = 10;
       double limit                = MaxSafeLength * safeLength;
@@ -145,22 +153,20 @@ static __device__ __forceinline__ void TransportElectrons(Track *electrons, cons
     vecgeom::NavStateIndex nextState;
     if (BzFieldValue != 0) {
       geometryStepLength = fieldPropagatorBz.ComputeStepAndNextVolume<BVHNavigator>(
-          currentTrack.energy, Mass, Charge, geometricalStepLengthFromPhysics, currentTrack.pos, currentTrack.dir,
-          currentTrack.navState, nextState, propagated, safety);
+          energy, Mass, Charge, geometricalStepLengthFromPhysics, pos, dir, navState, nextState, propagated, safety);
     } else {
-      geometryStepLength =
-          BVHNavigator::ComputeStepAndNextVolume(currentTrack.pos, currentTrack.dir, geometricalStepLengthFromPhysics,
-                                                 currentTrack.navState, nextState, kPush);
-      currentTrack.pos += geometryStepLength * currentTrack.dir;
+      geometryStepLength = BVHNavigator::ComputeStepAndNextVolume(pos, dir, geometricalStepLengthFromPhysics, navState,
+                                                                  nextState, kPush);
+      pos += geometryStepLength * dir;
     }
 
     // Set boundary state in navState so the next step and secondaries get the
-    // correct information (currentTrack.navState = nextState only if relocated
+    // correct information (navState = nextState only if relocated
     // in case of a boundary; see below)
-    currentTrack.navState.SetBoundaryState(nextState.IsOnBoundary());
+    navState.SetBoundaryState(nextState.IsOnBoundary());
 
     // Propagate information from geometrical step to MSC.
-    theTrack->SetDirection(currentTrack.dir.x(), currentTrack.dir.y(), currentTrack.dir.z());
+    theTrack->SetDirection(dir.x(), dir.y(), dir.z());
     theTrack->SetGStepLength(geometryStepLength);
     theTrack->SetOnBoundary(nextState.IsOnBoundary());
 
@@ -169,7 +175,7 @@ static __device__ __forceinline__ void TransportElectrons(Track *electrons, cons
 
     // Collect the direction change and displacement by MSC.
     const double *direction = theTrack->GetDirection();
-    currentTrack.dir.Set(direction[0], direction[1], direction[2]);
+    dir.Set(direction[0], direction[1], direction[2]);
     if (!nextState.IsOnBoundary()) {
       const double *mscDisplacement = mscData->GetDisplacement();
       vecgeom::Vector3D<Precision> displacement(mscDisplacement[0], mscDisplacement[1], mscDisplacement[2]);
@@ -186,18 +192,18 @@ static __device__ __forceinline__ void TransportElectrons(Track *electrons, cons
         // Apply displacement, depending on how close we are to a boundary.
         // 1a. Far away from geometry boundary:
         if (reducedSafety > 0.0 && dispR <= reducedSafety) {
-          currentTrack.pos += displacement;
+          pos += displacement;
         } else {
           // Recompute safety.
-          safety        = BVHNavigator::ComputeSafety(currentTrack.pos, currentTrack.navState);
+          safety        = BVHNavigator::ComputeSafety(pos, navState);
           reducedSafety = sFact * safety;
 
           // 1b. Far away from geometry boundary:
           if (reducedSafety > 0.0 && dispR <= reducedSafety) {
-            currentTrack.pos += displacement;
+            pos += displacement;
             // 2. Push to boundary:
           } else if (reducedSafety > kGeomMinLength) {
-            currentTrack.pos += displacement * (reducedSafety / dispR);
+            pos += displacement * (reducedSafety / dispR);
           }
           // 3. Very small safety: do nothing.
         }
@@ -209,7 +215,7 @@ static __device__ __forceinline__ void TransportElectrons(Track *electrons, cons
     atomicAdd(&scoringPerVolume->chargedTrackLength[volumeID], elTrack.GetPStepLength());
 
     // Collect the changes in energy and deposit.
-    currentTrack.energy  = theTrack->GetEKin();
+    energy               = theTrack->GetEKin();
     double energyDeposit = theTrack->GetEnergyDeposit();
     atomicAdd(&globalScoring->energyDeposit, energyDeposit);
     atomicAdd(&scoringPerVolume->energyDeposit[volumeID], energyDeposit);
@@ -234,13 +240,13 @@ static __device__ __forceinline__ void TransportElectrons(Track *electrons, cons
         double sinPhi, cosPhi;
         sincos(phi, &sinPhi, &cosPhi);
 
-        gamma1.InitAsSecondary(/*parent=*/currentTrack);
+        gamma1.InitAsSecondary(pos, navState);
         newRNG.Advance();
         gamma1.rngState = newRNG;
         gamma1.energy   = copcore::units::kElectronMassC2;
         gamma1.dir.Set(sint * cosPhi, sint * sinPhi, cost);
 
-        gamma2.InitAsSecondary(/*parent=*/currentTrack);
+        gamma2.InitAsSecondary(pos, navState);
         // Reuse the RNG state of the dying track.
         gamma2.rngState = currentTrack.rngState;
         gamma2.energy   = copcore::units::kElectronMassC2;
@@ -256,10 +262,10 @@ static __device__ __forceinline__ void TransportElectrons(Track *electrons, cons
 
       // Kill the particle if it left the world.
       if (nextState.Top() != nullptr) {
-        BVHNavigator::RelocateToNextVolume(currentTrack.pos, currentTrack.dir, nextState);
+        BVHNavigator::RelocateToNextVolume(pos, dir, nextState);
 
         // Move to the next boundary.
-        currentTrack.navState = nextState;
+        navState = nextState;
         survive();
       }
       continue;
@@ -313,14 +319,21 @@ __device__ void ElectronInteraction(int const globalSlot, SOAData const & /*soaD
                                     GlobalScoring *globalScoring, ScoringPerVolume *scoringPerVolume)
 {
   Track &currentTrack = particles[globalSlot];
-  const auto volume   = currentTrack.navState.Top();
+  auto energy         = currentTrack.energy;
+  const auto pos      = currentTrack.pos;
+  auto dir            = currentTrack.dir;
+  const auto navState = currentTrack.navState;
+  const auto volume   = navState.Top();
   // the MCC vector is indexed by the logical volume id
   const int lvolID     = volume->GetLogicalVolume()->id();
   const int theMCIndex = MCIndex[lvolID];
 
-  auto survive = [&] { activeQueue->push_back(globalSlot); };
+  auto survive = [&] {
+    currentTrack.dir    = dir;
+    currentTrack.energy = energy;
+    activeQueue->push_back(globalSlot);
+  };
 
-  const double energy   = currentTrack.energy;
   const double theElCut = g4HepEmData.fTheMatCutData->fMatCutData[theMCIndex].fSecElProdCutE;
 
   RanluxppDouble newRNG{currentTrack.rngState.Branch()};
@@ -331,20 +344,20 @@ __device__ void ElectronInteraction(int const globalSlot, SOAData const & /*soaD
     double deltaEkin = (IsElectron) ? G4HepEmElectronInteractionIoni::SampleETransferMoller(theElCut, energy, &rnge)
                                     : G4HepEmElectronInteractionIoni::SampleETransferBhabha(theElCut, energy, &rnge);
 
-    double dirPrimary[] = {currentTrack.dir.x(), currentTrack.dir.y(), currentTrack.dir.z()};
+    double dirPrimary[] = {dir.x(), dir.y(), dir.z()};
     double dirSecondary[3];
     G4HepEmElectronInteractionIoni::SampleDirections(energy, deltaEkin, dirSecondary, dirPrimary, &rnge);
 
     Track &secondary = secondaries.electrons.NextTrack();
     atomicAdd(&globalScoring->numElectrons, 1);
 
-    secondary.InitAsSecondary(/*parent=*/currentTrack);
+    secondary.InitAsSecondary(pos, navState);
     secondary.rngState = newRNG;
     secondary.energy   = deltaEkin;
     secondary.dir.Set(dirSecondary[0], dirSecondary[1], dirSecondary[2]);
 
-    currentTrack.energy = energy - deltaEkin;
-    currentTrack.dir.Set(dirPrimary[0], dirPrimary[1], dirPrimary[2]);
+    energy -= deltaEkin;
+    dir.Set(dirPrimary[0], dirPrimary[1], dirPrimary[2]);
     survive();
   } else if constexpr (ProcessIndex == 1) {
     // Invoke model for Bremsstrahlung: either SB- or Rel-Brem.
@@ -355,24 +368,24 @@ __device__ void ElectronInteraction(int const globalSlot, SOAData const & /*soaD
                            : G4HepEmElectronInteractionBrem::SampleETransferRB(&g4HepEmData, energy, logEnergy,
                                                                                theMCIndex, &rnge, IsElectron);
 
-    double dirPrimary[] = {currentTrack.dir.x(), currentTrack.dir.y(), currentTrack.dir.z()};
+    double dirPrimary[] = {dir.x(), dir.y(), dir.z()};
     double dirSecondary[3];
     G4HepEmElectronInteractionBrem::SampleDirections(energy, deltaEkin, dirSecondary, dirPrimary, &rnge);
 
     Track &gamma = secondaries.gammas.NextTrack();
     atomicAdd(&globalScoring->numGammas, 1);
 
-    gamma.InitAsSecondary(/*parent=*/currentTrack);
+    gamma.InitAsSecondary(pos, navState);
     gamma.rngState = newRNG;
     gamma.energy   = deltaEkin;
     gamma.dir.Set(dirSecondary[0], dirSecondary[1], dirSecondary[2]);
 
-    currentTrack.energy = energy - deltaEkin;
-    currentTrack.dir.Set(dirPrimary[0], dirPrimary[1], dirPrimary[2]);
+    energy -= deltaEkin;
+    dir.Set(dirPrimary[0], dirPrimary[1], dirPrimary[2]);
     survive();
   } else if constexpr (ProcessIndex == 2) {
     // Invoke annihilation (in-flight) for e+
-    double dirPrimary[] = {currentTrack.dir.x(), currentTrack.dir.y(), currentTrack.dir.z()};
+    double dirPrimary[] = {dir.x(), dir.y(), dir.z()};
     double theGamma1Ekin, theGamma2Ekin;
     double theGamma1Dir[3], theGamma2Dir[3];
     G4HepEmPositronInteractionAnnihilation::SampleEnergyAndDirectionsInFlight(
@@ -382,12 +395,12 @@ __device__ void ElectronInteraction(int const globalSlot, SOAData const & /*soaD
     Track &gamma2 = secondaries.gammas.NextTrack();
     atomicAdd(&globalScoring->numGammas, 2);
 
-    gamma1.InitAsSecondary(/*parent=*/currentTrack);
+    gamma1.InitAsSecondary(pos, navState);
     gamma1.rngState = newRNG;
     gamma1.energy   = theGamma1Ekin;
     gamma1.dir.Set(theGamma1Dir[0], theGamma1Dir[1], theGamma1Dir[2]);
 
-    gamma2.InitAsSecondary(/*parent=*/currentTrack);
+    gamma2.InitAsSecondary(pos, navState);
     // Reuse the RNG state of the dying track.
     gamma2.rngState = currentTrack.rngState;
     gamma2.energy   = theGamma2Ekin;
diff --git a/examples/Example19/example.cuh b/examples/Example19/example.cuh
@@ -36,7 +36,8 @@ struct Track {
 
   __host__ __device__ double Uniform() { return rngState.Rndm(); }
 
-  __host__ __device__ void InitAsSecondary(const Track &parent)
+  __host__ __device__ void InitAsSecondary(const vecgeom::Vector3D<Precision> &parentPos,
+                                           const vecgeom::NavStateIndex &parentNavState)
   {
     // The caller is responsible to branch a new RNG state and to set the energy.
     this->numIALeft[0] = -1.0;
@@ -49,8 +50,8 @@ struct Track {
 
     // A secondary inherits the position of its parent; the caller is responsible
     // to update the directions.
-    this->pos      = parent.pos;
-    this->navState = parent.navState;
+    this->pos      = parentPos;
+    this->navState = parentNavState;
   }
 };
 
diff --git a/examples/Example19/gammas.cu b/examples/Example19/gammas.cu
