Skip to content

Add atomicAdd for numerators, and remove denominators from kernel #59

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Qubitol merged 1 commit into from
Aug 7, 2026
Merged

Add atomicAdd for numerators, and remove denominators from kernel #59

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
28 changes: 11 additions & 17 deletions madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -693,35 +693,29 @@ 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
const fptype globaldenom) /* clang-format on */
{
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;
Expand Down
1 change: 1 addition & 0 deletions madgraph/iolibs/template_files/madmatrix/process_matrix.inc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,3 +47,4 @@
mgDebug( 1, __FUNCTION__ );
return;
}
#undef NUM_ATOMIC_ADD
30 changes: 22 additions & 8 deletions madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 ) );
Comment thread
roiser marked this conversation as resolved.
#else
// *** PART 0b - C++ ***
Expand Down Expand Up @@ -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 );
Expand Down Expand Up @@ -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++ )
//{
Expand Down
6 changes: 4 additions & 2 deletions madgraph/iolibs/template_files/madmatrix/umami.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -386,8 +386,10 @@ extern "C"
{reinterpret_cast<void**>(&matrix_elements), rounded_count * sizeof( fptype )},
{reinterpret_cast<void**>(&diagram_index), rounded_count * sizeof( unsigned int )},
{reinterpret_cast<void**>(&color_jamps), rounded_count * CPPProcess::ncolor * mgOnGpu::nx2 * sizeof( fptype )},
{reinterpret_cast<void**>(&numerators), rounded_count * CPPProcess::ndiagrams * CPPProcess::ncomb * sizeof( fptype )},
{reinterpret_cast<void**>(&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<void**>(&numerators), rounded_count * CPPProcess::ndiagrams * sizeof( fptype )},
{reinterpret_cast<void**>(&denominators), rounded_count * sizeof( fptype )},
{reinterpret_cast<void**>(&helicity_index), rounded_count * sizeof( int )},
{reinterpret_cast<void**>(&color_index), rounded_count * sizeof( int )},
{reinterpret_cast<void**>(&ghel_matrix_elements), rounded_count * CPPProcess::ncomb * sizeof( fptype )},
Expand Down
27 changes: 17 additions & 10 deletions madmatrix/model_handling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand All @@ -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;
Expand All @@ -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()
Expand Down Expand Up @@ -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 );
Expand All @@ -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<nIPF, nMF> flvCOUPs{ cIPF_partner1, cIPF_partner2, cIPF_value };
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading