[xpu] add ep custom ops (#3911)

custom_ops/xpu_ops/setup_ops.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+
+# 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.
+"""
+Copyright (c) 2025 Baidu.com, Inc. All Rights Reserved.
+
+Build and setup XPU custom ops for ERNIE Bot.
+"""
+import os
+import shutil
+import subprocess
+from pathlib import Path
+
+import paddle
+from paddle.utils.cpp_extension import CppExtension, setup
+
+current_file = Path(__file__).resolve()
+base_dir = os.path.join(current_file.parent, "src")
+
+
+def build_plugin(CLANG_PATH, XRE_INC_DIR, XRE_LIB_DIR, XDNN_INC_DIR, XDNN_LIB_DIR):
+    """
+    build xpu plugin
+    """
+    current_working_directory = base_dir
+    print(f"Current working directory: {current_working_directory}")
+
+    # 设置环境变量
+    os.environ["XRE_INC_DIR"] = XRE_INC_DIR
+    os.environ["XRE_LIB_DIR"] = XRE_LIB_DIR
+    os.environ["XDNN_INC_DIR"] = XDNN_INC_DIR
+    os.environ["XDNN_LIB_DIR"] = XDNN_LIB_DIR
+
+    # 设置 Clang 路径
+    os.environ["CLANG_PATH"] = CLANG_PATH
+
+    # 删除指定目录
+    dirs_to_remove = [
+        "dist",
+        "fastdeploy_ops.egg-info",
+        "build",
+        "plugin/build",
+    ]
+    for dir_name in dirs_to_remove:
+        if os.path.exists(dir_name):
+            shutil.rmtree(dir_name)
+            print(f"Removed directory: {dir_name}")
+
+    # 在 plugin 目录中执行构建脚本
+    plugin_dir = "plugin"
+    build_script = os.path.join(current_working_directory, plugin_dir, "build.sh")
+
+    print("build_script: ", build_script)
+
+    if not os.path.isfile(build_script):
+        print(f"Error: Build script not found at {build_script}")
+        return
+
+    # 赋予执行权限 (如果尚未设置)
+    if not os.access(build_script, os.X_OK):
+        os.chmod(build_script, 0o755)
+
+    # 执行构建脚本
+    try:
+        print("Running build script...")
+        subprocess.run(
+            [build_script],
+            check=True,
+            cwd=os.path.join(current_working_directory, plugin_dir),
+        )
+        print("Build completed successfully.")
+    except subprocess.CalledProcessError as e:
+        print(f"Build failed with error: {e}")
+    except Exception as e:
+        print(f"Unexpected error: {e!s}")
+
+
+def xpu_setup_ops():
+    """
+    setup xpu ops
+    """
+    PADDLE_PATH = os.path.dirname(paddle.__file__)
+    PADDLE_INCLUDE_PATH = os.path.join(PADDLE_PATH, "include")
+    PADDLE_LIB_PATH = os.path.join(PADDLE_PATH, "libs")
+
+    BKCL_PATH = os.getenv("BKCL_PATH")
+    if BKCL_PATH is None:
+        BKCL_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xpu")
+        BKCL_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libbkcl.so")
+    else:
+        BKCL_INC_PATH = os.path.join(BKCL_PATH, "include")
+        BKCL_LIB_PATH = os.path.join(BKCL_PATH, "so", "libbkcl.so")
+
+    CLANG_PATH = os.getenv("CLANG_PATH")
+    assert CLANG_PATH is not None, "CLANG_PATH is not set."
+
+    XRE_PATH = os.getenv("XRE_PATH")
+    if XRE_PATH is None:
+        XRE_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xre")
+        XRE_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libxpucuda.so")
+        XRE_LIB_DIR = os.path.join(PADDLE_LIB_PATH)
+    else:
+        XRE_INC_PATH = os.path.join(XRE_PATH, "include")
+        XRE_LIB_PATH = os.path.join(XRE_PATH, "so", "libxpucuda.so")
+        XRE_LIB_DIR = os.path.join(XRE_PATH, "so")
+
+    XDNN_PATH = os.getenv("XDNN_PATH")
+    if XDNN_PATH is None:
+        XDNN_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH)
+        XDNN_LIB_DIR = os.path.join(PADDLE_LIB_PATH)
+    else:
+        XDNN_INC_PATH = os.path.join(XDNN_PATH, "include")
+        XDNN_LIB_DIR = os.path.join(XDNN_PATH, "so")
+
+    XVLLM_PATH = os.getenv("XVLLM_PATH")
+    assert XVLLM_PATH is not None, "XVLLM_PATH is not set."
+    XVLLM_KERNEL_INC_PATH = os.path.join(XVLLM_PATH, "infer_ops", "include")
+    XVLLM_KERNEL_LIB_PATH = os.path.join(XVLLM_PATH, "infer_ops", "so", "libapiinfer.so")
+    XVLLM_KERNEL_LIB_DIR = os.path.join(XVLLM_PATH, "infer_ops", "so")
+    XVLLM_OP_INC_PATH = os.path.join(XVLLM_PATH, "xft_blocks", "include")
+    XVLLM_OP_LIB_PATH = os.path.join(XVLLM_PATH, "xft_blocks", "so", "libxft_blocks.so")
+    XVLLM_OP_LIB_DIR = os.path.join(XVLLM_PATH, "xft_blocks", "so")
+
+    # build plugin
+    build_plugin(CLANG_PATH, XRE_INC_PATH, XRE_LIB_DIR, XDNN_INC_PATH, XDNN_LIB_DIR)
+
+    ops = []
+    for root, dirs, files in os.walk(os.path.join(base_dir, "ops")):
+        for file in files:
+            if file.endswith(".cc"):
+                ops.append(os.path.join(root, file))
+
+    include_dirs = [
+        os.path.join(base_dir, "./"),
+        os.path.join(base_dir, "./plugin/include"),
+        BKCL_INC_PATH,
+        XRE_INC_PATH,
+        XVLLM_KERNEL_INC_PATH,
+        XVLLM_OP_INC_PATH,
+    ]
+    extra_objects = [
+        os.path.join(base_dir, "./plugin/build/libxpuplugin.a"),
+        BKCL_LIB_PATH,
+        XRE_LIB_PATH,
+        XVLLM_KERNEL_LIB_PATH,
+        XVLLM_OP_LIB_PATH,
+    ]
+
+    setup(
+        name="fastdeploy_ops",
+        ext_modules=CppExtension(
+            sources=ops,
+            include_dirs=include_dirs,
+            extra_objects=extra_objects,
+            extra_compile_args={
+                "cxx": [
+                    "-D_GLIBCXX_USE_CXX11_ABI=1",
+                    "-DPADDLE_WITH_XPU",
+                    "-DBUILD_MULTI_XPU",
+                ]
+            },
+            runtime_library_dirs=[XVLLM_KERNEL_LIB_DIR, XVLLM_OP_LIB_DIR],
+        ),
+    )
+
+
+if __name__ == "__main__":
+    xpu_setup_ops()