[CustomOp] Add new paddle custom op so (#31141)

* add new custom op so

* fix use new method error

* fix test failed

paddle/fluid/framework/CMakeLists.txt

@@ -320,9 +320,9 @@ message(STATUS "branch: ${PADDLE_BRANCH}")
 
 configure_file(commit.h.in commit.h)
 
-cc_library(custom_tensor SRCS ../extension/src/tensor.cc DEPS lod_tensor)
+cc_library(custom_tensor SRCS ../extension/src/tensor.cc DEPS lod_tensor memory enforce)
 cc_library(op_meta_info SRCS ../extension/src/op_meta_info.cc DEPS custom_tensor)
-cc_library(custom_operator SRCS custom_operator.cc DEPS operator op_registry device_context dynamic_loader custom_tensor op_meta_info)
+cc_library(custom_operator SRCS custom_operator.cc DEPS tensor attribute framework_proto op_registry operator dynamic_loader string_helper custom_tensor op_meta_info)
 cc_test(custom_tensor_test SRCS custom_tensor_test.cc DEPS custom_tensor glog)
 
 set(FLUID_FRAMEWORK_MODULES proto_desc memory lod_tensor executor data_feed_proto layer dynamic_loader custom_operator)
@@ -361,3 +361,37 @@ endif()
 if(WITH_TESTING AND TEST selected_rows_test)
   set_tests_properties(selected_rows_test PROPERTIES TIMEOUT 120)
 endif()
+
+# New custom op extension mechanism related
+
+# if not deps `layer`, will cause: undefined symbol: _ZN6paddle10imperative7VarBase9name_set_
+set(PADDLE_CUSTOM_OP_MODULES custom_tensor op_meta_info custom_operator layer)
+
+cc_library(paddle_custom_op_shared
+    SHARED SRCS custom_operator.cc ../extension/src/tensor.cc ../extension/src/op_meta_info.cc
+    ${CMAKE_SOURCE_DIR}/paddle/fluid/imperative/layer.cc
+    DEPS ${PADDLE_CUSTOM_OP_MODULES})
+get_property(os_dependency_modules GLOBAL PROPERTY OS_DEPENDENCY_MODULES)
+set_target_properties(paddle_custom_op_shared PROPERTIES OUTPUT_NAME paddle_custom_op)
+target_link_libraries(paddle_custom_op_shared ${os_dependency_modules})
+
+if (LINUX)
+  set(PADDLE_CUSTOM_OP_SHARED_LIB
+      ${PADDLE_BINARY_DIR}/paddle/fluid/framework/libpaddle_custom_op.so
+      CACHE INTERNAL "Paddle custom op lib")
+endif()
+
+if (WIN32)
+  set(PADDLE_CUSTOM_OP_SHARED_LIB
+      ${PADDLE_BINARY_DIR}/paddle/fluid/framework/${CMAKE_BUILD_TYPE}/paddle_custom_op.lib
+      CACHE INTERNAL "Paddle custom op lib")
+  set(PADDLE_CUSTOM_OP_SHARED_LIB
+      ${PADDLE_BINARY_DIR}/paddle/fluid/framework/${CMAKE_BUILD_TYPE}/paddle_custom_op.dll
+      CACHE INTERNAL "Paddle custom op dll")
+endif()
+
+if(APPLE)
+  set(PADDLE_CUSTOM_OP_SHARED_LIB
+      ${PADDLE_BINARY_DIR}/paddle/fluid/framework/paddle_custom_op.dylib
+      CACHE INTERNAL "Paddle custom op lib")
+endif()

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

@@ -24,8 +24,10 @@ from test_custom_relu_op_setup import custom_relu_dynamic, custom_relu_static
 
 # Because Windows don't use docker, the shared lib already exists in the 
 # cache dir, it will not be compiled again unless the shared lib is removed.
-if os.name == 'nt':
-    cmd = 'del {}\\custom_relu_module_jit.pyd'.format(get_build_directory())
+file = '{}\\custom_relu_module_jit\\custom_relu_module_jit.pyd'.format(
+    get_build_directory())
+if os.name == 'nt' and os.path.isfile(file):
+    cmd = 'del {}'.format(file)
     run_cmd(cmd, True)
 
 # Compile and load custom op Just-In-Time.

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

@@ -22,8 +22,9 @@ from paddle.utils.cpp_extension.extension_utils import run_cmd
 
 # Because Windows don't use docker, the shared lib already exists in the 
 # cache dir, it will not be compiled again unless the shared lib is removed.
-if os.name == 'nt':
-    cmd = 'del {}\\dispatch_op.pyd'.format(get_build_directory())
+file = '{}\\dispatch_op\\dispatch_op.pyd'.format(get_build_directory())
+if os.name == 'nt' and os.path.isfile(file):
+    cmd = 'del {}'.format(file)
     run_cmd(cmd, True)
 
 dispatch_op = load(

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

@@ -25,8 +25,9 @@ from utils import paddle_includes, extra_compile_args
 
 # Because Windows don't use docker, the shared lib already exists in the 
 # cache dir, it will not be compiled again unless the shared lib is removed.
-if os.name == 'nt':
-    cmd = 'del {}\\multi_out_jit.pyd'.format(get_build_directory())
+file = '{}\\multi_out_jit\\multi_out_jit.pyd'.format(get_build_directory())
+if os.name == 'nt' and os.path.isfile(file):
+    cmd = 'del {}'.format(file)
     run_cmd(cmd, True)
 
 # Compile and load custom op Just-In-Time.

python/paddle/utils/cpp_extension/cpp_extension.py

@@ -219,9 +219,6 @@ class BuildExtension(build_ext, object):
         super(BuildExtension, self).__init__(*args, **kwargs)
         self.no_python_abi_suffix = kwargs.get("no_python_abi_suffix", True)
         self.output_dir = kwargs.get("output_dir", None)
-        # for compatible two custom op define method
-        use_new_custom_op_load_method(
-            kwargs.get("use_new_method", use_new_custom_op_load_method()))
 
     def initialize_options(self):
         super(BuildExtension, self).initialize_options()

python/paddle/utils/cpp_extension/extension_utils.py

@@ -288,7 +288,10 @@ def normalize_extension_kwargs(kwargs, use_cuda=False):
 
         # append link flags
         extra_link_args = kwargs.get('extra_link_args', [])
-        extra_link_args.append('-lpaddle_framework')
+        if use_new_custom_op_load_method():
+            extra_link_args.append('-lpaddle_custom_op')
+        else:
+            extra_link_args.append('-lpaddle_framework')
         if use_cuda:
             extra_link_args.append('-lcudart')
 
@@ -592,6 +595,10 @@ def _write_setup_file(name,
     import os
     from paddle.utils.cpp_extension import CppExtension, CUDAExtension, BuildExtension, setup
     from paddle.utils.cpp_extension import get_build_directory
+    from paddle.utils.cpp_extension.extension_utils import use_new_custom_op_load_method
+
+    use_new_custom_op_load_method({use_new_method})
+
     setup(
         name='{name}',
         ext_modules=[
@@ -601,9 +608,8 @@ def _write_setup_file(name,
                 extra_compile_args={extra_compile_args},
                 extra_link_args={extra_link_args})],
         cmdclass={{"build_ext" : BuildExtension.with_options(
-            output_dir='{build_dir}',
-            no_python_abi_suffix=True,
-            use_new_method={use_new_method})
+            output_dir=r'{build_dir}',
+            no_python_abi_suffix=True)
         }})""").lstrip()
 
     with_cuda = False

python/setup.py.in

@@ -340,6 +340,11 @@ if sys.platform.startswith('linux'):
     shutil.copy('${FLUID_FRAMEWORK_SHARED_LIB}', libs_path)
     package_data['paddle.libs'] += ['libpaddle_framework.so']
 
+# copy libpaddle_custom_op.so to libs on linux
+if sys.platform.startswith('linux'):
+    shutil.copy('${PADDLE_CUSTOM_OP_SHARED_LIB}', libs_path)
+    package_data['paddle.libs'] += ['libpaddle_custom_op.so']
+
 # copy paddle_framework.lib/paddle_framework.dll to libs on windows
 if os.name == 'nt':
     shutil.copy('${FLUID_FRAMEWORK_IMPORT_LIB}', libs_path)