setup.py

# Copyright 2016 Sean Robertson

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import platform
import sys

from distutils.spawn import find_executable
from os import environ
from os import path
from os import walk
from re import findall
from setuptools import setup
from setuptools.extension import Extension
import numpy as np
import platform

IS_64_BIT = sys.maxsize > 2**32
ON_WINDOWS = platform.system() == "Windows"

if ON_WINDOWS:
    LIBRARY_SUFFIXES = {"lib"}
elif platform.system() == "Darwin":
    LIBRARY_SUFFIXES = {"a", "dylib"}
else:
    LIBRARY_SUFFIXES = {"a", "so"}


# modified this bad boy from
# https://stackoverflow.com/questions/33560364/python-windows-parsing-command-lines-with-shlex
def cmdline_split(s, platform=not ON_WINDOWS):
    """Multi-platform variant of shlex.split() for command-line splitting.
    For use with subprocess, for argv injection etc. Using fast REGEX.

    platform: 'this' = auto from current platform;
              1 = POSIX;
              0 = Windows/CMD
              (other values reserved)
    """
    if platform:
        RE_CMD_LEX = r""""((?:\\["\\]|[^"])*)"|'([^']*)'|(\\.)|(&&?|\|\|?|\d?'
            r'\>|[<])|([^\s'"\\&|<>]+)|(\s+)|(.)"""
    else:
        RE_CMD_LEX = r""""((?:""|\\["\\]|[^"])*)"?()|(\\\\(?=\\*")|\\")|(&&?|\|\|?'
            r'|\d?>|[<])|([^\s"&|<>]+)|(\s+)|(.)"""
    args = []
    accu = None  # collects pieces of one arg
    for qs, qss, esc, pipe, word, white, fail in findall(RE_CMD_LEX, s):
        if word:
            pass  # most frequent
        elif esc:
            word = esc[1]
        elif white or pipe:
            if accu is not None:
                args.append(accu)
            if pipe:
                args.append(pipe)
            accu = None
            continue
        elif fail:
            raise ValueError("invalid or incomplete shell string")
        elif qs:
            word = qs.replace('\\"', '"').replace("\\\\", "\\")
            if platform == 0:
                word = word.replace('""', '"')
        else:
            word = qss  # may be even empty; must be last
        accu = (accu or "") + word
    if accu is not None:
        args.append(accu)
    return args


def mkl_setup(roots, mkl_threading=None):
    found_mkl_libs = {
        "mkl_rt": False,
        "mkl_intel": False,
        "mkl_intel_lp64": False,
        "mkl_intel_thread": False,
        "mkl_gnu_thread": False,
        "mkl_sequential": False,
        "mkl_tbb_thread": False,
    }
    blas_library_dirs = set()
    blas_includes = set()
    for root in roots:
        root = path.abspath(root)
        for root_name, _, base_names in walk(root):
            for base_name in base_names:
                if base_name.split(".")[-1] in LIBRARY_SUFFIXES:
                    library_name = base_name.rsplit(".", maxsplit=1)[0]
                    if not ON_WINDOWS:
                        if library_name.startswith("lib"):
                            library_name = library_name[3:]
                        else:
                            library_name = None  # not sure what to do with this
                else:
                    library_name = None
                if library_name in found_mkl_libs and not found_mkl_libs[library_name]:
                    found_mkl_libs[library_name] = True
                    blas_library_dirs.add(root_name)
                elif base_name == "mkl.h" and not blas_includes:
                    blas_includes.add(root_name)
    if not blas_includes:
        raise Exception("Could not find mkl.h")
    blas_library_dirs = list(blas_library_dirs)
    blas_includes = list(blas_includes)
    if mkl_threading is None:
        if found_mkl_libs["mkl_rt"]:
            mkl_threading = "dynamic"
        elif found_mkl_libs["mkl_intel_thread"]:
            mkl_threading = "intel"
        elif found_mkl_libs["mkl_tbb_thread"]:
            mkl_threading = "tbb"
        elif found_mkl_libs["mkl_sequential"]:
            mkl_threading = "sequential"
        elif found_mkl_libs["mkl_gnu_thread"]:
            mkl_threading = "gnu"
        else:
            raise Exception("Could not find a threading library for MKL")
    if mkl_threading == "dynamic":
        blas_libraries = ["mkl_rt"]
    else:
        blas_libraries = ["mkl_intel_lp64"] if IS_64_BIT else ["mkl_intel"]
        if mkl_threading in ("intel", "iomp"):
            blas_libraries += ["mkl_intel_thread", "iomp5"]
        elif mkl_threading in ("gnu", "gomp"):
            blas_libraries += ["mkl_gnu_thread", "gomp"]
        elif mkl_threading == "tbb":
            blas_libraries.append("mkl_tbb_thread")
        elif mkl_threading == "sequential":
            blas_libraries.append("mkl_sequential")
        else:
            raise Exception("Invalid MKL_THREADING_TYPE: {}".format(mkl_threading))
        blas_libraries.append("mkl_core")
    if not all(found_mkl_libs[lib] for lib in blas_libraries):
        raise Exception(
            "MKL_THREADING_TYPE=={} requires {} libs".format(
                mkl_threading, blas_libraries
            )
        )
    return {
        "BLAS_LIBRARIES": blas_libraries,
        "BLAS_LIBRARY_DIRS": blas_library_dirs,
        "BLAS_INCLUDES": blas_includes,
        # 'LD_FLAGS': ['-Wl,--no-as-needed'],
        "DEFINES": [("HAVE_MKL", None)],
    }


def blas_setup(roots, library_names, headers, extra_entries_on_success):
    library_names = dict((lib, False) for lib in library_names)
    headers = dict((header, False) for header in headers)
    library_dirs = set()
    include_dirs = set()
    for root in roots:
        root = path.abspath(root)
        for root_name, _, base_names in walk(root):
            for base_name in base_names:
                if base_name.split(".")[-1] in LIBRARY_SUFFIXES:
                    library_name = base_name.rsplit(".", maxsplit=1)[0]
                    if not ON_WINDOWS:
                        if library_name.startswith("lib"):
                            library_name = library_name[3:]
                        else:
                            library_name = None  # not sure what to do with this
                else:
                    library_name = None
                if library_name in library_names and not library_names[library_name]:
                    library_names[library_name] = True
                    library_dirs.add(root_name)
                elif base_name in headers and not headers[base_name]:
                    headers[base_name] = True
                    include_dirs.add(root_name)
    if not all(library_names.values()):
        raise Exception(
            "Could not find {}".format(
                tuple(key for key, val in library_names.items() if not val)
            )
        )
    if not all(headers.values()):
        raise Exception(
            "Could not find {}".format(
                tuple(key for key, val in headers.items() if not val)
            )
        )
    ret = dict(extra_entries_on_success)
    ret["BLAS_LIBRARIES"] = list(library_names)
    ret["BLAS_LIBRARY_DIRS"] = list(library_dirs)
    ret["BLAS_INCLUDES"] = list(include_dirs)
    return ret


def openblas_setup(roots):
    return blas_setup(
        roots,
        ("openblas",),
        ("cblas.h", "lapacke.h"),
        {"DEFINES": [("HAVE_OPENBLAS", None)]},
    )


def atlas_setup(roots):
    # FIXME(sdrobert): There are a variety of ways that atlas could've been distributed.
    # I'm assuming Fedora b/c it works with my CI.
    return blas_setup(
        roots,
        ("satlas", "tatlas"),
        ("cblas.h", "clapack.h"),
        {"DEFINES": [("HAVE_ATLAS", None)]},
    )


def lapacke_setup(roots):
    # only relies on the routines from lapack, so don't link to lapacke
    return blas_setup(
        roots,
        ("blas", "cblas", "lapack"),
        ("cblas.h", "lapacke.h"),
        {"DEFINES": [("HAVE_OPENBLAS", None)]},
    )


def clapack_setup(roots):
    return blas_setup(
        roots,
        ("blas", "lapack", "f2c", "clapack"),
        ("cblas.h", "f2c.h", "clapack.h"),
        {"DEFINES": [("HAVE_CLAPACK", None)]},
    )


def accelerate_setup():
    return {
        "DEFINES": [("HAVE_CLAPACK", None)],
        "LD_FLAGS": ["-framework", "Accelerate"],
    }


def noblas_setup():
    return {"DEFINES": [("HAVE_NOBLAS", None)]}


PYTHON_DIR = "python"
SRC_DIR = "src"
SWIG_DIR = "swig"
DEFINES = [
    ("KALDI_DOUBLEPRECISION", environ.get("KALDI_DOUBLEPRECISION", "0")),
]
LD_FLAGS = []
if platform.system() != "Windows":
    FLAGS = ["-std=c++11", "-m64" if IS_64_BIT else "-m32", "-msse", "-msse2", "-FPIC"]
    if platform.system() == "Darwin":
        FLAGS += ["-flax-vector-conversions", "-stdlib=libc++"]
        LD_FLAGS += ["-stdlib=libc++"]
    DEFINES += [
        ("_GLIBCXX_USE_CXX11_ABI", "0"),
        ("HAVE_CXXABI_H", None),
    ]
    if not environ.get("AUDITWHEEL_PLAT", "").startswith("musl"):
        DEFINES.append(("HAVE_EXECINFO_H", "1"))
    LIBRARIES = ["m", "dl"]
else:
    FLAGS = []
    LIBRARIES = []
    DEFINES += []
SRC_FILES = []
INCLUDE_DIRS = [SRC_DIR, np.get_include()]
LIBRARY_DIRS = []

if find_executable("swig"):
    SRC_FILES.append(path.join(SWIG_DIR, "pydrobert", "kaldi.i"))
elif path.exists(path.join(SWIG_DIR, "pydrobert", "kaldi_wrap.cpp")):
    print(
        "SWIG could not be found, but kaldi_wrap.cpp exists. Using that",
        file=sys.stderr,
    )
    SRC_FILES.append(path.join(SWIG_DIR, "pydrobert", "kaldi_wrap.cpp"))
else:
    print(
        "SWIG could not be found and kaldi_wrap.cpp does not exist. Cannot " "succeed",
        file=sys.stderr,
    )
    sys.exit(1)
for base_dir, _, files in walk(SRC_DIR):
    SRC_FILES += [path.join(base_dir, f) for f in files if f.endswith(".cc")]

# Adds additional libraries. Primarily for alpine libc (musllinux build),
# which doesn't package execinfo by default.
LIBRARIES += cmdline_split(environ.get("ADDITIONAL_LIBS", ""))

MKL_ROOT = cmdline_split(environ.get("MKLROOT", ""))
OPENBLAS_ROOT = cmdline_split(environ.get("OPENBLASROOT", ""))
ATLAS_ROOT = cmdline_split(environ.get("ATLASROOT", ""))
CLAPACK_ROOT = cmdline_split(environ.get("CLAPACKROOT", ""))
LAPACKE_ROOT = cmdline_split(environ.get("LAPACKEROOT", ""))
USE_ACCELERATE = environ.get("ACCELERATE", None)
NUM_BLAS_OPTS = sum(
    bool(x)
    for x in (
        MKL_ROOT,
        OPENBLAS_ROOT,
        ATLAS_ROOT,
        USE_ACCELERATE,
        CLAPACK_ROOT,
        LAPACKE_ROOT,
    )
)
if NUM_BLAS_OPTS > 1:
    raise Exception("Only one BLAS type should be specified")
elif NUM_BLAS_OPTS:
    if MKL_ROOT:
        BLAS_DICT = mkl_setup(MKL_ROOT, environ.get("MKL_THREADING_TYPE", None))
    elif OPENBLAS_ROOT:
        BLAS_DICT = openblas_setup(OPENBLAS_ROOT)
    elif ATLAS_ROOT:
        BLAS_DICT = atlas_setup(ATLAS_ROOT)
    elif CLAPACK_ROOT:
        BLAS_DICT = clapack_setup(CLAPACK_ROOT)
    elif LAPACKE_ROOT:
        BLAS_DICT = lapacke_setup(LAPACKE_ROOT)
    elif platform.system() == "Darwin":
        BLAS_DICT = accelerate_setup()
    else:
        raise Exception("Accelerate is only available on OSX")
else:
    BLAS_DICT = noblas_setup()


LIBRARIES += BLAS_DICT.get("BLAS_LIBRARIES", [])
LIBRARY_DIRS += BLAS_DICT.get("BLAS_LIBRARY_DIRS", [])
INCLUDE_DIRS += BLAS_DICT.get("BLAS_INCLUDES", [])
LD_FLAGS += BLAS_DICT.get("LD_FLAGS", [])
DEFINES += BLAS_DICT.get("DEFINES", [])
FLAGS += BLAS_DICT.get("FLAGS", [])


KALDI_LIBRARY = Extension(
    "pydrobert.kaldi._internal",
    sources=SRC_FILES,
    libraries=LIBRARIES,
    library_dirs=LIBRARY_DIRS,
    include_dirs=INCLUDE_DIRS,
    extra_compile_args=FLAGS,
    extra_link_args=LD_FLAGS,
    define_macros=DEFINES,
    swig_opts=["-py3", "-c++", "-builtin", "-castmode", "-O", "-I{}".format(SWIG_DIR)],
    language="c++",
)


setup(ext_modules=[KALDI_LIBRARY])