|
1 | 1 | import os |
2 | | -import re |
3 | 2 | import shutil |
4 | 3 | import subprocess |
5 | | -import sysconfig |
6 | 4 | import sys |
7 | 5 |
|
8 | | -from setuptools import setup |
| 6 | +# TODO: update once there is replacement for clean: |
| 7 | +# https://github.com/pypa/setuptools/discussions/2838 |
| 8 | +from distutils import log # pylint: disable=[deprecated-module] |
| 9 | +from distutils.dir_util import remove_tree # pylint: disable=[deprecated-module] |
| 10 | +from distutils.command.clean import clean as _clean # pylint: disable=[deprecated-module] |
| 11 | + |
| 12 | +from setuptools import setup, Extension |
| 13 | +from setuptools.command.build_ext import build_ext as _build_ext |
9 | 14 |
|
10 | 15 | import torch |
11 | 16 |
|
12 | | -ipex_cmake_prefix_path = "" |
13 | | -USE_IPEX_OPTION = os.getenv("USE_IPEX", "1") |
14 | | -if USE_IPEX_OPTION == "1": |
15 | | - import intel_extension_for_pytorch |
16 | | - ipex_cmake_prefix_path = f";{intel_extension_for_pytorch.cmake_prefix_path}" |
| 17 | + |
| 18 | +class CMakeExtension(Extension): |
| 19 | + |
| 20 | + def __init__(self, name): |
| 21 | + # don't invoke the original build_ext for this special extension |
| 22 | + super().__init__(name, sources=[]) |
17 | 23 |
|
18 | 24 |
|
19 | 25 | class CMakeBuild(): |
20 | 26 |
|
21 | | - def __init__(self): |
| 27 | + def __init__(self, debug=False, dry_run=False): |
22 | 28 | self.current_dir = os.path.abspath(os.path.dirname(__file__)) |
23 | 29 | self.build_temp = self.current_dir + "/build/temp" |
24 | 30 | self.extdir = self.current_dir + "/triton_kernels_benchmark" |
| 31 | + self.build_type = self.get_build_type(debug) |
| 32 | + self.cmake_prefix_paths = [torch.utils.cmake_prefix_path] |
| 33 | + self.use_ipex = False |
| 34 | + self.dry_run = dry_run |
| 35 | + |
| 36 | + def get_build_type(self, debug): |
| 37 | + DEBUG_OPTION = os.getenv("DEBUG", "0") |
| 38 | + return "Debug" if debug or (DEBUG_OPTION == "1") else "Release" |
25 | 39 |
|
26 | 40 | def run(self): |
27 | | - try: |
28 | | - out = subprocess.check_output(["cmake", "--version"]) |
29 | | - except OSError as error: |
30 | | - raise RuntimeError("CMake must be installed") from error |
| 41 | + self.check_ipex() |
| 42 | + self.build_extension() |
31 | 43 |
|
32 | | - match = re.search(r"version\s*(?P<major>\d+)\.(?P<minor>\d+)([\d.]+)?", out.decode()) |
33 | | - cmake_major, cmake_minor = int(match.group("major")), int(match.group("minor")) |
34 | | - if (cmake_major, cmake_minor) < (3, 18): |
35 | | - raise RuntimeError("CMake >= 3.18.0 is required") |
| 44 | + def check_ipex(self): |
| 45 | + self.use_ipex = os.getenv("USE_IPEX", "1") == "1" |
| 46 | + if not self.use_ipex: |
| 47 | + return |
| 48 | + try: |
| 49 | + import intel_extension_for_pytorch |
| 50 | + except ImportError: |
| 51 | + log.warn("ipex is not installed trying to build without ipex") |
| 52 | + self.use_ipex = False |
| 53 | + return |
| 54 | + self.cmake_prefix_paths.append(intel_extension_for_pytorch.cmake_prefix_path) |
36 | 55 |
|
37 | | - self.build_extension() |
| 56 | + def check_call(self, *popenargs, **kwargs): |
| 57 | + log.info(" ".join(popenargs[0])) |
| 58 | + if not self.dry_run: |
| 59 | + subprocess.check_call(*popenargs, **kwargs) |
38 | 60 |
|
39 | 61 | def build_extension(self): |
40 | 62 | ninja_dir = shutil.which("ninja") |
41 | 63 | # create build directories |
42 | 64 | if not os.path.exists(self.build_temp): |
43 | 65 | os.makedirs(self.build_temp) |
44 | | - # python directories |
45 | | - python_include_dir = sysconfig.get_path("platinclude") |
46 | 66 | cmake_args = [ |
47 | 67 | "-G", |
48 | 68 | "Ninja", # Ninja is much faster than make |
49 | 69 | "-DCMAKE_MAKE_PROGRAM=" + |
50 | 70 | ninja_dir, # Pass explicit path to ninja otherwise cmake may cache a temporary path |
51 | | - f"-DCMAKE_PREFIX_PATH={torch.utils.cmake_prefix_path}{ipex_cmake_prefix_path}", |
52 | | - f"-DUSE_IPEX={USE_IPEX_OPTION}", |
53 | | - "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON", |
54 | | - "-DCMAKE_ARCHIVE_OUTPUT_DIRECTORY=" + self.extdir, |
55 | | - "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + self.extdir, |
56 | | - "-DPython3_EXECUTABLE:FILEPATH=" + sys.executable, |
57 | | - "-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON", |
58 | | - "-DPYTHON_INCLUDE_DIRS=" + python_include_dir, |
| 71 | + "-DCMAKE_PREFIX_PATH=" + ";".join(self.cmake_prefix_paths), |
| 72 | + f"-DUSE_IPEX={int(self.use_ipex)}", |
| 73 | + "-DCMAKE_INSTALL_PREFIX=" + self.extdir, |
| 74 | + "-DPython3_ROOT_DIR:FILEPATH=" + sys.exec_prefix, |
| 75 | + "-DCMAKE_VERBOSE_MAKEFILE=TRUE", |
59 | 76 | "-DCMAKE_C_COMPILER=icx", |
60 | 77 | "-DCMAKE_CXX_COMPILER=icpx", |
| 78 | + "-DCMAKE_BUILD_TYPE=" + self.build_type, |
| 79 | + "-S", |
| 80 | + self.current_dir, |
| 81 | + "-B", |
| 82 | + self.build_temp, |
61 | 83 | ] |
62 | 84 |
|
63 | | - # configuration |
64 | | - build_type = "Debug" |
65 | | - build_args = ["--config", build_type] |
66 | | - cmake_args += ["-DCMAKE_BUILD_TYPE=" + build_type] |
67 | 85 | max_jobs = os.getenv("MAX_JOBS", str(2 * os.cpu_count())) |
68 | | - build_args += ["-j" + max_jobs] |
| 86 | + build_args = [ |
| 87 | + "--build", |
| 88 | + self.build_temp, |
| 89 | + "-j" + max_jobs, |
| 90 | + ] |
| 91 | + |
| 92 | + install_args = [ |
| 93 | + "--build", |
| 94 | + self.build_temp, |
| 95 | + "--target", |
| 96 | + "install", |
| 97 | + ] |
69 | 98 |
|
70 | 99 | env = os.environ.copy() |
71 | | - cmake_dir = self.build_temp |
72 | | - subprocess.check_call(["cmake", self.current_dir] + cmake_args, cwd=cmake_dir, env=env) |
73 | | - subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=cmake_dir) |
| 100 | + self.check_call(["cmake"] + cmake_args, env=env) |
| 101 | + self.check_call(["cmake"] + build_args) |
| 102 | + self.check_call(["cmake"] + install_args) |
| 103 | + |
| 104 | + def clean(self): |
| 105 | + if os.path.exists(self.build_temp): |
| 106 | + remove_tree(self.build_temp, dry_run=self.dry_run) |
| 107 | + else: |
| 108 | + log.warn("'%s' does not exist -- can't clean it", os.path.relpath(self.build_temp, |
| 109 | + os.path.dirname(__file__))) |
| 110 | + |
74 | 111 |
|
| 112 | +class build_ext(_build_ext): |
| 113 | + |
| 114 | + def run(self): |
| 115 | + cmake = CMakeBuild(debug=self.debug, dry_run=self.dry_run) |
| 116 | + cmake.run() |
| 117 | + super().run() |
| 118 | + |
| 119 | + |
| 120 | +class clean(_clean): |
| 121 | + |
| 122 | + def run(self): |
| 123 | + cmake = CMakeBuild(dry_run=self.dry_run) |
| 124 | + cmake.clean() |
| 125 | + super().run() |
75 | 126 |
|
76 | | -cmake = CMakeBuild() |
77 | | -cmake.run() |
78 | 127 |
|
79 | 128 | setup(name="triton-kernels-benchmark", packages=[ |
80 | 129 | "triton_kernels_benchmark", |
81 | 130 | ], package_dir={ |
82 | 131 | "triton_kernels_benchmark": "triton_kernels_benchmark", |
83 | | -}, package_data={"triton_kernels_benchmark": ["xetla_kernel.so"]}) |
| 132 | +}, package_data={"triton_kernels_benchmark": ["xetla_kernel.cpython-*.so"]}, cmdclass={ |
| 133 | + "build_ext": build_ext, |
| 134 | + "clean": clean, |
| 135 | +}, ext_modules=[CMakeExtension("triton_kernels_benchmark")]) |
0 commit comments