citydreamer/extensions/extrude_tensor/setup.py

# -*- coding: utf-8 -*-
#
# @File:   setup.py
# @Author: Haozhe Xie
# @Date:   2023-03-24 20:35:43
# @Last Modified by: Haozhe Xie
# @Last Modified at: 2024-09-22 11:05:34
# @Email:  root@haozhexie.com

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

cxx_args = []
nvcc_args = ["-O3", "--gpu-architecture=compute_86", "--gpu-code=sm_86,compute_86"]

setup(
    name="extrude_tensor",
    version="1.0.0",
    ext_modules=[
        CUDAExtension(
            "extrude_tensor_ext",
            sources=[
                "bindings.cpp",
                "extrude_tensor_ext.cu",
            ],
            extra_compile_args={"cxx": cxx_args, "nvcc": nvcc_args},
        ),
    ],
    cmdclass={"build_ext": BuildExtension},
)