madspace GPU runtime: histogram scratch, matrix element stream, rank-4 kernels, view rank check

[oliviermattelaer]

Found by the review that followed the HIP cut crash (#254). These bugs are of the same class: silent or lucky on one platform, a fault on another.

Fixes

1. op_histogram (GPU): reduce_by_key writes n_bins + 2 keys (the bins, plus the underflow and overflow bins) into a scratch buffer of n_bins, so every histogram fill wrote two doubles past its end.
2. op_matrix_element (GPU): heap overflow, and the stream was never passed.
   • output_keys had output_count entries, but the stream key was written at index output_count.
   • The call passed output_count, so umami never received madspace's stream. Every matrix element has run on the default stream since ffb1610, which meant to pass it (it moved the stream from an input key to an output key).
   • Now passed, so the matrix element runs in stream order with the rest of the function. The unused extra input key is dropped.
3. run_kernel (GPU), rank 4: the bound check omitted size3, so only the first slice of the last dimension was computed. This affects elementwise kernels, copies, zero and add-in-place of rank-4 tensors (e.g. PDF coefficient grids).
4. op_histogram (GPU), squared values, found by the new madspace_ops check on lemaitre4 (GPU vs CPU up to 16% in a bin). The bin indices were sorted with the weights, then the already sorted indices were sorted again with the squared weights, which therefore stayed in event order. Each bin summed the squared weights of the wrong events, so the per-bin statistical errors of GPU histograms were wrong while the values and totals were right. One sort_by_key now carries both arrays (zip iterator).
5. Tensor::view / flat_view: a view with more dimensions than its tensor now throws, instead of reading strides and shapes that were never set. That was the mechanism of the cut-kernel crash. Sizes is now zero-initialised and push_back is bounded.

Merge order with #255. #255 (madmatrix GPU, stacked on #244) makes the madmatrix GPU template put all of its work on the stream it is given; before it, some of the work went to the default stream.

• With madspace's own streams, which are blocking, this PR is correct with both old and new process libraries.
• With a caller-supplied non-blocking stream (a torch side stream, e.g. MadNIS), only libraries generated with madmatrix GPU: running flavour couplings, no good helicity, stream order, unset selections #255 are race-free.
• Merging madmatrix GPU: running flavour couplings, no good helicity, stream order, unset selections #255 first is the safe order.

Not in this PR:

• From the same review (lower priority): the GPU VEGAS and discrete histograms have no range guard (the CPU skips out-of-range samples), and some 32-bit index products.
• On both CPU and GPU, a value in (min − bin width, min) lands in the first bin instead of the underflow, because the index is truncated toward zero instead of floored.

Tests

• madspace/tests/test_histogram_op.py: CPU semantics of the histogram (underflow, bins, overflow, per-event range).
• GPU CI job madspace_ops (madspace_ops_gpu.sh, madspace_checks.py) runs the histogram and a rank-4 elementwise kernel (x + x on batch × 2 × 3 × 4) on the GPU runtime and compares them with the CPU runtime. The inputs are torch device tensors, so torch must be importable after the modules are loaded; otherwise the job fails and says so.
• Local CPU results:
  • madspace suite: 1722 passed, 11 skipped, no view rank error;
  • mg7 CPU runs of p p > j j and p p > e+ e- j;
  • madspace_checks.py with a numpy stand-in for torch.
• GPU, all jobs green on both clusters (pp_ttx, pp_jj, madspace_ops):
  • lemaitre4 (HIP, MI300A): https://github.com/MadGraphTeam/MadGraph7/actions/runs/38048327359
    • histogram: GPU = CPU to 1.4e-13 (values, signed weights) and 2.5e-15 (squared values);
    • rank-4 x + x: 0 of 98304 elements differ.
  • manneback (CUDA): https://github.com/MadGraphTeam/MadGraph7/actions/runs/38048325027
    • madspace_ops is skipped: the default PyTorch module there (2.6.0-foss-2024a) is the CPU build;
    • the CUDA build with the changes, p p > t t~ and p p > j j all pass.

[theoheimel]

The fixes look good.

What I would change (but that's work for a separate PR): instead of adding GPU-specific unit tests for madspace as done here, we should move the madspace unit tests from numpy to torch (which we already use for some tests anyways), so that the same tests run on CPU and GPU, and then run the full madspace unit test suite on CUDA and HIP through the new GPU CI.