Skip to content

[Fix] Add PADDLE_WITH_DNNL macro detection in PaddleFleet build #691

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
risemeup1 wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

[Fix] Add PADDLE_WITH_DNNL macro detection in PaddleFleet build #691

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 25 additions & 3 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,11 +104,15 @@ def write_wheelfile(self, wheelfile_base, generator=None):


def setup_ops_extension():
from paddle.base.core import is_compiled_with_onednn
from paddle.utils.cpp_extension import CUDAExtension, setup

from build_utils import get_cuda_version

# 定义 NVCC 编译参数
# Detect Paddle build config to ensure ABI compatibility
paddle_compiled_with_onednn = is_compiled_with_onednn()

# Define NVCC compile flags
nvcc_args = [
"-O3",
"-U__CUDA_NO_HALF_OPERATORS__",
Expand All @@ -127,11 +131,29 @@ def setup_ops_extension():
"-gencode=arch=compute_100,code=sm_100",
"-DNDEBUG",
]

# Define PADDLE_WITH_DNNL if Paddle was built with OneDNN,
# to ensure ABI compatibility for DenseTensor, etc.
if paddle_compiled_with_onednn:
nvcc_args.append("-DPADDLE_WITH_DNNL")

cuda_major, cuda_minor = get_cuda_version()

if cuda_major == 12 and cuda_minor < 8:
nvcc_args = [arg for arg in nvcc_args if "compute_100" not in arg]

# Add oneDNN include path if Paddle was built with OneDNN,
# otherwise Paddle headers that #include "dnnl.hpp" will fail
extra_include_dirs = []
if paddle_compiled_with_onednn:
import paddle
paddle_include_dir = os.path.join(os.path.dirname(paddle.__file__), "include")
onednn_include_dir = os.path.join(
paddle_include_dir, "third_party", "install", "onednn", "include"
)
if os.path.exists(onednn_include_dir):
extra_include_dirs.append(onednn_include_dir)

ext_module = CUDAExtension(
sources=[
# cpp files
Expand All @@ -153,15 +175,15 @@ def setup_ops_extension():
],
include_dirs=[
os.path.join(os.getcwd(), "src/paddlefleet/_extensions"),
],
] + extra_include_dirs,
extra_compile_args={
"cxx": [
"-O3",
"-w",
"-Wno-abi",
"-fPIC",
"-std=c++17",
],
] + (["-DPADDLE_WITH_DNNL"] if paddle_compiled_with_onednn else []),
"nvcc": nvcc_args,
},
)
Expand Down
Loading