setup.py

from pathlib import Path
from setuptools import setup

import torch
from torch.utils.cpp_extension import BuildExtension, CppExtension
import os

ROOT = os.path.dirname(os.path.abspath(__file__))
has_cuda = torch.cuda.is_available()

# Backend paths
backend_dir = None
if os.path.exists(os.path.join(ROOT, "splatt3r_slam/backend")):
    backend_dir = "splatt3r_slam"

if backend_dir is None:
    print("Warning: No backend directory found in splatt3r_slam")
    ext_modules = []
else:
    include_dirs = [
        os.path.join(ROOT, f"{backend_dir}/backend/include"),
        os.path.join(ROOT, "thirdparty/eigen"),
    ]

    sources = [
        f"{backend_dir}/backend/src/gn.cpp",
    ]
    extra_compile_args = {
        "cxx": ["-O3"],
    }

    if has_cuda:
        from torch.utils.cpp_extension import CUDAExtension

        sources.append(f"{backend_dir}/backend/src/gn_kernels.cu")
        sources.append(f"{backend_dir}/backend/src/matching_kernels.cu")
        # CUDA 13 dropped sm_<75 (Maxwell/Pascal/Volta); this list covers
        # Turing through Hopper on CUDA 13.x toolchains. Developed with
        # nvcc 13.3 on dual RTX A6000 (sm_86).
        extra_compile_args["nvcc"] = [
            "-O3",
            "-gencode=arch=compute_75,code=sm_75",
            "-gencode=arch=compute_80,code=sm_80",
            "-gencode=arch=compute_86,code=sm_86",
            "-gencode=arch=compute_89,code=sm_89",
            "-gencode=arch=compute_90,code=sm_90",
        ]
        ext_modules = [
            CUDAExtension(
                "mast3r_slam_backends",
                include_dirs=include_dirs,
                sources=sources,
                extra_compile_args=extra_compile_args,
            )
        ]
    else:
        print("CUDA not found, cannot compile backend!")
        ext_modules = []

setup(
    ext_modules=ext_modules,
    cmdclass={"build_ext": BuildExtension} if ext_modules else {},
)