From 39f33456e4443cdeb4b56c0b87a667b6a0ee3699 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Sat, 9 Dec 2023 21:52:13 +0000
Subject: [PATCH] python3Packages.torch.gpuChecks: add rocm

---
 .../python-modules/torch/default.nix          |  2 +-
 .../python-modules/torch/gpu-checks.nix       | 50 +++++++++++++++++++
 .../python-modules/torch/test-cuda.nix        | 21 --------
 3 files changed, 51 insertions(+), 22 deletions(-)
 create mode 100644 pkgs/development/python-modules/torch/gpu-checks.nix
 delete mode 100644 pkgs/development/python-modules/torch/test-cuda.nix

diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index 62c69905609e..11a3b3df42ce 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -648,7 +648,7 @@ buildPythonPackage rec {
     blasProvider = blas.provider;
     # To help debug when a package is broken due to CUDA support
     inherit brokenConditions;
-    gpuChecks.cudaAvailable = callPackage ./test-cuda.nix { torch = torchWithCuda; };
+    gpuChecks = callPackage ./gpu-checks.nix { };
   };
 
   meta = {
diff --git a/pkgs/development/python-modules/torch/gpu-checks.nix b/pkgs/development/python-modules/torch/gpu-checks.nix
new file mode 100644
index 000000000000..71004d8457b2
--- /dev/null
+++ b/pkgs/development/python-modules/torch/gpu-checks.nix
@@ -0,0 +1,50 @@
+{
+  lib,
+  callPackage,
+  torchWithCuda,
+  torchWithRocm,
+}:
+
+let
+  accelAvailable =
+    {
+      feature,
+      versionAttr,
+      torch,
+      runCommandNoCC,
+      writers,
+    }:
+    let
+      name = "${torch.name}-${feature}-check";
+      unwrapped = writers.writePython3Bin "${name}-unwrapped" {libraries = [torch];} ''
+        import torch
+        message = f"{torch.cuda.is_available()=} and {torch.version.${versionAttr}=}"
+        assert torch.cuda.is_available() and torch.version.${versionAttr}, message
+        print(message)
+      '';
+    in
+    runCommandNoCC name
+      {
+        nativeBuildInputs = [unwrapped];
+        requiredSystemFeatures = [feature];
+        passthru = {
+          inherit unwrapped;
+        };
+      }
+      ''
+        ${name}-unwrapped
+        touch $out
+      '';
+in
+{
+  cudaAvailable = callPackage accelAvailable {
+    feature = "cuda";
+    versionAttr = "cuda";
+    torch = torchWithCuda;
+  };
+  rocmAvailable = callPackage accelAvailable {
+    feature = "rocm";
+    versionAttr = "hip";
+    torch = torchWithRocm;
+  };
+}
diff --git a/pkgs/development/python-modules/torch/test-cuda.nix b/pkgs/development/python-modules/torch/test-cuda.nix
deleted file mode 100644
index 22d73eeb9167..000000000000
--- a/pkgs/development/python-modules/torch/test-cuda.nix
+++ /dev/null
@@ -1,21 +0,0 @@
-{ runCommandNoCC
-, python
-, torch
-}:
-
-runCommandNoCC "${torch.name}-gpu-test"
-{
-  nativeBuildInputs = [
-    (python.withPackages (_: [ torch ]))
-  ];
-  requiredSystemFeatures = [
-    "cuda"
-  ];
-} ''
-  python3 << EOF
-  import torch
-  assert torch.cuda.is_available(), f"{torch.cuda.is_available()=}"
-  EOF
-
-  touch $out
-''