File size: 3,711 Bytes
38dbec8 |
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 |
import glob
import os
import platform
import torch
from setuptools import find_packages, setup
from torch.utils.cpp_extension import (
CUDA_HOME,
BuildExtension,
CppExtension,
CUDAExtension,
)
library_name = "texture_baker"
def get_extensions():
debug_mode = os.getenv("DEBUG", "0") == "1"
use_cuda = os.getenv("USE_CUDA", "1" if torch.cuda.is_available() else "0") == "1"
use_metal = (
os.getenv("USE_METAL", "1" if torch.backends.mps.is_available() else "0") == "1"
)
use_native_arch = os.getenv("USE_NATIVE_ARCH", "1") == "1"
if debug_mode:
print("Compiling in debug mode")
use_cuda = use_cuda and CUDA_HOME is not None
extension = CUDAExtension if use_cuda else CppExtension
extra_link_args = []
extra_compile_args = {
"cxx": [
"-O3" if not debug_mode else "-O0",
"-fdiagnostics-color=always",
"-fopenmp",
]
+ ["-march=native"]
if use_native_arch
else [],
"nvcc": [
"-O3" if not debug_mode else "-O0",
],
}
if debug_mode:
extra_compile_args["cxx"].append("-g")
if platform.system() == "Windows":
extra_compile_args["cxx"].append("/Z7")
extra_compile_args["cxx"].append("/Od")
extra_link_args.extend(["/DEBUG"])
extra_compile_args["cxx"].append("-UNDEBUG")
extra_compile_args["nvcc"].append("-UNDEBUG")
extra_compile_args["nvcc"].append("-g")
extra_link_args.extend(["-O0", "-g"])
define_macros = []
extensions = []
libraries = []
this_dir = os.path.dirname(os.path.curdir)
sources = glob.glob(
os.path.join(this_dir, library_name, "csrc", "**", "*.cpp"), recursive=True
)
if len(sources) == 0:
print("No source files found for extension, skipping extension compilation")
return None
if use_cuda:
define_macros += [
("THRUST_IGNORE_CUB_VERSION_CHECK", None),
]
sources += glob.glob(
os.path.join(this_dir, library_name, "csrc", "**", "*.cu"), recursive=True
)
libraries += ["cudart", "c10_cuda"]
if use_metal:
define_macros += [
("WITH_MPS", None),
]
sources += glob.glob(
os.path.join(this_dir, library_name, "csrc", "**", "*.mm"), recursive=True
)
extra_compile_args.update({"cxx": ["-O3", "-arch", "arm64"]})
extra_link_args += ["-arch", "arm64"]
extensions.append(
extension(
name=f"{library_name}._C",
sources=sources,
define_macros=define_macros,
extra_compile_args=extra_compile_args,
extra_link_args=extra_link_args,
libraries=libraries
+ [
"c10",
"torch",
"torch_cpu",
"torch_python",
],
)
)
for ext in extensions:
ext.libraries = ["cudart_static" if x == "cudart" else x for x in ext.libraries]
print(extensions)
return extensions
setup(
name=library_name,
version="0.0.1",
packages=find_packages(where="."),
package_dir={"": "."},
ext_modules=get_extensions(),
install_requires=[],
package_data={
library_name: [os.path.join("csrc", "*.h"), os.path.join("csrc", "*.metal")],
},
description="Small texture baker which rasterizes barycentric coordinates to a tensor.",
long_description=open("README.md").read(),
long_description_content_type="text/markdown",
url="https://github.com/Stability-AI/texture_baker",
cmdclass={"build_ext": BuildExtension},
)
|