paddle_cluster/setup_ops.py at paddle · PFCCLab/paddle_cluster · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
import glob
import os
import os.path as osp
import platform
import re
from itertools import product

import paddle
from paddle.utils.cpp_extension import CppExtension
from paddle.utils.cpp_extension import CUDAExtension
from paddle.utils.cpp_extension import setup
from paddle.utils.cpp_extension.cpp_extension import CUDA_HOME


def get_version():
    current_dir = osp.dirname(osp.abspath(__file__))
    with open(osp.join(current_dir, "paddle_cluster/__init__.py")) as f:
        content = f.read()
    version_match = re.search(r'__version__\s*=\s*["\']([^"\']+)["\']', content)
    if version_match:
        return version_match.group(1)

    raise RuntimeError("Cannot find __version__ in paddle_cluster/__init__.py")


__version__ = get_version()


WITH_CUDA = False
if paddle.device.cuda.device_count() > 0:
    WITH_CUDA = CUDA_HOME is not None

suffices = ["cpu", "cuda"] if WITH_CUDA else ["cpu"]
if os.getenv("FORCE_CUDA", "0") == "1":
    suffices = ["cuda", "cpu"]
if os.getenv("FORCE_ONLY_CUDA", "0") == "1":
    suffices = ["cuda"]
if os.getenv("FORCE_ONLY_CPU", "0") == "1":
    suffices = ["cpu"]

BUILD_DOCS = os.getenv("BUILD_DOCS", "0") == "1"

assert platform.system() == "Linux", "Only support build on linux now."


def set_cuda_archs():
    cuda_major, _ = paddle.version.cuda_version.split(".")
    if int(paddle.version.major) >= 3 and int(paddle.version.minor) >= 3:
        if int(cuda_major) >= 13:
            paddle_known_gpu_archs = ["Turing", "Ampere", "Ada", "Hopper", "Blackwell"]
        elif int(cuda_major) >= 12:
            paddle_known_gpu_archs = [
                "Pascal",
                "Volta",
                "Turing",
                "Ampere",
                "Ada",
                "Hopper",
            ]
        elif int(cuda_major) >= 11:
            paddle_known_gpu_archs = ["Pascal", "Volta", "Turing", "Ampere", "Ada"]
        elif int(cuda_major) >= 10:
            paddle_known_gpu_archs = ["Pascal", "Volta", "Turing"]
        else:
            raise ValueError("Not support cuda version.")

        os.environ["PADDLE_CUDA_ARCH_LIST"] = " ".join(
            [str(arch) for arch in paddle_known_gpu_archs]
        )
    else:
        # compatible with paddle < 3.3.0
        paddle_known_gpu_archs = paddle.version.compiled_cuda_archs

        os.environ["PADDLE_CUDA_ARCH_LIST"] = ",".join(
            [str(arch) for arch in paddle_known_gpu_archs]
        )


def get_extensions():
    extensions = []

    extensions_dir = osp.join("csrc")
    main_files = glob.glob(osp.join(extensions_dir, "*.cpp"))
    # remove generated 'hip' files, in case of rebuilds
    main_files = [path for path in main_files]

    define_macros = [("WITH_PYTHON", None)]
    undef_macros = []

    extra_compile_args = {"cxx": ["-O2"]}
    if not os.name == "nt":  # Not on Windows:
        extra_compile_args["cxx"] += ["-Wno-sign-compare"]
    extra_link_args = ["-s"]

    extra_compile_args["cxx"] += ["-fopenmp"]

    if "cuda" in suffices:
        set_cuda_archs()
        define_macros += [("WITH_CUDA", None)]
        nvcc_flags = os.getenv("NVCC_FLAGS", "")
        nvcc_flags = [] if nvcc_flags == "" else nvcc_flags.split(" ")
        nvcc_flags += ["-O2"]
        extra_compile_args["nvcc"] = nvcc_flags

        nvcc_flags += ["--expt-relaxed-constexpr"]

    sources = set()
    for main, suffix in product(main_files, suffices):
        name = main.split(os.sep)[-1][:-4]
        sources.add(main)

        path = osp.join(extensions_dir, "cpu", f"{name}_cpu.cpp")
        if osp.exists(path):
            sources.add(path)

        path = osp.join(extensions_dir, "cuda", f"{name}_cuda.cu")
        if suffix == "cuda" and osp.exists(path):
            sources.add(path)

    # Add csrc/cuda/utils.cu
    if "cuda" in suffices:
        path = osp.join(extensions_dir, "cuda/utils.cu")
        sources.add(path)

    Extension = CUDAExtension if "cuda" in suffices else CppExtension
    extension = Extension(
        list(sources),
        include_dirs=[extensions_dir],
        define_macros=define_macros,
        undef_macros=undef_macros,
        extra_compile_args=extra_compile_args,
        extra_link_args=extra_link_args,
    )
    extensions += [extension]

    return extensions


if __name__ == "__main__":
    setup(
        name="paddle_cluster_ops",
        version=__version__,
        description=(
            "Paddle Custom Operators Extension Library of Optimized Graph Cluster Algorithms"
        ),
        author="Ruibin Cheung",
        author_email="beinggod@foxmail.com",
        python_requires=">=3.8",
        ext_modules=get_extensions(),
    )