From fb197f15ea535453c21dfdebe7ef6d86682abd56 Mon Sep 17 00:00:00 2001 From: Daniele Massaro Date: Thu, 9 Jul 2026 19:41:11 +0200 Subject: [PATCH] Add atomicAdd for numerators, and remove denominators from kernel Adding atomic add in the main kernel allow to sum over the helicity, reducing the dimensionality of the numerators array. Denominators are calculated starting from numerators. --- .../madmatrix/MatrixElementKernels.cc | 9 +++--- .../process_function_definitions.inc | 28 +++++++---------- .../madmatrix/process_matrix.inc | 1 + .../madmatrix/process_sigmaKin_function.inc | 30 ++++++++++++++----- .../iolibs/template_files/madmatrix/umami.cc | 6 ++-- madmatrix/model_handling.py | 27 ++++++++++------- 6 files changed, 60 insertions(+), 41 deletions(-) diff --git a/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc b/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc index 872e4795e..a8e6a59d1 100644 --- a/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc +++ b/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc @@ -459,10 +459,11 @@ namespace mg5amcGpu // ... Create the "many-helicity" super-buffer of nGoodHel ME buffers (dynamically allocated because nGoodHel is determined at runtime) // ... (calling reset here deletes the previously created "one-helicity" buffers used for helicity filtering) m_pHelJamps.reset( new DeviceBufferSimple( nGoodHel * CPPProcess::ncolor * mgOnGpu::nx2 * nevt ) ); - // ... Create the "many-helicity" super-buffers of nGoodHel numerator and denominator buffers (dynamically allocated) - // ... (calling reset here deletes the previously created "one-helicity" buffers used for helicity filtering) - m_pHelNumerators.reset( new DeviceBufferSimple( nGoodHel * CPPProcess::ndiagrams * nevt ) ); - m_pHelDenominators.reset( new DeviceBufferSimple( nGoodHel * nevt ) ); + // ... Create the numerator and denominator buffers. These no longer carry a helicity dimension: + // ... the numerators are accumulated in place over all good helicities via atomicAdd in calculate_jamps + // ... ([nevt][ndiagrams]) and the denominators are derived from them ([nevt]). + m_pHelNumerators.reset( new DeviceBufferSimple( CPPProcess::ndiagrams * nevt ) ); + m_pHelDenominators.reset( new DeviceBufferSimple( nevt ) ); #ifndef MGONGPU_HAS_NO_BLAS // Create the "many-helicity" super-buffers of real/imag ncolor*nevt temporary buffers for cuBLAS/hipBLAS intermediate results in color_sum_blas #if defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT diff --git a/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc b/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc index 082c373aa..ee2923736 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc @@ -693,8 +693,8 @@ namespace mg5amcCpu __global__ void normalise_output( fptype* allMEs, // output: allMEs[nevt], |M|^2 running_sum_over_helicities const unsigned int* iflavorVec, - fptype* ghelAllNumerators, // input/tmp: allNumerators super-buffer for nGoodHel <= ncomb individual helicities (index is ighel) - fptype* ghelAllDenominators, // input/tmp: allNumerators super-buffer for nGoodHel <= ncomb individual helicities (index is ighel) + fptype* allNumerators, // input: multichannel numerators[nevt][ndiagrams], already summed over helicities (atomicAdd) + fptype* allDenominators, // output: multichannel denominators[nevt], derived here as the sum of numerators const unsigned int* allChannelIds, // input: multichannel channelIds[nevt] (1 to #diagrams); nullptr to disable SDE enhancement (fix #899/#911) bool storeChannelWeights, // if true, compute final multichannel weights bool mulChannelWeight, // if true, multiply matrix element by channel weight @@ -702,26 +702,20 @@ namespace mg5amcCpu { const int ievt = blockDim.x * blockIdx.x + threadIdx.x; // index of event (thread) allMEs[ievt] = allMEs[ievt] * broken_symmetry_factor(iflavorVec[ievt]) / globaldenom; - const int nevt = gridDim.x * blockDim.x; if( storeChannelWeights ) // fix segfault #892 (not 'channelIds[0] != 0') { - fptype* totAllNumerators = ghelAllNumerators; // reuse "helicity #0" buffer to compute the total over all helicities - fptype* totAllDenominators = ghelAllDenominators; // reuse "helicity #0" buffer to compute the total over all helicities - for( int ighel = 1; ighel < dcNGoodHel; ighel++ ) // NB: the loop starts at ighel=1 - { - fptype* hAllDenominators = ghelAllDenominators + ighel * nevt; - totAllDenominators[ievt] += hAllDenominators[ievt]; - fptype* hAllNumerators = ghelAllNumerators + ( ievt + ighel * nevt ) * processConfig::ndiagrams; - fptype* firstNumerator = ghelAllNumerators + ievt * processConfig::ndiagrams; - for( int idiag = 0; idiag < processConfig::ndiagrams; ++idiag ) - { - firstNumerator[idiag] += hAllNumerators[idiag]; - } - } + // The numerators have already been accumulated over all good helicities in place (atomicAdd in + // calculate_jamps), so there is no helicity dimension to sum here. The denominator is just the + // sum of all numerators for this event: derive it once and store it for the downstream consumers. + fptype* numerators = allNumerators + ievt * processConfig::ndiagrams; + fptype denominator = 0; + for( int idiag = 0; idiag < processConfig::ndiagrams; ++idiag ) + denominator += numerators[idiag]; + allDenominators[ievt] = denominator; if( mulChannelWeight ) { unsigned int channelId = allChannelIds[ievt]; - allMEs[ievt] *= totAllNumerators[channelId - 1 + ievt * processConfig::ndiagrams] / totAllDenominators[ievt]; + allMEs[ievt] *= numerators[channelId - 1] / denominator; } } return; diff --git a/madgraph/iolibs/template_files/madmatrix/process_matrix.inc b/madgraph/iolibs/template_files/madmatrix/process_matrix.inc index dff402fa3..9cf5d3a25 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_matrix.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_matrix.inc @@ -47,3 +47,4 @@ mgDebug( 1, __FUNCTION__ ); return; } +#undef NUM_ATOMIC_ADD diff --git a/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc b/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc index 227301a6e..aa64d8274 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc @@ -17,8 +17,11 @@ gpuMemset( allMEs, 0, nevt * sizeof( fptype ) ); gpuMemset( ghelAllJamps, 0, cNGoodHel * ncolor * mgOnGpu::nx2 * nevt * sizeof( fptype ) ); gpuMemset( colAllJamp2s, 0, ncolor * nevt * sizeof( fptype ) ); - gpuMemset( ghelAllNumerators, 0, cNGoodHel * processConfig::ndiagrams * nevt * sizeof( fptype ) ); - gpuMemset( ghelAllDenominators, 0, cNGoodHel * nevt * sizeof( fptype ) ); + // The numerators buffer has NO helicity dimension: all good helicities accumulate in place via + // atomicAdd, so it is zeroed once as [nevt][ndiagrams]. The denominators are derived from the + // numerators in normalise_output, so the buffer is just [nevt]. + gpuMemset( ghelAllNumerators, 0, processConfig::ndiagrams * nevt * sizeof( fptype ) ); + gpuMemset( ghelAllDenominators, 0, nevt * sizeof( fptype ) ); gpuMemset( ghelAllMEs, 0, cNGoodHel * nevt * sizeof( fptype ) ); #else // *** PART 0b - C++ *** @@ -60,9 +63,9 @@ { const int ihel = cGoodHel[ighel]; fptype* hAllJamps = ghelAllJamps + ighel * nevt; // HACK: bypass DeviceAccessJamp (consistent with layout defined there) - fptype* hAllNumerators = ghelAllNumerators + ighel * nevt * processConfig::ndiagrams; - fptype* hAllDenominators = ghelAllDenominators + ighel * nevt; - gpuLaunchKernelStream( calculate_jamps, gpublocks, gputhreads, ghelStreams[ighel], ihel, allmomenta, allcouplings, iflavorVec, hAllJamps, storeChannelWeights, hAllNumerators, hAllDenominators, colAllJamp2s, nevt, false ); + // NB: the numerators buffer has no helicity dimension: every helicity stream accumulates in place + // into the same [nevt][ndiagrams] slot via atomicAdd. The denominators are derived later. + gpuLaunchKernelStream( calculate_jamps, gpublocks, gputhreads, ghelStreams[ighel], ihel, allmomenta, allcouplings, iflavorVec, hAllJamps, storeChannelWeights, ghelAllNumerators, ghelAllDenominators, colAllJamp2s, nevt, false ); } // (2) Then compute the ME for that helicity from the color sum of QCD partial amplitudes jamps color_sum_gpu( ghelAllMEs, ghelAllJamps, ghelAllBlasTmp, pBlasHandle, ghelStreams, cNGoodHel, gpublocks, gputhreads, false ); @@ -286,20 +289,31 @@ // [NB 'sum over final spins, average over initial spins', eg see // https://www.uzh.ch/cmsssl/physik/dam/jcr:2e24b7b1-f4d7-4160-817e-47b13dbf1d7c/Handout_4_2016-UZH.pdf] #ifndef MGONGPUCPP_GPUIMPL + const bool storeChannelWeights = allChannelIds != nullptr || allrnddiagram != nullptr; for( int ipagV = 0; ipagV < npagV; ++ipagV ) { const int ievt0 = ipagV * neppV; fptype* MEs = E_ACCESS::ieventAccessRecord( allMEs, ievt0 ); fptype_sv& MEs_sv = E_ACCESS::kernelAccess( MEs ); MEs_sv = MEs_sv * broken_symmetry_factor(iflavorVec[ievt0]) / helcolDenominators[0]; - if( mulChannelWeight && allChannelIds != nullptr ) // fix segfault #892 (not 'channelIds[0] != 0') + if( storeChannelWeights ) // fix segfault #892 (not 'channelIds[0] != 0') { - const unsigned int channelId = getChannelId( allChannelIds, ievt0, false ); + // The numerators have already been accumulated over all good helicities in place (running sum + // over the helicity loop), so there is no helicity dimension to sum here. The denominator is + // just the sum of all numerators for this event page: derive it once and store it for the + // downstream consumers (e.g. the multichannel amp2 output). fptype* numerators = NUM_ACCESS::ieventAccessRecord( allNumerators, ievt0 * processConfig::ndiagrams ); fptype* denominators = DEN_ACCESS::ieventAccessRecord( allDenominators, ievt0 ); fptype_sv* numerators_sv = NUM_ACCESS::kernelAccessP( numerators ); fptype_sv& denominators_sv = DEN_ACCESS::kernelAccess( denominators ); - MEs_sv *= numerators_sv[channelId - 1] / denominators_sv; + denominators_sv = fptype_sv{ 0 }; + for( int idiag = 0; idiag < processConfig::ndiagrams; ++idiag ) + denominators_sv += numerators_sv[idiag]; + if( mulChannelWeight && allChannelIds != nullptr ) + { + const unsigned int channelId = getChannelId( allChannelIds, ievt0, false ); + MEs_sv *= numerators_sv[channelId - 1] / denominators_sv; + } } //for( int ieppV = 0; ieppV < neppV; ieppV++ ) //{ diff --git a/madgraph/iolibs/template_files/madmatrix/umami.cc b/madgraph/iolibs/template_files/madmatrix/umami.cc index d19c93bb9..362de9e16 100644 --- a/madgraph/iolibs/template_files/madmatrix/umami.cc +++ b/madgraph/iolibs/template_files/madmatrix/umami.cc @@ -386,8 +386,10 @@ extern "C" {reinterpret_cast(&matrix_elements), rounded_count * sizeof( fptype )}, {reinterpret_cast(&diagram_index), rounded_count * sizeof( unsigned int )}, {reinterpret_cast(&color_jamps), rounded_count * CPPProcess::ncolor * mgOnGpu::nx2 * sizeof( fptype )}, - {reinterpret_cast(&numerators), rounded_count * CPPProcess::ndiagrams * CPPProcess::ncomb * sizeof( fptype )}, - {reinterpret_cast(&denominators), rounded_count * CPPProcess::ncomb * sizeof( fptype )}, + // The numerators are accumulated in place over all helicities via atomicAdd (no helicity dimension), + // and the denominators are derived from them, so neither buffer carries the ncomb factor anymore. + {reinterpret_cast(&numerators), rounded_count * CPPProcess::ndiagrams * sizeof( fptype )}, + {reinterpret_cast(&denominators), rounded_count * sizeof( fptype )}, {reinterpret_cast(&helicity_index), rounded_count * sizeof( int )}, {reinterpret_cast(&color_index), rounded_count * sizeof( int )}, {reinterpret_cast(&ghel_matrix_elements), rounded_count * CPPProcess::ncomb * sizeof( fptype )}, diff --git a/madmatrix/model_handling.py b/madmatrix/model_handling.py index b4d1bac78..0025767cb 100644 --- a/madmatrix/model_handling.py +++ b/madmatrix/model_handling.py @@ -1805,6 +1805,17 @@ def get_all_sigmaKin_lines(self, color_amplitudes, class_name): // ** NB2: NEW Nov2024! in CUDA this now takes a channelId array as input (it used to take a scalar channelId as input) // In C++, this function processes a single event "page" or SIMD vector (or for two in "mixed" precision mode, nParity=2) // *** NB: in C++, calculate_jamps accepts a SCALAR channelId because it is GUARANTEED that all events in a SIMD vector have the same channelId #898 + + // Accumulate a multichannel numerator contribution in place. + // In CUDA all good-helicity blocks/streams for a given event race on the same numerator slot + // (the helicity dimension has been removed to save memory), so an atomicAdd is mandatory. + // In C++ each event page is processed serially within the helicity loop, so a plain sum suffices. +#ifdef MGONGPUCPP_GPUIMPL +#define NUM_ATOMIC_ADD( DST, VAL ) atomicAdd( &( DST ), VAL ) +#else +#define NUM_ATOMIC_ADD( DST, VAL ) ( DST ) += ( VAL ) +#endif + __global__ void /* clang-format off */ calculate_jamps( int ihel, const fptype* allmomenta, // input: momenta[nevt*npar*4] @@ -1838,7 +1849,6 @@ def get_all_sigmaKin_lines(self, color_amplitudes, class_name): using CI_ACCESS = DeviceAccessCouplingsFixed; // TRIVIAL access (independent couplings): buffer for one event using F_ACCESS = DeviceAccessIflavorVec; // non-trivial access: buffer includes all events using NUM_ACCESS = DeviceAccessNumerators; // non-trivial access: buffer includes all events - using DEN_ACCESS = DeviceAccessDenominators; // non-trivial access: buffer includes all events #else using namespace mg5amcCpu; using M_ACCESS = HostAccessMomenta; // non-trivial access: buffer includes all events @@ -1848,7 +1858,6 @@ def get_all_sigmaKin_lines(self, color_amplitudes, class_name): using CI_ACCESS = HostAccessCouplingsFixed; // TRIVIAL access (independent couplings): buffer for one event using F_ACCESS = HostAccessIflavorVec; // non-trivial access: buffer includes all events using NUM_ACCESS = HostAccessNumerators; // non-trivial access: buffer includes all events - using DEN_ACCESS = HostAccessDenominators; // non-trivial access: buffer includes all events #endif mgDebug( 0, __FUNCTION__ ); //bool debug = true; @@ -1863,8 +1872,9 @@ def get_all_sigmaKin_lines(self, color_amplitudes, class_name): int ighel = blockIdx.y; ihel = dcGoodHel[ighel]; allJamps = allJamps + ighel * nevt; - allNumerators = allNumerators + ighel * nevt * processConfig::ndiagrams; - allDenominators = allDenominators + ighel * nevt; + // NB: the numerators buffer has NO helicity dimension anymore: all good-helicity blocks + // for a given event accumulate in place into the same [nevt][ndiagrams] slot via atomicAdd. + // The denominators are no longer accumulated here (derived as the sum of numerators later). } #endif /* clang-format on */""") nwavefuncs = self.matrix_elements[0].get_number_of_wavefunctions() @@ -2449,7 +2459,6 @@ def super_get_matrix_element_calls(self, matrix_element, color_amplitudes, multi for( size_t ixcoup = 0; ixcoup < nxcoup; ixcoup++ ) COUPs[ixcoup] = allCOUPs[ixcoup]; const int ievt = blockDim.x * blockIdx.x + threadIdx.x; // index of event (thread) in grid fptype* numerators = &allNumerators[ievt * processConfig::ndiagrams]; - fptype* denominators = allDenominators; #else // C++ kernels take input/output buffers with momenta/MEs for one specific event (the first in the current event page) const fptype* momenta = M_ACCESS::ieventAccessRecordConst( allmomenta, ievt0 ); @@ -2460,7 +2469,6 @@ def super_get_matrix_element_calls(self, matrix_element, color_amplitudes, multi for( size_t iicoup = 0; iicoup < nIPC; iicoup++ ) // FIX #823 COUPs[ndcoup + iicoup] = allCOUPs[ndcoup + iicoup]; // independent couplings, fixed for all events fptype* numerators = NUM_ACCESS::ieventAccessRecord( allNumerators, ievt0 * processConfig::ndiagrams ); - fptype* denominators = DEN_ACCESS::ieventAccessRecord( allDenominators, ievt0 ); #endif // Create an array of views over the Flavor Couplings FLV_COUPLING_ARRAY flvCOUPs{ cIPF_partner1, cIPF_partner2, cIPF_value }; @@ -2494,9 +2502,9 @@ def super_get_matrix_element_calls(self, matrix_element, color_amplitudes, multi // Reset color flows (reset jamp_sv) at the beginning of a new event or event page for( int i = 0; i < ncolor; i++ ) { jamp_sv[i] = cxzero_sv(); } - // Numerators and denominators for the current event (CUDA) or SIMD event page (C++) + // Numerators for the current event (CUDA) or SIMD event page (C++) + // (denominators are no longer accumulated here: they are derived as the sum of numerators later) fptype_sv* numerators_sv = NUM_ACCESS::kernelAccessP( numerators ); - fptype_sv& denominators_sv = DEN_ACCESS::kernelAccess( denominators ); // Scalar iflavor for the current event // for GPU it is an int // for SIMD it is also an int, since it is constant across the SIMD vector @@ -2560,8 +2568,7 @@ def _guard_open(group_mask): diagnum = diagram.get('number') amp_block.append("if( storeChannelWeights )") amp_block.append("{") - amp_block.append(" numerators_sv[%i] += cxabs2( amp_sv[0] );" % (diagnum-1)) - amp_block.append(" denominators_sv += cxabs2( amp_sv[0] );") + amp_block.append(" NUM_ATOMIC_ADD( numerators_sv[%i], cxabs2( amp_sv[0] ) );" % (diagnum-1)) amp_block.append("}") for njamp, coeff in color[namp].items(): scoeff = OneProcessExporterMadMatrix.coeff(*coeff) # AV