swendel/nixpkgs
1
0
Fork 0
mirror of https://github.com/SebastianWendel/nixpkgs.git synced 2024-09-21 12:59:04 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #220017 from azuwis/torch

Browse source 
python3Packages.torch: Fix performance problem on darwin
This commit is contained in:
Weijia Wang 2023-03-14 09:52:49 +02:00 committed by GitHub
parent 39286f0b63 0d758c2aa1
commit 1474943fd9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 8 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
pkgs/development/python-modules/torch/default.nix
View file

@ -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,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

2
pkgs/top-level/python-packages.nix
View file

@ -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;
};