neural_codec/DCVC/src/layers/extensions/inference/setup.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import os
import glob
import sys
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension


cxx_flags = ["-O3"]
nvcc_flags = ["-O3", "--use_fast_math", "--extra-device-vectorization"]
# `-arch=native` detects the building machine's GPU at compile time, which is
# convenient interactively but (a) fails when no GPU is visible during the build
# and (b) can conflict with the `-gencode` flags torch appends from
# TORCH_CUDA_ARCH_LIST. So only use `native` when the caller has NOT pinned the
# arch via TORCH_CUDA_ARCH_LIST (e.g. cluster builds set it to "8.0").
if not os.environ.get("TORCH_CUDA_ARCH_LIST"):
    nvcc_flags.append("-arch=native")
if sys.platform == 'win32':
    cxx_flags = ["/O2"]


setup(
    name='inference_extensions_cuda',
    ext_modules=[
        CUDAExtension(
            name='inference_extensions_cuda',
            sources=glob.glob('*.cpp') + glob.glob('*.cu'),
            extra_compile_args={
                "cxx": cxx_flags,
                "nvcc": nvcc_flags,
            },
        ),
    ],
    cmdclass={
        'build_ext': BuildExtension
    }
)