【CustomOp】support setup.py to compile custom op (#30753)

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)

python/paddle/fluid/tests/custom_op/cpp_extension.py

+# 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

python/paddle/fluid/tests/custom_op/extension_utils.py

+# 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

python/paddle/fluid/tests/custom_op/setup.py

+# 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)
+    })

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()

python/paddle/fluid/tests/custom_op/test_custom_op_with_setup.py

+# 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()