Invert event-helicity parallelism on GPU

[Qubitol]

This is just an idea worth investigating.

On GPU, we are currently launching the calculate_jamps kernel with the following configuration:

• threads: (256, 1)
• blocks: (nevents // 256, nhelicities)

This means that each block gets a single value of helicity, so each event in a single block manages 1 helicity.

What if we invert this parallelism, and we make each block manage 1 event and all the helicities?

This might be useful, likely, only for those processes with many helicity combinations (read: high multiplicity), so we might need to have a (compile-time) switch (as soon as we understand what is the sweet spot) that selects if keeping launch parameters as they are, or inverting this parallelism.

This idea is justified by:

• the GPU kernel being branchless on the point of view of the helicity value, so it can run in perfect lockstep even if threads are managing different helicity values;
• the fact that we would have one event per block, that we can store in shared memory: each thread in a block will then just need to load the same nparticles * np4 * 2 double numbers from shared memory and 1 integer (the helicity value) from constant memory, instead of having different nparticles * np4 * 2 and anyway 1 integer per thread; at the block level this reduces the memory required to load all the momenta in a block (and consequently in a warp), possibly reducing the registers pressure.

Also intermediate scenarios could be studied, for example:

• threads: (256 // n, n)
• blocks: (nevents // (256 // n), nhelicities // n)

in which we would make use of a 2-dimensional grid approach, which might or not be more efficient in some scenarios.

I'd propose to study that on higher multiplicity processes such as g g > t t~ g g, and g g > t t~ g g g, by comparing the difference in throughput using the standalone exporter.
NSight Compute can then be used to study the warp scheduling and occupancy.

[theoheimel]

Related, but different: We can't guarantee that the batch size is divisible by block size because of cuts in the end-to-end GPU setup. In the worst case, that means we have one almost empty block for each helicity. I used the 2d grid mainly because it was the "minimally invasive" way to do it without touching all the memory management. But a better way to do it would probably be to flatten nevents * nhelicities into a single grid such that we are left with at most one padded block.

[stloufra]

After implementation wrap-up in block_helicities_test and some numbers for the performance:

Convention

Option	Launch bounds
EVENT	blocks=gpublocks, threads=gputhreads (1 thread = 1 event, fixed helicity per launch)
HELICITY	blocks=nevt, threads=cNGoodHel (1 block = 1 event, 1 thread = 1 helicity)

Selected via mgOnGpuConfig.h's MGONGPU_HELBLOCK_LAYOUT_EVENT

Verified identical ME values to reference commit 1a4f00d for both
gg_ttxgg and gg_ttxggg and both ways.

Throughput (evts/s, best config - launch bounds - each)

Measured on our a100 machine.

process	EVENT	HELICITY	HEL/EVT
gg_ttxgg	3.20e5	2.75e5	0.86
gg_ttxggg	6.85e3	6.18e3	0.90

Conclusion

HELICITY layout is slower (10-15%) than EVENT layout on both processes. calculate_jamps is memory (DRAM) bound mainly by the helicity amplitudes arrays. EVENT layout's coalesced per-thread momenta reads beat the shared-memory staging savings under the HELICITY. Might improve for higher multiplicity but IMO not something worth pursuing.

[Qubitol]

Thank you for checking this. It's very nice work. It's a bit unfortunate that it doesn't work better than the default, but that's why we are making tests.

[stloufra]

Following idea is to merge on the GPU side the two kernels color_sum and calculate_jamp so the array allJamp_sv[ncolor] for storing of the intermediate helicity combination color results could be eliminated and we would only atomically add to the M^2 in exchange for more compute done (color_sum quadratic form redone by each thread - for each of the helicity combination).

or

in combination with this development leave the amplitudes in shared memory as well and immediately perform the quadratic form within the same launch block/kernel.

[stloufra]

I have reopened the issue after discussion with Olivier. The subsequent step to examine gains/penalty of putting all of the momenta in shared memory - meaning fission of the ALOHAOBJ so wavefunctions and momenta are untangled and the momenta is reused directly from the L2 cache on GPU for all the threads computing the same event within the block with different helicity combinations.
or
having only the wave functions in memory and the momenta is recomputed on the fly from the phase space point.