Precompile renderer plugin during installation only (for faster start…

…-up time) (#43)

.gitignore

@@ -30,3 +30,10 @@ assets/*
 *.pdf
 *.pkl
 .DS_Store
+
+# C++ extensions
+*.so
+
+# Distribution / packaging
+build/
+dist/

bayes3d/__init__.py

@@ -2,6 +2,7 @@
 .. include:: ./documentation.md
 """
 
+from . import colmap, distributions, scene_graph, utils
 from .camera import *
 from .likelihood import *
 from .renderer import *
@@ -10,3 +11,5 @@
 from .viz import *
 
 RENDERER = None
+
+__all__ = ["colmap", "distributions", "scene_graph", "utils"]

bayes3d/rendering/nvdiffrast/common/__init__.py

@@ -6,6 +6,6 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-from .ops import RasterizeGLContext, _get_plugin, get_log_level, set_log_level
+from .ops import RasterizeGLContext, _get_plugin
 
-__all__ = ["RasterizeGLContext", "get_log_level", "set_log_level", "_get_plugin"]
+__all__ = ["RasterizeGLContext", "_get_plugin"]

bayes3d/rendering/nvdiffrast/common/ops.py

@@ -6,177 +6,19 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import importlib
-import logging
-import os
 
 import torch
-import torch.utils.cpp_extension
 
-# ----------------------------------------------------------------------------
-# C++/Cuda plugin compiler/loader.
-
-_cached_plugin = {}
-
-
-def _get_plugin(gl=False):
-    assert isinstance(gl, bool)
-
-    # Return cached plugin if already loaded.
-    if _cached_plugin.get(gl, None) is not None:
-        return _cached_plugin[gl]
-
-    # Make sure we can find the necessary compiler and libary binaries.
-    if os.name == "nt":
-        lib_dir = os.path.dirname(__file__) + r"\..\lib"
-
-        def find_cl_path():
-            import glob
-
-            for edition in ["Enterprise", "Professional", "BuildTools", "Community"]:
-                vs_relative_path = (
-                    r"\Microsoft Visual Studio\*\%s\VC\Tools\MSVC\*\bin\Hostx64\x64"
-                    % edition
-                )
-                paths = sorted(
-                    glob.glob(r"C:\Program Files" + vs_relative_path), reverse=True
-                )
-                paths += sorted(
-                    glob.glob(r"C:\Program Files (x86)" + vs_relative_path),
-                    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
-
-    # Compiler options.
-    opts = ["-DNVDR_TORCH"]
-
-    # Linker options for the GL-interfacing plugin.
-    ldflags = []
-    if gl:
-        if os.name == "posix":
-            ldflags = ["-lGL", "-lEGL"]
-        elif os.name == "nt":
-            libs = ["gdi32", "opengl32", "user32", "setgpu"]
-            ldflags = ["/LIBPATH:" + lib_dir] + ["/DEFAULTLIB:" + x for x in libs]
-
-    # List of source files.
-    if gl:
-        source_files = [
-            "common.cpp",
-            "glutil.cpp",
-            "rasterize_gl.cpp",
-        ]
-    else:
-        source_files = [
-            "../common/common.cpp",
-            "../common/rasterize.cu",
-            "../common/interpolate.cu",
-            "../common/texture.cu",
-            "../common/texture.cpp",
-            "../common/antialias.cu",
-            "torch_bindings.cpp",
-            "torch_rasterize.cpp",
-            "torch_interpolate.cpp",
-            "torch_texture.cpp",
-            "torch_antialias.cpp",
-        ]
-
-    # Some containers set this to contain old architectures that won't compile. We only need the one installed in the machine.
-    os.environ["TORCH_CUDA_ARCH_LIST"] = ""
-
-    # On Linux, show a warning if GLEW is being forcibly loaded when compiling the GL plugin.
-    if gl and (os.name == "posix") and ("libGLEW" in os.environ.get("LD_PRELOAD", "")):
-        logging.getLogger("nvdiffrast").warning(
-            "Warning: libGLEW is being loaded via LD_PRELOAD, and will probably conflict with the OpenGL plugin"
-        )
+import bayes3d.rendering.nvdiffrast.nvdiffrast_plugin_gl as plugin_gl
 
-    # Try to detect if a stray lock file is left in cache directory and show a warning. This sometimes happens on Windows if the build is interrupted at just the right moment.
-    plugin_name = "nvdiffrast_plugin" + ("_gl" if gl else "")
-    try:
-        lock_fn = os.path.join(
-            torch.utils.cpp_extension._get_build_directory(plugin_name, False), "lock"
-        )
-        if os.path.exists(lock_fn):
-            logging.getLogger("nvdiffrast").warning(
-                "Lock file exists in build directory: '%s'" % lock_fn
-            )
-    except Exception:
-        pass
-
-    # Speed up compilation on Windows.
-    if os.name == "nt":
-        # Skip telemetry sending step in vcvarsall.bat
-        os.environ["VSCMD_SKIP_SENDTELEMETRY"] = "1"
-
-        # Opportunistically patch distutils to cache MSVC environments.
-        try:
-            import distutils._msvccompiler
-            import functools
-
-            if not hasattr(distutils._msvccompiler._get_vc_env, "__wrapped__"):
-                distutils._msvccompiler._get_vc_env = functools.lru_cache()(
-                    distutils._msvccompiler._get_vc_env
-                )
-        except Exception:
-            pass
-
-    # Compile and load.
-    source_paths = [os.path.join(os.path.dirname(__file__), fn) for fn in source_files]
-    torch.utils.cpp_extension.load(
-        name=plugin_name,
-        sources=source_paths,
-        extra_cflags=opts,
-        extra_cuda_cflags=opts + ["-lineinfo"],
-        extra_ldflags=ldflags,
-        with_cuda=True,
-        verbose=False,
-    )
-
-    # Import, cache, and return the compiled module.
-    _cached_plugin[gl] = importlib.import_module(plugin_name)
-    return _cached_plugin[gl]
-
-
-# ----------------------------------------------------------------------------
-# Log level.
 # ----------------------------------------------------------------------------
+# C++/Cuda plugin loader.
 
 
-def get_log_level():
-    """Get current log level.
-
-    Returns:
-      Current log level in nvdiffrast. See `set_log_level()` for possible values.
-    """
-    return _get_plugin().get_log_level()
-
-
-def set_log_level(level):
-    """Set log level.
-
-    Log levels follow the convention on the C++ side of Torch:
-      0 = Info,
-      1 = Warning,
-      2 = Error,
-      3 = Fatal.
-    The default log level is 1.
-
-    Args:
-      level: New log level as integer. Internal nvdiffrast messages of this
-             severity or higher will be printed, while messages of lower
-             severity will be silent.
-    """
-    _get_plugin().set_log_level(level)
+def _get_plugin(gl=True):
+    # the gl flag is left here for backward compatibility
+    assert gl is True
+    return plugin_gl
 
 
 # ----------------------------------------------------------------------------

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=64", "setuptools_scm>=8"]
+requires = ["setuptools>=64", "setuptools_scm>=8", "torch>=2.0.0"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -24,7 +24,8 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12"
 ]
-dynamic = ["version"]
+dynamic = ["version", "dependencies", "optional-dependencies"]
+
 
 [tool.ruff]
 exclude = [
@@ -70,3 +71,15 @@ quote-style = "double"
 indent-style = "space"
 skip-magic-trailing-comma = false
 line-ending = "auto"
+
+
+[tool.setuptools_scm]
+
+
+[tool.setuptools.packages.find]
+include = ["bayes3d"]
+namespaces = false
+
+[tool.setuptools.dynamic]
+dependencies = {file = ["requirements.txt"]}
+optional-dependencies = {dev = { file = ["requirements-dev.txt"] }}

setup.py

@@ -1,10 +1,97 @@
+import os
+import warnings
+
 import setuptools
+from torch.utils import cpp_extension
+
+CPP_SRC_DIR = "bayes3d/rendering/nvdiffrast/common"
+
+
+# Nvdiffrast compilation setup
+if os.name == "nt":
+    lib_dir = os.path.dirname(__file__) + r"\..\lib"
+
+    def find_cl_path():
+        import glob
+
+        for edition in ["Enterprise", "Professional", "BuildTools", "Community"]:
+            vs_relative_path = (
+                r"\Microsoft Visual Studio\*\%s\VC\Tools\MSVC\*\bin\Hostx64\x64"
+                % edition
+            )
+            paths = sorted(
+                glob.glob(r"C:\Program Files" + vs_relative_path), reverse=True
+            )
+            paths += sorted(
+                glob.glob(r"C:\Program Files (x86)" + vs_relative_path), 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
+
+# Compiler options.
+opts = ["-DNVDR_TORCH"]
+
+# Linker options for the GL-interfacing plugin.
+ldflags = []
+if os.name == "posix":
+    ldflags = ["-lGL", "-lEGL"]
+elif os.name == "nt":
+    libs = ["gdi32", "opengl32", "user32", "setgpu"]
+    ldflags = ["/LIBPATH:" + lib_dir] + ["/DEFAULTLIB:" + x for x in libs]
 
-NAME = "bayes3d"
-VERSION = "0.0.1"
-if __name__ == "__main__":
-    setuptools.setup(
-        name=NAME,
-        version=VERSION,
-        packages=setuptools.find_namespace_packages(include=["bayes3d.*"]),
+# List of source files.
+source_files = [
+    "common.cpp",
+    "glutil.cpp",
+    "rasterize_gl.cpp",
+]
+source_files = [os.path.join(CPP_SRC_DIR, fn) for fn in source_files]
+
+# Some containers set this to contain old architectures that won't compile. We only need the one installed in the machine.
+os.environ["TORCH_CUDA_ARCH_LIST"] = ""
+
+# On Linux, show a warning if GLEW is being forcibly loaded when compiling the GL plugin.
+if (os.name == "posix") and ("libGLEW" in os.environ.get("LD_PRELOAD", "")):
+    warnings.warn(
+        "libGLEW is being loaded via LD_PRELOAD, and will probably conflict with the OpenGL plugin"
     )
+
+
+# Speed up compilation on Windows.
+if os.name == "nt":
+    # Skip telemetry sending step in vcvarsall.bat
+    os.environ["VSCMD_SKIP_SENDTELEMETRY"] = "1"
+
+    # Opportunistically patch distutils to cache MSVC environments.
+    try:
+        import distutils._msvccompiler
+        import functools
+
+        if not hasattr(distutils._msvccompiler._get_vc_env, "__wrapped__"):
+            distutils._msvccompiler._get_vc_env = functools.lru_cache()(
+                distutils._msvccompiler._get_vc_env
+            )
+    except Exception:
+        pass
+
+
+setuptools.setup(
+    ext_modules=[
+        cpp_extension.CUDAExtension(
+            name="bayes3d.rendering.nvdiffrast.nvdiffrast_plugin_gl",
+            sources=source_files,
+            extra_compile_args={"cxx": opts, "nvcc": opts + ["-lineinfo"]},
+            extra_link_args=ldflags,
+        ),
+    ],
+    cmdclass={"build_ext": cpp_extension.BuildExtension},
+)