From e92b304032eb7e9a46d18f38b3bc52ff00ee4701 Mon Sep 17 00:00:00 2001
From: Chen Weihang <chenweihang@baidu.com>
Date: Sat, 22 Jan 2022 12:02:40 +0800
Subject: [PATCH] [PTen] Auto generate include headers (#39123)

* auto gen include headers

* move to pten.cmake
---
 .gitignore                              |  2 ++
 cmake/{pten_kernel.cmake => pten.cmake} | 42 +++++++++++++++++++++++++
 paddle/pten/CMakeLists.txt              | 12 +++++++
 paddle/pten/kernels/CMakeLists.txt      |  2 --
 4 files changed, 56 insertions(+), 2 deletions(-)
 rename cmake/{pten_kernel.cmake => pten.cmake} (82%)

diff --git a/.gitignore b/.gitignore
index 6be36bf8c2..708126b3bb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,6 +5,8 @@ paddle/fluid/API_PR.spec
 paddle/fluid/op_use_default_grad_maker_DEV.spec
 paddle/fluid/op_use_default_grad_maker_PR.spec
 paddle/pten/api/*/api*
+paddle/pten/include/*
+paddle/pten/extension.h
 
 *.DS_Store
 *.vs
diff --git a/cmake/pten_kernel.cmake b/cmake/pten.cmake
similarity index 82%
rename from cmake/pten_kernel.cmake
rename to cmake/pten.cmake
index c2928376a0..70d61027da 100644
--- a/cmake/pten_kernel.cmake
+++ b/cmake/pten.cmake
@@ -12,6 +12,48 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+function(generate_unify_header DIR_NAME)
+    set(options "")
+    set(oneValueArgs HEADER_NAME SKIP_SUFFIX)
+    set(multiValueArgs "")
+    cmake_parse_arguments(generate_unify_header "${options}" "${oneValueArgs}"
+        "${multiValueArgs}" ${ARGN})
+
+    # get header name and suffix
+    set(header_name "${DIR_NAME}")
+    list(LENGTH generate_unify_header_HEADER_NAME generate_unify_header_HEADER_NAME_len)
+    if(${generate_unify_header_HEADER_NAME_len} GREATER 0)
+        set(header_name "${generate_unify_header_HEADER_NAME}")
+    endif()
+    set(skip_suffix "")
+    list(LENGTH generate_unify_header_SKIP_SUFFIX generate_unify_header_SKIP_SUFFIX_len)
+    if(${generate_unify_header_SKIP_SUFFIX_len} GREATER 0)
+        set(skip_suffix "${generate_unify_header_SKIP_SUFFIX}")
+    endif()
+
+    # generate target header file
+    set(header_file ${CMAKE_CURRENT_SOURCE_DIR}/include/${header_name}.h)
+    file(WRITE ${header_file} "// Header file generated by paddle/pten/CMakeLists.txt for external users,\n// DO NOT edit or include it within paddle.\n\n#pragma once\n\n")
+
+    # get all top-level headers and write into header file
+    file(GLOB HEADERS "${CMAKE_CURRENT_SOURCE_DIR}\/${DIR_NAME}\/*.h")
+    foreach(header ${HEADERS})
+        if("${skip_suffix}" STREQUAL "")
+            string(REPLACE "${PADDLE_SOURCE_DIR}\/" "" header "${header}")
+            file(APPEND ${header_file} "#include \"${header}\"\n")
+        else()
+            string(FIND "${header}" "${skip_suffix}.h" skip_suffix_found)
+            if(${skip_suffix_found} EQUAL -1)
+                string(REPLACE "${PADDLE_SOURCE_DIR}\/" "" header "${header}")
+                file(APPEND ${header_file} "#include \"${header}\"\n")
+            endif()
+        endif()
+    endforeach()
+    # append header into extension.h
+    string(REPLACE "${PADDLE_SOURCE_DIR}\/" "" header_file "${header_file}")
+    file(APPEND ${pten_extension_header_file} "#include \"${header_file}\"\n")
+endfunction()
+
 # call kernel_declare need to make sure whether the target of input exists
 function(kernel_declare TARGET_LIST)
     foreach(kernel_path ${TARGET_LIST})
diff --git a/paddle/pten/CMakeLists.txt b/paddle/pten/CMakeLists.txt
index cde5e719e3..671ed28313 100644
--- a/paddle/pten/CMakeLists.txt
+++ b/paddle/pten/CMakeLists.txt
@@ -1,3 +1,6 @@
+# pten auto cmake utils
+include(pten)
+
 # paddle experimental common components
 add_subdirectory(common)
 
@@ -25,3 +28,12 @@ message(STATUS "All standard pten kernels: ${pten_kernels}")
 set(PTEN_DEPS ${PTEN_DEPS} ${pten_kernels})
 
 cc_library(pten DEPS ${PTEN_DEPS})
+
+set(pten_extension_header_file ${CMAKE_CURRENT_SOURCE_DIR}/extension.h CACHE INTERNAL "pten/extension.h file")
+file(WRITE ${pten_extension_header_file} "// Header file generated by paddle/pten/CMakeLists.txt for external users,\n// DO NOT edit or include it within paddle.\n\n#pragma once\n\n")
+
+# generate inner headers include dir for users
+generate_unify_header(backends)
+generate_unify_header(core)
+generate_unify_header(infermeta)
+generate_unify_header(kernels SKIP_SUFFIX grad_kernel)
diff --git a/paddle/pten/kernels/CMakeLists.txt b/paddle/pten/kernels/CMakeLists.txt
index 76e1128088..f838c4d424 100644
--- a/paddle/pten/kernels/CMakeLists.txt
+++ b/paddle/pten/kernels/CMakeLists.txt
@@ -1,5 +1,3 @@
-include(pten_kernel)
-
 set(kernel_declare_file ${PADDLE_BINARY_DIR}/paddle/pten/kernels/declarations.h.tmp CACHE INTERNAL "declarations.h file")
 set(kernel_declare_file_final ${PADDLE_BINARY_DIR}/paddle/pten/kernels/declarations.h)
 file(WRITE ${kernel_declare_file} "// Generated by the paddle/pten/kernels/CMakeLists.txt.  DO NOT EDIT!\n\n#pragma once\n\n")
-- 
GitLab