diff --git a/pkgs/development/python-modules/torch/source/default.nix b/pkgs/development/python-modules/torch/source/default.nix index f16d1cb2cc82..ed985066ba62 100644 --- a/pkgs/development/python-modules/torch/source/default.nix +++ b/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 diff --git a/pkgs/development/python-modules/torchaudio/default.nix b/pkgs/development/python-modules/torchaudio/default.nix index 0245755df3e6..aa0235b46c4f 100644 --- a/pkgs/development/python-modules/torchaudio/default.nix +++ b/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 ]; };