From 8739497c6cbe3b52862d72b2023b4a765118274e Mon Sep 17 00:00:00 2001
From: JingZhuangzhuang <75348594+JZZ-NOTE@users.noreply.github.com>
Date: Fri, 21 Oct 2022 10:40:43 +0800
Subject: [PATCH] Add infer prune function (#47047)

* Add infer prune function

* add fusion op
---
 cmake/operators.cmake                 | 55 +++++++++++++++++++++++++++
 cmake/phi.cmake                       | 32 ++++++++++++++++
 paddle/fluid/inference/CMakeLists.txt |  8 ++++
 paddle/fluid/operators/CMakeLists.txt |  5 +++
 paddle/phi/kernels/CMakeLists.txt     |  7 ++++
 5 files changed, 107 insertions(+)

diff --git a/cmake/operators.cmake b/cmake/operators.cmake
index bbf77b6615..bb5722ae82 100644
--- a/cmake/operators.cmake
+++ b/cmake/operators.cmake
@@ -611,3 +611,58 @@ function(register_operators)
     endif()
   endif()
 endfunction()
+
+function(prune_pybind_h)
+  set(op_list ${OP_LIST})
+
+  list(APPEND op_list "load_combine")
+  list(APPEND op_list "tensorrt_engine")
+
+  # add fused_op in op_list
+  list(APPEND op_list "fc")
+  list(APPEND op_list "conv2d_fusion")
+  list(APPEND op_list "fusion_seqconv_eltadd_relu")
+  list(APPEND op_list "fusion_seqpool_cvm_concat")
+  list(APPEND op_list "fusion_gru")
+  list(APPEND op_list "fusion_seqexpand_concat_fc")
+  list(APPEND op_list "fusion_repeated_fc_relu")
+  list(APPEND op_list "fusion_squared_mat_sub")
+
+  # add plugin_op in op_list
+  list(APPEND op_list "anchor_generator")
+
+  file(STRINGS ${pybind_file} op_registry_list)
+
+  file(WRITE ${pybind_file_prune} "")
+  file(
+    APPEND ${pybind_file_prune}
+    "// Generated by the paddle/fluid/operators/CMakeLists.txt.  DO NOT EDIT!\n"
+  )
+
+  # add USE_OP_ITSELF for all op in op_list
+  foreach(op_name IN LISTS op_list)
+    file(APPEND ${pybind_file_prune} "USE_OP_ITSELF(${op_name});\n")
+  endforeach()
+
+  foreach(op_registry IN LISTS op_registry_list)
+    if(NOT ${op_registry} EQUAL "")
+      foreach(op_name IN LISTS op_list)
+        string(FIND ${op_registry} "(${op_name})" index1)
+        string(FIND ${op_registry} "(${op_name}," index2)
+        string(FIND ${op_registry} "USE_OP_ITSELF" index3)
+        if(((NOT ${index1} EQUAL "-1") OR (NOT ${index2} EQUAL "-1"))
+           AND (${index3} EQUAL "-1"))
+          file(APPEND ${pybind_file_prune} "${op_registry}\n")
+        endif()
+      endforeach()
+    endif()
+  endforeach()
+
+  file(WRITE ${pybind_file} "")
+  file(STRINGS ${pybind_file_prune} op_registry_list_tmp)
+  foreach(op_name IN LISTS op_registry_list_tmp)
+    if(NOT ${op_name} EQUAL "")
+      file(APPEND ${pybind_file} "${op_name}\n")
+    endif()
+  endforeach()
+endfunction()
diff --git a/cmake/phi.cmake b/cmake/phi.cmake
index 6886e65d37..04e16b9a4f 100644
--- a/cmake/phi.cmake
+++ b/cmake/phi.cmake
@@ -148,3 +148,35 @@ function(register_op_utils TARGET_NAME)
     SRCS ${utils_srcs}
     DEPS ${register_op_utils_DEPS})
 endfunction()
+
+function(prune_declaration_h)
+  set(kernel_list ${KERNEL_LIST})
+  file(STRINGS ${kernel_declare_file} kernel_registry_list)
+
+  file(WRITE ${kernel_declare_file_prune} "")
+  file(APPEND ${kernel_declare_file_prune}
+       "// Generated by the paddle/phi/kernels/CMakeLists.txt.  DO NOT EDIT!\n")
+  file(APPEND ${kernel_declare_file_prune} "#pragma once\n")
+  file(APPEND ${kernel_declare_file_prune}
+       "#include \"paddle/phi/core/kernel_registry.h\"\n")
+
+  foreach(kernel_registry IN LISTS kernel_registry_list)
+    if(NOT ${kernel_registry} EQUAL "")
+      foreach(kernel_name IN LISTS kernel_list)
+        string(FIND ${kernel_registry} "(${kernel_name})" index1)
+        string(FIND ${kernel_registry} "(${kernel_name}," index2)
+        if((NOT ${index1} EQUAL "-1") OR (NOT ${index2} EQUAL "-1"))
+          file(APPEND ${kernel_declare_file_prune} "${kernel_registry}\n")
+        endif()
+      endforeach()
+    endif()
+  endforeach()
+
+  file(WRITE ${kernel_declare_file} "")
+  file(STRINGS ${kernel_declare_file_prune} kernel_registry_list_tmp)
+  foreach(kernel_registry IN LISTS kernel_registry_list_tmp)
+    if(NOT ${kernel_registry} EQUAL "")
+      file(APPEND ${kernel_declare_file} "${kernel_registry}\n")
+    endif()
+  endforeach()
+endfunction()
diff --git a/paddle/fluid/inference/CMakeLists.txt b/paddle/fluid/inference/CMakeLists.txt
index 7f2daa942b..9e130170a3 100644
--- a/paddle/fluid/inference/CMakeLists.txt
+++ b/paddle/fluid/inference/CMakeLists.txt
@@ -58,6 +58,14 @@ set(STATIC_INFERENCE_API
     phi
     ${mkldnn_quantizer_cfg})
 
+set(OP_LIST
+    ""
+    CACHE STRING "The list of operators that will be compiled")
+
+set(KERNEL_LIST
+    ""
+    CACHE STRING "The list of phi kernels that will be compiled")
+
 #windows GPU static library over the limit, so not create_static_lib, and cc_library is dummy
 if(WIN32 AND WITH_GPU)
   cc_library(paddle_inference DEPS ${fluid_modules} phi ${STATIC_INFERENCE_API}
diff --git a/paddle/fluid/operators/CMakeLists.txt b/paddle/fluid/operators/CMakeLists.txt
index d985baf8c9..e4ad697015 100644
--- a/paddle/fluid/operators/CMakeLists.txt
+++ b/paddle/fluid/operators/CMakeLists.txt
@@ -8,6 +8,7 @@ add_definitions(-D_USE_MATH_DEFINES)
 unset(GLOB_OP_LIB CACHE)
 unset(OP_LIBRARY CACHE)
 set(pybind_file ${PADDLE_BINARY_DIR}/paddle/fluid/pybind/pybind.h.tmp CACHE INTERNAL "pybind.h file")
+set(pybind_file_prune ${PADDLE_BINARY_DIR}/paddle/fluid/pybind/pybind.h.prune CACHE INTERNAL "pybind.h file")
 set(pybind_file_final ${PADDLE_BINARY_DIR}/paddle/fluid/pybind/pybind.h)
 file(WRITE ${pybind_file} "// Generated by the paddle/fluid/operators/CMakeLists.txt.  DO NOT EDIT!\n\n")
 
@@ -258,3 +259,7 @@ copy_if_different(${pybind_file} ${pybind_file_final})
 if (WITH_CUSTOM_DEVICE)
 cc_library(custom_device_common_op_registry SRCS custom_device_common_op_registry.cc DEPS operator)
 endif()
+
+if(NOT "${OP_LIST}" STREQUAL "")
+prune_pybind_h()
+endif()
diff --git a/paddle/phi/kernels/CMakeLists.txt b/paddle/phi/kernels/CMakeLists.txt
index 7ea9041d77..7cbd218543 100644
--- a/paddle/phi/kernels/CMakeLists.txt
+++ b/paddle/phi/kernels/CMakeLists.txt
@@ -9,6 +9,9 @@ file(
 )
 file(APPEND ${kernel_declare_file}
      "#include \"paddle/phi/core/kernel_registry.h\"\n\n")
+set(kernel_declare_file_prune
+    ${PADDLE_BINARY_DIR}/paddle/phi/kernels/declarations.h.prune
+    CACHE INTERNAL "declarations.h file")
 
 # phi functors and functions called by kernels
 add_subdirectory(funcs)
@@ -171,3 +174,7 @@ if(WITH_XPU)
 endif()
 
 set_property(GLOBAL PROPERTY PHI_KERNELS ${ADD_PHI_KERNELS})
+
+if(NOT "${KERNEL_LIST}" STREQUAL "")
+  prune_declaration_h()
+endif()
-- 
GitLab