madspace: run the unit test suite on CUDA and HIP

.github/gpu_ci/README.md

@@ -22,7 +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 |
+| `madspace_tests` | cluster | the madspace unit tests, `madspace/tests`, with `MADSPACE_TEST_DEVICE=cuda` or `hip` (`madspace_tests_gpu.sh`): the tests put their torch tensors on the GPU, so madspace runs them on its GPU runtime. Host-side tests (marker `cpu_only`) are skipped. pytest and scipy are installed in the cache when the modules lack them, and so is a CUDA torch wheel when the PyTorch module is a CPU build (manneback). The log and `junit.xml` are uploaded as an artifact |
 | `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 |
 | `watchdog` | vmadgraph | from the start of the runner to the end of the allocation: if the allocation ends before `stop_runner` completed, cancel the run (`ceci-gate watch-runner`, `remote_watch_runner.sh`) |

.github/gpu_ci/madspace_tests_gpu.sh

@@ -0,0 +1,65 @@
+#!/bin/bash -l
+# Runs on the self-hosted runner (GPU node): job madspace_tests of gpu_runner_ci.yml.
+# The madspace unit tests (madspace/tests) on the GPU: MADSPACE_TEST_DEVICE=$BACKEND puts
+# their torch tensors on the GPU, so that madspace runs on its GPU runtime.
+# Environment: as pp_ttx_mg7.sh (BACKEND, MODULES, MADSPACE_PREFIX, WORKDIR).
+# 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)
+REPO=$(cd "$HERE/../.." && pwd)
+section() { echo; echo "=================== $* ($(date +%T))"; }
+START=$(date +%s)
+
+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 tests on the GPU need a GPU backend, not '$BACKEND'"; exit 2 ;;
+esac
+# test requirements the modules lack (pytest, scipy) are kept per python version, after
+# the packages of the modules in the path
+DEPS=${GLOBALSCRATCH:-$HOME}/mg7-gpu-ci/cache/test-deps-$(python3 -c 'import sys; print("%d.%d" % sys.version_info[:2])')
+export PYTHONPATH=$MADSPACE_PREFIX${PYTHONPATH:+:$PYTHONPATH}:$DEPS
+MISSING=$(python3 -c 'import importlib.util as u; print(" ".join(m for m in ("pytest", "scipy") if u.find_spec(m) is None))')
+if [ -n "$MISSING" ]; then
+    echo "Installing $MISSING into $DEPS"
+    python3 -m pip install --quiet --target="$DEPS" $MISSING
+fi
+# manneback has only a CPU build of PyTorch among its modules: then a CUDA wheel, kept in
+# the cache, goes before it in the path
+if [ "$BACKEND" = cuda ] && ! python3 -c 'import sys, torch; sys.exit(not torch.cuda.is_available())' 2> /dev/null; then
+    TORCH_WHEEL=${TORCH_WHEEL:-torch==2.8.0}
+    TORCH_INDEX=${TORCH_INDEX:-https://download.pytorch.org/whl/cu128}
+    TORCH_DIR=$DEPS-$(printf '%s %s' "$TORCH_WHEEL" "$TORCH_INDEX" | sha1sum | cut -c1-8)
+    if [ ! -e "$TORCH_DIR/.ready" ]; then
+        echo "The torch of the modules has no CUDA: installing $TORCH_WHEEL from $TORCH_INDEX into $TORCH_DIR"
+        rm -rf "$TORCH_DIR"
+        python3 -m pip install --quiet --target="$TORCH_DIR" --index-url "$TORCH_INDEX" "$TORCH_WHEEL"
+        touch "$TORCH_DIR/.ready"
+    fi
+    export PYTHONPATH=$TORCH_DIR:$PYTHONPATH
+fi
+python3 -c 'import torch; print("torch", torch.__version__, "GPU available:", torch.cuda.is_available())'
+
+rm -rf "$WORKDIR"
+mkdir -p "$WORKDIR"
+cd "$WORKDIR"
+section "madspace tests on $BACKEND ($GPU_NAME)"
+STATUS=0
+MADSPACE_TEST_DEVICE=$BACKEND python3 -m pytest "$REPO/madspace/tests" -ra -p no:cacheprovider \
+    --junitxml=junit.xml 2>&1 | tee tests.log || STATUS=1
+
+section "Summary"
+{
+    echo "tests=$(grep -E '^=+ .*(passed|failed|error|skipped).* =+$' tests.log | tail -1 | sed 's/^=* *//; s/ *=*$//')"
+    echo "node=$(uname -n)"
+    echo "gpu=${GPU_NAME:-none}"
+    echo "walltime=$(( $(date +%s) - START ))s"
+} > summary.txt
+cat summary.txt
+exit $STATUS

.github/workflows/gpu_runner_ci.yml

@@ -7,7 +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
+#   madspace_tests  cluster runner  the madspace unit tests on the GPU, .github/gpu_ci/madspace_tests_gpu.sh
 #   interference    cluster runner  squared split orders (interference), .github/gpu_ci/interference_gpu.sh
 #   ...             cluster runner  more CI jobs: same needs, if and runs-on, and add them to
 #                                   stop_runner's needs
@@ -301,7 +301,8 @@ jobs:
           path: results/
           if-no-files-found: warn
 
-  madspace_ops:
+
+  madspace_tests:
     # madspace GPU runtime operations against the CPU runtime (histogram, rank-4 kernels)
     needs: [start_runner, build_madspace]
     if: needs.start_runner.outputs.attempt == github.run_attempt
@@ -313,42 +314,48 @@ jobs:
     env:
       BACKEND: ${{ inputs.backend }}
       MODULES: ${{ inputs.modules }}
-      # no VENV: the python3 of the modules (see madspace_ops_gpu.sh)
+      # no VENV: the python3 of the modules (see madspace_tests_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
+        run: echo "WORKDIR=${GLOBALSCRATCH:-$HOME}/mg7-gpu-ci/runs/$RUN_TAG/madspace_tests" >> "$GITHUB_ENV"
+      - name: madspace tests on the GPU
+        run: .github/gpu_ci/madspace_tests_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
+          (cd "$WORKDIR" && tar -cf - --ignore-failed-read summary.txt tests.log junit.xml 2> /dev/null) | tar -xf - -C results || true
       - name: job summary
         if: always()
         run: |
           {
-            echo "### madspace operations on ${{ inputs.cluster }} ($BACKEND)"
+            echo "### madspace tests 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."
+              echo "No summary.txt: the job stopped early, see the log of the 'madspace tests on the GPU' step."
+            fi
+            if [ -f results/tests.log ]; then
+              echo
+              echo '```'
+              grep -E '^(FAILED|ERROR|XFAIL|SKIPPED) ' results/tests.log | head -100 || true
+              echo '```'
             fi
           } >> "$GITHUB_STEP_SUMMARY"
       - uses: actions/upload-artifact@v4
         if: always()
         with:
-          name: gpu-madspace_ops-${{ inputs.cluster }}-${{ inputs.backend }}
+          name: gpu-madspace_tests-${{ 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, madspace_ops, interference]
+    needs: [start_runner, build_madspace, pp_ttx, pp_jj, madspace_tests, interference]
     if: always() && needs.start_runner.result == 'success'
     # In an attempt that did not start the runner (a re-run of failed jobs, or of one job), a job
     # for the runner label would wait 24 hours, holding the concurrency group of the cluster.

.github/workflows/madspace.yml

@@ -87,6 +87,8 @@ jobs:
               scipy
               --extra-index-url https://download.pytorch.org/whl/cpu
               torch
+            # the GPU runs of these tests are in gpu_runner_ci.yml (job madspace_tests)
+            CIBW_TEST_ENVIRONMENT: MADSPACE_TEST_DEVICE=cpu
             CIBW_TEST_COMMAND: pytest {project}/madspace/tests
 
       - uses: actions/upload-artifact@v4

madspace/src/gpu/runtime.cu

@@ -40,8 +40,8 @@ void op_matrix_element(
     std::size_t output_count = locals[instruction.input_indices[2]].index_value();
     TensorVec contiguous_inputs(input_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
+    // 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) {
@@ -1408,7 +1408,8 @@ void op_histogram(
     Tensor square_weights_tmp(
         DataType::dt_float, {padded_size}, device, AllocHint::temporary
     );
-    // the keys reduce_by_key writes: the n_bins bins, the underflow and the overflow bin
+    // 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(
@@ -1583,7 +1584,8 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) :
     LastUseOfLocals last_use(function);
     InstructionDependencies dependencies(function);
 
-    std::size_t stream_count = 0, event_count = 0, backward_event_count = 0;
+    // at least the main stream, also for a function without instructions
+    std::size_t stream_count = 1, event_count = 0, backward_event_count = 0;
     for (auto& instr : function.instructions()) {
         if (instr.stream_index >= stream_count) {
             stream_count = instr.stream_index + 1;