"""
Build script for flux_ext PyTorch C++/CUDA extension.

Usage:
    py -3.10 setup_ext.py install

This compiles the CUTLASS FP8 kernels (from flux_opt.cu) into a proper
PyTorch extension, eliminating all ctypes overhead.
"""

import os

# Use CUDA 12.8 (compatible with PyTorch 2.5.1 which was built with CUDA 12.1)
os.environ["CUDA_HOME"] = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8"
os.environ["CUDA_PATH"] = os.environ["CUDA_HOME"]

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

ROOT = os.path.dirname(os.path.abspath(__file__))

setup(
    name="flux_ext",
    ext_modules=[
        CUDAExtension(
            name="flux_ext",
            sources=[
                os.path.join(ROOT, "csrc", "FLUXPATCH.cpp"),
                os.path.join(ROOT, "csrc", "FLUXNATION.cu"),
                os.path.join(ROOT, "flux_opt.cu"),
            ],
            include_dirs=[
                ROOT,                                          # SINGLE_STREAM_BLOCK.H
                os.path.join(ROOT, "cutlass", "include"),      # CUTLASS headers
            ],
            extra_compile_args={
                "cxx": ["/O2", "/std:c++17", "/EHsc", "/MD"],
                "nvcc": [
                    "-O3",
                    "-std=c++17",
                    "-arch=sm_89",
                    "-DFLUX_TORCH_EXT=1",
                    "--allow-unsupported-compiler",
                    "-diag-suppress=20013",
                    "-diag-suppress=20015",
                    "-Xcompiler", "/MD",
                    "-Xcompiler", "/EHsc",
                    "-Xcompiler", "/Zc:preprocessor",
                ],
            },
            libraries=["cublas"],
        )
    ],
    cmdclass={"build_ext": BuildExtension},
)