pythonPackages.{numpy,scipy,numexpr}: support MKL as BLAS

This adds support building with MKL.
NixOS · Oct 20, 2018 · 277b73a · 277b73a
1 parent eee82ae
commit 277b73a
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 16 deletions.
diff --git a/doc/languages-frameworks/python.section.md b/doc/languages-frameworks/python.section.md
@@ -1079,8 +1079,7 @@ To modify only a Python package set instead of a whole Python derivation, use th
 Use the following overlay template:
 
 ```nix
-self: super:
-{
+self: super: {
   python = super.python.override {
     packageOverrides = python-self: python-super: {
       zerobin = python-super.zerobin.overrideAttrs (oldAttrs: {
@@ -1095,6 +1094,25 @@ self: super:
 }
 ```
 
+### How to use Intel's MKL with numpy and scipy?
+
+A `site.cfg` is created that configures BLAS based on the `blas` parameter
+of the `numpy` derivation. By passing in `mkl`, `numpy` and packages depending
+on `numpy` will be built with `mkl`.
+
+The following is an overlay that configures `numpy` to use `mkl`:
+```nix
+self: super: {
+  python36 = super.python36.override {
+    packageOverrides = python-self: python-super: {
+      numpy = python-super.numpy.override {
+        blas = super.pkgs.mkl;
+      };
+    };
+  };
+}
+```
+
 ## Contributing
 
 ### Contributing guidelines

diff --git a/pkgs/development/python-modules/numexpr/default.nix b/pkgs/development/python-modules/numexpr/default.nix
@@ -3,6 +3,7 @@
 , fetchPypi
 , python
 , numpy
+, llvmPackages ? null
 }:
 
 buildPythonPackage rec {
@@ -14,6 +15,17 @@ buildPythonPackage rec {
     sha256 = "ee8bc7201aa2f1962c67d27c326a11eef9df887d7b87b1278a1d4e722bf44375";
   };
 
+  # Remove existing site.cfg, use the one we built for numpy.
+  # Somehow openmp needs to be added to LD_LIBRARY_PATH
+  # https://software.intel.com/en-us/forums/intel-system-studio/topic/611682
+  preBuild = ''
+    rm site.cfg
+    ln -s ${numpy.cfg} site.cfg
+    export LD_LIBRARY_PATH=${llvmPackages.openmp}/lib
+  '';
+
+  buildInputs = [] ++ lib.optional (numpy.blasImplementation == "mkl") llvmPackages.openmp;
+
   propagatedBuildInputs = [ numpy ];
 
   # Run the test suite.

diff --git a/pkgs/development/python-modules/numpy/default.nix b/pkgs/development/python-modules/numpy/default.nix
@@ -1,6 +1,20 @@
-{ stdenv, lib, fetchPypi, python, buildPythonPackage, isPyPy, gfortran, pytest, blas }:
+{ stdenv, lib, fetchPypi, python, buildPythonPackage, isPyPy, gfortran, pytest, blas, writeTextFile }:
 
-buildPythonPackage rec {
+let
+  blasImplementation = lib.nameFromURL blas.name "-";
+  cfg = writeTextFile {
+    name = "site.cfg";
+    text = (lib.generators.toINI {} {
+      "${blasImplementation}" = {
+        include_dirs = "${blas}/include";
+        library_dirs = "${blas}/lib";
+      } // lib.optionalAttrs (blasImplementation == "mkl") {
+        mkl_libs = "mkl_rt";
+        lapack_libs = "";
+      };
+    });
+  };
+in buildPythonPackage rec {
   pname = "numpy";
   version = "1.15.2";
 
@@ -39,12 +53,7 @@ buildPythonPackage rec {
   '';
 
   preBuild = ''
-    echo "Creating site.cfg file..."
-    cat << EOF > site.cfg
-    [openblas]
-    include_dirs = ${blas}/include
-    library_dirs = ${blas}/lib
-    EOF
+    ln -s ${cfg} site.cfg
   '';
 
   enableParallelBuilding = true;
@@ -59,8 +68,11 @@ buildPythonPackage rec {
 
   passthru = {
     blas = blas;
+    inherit blasImplementation cfg;
   };
 
+  doCheck = blasImplementation != "mkl";
+
   # Disable two tests
   # - test_f2py: f2py isn't yet on path.
   # - test_large_file_support: takes a long time and can cause the machine to run out of disk space

diff --git a/pkgs/development/python-modules/scipy/default.nix b/pkgs/development/python-modules/scipy/default.nix
@@ -29,12 +29,7 @@ buildPythonPackage rec {
   '';
 
   preBuild = ''
-    echo "Creating site.cfg file..."
-    cat << EOF > site.cfg
-    [openblas]
-    include_dirs = ${numpy.blas}/include
-    library_dirs = ${numpy.blas}/lib
-    EOF
+    ln -s ${numpy.cfg} site.cfg
   '';
 
   enableParallelBuilding = true;