From 2c974cc316bce4054bdf28d1f6b4c3bb8bd99d75 Mon Sep 17 00:00:00 2001 From: Aurelius84 Date: Fri, 29 Jan 2021 16:51:33 +0800 Subject: [PATCH] =?UTF-8?q?=E3=80=90CustomOp=E3=80=91support=20setup.py=20?= =?UTF-8?q?to=20compile=20custom=20op=20(#30753)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../fluid/tests/custom_op/CMakeLists.txt | 6 +- .../fluid/tests/custom_op/cpp_extension.py | 179 +++++++++++++++ .../fluid/tests/custom_op/extension_utils.py | 216 ++++++++++++++++++ python/paddle/fluid/tests/custom_op/setup.py | 49 ++++ .../fluid/tests/custom_op/test_custom_op.py | 12 +- .../custom_op/test_custom_op_with_setup.py | 33 +++ 6 files changed, 489 insertions(+), 6 deletions(-) create mode 100644 python/paddle/fluid/tests/custom_op/cpp_extension.py create mode 100644 python/paddle/fluid/tests/custom_op/extension_utils.py create mode 100644 python/paddle/fluid/tests/custom_op/setup.py create mode 100644 python/paddle/fluid/tests/custom_op/test_custom_op_with_setup.py diff --git a/python/paddle/fluid/tests/custom_op/CMakeLists.txt b/python/paddle/fluid/tests/custom_op/CMakeLists.txt index ef3b39ef5c5..85d38c7548b 100644 --- a/python/paddle/fluid/tests/custom_op/CMakeLists.txt +++ b/python/paddle/fluid/tests/custom_op/CMakeLists.txt @@ -22,9 +22,9 @@ set_property(TARGET relu_op_shared PROPERTY LINK_LIBRARIES ${TARGET_LIBRARIES} file(GLOB TEST_OPS RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}" "test_*.py") string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}") -# for coverage -LIST(REMOVE_ITEM TEST_OPS test_custom_op) - foreach(src ${TEST_OPS}) py_test(${src} SRCS ${src}.py) endforeach() + +# Compiling .so will cost some time, but running process is very fast. +set_tests_properties(test_custom_op_with_setup PROPERTIES TIMEOUT 180) diff --git a/python/paddle/fluid/tests/custom_op/cpp_extension.py b/python/paddle/fluid/tests/custom_op/cpp_extension.py new file mode 100644 index 00000000000..e1243f00185 --- /dev/null +++ b/python/paddle/fluid/tests/custom_op/cpp_extension.py @@ -0,0 +1,179 @@ +# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. +# +# 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 os +import six +import sys +import copy +import setuptools +from setuptools.command.build_ext import build_ext + +from extension_utils import find_cuda_home, normalize_extension_kwargs, add_compile_flag +from extension_utils import is_cuda_file, prepare_unix_cflags, add_std_without_repeat, get_build_directory + +IS_WINDOWS = os.name == 'nt' +CUDA_HOME = find_cuda_home() + + +def CppExtension(name, sources, *args, **kwargs): + """ + Returns setuptools.CppExtension instance for setup.py to make it easy + to specify compile flags while build C++ custommed op kernel. + """ + kwargs = normalize_extension_kwargs(kwargs, use_cuda=False) + + return setuptools.Extension(name, sources, *args, **kwargs) + + +def CUDAExtension(name, sources, *args, **kwargs): + """ + Returns setuptools.CppExtension instance for setup.py to make it easy + to specify compile flags while build CUDA custommed op kernel. + """ + kwargs = normalize_extension_kwargs(kwargs, use_cuda=True) + + return setuptools.Extension(name, sources, *args, **kwargs) + + +class BuildExtension(build_ext, object): + """ + For setuptools.cmd_class. + """ + + @classmethod + def with_options(cls, **options): + ''' + Returns a BuildExtension subclass that support to specific use-defined options. + ''' + + class cls_with_options(cls): + def __init__(self, *args, **kwargs): + kwargs.update(options) + cls.__init__(self, *args, **kwargs) + + return cls_with_options + + def __init__(self, *args, **kwargs): + super(BuildExtension, self).__init__(*args, **kwargs) + self.no_python_abi_suffix = kwargs.get("no_python_abi_suffix", False) + + def initialize_options(self): + super(BuildExtension, self).initialize_options() + # update options here + # FIXME(Aurelius84): for unittest + self.build_lib = './' + + def finalize_options(self): + super(BuildExtension, self).finalize_options() + + def build_extensions(self): + self._check_abi() + for extension in self.extensions: + # check settings of compiler + if isinstance(extension.extra_compile_args, dict): + for compiler in ['cxx', 'nvcc']: + if compiler not in extension.extra_compile_args: + extension.extra_compile_args[compiler] = [] + # add determine compile flags + add_compile_flag(extension, '-std=c++11') + # add_compile_flag(extension, '-lpaddle_framework') + + # Consider .cu, .cu.cc as valid source extensions. + self.compiler.src_extensions += ['.cu', '.cu.cc'] + # Save the original _compile method for later. + if self.compiler.compiler_type == 'msvc' or IS_WINDOWS: + raise NotImplementedError("Not support on MSVC currently.") + else: + original_compile = self.compiler._compile + + def unix_custom_single_compiler(obj, src, ext, cc_args, extra_postargs, + pp_opts): + """ + Monkey patch machanism to replace inner compiler to custom complie process on Unix platform. + """ + # use abspath to ensure no warning + src = os.path.abspath(src) + cflags = copy.deepcopy(extra_postargs) + + try: + original_compiler = self.compiler.compiler_so + # ncvv compile CUDA source + if is_cuda_file(src): + assert CUDA_HOME is not None + nvcc_cmd = os.path.join(CUDA_HOME, 'bin', 'nvcc') + self.compiler.set_executable('compiler_so', nvcc_cmd) + # {'nvcc': {}, 'cxx: {}} + if isinstance(cflags, dict): + cflags = cflags['nvcc'] + else: + cflags = prepare_unix_cflags(cflags) + # cxx compile Cpp source + elif isinstance(cflags, dict): + cflags = cflags['cxx'] + + add_std_without_repeat( + cflags, self.compiler.compiler_type, use_std14=False) + original_compile(obj, src, ext, cc_args, cflags, pp_opts) + finally: + # restore original_compiler + self.compiler.compiler_so = original_compiler + + def object_filenames_with_cuda(origina_func): + """ + Decorated the function to add customized naming machanism. + """ + + def wrapper(source_filenames, strip_dir=0, output_dir=''): + try: + objects = origina_func(source_filenames, strip_dir, + output_dir) + for i, source in enumerate(source_filenames): + # modify xx.o -> xx.cu.o + if is_cuda_file(source): + old_obj = objects[i] + objects[i] = old_obj[:-1] + 'cu.o' + # ensure to use abspath + objects = [os.path.abspath(obj) for obj in objects] + finally: + self.compiler.object_filenames = origina_func + + return objects + + return wrapper + + # customized compile process + self.compiler._compile = unix_custom_single_compiler + self.compiler.object_filenames = object_filenames_with_cuda( + self.compiler.object_filenames) + + build_ext.build_extensions(self) + + def get_ext_filename(self, fullname): + # for example: custommed_extension.cpython-37m-x86_64-linux-gnu.so + ext_name = super(BuildExtension, self).get_ext_filename(fullname) + if self.no_python_abi_suffix and six.PY3: + split_str = '.' + name_items = ext_name.split(split_str) + assert len( + name_items + ) > 2, "Expected len(name_items) > 2, but received {}".format( + len(name_items)) + name_items.pop(-2) + # custommed_extension.so + ext_name = split_str.join(name_items) + + return ext_name + + def _check_abi(self): + pass diff --git a/python/paddle/fluid/tests/custom_op/extension_utils.py b/python/paddle/fluid/tests/custom_op/extension_utils.py new file mode 100644 index 00000000000..c2683140e8e --- /dev/null +++ b/python/paddle/fluid/tests/custom_op/extension_utils.py @@ -0,0 +1,216 @@ +# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. +# +# 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 os +import six +import sys +import copy +import glob +import warnings +import subprocess + +import paddle + +IS_WINDOWS = os.name == 'nt' +# TODO(Aurelius84): Need check version of gcc and g++ is same. +# After CI path is fixed, we will modify into cc. +NVCC_COMPILE_FLAGS = [ + '-ccbin', 'gcc', '-DPADDLE_WITH_CUDA', '-DEIGEN_USE_GPU', + '-DPADDLE_USE_DSO', '-Xcompiler', '-fPIC', '-w', '--expt-relaxed-constexpr', + '-O3', '-DNVCC' +] + + +def prepare_unix_cflags(cflags): + """ + Prepare all necessary compiled flags for nvcc compiling CUDA files. + """ + cflags = NVCC_COMPILE_FLAGS + cflags + get_cuda_arch_flags(cflags) + + return cflags + + +def add_std_without_repeat(cflags, compiler_type, use_std14=False): + """ + Append -std=c++11/14 in cflags if without specific it before. + """ + cpp_flag_prefix = '/std:' if compiler_type == 'msvc' else '-std=' + if not any(cpp_flag_prefix in flag for flag in cflags): + suffix = 'c++14' if use_std14 else 'c++11' + cpp_flag = cpp_flag_prefix + suffix + cflags.append(cpp_flag) + + +def get_cuda_arch_flags(cflags): + """ + For an arch, say "6.1", the added compile flag will be + ``-gencode=arch=compute_61,code=sm_61``. + For an added "+PTX", an additional + ``-gencode=arch=compute_xx,code=compute_xx`` is added. + """ + # TODO(Aurelius84): + return [] + + +def normalize_extension_kwargs(kwargs, use_cuda=False): + """ + Normalize include_dirs, library_dir and other attributes in kwargs. + """ + assert isinstance(kwargs, dict) + # append necessary include dir path of paddle + include_dirs = kwargs.get('include_dirs', []) + include_dirs.extend(find_paddle_includes(use_cuda)) + kwargs['include_dirs'] = include_dirs + + # append necessary lib path of paddle + library_dirs = kwargs.get('library_dirs', []) + library_dirs.extend(find_paddle_libraries(use_cuda)) + kwargs['library_dirs'] = library_dirs + + # add runtime library dirs + runtime_library_dirs = kwargs.get('runtime_library_dirs', []) + runtime_library_dirs.extend(find_paddle_libraries(use_cuda)) + kwargs['runtime_library_dirs'] = runtime_library_dirs + + # append compile flags + extra_compile_args = kwargs.get('extra_compile_args', []) + extra_compile_args.extend(['-g']) + kwargs['extra_compile_args'] = extra_compile_args + + # append link flags + extra_link_args = kwargs.get('extra_link_args', []) + extra_link_args.extend(['-lpaddle_framework', '-lcudart']) + kwargs['extra_link_args'] = extra_link_args + + kwargs['language'] = 'c++' + return kwargs + + +def find_paddle_includes(use_cuda=False): + """ + Return Paddle necessary include dir path. + """ + # pythonXX/site-packages/paddle/include + paddle_include_dir = paddle.sysconfig.get_include() + third_party_dir = os.path.join(paddle_include_dir, 'third_party') + + include_dirs = [paddle_include_dir, third_party_dir] + + return include_dirs + + +def find_cuda_includes(): + + cuda_home = find_cuda_home() + if cuda_home is None: + raise ValueError( + "Not found CUDA runtime, please use `export CUDA_HOME=XXX` to specific it." + ) + + return [os.path.join(cuda_home, 'lib64')] + + +def find_cuda_home(): + """ + Use heuristic method to find cuda path + """ + # step 1. find in $CUDA_HOME or $CUDA_PATH + cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH') + + # step 2. find path by `which nvcc` + if cuda_home is None: + which_cmd = 'where' if IS_WINDOWS else 'which' + try: + with open(os.devnull, 'w') as devnull: + nvcc_path = subprocess.check_output( + [which_cmd, 'nvcc'], stderr=devnull) + if six.PY3: + nvcc_path = nvcc_path.decode() + nvcc_path = nvcc_path.rstrip('\r\n') + # for example: /usr/local/cuda/bin/nvcc + cuda_home = os.path.dirname(os.path.dirname(nvcc_path)) + except: + if IS_WINDOWS: + # search from default NVIDIA GPU path + candidate_paths = glob.glob( + 'C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*') + if len(candidate_paths) > 0: + cuda_home = candidate_paths[0] + else: + cuda_home = "/usr/local/cuda" + # step 3. check whether path is valid + if not os.path.exists(cuda_home) and paddle.is_compiled_with_cuda(): + cuda_home = None + warnings.warn( + "Not found CUDA runtime, please use `export CUDA_HOME= XXX` to specific it." + ) + + return cuda_home + + +def find_paddle_libraries(use_cuda=False): + """ + Return Paddle necessary library dir path. + """ + # pythonXX/site-packages/paddle/libs + paddle_lib_dirs = [paddle.sysconfig.get_lib()] + if use_cuda: + cuda_dirs = find_cuda_includes() + paddle_lib_dirs.extend(cuda_dirs) + return paddle_lib_dirs + + +def append_necessary_flags(extra_compile_args, use_cuda=False): + """ + Add necessary compile flags for gcc/nvcc compiler. + """ + necessary_flags = ['-std=c++11'] + + if use_cuda: + necessary_flags.extend(NVCC_COMPILE_FLAGS) + + +def add_compile_flag(extension, flag): + extra_compile_args = copy.deepcopy(extension.extra_compile_args) + if isinstance(extra_compile_args, dict): + for args in extra_compile_args.values(): + args.append(flag) + else: + extra_compile_args.append(flag) + + extension.extra_compile_args = extra_compile_args + + +def is_cuda_file(path): + + cuda_suffix = set(['.cu']) + items = os.path.splitext(path) + assert len(items) > 1 + return items[-1] in cuda_suffix + + +def get_build_directory(name): + """ + Return paddle extension root directory, default specific by `PADDLE_EXTENSION_DIR` + """ + root_extensions_directory = os.envsiron.get('PADDLE_EXTENSION_DIR') + if root_extensions_directory is None: + # TODO(Aurelius84): consider wind32/macOs + here = os.path.abspath(__file__) + root_extensions_directory = os.path.realpath(here) + warnings.warn( + "$PADDLE_EXTENSION_DIR is not set, using path: {} by default." + .format(root_extensions_directory)) + + return root_extensions_directory diff --git a/python/paddle/fluid/tests/custom_op/setup.py b/python/paddle/fluid/tests/custom_op/setup.py new file mode 100644 index 00000000000..b61b745508d --- /dev/null +++ b/python/paddle/fluid/tests/custom_op/setup.py @@ -0,0 +1,49 @@ +# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. +# +# 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 os +import six +from distutils.sysconfig import get_python_lib +from setuptools import setup +from cpp_extension import CppExtension, CUDAExtension, BuildExtension, IS_WINDOWS +from setuptools import Extension + +file_dir = os.path.dirname(os.path.abspath(__file__)) +site_packages_path = get_python_lib() +# Note(Aurelius84): We use `add_test` in Cmake to config how to run unittest in CI. +# `PYTHONPATH` will be set as `build/python/paddle` that will make no way to find +# paddle include directory. Because the following path is generated after insalling +# PaddlePaddle whl. So here we specific `include_dirs` to avoid errors in CI. +paddle_includes = [ + os.path.join(site_packages_path, 'paddle/include'), + os.path.join(site_packages_path, 'paddle/include/third_party') +] + +# TODO(Aurelius84): Memory layout is different if build paddle with PADDLE_WITH_MKLDNN=ON, +# and will lead to ABI problem on Coverage CI. We will handle it in next PR. +extra_compile_args = ['-DPADDLE_WITH_MKLDNN' + ] if six.PY2 and not IS_WINDOWS else [] + +setup( + name='relu_op_shared', + ext_modules=[ + CUDAExtension( + name='librelu2_op_from_setup', + sources=['relu_op.cc', 'relu_op.cu'], + include_dirs=paddle_includes, + extra_compile_args=extra_compile_args, + output_dir=file_dir) + ], + cmdclass={ + 'build_ext': BuildExtension.with_options(no_python_abi_suffix=True) + }) diff --git a/python/paddle/fluid/tests/custom_op/test_custom_op.py b/python/paddle/fluid/tests/custom_op/test_custom_op.py index c9f7d0b7c96..1c0db0be154 100644 --- a/python/paddle/fluid/tests/custom_op/test_custom_op.py +++ b/python/paddle/fluid/tests/custom_op/test_custom_op.py @@ -20,11 +20,16 @@ import contextlib import paddle import paddle.fluid as fluid - paddle.enable_static() -file_dir = os.path.dirname(os.path.abspath(__file__)) -fluid.load_op_library(os.path.join(file_dir, 'librelu2_op.so')) + +def load_so(so_name): + """ + Load .so file and parse custom op into OpInfoMap. + """ + file_dir = os.path.dirname(os.path.abspath(__file__)) + fluid.load_op_library(os.path.join(file_dir, so_name)) + from paddle.fluid.layer_helper import LayerHelper @@ -111,4 +116,5 @@ class CustomOpTest(unittest.TestCase): if __name__ == '__main__': + load_so(so_name='librelu2_op.so') unittest.main() diff --git a/python/paddle/fluid/tests/custom_op/test_custom_op_with_setup.py b/python/paddle/fluid/tests/custom_op/test_custom_op_with_setup.py new file mode 100644 index 00000000000..be9442cc71a --- /dev/null +++ b/python/paddle/fluid/tests/custom_op/test_custom_op_with_setup.py @@ -0,0 +1,33 @@ +# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved. +# +# 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 os +import unittest + +from test_custom_op import CustomOpTest, load_so + + +def compile_so(): + """ + Compile .so file by running setup.py config. + """ + # build .so with setup.py + file_dir = os.path.dirname(os.path.abspath(__file__)) + os.system('cd {} && python setup.py build'.format(file_dir)) + + +if __name__ == '__main__': + compile_so() + load_so(so_name='librelu2_op_from_setup.so') + unittest.main() -- GitLab