better handling of compiler arguments

.Rbuildignore

@@ -21,3 +21,4 @@
 ^requirements\.txt$
 ^pyproject\.toml$
 ^a\.out$
+^a\.exe

.gitignore

@@ -22,6 +22,10 @@ poismf.egg-info/*
 
 *.so
 *.o
+*.a
+*.out
+*.exe
+*.dylib
 *.Rhistory
 *.directory
 

README.md

@@ -35,14 +35,14 @@ or if that fails:
 pip install --no-use-pep517 poismf
 ```
 
-Requires some BLAS library such as MKL (`pip install mkl-devel`) or OpenBLAS - will attempt to use the same as NumPy is using. Also requires a C compiler such as GCC or Visual Studio (in windows + conda, install Visual Studio Build Tools, and select package MSVC140 in the install options).
-
-**Note for macOS users:** on macOS, the Python version of this package will compile **without** multi-threading capabilities. This is due to default apple's redistribution of clang not providing OpenMP modules, and aliasing it to gcc which causes confusions in build scripts. If you have a non-apple version of clang with the OpenMP modules, or if you have gcc installed, you can compile this package with multi-threading enabled by setting up an environment variable `ENABLE_OMP=1`:
+**Note for macOS users:** on macOS, the Python version of this package might compile **without** multi-threading capabilities. In order to enable multi-threading support, first install OpenMP:
 ```
-export ENABLE_OMP=1
-pip install isotree
+brew install libomp
 ```
-(Alternatively, can also pass argument enable-omp to the setup.py file: `python setup.py install enable-omp`)
+And then reinstall this package: `pip install --force-reinstall poismf`.
+
+
+Requires some BLAS library such as MKL (`pip install mkl-devel`) or OpenBLAS, and speed will depend mostly on the BLAS implementation. Also requires a C compiler such as GCC or Visual Studio (in windows + conda, install Visual Studio Build Tools, and select package MSVC140 in the install options).
 
 For any installation problems, please open an issue in GitHub providing information about your system (OS, BLAS, C compiler) and Python installation.
 

poismf/__init__.py

@@ -221,6 +221,13 @@ def __init__(self, k = 50, method = "tncg",
         assert nthreads > 0
         assert isinstance(nthreads, int) or isinstance(nthreads, np.int64)
 
+        if (nthreads > 1) and not (c_funs_double._get_has_openmp()):
+            msg_omp  = "Attempting to use more than 1 thread, but "
+            msg_omp += "package was built without multi-threading "
+            msg_omp += "support - see the project's GitHub page for "
+            msg_omp += "more information."
+            warnings.warn(msg_omp)
+
         if isinstance(random_state, np.random.RandomState):
             random_state = random_state.randint(np.iinfo(np.int32).max)
         elif random_state is None:

poismf/poismf_c_wrapper.pxi

@@ -9,6 +9,7 @@ cdef extern from "../src/poismf.h":
         tncg = 1
         cg = 2
         pg = 3
+    bint get_has_openmp() nogil
     int run_poismf(
         real_t *A, real_t *Xr, sparse_ix *Xr_indptr, sparse_ix *Xr_indices,
         real_t *B, real_t *Xc, sparse_ix *Xc_indptr, sparse_ix *Xc_indices,
@@ -61,6 +62,9 @@ cdef extern from "../src/poismf.h":
         size_t n_top, size_t n, int nthreads
     ) nogil
 
+def _get_has_openmp():
+    return get_has_openmp()
+
 def _run_poismf(
     np.ndarray[real_t, ndim=1] Xr,
     np.ndarray[size_t, ndim=1] Xr_indices,

setup.py

@@ -6,52 +6,98 @@
     from distutils.core import setup
     from distutils.extension import Extension
 import numpy
-from sys import platform
-import os
+import sys, os, subprocess, warnings, re
 from Cython.Distutils import build_ext
-from sys import platform
-import sys, os
 
+found_omp = True
+def set_omp_false():
+    global found_omp
+    found_omp = False
 
 ## https://stackoverflow.com/questions/724664/python-distutils-how-to-get-a-compiler-that-is-going-to-be-used
 class build_ext_subclass( build_ext ):
     def build_extensions(self):
-        compiler = self.compiler.compiler_type
-        if compiler == 'msvc': # visual studio is 'special'
+        if self.compiler.compiler_type == 'msvc':
             for e in self.extensions:
                 e.extra_compile_args += ['/O2', '/openmp']
         else:
+            self.add_march_native()
+            self.add_openmp_linkage()
+
             for e in self.extensions:
-                e.extra_compile_args += ['-O3', '-fopenmp', '-march=native', '-std=c99']
-                e.extra_link_args += ['-fopenmp']
+                # e.extra_compile_args += ['-O3', '-fopenmp', '-march=native', '-std=c99']
+                # e.extra_link_args += ['-fopenmp']
+
                 # e.extra_compile_args += ["-fsanitize=address", "-static-libasan", "-ggdb"]
                 # e.extra_link_args += ["-fsanitize=address", "-static-libasan"]
                 # e.extra_compile_args += ["-ggdb"]
-
-        ## Note: apple will by default alias 'gcc' to 'clang', and will ship its own "special"
-        ## 'clang' which has no OMP support and nowadays will purposefully fail to compile when passed
-        ## '-fopenmp' flags. If you are using mac, and have an OMP-capable compiler,
-        ## comment out the code below, or set 'use_omp' to 'True'.
-        if not use_omp:
-            for e in self.extensions:
-                e.extra_compile_args = [arg for arg in e.extra_compile_args if arg != '-fopenmp']
-                e.extra_link_args    = [arg for arg in e.extra_link_args    if arg != '-fopenmp']
+
+                e.extra_compile_args += ['-O3', '-std=c99']
 
         build_ext.build_extensions(self)
 
-use_omp = (("enable-omp" in sys.argv)
-           or ("-enable-omp" in sys.argv)
-           or ("--enable-omp" in sys.argv))
-if use_omp:
-    sys.argv = [a for a in sys.argv if a not in ("enable-omp", "-enable-omp", "--enable-omp")]
-if os.environ.get('ENABLE_OMP') is not None:
-    use_omp = True
-if platform[:3] != "dar":
-    use_omp = True
+    def add_march_native(self):
+        arg_march_native = "-march=native"
+        arg_mcpu_native = "-mcpu=native"
+        if self.test_supports_compile_arg(arg_march_native):
+            for e in self.extensions:
+                e.extra_compile_args.append(arg_march_native)
+        elif self.test_supports_compile_arg(arg_mcpu_native):
+            for e in self.extensions:
+                e.extra_compile_args.append(arg_mcpu_native)
+
+    def add_openmp_linkage(self):
+        arg_omp1 = "-fopenmp"
+        arg_omp2 = "-qopenmp"
+        arg_omp3 = "-xopenmp"
+        args_apple_omp = ["-Xclang", "-fopenmp", "-lomp"]
+        if self.test_supports_compile_arg(arg_omp1):
+            for e in self.extensions:
+                e.extra_compile_args.append(arg_omp1)
+                e.extra_link_args.append(arg_omp1)
+        elif (sys.platform[:3].lower() == "dar") and self.test_supports_compile_arg(args_apple_omp):
+            for e in self.extensions:
+                e.extra_compile_args += ["-Xclang", "-fopenmp"]
+                e.extra_link_args += ["-lomp"]
+        elif self.test_supports_compile_arg(arg_omp2):
+            for e in self.extensions:
+                e.extra_compile_args.append(arg_omp2)
+                e.extra_link_args.append(arg_omp2)
+        elif self.test_supports_compile_arg(arg_omp3):
+            for e in self.extensions:
+                e.extra_compile_args.append(arg_omp3)
+                e.extra_link_args.append(arg_omp3)
+        else:
+            set_omp_false()
 
-### Shorthand for apple computer:
-### uncomment line below
-# use_omp = True
+    def test_supports_compile_arg(self, comm):
+        is_supported = False
+        try:
+            if not hasattr(self.compiler, "compiler"):
+                return False
+            if not isinstance(comm, list):
+                comm = [comm]
+            print("--- Checking compiler support for option '%s'" % " ".join(comm))
+            fname = "poismf_compiler_testing.c"
+            with open(fname, "w") as ftest:
+                ftest.write(u"int main(int argc, char**argv) {return 0;}\n")
+            try:
+                cmd = [self.compiler.compiler[0]]
+            except:
+                cmd = list(self.compiler.compiler)
+            val_good = subprocess.call(cmd + [fname])
+            try:
+                val = subprocess.call(cmd + comm + [fname])
+                is_supported = (val == val_good)
+            except:
+                is_supported = False
+        except:
+            pass
+        try:
+            os.remove(fname)
+        except:
+            pass
+        return is_supported
 
 
 from_rtd = os.environ.get('READTHEDOCS') == 'True'
@@ -62,7 +108,7 @@ def build_extensions(self):
         author = 'David Cortes',
         author_email = '[email protected]',
         url = 'https://github.com/david-cortes/poismf',
-        version = '0.3.1-1',
+        version = '0.3.1-2',
         install_requires = ['numpy', 'pandas>=0.24', 'cython', 'scipy'],
         description = 'Fast and memory-efficient Poisson factorization for sparse count matrices',
         cmdclass = {'build_ext': build_ext_subclass},
@@ -82,22 +128,24 @@ def build_extensions(self):
             ]
     )
 
-    if not use_omp:
-        import warnings
-        apple_msg  = "\n\n\nMacOS detected. Package will be built without multi-threading capabilities, "
-        apple_msg += "due to Apple's lack of OpenMP support in default clang installs. In order to enable it, "
-        apple_msg += "install the package directly from GitHub: https://www.github.com/david-cortes/poismf\n"
-        apple_msg += "Using 'python setup.py install enable-omp'. "
-        apple_msg += "You'll also need an OpenMP-capable compiler.\n\n\n"
-        warnings.warn(apple_msg)
+    if not found_omp:
+        omp_msg  = "\n\n\nCould not detect OpenMP. Package will be built without multi-threading capabilities. "
+        omp_msg += " To enable multi-threading, first install OpenMP"
+        if (sys.platform[:3] == "dar"):
+            omp_msg += " - for macOS: 'brew install libomp'\n"
+        else:
+            omp_msg += " modules for your compiler. "
+
+        omp_msg += "Then reinstall this package from scratch: 'pip install --force-reinstall poismf'.\n"
+        warnings.warn(omp_msg)
 else:
     setup(
         name  = "poismf",
         packages = ["poismf"],
         author = 'David Cortes',
         author_email = '[email protected]',
         url = 'https://github.com/david-cortes/poismf',
-        version = '0.3.1-1',
+        version = '0.3.1-2',
         install_requires = ['numpy', 'scipy', 'pandas>=0.24', 'cython'],
         description = 'Fast and memory-efficient Poisson factorization for sparse count matrices',
     )

src/poismf.c

@@ -51,6 +51,16 @@ void set_interrup_global_variable(int s)
     }
 }
 
+/* For making sure that it's multi-threaded */
+bool get_has_openmp()
+{
+    #ifdef _OPENMP
+    return true;
+    #else
+    return false;
+    #endif
+}
+
 /* Helper functions */
 #define nonneg(x) (((x) > 0.)? (x) : 0.)