From ac43ec34d9181114a290ca734d13f16e25229d7f Mon Sep 17 00:00:00 2001
From: Yaroslav Bolyukin <iam@lach.pw>
Date: Tue, 6 Feb 2024 16:44:11 +0100
Subject: [PATCH 1/2] torch: passthru rocmSupport flag

---
 pkgs/development/python-modules/torch/default.nix | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index d6c51904bd9d..c67be4cb9233 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -492,7 +492,7 @@ in buildPythonPackage rec {
   requiredSystemFeatures = [ "big-parallel" ];
 
   passthru = {
-    inherit cudaSupport cudaPackages;
+    inherit cudaSupport cudaPackages rocmSupport rocmPackages;
     # At least for 1.10.2 `torch.fft` is unavailable unless BLAS provider is MKL. This attribute allows for easy detection of its availability.
     blasProvider = blas.provider;
     # To help debug when a package is broken due to CUDA support

From 2eedfae46b0e632089aefad388cbc7093e7eea0b Mon Sep 17 00:00:00 2001
From: Yaroslav Bolyukin <iam@lach.pw>
Date: Sat, 2 Mar 2024 04:08:30 +0100
Subject: [PATCH 2/2] torchaudio: add rocm support

---
 .../python-modules/torchaudio/default.nix     | 52 ++++++++++++++++++-
 1 file changed, 51 insertions(+), 1 deletion(-)

diff --git a/pkgs/development/python-modules/torchaudio/default.nix b/pkgs/development/python-modules/torchaudio/default.nix
index 73aec87cca61..f9d4945ef6c0 100644
--- a/pkgs/development/python-modules/torchaudio/default.nix
+++ b/pkgs/development/python-modules/torchaudio/default.nix
@@ -9,10 +9,46 @@
 , pybind11
 , sox
 , torch
+
 , cudaSupport ? torch.cudaSupport
 , cudaPackages
+, rocmSupport ? torch.rocmSupport
+, rocmPackages
+
+, gpuTargets ? []
 }:
 
+let
+  # TODO: Reuse one defined in torch?
+  # Some of those dependencies are probbly not required,
+  # but it breaks when the store path is different between torch and torchaudio
+  rocmtoolkit_joined = symlinkJoin {
+    name = "rocm-merged";
+
+    paths = with rocmPackages; [
+      rocm-core clr rccl miopen miopengemm rocrand rocblas
+      rocsparse hipsparse rocthrust rocprim hipcub roctracer
+      rocfft rocsolver hipfft hipsolver hipblas
+      rocminfo rocm-thunk rocm-comgr rocm-device-libs
+      rocm-runtime clr.icd hipify
+    ];
+
+    # Fix `setuptools` not being found
+    postBuild = ''
+      rm -rf $out/nix-support
+    '';
+  };
+  # Only used for ROCm
+  gpuTargetString = lib.strings.concatStringsSep ";" (
+    if gpuTargets != [ ] then
+    # If gpuTargets is specified, it always takes priority.
+      gpuTargets
+    else if rocmSupport then
+      rocmPackages.clr.gpuTargets
+    else
+      throw "No GPU targets specified"
+  );
+in
 buildPythonPackage rec {
   pname = "torchaudio";
   version = "2.2.1";
@@ -33,6 +69,11 @@ buildPythonPackage rec {
     substituteInPlace setup.py \
       --replace 'print(" --- Initializing submodules")' "return" \
       --replace "_fetch_archives(_parse_sources())" "pass"
+  ''
+  + lib.optionalString rocmSupport ''
+    # There is no .info/version-dev, only .info/version
+    substituteInPlace cmake/LoadHIP.cmake \
+      --replace "/.info/version-dev" "/.info/version"
   '';
 
   env = {
@@ -55,7 +96,11 @@ buildPythonPackage rec {
     ninja
   ] ++ lib.optionals cudaSupport [
     cudaPackages.cuda_nvcc
-  ];
+  ] ++ lib.optionals rocmSupport (with rocmPackages; [
+    clr
+    rocblas
+    hipblas
+  ]);
 
   buildInputs = [
     ffmpeg-full
@@ -73,6 +118,11 @@ buildPythonPackage rec {
   BUILD_RNNT=0;
   BUILD_CTC_DECODER=0;
 
+  preConfigure = lib.optionalString rocmSupport ''
+    export ROCM_PATH=${rocmtoolkit_joined}
+    export PYTORCH_ROCM_ARCH="${gpuTargetString}"
+  '';
+
   dontUseCmakeConfigure = true;
 
   doCheck = false; # requires sox backend