22import platform
33import subprocess
44from pathlib import Path
5+ import distutils .sysconfig
56
7+ from setuptools import Extension
8+ from setuptools .command .build_ext import build_ext
69import torch
7- from torch .utils .cpp_extension import (
8- CppExtension ,
9- BuildExtension as TorchBuildExtension
10- )
1110
1211__all__ = [
1312 'get_ext_modules' ,
14- 'BuildExtension ' ,
13+ 'CMakeBuild ' ,
1514]
1615
1716_THIS_DIR = Path (__file__ ).parent .resolve ()
1817_ROOT_DIR = _THIS_DIR .parent .parent .resolve ()
19- _CSRC_DIR = _ROOT_DIR / 'torchaudio' / 'csrc'
20- _TP_BASE_DIR = _ROOT_DIR / 'third_party'
21- _TP_INSTALL_DIR = _TP_BASE_DIR / 'install'
18+ _TORCHAUDIO_DIR = _ROOT_DIR / 'torchaudio'
2219
2320
2421def _get_build (var ):
@@ -38,132 +35,71 @@ def _get_build(var):
3835_BUILD_TRANSDUCER = _get_build ("BUILD_TRANSDUCER" )
3936
4037
41- def _get_eca (debug ):
42- eca = []
43- if debug :
44- eca += ["-O0" , "-g" ]
45- else :
46- eca += ["-O3" ]
47- if _BUILD_TRANSDUCER :
48- eca += ['-DBUILD_TRANSDUCER' ]
49- return eca
50-
51-
52- def _get_ela (debug ):
53- ela = []
54- if debug :
55- if platform .system () == "Windows" :
56- ela += ["/DEBUG:FULL" ]
57- else :
58- ela += ["-O0" , "-g" ]
59- else :
60- ela += ["-O3" ]
61- return ela
62-
63-
64- def _get_srcs ():
65- srcs = [_CSRC_DIR / 'pybind.cpp' ]
66- srcs += list (_CSRC_DIR .glob ('sox/**/*.cpp' ))
67- if _BUILD_TRANSDUCER :
68- srcs += [_CSRC_DIR / 'transducer.cpp' ]
69- return [str (path ) for path in srcs ]
70-
71-
72- def _get_include_dirs ():
73- dirs = [
74- str (_ROOT_DIR ),
75- ]
76- if _BUILD_SOX or _BUILD_TRANSDUCER :
77- dirs .append (str (_TP_INSTALL_DIR / 'include' ))
78- return dirs
79-
80-
81- def _get_extra_objects ():
82- libs = []
83- if _BUILD_SOX :
84- # NOTE: The order of the library listed bellow matters.
85- #
86- # (the most important thing is that dependencies come after a library
87- # e.g., sox comes first, flac/vorbis comes before ogg, and
88- # vorbisenc/vorbisfile comes before vorbis
89- libs += [
90- 'libsox.a' ,
91- 'libmad.a' ,
92- 'libFLAC.a' ,
93- 'libmp3lame.a' ,
94- 'libopusfile.a' ,
95- 'libopus.a' ,
96- 'libvorbisenc.a' ,
97- 'libvorbisfile.a' ,
98- 'libvorbis.a' ,
99- 'libogg.a' ,
100- 'libopencore-amrnb.a' ,
101- 'libopencore-amrwb.a' ,
102- ]
103- if _BUILD_TRANSDUCER :
104- libs += ['libwarprnnt.a' ]
105-
106- return [str (_TP_INSTALL_DIR / 'lib' / lib ) for lib in libs ]
107-
108-
109- def _get_libraries ():
110- return [] if _BUILD_SOX else ['sox' ]
111-
112-
113- def _get_cxx11_abi ():
114- try :
115- value = int (torch ._C ._GLIBCXX_USE_CXX11_ABI )
116- except ImportError :
117- value = 0
118- return f'-D_GLIBCXX_USE_CXX11_ABI={ value }
119-
38+ def get_ext_modules ():
39+ if platform .system () == 'Windows' :
40+ return None
41+ return [Extension (name = 'torchaudio._torchaudio' , sources = [])]
12042
121- def _build_third_party (base_build_dir ):
122- build_dir = os .path .join (base_build_dir , 'third_party' )
123- os .makedirs (build_dir , exist_ok = True )
124- subprocess .run (
125- args = [
126- 'cmake' ,
127- f"-DCMAKE_CXX_FLAGS='{ _get_cxx11_abi ()} ,
128- '-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON' ,
129- f'-DCMAKE_INSTALL_PREFIX={ _TP_INSTALL_DIR } ,
130- f'-DBUILD_SOX={ "ON" if _BUILD_SOX else "OFF" } ,
131- f'-DBUILD_TRANSDUCER={ "ON" if _BUILD_TRANSDUCER else "OFF" } ,
132- f'{ _TP_BASE_DIR } ],
133- cwd = build_dir ,
134- check = True ,
135- )
136- command = ['cmake' , '--build' , '.' ]
137- if _BUILD_TRANSDUCER :
138- command += ['--target' , 'install' ]
139- subprocess .run (
140- args = command ,
141- cwd = build_dir ,
142- check = True ,
143- )
14443
44+ # Based off of
45+ # https://github.com/pybind/cmake_example/blob/580c5fd29d4651db99d8874714b07c0c49a53f8a/setup.py
46+ class CMakeBuild (build_ext ):
47+ def run (self ):
48+ try :
49+ subprocess .check_output (['cmake' , '--version' ])
50+ except OSError :
51+ raise RuntimeError ("CMake is not available." )
52+ super ().run ()
14553
146- _EXT_NAME = 'torchaudio._torchaudio'
54+ def build_extension (self , ext ):
55+ extdir = os .path .abspath (
56+ os .path .dirname (self .get_ext_fullpath (ext .name )))
57+
58+ # required for auto-detection of auxiliary "native" libs
59+ if not extdir .endswith (os .path .sep ):
60+ extdir += os .path .sep
61+
62+ cfg = "Debug" if self .debug else "Release"
63+
64+ cmake_args = [
65+ f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={ extdir } ,
66+ f"-DCMAKE_BUILD_TYPE={ cfg } ,
67+ f"-DCMAKE_PREFIX_PATH={ torch .utils .cmake_prefix_path } ,
68+ '-DCMAKE_VERBOSE_MAKEFILE=ON' ,
69+ f"-DPython_INCLUDE_DIR={ distutils .sysconfig .get_python_inc ()} ,
70+ f"-DBUILD_SOX:BOOL={ 'ON' if _BUILD_SOX else 'OFF' } ,
71+ f"-DBUILD_TRANSDUCER:BOOL={ 'ON' if _BUILD_TRANSDUCER else 'OFF' } ,
72+ "-DBUILD_PYTHON_EXTENSION:BOOL=ON" ,
73+ "-DBUILD_LIBTORCHAUDIO:BOOL=OFF" ,
74+ ]
75+ build_args = [
14776
77+ ]
14878
149- def get_ext_modules (debug = False ):
150- if platform .system () == 'Windows' :
151- return None
152- return [
153- CppExtension (
154- _EXT_NAME ,
155- _get_srcs (),
156- libraries = _get_libraries (),
157- include_dirs = _get_include_dirs (),
158- extra_compile_args = _get_eca (debug ),
159- extra_objects = _get_extra_objects (),
160- extra_link_args = _get_ela (debug ),
161- ),
162- ]
163-
164-
165- class BuildExtension (TorchBuildExtension ):
166- def build_extension (self , ext ):
167- if ext .name == _EXT_NAME and _BUILD_SOX :
168- _build_third_party (self .build_temp )
169- super ().build_extension (ext )
79+ # Default to Ninja
80+ if 'CMAKE_GENERATOR' not in os .environ :
81+ cmake_args += ["-GNinja" ]
82+
83+ # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level
84+ # across all generators.
85+ if "CMAKE_BUILD_PARALLEL_LEVEL" not in os .environ :
86+ # self.parallel is a Python 3 only way to set parallel jobs by hand
87+ # using -j in the build_ext call, not supported by pip or PyPA-build.
88+ if hasattr (self , "parallel" ) and self .parallel :
89+ # CMake 3.12+ only.
90+ build_args += ["-j{}" .format (self .parallel )]
91+
92+ if not os .path .exists (self .build_temp ):
93+ os .makedirs (self .build_temp )
94+
95+ subprocess .check_call (
96+ ["cmake" , str (_ROOT_DIR )] + cmake_args , cwd = self .build_temp )
97+ subprocess .check_call (
98+ ["cmake" , "--build" , "." ] + build_args , cwd = self .build_temp )
99+
100+ def get_ext_filename (self , fullname ):
101+ ext_filename = super ().get_ext_filename (fullname )
102+ ext_filename_parts = ext_filename .split ('.' )
103+ without_abi = ext_filename_parts [:- 2 ] + ext_filename_parts [- 1 :]
104+ ext_filename = '.' .join (without_abi )
105+ return ext_filename
0 commit comments