[BUILD] Fix Warnings and Enable Warnings as Errors (#794)

python/setup.py

@@ -1,11 +1,12 @@
 import distutils
-import distutils.spawn
+import itertools
 import os
 import platform
 import re
 import shutil
 import subprocess
 import sys
+import sysconfig
 import tarfile
 import tempfile
 import urllib.request
@@ -16,6 +17,74 @@ from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
 
 
+# Logic from https://github.com/Kitware/CMake/blob/master/Modules/FindPythonLibs.cmake
+# Code from https://stackoverflow.com/questions/47423246
+def get_python_library():
+    """Get path to the python library associated with the current python
+    interpreter."""
+    # determine direct path to libpython
+    python_version = sysconfig.get_python_version()
+    python_library = sysconfig.get_config_var('LIBRARY')
+
+    # if static (or nonexistent), try to find a suitable dynamic libpython
+    if (python_library is None or os.path.splitext(python_library)[1][-2:] == '.a'):
+
+        candidate_lib_prefixes = ['', 'lib']
+
+        candidate_extensions = ['.lib', '.so', '.a']
+        if sysconfig.get_config_var('WITH_DYLD'):
+            candidate_extensions.insert(0, '.dylib')
+
+        candidate_versions = [python_version]
+        if python_version:
+            candidate_versions.append('')
+            candidate_versions.insert(
+                0, "".join(python_version.split(".")[:2]))
+
+        abiflags = getattr(sys, 'abiflags', '')
+        candidate_abiflags = [abiflags]
+        if abiflags:
+            candidate_abiflags.append('')
+
+        # Ensure the value injected by virtualenv is
+        # returned on windows.
+        # Because calling `sysconfig.get_config_var('multiarchsubdir')`
+        # returns an empty string on Linux, `du_sysconfig` is only used to
+        # get the value of `LIBDIR`.
+        libdir = distutils.sysconfig.get_config_var('LIBDIR')
+        if sysconfig.get_config_var('MULTIARCH'):
+            masd = sysconfig.get_config_var('multiarchsubdir')
+            if masd:
+                if masd.startswith(os.sep):
+                    masd = masd[len(os.sep):]
+                libdir = os.path.join(libdir, masd)
+
+        if libdir is None:
+            libdir = os.path.abspath(os.path.join(
+                sysconfig.get_config_var('LIBDEST'), "..", "libs"))
+
+        candidates = (
+            os.path.join(
+                libdir,
+                ''.join((pre, 'python', ver, abi, ext))
+            )
+            for (pre, ext, ver, abi) in itertools.product(
+                candidate_lib_prefixes,
+                candidate_extensions,
+                candidate_versions,
+                candidate_abiflags
+            )
+        )
+
+        for candidate in candidates:
+            if os.path.exists(candidate):
+                # we found a (likely alternate) libpython
+                python_library = candidate
+                break
+
+    return python_library
+
+
 # Taken from https://github.com/pytorch/pytorch/blob/master/tools/setup_helpers/env.py
 def check_env_flag(name: str, default: str = "") -> bool:
     return os.getenv(name, default).upper() in ["ON", "1", "YES", "TRUE", "Y"]
@@ -136,14 +205,19 @@ class CMakeBuild(build_ext):
         if not os.path.exists(llvm_build_dir):
             os.makedirs(llvm_build_dir)
         # python directories
-        python_include_dirs = [distutils.sysconfig.get_python_inc()] + ['/usr/local/cuda/include']
+        python_include_dir = distutils.sysconfig.get_python_inc()
+        python_link_dir = distutils.sysconfig.get_python_lib()
+        python_library = get_python_library()
         cmake_args = [
+            "-DLLVM_ENABLE_WERROR=ON",
             "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + extdir,
             "-DTRITON_BUILD_TUTORIALS=OFF",
             "-DTRITON_BUILD_PYTHON_MODULE=ON",
             # '-DPYTHON_EXECUTABLE=' + sys.executable,
             # '-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON',
-            "-DPYTHON_INCLUDE_DIRS=" + ";".join(python_include_dirs),
+            "-DPYTHON_INCLUDE_DIRS=" + python_include_dir,
+            "-DPYTHON_LINK_DIRS=" + python_link_dir,
+            "-DPYTHON_LIBRARIES=" + python_library,
             "-DLLVM_EXTERNAL_LIT=" + lit_dir
         ] + thirdparty_cmake_args