Merge pull request NixOS#220017 from azuwis/torch

python3Packages.torch: Fix performance problem on darwin

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

@@ -1,6 +1,6 @@
 { stdenv, lib, fetchFromGitHub, fetchpatch, buildPythonPackage, python,
   cudaSupport ? false, cudaPackages, magma,
-  mklDnnSupport ? true, useSystemNccl ? true,
+  useSystemNccl ? true,
   MPISupport ? false, mpi,
   buildDocs ? false,
 
@@ -9,14 +9,18 @@
 
   # Build inputs
   numactl,
-  CoreServices, libobjc,
+  Accelerate, CoreServices, libobjc,
 
   # Propagated build inputs
   numpy, pyyaml, cffi, click, typing-extensions,
 
   # Unit tests
   hypothesis, psutil,
 
+  # Disable MKLDNN on aarch64-darwin, it negatively impacts performance,
+  # this is also what official pytorch build does
+  mklDnnSupport ? !(stdenv.isDarwin && stdenv.isAarch64),
+
   # virtual pkg that consistently instantiates blas across nixpkgs
   # See https://github.com/NixOS/nixpkgs/pull/83888
   blas,
@@ -275,7 +279,7 @@ in buildPythonPackage rec {
     ++ lib.optionals rocmSupport [ openmp ]
     ++ lib.optionals (cudaSupport || rocmSupport) [ magma ]
     ++ lib.optionals stdenv.isLinux [ numactl ]
-    ++ lib.optionals stdenv.isDarwin [ CoreServices libobjc ];
+    ++ lib.optionals stdenv.isDarwin [ Accelerate CoreServices libobjc ];
 
   propagatedBuildInputs = [
     cffi

pkgs/top-level/python-packages.nix

@@ -11715,7 +11715,7 @@ self: super: with self; {
 
   torch = callPackage ../development/python-modules/torch {
     cudaSupport = pkgs.config.cudaSupport or false;
-    inherit (pkgs.darwin.apple_sdk.frameworks) CoreServices;
+    inherit (pkgs.darwin.apple_sdk.frameworks) Accelerate CoreServices;
     inherit (pkgs.darwin) libobjc;
     inherit (pkgs.llvmPackages_rocm) openmp;
   };