[XPU]add support for XPUFFT (#72477)

fangfangssj · web-flow · commit 1943d976738e · 2025-05-07T16:18:42.000+08:00
* add cmake

* fix

* fix

* fix setup.in
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -56,6 +56,7 @@ option(WITH_XPU_KP "Compile PaddlePaddle with BAIDU XPU compiler " OFF)
 option(WITH_XPU_XFT "Compile PaddlePaddle with BAIDU XPU-XFT" OFF)
 option(WITH_XPU_PLUGIN "Compile PaddlePaddle with BAIDU XPU plugin" OFF)
 option(WITH_XPU_XRE5 "Compile PaddlePaddle with BAIDU XPU XRE 5" OFF)
+option(WITH_XPU_FFT "Compile PaddlePaddle with BAIDU XPU FFT" OFF)
 option(WITH_WIN_DUMP_DBG "Compile with windows core dump debug mode" OFF)
 option(WITH_ROCM "Compile PaddlePaddle with ROCM platform" OFF)
 option(WITH_IPU "Compile PaddlePaddle with Graphcore IPU" OFF)
@@ -95,6 +96,9 @@ endif()
 if(WITH_GPU AND WITH_XPU_XRE5)
   message(FATAL_ERROR "Error when compile GPU and XPU-XRE5 at the same time")
 endif()
+if(WITH_GPU AND WITH_XPU_FFT)
+  message(FATAL_ERROR "Error when compile GPU and XPU-FFT at the same time")
+endif()
 if(WITH_GPU AND WITH_ROCM)
   message(FATAL_ERROR "Error when compile CUDA and ROCM at the same time")
 endif()
@@ -478,6 +482,14 @@ if(NOT WITH_XPU AND WITH_XPU_XRE5)
       CACHE STRING "Disable XRE5 when compiling without XPU" FORCE)
 endif()
 
+if(NOT WITH_XPU AND WITH_XPU_FFT)
+  message((WARNING
+           "Disable FFT when compiling without XPU. Force WITH_XPU_FFT=OFF"))
+  set(WITH_XPU_FFT
+      OFF
+      CACHE STRING "Disable FFT when compiling without XPU" FORCE)
+endif()
+
 if(WITH_NCCL)
   add_definitions("-DPADDLE_WITH_NCCL")
   include(nccl)
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -28,6 +28,7 @@ set(XPU_XPTI_LIB_NAME "libxpti.so")
 set(XPU_XBLAS_LIB_NAME "libxpu_blas.so")
 set(XPU_XFA_LIB_NAME "libxpu_flash_attention.so")
 set(XPU_XPUDNN_LIB_NAME "libxpu_dnn.so")
+set(XPU_FFT_LIB_NAME "libcufft.so")
 
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
   set(XPU_XHPC_BASE_DATE "dev/20250417")
@@ -51,6 +52,10 @@ else()
   set(XPU_XPTI_BASE_VERSION "0.0.1")
 endif()
 
+if(NOT DEFINED XPU_FFT_BASE_DATE)
+  set(XPU_FFT_BASE_DATE "20250425")
+endif()
+
 set(XPU_XRE_BASE_URL
     "https://klx-sdk-release-public.su.bcebos.com/xre/release/${XPU_XRE_BASE_VERSION}"
 )
@@ -81,6 +86,13 @@ if(WITH_XPU_XRE5)
   )
 endif()
 
+if(WITH_XPU_FFT)
+  set(XPU_FFT_BASE_URL
+      "https://klx-sdk-release-public.su.bcebos.com/xpufft/kl3/${XPU_FFT_BASE_DATE}"
+  )
+  set(XPU_FFT_DIR_NAME "xpufft_ubuntu2004-x86_64")
+endif()
+
 if(WITH_AARCH64)
   set(XPU_XRE_DIR_NAME "xre-kylin_aarch64")
   set(XPU_XCCL_DIR_NAME "") # TODO: xccl has no kylin output now.
@@ -130,6 +142,10 @@ if(WITH_XPTI)
   set(XPU_XPTI_URL "${XPU_XPTI_BASE_URL}/${XPU_XPTI_DIR_NAME}.tar.gz")
 endif()
 
+if(WITH_XPU_FFT)
+  set(XPU_FFT_URL "${XPU_FFT_BASE_URL}/${XPU_FFT_DIR_NAME}.tar.gz")
+endif()
+
 set(XPU_XHPC_URL
     "https://klx-sdk-release-public.su.bcebos.com/xhpc/${XPU_XHPC_BASE_DATE}/${XPU_XHPC_DIR_NAME}.tar.gz"
     CACHE STRING "" FORCE)
@@ -205,6 +221,12 @@ if(DEFINED ENV{XPU_LIB_ROOT})
     set(XFT_COMMAND
         "${CMAKE_SOURCE_DIR}/tools/xpu/get_xft_dependence_from_custom_path.sh")
   endif()
+
+  # FFT
+  if(DEFINED ENV{XPU_FFT_DIR_NAME})
+    set(XPU_FFT_URL "${XPU_LIB_ROOT}/$ENV{XPU_FFT_DIR_NAME}")
+    set(XPU_FFT_DIR_NAME "$ENV{XPU_FFT_DIR_NAME}")
+  endif()
 endif()
 
 if(WITH_XPU_XRE5)
@@ -219,7 +241,9 @@ if(WITH_XPU_XRE5)
       ${XPU_XCCL_URL} ${XPU_XCCL_DIR_NAME} 1 && wget ${XPU_XFT_GET_DEPENCE_URL}
       && bash ${XFT_COMMAND} ${XPU_XFT_URL} ${XPU_XFT_DIR_NAME} && bash
       ${CMAKE_SOURCE_DIR}/tools/xpu/get_xpti_dependence.sh ${XPU_XPTI_URL}
-      ${XPU_XPTI_DIR_NAME}
+      ${XPU_XPTI_DIR_NAME} && bash
+      ${CMAKE_SOURCE_DIR}/tools/xpu/get_xpufft_dependence.sh ${XPU_FFT_URL}
+      ${XPU_FFT_DIR_NAME}
     DOWNLOAD_NO_PROGRESS 1
     UPDATE_COMMAND ""
     CMAKE_ARGS -DCMAKE_INSTALL_PREFIX=${XPU_INSTALL_ROOT}
@@ -244,7 +268,9 @@ else()
       ${XPU_XCCL_URL} ${XPU_XCCL_DIR_NAME} 0 && wget ${XPU_XFT_GET_DEPENCE_URL}
       && bash get_xft_dependence.sh ${XPU_XFT_URL} ${XPU_XFT_DIR_NAME} && bash
       ${CMAKE_SOURCE_DIR}/tools/xpu/get_xpti_dependence.sh ${XPU_XPTI_URL}
-      ${XPU_XPTI_DIR_NAME}
+      ${XPU_XPTI_DIR_NAME} && bash
+      ${CMAKE_SOURCE_DIR}/tools/xpu/get_xpufft_dependence.sh ${XPU_FFT_URL}
+      ${XPU_FFT_DIR_NAME}
     DOWNLOAD_NO_PROGRESS 1
     UPDATE_COMMAND ""
     CMAKE_ARGS -DCMAKE_INSTALL_PREFIX=${XPU_INSTALL_ROOT}
@@ -272,6 +298,16 @@ if(WITH_XPU_XFT)
   target_link_libraries(xpulib ${XPU_XFT_LIB})
 endif()
 
+if(WITH_XPU_FFT)
+  message(STATUS "Compile with XPU FFT!")
+  add_definitions(-DPADDLE_WITH_XPU_FFT)
+
+  set(XPU_FFT_INC_DIR "${XPU_INC_DIR}/fft")
+  include_directories(${XPU_FFT_INC_DIR})
+  set(XPU_FFT_LIB "${XPU_LIB_DIR}/${XPU_FFT_LIB_NAME}")
+  target_link_libraries(xpulib ${XPU_FFT_LIB})
+endif()
+
 set(XPU_XHPC_INC_DIR "${XPU_INC_DIR}/xhpc")
 include_directories(${XPU_XHPC_INC_DIR})
 set(XPU_XRE_INC_DIR "${XPU_INC_DIR}/xre")
diff --git a/python/env_dict.py.in b/python/env_dict.py.in
@@ -71,6 +71,9 @@ env_dict={
     'XPU_CUDA_LIB':'@XPU_CUDA_LIB@',
     'XPU_CUDA_LIB_NAME':'@XPU_CUDA_LIB_NAME@',
     'WITH_XPU_XRE5':'@WITH_XPU_XRE5@',
+    'WITH_XPU_FFT':'@WITH_XPU_FFT@',
+    'XPU_FFT_LIB':'@XPU_FFT_LIB@',
+    'XPU_FFT_LIB_NAME':'@XPU_FFT_LIB_NAME@',
     'XPU_CUDA_RT_LIB':'@XPU_CUDA_RT_LIB@',
     'XPU_CUDA_RT_LIB_NAME':'@XPU_CUDA_RT_LIB_NAME@',
     'XPU_ML_LIB':'@XPU_ML_LIB@',
diff --git a/python/setup.py.in b/python/setup.py.in
@@ -1155,6 +1155,12 @@ if '${WITH_XPU_BKCL}' == 'ON':
     shutil.copy('${XPU_BKCL_LIB}', libs_path)
     package_data['paddle.libs']+=['${XPU_BKCL_LIB_NAME}']
 
+if '${WITH_XPU_FFT}' == 'ON':
+    xpu_fft_lib_list = glob.glob('${XPU_FFT_LIB}*')
+    for xpu_fft_lib_file in xpu_fft_lib_list:
+        shutil.copy(xpu_fft_lib_file, libs_path)
+        package_data['paddle.libs']+=[os.path.basename(xpu_fft_lib_file)]
+
 if '${WITH_XPU_XFT}' == 'ON':
     shutil.copy('${XPU_XFT_LIB}', libs_path)
     package_data['paddle.libs']+=['${XPU_XFT_LIB_NAME}']
diff --git a/setup.py b/setup.py
@@ -1619,6 +1619,12 @@ def get_package_data_and_package_dir():
         shutil.copy(env_dict.get("XPU_BKCL_LIB"), libs_path)
         package_data['paddle.libs'] += [env_dict.get("XPU_BKCL_LIB_NAME")]
 
+    if env_dict.get("WITH_XPU_FFT") == 'ON':
+        xpu_fft_lib_list = glob.glob(env_dict.get("XPU_FFT_LIB") + '*')
+        for xpu_fft_lib_file in xpu_fft_lib_list:
+            shutil.copy(xpu_fft_lib_file, libs_path)
+            package_data['paddle.libs'] += [os.path.basename(xpu_fft_lib_file)]
+
     if env_dict.get("WITH_XPU_XFT") == 'ON':
         shutil.copy(env_dict.get("XPU_XFT_LIB"), libs_path)
         package_data['paddle.libs'] += [env_dict.get("XPU_XFT_LIB_NAME")]
diff --git a/tools/xpu/get_xpufft_dependence.sh b/tools/xpu/get_xpufft_dependence.sh
@@ -0,0 +1,60 @@
+#!/bin/bash
+
+# Copyright (c) 2025 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.
+
+set -ex
+
+FFT_URL=$1
+FFT_DIR_NAME=$2
+
+if ! [ -n "$FFT_URL" ]; then
+  exit 0
+fi
+
+mkdir -p xpu/include/fft
+
+function download_from_bos() {
+  local url=$1
+  wget --no-check-certificate ${url} -q -O tmp.tar.gz
+  if [[ $? -ne 0 ]]; then
+    echo "downloading failed: ${url}"
+    exit 1
+  fi
+  tar xvf tmp.tar.gz
+  rm -f tmp.tar.gz
+}
+
+function check_files() {
+  local files=("$@")
+  for file in "${files[@]}";
+  do
+    echo "checking $file"
+    if [[ ! -f $file ]]; then
+        echo "checking failed: $file"
+        exit 1
+    else
+        echo "checking ok: $file"
+    fi
+  done
+}
+
+download_from_bos ${FFT_URL}
+check_files ${FFT_DIR_NAME}/include/cufft.h ${FFT_DIR_NAME}/lib64/libcufft.so
+cp -r ${FFT_DIR_NAME}/include/* xpu/include/xpu/
+cp -r ${FFT_DIR_NAME}/lib64/* xpu/lib/
+cp -r ${FFT_DIR_NAME}/include/* xpu/include/fft
+patchelf --set-rpath '$ORIGIN/' xpu/lib/libcufft.so
+patchelf --set-rpath '$ORIGIN/' xpu/lib/libcufft.so.10
+patchelf --set-rpath '$ORIGIN/' xpu/lib/libcufft.so.10.7.2.91
