Allow setting OpenMP flags using env var

Author: thilinarmtb

setup.py

@@ -29,6 +29,15 @@
 BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1'
 
 
+def get_openmp_flag():
+    if os.getenv("OPENMP_FLAG") is not None:
+        return os.getenv("OPENMP_FLAG")
+    if sys.platform == 'win32':
+        return '/openmp'
+    else:
+        return '-fopenmp'
+
+
 def get_extensions():
     extensions = []
 
@@ -52,11 +61,8 @@ def get_extensions():
         info = parallel_info()
         if ('backend: OpenMP' in info and 'OpenMP not found' not in info
                 and sys.platform != 'darwin'):
-            extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP']
-            if sys.platform == 'win32':
-                extra_compile_args['cxx'] += ['/openmp']
-            else:
-                extra_compile_args['cxx'] += ['-fopenmp']
+            extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP',
+                                          get_openmp_flag()]
         else:
             print('Compiling without OpenMP...')