Skip to content

Commit 655f663

Browse files
mikethegoblinmikethegoblin
authored
[UIL] Enable installing flagcx as a python package (#405)
1 parent ddb907d commit 655f663

5 files changed

Lines changed: 391 additions & 122 deletions

File tree

‎.gitignore‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,11 @@ debian-packages
77
# Ignore compiled Python files and shared object files
88
plugin/*/*.so
99
plugin/*/*/*.so
10-
plugin/*/*/__pycache__
10+
__pycache__/
11+
**/__pycache__/
1112

12-
# Ignore egg-info directory
13+
# Ignore egg-info directories
14+
*.egg-info/
1315
plugin/*/*.egg-info
1416

1517
# Ignore vscode workspace folder

‎plugin/torch/_build_config.py‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
"""
2+
Shared build configuration for the flagcx torch plugin.
3+
4+
Used by both the root setup.py and plugin/torch/setup.py to avoid
5+
duplicating adaptor detection, device-specific paths, and extension
6+
class selection logic.
7+
"""
8+
9+
import os
10+
import sys
11+
12+
from packaging.version import Version, parse as vparse
13+
14+
# ---------------------------------------------------------------------------
15+
# Adaptor name -> C++ define flag
16+
# ---------------------------------------------------------------------------
17+
18+
ADAPTOR_MAP = {
19+
"nvidia": "-DUSE_NVIDIA_ADAPTOR",
20+
"iluvatar_corex": "-DUSE_ILUVATAR_COREX_ADAPTOR",
21+
"cambricon": "-DUSE_CAMBRICON_ADAPTOR",
22+
"metax": "-DUSE_METAX_ADAPTOR",
23+
"musa": "-DUSE_MUSA_ADAPTOR",
24+
"du": "-DUSE_DU_ADAPTOR",
25+
"klx": "-DUSE_KUNLUNXIN_ADAPTOR",
26+
"ascend": "-DUSE_ASCEND_ADAPTOR",
27+
"amd": "-DUSE_AMD_ADAPTOR",
28+
"tsm": "-DUSE_TSM_ADAPTOR",
29+
"enflame": "-DUSE_ENFLAME_ADAPTOR",
30+
}
31+
32+
# Adaptor name -> Make variable (for root setup.py make invocation)
33+
ADAPTOR_TO_MAKE_FLAG = {
34+
"nvidia": "USE_NVIDIA",
35+
"ascend": "USE_ASCEND",
36+
"iluvatar_corex": "USE_ILUVATAR_COREX",
37+
"cambricon": "USE_CAMBRICON",
38+
"metax": "USE_METAX",
39+
"musa": "USE_MUSA",
40+
"klx": "USE_KUNLUNXIN",
41+
"du": "USE_DU",
42+
"amd": "USE_AMD",
43+
"tsm": "USE_TSM",
44+
"enflame": "USE_ENFLAME",
45+
}
46+
47+
VALID_ADAPTORS = list(ADAPTOR_MAP.keys())
48+
49+
50+
def detect_adaptor():
51+
"""Detect the adaptor from FLAGCX_ADAPTOR env var, --adaptor CLI arg, or
52+
USE_* env vars. Returns the adaptor name string. Defaults to 'nvidia'."""
53+
adaptor = os.environ.get("FLAGCX_ADAPTOR", "").strip()
54+
55+
# Check --adaptor CLI argument (consumed from sys.argv)
56+
if not adaptor and "--adaptor" in sys.argv:
57+
arg_index = sys.argv.index("--adaptor")
58+
sys.argv.remove("--adaptor")
59+
if arg_index < len(sys.argv):
60+
adaptor = sys.argv[arg_index]
61+
sys.argv.remove(adaptor)
62+
else:
63+
print("No adaptor provided after '--adaptor'. Using default nvidia adaptor")
64+
65+
# Check USE_* env vars
66+
if not adaptor:
67+
for name, make_flag in ADAPTOR_TO_MAKE_FLAG.items():
68+
if os.environ.get(make_flag, "0") == "1":
69+
adaptor = name
70+
break
71+
72+
# Default
73+
if not adaptor:
74+
adaptor = "nvidia"
75+
76+
assert adaptor in VALID_ADAPTORS, f"Invalid adaptor: {adaptor}. Valid: {VALID_ADAPTORS}"
77+
return adaptor
78+
79+
80+
def detect_torch_flag():
81+
"""Detect the torch version flag for conditional compilation."""
82+
torch_flag = "-DTORCH_VER_LT_250"
83+
try:
84+
import torch
85+
torch_version = vparse(torch.__version__.split("+")[0])
86+
if torch_version >= Version("2.5.0"):
87+
print("torch version >= 2.5.0, set TORCH_VER_GE_250 flag")
88+
torch_flag = "-DTORCH_VER_GE_250"
89+
except ImportError:
90+
print("Warning: torch not found.")
91+
return torch_flag
92+
93+
94+
def get_device_config(adaptor_flag):
95+
"""Return (extra_include_dirs, extra_library_dirs, extra_libs) for the
96+
given adaptor define flag."""
97+
include_dirs = []
98+
library_dirs = []
99+
libs = []
100+
101+
if adaptor_flag == "-DUSE_NVIDIA_ADAPTOR":
102+
include_dirs += ["/usr/local/cuda/include"]
103+
library_dirs += ["/usr/local/cuda/lib64"]
104+
libs += ["cuda", "cudart", "c10_cuda", "torch_cuda"]
105+
elif adaptor_flag == "-DUSE_ILUVATAR_COREX_ADAPTOR":
106+
include_dirs += ["/usr/local/corex/include"]
107+
library_dirs += ["/usr/local/corex/lib64"]
108+
libs += ["cuda", "cudart", "c10_cuda", "torch_cuda"]
109+
elif adaptor_flag == "-DUSE_CAMBRICON_ADAPTOR":
110+
import torch_mlu
111+
neuware_home_path = os.getenv("NEUWARE_HOME")
112+
torch_mlu_path = torch_mlu.__file__.split("__init__")[0]
113+
torch_mlu_lib_dir = os.path.join(torch_mlu_path, "csrc/lib/")
114+
torch_mlu_include_dir = os.path.join(torch_mlu_path, "csrc/")
115+
torch_mlu_include_dir2 = os.path.join(torch_mlu_path, "csrc", "include")
116+
include_dirs += [f"{neuware_home_path}/include", torch_mlu_include_dir, torch_mlu_include_dir2]
117+
library_dirs += [f"{neuware_home_path}/lib64", torch_mlu_lib_dir]
118+
libs += ["cnrt", "cncl", "torch_mlu"]
119+
elif adaptor_flag == "-DUSE_METAX_ADAPTOR":
120+
include_dirs += ["/opt/maca/include"]
121+
library_dirs += ["/opt/maca/lib64"]
122+
libs += ["cuda", "cudart", "c10_cuda", "torch_cuda"]
123+
elif adaptor_flag == "-DUSE_MUSA_ADAPTOR":
124+
import torch_musa
125+
pytorch_musa_install_path = os.path.dirname(os.path.abspath(torch_musa.__file__))
126+
pytorch_library_path = os.path.join(pytorch_musa_install_path, "lib")
127+
library_dirs += ["/usr/local/musa/lib/", pytorch_library_path]
128+
libs += ["musa", "musart"]
129+
elif adaptor_flag == "-DUSE_DU_ADAPTOR":
130+
include_dirs += ["${CUDA_PATH}/include"]
131+
library_dirs += ["${CUDA_PATH}/lib64"]
132+
libs += ["cuda", "cudart", "c10_cuda", "torch_cuda"]
133+
elif adaptor_flag == "-DUSE_KUNLUNXIN_ADAPTOR":
134+
include_dirs += ["/opt/kunlun/include"]
135+
library_dirs += ["/opt/kunlun/lib"]
136+
libs += ["cuda", "cudart", "c10_cuda", "torch_cuda"]
137+
elif adaptor_flag == "-DUSE_ASCEND_ADAPTOR":
138+
import torch_npu
139+
pytorch_npu_install_path = os.path.dirname(os.path.abspath(torch_npu.__file__))
140+
pytorch_library_path = os.path.join(pytorch_npu_install_path, "lib")
141+
include_dirs += [os.path.join(pytorch_npu_install_path, "include")]
142+
library_dirs += [pytorch_library_path]
143+
libs += ["torch_npu"]
144+
elif adaptor_flag == "-DUSE_AMD_ADAPTOR":
145+
include_dirs += ["/opt/rocm/include"]
146+
library_dirs += ["/opt/rocm/lib"]
147+
libs += ["hiprtc", "c10_hip", "torch_hip"]
148+
elif adaptor_flag == "-DUSE_TSM_ADAPTOR":
149+
import torch_txda
150+
txda_install_path = os.path.dirname(os.path.abspath(torch_txda.__file__))
151+
txda_library_path = os.path.join(txda_install_path, "lib")
152+
include_dirs += ["/usr/local/kuiper/include", os.path.join(txda_install_path, "include")]
153+
library_dirs += ["/usr/local/kuiper/lib", txda_library_path]
154+
libs += ["torch_txda", "hpgr"]
155+
elif adaptor_flag == "-DUSE_ENFLAME_ADAPTOR":
156+
import torch_gcu
157+
pytorch_gcu_install_path = os.path.dirname(os.path.abspath(torch_gcu.__file__))
158+
pytorch_library_path = os.path.join(pytorch_gcu_install_path, "lib")
159+
include_dirs += ["/opt/tops/include", os.path.join(pytorch_gcu_install_path, "include")]
160+
library_dirs += ["/opt/tops/lib", pytorch_library_path]
161+
libs += ["topsrt", "torch_gcu"]
162+
163+
return include_dirs, library_dirs, libs
164+
165+
166+
def get_ext_classes(adaptor_flag):
167+
"""Return (CppExtension, BuildExtension) for the given adaptor, or
168+
(None, None) if unavailable."""
169+
try:
170+
if adaptor_flag == "-DUSE_MUSA_ADAPTOR":
171+
from torch_musa.utils.musa_extension import MUSAExtension as CppExtension
172+
from torch_musa.utils.musa_extension import BuildExtension
173+
elif adaptor_flag == "-DUSE_ENFLAME_ADAPTOR":
174+
from tops_extension import TopsBuildExtension as BuildExtension
175+
from tops_extension.torch import TopsTorchExtension as CppExtension
176+
else:
177+
from torch.utils.cpp_extension import CppExtension, BuildExtension
178+
return CppExtension, BuildExtension
179+
except ImportError:
180+
print("Warning: CppExtension or BuildExtension not found.")
181+
return None, None

‎plugin/torch/setup.py‎

Lines changed: 26 additions & 120 deletions
Original file line numberDiff line numberDiff line change
@@ -5,143 +5,49 @@
55
os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
66

77
# Modern setuptools (>=64) uses pip for 'develop' which creates isolated build envs.
8-
# For packages depending on torch, this often fails.
8+
# For packages depending on torch, this often fails.
99
# We try to disable build isolation if not explicitly set.
1010
if "PIP_NO_BUILD_ISOLATION" not in os.environ:
1111
os.environ["PIP_NO_BUILD_ISOLATION"] = "1"
1212

1313
from setuptools import setup, find_packages
14-
from packaging.version import Version, parse as vparse
15-
16-
adaptor = os.environ.get("FLAGCX_ADAPTOR", "nvidia")
17-
if '--adaptor' in sys.argv:
18-
arg_index = sys.argv.index('--adaptor')
19-
sys.argv.remove("--adaptor")
20-
if arg_index < len(sys.argv):
21-
adaptor = sys.argv[arg_index]
22-
sys.argv.remove(adaptor)
23-
else:
24-
print("No adaptor provided after '--adaptor'. Using default nvidia adaptor")
14+
from _build_config import (
15+
ADAPTOR_MAP,
16+
detect_adaptor,
17+
detect_torch_flag,
18+
get_device_config,
19+
get_ext_classes,
20+
)
2521

26-
valid_adaptors = ["nvidia", "iluvatar_corex", "cambricon", "metax", "du", "klx", "ascend", "musa", "amd", "enflame"]
27-
assert adaptor in valid_adaptors, f"Invalid adaptor: {adaptor}"
22+
adaptor = detect_adaptor()
2823
print(f"Using {adaptor} adaptor")
2924

30-
adaptor_map = {
31-
"nvidia": "-DUSE_NVIDIA_ADAPTOR",
32-
"iluvatar_corex": "-DUSE_ILUVATAR_COREX_ADAPTOR",
33-
"cambricon": "-DUSE_CAMBRICON_ADAPTOR",
34-
"metax": "-DUSE_METAX_ADAPTOR",
35-
"musa": "-DUSE_MUSA_ADAPTOR",
36-
"du": "-DUSE_DU_ADAPTOR",
37-
"klx": "-DUSE_KUNLUNXIN_ADAPTOR",
38-
"ascend": "-DUSE_ASCEND_ADAPTOR",
39-
"amd": "-DUSE_AMD_ADAPTOR",
40-
"tsm": "-DUSE_TSM_ADAPTOR",
41-
"enflame": "-DUSE_ENFLAME_ADAPTOR"
42-
}
43-
adaptor_flag = adaptor_map[adaptor]
44-
torch_flag = "-DTORCH_VER_LT_250"
25+
adaptor_flag = ADAPTOR_MAP[adaptor]
26+
torch_flag = detect_torch_flag()
27+
28+
plugin_dir = os.path.dirname(os.path.abspath(__file__))
29+
repo_root = os.path.join(plugin_dir, "..", "..")
4530

4631
sources = ["flagcx/src/backend_flagcx.cpp", "flagcx/src/utils_flagcx.cpp"]
4732
include_dirs = [
48-
f"{os.path.dirname(os.path.abspath(__file__))}/flagcx/include",
49-
f"{os.path.dirname(os.path.abspath(__file__))}/../../flagcx/include",
50-
f"{os.path.dirname(os.path.abspath(__file__))}/../../third-party/json/single_include",
33+
os.path.join(plugin_dir, "flagcx", "include"),
34+
os.path.join(repo_root, "flagcx", "include"),
35+
os.path.join(repo_root, "third-party", "json", "single_include"),
5136
]
5237

5338
library_dirs = [
54-
f"{os.path.dirname(os.path.abspath(__file__))}/../../build/lib",
39+
os.path.join(repo_root, "build", "lib"),
5540
]
5641

5742
libs = ["flagcx"]
5843

59-
try:
60-
import torch
61-
torch_version = vparse(torch.__version__.split("+")[0])
62-
if torch_version >= Version("2.5.0"):
63-
print("torch version >= 2.5.0, set TORCH_VER_GE_250 flag")
64-
torch_flag = "-DTORCH_VER_GE_250"
65-
except ImportError:
66-
torch = None
67-
print("Warning: torch not found.")
68-
69-
if adaptor_flag == "-DUSE_NVIDIA_ADAPTOR":
70-
include_dirs += ["/usr/local/cuda/include"]
71-
library_dirs += ["/usr/local/cuda/lib64"]
72-
libs += ["cuda", "cudart", "c10_cuda", "torch_cuda"]
73-
elif adaptor_flag == "-DUSE_ILUVATAR_COREX_ADAPTOR":
74-
include_dirs += ["/usr/local/corex/include"]
75-
library_dirs += ["/usr/local/corex/lib64"]
76-
libs += ["cuda", "cudart", "c10_cuda", "torch_cuda"]
77-
elif adaptor_flag == "-DUSE_CAMBRICON_ADAPTOR":
78-
import torch_mlu
79-
neuware_home_path=os.getenv("NEUWARE_HOME")
80-
pytorch_home_path=os.getenv("PYTORCH_HOME")
81-
torch_mlu_path = torch_mlu.__file__.split("__init__")[0]
82-
torch_mlu_lib_dir = os.path.join(torch_mlu_path, "csrc/lib/")
83-
torch_mlu_include_dir = os.path.join(torch_mlu_path, "csrc/")
84-
torch_mlu_include_dir2 = os.path.join(torch_mlu_path, "csrc", "include")
85-
include_dirs += [f"{neuware_home_path}/include", torch_mlu_include_dir, torch_mlu_include_dir2]
86-
library_dirs += [f"{neuware_home_path}/lib64", torch_mlu_lib_dir]
87-
libs += ["cnrt", "cncl", "torch_mlu"]
88-
elif adaptor_flag == "-DUSE_METAX_ADAPTOR":
89-
include_dirs += ["/opt/maca/include"]
90-
library_dirs += ["/opt/maca/lib64"]
91-
libs += ["cuda", "cudart", "c10_cuda", "torch_cuda"]
92-
elif adaptor_flag == "-DUSE_MUSA_ADAPTOR":
93-
import torch_musa
94-
pytorch_musa_install_path = os.path.dirname(os.path.abspath(torch_musa.__file__))
95-
pytorch_library_path = os.path.join(pytorch_musa_install_path, "lib")
96-
library_dirs += ['/usr/local/musa/lib/',pytorch_library_path]
97-
libs += ["musa","musart"]
98-
elif adaptor_flag == "-DUSE_DU_ADAPTOR":
99-
include_dirs += ["${CUDA_PATH}/include"]
100-
library_dirs += ["${CUDA_PATH}/lib64"]
101-
libs += ["cuda", "cudart", "c10_cuda", "torch_cuda"]
102-
elif adaptor_flag == "-DUSE_KUNLUNXIN_ADAPTOR":
103-
include_dirs += ["/opt/kunlun/include"]
104-
library_dirs += ["/opt/kunlun/lib"]
105-
libs += ["cuda", "cudart", "c10_cuda", "torch_cuda"]
106-
elif adaptor_flag == "-DUSE_ASCEND_ADAPTOR":
107-
import torch_npu
108-
pytorch_npu_install_path = os.path.dirname(os.path.abspath(torch_npu.__file__))
109-
pytorch_library_path = os.path.join(pytorch_npu_install_path, "lib")
110-
include_dirs += [os.path.join(pytorch_npu_install_path, "include")]
111-
library_dirs += [pytorch_library_path]
112-
libs += ["torch_npu"]
113-
elif adaptor_flag == "-DUSE_AMD_ADAPTOR":
114-
include_dirs += ["/opt/rocm/include"]
115-
library_dirs += ["/opt/rocm/lib"]
116-
libs += ["hiprtc", "c10_hip", "torch_hip"]
117-
elif adaptor_flag == "-DUSE_TSM_ADAPTOR":
118-
import torch_txda
119-
txda_install_path = os.path.dirname(os.path.abspath(torch_txda.__file__))
120-
txda_library_path = os.path.join(txda_install_path, "lib")
121-
include_dirs += ["/usr/local/kuiper/include", os.path.join(txda_install_path, "include")]
122-
library_dirs += ["/usr/local/kuiper/lib",txda_library_path]
123-
libs += ["torch_txda", "hpgr"]
124-
elif adaptor_flag == "-DUSE_ENFLAME_ADAPTOR":
125-
import torch_gcu
126-
pytorch_gcu_install_path = os.path.dirname(os.path.abspath(torch_gcu.__file__))
127-
pytorch_library_path = os.path.join(pytorch_gcu_install_path, "lib")
128-
include_dirs += ["/opt/tops/include", os.path.join(pytorch_gcu_install_path, "include")]
129-
library_dirs += ["/opt/tops/lib", pytorch_library_path]
130-
libs += ["topsrt", "torch_gcu"]
44+
# Add device-specific paths
45+
dev_includes, dev_libdirs, dev_libs = get_device_config(adaptor_flag)
46+
include_dirs += dev_includes
47+
library_dirs += dev_libdirs
48+
libs += dev_libs
13149

132-
try:
133-
if adaptor_flag == "-DUSE_MUSA_ADAPTOR":
134-
from torch_musa.utils.musa_extension import MUSAExtension as CppExtension
135-
from torch_musa.utils.musa_extension import BuildExtension
136-
elif adaptor_flag == "-DUSE_ENFLAME_ADAPTOR":
137-
from tops_extension import TopsBuildExtension as BuildExtension
138-
from tops_extension.torch import TopsTorchExtension as CppExtension
139-
else:
140-
from torch.utils.cpp_extension import CppExtension, BuildExtension
141-
except ImportError:
142-
CppExtension = None
143-
BuildExtension = None
144-
print("Warning: CppExtension or BuildExtension not found.")
50+
CppExtension, BuildExtension = get_ext_classes(adaptor_flag)
14551

14652
ext_modules = []
14753
if CppExtension is not None:
@@ -152,7 +58,7 @@
15258
extra_compile_args={
15359
'cxx': [adaptor_flag, torch_flag]
15460
},
155-
extra_link_args=["-Wl,-rpath,"+f"{os.path.dirname(os.path.abspath(__file__))}/../../build/lib"],
61+
extra_link_args=["-Wl,-rpath," + os.path.join(repo_root, "build", "lib")],
15662
library_dirs=library_dirs,
15763
libraries=libs,
15864
)
@@ -164,7 +70,7 @@
16470

16571
setup(
16672
name="flagcx",
167-
version="0.8.0",
73+
version="0.10.0",
16874
ext_modules=ext_modules,
16975
cmdclass=cmdclass,
17076
packages=find_packages(),

‎pyproject.toml‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
[build-system]
2+
requires = ["setuptools", "wheel", "torch", "packaging"]
3+
build-backend = "setuptools.build_meta"
4+
5+
[project]
6+
name = "flagcx"
7+
version = "0.10.0"
8+
description = "FlagCX: A unified collective communication library"
9+
requires-python = ">=3.8"

0 commit comments

Comments
 (0)