From 883b6aba8b695994336504577578cad95abe7c1f Mon Sep 17 00:00:00 2001
From: haosicheng <47998305+HarperCy@users.noreply.github.com>
Date: Tue, 7 Mar 2023 15:30:14 +0800
Subject: [PATCH] Expose xdnn headers files for custom op (#51055)

only when -DWITH_XDNN_API=ON, will those files be packed into whl
package
---
 python/env_dict.py.in |  3 ++-
 python/setup.py.in    |  5 ++++-
 setup.py              | 10 ++++++++++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/python/env_dict.py.in b/python/env_dict.py.in
index e5bd6a0619d..d06a9d1af3a 100644
--- a/python/env_dict.py.in
+++ b/python/env_dict.py.in
@@ -71,5 +71,6 @@ env_dict={
     'WIN32':'@WIN32@',
     'JIT_RELEASE_WHL':'@JIT_RELEASE_WHL@',
     'WITH_PSLIB':'@WITH_PSLIB@',
-    'PYBIND_INCLUDE_DIR':'@PYBIND_INCLUDE_DIR@'
+    'PYBIND_INCLUDE_DIR':'@PYBIND_INCLUDE_DIR@',
+    'WITH_XDNN_API':'@WITH_XDNN_API@'
 }
diff --git a/python/setup.py.in b/python/setup.py.in
index f92eeb39c1d..452ca01e2b6 100644
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -728,6 +728,9 @@ if '${WITH_GPU}' == 'ON' or '${WITH_ROCM}' == 'ON':
     # externalErrorMsg.pb for External Error message
     headers += list(find_files('*.pb', '${externalError_INCLUDE_DIR}'))
 
+if '${WITH_XDNN_API}' == 'ON':
+    headers += list(find_files('*.h', '@PADDLE_BINARY_DIR@/third_party/xpu/src/extern_xpu/xpu', recursive=True)) # xdnn api headers
+
 headers += list(find_files('*.h', '${PYBIND_INCLUDE_DIR}', True)) # pybind headers
 
 class InstallCommand(InstallCommandBase):
@@ -776,7 +779,7 @@ class InstallHeaders(Command):
         else:
             # third_party
             install_dir = re.sub('${THIRD_PARTY_PATH}', 'third_party', header)
-            patterns = ['install/mkldnn/include', 'pybind/src/extern_pybind/include']
+            patterns = ['install/mkldnn/include', 'pybind/src/extern_pybind/include', 'third_party/xpu/src/extern_xpu/xpu/include/']
             for pattern in patterns:
                 install_dir = re.sub(pattern, '', install_dir)
         install_dir = os.path.join(self.install_dir, os.path.dirname(install_dir))
diff --git a/setup.py b/setup.py
index 4e6c0436f0a..07c96b720fa 100644
--- a/setup.py
+++ b/setup.py
@@ -155,6 +155,7 @@ def get_header_install_dir(header):
         patterns = [
             'install/mkldnn/include',
             'pybind/src/extern_pybind/include',
+            'third_party/xpu/src/extern_xpu/xpu/include/',
         ]
         for pattern in patterns:
             install_dir = re.sub(pattern, '', install_dir)
@@ -1222,6 +1223,15 @@ def get_headers():
             find_files('*.pb', env_dict.get("externalError_INCLUDE_DIR"))
         )
 
+    if env_dict.get("WITH_XDNN_API") == 'ON':
+        headers += list(
+            find_files(
+                '*.h',
+                paddle_binary_dir + '/third_party/xpu/src/extern_xpu/xpu',
+                recursive=True,
+            )
+        )  # xdnn api headers
+
     # pybind headers
     headers += list(find_files('*.h', env_dict.get("PYBIND_INCLUDE_DIR"), True))
     return headers
-- 
GitLab