34 lines
940 B
Python
34 lines
940 B
Python
import os
|
|
import glob
|
|
import paddle
|
|
from paddle.utils.cpp_extension import CppExtension, CUDAExtension, setup
|
|
|
|
|
|
def get_extensions():
|
|
root_dir = os.path.dirname(os.path.abspath(__file__))
|
|
ext_root_dir = os.path.join(root_dir, 'csrc')
|
|
sources = []
|
|
for ext_name in os.listdir(ext_root_dir):
|
|
ext_dir = os.path.join(ext_root_dir, ext_name)
|
|
source = glob.glob(os.path.join(ext_dir, '*.cc'))
|
|
kwargs = dict()
|
|
if paddle.device.is_compiled_with_cuda():
|
|
source += glob.glob(os.path.join(ext_dir, '*.cu'))
|
|
|
|
if not source:
|
|
continue
|
|
|
|
sources += source
|
|
|
|
if paddle.device.is_compiled_with_cuda():
|
|
extension = CUDAExtension(
|
|
sources, extra_compile_args={'cxx': ['-DPADDLE_WITH_CUDA']})
|
|
else:
|
|
extension = CppExtension(sources)
|
|
|
|
return extension
|
|
|
|
|
|
if __name__ == "__main__":
|
|
setup(name='ext_op', ext_modules=get_extensions())
|