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