Skip to content

Commit

Permalink
improve python setup script
Browse files Browse the repository at this point in the history
  • Loading branch information
@david-cortes
david-cortes committed Mar 19, 2023
1 parent af69f65 commit 01de61f
Show file tree
Hide file tree
Showing 4 changed files with 89 additions and 35 deletions.
2 changes: 1 addition & 1 deletion README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ pip install --no-use-pep517 poismf
```
brew install libomp
```
And then reinstall this package: `pip install --force-reinstall poismf`.
And then reinstall this package: `pip install --upgrade --no-deps --force-reinstall poismf`.
** *
**IMPORTANT:** the setup script will try to add compilation flag `-march=native`. This instructs the compiler to tune the package for the CPU in which it is being installed (by e.g. using AVX instructions if available), but the result might not be usable in other computers. If building a binary wheel of this package or putting it into a docker image which will be used in different machines, this can be overriden either by (a) defining an environment variable `DONT_SET_MARCH=1`, or by (b) manually supplying compilation `CFLAGS` as an environment variable with something related to architecture. For maximum compatibility (but slowest speed), it's possible to do something like this:

Expand Down
1 change: 1 addition & 0 deletions poismf/cfuns_double.pyx
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
#cython: language_level=2
import ctypes

ctypedef double real_t
Expand Down
1 change: 1 addition & 0 deletions poismf/cfuns_float.pyx
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
#cython: language_level=2
import ctypes

ctypedef float real_t
Expand Down
120 changes: 86 additions & 34 deletions setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,37 +14,51 @@ def set_omp_false():
global found_omp
found_omp = False

## Modify this to make the output of the compilation tests more verbose
silent_tests = not (("verbose" in sys.argv)
or ("-verbose" in sys.argv)
or ("--verbose" in sys.argv))

## Workaround for python<=3.9 on windows
try:
EXIT_SUCCESS = os.EX_OK
except AttributeError:
EXIT_SUCCESS = 0

## 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):
if self.compiler.compiler_type == 'msvc':
for e in self.extensions:
e.extra_compile_args += ['/O2', '/openmp', '/fp:fast']
e.extra_compile_args += ['/O2', '/openmp', '/GL', '/fp:fast']
else:
if not self.check_for_variable_dont_set_march() and not self.check_cflags_contain_arch():
self.add_march_native()
self.add_openmp_linkage()
self.add_no_math_errno()
self.add_no_trapping_math()
self.add_ffp_contract_fast()
if sys.platform[:3].lower() != "win":
self.add_link_time_optimization()
self.add_O3()
self.add_std_c99()
self.add_link_time_optimization()

for e in self.extensions:
# for e in self.extensions:
# 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"]

e.extra_compile_args += ['-O3', '-std=c99']

build_ext.build_extensions(self)

def check_cflags_contain_arch(self):
if "CFLAGS" in os.environ:
arch_list = ["-march", "-mcpu", "-mtune", "-msse", "-msse2", "-msse3", "-mssse3", "-msse4", "-msse4a", "-msse4.1", "-msse4.2", "-mavx", "-mavx2"]
arch_list = [
"-march", "-mcpu", "-mtune", "-msse", "-msse2", "-msse3",
"-mssse3", "-msse4", "-msse4a", "-msse4.1", "-msse4.2",
"-mavx", "-mavx2", "-mavx512"
]
for flag in arch_list:
if flag in os.environ["CFLAGS"]:
return True
Expand All @@ -54,21 +68,21 @@ def check_for_variable_dont_set_march(self):
return "DONT_SET_MARCH" in os.environ

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)
args_march_native = ["-march=native", "-mcpu=native"]
for arg_march_native in args_march_native:
if self.test_supports_compile_arg(arg_march_native):
for e in self.extensions:
e.extra_compile_args.append(arg_march_native)
break

def add_link_time_optimization(self):
arg_lto = "-flto"
if self.test_supports_compile_arg(arg_lto):
for e in self.extensions:
e.extra_compile_args.append(arg_lto)
e.extra_link_args.append(arg_lto)
args_lto = ["-flto=auto", "-flto"]
for arg_lto in args_lto:
if self.test_supports_compile_arg(arg_lto):
for e in self.extensions:
e.extra_compile_args.append(arg_lto)
e.extra_link_args.append(arg_lto)
break

def add_no_math_errno(self):
arg_fnme = "-fno-math-errno"
Expand All @@ -91,33 +105,72 @@ def add_ffp_contract_fast(self):
e.extra_compile_args.append(arg_ffpc)
e.extra_link_args.append(arg_ffpc)

def add_O3(self):
arg_O3 = "-O3"
if self.test_supports_compile_arg(arg_O3):
for e in self.extensions:
e.extra_compile_args.append(arg_O3)
e.extra_link_args.append(arg_O3)

def add_std_c99(self):
arg_std_c99 = "-std=c99"
if self.test_supports_compile_arg(arg_std_c99):
for e in self.extensions:
e.extra_compile_args.append(arg_std_c99)
e.extra_link_args.append(arg_std_c99)

def add_openmp_linkage(self):
arg_omp1 = "-fopenmp"
arg_omp2 = "-qopenmp"
arg_omp3 = "-xopenmp"
arg_omp2 = "-fopenmp=libomp"
args_omp2 = ["-fopenmp=libomp", "-lomp"]
arg_omp4 = "-qopenmp"
arg_omp5 = "-xopenmp"
is_apple = sys.platform[:3].lower() == "dar"
args_apple_omp = ["-Xclang", "-fopenmp", "-lomp"]
args_apple_omp2 = ["-Xclang", "-fopenmp", "-L/usr/local/lib", "-lomp", "-I/usr/local/include"]
has_brew_omp = False
if is_apple:
res_brew_pref = subprocess.run(["brew", "--prefix", "libomp"], capture_output=silent_tests)
if res_brew_pref.returncode == EXIT_SUCCESS:
has_brew_omp = True
brew_omp_prefix = res_brew_pref.stdout.decode().strip()
args_apple_omp3 = ["-Xclang", "-fopenmp", f"-L{brew_omp_prefix}/lib", "-lomp", f"-I{brew_omp_prefix}/include"]


if self.test_supports_compile_arg(arg_omp1, with_omp=True):
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, with_omp=True):
elif is_apple and self.test_supports_compile_arg(args_apple_omp, with_omp=True):
for e in self.extensions:
e.extra_compile_args += ["-Xclang", "-fopenmp"]
e.extra_link_args += ["-lomp"]
elif (sys.platform[:3].lower() == "dar") and self.test_supports_compile_arg(args_apple_omp2, with_omp=True):
elif is_apple and self.test_supports_compile_arg(args_apple_omp2, with_omp=True):
for e in self.extensions:
e.extra_compile_args += ["-Xclang", "-fopenmp"]
e.extra_link_args += ["-L/usr/local/lib", "-lomp"]
e.include_dirs += ["/usr/local/include"]
elif is_apple and has_brew_omp and self.test_supports_compile_arg(args_apple_omp3, with_omp=True):
for e in self.extensions:
e.extra_compile_args += ["-Xclang", "-fopenmp"]
e.extra_link_args += [f"-L{brew_omp_prefix}/lib", "-lomp"]
e.include_dirs += [f"{brew_omp_prefix}/include"]
elif self.test_supports_compile_arg(arg_omp2, with_omp=True):
for e in self.extensions:
e.extra_compile_args.append(arg_omp2)
e.extra_link_args.append(arg_omp2)
e.extra_compile_args += ["-fopenmp=libomp"]
e.extra_link_args += ["-fopenmp"]
elif self.test_supports_compile_arg(arg_omp3, with_omp=True):
for e in self.extensions:
e.extra_compile_args.append(arg_omp3)
e.extra_link_args.append(arg_omp3)
e.extra_compile_args += ["-fopenmp=libomp"]
e.extra_link_args += ["-fopenmp", "-lomp"]
elif self.test_supports_compile_arg(arg_omp4, with_omp=True):
for e in self.extensions:
e.extra_compile_args.append(arg_omp4)
e.extra_link_args.append(arg_omp4)
elif self.test_supports_compile_arg(arg_omp5, with_omp=True):
for e in self.extensions:
e.extra_compile_args.append(arg_omp5)
e.extra_link_args.append(arg_omp5)
else:
set_omp_false()

Expand All @@ -139,13 +192,12 @@ def test_supports_compile_arg(self, comm, with_omp=False):
cmd = self.compiler.compiler
except Exception:
cmd = self.compiler.compiler
val_good = subprocess.call(cmd + [fname])
if with_omp:
with open(fname, "w") as ftest:
ftest.write(u"#include <omp.h>\nint main(int argc, char**argv) {return 0;}\n")
try:
val = subprocess.call(cmd + comm + [fname])
is_supported = (val == val_good)
val = subprocess.run(cmd + comm + [fname], capture_output=silent_tests).returncode
is_supported = (val == EXIT_SUCCESS)
except Exception:
is_supported = False
except Exception:
Expand All @@ -165,7 +217,7 @@ def test_supports_compile_arg(self, comm, with_omp=False):
author = 'David Cortes',
author_email = 'david.cortes.rivera@gmail.com',
url = 'https://github.com/david-cortes/poismf',
version = '0.4.0-2',
version = '0.4.0-3',
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},
Expand Down Expand Up @@ -193,7 +245,7 @@ def test_supports_compile_arg(self, comm, with_omp=False):
else:
omp_msg += " modules for your compiler. "

omp_msg += "Then reinstall this package from scratch: 'pip install --force-reinstall poismf'.\n"
omp_msg += "Then reinstall this package from scratch: 'pip install --upgrade --no-deps --force-reinstall poismf'.\n"
warnings.warn(omp_msg)
else:
setup(
Expand All @@ -202,7 +254,7 @@ def test_supports_compile_arg(self, comm, with_omp=False):
author = 'David Cortes',
author_email = 'david.cortes.rivera@gmail.com',
url = 'https://github.com/david-cortes/poismf',
version = '0.4.0-2',
version = '0.4.0-3',
install_requires = ['numpy', 'scipy', 'pandas>=0.24', 'cython'],
description = 'Fast and memory-efficient Poisson factorization for sparse count matrices',
)

0 comments on commit 01de61f

Please sign in to comment.