From be7e126e474c29cefe66070743d001404b257fe0 Mon Sep 17 00:00:00 2001
From: Connor Baker <ConnorBaker01@gmail.com>
Date: Fri, 6 Jun 2025 13:29:59 -0700
Subject: [PATCH] magma: add passthru.testers and passthru.tests

Signed-off-by: Connor Baker <ConnorBaker01@gmail.com>
---
 .../libraries/science/math/magma/generic.nix  | 143 +++++++++++++++++-
 1 file changed, 141 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/libraries/science/math/magma/generic.nix b/pkgs/development/libraries/science/math/magma/generic.nix
index c7938b98368a..394e249abc7e 100644
--- a/pkgs/development/libraries/science/math/magma/generic.nix
+++ b/pkgs/development/libraries/science/math/magma/generic.nix
@@ -27,8 +27,10 @@
   # At least one back-end has to be enabled,
   # and we can't default to CUDA since it's unfree
   rocmSupport ? !cudaSupport,
+  runCommand,
   static ? stdenv.hostPlatform.isStatic,
   stdenv,
+  writeShellApplication,
 }:
 
 let
@@ -104,7 +106,7 @@ in
 
 assert (builtins.match "[^[:space:]]*" gpuTargetString) != null;
 
-stdenv.mkDerivation {
+stdenv.mkDerivation (finalAttrs: {
   pname = "magma";
   inherit version;
 
@@ -241,6 +243,143 @@ stdenv.mkDerivation {
   passthru = {
     inherit cudaSupport rocmSupport gpuTargets;
     cudaPackages = effectiveCudaPackages;
+    testers = {
+      all =
+        let
+          magma = finalAttrs.finalPackage;
+        in
+        writeShellApplication {
+          derivationArgs = {
+            __structuredAttrs = true;
+            strictDeps = true;
+          };
+          name = "magma-testers-all";
+          text = ''
+            logWithDate() {
+              printf "%s: %s\n" "$(date --utc --iso-8601=seconds)" "$*"
+            }
+
+            isIgnoredTest() {
+              case $1 in
+                # Skip the python scripts
+                *.py) return 0 ;;
+
+                # These test require files, so we skip them
+                testing_?io) ;&
+                testing_?madd) ;&
+                testing_?matrix) ;&
+                testing_?matrixcapcup) ;&
+                testing_?matrixinfo) ;&
+                testing_?mcompressor) ;&
+                testing_?mconverter) ;&
+                testing_?preconditioner) ;&
+                testing_?solver) ;&
+                testing_?solver_rhs) ;&
+                testing_?solver_rhs_scaling) ;&
+                testing_?sort) ;&
+                testing_?spmm) ;&
+                testing_?spmv) ;&
+                testing_?spmv_check) ;&
+                testing_?sptrsv) ;&
+                testing_dsspmv_mixed) ;&
+                testing_zcspmv_mixed)
+                  logWithDate "skipping $1 because it requires input"
+                  return 0
+                  ;;
+
+                # These test require outputing to files, so we skip them
+                testing_?print)
+                  logWithDate "skipping $1 because it requires creating output"
+                  return 0
+                  ;;
+
+                # These test succeed but exit with a non-zero code
+                testing_[cdz]gglse) ;&
+                testing_sgemm_fp16)
+                  logWithDate "skipping $1 because has a non-zero exit code"
+                  return 0
+                  ;;
+
+                # These test have memory freeing/allocation errors:
+                testing_?mdotc)
+                  logWithDate "skipping $1 because it fails to allocate or free memory"
+                  return 0
+                  ;;
+
+                # Test is not ignored otherwise.
+                *) return 1 ;;
+              esac
+            }
+
+            runTests() {
+              local -nr outputArray="$1"
+              local -i programExitCode=0
+              local file
+
+              # TODO: Collect and sort filenames prior to iterating so the order isn't dependent on the filesystem.
+              for file in "${magma.test}"/bin/*; do
+                if isIgnoredTest "$(basename "$file")"; then
+                  continue
+                fi
+
+                logWithDate "Starting $file"
+
+                # Since errexit is set, we need to reset programExitCode every iteration and use an OR
+                # to set it only when the test fails (which should not fail, avoiding tripping errexit).
+                programExitCode=0
+
+                # A number of test cases require an input <=128, so we set the range to include [128, 1024].
+                # Batch is kept small to keep tests fast.
+                "$file" --range 128:1024:896 --batch 32 || programExitCode=$?
+
+                logWithDate "Finished $file with exit code $programExitCode"
+
+                if ((programExitCode)); then
+                  outputArray+=("$file")
+                fi
+              done
+            }
+
+            main() {
+              local -a failedPrograms=()
+              runTests failedPrograms
+
+              if ((''${#failedPrograms[@]})); then
+                logWithDate "The following programs had non-zero exit codes:"
+                for file in "''${failedPrograms[@]}"; do
+                  # Using echo to avoid printing the date
+                  echo "- $file"
+                done
+                logWithDate "Exiting with code 1 because at least one test failed."
+                exit 1
+              fi
+
+              logWithDate "All tests passed!"
+              exit 0
+            }
+
+            main
+          '';
+          runtimeInputs = [ magma.test ];
+        };
+    };
+    tests = {
+      all =
+        runCommand "magma-tests-all"
+          {
+            __structuredAttrs = true;
+            strictDeps = true;
+            nativeBuildInputs = [ finalAttrs.passthru.testers.all ];
+            requiredSystemFeatures = lib.optionals cudaSupport [ "cuda" ];
+          }
+          ''
+            if magma-testers-all; then
+              touch "$out"
+            else
+              exit 1
+            fi
+          '';
+    };
   };
 
   meta = with lib; {
@@ -260,4 +399,4 @@ stdenv.mkDerivation {
       || (cudaSupport && strings.versionOlder version "2.7.1" && cudaPackages_11 == null)
       || (rocmSupport && strings.versionOlder version "2.8.0");
   };
-}
+})