diff --git a/setup-win.py b/setup-win.py new file mode 100644 index 0000000..fcf1921 --- /dev/null +++ b/setup-win.py @@ -0,0 +1,159 @@ +import os +import sys +import platform +import subprocess +from pathlib import Path + +import av +from setuptools import Extension, find_packages, setup +from setuptools.command.build_ext import build_ext as _build_ext +from Cython.Build import cythonize + + +# ----------------------------------------------------------------------------- +# Windows-specific configuration +# ----------------------------------------------------------------------------- + +IS_WINDOWS = platform.system() == "Windows" + +# On Windows, CUDA is often found via CUDA_PATH instead of CUDA_HOME +CUDA_HOME = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH", None) +CUDA_ARCH_LIST = os.environ.get("CUDA_ARCH_LIST", "75,86") + +# Manually specify where your FFmpeg (headers & libs) is installed on Windows. +# Adjust these paths to your local FFmpeg installation. For instance, if you +# installed FFmpeg with vcpkg, set these accordingly. +FFMPEG_INCLUDE_DIR = r"C:\Users\Administrator\Downloads\Programs\ffmpeg-n6.1-latest-win64-gpl-shared-6.1\include" +FFMPEG_LIB_DIR = r"C:\Users\Administrator\Downloads\Programs\ffmpeg-n6.1-latest-win64-gpl-shared-6.1\lib" + +# Which FFmpeg libraries you want to link against (the .lib import libs) +FFMPEG_LIBRARIES = [ + "avcodec", + "avutil", + # Add "avformat", "swscale", etc. if you need them +] + +# ----------------------------------------------------------------------------- +# get_include_dirs / library logic +# ----------------------------------------------------------------------------- + +def get_extension_config_windows(): + """ + On Windows, we won't use pkg-config. Instead, manually specify + include_dirs, library_dirs, and libraries for CUDA and FFmpeg. + """ + include_dirs = [] + library_dirs = [] + libraries = [] + + # Add FFmpeg includes/libs + include_dirs.append(FFMPEG_INCLUDE_DIR) + library_dirs.append(FFMPEG_LIB_DIR) + libraries.extend(FFMPEG_LIBRARIES) + + # Add CUDA includes/libs if CUDA is found + if CUDA_HOME: + include_dirs.append(str(Path(CUDA_HOME) / "include")) + # Typically on Windows x64: /lib/x64 + library_dirs.append(str(Path(CUDA_HOME) / "lib" / "x64")) + # Example: link the needed libraries. Adjust as needed (e.g. nppicc, cudart, etc.). + libraries.append("nppicc") + libraries.append("cudart") + libraries.append("nppif") + libraries.append("nppc") + libraries.append("nppig") + libraries.append("nppim") + libraries.append("nppist") + libraries.append("nppisu") + + return include_dirs, library_dirs, libraries + + +class CustomBuildExt(_build_ext): + """ + Custom build_ext class to handle .cu files via NVCC on Windows. + """ + def build_extensions(self): + if not CUDA_HOME: + raise ValueError("Couldn't find CUDA. Please set `CUDA_HOME` or `CUDA_PATH` on Windows.") + nvcc_path = str(Path(CUDA_HOME) / "bin" / "nvcc.exe") + + # Add .cu source extension + self.compiler.src_extensions.append(".cu") + default_compile = self.compiler._compile + + def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts): + # If we're building .cu files, switch to NVCC + if Path(src).suffix == ".cu": + self.compiler.set_executable("compiler_so", nvcc_path) + self.compiler.set_executable("compiler_cxx", nvcc_path) + self.compiler.set_executable("compiler", nvcc_path) + postargs = extra_postargs["nvcc"] + else: + # Otherwise, MSVC + postargs = extra_postargs["msvc"] + default_compile(obj, src, ext, cc_args, postargs, pp_opts) + + self.compiler._compile = _compile + super().build_extensions() + + +# ----------------------------------------------------------------------------- +# Gather sources & compile args +# ----------------------------------------------------------------------------- + +cuda_filepaths = [str(path) for path in Path("avcuda/cuda").glob("**/*.cu")] +cuda_arch_list = CUDA_ARCH_LIST.split(",") +cuda_arch_flags = [] +for arch in cuda_arch_list: + cuda_arch_flags.extend(["-gencode", f"arch=compute_{arch},code=sm_{arch}"]) + +# For safety, add a fallback arch. You can remove or adjust. +# (Compute_86 is Ampere, for example.) +cuda_arch_flags.extend(["-gencode", "arch=compute_86,code=compute_86"]) + +include_dirs, library_dirs, libraries = get_extension_config_windows() + +# Example compiler args for NVCC on Windows (MSVC) +extra_compile_args = { + "msvc": ["/std:c++17"], # MSVC flags + "nvcc": [ + "-c", + "--compiler-options=/std:c++17", # compile as C++17 + "--ptxas-options=-v", + *cuda_arch_flags + ] +} + +# ----------------------------------------------------------------------------- +# Build extension modules +# ----------------------------------------------------------------------------- + +extensions = [] +for pyx in Path("avcuda").glob("**/*.pyx"): + module_name = ( + str(pyx.with_suffix("")) + .replace("/", ".") + .replace("\\", ".") # Windows path fix + ) + extensions.extend(cythonize( + Extension( + module_name, + sources=[str(pyx)] + cuda_filepaths, + include_dirs=["avcuda", av.get_include()] + include_dirs, + libraries=libraries, + library_dirs=library_dirs, + # No runtime_library_dirs on Windows + extra_compile_args=extra_compile_args, + ), + language_level=3, + build_dir="build", + include_path=[av.get_include()], + )) + +setup( + name="avcuda", + packages=find_packages(exclude=["build*"]), + ext_modules=extensions, + cmdclass={"build_ext": CustomBuildExt}, +)