Optional direct linkage to Intel MKL with fine-grained multithreading control

.github/workflows/ci.yml

@@ -41,6 +41,27 @@ jobs:
         with:
           token: ${{secrets.CODECOV_TOKEN}}
 
+  linux-build-mkl:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-20.04]
+        python-version: ["3.12"]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install and Test
+        run: ./.github/workflows/run_ci_mkl.sh
+      - name: Upload to codecov
+        if: matrix.python-version == '3.8'
+        uses: codecov/[email protected]
+        with:
+          token: ${{secrets.CODECOV_TOKEN}}
+
   linux-build-aarch64:
     runs-on: ${{ matrix.os }}
     strategy:

.github/workflows/ci_linux_mkl/build_pyscf.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+set -e
+
+cd ./pyscf/lib
+curl -L "https://github.com/pyscf/pyscf-build-deps/blob/master/pyscf-2.4a-deps.tar.gz?raw=true" | tar xzf -
+mkdir build; cd build
+cmake -DUSE_MKL=ON -DBUILD_LIBXC=OFF -DBUILD_XCFUN=ON -DBUILD_LIBCINT=OFF -DXCFUN_MAX_ORDER=4 ..
+make -j4
+cd ..
+rm -Rf build
+cd ../..

.github/workflows/ci_linux_mkl/deps_apt.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+sudo apt-get -qq install \
+    gcc \
+    libblas-dev \
+    cmake \
+    curl

.github/workflows/ci_linux_mkl/python_deps.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+python -m pip install --upgrade pip
+pip install "numpy!=1.16,!=1.17" "scipy!=1.5" h5py pytest pytest-cov pytest-timer
+pip install mkl mkl-devel mkl-include
+pip install pyberny
+pip install --no-deps pyscf-dispersion
+
+version=$(python -c 'import sys; version=sys.version_info[:2]; print("{0}.{1}".format(*version))')
+if [ $version != '3.12' ]; then
+    pip install geometric
+    pip install spglib
+fi
+
+#cppe
+if [ "$RUNNER_OS" == "Linux" ] && [ $version != "3.12" ]; then
+    pip install cppe
+fi

.github/workflows/run_ci_mkl.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+set -e
+
+if [ "$RUNNER_OS" == "Linux" ]; then
+    os='linux'
+elif [ "$RUNNER_OS" == "macOS" ]; then
+    os='macos'
+else
+    echo "$RUNNER_OS not supported"
+    exit 1
+fi
+
+./.github/workflows/ci_"$os"_mkl/deps_apt.sh
+./.github/workflows/ci_"$os"_mkl/python_deps.sh
+./.github/workflows/ci_"$os"_mkl/build_pyscf.sh
+./.github/workflows/run_tests.sh

pyscf/lib/CMakeLists.txt

@@ -66,17 +66,39 @@ if(EXISTS "${PROJECT_SOURCE_DIR}/cmake.arch.inc")
   include("${PROJECT_SOURCE_DIR}/cmake.arch.inc")
 endif()
 
-if (NOT BLAS_LIBRARIES)
-#enable_language(Fortran)
-find_package(BLAS)
-check_function_exists(ffsll HAVE_FFS)
-endif()
+option(USE_MKL, "Use Intel MKL API and headers" OFF)
 
-if (NOT BLAS_LIBRARIES)
-  message(FATAL_ERROR "A required library with BLAS API not found.")
+if(USE_MKL)
+  # Finer threading control is available thanks to
+  # MKL threading control support functions, so it's okay
+  # not to use sequential BLAS.
+  if(CMAKE_C_COMPILER_ID STREQUAL "Intel")
+    set(MKL_THREADING "intel_thread")
+  elseif(CMAKE_C_COMPILER_ID STREQUAL "GNU")
+    set(MKL_THREADING "gnu_thread")
+  endif()
+  set(MKL_INTERFACE "lp64")
+  find_package(MKL CONFIG REQUIRED)
+  if(NOT MKL_FOUND)
+    message(FATAL_ERROR "MKLConfig.cmake not found. Make sure CMake can find it.")
+  endif()
+  set(BLAS_LIBRARIES MKL::MKL)
+  add_compile_definitions(PYSCF_USE_MKL)
 else()
-  message(STATUS "BLAS libraries: ${BLAS_LIBRARIES}")
-endif()
+  # If MKL is not used, try to find BLAS.
+  if (NOT BLAS_LIBRARIES)
+  #enable_language(Fortran)
+  find_package(BLAS)
+  check_function_exists(ffsll HAVE_FFS)
+  endif()
+
+  if (NOT BLAS_LIBRARIES)
+    message(FATAL_ERROR "A required library with BLAS API not found.")
+  else()
+    message(STATUS "BLAS libraries: ${BLAS_LIBRARIES}")
+  endif()
+endif() # USE_MKL
+
 # if unable to find mkl library, just create BLAS_LIBRARIES here, e.g.
 # set(BLAS_LIBRARIES "-L/path/to/mkl/lib -lmkl_intel_lp64 -lmkl_sequential -lmkl_core -lpthread -lmkl_avx -lm")
 # or