Commit 0d758c2a authored by Zhong Jianxin's avatar Zhong Jianxin
Browse files

python3Packages.torch: Disable MKLDNN on aarch64-darwin 

It negatively impacts performance, this is also what official pytorch
build does.

In my test in a Macbook Pro M1 2020, using [pianotrans][1] to transcribe
[cut_liszt.opus][2], transcription time:

Enable mkldnn and disable Accelerate.framework: ~88s

Disable mkldnn and disable Accelerate.framework: ~21s

Disable mkldnn and enable Accelerate.framework: ~9s

The final result is close to using torch-bin.

See also https://github.com/NixOS/nixpkgs/pull/219104

[1]: https://github.com/azuwis/pianotrans
[2]: https://github.com/azuwis/pianotrans/raw/master/test/cut_liszt.opus
parent 4e59252a

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

+5 −1
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,
@@ -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,