Unverified Commit 1474943f authored by Weijia Wang's avatar Weijia Wang Committed by GitHub
Browse files

Merge pull request #220017 from azuwis/torch 

python3Packages.torch: Fix performance problem on darwin
parents 39286f0b 0d758c2a

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

+7 −3
Original line number Diff line number Diff line 
{ stdenv, lib, fetchFromGitHub, fetchpatch, buildPythonPackage, python,

  cudaSupport ? false, cudaPackages, magma,

  mklDnnSupport ? true, useSystemNccl ? true,

  useSystemNccl ? true,

  MPISupport ? false, mpi,

  buildDocs ? false,
@@ -9,7 +9,7 @@ 


  # Build inputs

  numactl,

  CoreServices, libobjc,

  Accelerate, CoreServices, libobjc,



  # Propagated build inputs

  numpy, pyyaml, cffi, click, typing-extensions,
@@ -17,6 +17,10 @@ 
  # 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

+1 −1
Original line number Diff line number Diff line
@@ -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;

  };