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

.github/gpu_ci/README.md

@@ -22,6 +22,7 @@ allocation per workflow run, and every CI job of that run executes in it:
 | `build_madspace` | cluster | build madspace with `ENABLE_CUDA`/`ENABLE_HIP` for the GPU of the node (`build_madspace_gpu.sh`). The build is kept in `$GLOBALSCRATCH/mg7-gpu-ci/cache` and redone only when the madspace sources, the modules or the GPU architecture change |
 | `pp_ttx` | cluster | `generate p p > t t~`, `output mg7`, `device = ["cuda"]` or `["hip"]` in `run_card.toml`, `bin/generate_events -f`, check the cross section in `info.json` (`pp_ttx_mg7.sh`). The logs and cards are uploaded as an artifact, and the cross section is shown in the run summary |
 | `pp_jj` | cluster | the same with `generate p p > j j` (`PROCESS`/`TAG` of `pp_ttx_mg7.sh`): its jet cuts act on several objects, which `p p > t t~` has none of |
+| `madspace_ops` | cluster | madspace GPU runtime operations against the CPU runtime on the same inputs (`madspace_ops_gpu.sh`, `madspace_checks.py`; torch puts the inputs on the GPU): a histogram with underflow and overflow bins, an elementwise kernel on a rank-4 tensor |
 | `interference` | cluster | squared split orders on the GPU backend of madmatrix (`interference_gpu.sh`, `interference_checks.py`): standalone `u u~ > u u~` constraints, the helicity choice of an interference matrix element through umami (exact and summed), and `p p > u u~ QCD^2==2` with mg7 against its CPU reference. Input `interference`, on by default and on push |
 | `stop_runner` | cluster | clean up, then create the stop file: `runner_batch.sh` stops the runner and the allocation ends |
 

.github/gpu_ci/madspace_checks.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+"""madspace GPU runtime operations against the CPU runtime, on the same inputs.
+
+Run by madspace_ops_gpu.sh (job madspace_ops of gpu_runner_ci.yml) in its run
+directory. The inputs go to the GPU as torch tensors (dlpack), so that madspace runs
+the function on its GPU runtime; the same function on numpy inputs runs on the CPU.
+
+1. histogram: a weighted histogram with underflow and overflow bins, its values and
+   squared values (the GPU reduction used to write its n_bins + 2 keys into a scratch
+   buffer of n_bins, and summed the squared weights of the wrong events in each bin);
+2. rank4_add: an elementwise kernel on a rank-4 tensor, batch x 2 x 3 x 4 (the GPU
+   launcher computed only the first slice of the last dimension).
+
+Without a GPU-enabled torch among the modules (manneback: the default PyTorch module is
+the CPU build) the checks are reported as skipped. Writes summary.txt (key=value lines)
+in the current directory; exit status 1 if a check fails.
+"""
+
+import os
+import sys
+import time
+
+import numpy as np
+
+import madspace as ms
+
+results = {}
+
+
+def check(name, ok, detail):
+    print('%s %s: %s' % ('PASS' if ok else 'FAIL', name, detail), flush=True)
+    results[name] = (ok, detail)
+
+
+def to_numpy(tensor, torch):
+    return torch.from_dlpack(tensor).cpu().numpy()
+
+
+def check_histogram(torch, device, n=100000, n_bins=40):
+    out_type = ms.Type(ms.DataType.float, ms.BatchSize.one, [n_bins + 2])
+    fb = ms.FunctionBuilder(
+        ms.NamedTypes([(name, ms.batch_float) for name in ('x', 'w', 'lo', 'hi')]),
+        ms.NamedTypes([('values', out_type), ('square_values', out_type)]),
+    )
+    values, square_values = fb.histogram(*(fb.input(i) for i in range(4)), n_bins)
+    fb.output(0, values)
+    fb.output(1, square_values)
+    function = fb.function()
+    rng = np.random.default_rng(17)
+    inputs = [rng.uniform(-0.5, 1.5, n), rng.normal(size=n), np.zeros(n), np.ones(n)]
+    cpu = [np.asarray(ms.Tensor.numpy(t))
+           for t in ms.FunctionRuntime(function, ms.default_context()).call(inputs)]
+    gpu = [to_numpy(t, torch) for t in ms.FunctionRuntime(function).call(
+        [torch.from_numpy(a).to(device) for a in inputs])]
+    worst = [float(np.max(np.abs(g - c) / np.maximum(np.abs(c), 1e-300)))
+             for g, c in zip(gpu, cpu)]
+    under, over = cpu[0][0][0], cpu[0][0][-1]
+    check('histogram', max(worst) < 1e-9 and under != 0 and over != 0,
+          '%d bins + underflow/overflow, max rel diff GPU vs CPU %.1e (values), %.1e '
+          '(squared values); underflow %.4g, overflow %.4g'
+          % (n_bins, worst[0], worst[1], under, over))
+
+
+def check_rank4_add(torch, device, n=4096):
+    rank4 = ms.Type(ms.DataType.float, ms.batch_size, [2, 3, 4])
+    fb = ms.FunctionBuilder(ms.NamedTypes([('x', rank4)]), ms.NamedTypes([('y', rank4)]))
+    x = fb.input(0)
+    fb.output(0, fb.add(x, x))
+    a = np.random.default_rng(23).normal(size=(n, 2, 3, 4))
+    (y,) = ms.FunctionRuntime(fb.function()).call([torch.from_numpy(a).to(device)])
+    y = to_numpy(y, torch)
+    wrong = int(np.sum(y != 2 * a))
+    check('rank4_add', wrong == 0, '%d of %d elements of x + x differ from 2 x'
+          % (wrong, a.size))
+
+
+def main():
+    start = time.time()
+    try:
+        import torch
+    except ImportError as error:
+        check('torch', False, 'cannot import torch (%s): no GPU inputs' % error)
+        torch = None
+    skipped = None
+    if torch is not None and not torch.cuda.is_available():
+        skipped = 'torch %s of the modules has no GPU support' % torch.__version__
+        print('SKIP all checks: %s' % skipped, flush=True)
+        torch = None
+    if torch is not None:
+        for name, step in (('histogram', check_histogram), ('rank4_add', check_rank4_add)):
+            try:
+                step(torch, 'cuda')  # also the ROCm devices in torch
+            except Exception as error:  # one broken step must not hide the others
+                check(name, False, 'error: %s' % error)
+    failures = [name for name, (ok, _detail) in results.items() if not ok]
+    with open('summary.txt', 'w') as f:
+        for name, (ok, detail) in results.items():
+            f.write('%s=%s (%s)\n' % (name, 'ok' if ok else 'FAILED', detail))
+        if skipped:
+            f.write('checks=skipped (%s)\n' % skipped)
+        f.write('node=%s\ngpu=%s\nwalltime=%ds\n' % (
+            os.uname().nodename, os.environ.get('GPU_NAME', 'none'), time.time() - start))
+    print('\n%d check(s) failed: %s' % (len(failures), ', '.join(failures)) if failures
+          else '\nskipped: %s' % skipped if skipped else '\nall checks passed')
+    return 1 if failures else 0
+
+
+if __name__ == '__main__':
+    sys.exit(main())

.github/gpu_ci/madspace_ops_gpu.sh

@@ -0,0 +1,33 @@
+#!/bin/bash -l
+# Runs on the self-hosted runner (GPU node): job madspace_ops of gpu_runner_ci.yml.
+# madspace GPU runtime operations against the CPU runtime: see madspace_checks.py.
+# Environment: as pp_ttx_mg7.sh (BACKEND, MODULES, MADSPACE_PREFIX, WORKDIR).
+# torch (from the PyTorch module) puts the inputs on the GPU. This runs the python3 of
+# the modules, not the venv of build_madspace (same interpreter, which madspace was built
+# with): EasyBuild finds some packages of the modules (e.g. typing_extensions of
+# Python-bundle-PyPI, which torch imports) through EBPYTHONPREFIXES, which only the
+# python3 of the modules reads (its sitecustomize), not a venv made from it.
+set -eo pipefail
+HERE=$(cd "$(dirname "$0")" && pwd)
+section() { echo; echo "=================== $* ($(date +%T))"; }
+
+section "Environment"
+if [ -n "$MODULES" ]; then source "$HERE/load_modules.sh"; fi
+python3 --version
+case $BACKEND in
+    cuda) GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | sed -n 1p) ;;
+    hip)  GPU_NAME=$(rocminfo 2> /dev/null | awk -F': *' '/Marketing Name/ {name=$2} /Device Type: *GPU/ {print name; exit}' || true) ;;
+    *)    echo "::error::the madspace checks need a GPU backend, not '$BACKEND'"; exit 2 ;;
+esac
+export GPU_NAME
+export PYTHONPATH=$MADSPACE_PREFIX${PYTHONPATH:+:$PYTHONPATH}
+
+rm -rf "$WORKDIR"
+mkdir -p "$WORKDIR"
+cd "$WORKDIR"
+STATUS=0
+python3 "$HERE/madspace_checks.py" 2>&1 | tee checks.log || STATUS=1
+
+section "Summary"
+cat summary.txt 2> /dev/null || echo "no summary.txt"
+exit $STATUS

.github/workflows/gpu_runner_ci.yml

@@ -7,6 +7,7 @@
 #   build_madspace  cluster runner  madspace with the CUDA/HIP backend (cached on the cluster)
 #   pp_ttx          cluster runner  p p > t t~ with mg7 on the GPU
 #   pp_jj           cluster runner  p p > j j with mg7 on the GPU (cuts on several jets)
+#   madspace_ops    cluster runner  madspace GPU operations against the CPU runtime, .github/gpu_ci/madspace_checks.py
 #   interference    cluster runner  squared split orders (interference), .github/gpu_ci/interference_gpu.sh
 #   ...             cluster runner  more CI jobs: same runs-on, and add them to stop_runner's needs
 #   stop_runner     cluster runner  clean up and end the allocation (no ssh key needed)
@@ -291,9 +292,53 @@ jobs:
           path: results/
           if-no-files-found: warn
 
+  madspace_ops:
+    # madspace GPU runtime operations against the CPU runtime (histogram, rank-4 kernels)
+    needs: build_madspace
+    runs-on: [self-hosted, "${{ inputs.runner_label }}"]
+    timeout-minutes: 30
+    defaults:
+      run:
+        shell: bash -l -eo pipefail {0}
+    env:
+      BACKEND: ${{ inputs.backend }}
+      MODULES: ${{ inputs.modules }}
+      # no VENV: the python3 of the modules (see madspace_ops_gpu.sh)
+      MADSPACE_PREFIX: ${{ needs.build_madspace.outputs.madspace_prefix }}
+    steps:
+      - uses: actions/checkout@v5
+      - name: run directory
+        run: echo "WORKDIR=${GLOBALSCRATCH:-$HOME}/mg7-gpu-ci/runs/$RUN_TAG/madspace_ops" >> "$GITHUB_ENV"
+      - name: madspace operations on the GPU
+        run: .github/gpu_ci/madspace_ops_gpu.sh
+      - name: collect logs
+        if: always()
+        run: |
+          mkdir -p results
+          (cd "$WORKDIR" && tar -cf - --ignore-failed-read summary.txt checks.log 2> /dev/null) | tar -xf - -C results || true
+      - name: job summary
+        if: always()
+        run: |
+          {
+            echo "### madspace operations on ${{ inputs.cluster }} ($BACKEND)"
+            if [ -f results/summary.txt ]; then
+              echo "| | |"
+              echo "|---|---|"
+              sed 's/^\([^=]*\)=\(.*\)$/| \1 | \2 |/' results/summary.txt
+            else
+              echo "No summary.txt: the job stopped early, see the log of the 'madspace operations on the GPU' step."
+            fi
+          } >> "$GITHUB_STEP_SUMMARY"
+      - uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          name: gpu-madspace_ops-${{ inputs.cluster }}-${{ inputs.backend }}
+          path: results/
+          if-no-files-found: warn
+
   stop_runner:
     # must list every job that runs on the self-hosted runner
-    needs: [start_runner, build_madspace, pp_ttx, pp_jj, interference]
+    needs: [start_runner, build_madspace, pp_ttx, pp_jj, madspace_ops, interference]
     if: always() && needs.start_runner.result == 'success'
     runs-on: [self-hosted, "${{ inputs.runner_label }}"]
     defaults:

madspace/include/madspace/driver/tensor.hpp

@@ -10,6 +10,8 @@
 #include <functional>
 #include <initializer_list>
 #include <optional>
+#include <stdexcept>
+#include <string>
 #include <vector>
 
 namespace madspace {
@@ -74,6 +76,9 @@ class Sizes {
     const std::size_t* end() const { return &_values[_size]; }
     /// Append one more dimension.
     void push_back(std::size_t item) {
+        if (_size >= max_size) {
+            throw std::out_of_range("maximum dimension exceeded");
+        }
         _values[_size] = item;
         ++_size;
     }
@@ -95,7 +100,9 @@ class Sizes {
     }
 
 private:
-    std::size_t _values[max_size];
+    // zero beyond size(): a view with more dimensions than its tensor (see
+    // Tensor::view) must not read stale strides and shapes
+    std::size_t _values[max_size] = {};
     std::size_t _size;
 };
 
@@ -514,6 +521,7 @@ class Tensor {
     /// Typed, dimension-checked view onto the data for direct element access.
     TensorView<T, dim> view() {
         check_impl();
+        check_view_dim(dim, impl->shape.size());
         T* data = static_cast<T*>(impl->data);
         return TensorView<T, dim>(data, impl->stride.data(), impl->shape.data());
     }
@@ -522,6 +530,7 @@ class Tensor {
     /// Typed, dimension-checked view onto the data for direct element access.
     const TensorView<T, dim> view() const {
         check_impl();
+        check_view_dim(dim, impl->shape.size());
         T* data = static_cast<T*>(impl->data);
         return TensorView<T, dim>(data, impl->stride.data(), impl->shape.data());
     }
@@ -532,11 +541,13 @@ class Tensor {
         check_impl();
         T* data = static_cast<T*>(impl->data);
         if (flatten_count <= 1) {
+            check_view_dim(dim, impl->shape.size());
             return {data, impl->stride, impl->shape};
         }
         if (flatten_count > impl->contiguous_dims) {
             throw std::invalid_argument("can only flatten contiguous dimensions");
         }
+        check_view_dim(dim, impl->shape.size() - flatten_count + 1);
         Sizes stride{1}, shape{1};
         std::size_t i = 0;
         for (; i < flatten_count; ++i) {
@@ -859,6 +870,18 @@ class Tensor {
         }
     }
 
+    // a view of more dimensions than the tensor has reads strides and shapes that were
+    // never set: out-of-bounds accesses, silent on one platform, a memory fault on
+    // another (e.g. a kernel taking FIn<T, 1> for a tensor of declared rank 0)
+    static void check_view_dim(std::size_t view_dim, std::size_t tensor_dim) {
+        if (view_dim > tensor_dim) {
+            throw std::invalid_argument(
+                "view of dimension " + std::to_string(view_dim) +
+                " of a tensor of dimension " + std::to_string(tensor_dim)
+            );
+        }
+    }
+
     template <typename D>
     void allocate(std::size_t size, const D& device, AllocHint hint) {
         auto [data, parent] = device.allocate(size, hint);

madspace/src/gpu/runtime.cu

@@ -12,6 +12,7 @@
 #include <thrust/fill.h>
 #include <thrust/gather.h>
 #include <thrust/iterator/constant_iterator.h>
+#include <thrust/iterator/zip_iterator.h>
 #include <thrust/sequence.h>
 #include <thrust/sort.h>
 
@@ -38,8 +39,10 @@ void op_matrix_element(
     std::size_t input_count = locals[instruction.input_indices[1]].index_value();
     std::size_t output_count = locals[instruction.input_indices[2]].index_value();
     TensorVec contiguous_inputs(input_count);
-    std::vector<UmamiInputKey> input_keys(input_count + 1);
-    std::vector<UmamiOutputKey> output_keys(output_count);
+    std::vector<UmamiInputKey> input_keys(input_count);
+    // one more output: the stream (UMAMI_OUT_GPU_STREAM), so that the matrix element runs
+    // in stream order with the rest of the function, not on the default stream
+    std::vector<UmamiOutputKey> output_keys(output_count + 1);
     std::vector<void*> input_ptrs(input_count), output_ptrs(output_count + 1);
     for (std::size_t i = 0; i < input_count; ++i) {
         input_keys[i] = static_cast<UmamiInputKey>(
@@ -113,7 +116,7 @@ void op_matrix_element(
         input_count,
         input_keys.data(),
         input_ptrs.data(),
-        output_count,
+        output_count + 1,
         output_keys.data(),
         output_ptrs.data()
     );
@@ -1405,7 +1408,8 @@ void op_histogram(
     Tensor square_weights_tmp(
         DataType::dt_float, {padded_size}, device, AllocHint::temporary
     );
-    Tensor reduce_tmp(DataType::dt_float, {n_bins}, device, AllocHint::temporary);
+    // the keys reduce_by_key writes: the n_bins bins, the underflow and the overflow bin
+    Tensor reduce_tmp(DataType::dt_float, {n_bins + 2}, device, AllocHint::temporary);
 
     launch_kernel(
         kernel_prepare_hist,
@@ -1431,7 +1435,19 @@ void op_histogram(
     auto weights_ptr =
         thrust::device_pointer_cast(static_cast<double*>(weights_tmp.data()));
     auto values_ptr = thrust::device_pointer_cast(static_cast<double*>(values.data()));
-    thrust::sort_by_key(policy, indices_ptr, indices_ptr + padded_size, weights_ptr);
+    auto square_weights_ptr =
+        thrust::device_pointer_cast(static_cast<double*>(square_weights_tmp.data()));
+    auto square_values_ptr =
+        thrust::device_pointer_cast(static_cast<double*>(square_values.data()));
+    // one sort of the bin indices, carrying the weights AND the squared weights along:
+    // sorting the indices again for the squared weights (as before) left these in the
+    // event order, so the squared weights of the wrong events were summed in each bin
+    thrust::sort_by_key(
+        policy,
+        indices_ptr,
+        indices_ptr + padded_size,
+        thrust::make_zip_iterator(thrust::make_tuple(weights_ptr, square_weights_ptr))
+    );
     thrust::reduce_by_key(
         policy,
         indices_ptr,
@@ -1440,14 +1456,6 @@ void op_histogram(
         reduce_tmp_ptr,
         values_ptr
     );
-
-    auto square_weights_ptr =
-        thrust::device_pointer_cast(static_cast<double*>(square_weights_tmp.data()));
-    auto square_values_ptr =
-        thrust::device_pointer_cast(static_cast<double*>(square_values.data()));
-    thrust::sort_by_key(
-        policy, indices_ptr, indices_ptr + padded_size, square_weights_ptr
-    );
     thrust::reduce_by_key(
         policy,
         indices_ptr,

madspace/src/gpu/tensor.cuh

@@ -216,7 +216,7 @@ __global__ void run_kernel(std::size_t batch_size, V... views) {
         std::size_t size1 = first_view.size(1);
         std::size_t size2 = first_view.size(2);
         std::size_t size3 = first_view.size(3);
-        if (index >= batch_size * size1 * size2) {
+        if (index >= batch_size * size1 * size2 * size3) {
             return;
         }
         std::size_t i = index % batch_size;