python3Packages.{torch,torchaudio}: DRY shared ROCm infrastructure (#498204)

pkgs/development/python-modules/torch/source/default.nix

@@ -746,6 +746,8 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
       rocmSupport
       rocmPackages
       unroll-src
+      gpuTargetString
+      rocmtoolkit_joined
       ;
     cudaCapabilities = if cudaSupport then supportedCudaCapabilities else [ ];
     # At least for 1.10.2 `torch.fft` is unavailable unless BLAS provider is MKL. This attribute allows for easy detection of its availability.
@@ -762,6 +764,7 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
     homepage = "https://pytorch.org/";
     license = lib.licenses.bsd3;
     maintainers = with lib.maintainers; [
+      caniko
       GaetanLepage
       LunNova # esp. for ROCm
       teh

pkgs/development/python-modules/torchaudio/default.nix

@@ -1,6 +1,5 @@
 {
   lib,
-  symlinkJoin,
   buildPythonPackage,
   fetchFromGitHub,
 
@@ -19,62 +18,9 @@
   cudaPackages,
   rocmSupport ? torch.rocmSupport,
   rocmPackages,
-
-  gpuTargets ? [ ],
 }:
 
-let
-  # TODO: Reuse one defined in torch?
-  # Some of those dependencies are probably 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
-      rocrand
-      rocblas
-      rocsparse
-      hipsparse
-      rocthrust
-      rocprim
-      hipcub
-      roctracer
-      rocfft
-      rocsolver
-      hipfft
-      hipsolver
-      hipblas-common
-      hipblas
-      rocminfo
-      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"
-  );
-  stdenv = torch.stdenv;
-in
-buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
+buildPythonPackage.override { stdenv = torch.stdenv; } (finalAttrs: {
   pname = "torchaudio";
   version = "2.10.0";
   pyproject = true;
@@ -124,13 +70,13 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
     sox
     torch.cxxdev
   ]
-  ++ lib.optionals stdenv.cc.isClang [ llvmPackages.openmp ];
+  ++ lib.optionals torch.stdenv.cc.isClang [ llvmPackages.openmp ];
 
   dependencies = [ torch ];
 
   preConfigure = lib.optionalString rocmSupport ''
-    export ROCM_PATH=${rocmtoolkit_joined}
-    export PYTORCH_ROCM_ARCH="${gpuTargetString}"
+    export ROCM_PATH=${torch.rocmtoolkit_joined}
+    export PYTORCH_ROCM_ARCH="${torch.gpuTargetString}"
   '';
 
   dontUseCmakeConfigure = true;
@@ -148,6 +94,7 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: {
       lib.platforms.linux ++ lib.optionals (!cudaSupport && !rocmSupport) lib.platforms.darwin;
     maintainers = with lib.maintainers; [
       GaetanLepage
+      caniko
       junjihashimoto
     ];
   };