Re-enable MKL support for FFTs

CMakeLists.txt

@@ -53,6 +53,7 @@ if(NOT TOAST_STATIC_DEPS AND NOT $ENV{TOAST_STATIC_DEPS} STREQUAL "")
   set(TOAST_STATIC_DEPS $ENV{TOAST_STATIC_DEPS})
 endif()
 
+# OpenMP
 find_package(OpenMP)
 
 if(TOAST_STATIC_DEPS)
@@ -62,14 +63,26 @@ if(TOAST_STATIC_DEPS)
     set(SUITESPARSE_USE_STATIC_LIBS TRUE)
 endif()
 
-# Look for FFTW first, in case we are using MKL for BLAS/LAPACK.  MKL sometimes
-# includes an FFTW interface that we want to avoid using.
-find_package(FFTW)
+# First look for MKL, since that will provide both FFT support and BLAS/LAPACK
+find_package(MKL)
 
-find_package(BLAS)
+if(MKL_FOUND)
+    # Use MKL for BLAS / LAPACK
+    set(BLAS_LIBRARIES "${MKL_LIBRARIES}")
+    set(LAPACK_LIBRARIES "${MKL_LIBRARIES}")
+    set(BLAS_FOUND TRUE)
+    set(LAPACK_FOUND TRUE)
+else()
+    # Search for FFTW instead
+    find_package(FFTW)
+    if(NOT FFTW_FOUND)
+        message(FATAL_ERROR "Could not find a supported FFT library (MKL or FFTW)")
+    endif()
+    find_package(BLAS)
+    find_package(LAPACK)
+endif()
 
 if(BLAS_FOUND)
-    find_package(LAPACK)
     if(LAPACK_FOUND)
         find_package(LAPACKnames)
     else()

cmake/FindMKL.cmake

@@ -0,0 +1,99 @@
+# - Find INTEL MKL Single Dynamic Library
+#
+# First version:
+#   Copyright (c) 2020, Theodore Kisner
+#
+# Usage:
+#   find_package(MKL [REQUIRED])
+#
+# It sets the following variables:
+#   MKL_FOUND                  ... true if MKL is found on the system
+#   MKL_LIBRARIES              ... full paths to all found MKL libraries
+#   MKL_INCLUDE_DIRS           ... MKL include directory paths
+#
+# The following variables will be checked by the function
+#   MKL_DIR                    ... if set, look for MKL in this directory.
+#   environment $MKLROOT       ... if set, look for MKL in this directory.
+#
+
+# Set default value of MKL_DIR.
+
+if(NOT MKL_DIR)
+    if(WIN32)
+        if(DEFINED ENV{MKLProductDir})
+            set(MKL_DIR "$ENV{MKLProductDir}/mkl")
+        else()
+            set(
+                MKL_DIR
+                "C:/Program Files (x86)/IntelSWTools/compilers_and_libraries/windows/mkl"
+            )
+        endif()
+    else()
+        set(MKL_DIR "/opt/intel/mkl")
+        if(DEFINED ENV{MKLROOT})
+            # This is the standard MKL environment variable.
+            set(MKL_DIR "$ENV{MKLROOT}")
+        endif()
+    endif()
+endif()
+
+find_package(Threads)
+
+# First look in user-specified or default location
+
+if(NOT MKL_FOUND)
+    set(MKL_LIB_DIR "${MKL_DIR}/lib")
+    find_library(
+        MKL_RT_LIB
+        NAMES "mkl_rt"
+        PATHS "${MKL_LIB_DIR}"
+        PATH_SUFFIXES "intel64"
+        NO_DEFAULT_PATH
+    )
+    # find includes
+    find_path(MKL_INCLUDE_DIRS
+        NAMES "mkl_dfti.h"
+        PATHS ${MKL_DIR}
+        PATH_SUFFIXES "include"
+        NO_DEFAULT_PATH
+    )
+    if(MKL_RT_LIB)
+        set(MKL_FOUND TRUE)
+    endif()
+endif()
+
+# If we did not find it, search everywhere
+
+if(NOT MKL_FOUND)
+    find_library(
+        MKL_RT_LIB
+        NAMES "mkl_rt"
+    )
+    # find includes
+    find_path(MKL_INCLUDE_DIRS
+        NAMES "mkl_dfti.h"
+    )
+    if(MKL_RT_LIB)
+        set(MKL_FOUND TRUE)
+    endif()
+endif()
+
+if(MKL_FOUND)
+    set(MKL_LIBRARIES
+	${MKL_RT_LIB} ${CMAKE_THREAD_LIBS_INIT} -lm -l${CMAKE_DL_LIBS}
+	CACHE STRING "MKL dynamic libraries"
+    )
+endif()
+
+include(FindPackageHandleStandardArgs)
+
+find_package_handle_standard_args(MKL
+    REQUIRED_VARS MKL_INCLUDE_DIRS MKL_LIBRARIES
+    HANDLE_COMPONENTS
+)
+
+mark_as_advanced(
+    MKL_INCLUDE_DIRS
+    MKL_LIBRARIES
+)
+

cmake/FindSuiteSparse.cmake

@@ -291,19 +291,23 @@ endmacro()
 unset(SUITESPARSE_FOUND_REQUIRED_VARS)
 
 # BLAS.
-find_package(BLAS QUIET)
 if (NOT BLAS_FOUND)
-  suitesparse_report_not_found(
-    "Did not find BLAS library (required for SuiteSparse).")
-endif (NOT BLAS_FOUND)
+  find_package(BLAS QUIET)
+  if (NOT BLAS_FOUND)
+    suitesparse_report_not_found(
+      "Did not find BLAS library (required for SuiteSparse).")
+  endif (NOT BLAS_FOUND)
+endif()
 list(APPEND SUITESPARSE_FOUND_REQUIRED_VARS BLAS_FOUND)
 
 # LAPACK.
-find_package(LAPACK QUIET)
 if (NOT LAPACK_FOUND)
-  suitesparse_report_not_found(
-    "Did not find LAPACK library (required for SuiteSparse).")
-endif (NOT LAPACK_FOUND)
+  find_package(LAPACK QUIET)
+  if (NOT LAPACK_FOUND)
+    suitesparse_report_not_found(
+      "Did not find LAPACK library (required for SuiteSparse).")
+  endif (NOT LAPACK_FOUND)
+endif()
 list(APPEND SUITESPARSE_FOUND_REQUIRED_VARS LAPACK_FOUND)
 
 suitesparse_find_component(AMD REQUIRED FILES amd.h LIBRARIES amd)

src/libtoast/CMakeLists.txt

@@ -105,18 +105,20 @@ if(CHOLMOD_FOUND)
     endif (METIS_FOUND)
 endif(CHOLMOD_FOUND)
 
-if(LAPACK_FOUND)
-    target_compile_definitions(toast PRIVATE HAVE_LAPACK=1)
-    target_compile_definitions(toast PRIVATE "LAPACK_NAMES_${LAPACK_NAMES}")
-    set_target_properties(toast PROPERTIES LINK_FLAGS
-        "${LAPACK_LINKER_FLAGS} ${BLAS_LINKER_FLAGS}"
-    )
-    target_link_libraries(toast "${LAPACK_LIBRARIES}")
-    target_link_libraries(toast "${BLAS_LIBRARIES}")
-endif(LAPACK_FOUND)
-
-# Link to FFTW *after* BLAS / LAPACK.  This should override any FFTW symbols contained
-# in MKL if MKL used for BLAS / LAPACK.
+# Link to our FFT libraries
+
+if(MKL_FOUND)
+    target_compile_definitions(toast PRIVATE HAVE_MKL=1)
+    target_include_directories(toast PUBLIC "${MKL_INCLUDE_DIRS}")
+    target_link_libraries(toast "${MKL_LIBRARIES}")
+    # If we are using GNU compilers, we will need to internally set the threading layer
+    # to use GNU.  Otherwise, assume we are using Intel threads.
+    if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+        target_compile_definitions(toast PRIVATE USE_MKL_GNU_THREADS=1)
+    else()
+        target_compile_definitions(toast PRIVATE USE_MKL_INTEL_THREADS=1)
+    endif()
+endif(MKL_FOUND)
 
 if(FFTW_FOUND)
     target_compile_definitions(toast PRIVATE HAVE_FFTW=1)
@@ -127,6 +129,18 @@ if(FFTW_FOUND)
     endif(FFTW_DOUBLE_THREADS_LIB_FOUND)
 endif(FFTW_FOUND)
 
+# LAPACK / BLAS
+
+if(LAPACK_FOUND)
+    target_compile_definitions(toast PRIVATE HAVE_LAPACK=1)
+    target_compile_definitions(toast PRIVATE "LAPACK_NAMES_${LAPACK_NAMES}")
+    set_target_properties(toast PROPERTIES LINK_FLAGS
+        "${LAPACK_LINKER_FLAGS} ${BLAS_LINKER_FLAGS}"
+    )
+    target_link_libraries(toast "${LAPACK_LIBRARIES}")
+    target_link_libraries(toast "${BLAS_LIBRARIES}")
+endif(LAPACK_FOUND)
+
 # Installation
 
 #install(DIRECTORY include/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})

src/libtoast/src/toast_math_fft.cpp

@@ -9,6 +9,7 @@
 #include <toast/math_fft_fftw.hpp>
 #include <toast/math_fft_mkl.hpp>
 
+#include <cstring>
 #include <cmath>
 #include <vector>
 

src/libtoast/src/toast_sys_environment.cpp

@@ -15,6 +15,12 @@ extern "C" {
 # include <omp.h>
 #endif // ifdef _OPENMP
 
+#ifdef HAVE_MKL
+extern "C" {
+#include <mkl.h>
+}
+#endif
+
 // These variables are autogenerated and compiled
 // into the library by the version.cmake script
 extern "C" {
@@ -147,6 +153,15 @@ toast::Environment::Environment() {
         func_timers_ = true;
     }
 
+    // Select MKL threading layer
+    #ifdef HAVE_MKL
+    # ifdef USE_MKL_GNU_THREADS
+    auto td = mkl_set_threading_layer(MKL_THREADING_GNU);
+    # else // ifdef USE_MKL_GNU_THREADS
+    auto td = mkl_set_threading_layer(MKL_THREADING_INTEL);
+    # endif // ifdef USE_MKL_GNU_THREADS
+    #endif // ifdef HAVE_MKL
+
     // OpenMP
     max_threads_ = 1;
     #ifdef _OPENMP

src/toast/_libtoast_sys.cpp

@@ -456,4 +456,6 @@ void init_sys(py::module & m) {
                 None
 
         )");
+
+    auto env = toast::Environment::get();
 }