From 18a3a2041667332bc2ea8a76d349230b6b785a1f Mon Sep 17 00:00:00 2001
From: Mikalai Chaly <mikalai_chaly@epam.com>
Date: Mon, 9 Mar 2020 10:40:45 +0300
Subject: [PATCH] Turn on OpenMP compiler option in Windows.

---
 README.md |  8 +++++++
 setup.py  | 67 ++++++++++++++++++++++++++++++++++++++++---------------
 2 files changed, 57 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index 23732176..c3494453 100644
--- a/README.md
+++ b/README.md
@@ -83,3 +83,11 @@ Pull requests are welcome. To install for development:
 4. LightFM uses [black](https://github.com/ambv/black) (version `18.6b4`) to enforce code formatting.
 
 When making changes to the `.pyx` extension files, you'll need to run `python setup.py cythonize` in order to produce the extension `.c` files before running `pip install -e .`.
+
+### Windows specific
+Note for Windows users: to build LightFM from source code,
+[Microsoft Visual C++ build tools](http://go.microsoft.com/fwlink/?LinkId=691126) 
+are required. 
+
+MSVC compiler [supports OpenMP](https://docs.microsoft.com/en-us/cpp/parallel/openmp/openmp-in-visual-cpp?view=vs-2019),
+so this option is turned on by default in Windows.
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 5c27b207..9176ea32 100644
--- a/setup.py
+++ b/setup.py
@@ -6,33 +6,65 @@
 
 from setuptools import Command, Extension, setup
 from setuptools.command.test import test as TestCommand
+from setuptools.command.build_ext import build_ext
 
 # Import version even when extensions are not yet built
 __builtins__.__LIGHTFM_SETUP__ = True
 from lightfm import __version__ as version  # NOQA
 
 
-def define_extensions(use_openmp):
-    compile_args = ['-ffast-math']
+class LightfmBuildExt(build_ext):
+    """
+    Configures compilation options depending on compiler.
+    """
+
+    def make_compiler_args(self):
+        compiler = self.compiler.compiler_type
+        compile_args = []
+        link_args = []
+
+        if compiler == 'msvc':
+            if use_openmp:
+                compile_args.append('-openmp')
+
+        elif compiler in ('gcc', 'unix'):
+            compile_args.append('-ffast-math')
 
-    # There are problems with illegal ASM instructions
-    # when using the Anaconda distribution (at least on OSX).
-    # This could be because Anaconda uses its own assembler?
-    # To work around this we do not add -march=native if we
-    # know we're dealing with Anaconda
-    if 'anaconda' not in sys.version.lower():
-        compile_args.append('-march=native')
+            # There are problems with illegal ASM instructions
+            # when using the Anaconda distribution (at least on OSX).
+            # This could be because Anaconda uses its own assembler?
+            # To work around this we do not add -march=native if we
+            # know we're dealing with Anaconda
+            if 'anaconda' not in sys.version.lower():
+                compile_args.append('-march=native')
 
+            if use_openmp:
+                compile_args.append('-fopenmp')
+                link_args.append('-fopenmp')
+
+        print('Use openmp:', use_openmp)
+        print('Compiler:', compiler)
+        print('Compile args:', compile_args)
+        print('Link args:', link_args)
+
+        return compile_args, link_args
+
+    def build_extensions(self):
+        compile_args, link_args = self.make_compiler_args()
+        for extension in self.extensions:
+            extension.extra_compile_args.extend(compile_args)
+            extension.extra_link_args.extend(link_args)
+        super().build_extensions()
+
+
+def define_extensions(use_openmp):
     if not use_openmp:
         print('Compiling without OpenMP support.')
         return [Extension("lightfm._lightfm_fast_no_openmp",
-                          ['lightfm/_lightfm_fast_no_openmp.c'],
-                          extra_compile_args=compile_args)]
+                          ['lightfm/_lightfm_fast_no_openmp.c'])]
     else:
         return [Extension("lightfm._lightfm_fast_openmp",
-                          ['lightfm/_lightfm_fast_openmp.c'],
-                          extra_link_args=["-fopenmp"],
-                          extra_compile_args=compile_args + ['-fopenmp'])]
+                          ['lightfm/_lightfm_fast_openmp.c'])]
 
 
 class Cythonize(Command):
@@ -95,8 +127,7 @@ def run(self):
         cythonize([Extension("lightfm._lightfm_fast_no_openmp",
                              ['lightfm/_lightfm_fast_no_openmp.pyx']),
                    Extension("lightfm._lightfm_fast_openmp",
-                             ['lightfm/_lightfm_fast_openmp.pyx'],
-                             extra_link_args=['-fopenmp'])])
+                             ['lightfm/_lightfm_fast_openmp.pyx'])])
 
 
 class Clean(Command):
@@ -144,7 +175,7 @@ def run_tests(self):
         sys.exit(errno)
 
 
-use_openmp = not sys.platform.startswith('darwin') and not sys.platform.startswith('win')
+use_openmp = not sys.platform.startswith('darwin')
 
 setup(
     name='lightfm',
@@ -157,7 +188,7 @@ def run_tests(self):
     package_data={'': ['*.c']},
     install_requires=['numpy', 'scipy>=0.17.0', 'requests', 'scikit-learn'],
     tests_require=['pytest', 'requests', 'scikit-learn'],
-    cmdclass={'test': PyTest, 'cythonize': Cythonize, 'clean': Clean},
+    cmdclass={'test': PyTest, 'cythonize': Cythonize, 'clean': Clean, 'build_ext': LightfmBuildExt},
     author='Lyst Ltd (Maciej Kula)',
     author_email='data@ly.st',
     license='MIT',