Skip to content

GitLab

Commit 7d436274 authored by ashawkey's avatar ashawkey
Browse files

fix building for torch 2

parent d3e923c9
Hide whitespace changes
Inline Side-by-side
Showing with 172 additions and 135 deletions
setup.py
View file @ 7d436274
import os import os
import re
import subprocess
from pkg_resources import parse_version
from setuptools import setup from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension from torch.utils.cpp_extension import BuildExtension, CUDAExtension
...@@ -6,106 +9,139 @@ _src_path = os.path.dirname(os.path.abspath(__file__)) ...@@ -6,106 +9,139 @@ _src_path = os.path.dirname(os.path.abspath(__file__))
# ref: https://github.com/sxyu/sdf/blob/master/setup.py # ref: https://github.com/sxyu/sdf/blob/master/setup.py
def find_eigen(min_ver=(3, 3, 0)): def find_eigen(min_ver=(3, 3, 0)):
"""Helper to find or download the Eigen C++ library""" """Helper to find or download the Eigen C++ library"""
import re, os import re, os
try_paths = [ try_paths = [
'/usr/include/eigen3', '/usr/include/eigen3',
'/usr/local/include/eigen3', '/usr/local/include/eigen3',
os.path.expanduser('~/.local/include/eigen3'), os.path.expanduser('~/.local/include/eigen3'),
'C:/Program Files/eigen3', 'C:/Program Files/eigen3',
'C:/Program Files (x86)/eigen3', 'C:/Program Files (x86)/eigen3',
] ]
WORLD_VER_STR = "#define EIGEN_WORLD_VERSION" WORLD_VER_STR = "#define EIGEN_WORLD_VERSION"
MAJOR_VER_STR = "#define EIGEN_MAJOR_VERSION" MAJOR_VER_STR = "#define EIGEN_MAJOR_VERSION"
MINOR_VER_STR = "#define EIGEN_MINOR_VERSION" MINOR_VER_STR = "#define EIGEN_MINOR_VERSION"
EIGEN_WEB_URL = 'https://gitlab.com/libeigen/eigen/-/archive/3.3.7/eigen-3.3.7.tar.bz2' EIGEN_WEB_URL = 'https://gitlab.com/libeigen/eigen/-/archive/3.3.7/eigen-3.3.7.tar.bz2'
TMP_EIGEN_FILE = 'tmp_eigen.tar.bz2' TMP_EIGEN_FILE = 'tmp_eigen.tar.bz2'
TMP_EIGEN_DIR = 'eigen-3.3.7' TMP_EIGEN_DIR = './eigen-3.3.7'
min_ver_str = '.'.join(map(str, min_ver)) min_ver_str = '.'.join(map(str, min_ver))
eigen_path = None eigen_path = None
for path in try_paths: for path in try_paths:
macros_path = os.path.join(path, 'Eigen/src/Core/util/Macros.h') macros_path = os.path.join(path, 'Eigen/src/Core/util/Macros.h')
if os.path.exists(macros_path): if os.path.exists(macros_path):
macros = open(macros_path, 'r').read().split('\n') macros = open(macros_path, 'r').read().split('\n')
world_ver, major_ver, minor_ver = None, None, None world_ver, major_ver, minor_ver = None, None, None
for line in macros: for line in macros:
if line.startswith(WORLD_VER_STR): if line.startswith(WORLD_VER_STR):
world_ver = int(line[len(WORLD_VER_STR):]) world_ver = int(line[len(WORLD_VER_STR):])
elif line.startswith(MAJOR_VER_STR): elif line.startswith(MAJOR_VER_STR):
major_ver = int(line[len(MAJOR_VER_STR):]) major_ver = int(line[len(MAJOR_VER_STR):])
elif line.startswith(MINOR_VER_STR): elif line.startswith(MINOR_VER_STR):
minor_ver = int(line[len(MINOR_VER_STR):]) minor_ver = int(line[len(MINOR_VER_STR):])
if world_ver is None or major_ver is None or minor_ver is None: if world_ver is None or major_ver is None or minor_ver is None:
print('Failed to parse macros file', macros_path) print('Failed to parse macros file', macros_path)
else: else:
ver = (world_ver, major_ver, minor_ver) ver = (world_ver, major_ver, minor_ver)
ver_str = '.'.join(map(str, ver)) ver_str = '.'.join(map(str, ver))
if ver < min_ver: if ver < min_ver:
print('Found unsuitable Eigen version', ver_str, 'at', print('Found unsuitable Eigen version', ver_str, 'at',
path, '(need >= ' + min_ver_str + ')') path, '(need >= ' + min_ver_str + ')')
else: else:
print('Found Eigen version', ver_str, 'at', path) print('Found Eigen version', ver_str, 'at', path)
eigen_path = path eigen_path = path
break break
if eigen_path is None: if eigen_path is None:
try: try:
import urllib.request import urllib.request
print("Couldn't find Eigen locally, downloading...") print("Couldn't find Eigen locally, downloading...")
req = urllib.request.Request( req = urllib.request.Request(
EIGEN_WEB_URL, EIGEN_WEB_URL,
data=None, data=None,
headers={ headers={
'User-Agent': 'User-Agent':
'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4147.135 Safari/537.36' 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4147.135 Safari/537.36'
}) })
with urllib.request.urlopen(req) as resp,\ with urllib.request.urlopen(req) as resp,\
open(TMP_EIGEN_FILE, 'wb') as file: open(TMP_EIGEN_FILE, 'wb') as file:
data = resp.read() data = resp.read()
file.write(data) file.write(data)
import tarfile import tarfile
tar = tarfile.open(TMP_EIGEN_FILE) tar = tarfile.open(TMP_EIGEN_FILE)
tar.extractall() tar.extractall()
tar.close() tar.close()
eigen_path = TMP_EIGEN_DIR eigen_path = TMP_EIGEN_DIR
os.remove(TMP_EIGEN_FILE) os.remove(TMP_EIGEN_FILE)
except: except:
print('Download failed, failed to find Eigen') print('Download failed, failed to find Eigen')
if eigen_path is not None: if eigen_path is not None:
print('Found eigen at', eigen_path) print('Found eigen at', eigen_path)
return eigen_path return eigen_path
nvcc_flags = [
'-O3', '-std=c++14', if os.name == "nt":
"--expt-extended-lambda",
# find cl.exe
def find_cl_path():
import glob
for executable in ["Program Files (x86)", "Program Files"]:
for edition in ["Enterprise", "Professional", "BuildTools", "Community"]:
paths = sorted(glob.glob(f"C:\\{executable}\\Microsoft Visual Studio\\*\\{edition}\\VC\\Tools\\MSVC\\*\\bin\\Hostx64\\x64"), reverse=True)
if paths:
return paths[0]
# If cl.exe is not on path, try to find it.
if os.system("where cl.exe >nul 2>nul") != 0:
cl_path = find_cl_path()
if cl_path is None:
raise RuntimeError("Could not locate a supported Microsoft Visual C++ installation")
os.environ["PATH"] += ";" + cl_path
else:
# cl.exe was found in PATH, so we can assume that the user is already in a developer command prompt
# In this case, BuildExtensions requires the following environment variable to be set such that it
# won't try to activate a developer command prompt a second time.
os.environ["DISTUTILS_USE_SDK"] = "1"
cpp_standard = 14
# Get CUDA version and make sure the targeted compute capability is compatible
if os.system("nvcc --version") == 0:
nvcc_out = subprocess.check_output(["nvcc", "--version"]).decode()
cuda_version = re.search(r"release (\S+),", nvcc_out)
if cuda_version:
cuda_version = parse_version(cuda_version.group(1))
print(f"Detected CUDA version {cuda_version}")
if cuda_version >= parse_version("11.0"):
cpp_standard = 17
print(f"Targeting C++ standard {cpp_standard}")
base_nvcc_flags = [
f"-std=c++{cpp_standard}",
"--extended-lambda",
"--expt-relaxed-constexpr", "--expt-relaxed-constexpr",
'-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__', # The following definitions must be undefined
# since TCNN requires half-precision operation.
"-U__CUDA_NO_HALF_OPERATORS__",
"-U__CUDA_NO_HALF_CONVERSIONS__",
"-U__CUDA_NO_HALF2_OPERATORS__",
] ]
if os.name == "posix": if os.name == "posix":
c_flags = ['-O3', '-std=c++14'] base_cflags = [f"-std=c++{cpp_standard}"]
base_nvcc_flags += [
"-Xcompiler=-Wno-float-conversion",
"-Xcompiler=-fno-strict-aliasing",
]
elif os.name == "nt": elif os.name == "nt":
c_flags = ['/O2', '/std:c++17'] base_cflags = [f"/std:c++{cpp_standard}"]
# find cl.exe
def find_cl_path():
import glob
for edition in ["Enterprise", "Professional", "BuildTools", "Community"]:
paths = sorted(glob.glob(r"C:\\Program Files (x86)\\Microsoft Visual Studio\\*\\%s\\VC\\Tools\\MSVC\\*\\bin\\Hostx64\\x64" % edition), reverse=True)
if paths:
return paths[0]
# If cl.exe is not on path, try to find it.
if os.system("where cl.exe >nul 2>nul") != 0:
cl_path = find_cl_path()
if cl_path is None:
raise RuntimeError("Could not locate a supported Microsoft Visual C++ installation")
os.environ["PATH"] += ";" + cl_path
''' '''
Usage: Usage:
...@@ -116,41 +152,42 @@ python setup.py develop # build extensions and install (symbolic) to PATH. ...@@ -116,41 +152,42 @@ python setup.py develop # build extensions and install (symbolic) to PATH.
pip install -e . # ditto but better (e.g., dependency & metadata handling) pip install -e . # ditto but better (e.g., dependency & metadata handling)
''' '''
setup( setup(
name='cubvh', # package name, import this to use python API name='cubvh', # package name, import this to use python API
version='0.1.0', version='0.1.0',
description='CUDA BVH implementation', description='CUDA BVH implementation',
url='https://github.com/ashawkey/cubvh', url='https://github.com/ashawkey/cubvh',
author='kiui', author='kiui',
author_email='ashawkey1999@gmail.com', author_email='ashawkey1999@gmail.com',
packages=['cubvh'], packages=['cubvh'],
ext_modules=[ ext_modules=[
CUDAExtension( CUDAExtension(
name='_cubvh', # extension name, import this to use CUDA API name='_cubvh', # extension name, import this to use CUDA API
sources=[os.path.join(_src_path, 'src', f) for f in [ sources=[os.path.join(_src_path, 'src', f) for f in [
'bvh.cu', 'bvh.cu',
'api.cu', 'api.cu',
'bindings.cpp', 'bindings.cpp',
]], ]],
include_dirs=[ include_dirs=[
os.path.join(_src_path, 'include'), os.path.join(_src_path, 'include'),
find_eigen(), find_eigen(),
], ],
extra_compile_args={ extra_compile_args={
'cxx': c_flags, 'cxx': base_cflags,
'nvcc': nvcc_flags, 'nvcc': base_nvcc_flags,
} }
), ),
], ],
cmdclass={ cmdclass={
'build_ext': BuildExtension, 'build_ext': BuildExtension,
}, },
install_requires=[ install_requires=[
'ninja', 'ninja',
'trimesh', 'pybind11',
'opencv-python', 'trimesh',
'torch', 'opencv-python',
'numpy', 'torch',
'tqdm', 'numpy',
'matplotlib', 'tqdm',
], 'matplotlib',
],
) )
src/bindings.cpp
View file @ 7d436274
#include <pybind11/pybind11.h> // #include <pybind11/pybind11.h>
#include <pybind11/numpy.h> // #include <pybind11/numpy.h>
#include <pybind11/eigen.h> // #include <pybind11/eigen.h>
#include <torch/extension.h> #include <torch/extension.h>
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment