98 lines
2.9 KiB
Python
98 lines
2.9 KiB
Python
from setuptools import setup
|
|
import os
|
|
import platform
|
|
import sys
|
|
|
|
ROOT = os.path.dirname(os.path.abspath(__file__))
|
|
BUILD_TARGET = os.environ.get("BUILD_TARGET", "auto")
|
|
|
|
ext_modules = []
|
|
cmdclass = {}
|
|
|
|
if BUILD_TARGET == "cpu" or (BUILD_TARGET == "auto" and platform.system() == "Darwin"):
|
|
# CPU-only build: compile only C++ sources (Eigen-based), skip CUDA
|
|
from torch.utils.cpp_extension import CppExtension, BuildExtension
|
|
|
|
cpu_sources = [
|
|
"src/convert/flexible_dual_grid.cpp",
|
|
"src/convert/volumetic_attr.cpp",
|
|
"src/io/svo.cpp",
|
|
"src/io/filter_parent.cpp",
|
|
"src/io/filter_neighbor.cpp",
|
|
"src/ext_cpu.cpp",
|
|
]
|
|
# Only build if the CPU ext entry point exists
|
|
ext_cpu_path = os.path.join(ROOT, "src/ext_cpu.cpp")
|
|
if os.path.exists(ext_cpu_path):
|
|
ext_modules = [
|
|
CppExtension(
|
|
name="o_voxel._C",
|
|
sources=cpu_sources,
|
|
include_dirs=[
|
|
os.path.join(ROOT, "third_party/eigen"),
|
|
],
|
|
extra_compile_args={
|
|
"cxx": ["-O3", "-std=c++17"],
|
|
},
|
|
)
|
|
]
|
|
else:
|
|
# No C++ extension — pure Python fallback
|
|
print("[o_voxel] Building without C++ extension (pure Python fallback)")
|
|
ext_modules = []
|
|
cmdclass = {'build_ext': BuildExtension}
|
|
else:
|
|
# CUDA/ROCm build (original)
|
|
from torch.utils.cpp_extension import CUDAExtension, BuildExtension, IS_HIP_EXTENSION
|
|
|
|
if BUILD_TARGET == "auto":
|
|
IS_HIP = IS_HIP_EXTENSION
|
|
elif BUILD_TARGET == "cuda":
|
|
IS_HIP = False
|
|
elif BUILD_TARGET == "rocm":
|
|
IS_HIP = True
|
|
|
|
if not IS_HIP:
|
|
cc_flag = []
|
|
else:
|
|
archs = os.getenv("GPU_ARCHS", "native").split(";")
|
|
cc_flag = [f"--offload-arch={arch}" for arch in archs]
|
|
|
|
ext_modules = [
|
|
CUDAExtension(
|
|
name="o_voxel._C",
|
|
sources=[
|
|
"src/hash/hash.cu",
|
|
"src/convert/flexible_dual_grid.cpp",
|
|
"src/convert/volumetic_attr.cpp",
|
|
"src/serialize/api.cu",
|
|
"src/serialize/hilbert.cu",
|
|
"src/serialize/z_order.cu",
|
|
"src/io/svo.cpp",
|
|
"src/io/filter_parent.cpp",
|
|
"src/io/filter_neighbor.cpp",
|
|
"src/rasterize/rasterize.cu",
|
|
"src/ext.cpp",
|
|
],
|
|
include_dirs=[
|
|
os.path.join(ROOT, "third_party/eigen"),
|
|
],
|
|
extra_compile_args={
|
|
"cxx": ["-O3", "-std=c++17"],
|
|
"nvcc": ["-O3", "-std=c++17"] + cc_flag,
|
|
},
|
|
)
|
|
]
|
|
cmdclass = {'build_ext': BuildExtension}
|
|
|
|
setup(
|
|
name="o_voxel",
|
|
packages=[
|
|
'o_voxel',
|
|
'o_voxel.convert',
|
|
'o_voxel.io',
|
|
],
|
|
ext_modules=ext_modules,
|
|
cmdclass=cmdclass,
|
|
)
|