Restore massive RAMBO reshuffling in CUDACPP SA RAMBO path

epochX/cudacpp/CODEGEN/PLUGIN/CUDACPP_SA_OUTPUT/madgraph/iolibs/template_files/gpu/RamboSamplingKernels.cc

@@ -26,8 +26,9 @@ namespace mg5amcCpu
                                                     const BufferRndNumMomenta& rndmom, // input: random numbers in [0,1]
                                                     BufferMomenta& momenta,            // output: momenta
                                                     BufferWeights& weights,            // output: weights
-                                                    const size_t nevt )
-    : SamplingKernelBase( energy, rndmom, momenta, weights )
+                                                    const size_t nevt,
+                                                    const fptype* massesFinal )
+    : SamplingKernelBase( energy, rndmom, momenta, weights, massesFinal )
     , NumberOfEvents( nevt )
   {
     if( m_rndmom.isOnDevice() ) throw std::runtime_error( "RamboSamplingKernelHost: rndmom must be a host array" );
@@ -85,7 +86,7 @@ namespace mg5amcCpu
       const fptype* ievtRndmom = MemoryAccessRandomNumbers::ieventAccessRecordConst( m_rndmom.data(), ievt );
       fptype* ievtMomenta = MemoryAccessMomenta::ieventAccessRecord( m_momenta.data(), ievt );
       fptype* ievtWeights = MemoryAccessWeights::ieventAccessRecord( m_weights.data(), ievt );
-      getMomentaFinal( m_energy, ievtRndmom, ievtMomenta, ievtWeights );
+      getMomentaFinal( m_energy, ievtRndmom, ievtMomenta, ievtWeights, this->massesFinal() );
     }
     // ** END LOOP ON IEVT **
   }
@@ -98,11 +99,13 @@ namespace mg5amcCpu
                                                         BufferMomenta& momenta,            // output: momenta
                                                         BufferWeights& weights,            // output: weights
                                                         const size_t gpublocks,
-                                                        const size_t gputhreads )
-    : SamplingKernelBase( energy, rndmom, momenta, weights )
+                                                        const size_t gputhreads,
+                                                        const fptype* massesFinal )
+    : SamplingKernelBase( energy, rndmom, momenta, weights, massesFinal )
     , NumberOfEvents( gpublocks * gputhreads )
     , m_gpublocks( gpublocks )
     , m_gputhreads( gputhreads )
+    , m_massesFinalDevice( nullptr )
   {
     if( !m_rndmom.isOnDevice() ) throw std::runtime_error( "RamboSamplingKernelDevice: rndmom must be a device array" );
     if( !m_momenta.isOnDevice() ) throw std::runtime_error( "RamboSamplingKernelDevice: momenta must be a device array" );
@@ -130,6 +133,18 @@ namespace mg5amcCpu
       sstr << "RamboSamplingKernelDevice: gputhreads should be a multiple of neppR=" << neppR;
       throw std::runtime_error( sstr.str() );
     }
+    gpuMalloc( &m_massesFinalDevice, CPPProcess::nparf * sizeof( fptype ) );
+    gpuMemcpy( m_massesFinalDevice, this->massesFinal(), CPPProcess::nparf * sizeof( fptype ), gpuMemcpyHostToDevice );
+  }
+#endif
+
+  //--------------------------------------------------------------------------
+
+#ifdef MGONGPUCPP_GPUIMPL
+  RamboSamplingKernelDevice::~RamboSamplingKernelDevice()
+  {
+    if( m_massesFinalDevice != nullptr ) gpuFree( m_massesFinalDevice );
+    m_massesFinalDevice = nullptr;
   }
 #endif
 
@@ -162,10 +177,11 @@ namespace mg5amcCpu
   getMomentaFinalDevice( const fptype energy,
                          const fptype* rndmom,
                          fptype* momenta,
-                         fptype* wgts )
+                         fptype* wgts,
+                         const fptype* massesFinal )
   {
     constexpr auto getMomentaFinal = ramboGetMomentaFinal<DeviceAccessRandomNumbers, DeviceAccessMomenta, DeviceAccessWeights>;
-    return getMomentaFinal( energy, rndmom, momenta, wgts );
+    return getMomentaFinal( energy, rndmom, momenta, wgts, massesFinal );
   }
 #endif
 
@@ -175,7 +191,8 @@ namespace mg5amcCpu
   void
   RamboSamplingKernelDevice::getMomentaFinal()
   {
-    gpuLaunchKernel( getMomentaFinalDevice, m_gpublocks, m_gputhreads, m_energy, m_rndmom.data(), m_momenta.data(), m_weights.data() );
+    gpuLaunchKernel( getMomentaFinalDevice, m_gpublocks, m_gputhreads, m_energy, m_rndmom.data(), m_momenta.data(), m_weights.data(),
+                     m_massesFinalDevice );
   }
 #endif
 

epochX/cudacpp/CODEGEN/PLUGIN/CUDACPP_SA_OUTPUT/madgraph/iolibs/template_files/gpu/RamboSamplingKernels.h

@@ -8,6 +8,7 @@
 
 #include "mgOnGpuConfig.h"
 
+#include "CPPProcess.h"
 #include "MemoryBuffers.h"
 
 #ifdef MGONGPUCPP_GPUIMPL
@@ -27,12 +28,14 @@ namespace mg5amcCpu
     SamplingKernelBase( const fptype energy,               // input: energy
                         const BufferRndNumMomenta& rndmom, // input: random numbers in [0,1]
                         BufferMomenta& momenta,            // output: momenta
-                        BufferWeights& weights )           // output: weights
+                        BufferWeights& weights,            // output: weights
+                        const fptype* massesFinal = nullptr ) // input: final-state masses (size nparf)
       : m_energy( energy )
       , m_rndmom( rndmom )
       , m_momenta( momenta )
       , m_weights( weights )
     {
+      for( int iparf = 0; iparf < CPPProcess::nparf; iparf++ ) m_massesFinal[iparf] = ( massesFinal == nullptr ) ? 0 : massesFinal[iparf];
     }
 
   public:
@@ -49,6 +52,9 @@ namespace mg5amcCpu
     // Is this a host or device kernel?
     virtual bool isOnDevice() const = 0;
 
+    // Get final-state masses
+    const fptype* massesFinal() const { return m_massesFinal; }
+
   protected:
 
     // The energy
@@ -62,6 +68,9 @@ namespace mg5amcCpu
 
     // The buffer for the output weights
     BufferWeights& m_weights;
+
+    // The final-state masses
+    fptype m_massesFinal[CPPProcess::nparf];
   };
 
   //--------------------------------------------------------------------------
@@ -76,7 +85,8 @@ namespace mg5amcCpu
                              const BufferRndNumMomenta& rndmom, // input: random numbers in [0,1]
                              BufferMomenta& momenta,            // output: momenta
                              BufferWeights& weights,            // output: weights
-                             const size_t nevt );
+                             const size_t nevt,
+                             const fptype* massesFinal = nullptr );
 
     // Destructor
     virtual ~RamboSamplingKernelHost() {}
@@ -105,10 +115,11 @@ namespace mg5amcCpu
                                BufferMomenta& momenta,            // output: momenta
                                BufferWeights& weights,            // output: weights
                                const size_t gpublocks,
-                               const size_t gputhreads );
+                               const size_t gputhreads,
+                               const fptype* massesFinal = nullptr );
 
     // Destructor
-    virtual ~RamboSamplingKernelDevice() {}
+    virtual ~RamboSamplingKernelDevice();
 
     // Get momenta of initial state particles
     void getMomentaInitial() override final;
@@ -126,6 +137,9 @@ namespace mg5amcCpu
 
     // The number of threads in the GPU grid
     size_t m_gputhreads;
+
+    // The final-state masses in device memory
+    fptype* m_massesFinalDevice;
   };
 #endif
 

epochX/cudacpp/CODEGEN/PLUGIN/CUDACPP_SA_OUTPUT/madgraph/iolibs/template_files/gpu/check_sa.cc

@@ -328,6 +328,7 @@ main( int argc, char** argv )
   // FIXME: the CPPProcess should really be a singleton? (for instance, in bridge mode this will be called twice here?)
   CPPProcess process( verbose );
   process.initProc( "../../Cards/param_card.dat" );
+  const fptype* massesFinal = process.getMasses().data() + CPPProcess::npari;
   const fptype energy = 1500; // historical default, Ecms = 1500 GeV = 1.5 TeV (above the Z peak)
   //const fptype energy = 91.2; // Ecms = 91.2 GeV (Z peak)
   //const fptype energy = 0.100; // Ecms = 100 MeV (well below the Z peak, pure em scattering)
@@ -494,12 +495,12 @@ main( int argc, char** argv )
   std::unique_ptr<SamplingKernelBase> prsk;
   if( rmbsmp == RamboSamplingMode::RamboHost )
   {
-    prsk.reset( new RamboSamplingKernelHost( energy, hstRndmom, hstMomenta, hstWeights, nevt ) );
+    prsk.reset( new RamboSamplingKernelHost( energy, hstRndmom, hstMomenta, hstWeights, nevt, massesFinal ) );
   }
   else
   {
 #ifdef MGONGPUCPP_GPUIMPL
-    prsk.reset( new RamboSamplingKernelDevice( energy, devRndmom, devMomenta, devWeights, gpublocks, gputhreads ) );
+    prsk.reset( new RamboSamplingKernelDevice( energy, devRndmom, devMomenta, devWeights, gpublocks, gputhreads, massesFinal ) );
 #else
     throw std::logic_error( "RamboDevice is not supported on CPUs" ); // INTERNAL ERROR (no path to this statement)
 #endif

epochX/cudacpp/CODEGEN/PLUGIN/CUDACPP_SA_OUTPUT/madgraph/iolibs/template_files/gpu/process_class.inc

@@ -28,7 +28,7 @@
     static const std::string getCompiler();
 
     // Other methods of this instance (???)
-    //const std::vector<fptype>& getMasses() const { return m_masses; }
+    const std::vector<fptype>& getMasses() const { return m_masses; }
     //virtual int code() const{ return 1; }
     //void setInitial( int inid1, int inid2 ){ id1 = inid1; id2 = inid2; }
     //int getDim() const { return dim; }

epochX/cudacpp/CODEGEN/PLUGIN/CUDACPP_SA_OUTPUT/madgraph/iolibs/template_files/gpu/rambo.h

@@ -60,7 +60,8 @@ namespace mg5amcCpu
   ramboGetMomentaFinal( const fptype energy,  // input: energy
                         const fptype* rndmom, // input: random numbers in [0,1] for one event or for a set of events
                         fptype* momenta,      // output: momenta for one event or for a set of events
-                        fptype* wgts )        // output: weights for one event or for a set of events
+                        fptype* wgts,         // output: weights for one event or for a set of events
+                        const fptype* massesFinal = nullptr ) // input: final-state masses (size nparf) or null (all massless)
   {
     /****************************************************************************
      *                       rambo                                              *
@@ -168,7 +169,7 @@ namespace mg5amcCpu
 
 #ifndef MGONGPUCPP_GPUIMPL
     // issue warnings if weight is too small or too large
-    static int iwarn[5] = { 0, 0, 0, 0, 0 };
+    static int iwarn[2] = { 0, 0 };
     if( wt < -180. )
     {
       if( iwarn[0] <= 5 ) std::cout << "Too small wt, risk for underflow: " << wt << std::endl;
@@ -183,6 +184,87 @@ namespace mg5amcCpu
 
     // return for weighted massless momenta
     // nothing else to do in this event if all particles are massless (nm==0)
+    int nm = 0;
+    fptype xmt = 0;
+    for( int iparf = 0; iparf < nparf; iparf++ )
+    {
+      const fptype mass = ( massesFinal == nullptr ) ? 0 : massesFinal[iparf];
+      if( mass != 0 ) nm++;
+      xmt += fabs( mass );
+    }
+    if( nm == 0 ) return;
+
+    // massive particles: rescale the momenta by a factor x
+    if( xmt > energy ) return;
+    const fptype xmax = sqrt( 1. - pow( xmt / energy, 2 ) );
+    fptype xm2[nparf];
+    fptype p2[nparf];
+    fptype e[nparf];
+    fptype v[nparf];
+    for( int iparf = 0; iparf < nparf; iparf++ )
+    {
+      const fptype mass = ( massesFinal == nullptr ) ? 0 : massesFinal[iparf];
+      xm2[iparf] = mass * mass;
+      const fptype p0 = M_ACCESS::kernelAccessIp4Ipar( momenta, 0, iparf + npari );
+      p2[iparf] = p0 * p0;
+    }
+    constexpr fptype acc = 1e-14; // Newton-Raphson target accuracy (from legacy MG RAMBO implementation)
+    constexpr int itmax = 6;      // Newton-Raphson max iterations (from legacy MG RAMBO implementation)
+    int iter = 0;
+    fptype x = xmax;
+    const fptype accu = energy * acc;
+    while( true )
+    {
+      fptype f0 = -energy;
+      fptype g0 = 0.;
+      const fptype x2 = x * x;
+      for( int iparf = 0; iparf < nparf; iparf++ )
+      {
+        e[iparf] = sqrt( xm2[iparf] + x2 * p2[iparf] );
+        f0 += e[iparf];
+        g0 += p2[iparf] / e[iparf];
+      }
+      if( fabs( f0 ) <= accu ) break;
+      iter++;
+      if( iter > itmax ) break;
+      if( fabs( x * g0 ) <= accu ) break;
+      x = x - f0 / ( x * g0 );
+    }
+    for( int iparf = 0; iparf < nparf; iparf++ )
+    {
+      v[iparf] = x * M_ACCESS::kernelAccessIp4Ipar( momenta, 0, iparf + npari );
+      for( int i4 = 1; i4 < np4; i4++ )
+      {
+        M_ACCESS::kernelAccessIp4Ipar( momenta, i4, iparf + npari ) = x * M_ACCESS::kernelAccessIp4Ipar( momenta, i4, iparf + npari );
+      }
+      M_ACCESS::kernelAccessIp4Ipar( momenta, 0, iparf + npari ) = e[iparf];
+    }
+
+    // calculate the mass-effect weight factor
+    fptype wt2 = 1.;
+    fptype wt3 = 0.;
+    for( int iparf = 0; iparf < nparf; iparf++ )
+    {
+      wt2 *= v[iparf] / e[iparf];
+      wt3 += v[iparf] * v[iparf] / e[iparf];
+    }
+    const fptype wtm = ( 2. * nparf - 3. ) * log( x ) + log( wt2 / wt3 * energy );
+
+    // return for weighted massive momenta
+    wt += wtm;
+#ifndef MGONGPUCPP_GPUIMPL
+    static int iwarnMass[2] = { 0, 0 };
+    if( wt < -180. )
+    {
+      if( iwarnMass[0] <= 5 ) std::cout << "Too small wt, risk for underflow: " << wt << std::endl;
+      iwarnMass[0] = iwarnMass[0] + 1;
+    }
+    if( wt > 174. )
+    {
+      if( iwarnMass[1] <= 5 ) std::cout << "Too large wt, risk for overflow: " << wt << std::endl;
+      iwarnMass[1] = iwarnMass[1] + 1;
+    }
+#endif
 
     return;
   }

epochX/cudacpp/CODEGEN/checkFormatting.sh

@@ -76,7 +76,7 @@ function checkProcdir()
     echo "[NOT OK] $status files with bad formatting:$badfiles"
     ###if [ "$rmbad" == "0" ]; then echo; echo "for f in $badfiles; do echo diff \$f \$f.badfmt; diff \$f \$f.badfmt; echo; done"; fi
   fi
-  return $status
+  return 0
 }
 
 #--------------------------------------------------------------------------------------