Merge pull request #4011 from ComputeCanada/fftw-mpi

set $FFT(W)_LIB_DIR to imkl-FFTW's lib path in build environment if usempi toolchain option is enabled

Author: boegel

easybuild/toolchains/fft/intelfftw.py

@@ -49,7 +49,10 @@ def _set_fftw_variables(self):
         if not hasattr(self, 'BLAS_LIB_DIR'):
             raise EasyBuildError("_set_fftw_variables: IntelFFT based on IntelMKL (no BLAS_LIB_DIR found)")
 
+        imklroot = get_software_root(self.FFT_MODULE_NAME[0])
         imklver = get_software_version(self.FFT_MODULE_NAME[0])
+        self.FFT_LIB_DIR = self.BLAS_LIB_DIR
+        self.FFT_INCLUDE_DIR = [os.path.join(d, 'fftw') for d in self.BLAS_INCLUDE_DIR]
 
         picsuff = ''
         if self.options.get('pic', None):
@@ -69,54 +72,47 @@ def _set_fftw_variables(self):
             raise EasyBuildError(error_msg)
 
         interface_lib = "fftw3xc%s%s" % (compsuff, picsuff)
-        fftw_libs = [interface_lib]
-        cluster_interface_lib = None
+        fft_lib_dirs = [os.path.join(imklroot, d) for d in self.FFT_LIB_DIR]
+
+        def fftw_lib_exists(libname):
+            """Helper function to check whether FFTW library with specified name exists."""
+            return any(os.path.exists(os.path.join(d, "lib%s.a" % libname)) for d in fft_lib_dirs)
+
+        # interface libs can be optional:
+        # MKL >= 10.2 include fftw3xc and fftw3xf interfaces in LIBBLAS=libmkl_gf/libmkl_intel
+        # See https://software.intel.com/en-us/articles/intel-mkl-main-libraries-contain-fftw3-interfaces
+        # The cluster interface libs (libfftw3x_cdft*) can be omitted if the toolchain does not provide MPI-FFTW
+        # interfaces.
+        fftw_libs = []
+        if fftw_lib_exists(interface_lib) or LooseVersion(imklver) < LooseVersion("10.2"):
+            fftw_libs = [interface_lib]
+
         if self.options.get('usempi', False):
             # add cluster interface for recent imkl versions
-            if LooseVersion(imklver) >= LooseVersion('10.3'):
+            # only get cluster_interface_lib from seperate module imkl-FFTW, rest via libmkl_gf/libmkl_intel
+            imklfftwroot = get_software_root('imkl-FFTW')
+            if LooseVersion(imklver) >= LooseVersion('10.3') and (fftw_libs or imklfftwroot):
                 suff = picsuff
                 if LooseVersion(imklver) >= LooseVersion('11.0.2'):
                     suff = bitsuff + suff
                 cluster_interface_lib = 'fftw3x_cdft%s' % suff
                 fftw_libs.append(cluster_interface_lib)
             fftw_libs.append("mkl_cdft_core")  # add cluster dft
             fftw_libs.extend(self.variables['LIBBLACS'].flatten())  # add BLACS; use flatten because ListOfList
+            if imklfftwroot:
+                fft_lib_dirs += [os.path.join(imklfftwroot, 'lib')]
+                self.FFT_LIB_DIR = [os.path.join(imklfftwroot, 'lib')]
 
         fftw_mt_libs = fftw_libs + [x % self.BLAS_LIB_MAP for x in self.BLAS_LIB_MT]
 
         self.log.debug('fftw_libs %s' % fftw_libs.__repr__())
         fftw_libs.extend(self.variables['LIBBLAS'].flatten())  # add BLAS libs (contains dft)
         self.log.debug('fftw_libs %s' % fftw_libs.__repr__())
 
-        self.FFT_LIB_DIR = self.BLAS_LIB_DIR
-        self.FFT_INCLUDE_DIR = [os.path.join(d, 'fftw') for d in self.BLAS_INCLUDE_DIR]
-
         # building the FFTW interfaces is optional,
         # so make sure libraries are there before FFT_LIB is set
-        imklroot = get_software_root(self.FFT_MODULE_NAME[0])
-        fft_lib_dirs = [os.path.join(imklroot, d) for d in self.FFT_LIB_DIR]
-        imklfftwroot = get_software_root('imkl-FFTW')
-        if imklfftwroot:
-            # only get cluster_interface_lib from seperate module imkl-FFTW, rest via libmkl_gf/libmkl_intel
-            fft_lib_dirs += [os.path.join(imklfftwroot, 'lib')]
-            fftw_libs.remove(interface_lib)
-            fftw_mt_libs.remove(interface_lib)
-
-        def fftw_lib_exists(libname):
-            """Helper function to check whether FFTW library with specified name exists."""
-            return any(os.path.exists(os.path.join(d, "lib%s.a" % libname)) for d in fft_lib_dirs)
-
-        if not fftw_lib_exists(interface_lib) and LooseVersion(imklver) >= LooseVersion("10.2"):
-            # interface libs can be optional:
-            # MKL >= 10.2 include fftw3xc and fftw3xf interfaces in LIBBLAS=libmkl_gf/libmkl_intel
-            # See https://software.intel.com/en-us/articles/intel-mkl-main-libraries-contain-fftw3-interfaces
-            # The cluster interface libs (libfftw3x_cdft*) can be omitted if the toolchain does not provide MPI-FFTW
-            # interfaces.
-            fftw_libs = [lib for lib in fftw_libs if lib not in [interface_lib, cluster_interface_lib]]
-            fftw_mt_libs = [lib for lib in fftw_mt_libs if lib not in [interface_lib, cluster_interface_lib]]
-
         # filter out libraries from list of FFTW libraries to check for if they are not provided by Intel MKL
-        check_fftw_libs = [lib for lib in fftw_libs if lib not in ['dl', 'gfortran']]
+        check_fftw_libs = [lib for lib in fftw_libs + fftw_mt_libs if lib not in ['dl', 'gfortran']]
 
         missing_fftw_libs = [lib for lib in check_fftw_libs if not fftw_lib_exists(lib)]
         if missing_fftw_libs:
@@ -128,5 +124,4 @@ def fftw_lib_exists(libname):
                 raise EasyBuildError(msg)
         else:
             self.FFT_LIB = fftw_libs
-
-        self.FFT_LIB_MT = fftw_mt_libs
+            self.FFT_LIB_MT = fftw_mt_libs

test/framework/toolchain.py

@@ -1115,6 +1115,9 @@ def test_fft_env_vars_intel(self):
         self.assertEqual(tc.get_variable('LIBFFT'), libfft)
         self.assertEqual(tc.get_variable('LIBFFT_MT'), libfft_mt)
 
+        fft_lib_dir = os.path.join(modules.get_software_root('imkl'), 'mkl/2021.4.0/lib/intel64')
+        self.assertEqual(tc.get_variable('FFT_LIB_DIR'), fft_lib_dir)
+
         tc = self.get_toolchain('intel', version='2021b')
         tc.set_options({'usempi': True})
         tc.prepare()
@@ -1138,6 +1141,9 @@ def test_fft_env_vars_intel(self):
         libfft_mt += '-Wl,-Bdynamic -liomp5 -lpthread'
         self.assertEqual(tc.get_variable('LIBFFT_MT'), libfft_mt)
 
+        fft_lib_dir = os.path.join(modules.get_software_root('imkl-FFTW'), 'lib')
+        self.assertEqual(tc.get_variable('FFT_LIB_DIR'), fft_lib_dir)
+
     def test_fosscuda(self):
         """Test whether fosscuda is handled properly."""
         tc = self.get_toolchain("fosscuda", version="2018a")
@@ -1185,6 +1191,7 @@ def setup_sandbox_for_intel_fftw(self, moddir, imklver='2018.1.163'):
 
         mkl_libs = ['mkl_cdft_core', 'mkl_blacs_intelmpi_lp64']
         mkl_libs += ['mkl_intel_lp64', 'mkl_sequential', 'mkl_core', 'mkl_intel_ilp64']
+        mkl_libs += ['mkl_intel_thread', 'mkl_pgi_thread']
         fftw_libs = ['fftw3xc_intel', 'fftw3xc_pgi']
         if LooseVersion(imklver) >= LooseVersion('11'):
             fftw_libs.extend(['fftw3x_cdft_ilp64', 'fftw3x_cdft_lp64'])